Add MCP hook override functionality in Assistant tests and implementation

- Introduced a new test case for MCP hook overrides in build_mcp_test.go, validating that hooks can successfully modify MCP server configurations.
- Updated the applyMCPTools and buildMCPTools methods to prioritize hook-provided MCP servers over assistant configurations, enhancing flexibility in tool management.
- Enhanced the HookCreateResponse struct to include MCPServers, allowing hooks to specify server configurations for requests.
- Improved error handling and logging throughout the MCP tool application process, ensuring better traceability and debugging capabilities.

agent/assistant/build.go

@@ -92,7 +92,7 @@ func (ast *Assistant) buildCompletionOptions(ctx *context.Context, createRespons
 	}
 
 	// Add MCP tools if configured and get samples prompt
-	mcpSamplesPrompt, err := ast.applyMCPTools(ctx, options)
+	mcpSamplesPrompt, err := ast.applyMCPTools(ctx, options, createResponse)
 	if err != nil {
 		return nil, "", fmt.Errorf("failed to apply MCP tools: %w", err)
 	}
@@ -361,14 +361,26 @@ func (ast *Assistant) getUses() *context.Uses {
 
 // applyMCPTools adds MCP tools to completion options and returns samples prompt
 // Returns (samplesPrompt, error)
-func (ast *Assistant) applyMCPTools(ctx *context.Context, options *context.CompletionOptions) (string, error) {
+func (ast *Assistant) applyMCPTools(ctx *context.Context, options *context.CompletionOptions, createResponse *context.HookCreateResponse) (string, error) {
 
-	if ast.MCP == nil || len(ast.MCP.Servers) == 0 {
-		return "", nil
+	// Priority 1: Check if hook provides MCP servers
+	if createResponse != nil && len(createResponse.MCPServers) > 0 {
+		return ast.buildAndApplyMCPTools(ctx, options, createResponse)
 	}
 
+	// Priority 2: Check if assistant has MCP config
+	if ast.MCP != nil && len(ast.MCP.Servers) > 0 {
+		return ast.buildAndApplyMCPTools(ctx, options, nil)
+	}
+
+	// No MCP config
+	return "", nil
+}
+
+// buildAndApplyMCPTools builds MCP tools and applies them to options
+func (ast *Assistant) buildAndApplyMCPTools(ctx *context.Context, options *context.CompletionOptions, createResponse *context.HookCreateResponse) (string, error) {
 	// Build MCP tools and get samples prompt
-	mcpTools, samplesPrompt, err := ast.buildMCPTools(ctx)
+	mcpTools, samplesPrompt, err := ast.buildMCPTools(ctx, createResponse)
 	if err != nil {
 		return "", fmt.Errorf("failed to build MCP tools: %w", err)
 	}

agent/assistant/build_mcp_test.go

@@ -200,6 +200,88 @@ func TestBuildRequest_MCP(t *testing.T) {
 			t.Logf("✓ Tool schema valid: %v", fn["name"])
 		}
 	})
+
+	t.Run("MCPHookOverride", func(t *testing.T) {
+		// Test that hook can override MCP servers
+		// Use tests.mcptest-hook which has a create hook that returns only ["ping"]
+		hookAgent, err := assistant.Get("tests.mcptest-hook")
+		if err != nil {
+			t.Fatalf("Failed to get tests.mcptest-hook assistant: %s", err.Error())
+		}
+
+		hookCtx := newTestContext("chat-test-mcp-hook", "tests.mcptest-hook")
+		inputMessages := []context.Message{{Role: context.RoleUser, Content: "test hook override"}}
+
+		// Call create hook to get createResponse
+		var createResponse *context.HookCreateResponse
+		if hookAgent.Script != nil {
+			createResponse, err = hookAgent.Script.Create(hookCtx, inputMessages)
+			if err != nil {
+				t.Fatalf("Failed to call create hook: %s", err.Error())
+			}
+
+			t.Logf("Create hook response: %+v", createResponse)
+			if createResponse != nil && len(createResponse.MCPServers) > 0 {
+				t.Logf("Hook MCP servers: %+v", createResponse.MCPServers)
+			}
+		} else {
+			t.Fatal("Expected hookAgent to have Script/hook configured")
+		}
+
+		// Build LLM request with create hook response
+		_, options, err := hookAgent.BuildRequest(hookCtx, inputMessages, createResponse)
+		if err != nil {
+			t.Fatalf("Failed to build LLM request: %s", err.Error())
+		}
+
+		// Verify that tools are loaded
+		if options.Tools == nil {
+			t.Fatal("Expected tools to be loaded, got nil")
+		}
+
+		// Count MCP tools
+		mcpToolCount := 0
+		var toolNames []string
+		for _, toolMap := range options.Tools {
+			fn, ok := toolMap["function"].(map[string]interface{})
+			if !ok {
+				continue
+			}
+			name, ok := fn["name"].(string)
+			if ok {
+				toolNames = append(toolNames, name)
+				mcpToolCount++
+			}
+		}
+
+		t.Logf("Found %d MCP tools after hook override: %v", mcpToolCount, toolNames)
+
+		// Verify tool count (hook should override to only 1: ping)
+		if mcpToolCount != 1 {
+			t.Errorf("Expected 1 MCP tool (ping only), got %d: %v", mcpToolCount, toolNames)
+		}
+
+		// Verify only ping tool exists
+		hasEchoPing := false
+		hasEchoEcho := false
+		for _, name := range toolNames {
+			if name == "echo__ping" {
+				hasEchoPing = true
+			}
+			if name == "echo__echo" {
+				hasEchoEcho = true
+			}
+		}
+
+		if !hasEchoPing {
+			t.Error("Expected 'echo__ping' tool to be present")
+		}
+		if hasEchoEcho {
+			t.Error("Tool 'echo__echo' should be filtered out by hook override but was found")
+		}
+
+		t.Log("✓ Hook successfully overrode MCP servers configuration")
+	})
 }
 
 // Helper function to check if string contains substring

agent/assistant/mcp.go

@@ -11,6 +11,7 @@ import (
 	mcpTypes "github.com/yaoapp/gou/mcp/types"
 	"github.com/yaoapp/kun/log"
 	agentContext "github.com/yaoapp/yao/agent/context"
+	storeTypes "github.com/yaoapp/yao/agent/store/types"
 	"github.com/yaoapp/yao/trace/types"
 )
 
@@ -72,8 +73,26 @@ func ParseMCPToolName(formattedName string) (string, string, bool) {
 
 // buildMCPTools builds tool definitions and samples system prompt from MCP servers
 // Returns (tools, samplesPrompt, error)
-func (ast *Assistant) buildMCPTools(ctx *agentContext.Context) ([]MCPTool, string, error) {
-	if ast.MCP == nil || len(ast.MCP.Servers) == 0 {
+func (ast *Assistant) buildMCPTools(ctx *agentContext.Context, createResponse *agentContext.HookCreateResponse) ([]MCPTool, string, error) {
+	// Determine which MCP servers to use: hook's or assistant's (hook takes precedence)
+	var servers []storeTypes.MCPServerConfig
+
+	// If hook provides MCP servers, use those (override)
+	if createResponse != nil && len(createResponse.MCPServers) > 0 {
+		servers = make([]storeTypes.MCPServerConfig, len(createResponse.MCPServers))
+		for i, hookServer := range createResponse.MCPServers {
+			// Convert context.MCPServerConfig to storeTypes.MCPServerConfig
+			servers[i] = storeTypes.MCPServerConfig{
+				ServerID:  hookServer.ServerID,
+				Tools:     hookServer.Tools,
+				Resources: hookServer.Resources,
+			}
+		}
+	} else if ast.MCP != nil && len(ast.MCP.Servers) > 0 {
+		// Otherwise, use assistant's configured servers
+		servers = ast.MCP.Servers
+	} else {
+		// No servers configured
 		return nil, "", nil
 	}
 
@@ -88,7 +107,7 @@ func (ast *Assistant) buildMCPTools(ctx *agentContext.Context) ([]MCPTool, strin
 	hasSamples := false
 
 	// Process each MCP server in order
-	for _, serverConfig := range ast.MCP.Servers {
+	for _, serverConfig := range servers {
 		if len(allTools) >= MaxMCPTools {
 			log.Warn("[Assistant MCP] Reached maximum tool limit (%d), skipping remaining servers", MaxMCPTools)
 			break

agent/context/types.go

@@ -299,6 +299,9 @@ type HookCreateResponse struct {
 	MaxTokens           *int     `json:"max_tokens,omitempty"`
 	MaxCompletionTokens *int     `json:"max_completion_tokens,omitempty"`
 
+	// MCP configuration - allow hook to add/override MCP servers for this request
+	MCPServers []MCPServerConfig `json:"mcp_servers,omitempty"`
+
 	// Context adjustments - allow hook to modify context fields
 	AssistantID string                 `json:"assistant_id,omitempty"` // Override assistant ID
 	Connector   string                 `json:"connector,omitempty"`    // Override connector
@@ -467,3 +470,11 @@ type AudioConfig struct {
 type StreamOptions struct {
 	IncludeUsage bool `json:"include_usage,omitempty"` // If true, include usage statistics in the final chunk
 }
+
+// MCPServerConfig represents an MCP server configuration
+// This mirrors agent/store/types.MCPServerConfig to avoid import cycles
+type MCPServerConfig struct {
+	ServerID  string   `json:"server_id"`           // MCP server ID (required)
+	Tools     []string `json:"tools,omitempty"`     // Tool name filter (empty = all tools)
+	Resources []string `json:"resources,omitempty"` // Resource URI filter (empty = all resources)
+}