From 9acf6b64e139972deffda35beb0e8b2c498e89a3 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 29 Nov 2025 20:45:39 +0800 Subject: [PATCH] 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 | 22 +++++++-- agent/assistant/build_mcp_test.go | 82 +++++++++++++++++++++++++++++++ agent/assistant/mcp.go | 25 ++++++++-- agent/context/types.go | 11 +++++ 4 files changed, 132 insertions(+), 8 deletions(-) diff --git a/agent/assistant/build.go b/agent/assistant/build.go index 375c50d4..3a7fd700 100644 --- a/agent/assistant/build.go +++ b/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) } diff --git a/agent/assistant/build_mcp_test.go b/agent/assistant/build_mcp_test.go index 3927a089..aa55b581 100644 --- a/agent/assistant/build_mcp_test.go +++ b/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 diff --git a/agent/assistant/mcp.go b/agent/assistant/mcp.go index 00f5e246..e941fef7 100644 --- a/agent/assistant/mcp.go +++ b/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 diff --git a/agent/context/types.go b/agent/context/types.go index d1f4a554..8cdefe61 100644 --- a/agent/context/types.go +++ b/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) +}