From a66c2c76660e40fd05e01d31902d242a5f877ed7 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 5 Dec 2025 20:09:38 +0800 Subject: [PATCH] Implement context passing for MCP tool calls - Added tests to verify that agent context is correctly passed to MCP tools during single and parallel calls. - Updated the MCP client methods to accept agent context as an additional argument, enhancing the context management during tool execution. - Improved test coverage for context handling, ensuring that context data is accurately received and validated in tool responses. --- agent/assistant/mcp.go | 24 +++--- agent/assistant/mcp_test.go | 143 ++++++++++++++++++++++++++++++++++++ 2 files changed, 157 insertions(+), 10 deletions(-) diff --git a/agent/assistant/mcp.go b/agent/assistant/mcp.go index e941fef7..53c5df85 100644 --- a/agent/assistant/mcp.go +++ b/agent/assistant/mcp.go @@ -372,10 +372,12 @@ func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall } } - // Call the tool + // Call the tool with agent context as extra argument log.Trace("[Assistant MCP] Calling tool: %s (server: %s)", toolName, serverID) fmt.Printf(">>> executeSingleToolCall: CALLING client.CallTool (tool: %s, server: %s)\n", toolName, serverID) - callResult, err := client.CallTool(mcpCtx, toolName, args) + + // Pass agent context as extra argument (only used for Process transport) + callResult, err := client.CallTool(mcpCtx, toolName, args, ctx) fmt.Printf(">>> executeSingleToolCall: client.CallTool RETURNED (err: %v)\n", err) if err != nil { result.Error = fmt.Errorf("tool call failed: %w", err) @@ -468,14 +470,14 @@ func (ast *Assistant) executeMultipleToolCallsParallel(ctx *agentContext.Context // Try parallel execution serverResults, serverHasErrors := ast.executeServerToolsParallelWithTrace( - mcpCtx, trace, client, serverID, calls, + mcpCtx, ctx, trace, client, serverID, calls, ) // If parallel execution failed with retryable error, try sequential if serverHasErrors && ast.shouldRetrySequential(serverResults) { log.Warn("[Assistant MCP] Parallel execution had parameter errors for server '%s', retrying sequentially", serverID) serverResults, serverHasErrors = ast.executeServerToolsSequentialWithTrace( - mcpCtx, trace, client, serverID, calls, + mcpCtx, ctx, trace, client, serverID, calls, ) } @@ -558,7 +560,7 @@ func (ast *Assistant) shouldRetrySequential(results []ToolCallResult) bool { } // executeServerToolsParallelWithTrace executes tools for a single server in parallel with trace -func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context, trace types.Manager, client mcp.Client, serverID string, toolCalls []agentContext.ToolCall) ([]ToolCallResult, bool) { +func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context, ctx *agentContext.Context, trace types.Manager, client mcp.Client, serverID string, toolCalls []agentContext.ToolCall) ([]ToolCallResult, bool) { // Prepare parallel trace inputs var parallelInputs []types.TraceParallelInput mcpCalls := make([]mcpTypes.ToolCall, 0, len(toolCalls)) @@ -616,9 +618,11 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context fmt.Printf(">>> executeServerToolsParallelWithTrace: NOT creating trace nodes (trace: %v, inputs: %d)\n", trace != nil, len(parallelInputs)) } - // Call tools in parallel + // Call tools in parallel with agent context as extra argument log.Trace("[Assistant MCP] Calling %d tools in parallel on server '%s'", len(mcpCalls), serverID) - mcpResponse, err := client.CallToolsParallel(mcpCtx, mcpCalls) + + // Pass agent context as extra argument (only used for Process transport) + mcpResponse, err := client.CallToolsParallel(mcpCtx, mcpCalls, ctx) if err != nil { log.Error("[Assistant MCP] Parallel call failed: %v", err) // Mark all trace nodes as failed @@ -690,7 +694,7 @@ func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context } // executeServerToolsSequentialWithTrace executes tools for a single server sequentially with trace -func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Context, trace types.Manager, client mcp.Client, serverID string, toolCalls []agentContext.ToolCall) ([]ToolCallResult, bool) { +func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Context, ctx *agentContext.Context, trace types.Manager, client mcp.Client, serverID string, toolCalls []agentContext.ToolCall) ([]ToolCallResult, bool) { results := make([]ToolCallResult, 0, len(toolCalls)) hasErrors := false @@ -800,9 +804,9 @@ func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Conte } } - // Call single tool + // Call single tool with agent context as extra argument log.Trace("[Assistant MCP] Calling tool: %s", toolName) - mcpResult, err := client.CallTool(mcpCtx, toolName, args) + mcpResult, err := client.CallTool(mcpCtx, toolName, args, ctx) result := ToolCallResult{ ToolCallID: tc.ID, diff --git a/agent/assistant/mcp_test.go b/agent/assistant/mcp_test.go index 30edf3b1..2b2339cd 100644 --- a/agent/assistant/mcp_test.go +++ b/agent/assistant/mcp_test.go @@ -1,10 +1,17 @@ package assistant_test import ( + "context" "testing" + jsoniter "github.com/json-iterator/go" + "github.com/stretchr/testify/assert" + "github.com/yaoapp/gou/mcp" + mcpTypes "github.com/yaoapp/gou/mcp/types" "github.com/yaoapp/yao/agent/assistant" + agentContext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/testutils" + "github.com/yaoapp/yao/openapi/oauth/types" ) func TestMCPToolName(t *testing.T) { @@ -234,3 +241,139 @@ func TestMCPToolName_RoundTrip(t *testing.T) { }) } } + +// TestMCPToolContextPassing tests that agent context is correctly passed to MCP tools +func TestMCPToolContextPassing(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + // Get the echo MCP client + client, err := mcp.Select("echo") + assert.NoError(t, err, "Failed to select echo MCP client") + assert.NotNil(t, client, "MCP client should not be nil") + + // Create a test agent context + authorized := &types.AuthorizedInfo{ + UserID: "test-user-123", + TenantID: "test-tenant-456", + } + ctx := agentContext.New(context.Background(), authorized, "test-chat-789") + ctx.AssistantID = "test-assistant-mcptest" + ctx.Locale = "en" + ctx.Theme = "dark" + + // Call the echo tool with context + args := map[string]interface{}{ + "message": "test message from context test", + } + + // Call the tool - the agent context will be passed as extra parameter + result, err := client.CallTool(ctx.Context, "echo", args, ctx) + assert.NoError(t, err, "CallTool should not return error") + assert.NotNil(t, result, "Result should not be nil") + assert.False(t, result.IsError, "Result should not be an error") + assert.Greater(t, len(result.Content), 0, "Result should have content") + + // Parse the result content + var echoResult map[string]interface{} + err = jsoniter.Unmarshal([]byte(result.Content[0].Text), &echoResult) + assert.NoError(t, err, "Failed to parse result content") + + t.Logf("Echo result: %+v", echoResult) + + // Verify the context was received + contextData, ok := echoResult["context"].(map[string]interface{}) + assert.True(t, ok, "Result should contain context field") + assert.NotNil(t, contextData, "Context data should not be nil") + + // Verify context has_context flag + hasContext, ok := contextData["has_context"].(bool) + assert.True(t, ok, "Context should have has_context field") + assert.True(t, hasContext, "Context should indicate it has context") + + // Verify chat_id and assistant_id have values (main verification) + chatID, ok := contextData["chat_id"].(string) + assert.True(t, ok, "Context should have chat_id field") + assert.NotEmpty(t, chatID, "chat_id should have a value") + assert.Equal(t, "test-chat-789", chatID, "chat_id should match") + + assistantID, ok := contextData["assistant_id"].(string) + assert.True(t, ok, "Context should have assistant_id field") + assert.NotEmpty(t, assistantID, "assistant_id should have a value") + assert.Equal(t, "test-assistant-mcptest", assistantID, "assistant_id should match") + + t.Logf("✓ Context successfully passed to MCP tool") + t.Logf(" - ChatID: %s", chatID) + t.Logf(" - AssistantID: %s", assistantID) +} + +// TestMCPToolContextPassingParallel tests that agent context is correctly passed in parallel calls +func TestMCPToolContextPassingParallel(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + // Get the echo MCP client + client, err := mcp.Select("echo") + assert.NoError(t, err, "Failed to select echo MCP client") + assert.NotNil(t, client, "MCP client should not be nil") + + // Create a test agent context + authorized := &types.AuthorizedInfo{ + UserID: "parallel-user-123", + TenantID: "parallel-tenant-456", + } + ctx := agentContext.New(context.Background(), authorized, "parallel-chat-789") + ctx.AssistantID = "test-assistant-parallel" + ctx.Locale = "zh-CN" + + // Call multiple echo tools in parallel + toolCalls := []mcpTypes.ToolCall{ + { + Name: "echo", + Arguments: map[string]interface{}{ + "message": "parallel message 1", + }, + }, + { + Name: "echo", + Arguments: map[string]interface{}{ + "message": "parallel message 2", + }, + }, + } + + // Call tools in parallel - the agent context will be passed as extra parameter + results, err := client.CallToolsParallel(ctx.Context, toolCalls, ctx) + assert.NoError(t, err, "CallToolsParallel should not return error") + assert.NotNil(t, results, "Results should not be nil") + assert.Equal(t, 2, len(results.Results), "Should have 2 results") + + // Verify both results received the context + for i, result := range results.Results { + assert.False(t, result.IsError, "Result %d should not be an error", i) + assert.Greater(t, len(result.Content), 0, "Result %d should have content", i) + + // Parse the result content + var echoResult map[string]interface{} + err = jsoniter.Unmarshal([]byte(result.Content[0].Text), &echoResult) + assert.NoError(t, err, "Failed to parse result %d content", i) + + // Verify the context was received + contextData, ok := echoResult["context"].(map[string]interface{}) + assert.True(t, ok, "Result %d should contain context field", i) + assert.NotNil(t, contextData, "Context data %d should not be nil", i) + + hasContext, ok := contextData["has_context"].(bool) + assert.True(t, ok, "Context %d should have has_context field", i) + assert.True(t, hasContext, "Context %d should indicate it has context", i) + + // Verify chat_id in parallel call + chatID, ok := contextData["chat_id"].(string) + assert.True(t, ok, "Context %d should have chat_id field", i) + assert.Equal(t, "parallel-chat-789", chatID, "Chat ID in result %d should match", i) + + t.Logf("✓ Result %d successfully received context", i) + } + + t.Log("✓ Context successfully passed to all parallel MCP tool calls") +}