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.
This commit is contained in:
Max
parent
f0e445862f
commit
a66c2c7666
2 changed files with 157 additions and 10 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue