diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index 0b8bc776..bbb315bb 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -7,12 +7,11 @@ import ( jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/connector" "github.com/yaoapp/kun/log" + "github.com/yaoapp/kun/utils" "github.com/yaoapp/yao/agent/assistant/handlers" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" - "github.com/yaoapp/yao/agent/llm" "github.com/yaoapp/yao/agent/output/message" - "github.com/yaoapp/yao/trace/types" ) // Stream stream the agent @@ -34,174 +33,198 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa }) } + // ================================================ + // Initialize + // ================================================ + // Initialize stack and auto-handle completion/failure/restore - _, traceID, done := context.EnterStack(ctx, ast.ID, ctx.Referer) + _, _, done := context.EnterStack(ctx, ast.ID, ctx.Referer) defer done() - _ = traceID // traceID is available for trace logging - - // Get connector and capabilities early (before sending stream_start) - // so that output adapters can use them when converting stream_start event - if ast.Prompts != nil || ast.MCP != nil { - _, capabilities, err := ast.GetConnector(ctx) - if err != nil { - streamHandler := ast.getStreamHandler(ctx, handler...) - ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) - return nil, err - } - - // Set capabilities in context for output adapters to use - if capabilities != nil { - ctx.Capabilities = capabilities - } - } - // Determine stream handler streamHandler := ast.getStreamHandler(ctx, handler...) + // Get connector and capabilities early (before sending stream_start) + // so that output adapters can use them when converting stream_start event + err = ast.initializeCapabilities(ctx) + if err != nil { + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) + return nil, err + } + // Send ChunkStreamStart only for root stack (agent-level stream start) // Now ctx.Capabilities is set, so output adapters can use it ast.sendAgentStreamStart(ctx, streamHandler, streamStartTime) - // Trace Add - trace, _ := ctx.Trace() - var agentNode types.Node = nil - if trace != nil { - agentNode, _ = trace.Add(inputMessages, types.TraceNodeOption{ - Label: i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.label"), // "Assistant {{name}}" - Type: "agent", - Icon: "assistant", - Description: i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.description"), // "Assistant {{name}} is processing the request" - }) - } + // Initialize agent trace node + agentNode := ast.initAgentTraceNode(ctx, inputMessages) - // Full input messages with chat history - fullMessages, err := ast.WithHistory(ctx, inputMessages) + // ================================================ + // Get Full Messages with chat history + // ================================================ + fullMessages, err := ast.WithHistory(ctx, inputMessages, agentNode) if err != nil { - if agentNode != nil { - agentNode.Fail(err) - } - // Send error stream_end for root stack + ast.traceAgentFail(agentNode, err) ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } - // Log the chat history - if agentNode != nil { - agentNode.Info(i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.history"), map[string]any{"messages": fullMessages}) // "Get Chat History" - } - + // ================================================ + // Execute Create Hook + // ================================================ // Request Create hook ( Optional ) var createResponse *context.HookCreateResponse if ast.Script != nil { var err error createResponse, err = ast.Script.Create(ctx, fullMessages) if err != nil { - if agentNode != nil { - agentNode.Fail(err) - } + ast.traceAgentFail(agentNode, err) // Send error stream_end for root stack ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } // Log the create response - if agentNode != nil { - agentNode.Debug("Call Create Hook", map[string]any{"response": createResponse}) - } + ast.traceCreateHook(agentNode, createResponse) } - var completionOptions *context.CompletionOptions // default is nil - + // ================================================ + // Execute LLM Call Stream + // ================================================ // LLM Call Stream ( Optional ) - var completionMessages []context.Message var completionResponse *context.CompletionResponse + var completionMessages []context.Message + var completionOptions *context.CompletionOptions if ast.Prompts != nil || ast.MCP != nil { // Build the LLM request first completionMessages, completionOptions, err = ast.BuildRequest(ctx, inputMessages, createResponse) if err != nil { - if agentNode != nil { - agentNode.Fail(err) - } + ast.traceAgentFail(agentNode, err) // Send error stream_end for root stack ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } - // Get connector object (capabilities were already set above, before stream_start) - conn, capabilities, err := ast.GetConnector(ctx) + // Execute the LLM streaming call + completionResponse, err = ast.executeLLMStream(ctx, completionMessages, completionOptions, agentNode, streamHandler) if err != nil { - if agentNode != nil { - agentNode.Fail(err) - } + ast.traceAgentFail(agentNode, err) // Send error stream_end for root stack ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } - // Set capabilities in options if not already set - if completionOptions.Capabilities == nil && capabilities != nil { - completionOptions.Capabilities = capabilities - } - - // Log the capabilities - if agentNode != nil { - agentNode.Debug("Get Connector Capabilities", map[string]any{"capabilities": capabilities}) - } - - // Trace Add - if trace != nil { - trace.Add( - map[string]any{"messages": completionMessages, "options": completionOptions}, - types.TraceNodeOption{ - Label: fmt.Sprintf(i18n.Tr(ast.ID, ctx.Locale, "llm.openai.stream.label"), conn.ID()), // "LLM %s" - Type: "llm", - Icon: "psychology", - Description: fmt.Sprintf(i18n.Tr(ast.ID, ctx.Locale, "llm.openai.stream.description"), conn.ID()), // "LLM %s is processing the request" - }, - ) - } - - // Create LLM instance with connector and options - llmInstance, err := llm.New(conn, completionOptions) - if err != nil { - // Send error stream_end for root stack - ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) - return nil, err - } - - // Call the LLM Completion Stream (streamHandler was set earlier) - log.Trace("[AGENT] Calling LLM Stream: assistant=%s", ast.ID) - completionResponse, err = llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler) - log.Trace("[AGENT] LLM Stream returned: assistant=%s, err=%v", ast.ID, err) - if err != nil { - // Send error stream_end for root stack - log.Trace("[AGENT] Calling sendStreamEndOnError") - ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) - log.Trace("[AGENT] sendStreamEndOnError returned") - return nil, err - } - - // Mark LLM Request Complete - if trace != nil { - trace.Complete(completionResponse) + // === Debug Completion Response === + fmt.Println("--- Debug Completion Response ----------------------") + fmt.Printf("completionResponse: %+v\n", completionResponse) + if completionResponse != nil { + fmt.Printf("ToolCalls: %+v\n", completionResponse.ToolCalls) } + fmt.Println("----------------------------------------------------") + // === End Debug === } - // Request MCP hook ( Optional ) - var mcpResponse *context.ResponseHookMCP - if ast.MCP != nil { - _ = mcpResponse // mcpResponse is available for further processing + // ================================================ + // Execute tool calls with retry + // ================================================ + if completionResponse != nil && completionResponse.ToolCalls != nil { - // MCP Execution Loop + // === Debug Tool Calls === + fmt.Println("--- Debug Tool Calls --------------------------------") + utils.Dump(completionResponse.ToolCalls) + + // === End Debug Tool Calls === + + maxToolRetries := 3 + currentMessages := completionMessages + currentResponse := completionResponse + + for attempt := 0; attempt < maxToolRetries; attempt++ { + + fmt.Println("attempt", attempt) + // Execute all tool calls + toolResults, hasErrors := ast.executeToolCalls(ctx, currentResponse.ToolCalls, attempt) + // If all successful, break out + if !hasErrors { + log.Trace("[AGENT] All tool calls succeeded (attempt %d)", attempt) + for _, result := range toolResults { + fmt.Println("--") + fmt.Printf("Result :%s %s %s\n", result.ToolCallID, result.Server(), result.Tool()) + res, err := result.ParsedContent() + if err != nil { + fmt.Println("Error: ", err) + } + utils.Dump(res) + fmt.Println("--") + } + fmt.Println("--------------------------------") + break + } + + // Check if any errors are retryable (parameter/validation issues) + hasRetryableErrors := false + for _, result := range toolResults { + if result.Error != nil && result.IsRetryableError { + hasRetryableErrors = true + break + } + } + + // If no retryable errors, don't retry (MCP internal issues) + if !hasRetryableErrors { + err := fmt.Errorf("tool calls failed with non-retryable errors (MCP internal issues)") + log.Error("[AGENT] %v", err) + ast.traceAgentFail(agentNode, err) + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) + return nil, err + } + + // If it's the last attempt, return error + if attempt == maxToolRetries-1 { + err := fmt.Errorf("tool calls failed after %d attempts", maxToolRetries) + log.Error("[AGENT] %v", err) + ast.traceAgentFail(agentNode, err) + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) + return nil, err + } + + // Build retry messages with tool call results (including errors) + retryMessages := ast.buildToolRetryMessages(currentMessages, currentResponse, toolResults) + + // Retry LLM call (streaming to keep user informed) + log.Trace("[AGENT] Retrying LLM for tool call correction (attempt %d/%d)", attempt+1, maxToolRetries-1) + currentResponse, err = ast.executeLLMForToolRetry(ctx, retryMessages, completionOptions, agentNode, streamHandler) + if err != nil { + log.Error("[AGENT] LLM retry failed: %v", err) + ast.traceAgentFail(agentNode, err) + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) + return nil, err + } + + // If LLM didn't return tool calls, it might have given up + if currentResponse.ToolCalls == nil { + err := fmt.Errorf("LLM did not return tool calls in retry attempt %d", attempt+1) + log.Error("[AGENT] %v", err) + ast.traceAgentFail(agentNode, err) + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) + return nil, err + } + + // Update messages for next iteration + currentMessages = retryMessages + } + + // Update completionResponse with the final successful response + completionResponse = currentResponse } // Request Done hook ( Optional ) var doneResponse *context.ResponseHookDone if ast.Script != nil { var err error - doneResponse, err = ast.Script.Done(ctx, fullMessages, completionResponse, mcpResponse) + doneResponse, err = ast.Script.Done(ctx, fullMessages, completionResponse, nil) if err != nil { + ast.traceAgentFail(agentNode, err) // Send error stream_end for root stack ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err @@ -211,9 +234,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa _ = doneResponse // doneResponse is available for further processing // Set the output of the agent node - if agentNode != nil { - agentNode.SetOutput(context.Response{Create: createResponse, Done: doneResponse, Completion: completionResponse}) - } + ast.traceAgentOutput(agentNode, createResponse, doneResponse, completionResponse) // Only close output and send stream_end if this is the root call (entry point) // Nested calls (from MCP, hooks, etc.) should not close the output or send stream_end @@ -358,11 +379,6 @@ func (ast *Assistant) Info(locale ...string) *message.AssistantInfo { } } -// WithHistory with the history messages -func (ast *Assistant) WithHistory(ctx *context.Context, messages []context.Message) ([]context.Message, error) { - return messages, nil -} - // getStreamHandler returns the stream handler from the provided handlers or a default one func (ast *Assistant) getStreamHandler(ctx *context.Context, handler ...message.StreamFunc) message.StreamFunc { if len(handler) > 0 && handler[0] != nil { @@ -464,3 +480,72 @@ func (ast *Assistant) handleInterrupt(ctx *context.Context, signal *context.Inte return nil } + +// initializeCapabilities gets connector and capabilities, then sets them in context +// This should be called early (before sending stream_start) so that output adapters +// can use capabilities when converting stream_start event +func (ast *Assistant) initializeCapabilities(ctx *context.Context) error { + if ast.Prompts == nil && ast.MCP == nil { + return nil + } + + _, capabilities, err := ast.GetConnector(ctx) + if err != nil { + return err + } + + // Set capabilities in context for output adapters to use + if capabilities != nil { + ctx.Capabilities = capabilities + } + + return nil +} + +// buildToolRetryMessages builds messages for LLM retry with tool call results +// Format follows OpenAI's tool call response pattern: +// 1. Assistant message with tool calls +// 2. Tool messages with results (one per tool call) +// 3. System message explaining the retry +func (ast *Assistant) buildToolRetryMessages( + previousMessages []context.Message, + completionResponse *context.CompletionResponse, + toolResults []ToolCallResult, +) []context.Message { + retryMessages := make([]context.Message, 0, len(previousMessages)+len(toolResults)+2) + + // Add all previous messages + retryMessages = append(retryMessages, previousMessages...) + + // Add assistant message with tool calls + assistantMsg := context.Message{ + Role: context.RoleAssistant, + Content: completionResponse.Content, + ToolCalls: completionResponse.ToolCalls, + } + retryMessages = append(retryMessages, assistantMsg) + + // Add tool result messages (one per tool call) + for _, result := range toolResults { + toolMsg := context.Message{ + Role: context.RoleTool, + Content: result.Content, + ToolCallID: &result.ToolCallID, + } + // Add tool name if available + if result.Name != "" { + name := result.Name + toolMsg.Name = &name + } + retryMessages = append(retryMessages, toolMsg) + } + + // Add system message explaining the retry (optional, helps LLM understand context) + systemMsg := context.Message{ + Role: context.RoleSystem, + Content: i18n.Tr(ast.ID, "en", "assistant.agent.tool_retry_prompt"), + } + retryMessages = append(retryMessages, systemMsg) + + return retryMessages +} diff --git a/agent/assistant/assistant.go b/agent/assistant/assistant.go index 2053a818..b2a7b8ef 100644 --- a/agent/assistant/assistant.go +++ b/agent/assistant/assistant.go @@ -152,8 +152,22 @@ func (ast *Assistant) Clone() *Assistant { if ast.MCP != nil { clone.MCP = &store.MCPServers{} if ast.MCP.Servers != nil { - clone.MCP.Servers = make([]string, len(ast.MCP.Servers)) - copy(clone.MCP.Servers, ast.MCP.Servers) + clone.MCP.Servers = make([]store.MCPServerConfig, len(ast.MCP.Servers)) + for i, server := range ast.MCP.Servers { + clone.MCP.Servers[i] = store.MCPServerConfig{ + ServerID: server.ServerID, + } + // Deep copy Resources slice + if server.Resources != nil { + clone.MCP.Servers[i].Resources = make([]string, len(server.Resources)) + copy(clone.MCP.Servers[i].Resources, server.Resources) + } + // Deep copy Tools slice + if server.Tools != nil { + clone.MCP.Servers[i].Tools = make([]string, len(server.Tools)) + copy(clone.MCP.Servers[i].Tools, server.Tools) + } + } } if ast.MCP.Options != nil { clone.MCP.Options = make(map[string]interface{}) diff --git a/agent/assistant/build.go b/agent/assistant/build.go index 20939045..375c50d4 100644 --- a/agent/assistant/build.go +++ b/agent/assistant/build.go @@ -9,14 +9,14 @@ import ( // BuildRequest build the LLM request func (ast *Assistant) BuildRequest(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse) ([]context.Message, *context.CompletionOptions, error) { - // Build final messages with proper priority - finalMessages, err := ast.buildMessages(ctx, messages, createResponse) + // Build completion options from createResponse and ctx (includes MCP tools) + options, mcpSamplesPrompt, err := ast.buildCompletionOptions(ctx, createResponse) if err != nil { return nil, nil, err } - // Build completion options from createResponse and ctx - options, err := ast.buildCompletionOptions(ctx, createResponse) + // Build final messages with proper priority (includes MCP samples if available) + finalMessages, err := ast.buildMessages(ctx, messages, createResponse, mcpSamplesPrompt) if err != nil { return nil, nil, err } @@ -25,9 +25,9 @@ func (ast *Assistant) BuildRequest(ctx *context.Context, messages []context.Mess } // buildMessages builds the final message list with proper priority -// Priority: Prompts > createResponse.Messages > input messages +// Priority: Prompts > MCP Samples > createResponse.Messages > input messages // If createResponse is nil or has no messages, use input messages -func (ast *Assistant) buildMessages(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse) ([]context.Message, error) { +func (ast *Assistant) buildMessages(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse, mcpSamplesPrompt string) ([]context.Message, error) { var finalMessages []context.Message // If createResponse is nil or has no messages, use input messages @@ -38,6 +38,16 @@ func (ast *Assistant) buildMessages(ctx *context.Context, messages []context.Mes finalMessages = createResponse.Messages } + // Add MCP samples prompt as a system message (if available) + if mcpSamplesPrompt != "" { + mcpSamplesMsg := context.Message{ + Role: context.RoleSystem, + Content: mcpSamplesPrompt, + } + // Prepend MCP samples before other messages + finalMessages = append([]context.Message{mcpSamplesMsg}, finalMessages...) + } + // ⚠️ Just for testing, will remove later // If we have prompts, prepend them to the beginning if len(ast.Prompts) > 0 { @@ -64,12 +74,13 @@ func (ast *Assistant) buildMessages(ctx *context.Context, messages []context.Mes // buildCompletionOptions builds completion options from multiple sources // Priority (lowest to highest, later overrides earlier): ast > ctx > createResponse // The priority means: if createResponse has a value, use it; else use ctx; else use ast -func (ast *Assistant) buildCompletionOptions(ctx *context.Context, createResponse *context.HookCreateResponse) (*context.CompletionOptions, error) { +// Returns (options, mcpSamplesPrompt, error) +func (ast *Assistant) buildCompletionOptions(ctx *context.Context, createResponse *context.HookCreateResponse) (*context.CompletionOptions, string, error) { options := &context.CompletionOptions{} // Layer 1 (base): Apply ast - Assistant configuration if err := ast.applyAssistantOptions(options); err != nil { - return nil, err + return nil, "", err } // Layer 2 (middle): Apply ctx - Context configuration (overrides ast) @@ -80,7 +91,24 @@ func (ast *Assistant) buildCompletionOptions(ctx *context.Context, createRespons ast.applyCreateResponseOptions(options, createResponse) } - return options, nil + // Add MCP tools if configured and get samples prompt + mcpSamplesPrompt, err := ast.applyMCPTools(ctx, options) + if err != nil { + return nil, "", fmt.Errorf("failed to apply MCP tools: %w", err) + } + + // === Debug MCP Tools === + fmt.Println("--- Debug MCP Tools after applyMCPTools ---------------") + fmt.Printf("options.Tools count: %d\n", len(options.Tools)) + if len(options.Tools) > 0 { + for i, tool := range options.Tools { + fmt.Printf("Tool %d: %+v\n", i, tool) + } + } + fmt.Println("-------------------------------------------------------") + // === End Debug === + + return options, mcpSamplesPrompt, nil } // applyAssistantOptions applies options from ast.Options to CompletionOptions @@ -330,3 +358,42 @@ func (ast *Assistant) getUses() *context.Uses { // Priority 2: Global settings only return globalUses } + +// 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) { + + if ast.MCP == nil || len(ast.MCP.Servers) == 0 { + return "", nil + } + + // Build MCP tools and get samples prompt + mcpTools, samplesPrompt, err := ast.buildMCPTools(ctx) + if err != nil { + return "", fmt.Errorf("failed to build MCP tools: %w", err) + } + + // Convert mcpTools to map format for CompletionOptions.Tools + if len(mcpTools) > 0 { + toolMaps := make([]map[string]interface{}, len(mcpTools)) + for i, tool := range mcpTools { + toolMaps[i] = map[string]interface{}{ + "type": "function", + "function": map[string]interface{}{ + "name": tool.Name, + "description": tool.Description, + "parameters": tool.Parameters, + }, + } + } + + // Add MCP tools to existing tools (append to preserve existing tools) + if options.Tools == nil { + options.Tools = toolMaps + } else { + options.Tools = append(options.Tools, toolMaps...) + } + } + + return samplesPrompt, nil +} diff --git a/agent/assistant/build_mcp_test.go b/agent/assistant/build_mcp_test.go new file mode 100644 index 00000000..3927a089 --- /dev/null +++ b/agent/assistant/build_mcp_test.go @@ -0,0 +1,222 @@ +package assistant_test + +import ( + "testing" + + "github.com/yaoapp/yao/agent/assistant" + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/testutils" +) + +// TestBuildRequest_MCP tests MCP tool integration in BuildRequest +func TestBuildRequest_MCP(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + agent, err := assistant.Get("tests.mcptest") + if err != nil { + t.Fatalf("Failed to get tests.mcptest assistant: %s", err.Error()) + } + + ctx := newTestContext("chat-test-mcp", "tests.mcptest") + + t.Run("MCPToolsLoaded", func(t *testing.T) { + inputMessages := []context.Message{{Role: context.RoleUser, Content: "test mcp tools"}} + + // Build LLM request + _, options, err := agent.BuildRequest(ctx, inputMessages, nil) + 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") + } + + if len(options.Tools) == 0 { + t.Fatal("Expected at least some MCP tools, got empty list") + } + + // Count MCP tools (should be filtered to only ping and echo) + 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: %v", mcpToolCount, toolNames) + + // Verify tool count (should be exactly 2: ping and echo) + if mcpToolCount != 2 { + t.Errorf("Expected 2 MCP tools (ping, echo), got %d: %v", mcpToolCount, toolNames) + } + + // Verify specific tools exist + 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("Expected 'echo__echo' tool to be present") + } + + // Verify that 'status' tool is NOT included (filtered out) + for _, name := range toolNames { + if name == "echo__status" { + t.Error("Tool 'echo__status' should be filtered out but was found") + } + } + + t.Log("✓ MCP tools loaded and filtered correctly") + }) + + t.Run("MCPSamplesPrompt", func(t *testing.T) { + inputMessages := []context.Message{{Role: context.RoleUser, Content: "test mcp samples"}} + + // Build LLM request + finalMessages, _, err := agent.BuildRequest(ctx, inputMessages, nil) + if err != nil { + t.Fatalf("Failed to build LLM request: %s", err.Error()) + } + + // Check if messages contain MCP samples prompt + // The samples prompt should be added as a system message + hasMCPSamples := false + for _, msg := range finalMessages { + if msg.Role == context.RoleSystem { + if content, ok := msg.Content.(string); ok { + if len(content) > 50 && + (contains(content, "MCP Tool Usage Examples") || + contains(content, "echo.ping") || + contains(content, "echo.echo")) { + hasMCPSamples = true + t.Logf("Found MCP samples prompt (length: %d chars)", len(content)) + break + } + } + } + } + + // Note: samples may not exist for echo tools, so this is informational + if hasMCPSamples { + t.Log("✓ MCP samples prompt included in messages") + } else { + t.Log("ℹ No MCP samples prompt found (may not have sample files)") + } + }) + + t.Run("MCPToolNameFormat", func(t *testing.T) { + inputMessages := []context.Message{{Role: context.RoleUser, Content: "test tool format"}} + + // Build LLM request + _, options, err := agent.BuildRequest(ctx, inputMessages, nil) + if err != nil { + t.Fatalf("Failed to build LLM request: %s", err.Error()) + } + + // Verify tool name format: server_id.tool_name + for _, toolMap := range options.Tools { + fn, ok := toolMap["function"].(map[string]interface{}) + if !ok { + continue + } + name, ok := fn["name"].(string) + if ok { + // Parse tool name + serverID, toolName, ok := assistant.ParseMCPToolName(name) + if !ok { + t.Errorf("Tool name '%s' is not in correct format (server_id.tool_name)", name) + continue + } + + // Verify server ID + if serverID != "echo" { + t.Errorf("Expected server_id 'echo', got '%s' for tool '%s'", serverID, name) + } + + // Verify tool name is either ping or echo + if toolName != "ping" && toolName != "echo" { + t.Errorf("Expected tool name 'ping' or 'echo', got '%s'", toolName) + } + + t.Logf("✓ Tool name format correct: %s → (%s, %s)", name, serverID, toolName) + } + } + }) + + t.Run("MCPToolSchema", func(t *testing.T) { + inputMessages := []context.Message{{Role: context.RoleUser, Content: "test tool schema"}} + + // Build LLM request + _, options, err := agent.BuildRequest(ctx, inputMessages, nil) + if err != nil { + t.Fatalf("Failed to build LLM request: %s", err.Error()) + } + + // Verify tool schema structure + for _, toolMap := range options.Tools { + // Verify type field + if toolType, ok := toolMap["type"].(string); !ok || toolType != "function" { + t.Errorf("Expected tool type 'function', got: %v", toolMap["type"]) + } + + // Verify function field exists + fn, ok := toolMap["function"].(map[string]interface{}) + if !ok { + t.Error("Tool missing 'function' field or wrong type") + continue + } + + // Verify required fields + if _, hasName := fn["name"]; !hasName { + t.Error("Tool function missing 'name' field") + } + if _, hasDesc := fn["description"]; !hasDesc { + t.Error("Tool function missing 'description' field") + } + if _, hasParams := fn["parameters"]; !hasParams { + t.Error("Tool function missing 'parameters' field") + } + + t.Logf("✓ Tool schema valid: %v", fn["name"]) + } + }) +} + +// Helper function to check if string contains substring +func contains(s, substr string) bool { + return len(s) >= len(substr) && + (s == substr || + len(s) > len(substr) && + (s[:len(substr)] == substr || + s[len(s)-len(substr):] == substr || + findSubstring(s, substr))) +} + +func findSubstring(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} diff --git a/agent/assistant/history.go b/agent/assistant/history.go new file mode 100644 index 00000000..bbecb73c --- /dev/null +++ b/agent/assistant/history.go @@ -0,0 +1,31 @@ +package assistant + +import ( + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/trace/types" +) + +// WithHistory merges the input messages with chat history and traces it +// This method can be overridden or extended to implement actual history loading +func (ast *Assistant) WithHistory( + ctx *context.Context, + inputMessages []context.Message, + agentNode types.Node, +) ([]context.Message, error) { + + // TODO: Implement actual history loading logic here + // For now, just simulate a check and return the input messages as is + + // Simulate error check (this is where actual history loading would happen) + // if some_condition { + // ast.traceAgentFail(agentNode, err) + // return nil, err + // } + + fullMessages := inputMessages + + // Log the chat history + ast.traceAgentHistory(ctx, agentNode, fullMessages) + + return fullMessages, nil +} diff --git a/agent/assistant/hook/mcp.go b/agent/assistant/hook/mcp.go deleted file mode 100644 index 0e0ba6cc..00000000 --- a/agent/assistant/hook/mcp.go +++ /dev/null @@ -1,8 +0,0 @@ -package hook - -import "github.com/yaoapp/yao/agent/context" - -// MCP MCP hook -func (s *Script) MCP(ctx *context.Context, messages []context.Message) (*context.ResponseHookMCP, error) { - return &context.ResponseHookMCP{}, nil -} diff --git a/agent/assistant/llm.go b/agent/assistant/llm.go new file mode 100644 index 00000000..de58aac0 --- /dev/null +++ b/agent/assistant/llm.go @@ -0,0 +1,126 @@ +package assistant + +import ( + "fmt" + + "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/llm" + "github.com/yaoapp/yao/agent/output/message" + "github.com/yaoapp/yao/trace/types" +) + +// executeLLMStream executes the LLM streaming call with pre-built request +// Returns completionResponse and error +func (ast *Assistant) executeLLMStream( + ctx *context.Context, + completionMessages []context.Message, + completionOptions *context.CompletionOptions, + agentNode types.Node, + streamHandler message.StreamFunc, +) (*context.CompletionResponse, error) { + + // === Debug LLM Stream Start === + fmt.Println(">>> executeLLMStream: STARTING") + fmt.Printf(">>> Messages count: %d\n", len(completionMessages)) + fmt.Printf(">>> Tools count: %d\n", len(completionOptions.Tools)) + // === End Debug === + + // Get connector object (capabilities were already set above, before stream_start) + conn, capabilities, err := ast.GetConnector(ctx) + if err != nil { + ast.traceAgentFail(agentNode, err) + return nil, err + } + + // Set capabilities in options if not already set + if completionOptions.Capabilities == nil && capabilities != nil { + completionOptions.Capabilities = capabilities + } + + // Log the capabilities + ast.traceConnectorCapabilities(agentNode, capabilities) + + // Trace Add LLM request + ast.traceLLMRequest(ctx, conn.ID(), completionMessages, completionOptions) + + // Create LLM instance with connector and options + llmInstance, err := llm.New(conn, completionOptions) + if err != nil { + return nil, err + } + + // Call the LLM Completion Stream (streamHandler was set earlier) + log.Trace("[AGENT] Calling LLM Stream: assistant=%s", ast.ID) + + // === Debug LLM Stream Call === + fmt.Println(">>> executeLLMStream: CALLING llmInstance.Stream()") + // === End Debug === + + completionResponse, err := llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler) + + // === Debug LLM Stream Return === + fmt.Println(">>> executeLLMStream: llmInstance.Stream() RETURNED") + fmt.Printf(">>> err: %v\n", err) + if completionResponse != nil { + fmt.Printf(">>> ToolCalls: %d\n", len(completionResponse.ToolCalls)) + } + // === End Debug === + + log.Trace("[AGENT] LLM Stream returned: assistant=%s, err=%v", ast.ID, err) + if err != nil { + log.Trace("[AGENT] Calling sendStreamEndOnError") + return nil, err + } + + // Mark LLM Request Complete + ast.traceLLMComplete(ctx, completionResponse) + + return completionResponse, nil +} + +// executeLLMForToolRetry executes LLM call for tool retry with streaming output +// This is used when retrying tool calls - we still want to show LLM's response to users +// Returns completionResponse and error +func (ast *Assistant) executeLLMForToolRetry( + ctx *context.Context, + completionMessages []context.Message, + completionOptions *context.CompletionOptions, + agentNode types.Node, + streamHandler message.StreamFunc, +) (*context.CompletionResponse, error) { + + // Get connector object + conn, capabilities, err := ast.GetConnector(ctx) + if err != nil { + ast.traceAgentFail(agentNode, err) + return nil, err + } + + // Set capabilities in options if not already set + if completionOptions.Capabilities == nil && capabilities != nil { + completionOptions.Capabilities = capabilities + } + + // Trace Add LLM retry request + ast.traceLLMRetryRequest(ctx, conn.ID(), completionMessages, completionOptions) + + // Create LLM instance with connector and options + llmInstance, err := llm.New(conn, completionOptions) + if err != nil { + return nil, err + } + + // Call the LLM Completion Stream (still streaming for tool retry) + log.Trace("[AGENT] Calling LLM Stream for tool retry: assistant=%s", ast.ID) + completionResponse, err := llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler) + log.Trace("[AGENT] LLM tool retry stream returned: assistant=%s, err=%v", ast.ID, err) + if err != nil { + return nil, err + } + + // Mark LLM Request Complete + ast.traceLLMComplete(ctx, completionResponse) + + return completionResponse, nil +} diff --git a/agent/assistant/mcp.go b/agent/assistant/mcp.go new file mode 100644 index 00000000..00f5e246 --- /dev/null +++ b/agent/assistant/mcp.go @@ -0,0 +1,834 @@ +package assistant + +import ( + "context" + "fmt" + "strings" + + jsoniter "github.com/json-iterator/go" + gouJson "github.com/yaoapp/gou/json" + "github.com/yaoapp/gou/mcp" + mcpTypes "github.com/yaoapp/gou/mcp/types" + "github.com/yaoapp/kun/log" + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/trace/types" +) + +const ( + // MaxMCPTools maximum number of MCP tools to include (to avoid overwhelming the LLM) + MaxMCPTools = 20 +) + +// MCPToolName formats a tool name with MCP server prefix +// Format: server_id__tool_name (double underscore separator) +// Dots in server_id are replaced with single underscores +// Examples: +// - ("echo", "ping") → "echo__ping" +// - ("github.enterprise", "search") → "github_enterprise__search" +// +// Naming constraint: MCP server_id MUST NOT contain underscores (_) +// Only dots (.), letters, numbers, and hyphens (-) are allowed in server_id +func MCPToolName(serverID, toolName string) string { + if serverID == "" || toolName == "" { + return "" + } + // Replace dots with single underscores in server_id + cleanServerID := strings.ReplaceAll(serverID, ".", "_") + // Use double underscore as separator + return fmt.Sprintf("%s__%s", cleanServerID, toolName) +} + +// ParseMCPToolName parses a formatted MCP tool name into server ID and tool name +// Splits by double underscore (__), then restores dots in server_id +// Examples: +// - "echo__ping" → ("echo", "ping") +// - "github_enterprise__search" → ("github.enterprise", "search") +// +// Returns (serverID, toolName, true) if valid format, ("", "", false) otherwise +func ParseMCPToolName(formattedName string) (string, string, bool) { + if formattedName == "" { + return "", "", false + } + + // Split by double underscore + parts := strings.Split(formattedName, "__") + if len(parts) != 2 { + return "", "", false + } + + cleanServerID := parts[0] + toolName := parts[1] + + // Validate that both parts are non-empty + if cleanServerID == "" || toolName == "" { + return "", "", false + } + + // Restore dots in server_id (replace single underscores back to dots) + serverID := strings.ReplaceAll(cleanServerID, "_", ".") + + return serverID, toolName, true +} + +// 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 { + return nil, "", nil + } + + // Use the agent context for cancellation and timeout control + mcpCtx := ctx.Context + if mcpCtx == nil { + mcpCtx = context.Background() + } + + allTools := make([]MCPTool, 0) + samplesBuilder := strings.Builder{} + hasSamples := false + + // Process each MCP server in order + for _, serverConfig := range ast.MCP.Servers { + if len(allTools) >= MaxMCPTools { + log.Warn("[Assistant MCP] Reached maximum tool limit (%d), skipping remaining servers", MaxMCPTools) + break + } + + // Get MCP client + client, err := mcp.Select(serverConfig.ServerID) + if err != nil { + log.Warn("[Assistant MCP] Failed to select MCP client '%s': %v", serverConfig.ServerID, err) + continue + } + + // Get tools list (filter by serverConfig.Tools if specified) + toolsResponse, err := client.ListTools(mcpCtx, "") + if err != nil { + log.Warn("[Assistant MCP] Failed to list tools for '%s': %v", serverConfig.ServerID, err) + continue + } + + // Build tool filter map if specified + toolFilter := make(map[string]bool) + if len(serverConfig.Tools) > 0 { + for _, toolName := range serverConfig.Tools { + toolFilter[toolName] = true + } + } + + // Process each tool + for _, tool := range toolsResponse.Tools { + // Check tool limit + if len(allTools) >= MaxMCPTools { + break + } + + // Apply tool filter if specified + if len(toolFilter) > 0 && !toolFilter[tool.Name] { + continue + } + + // Format tool name with server prefix + formattedName := MCPToolName(serverConfig.ServerID, tool.Name) + + // Convert MCP tool to MCPTool format + mcpTool := MCPTool{ + Name: formattedName, + Description: tool.Description, + Parameters: tool.InputSchema, + } + + allTools = append(allTools, mcpTool) + + // Try to get samples for this tool + samples, err := client.ListSamples(mcpCtx, mcpTypes.SampleTool, tool.Name) + if err == nil && len(samples.Samples) > 0 { + if !hasSamples { + samplesBuilder.WriteString("\n\n## MCP Tool Usage Examples\n\n") + samplesBuilder.WriteString("The following examples demonstrate how to use MCP tools correctly:\n\n") + hasSamples = true + } + + samplesBuilder.WriteString(fmt.Sprintf("### %s\n\n", formattedName)) + if tool.Description != "" { + samplesBuilder.WriteString(fmt.Sprintf("**Description**: %s\n\n", tool.Description)) + } + + for i, sample := range samples.Samples { + if i >= 3 { // Limit to 3 examples per tool + break + } + + samplesBuilder.WriteString(fmt.Sprintf("**Example %d", i+1)) + if sample.Name != "" { + samplesBuilder.WriteString(fmt.Sprintf(" - %s", sample.Name)) + } + samplesBuilder.WriteString("**:\n") + + // Check metadata for description + if sample.Metadata != nil { + if desc, ok := sample.Metadata["description"].(string); ok && desc != "" { + samplesBuilder.WriteString(fmt.Sprintf("- Description: %s\n", desc)) + } + } + + if sample.Input != nil { + samplesBuilder.WriteString(fmt.Sprintf("- Input: `%v`\n", sample.Input)) + } + + if sample.Output != nil { + samplesBuilder.WriteString(fmt.Sprintf("- Output: `%v`\n", sample.Output)) + } + + samplesBuilder.WriteString("\n") + } + } + } + + log.Trace("[Assistant MCP] Loaded %d tools from server '%s'", len(toolsResponse.Tools), serverConfig.ServerID) + } + + samplesPrompt := "" + if hasSamples { + samplesPrompt = samplesBuilder.String() + } + + log.Trace("[Assistant MCP] Total MCP tools loaded: %d", len(allTools)) + return allTools, samplesPrompt, nil +} + +// ToolCallResult represents the result of a tool call execution +// executeToolCalls executes tool calls with intelligent strategy and trace logging: +// - Single tool: use CallTool, single trace node +// - Multiple tools: use CallToolsParallel with parallel trace nodes, fallback to sequential on certain errors +// Returns (results, hasErrors) +func (ast *Assistant) executeToolCalls(ctx *agentContext.Context, toolCalls []agentContext.ToolCall, attempt int) ([]ToolCallResult, bool) { + if len(toolCalls) == 0 { + return nil, false + } + + // === Debug === + fmt.Printf(">>> executeToolCalls: START (attempt %d, toolCalls count: %d)\n", attempt, len(toolCalls)) + // === End Debug === + + log.Trace("[Assistant MCP] Executing %d tool calls (attempt %d)", len(toolCalls), attempt) + + // Single tool call + if len(toolCalls) == 1 { + fmt.Println(">>> executeToolCalls: Calling executeSingleToolCall") + results, hasErrors := ast.executeSingleToolCall(ctx, toolCalls[0]) + fmt.Printf(">>> executeToolCalls: executeSingleToolCall RETURNED (hasErrors: %v)\n", hasErrors) + return results, hasErrors + } + + // Multiple tool calls - try parallel first + fmt.Println(">>> executeToolCalls: Calling executeMultipleToolCallsParallel") + results, hasErrors := ast.executeMultipleToolCallsParallel(ctx, toolCalls) + fmt.Printf(">>> executeToolCalls: executeMultipleToolCallsParallel RETURNED (hasErrors: %v)\n", hasErrors) + return results, hasErrors +} + +// executeSingleToolCall executes a single tool call with trace logging +func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall agentContext.ToolCall) ([]ToolCallResult, bool) { + // === Debug === + fmt.Printf(">>> executeSingleToolCall: START (tool: %s)\n", toolCall.Function.Name) + // === End Debug === + + trace, _ := ctx.Trace() + + // Use the agent context for cancellation and timeout control + mcpCtx := ctx.Context + if mcpCtx == nil { + mcpCtx = context.Background() + } + + result := ToolCallResult{ + ToolCallID: toolCall.ID, + Name: toolCall.Function.Name, + } + + // Parse tool name + fmt.Println(">>> executeSingleToolCall: Parsing tool name") + serverID, toolName, ok := ParseMCPToolName(toolCall.Function.Name) + if !ok { + result.Error = fmt.Errorf("invalid MCP tool name format: %s", toolCall.Function.Name) + result.Content = result.Error.Error() + log.Error("[Assistant MCP] %v", result.Error) + return []ToolCallResult{result}, true + } + + // Get MCP client + client, err := mcp.Select(serverID) + if err != nil { + result.Error = fmt.Errorf("failed to select MCP client '%s': %w", serverID, err) + result.Content = result.Error.Error() + result.IsRetryableError = false // MCP client selection error is not retryable + log.Error("[Assistant MCP] %v", result.Error) + return []ToolCallResult{result}, true + } + + // Get tool info for description and schema + toolsResponse, err := client.ListTools(mcpCtx, "") + var toolDescription string + var toolSchema interface{} + if err == nil { + for _, t := range toolsResponse.Tools { + if t.Name == toolName { + toolDescription = t.Description + toolSchema = t.InputSchema + break + } + } + } + if toolDescription == "" { + toolDescription = fmt.Sprintf("MCP tool '%s'", toolName) + } + + // Add trace node for this tool call + var toolNode types.Node + if trace != nil { + toolNode, _ = trace.Add( + map[string]any{ + "tool_call_id": toolCall.ID, + "server": serverID, + "tool": toolName, + "arguments": toolCall.Function.Arguments, + }, + types.TraceNodeOption{ + Label: toolDescription, + Type: "mcp_tool", + Icon: "build", + Description: fmt.Sprintf("Calling '%s' on server '%s'", toolName, serverID), + }, + ) + } + + // Parse arguments with repair support for better tolerance + var args map[string]interface{} + if toolCall.Function.Arguments != "" { + parsed, err := gouJson.Parse(toolCall.Function.Arguments) + if err != nil { + result.Error = fmt.Errorf("failed to parse arguments: %w", err) + result.Content = result.Error.Error() + result.IsRetryableError = true // Argument parsing error is retryable by LLM + log.Error("[Assistant MCP] %v", result.Error) + if toolNode != nil { + toolNode.Fail(result.Error) + } + return []ToolCallResult{result}, true + } + + // Convert to map + if argsMap, ok := parsed.(map[string]interface{}); ok { + args = argsMap + } else { + result.Error = fmt.Errorf("arguments must be an object, got %T", parsed) + result.Content = result.Error.Error() + result.IsRetryableError = true // Type error is retryable by LLM + log.Error("[Assistant MCP] %v", result.Error) + if toolNode != nil { + toolNode.Fail(result.Error) + } + return []ToolCallResult{result}, true + } + + // Validate arguments against tool schema if available + if toolSchema != nil { + fmt.Println(">>> executeSingleToolCall: Validating arguments against schema") + if err := gouJson.Validate(args, toolSchema); err != nil { + fmt.Printf(">>> executeSingleToolCall: Validation FAILED: %v\n", err) + result.Error = fmt.Errorf("argument validation failed: %w", err) + result.Content = result.Error.Error() + result.IsRetryableError = true // Validation error is retryable by LLM + log.Error("[Assistant MCP] %v", result.Error) + if toolNode != nil { + fmt.Println(">>> executeSingleToolCall: Failing toolNode due to validation error") + toolNode.Fail(result.Error) + fmt.Println(">>> executeSingleToolCall: toolNode.Fail() finished") + } + fmt.Println(">>> executeSingleToolCall: RETURNING with validation error") + return []ToolCallResult{result}, true + } + fmt.Println(">>> executeSingleToolCall: Validation PASSED") + } + } + + // Call the tool + 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) + fmt.Printf(">>> executeSingleToolCall: client.CallTool RETURNED (err: %v)\n", err) + if err != nil { + result.Error = fmt.Errorf("tool call failed: %w", err) + result.Content = result.Error.Error() + // Check if error is retryable (parameter/validation errors) + result.IsRetryableError = isRetryableToolError(err) + log.Error("[Assistant MCP] Tool call failed: %v (retryable: %v)", err, result.IsRetryableError) + if toolNode != nil { + toolNode.Fail(result.Error) + } + fmt.Println(">>> executeSingleToolCall: RETURNING with error") + return []ToolCallResult{result}, true + } + + // Check if result is an error + if callResult.IsError { + result.Error = fmt.Errorf("MCP tool error") + result.IsRetryableError = false // MCP internal error is not retryable + } + + // Serialize the Content field only ([]ToolContent) + contentBytes, err := jsoniter.Marshal(callResult.Content) + if err != nil { + result.Error = fmt.Errorf("failed to serialize result: %w", err) + result.Content = result.Error.Error() + result.IsRetryableError = false + log.Error("[Assistant MCP] %v", result.Error) + if toolNode != nil { + toolNode.Fail(result.Error) + } + return []ToolCallResult{result}, true + } + + result.Content = string(contentBytes) + log.Trace("[Assistant MCP] Tool call succeeded: %s", toolName) + + if toolNode != nil { + fmt.Println(">>> executeSingleToolCall: Completing toolNode") + toolNode.Complete(map[string]any{ + "result": callResult, + }) + fmt.Println(">>> executeSingleToolCall: toolNode.Complete() finished") + } + + fmt.Println(">>> executeSingleToolCall: RETURNING success") + return []ToolCallResult{result}, false +} + +// executeMultipleToolCallsParallel executes multiple tool calls in parallel with trace logging +func (ast *Assistant) executeMultipleToolCallsParallel(ctx *agentContext.Context, toolCalls []agentContext.ToolCall) ([]ToolCallResult, bool) { + trace, _ := ctx.Trace() + + // Use the agent context for cancellation and timeout control + mcpCtx := ctx.Context + if mcpCtx == nil { + mcpCtx = context.Background() + } + + // Group tool calls by server + serverGroups := make(map[string][]agentContext.ToolCall) + for _, tc := range toolCalls { + serverID, _, ok := ParseMCPToolName(tc.Function.Name) + if !ok { + log.Warn("[Assistant MCP] Invalid tool name format: %s", tc.Function.Name) + continue + } + serverGroups[serverID] = append(serverGroups[serverID], tc) + } + + results := make([]ToolCallResult, 0, len(toolCalls)) + hasErrors := false + + // Process each server's tools + for serverID, calls := range serverGroups { + client, err := mcp.Select(serverID) + if err != nil { + log.Error("[Assistant MCP] Failed to select MCP client '%s': %v", serverID, err) + // Add error results for all calls to this server + for _, tc := range calls { + results = append(results, ToolCallResult{ + ToolCallID: tc.ID, + Name: tc.Function.Name, + Content: fmt.Sprintf("Failed to select MCP client: %v", err), + Error: err, + }) + } + hasErrors = true + continue + } + + // Try parallel execution + serverResults, serverHasErrors := ast.executeServerToolsParallelWithTrace( + mcpCtx, 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, + ) + } + + results = append(results, serverResults...) + if serverHasErrors { + hasErrors = true + } + } + + return results, hasErrors +} + +// isRetryableToolError checks if an error is retryable by LLM (parameter/validation errors) +// Returns true for errors that LLM can potentially fix by adjusting parameters +// Returns false for MCP internal errors (network, auth, service unavailable, etc.) +func isRetryableToolError(err error) bool { + if err == nil { + return false + } + + errMsg := strings.ToLower(err.Error()) + + // These are NOT retryable (MCP internal issues) + nonRetryablePatterns := []string{ + "network", + "timeout", + "connection", + "unauthorized", + "forbidden", + "unavailable", + "failed to select", + "context canceled", + "context deadline", + "server error", + "internal error", + } + + for _, pattern := range nonRetryablePatterns { + if strings.Contains(errMsg, pattern) { + return false + } + } + + // These ARE retryable (parameter/validation issues LLM can fix) + retryablePatterns := []string{ + "invalid", + "required", + "missing", + "validation", + "schema", + "type", + "format", + "parse", + "argument", + "parameter", + } + + for _, pattern := range retryablePatterns { + if strings.Contains(errMsg, pattern) { + return true + } + } + + // Default: assume it's retryable unless proven otherwise + // This allows LLM to attempt fixes for unknown error types + return true +} + +// shouldRetrySequential checks if errors are retryable (parameter issues, not network/service issues) +func (ast *Assistant) shouldRetrySequential(results []ToolCallResult) bool { + // Check if any result has a retryable error + hasRetryable := false + for _, result := range results { + if result.Error != nil && result.IsRetryableError { + hasRetryable = true + break + } + } + return hasRetryable +} + +// 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) { + // Prepare parallel trace inputs + var parallelInputs []types.TraceParallelInput + mcpCalls := make([]mcpTypes.ToolCall, 0, len(toolCalls)) + callMap := make(map[string]agentContext.ToolCall) + + for _, tc := range toolCalls { + _, toolName, ok := ParseMCPToolName(tc.Function.Name) + if !ok { + continue + } + + var args map[string]interface{} + if tc.Function.Arguments != "" { + if err := jsoniter.UnmarshalFromString(tc.Function.Arguments, &args); err != nil { + log.Error("[Assistant MCP] Failed to parse arguments for %s: %v", toolName, err) + continue + } + } + + mcpCalls = append(mcpCalls, mcpTypes.ToolCall{ + Name: toolName, + Arguments: args, + }) + callMap[toolName] = tc + + // Add trace input for this tool + parallelInputs = append(parallelInputs, types.TraceParallelInput{ + Input: map[string]any{ + "tool_call_id": tc.ID, + "server": serverID, + "tool": toolName, + "arguments": tc.Function.Arguments, + }, + Option: types.TraceNodeOption{ + Label: fmt.Sprintf("Tool: %s", toolName), + Type: "mcp_tool", + Icon: "build", + Description: fmt.Sprintf("Calling MCP tool '%s' on server '%s'", toolName, serverID), + }, + }) + } + + // Create parallel trace nodes + var toolNodes []types.Node + if trace != nil && len(parallelInputs) > 0 { + fmt.Printf(">>> executeServerToolsParallelWithTrace: Creating %d parallel trace nodes\n", len(parallelInputs)) + var err error + toolNodes, err = trace.Parallel(parallelInputs) + if err != nil { + fmt.Printf(">>> executeServerToolsParallelWithTrace: trace.Parallel() FAILED: %v\n", err) + } else { + fmt.Printf(">>> executeServerToolsParallelWithTrace: Created %d trace nodes\n", len(toolNodes)) + } + } else { + fmt.Printf(">>> executeServerToolsParallelWithTrace: NOT creating trace nodes (trace: %v, inputs: %d)\n", trace != nil, len(parallelInputs)) + } + + // Call tools in parallel + log.Trace("[Assistant MCP] Calling %d tools in parallel on server '%s'", len(mcpCalls), serverID) + mcpResponse, err := client.CallToolsParallel(mcpCtx, mcpCalls) + if err != nil { + log.Error("[Assistant MCP] Parallel call failed: %v", err) + // Mark all trace nodes as failed + for _, node := range toolNodes { + if node != nil { + node.Fail(err) + } + } + return nil, true + } + + // Process results + results := make([]ToolCallResult, 0, len(mcpResponse.Results)) + hasErrors := false + + for i, mcpResult := range mcpResponse.Results { + toolName := mcpCalls[i].Name + originalCall := callMap[toolName] + var toolNode types.Node + if i < len(toolNodes) { + toolNode = toolNodes[i] + } + + result := ToolCallResult{ + ToolCallID: originalCall.ID, + Name: originalCall.Function.Name, + } + + // Serialize content + contentBytes, err := jsoniter.Marshal(mcpResult.Content) + if err != nil { + result.Error = fmt.Errorf("failed to serialize result: %w", err) + result.Content = result.Error.Error() + result.IsRetryableError = false // Serialization error is not retryable + hasErrors = true + if toolNode != nil { + toolNode.Fail(result.Error) + } + } else { + result.Content = string(contentBytes) + + // Check if it's an error result + if mcpResult.IsError { + result.Error = fmt.Errorf("tool call error: %s", result.Content) + result.IsRetryableError = isRetryableToolError(result.Error) + hasErrors = true + log.Error("[Assistant MCP] Tool call failed: %s - %s (retryable: %v)", toolName, result.Content, result.IsRetryableError) + if toolNode != nil { + toolNode.Fail(result.Error) + } + } else { + // Success + if toolNode != nil { + fmt.Printf(">>> executeServerToolsParallelWithTrace: Completing toolNode %d\n", i) + toolNode.Complete(map[string]any{ + "result": mcpResult.Content, + }) + fmt.Printf(">>> executeServerToolsParallelWithTrace: toolNode %d completed\n", i) + } else { + fmt.Printf(">>> executeServerToolsParallelWithTrace: toolNode %d is nil!\n", i) + } + } + } + + results = append(results, result) + } + + return results, hasErrors +} + +// 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) { + results := make([]ToolCallResult, 0, len(toolCalls)) + hasErrors := false + + log.Trace("[Assistant MCP] Calling %d tools sequentially on server '%s'", len(toolCalls), serverID) + + for _, tc := range toolCalls { + _, toolName, ok := ParseMCPToolName(tc.Function.Name) + if !ok { + results = append(results, ToolCallResult{ + ToolCallID: tc.ID, + Name: tc.Function.Name, + Content: fmt.Sprintf("Invalid tool name format: %s", tc.Function.Name), + Error: fmt.Errorf("invalid tool name format"), + }) + hasErrors = true + continue + } + + // Get tool schema for validation + toolsResponse, err := client.ListTools(mcpCtx, "") + var toolSchema interface{} + if err == nil { + for _, t := range toolsResponse.Tools { + if t.Name == toolName { + toolSchema = t.InputSchema + break + } + } + } + + // Add trace node for this tool call + var toolNode types.Node + if trace != nil { + toolNode, _ = trace.Add( + map[string]any{ + "tool_call_id": tc.ID, + "server": serverID, + "tool": toolName, + "arguments": tc.Function.Arguments, + }, + types.TraceNodeOption{ + Label: fmt.Sprintf("Tool: %s (sequential retry)", toolName), + Type: "mcp_tool", + Icon: "build", + Description: fmt.Sprintf("Retrying MCP tool '%s' on server '%s' sequentially", toolName, serverID), + }, + ) + } + + // Parse arguments with repair support + var args map[string]interface{} + if tc.Function.Arguments != "" { + parsed, err := gouJson.Parse(tc.Function.Arguments) + if err != nil { + result := ToolCallResult{ + ToolCallID: tc.ID, + Name: tc.Function.Name, + Content: fmt.Sprintf("Failed to parse arguments: %v", err), + Error: err, + IsRetryableError: true, // Parsing error is retryable + } + results = append(results, result) + hasErrors = true + if toolNode != nil { + toolNode.Fail(err) + } + continue + } + + // Convert to map + if argsMap, ok := parsed.(map[string]interface{}); ok { + args = argsMap + } else { + err := fmt.Errorf("arguments must be an object, got %T", parsed) + result := ToolCallResult{ + ToolCallID: tc.ID, + Name: tc.Function.Name, + Content: err.Error(), + Error: err, + IsRetryableError: true, // Type error is retryable + } + results = append(results, result) + hasErrors = true + if toolNode != nil { + toolNode.Fail(err) + } + continue + } + + // Validate arguments against tool schema if available + if toolSchema != nil { + if err := gouJson.Validate(args, toolSchema); err != nil { + result := ToolCallResult{ + ToolCallID: tc.ID, + Name: tc.Function.Name, + Content: fmt.Sprintf("Argument validation failed: %v", err), + Error: err, + IsRetryableError: true, // Validation error is retryable + } + results = append(results, result) + hasErrors = true + if toolNode != nil { + toolNode.Fail(err) + } + continue + } + } + } + + // Call single tool + log.Trace("[Assistant MCP] Calling tool: %s", toolName) + mcpResult, err := client.CallTool(mcpCtx, toolName, args) + + result := ToolCallResult{ + ToolCallID: tc.ID, + Name: tc.Function.Name, + } + + if err != nil { + result.Error = err + result.Content = fmt.Sprintf("Tool call failed: %v", err) + result.IsRetryableError = isRetryableToolError(err) + hasErrors = true + log.Error("[Assistant MCP] Tool call failed: %s - %v (retryable: %v)", toolName, err, result.IsRetryableError) + if toolNode != nil { + toolNode.Fail(err) + } + } else { + // Check if result is an error + if mcpResult.IsError { + result.Error = fmt.Errorf("MCP tool error") + result.IsRetryableError = false // MCP internal error is not retryable + hasErrors = true + } + + // Serialize the Content field only ([]ToolContent) + contentBytes, err := jsoniter.Marshal(mcpResult.Content) + if err != nil { + result.Error = err + result.Content = fmt.Sprintf("Failed to serialize result: %v", err) + result.IsRetryableError = false // Serialization error is not retryable + hasErrors = true + if toolNode != nil { + toolNode.Fail(err) + } + } else { + result.Content = string(contentBytes) + if toolNode != nil { + toolNode.Complete(map[string]any{ + "result": mcpResult.Content, + }) + } + } + } + + results = append(results, result) + } + + return results, hasErrors +} diff --git a/agent/assistant/mcp_test.go b/agent/assistant/mcp_test.go new file mode 100644 index 00000000..30edf3b1 --- /dev/null +++ b/agent/assistant/mcp_test.go @@ -0,0 +1,236 @@ +package assistant_test + +import ( + "testing" + + "github.com/yaoapp/yao/agent/assistant" + "github.com/yaoapp/yao/agent/testutils" +) + +func TestMCPToolName(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + tests := []struct { + name string + serverID string + toolName string + wantResult string + }{ + { + name: "Simple tool name", + serverID: "github", + toolName: "search_repos", + wantResult: "github__search_repos", + }, + { + name: "Server with dots", + serverID: "github.enterprise", + toolName: "search_repos", + wantResult: "github_enterprise__search_repos", + }, + { + name: "Tool with underscores", + serverID: "customer-db", + toolName: "create_customer", + wantResult: "customer-db__create_customer", + }, + { + name: "Complex server with multiple dots", + serverID: "com.example.mcp", + toolName: "tool_name", + wantResult: "com_example_mcp__tool_name", + }, + { + name: "Empty server ID", + serverID: "", + toolName: "tool", + wantResult: "", + }, + { + name: "Empty tool name", + serverID: "server", + toolName: "", + wantResult: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := assistant.MCPToolName(tt.serverID, tt.toolName) + if result != tt.wantResult { + t.Errorf("MCPToolName() = %v, want %v", result, tt.wantResult) + } + }) + } +} + +func TestParseMCPToolName(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + tests := []struct { + name string + formattedName string + wantServerID string + wantToolName string + wantOK bool + }{ + { + name: "Valid simple format", + formattedName: "github__search_repos", + wantServerID: "github", + wantToolName: "search_repos", + wantOK: true, + }, + { + name: "Server with dots restored", + formattedName: "github_enterprise__search_repos", + wantServerID: "github.enterprise", + wantToolName: "search_repos", + wantOK: true, + }, + { + name: "Complex server ID with multiple dots", + formattedName: "com_example_mcp_server__tool_name", + wantServerID: "com.example.mcp.server", + wantToolName: "tool_name", + wantOK: true, + }, + { + name: "Tool name with underscores", + formattedName: "server__create_new_user", + wantServerID: "server", + wantToolName: "create_new_user", + wantOK: true, + }, + { + name: "Server with hyphens", + formattedName: "mcp-server__tool", + wantServerID: "mcp-server", + wantToolName: "tool", + wantOK: true, + }, + { + name: "Invalid format - no double underscore", + formattedName: "invalid", + wantServerID: "", + wantToolName: "", + wantOK: false, + }, + { + name: "Invalid format - empty string", + formattedName: "", + wantServerID: "", + wantToolName: "", + wantOK: false, + }, + { + name: "Invalid format - only double underscore", + formattedName: "__", + wantServerID: "", + wantToolName: "", + wantOK: false, + }, + { + name: "Invalid format - ends with double underscore", + formattedName: "server__", + wantServerID: "", + wantToolName: "", + wantOK: false, + }, + { + name: "Invalid format - starts with double underscore", + formattedName: "__tool", + wantServerID: "", + wantToolName: "", + wantOK: false, + }, + { + name: "Invalid format - multiple double underscores", + formattedName: "server__middle__tool", + wantServerID: "", + wantToolName: "", + wantOK: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + serverID, toolName, ok := assistant.ParseMCPToolName(tt.formattedName) + if serverID != tt.wantServerID { + t.Errorf("ParseMCPToolName() serverID = %v, want %v", serverID, tt.wantServerID) + } + if toolName != tt.wantToolName { + t.Errorf("ParseMCPToolName() toolName = %v, want %v", toolName, tt.wantToolName) + } + if ok != tt.wantOK { + t.Errorf("ParseMCPToolName() ok = %v, want %v", ok, tt.wantOK) + } + }) + } +} + +func TestMCPToolName_RoundTrip(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + tests := []struct { + name string + serverID string + toolName string + }{ + { + name: "Simple IDs", + serverID: "github", + toolName: "search_repos", + }, + { + name: "Server with dots", + serverID: "github.enterprise", + toolName: "search", + }, + { + name: "Complex server ID", + serverID: "com.example.mcp.server", + toolName: "tool_name", + }, + { + name: "Server with dashes", + serverID: "mcp-server-123", + toolName: "tool_with_underscores", + }, + { + name: "Mixed dots and dashes", + serverID: "github.enterprise-prod", + toolName: "api_call", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Format + formatted := assistant.MCPToolName(tt.serverID, tt.toolName) + if formatted == "" { + t.Fatal("MCPToolName() returned empty string") + } + + // Parse + serverID, toolName, ok := assistant.ParseMCPToolName(formatted) + + // Verify round-trip + if !ok { + t.Fatal("ParseMCPToolName() failed") + } + if serverID != tt.serverID { + t.Errorf("Round-trip failed: serverID = %v, want %v", serverID, tt.serverID) + } + if toolName != tt.toolName { + t.Errorf("Round-trip failed: toolName = %v, want %v", toolName, tt.toolName) + } + + t.Logf("✓ Round-trip successful: (%s, %s) → %s → (%s, %s)", + tt.serverID, tt.toolName, formatted, serverID, toolName) + }) + } +} diff --git a/agent/assistant/trace.go b/agent/assistant/trace.go new file mode 100644 index 00000000..f025d8a6 --- /dev/null +++ b/agent/assistant/trace.go @@ -0,0 +1,126 @@ +package assistant + +import ( + "fmt" + + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/i18n" + "github.com/yaoapp/yao/trace/types" +) + +// initAgentTraceNode creates and returns the agent trace node +func (ast *Assistant) initAgentTraceNode(ctx *context.Context, inputMessages []context.Message) types.Node { + trace, _ := ctx.Trace() + if trace == nil { + return nil + } + + agentNode, _ := trace.Add(inputMessages, types.TraceNodeOption{ + Label: i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.label"), // "Assistant {{name}}" + Type: "agent", + Icon: "assistant", + Description: i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.description"), // "Assistant {{name}} is processing the request" + }) + + return agentNode +} + +// traceAgentHistory logs the chat history to the agent trace node +func (ast *Assistant) traceAgentHistory(ctx *context.Context, agentNode types.Node, fullMessages []context.Message) { + if agentNode == nil { + return + } + + agentNode.Info( + i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.history"), // "Get Chat History" + map[string]any{"messages": fullMessages}, + ) +} + +// traceCreateHook logs the create hook response to the agent trace node +func (ast *Assistant) traceCreateHook(agentNode types.Node, createResponse *context.HookCreateResponse) { + if agentNode == nil { + return + } + + agentNode.Debug("Call Create Hook", map[string]any{"response": createResponse}) +} + +// traceConnectorCapabilities logs the connector capabilities to the agent trace node +func (ast *Assistant) traceConnectorCapabilities(agentNode types.Node, capabilities *context.ModelCapabilities) { + if agentNode == nil { + return + } + + agentNode.Debug("Get Connector Capabilities", map[string]any{"capabilities": capabilities}) +} + +// traceLLMRequest adds a LLM trace node to the trace +func (ast *Assistant) traceLLMRequest(ctx *context.Context, connID string, completionMessages []context.Message, completionOptions *context.CompletionOptions) { + trace, _ := ctx.Trace() + if trace == nil { + return + } + + trace.Add( + map[string]any{"messages": completionMessages, "options": completionOptions}, + types.TraceNodeOption{ + Label: fmt.Sprintf(i18n.Tr(ast.ID, ctx.Locale, "llm.openai.stream.label"), connID), // "LLM %s" + Type: "llm", + Icon: "psychology", + Description: fmt.Sprintf(i18n.Tr(ast.ID, ctx.Locale, "llm.openai.stream.description"), connID), // "LLM %s is processing the request" + }, + ) +} + +// traceLLMComplete marks the LLM request as complete in the trace +func (ast *Assistant) traceLLMComplete(ctx *context.Context, completionResponse *context.CompletionResponse) { + trace, _ := ctx.Trace() + if trace == nil { + return + } + + trace.Complete(completionResponse) +} + +// traceAgentOutput sets the output of the agent trace node +func (ast *Assistant) traceAgentOutput(agentNode types.Node, createResponse *context.HookCreateResponse, doneResponse *context.ResponseHookDone, completionResponse *context.CompletionResponse) { + if agentNode == nil { + return + } + + output := context.Response{ + Create: createResponse, + Done: doneResponse, + Completion: completionResponse, + } + + agentNode.Complete(output) +} + +// traceAgentFail marks the agent trace node as failed +func (ast *Assistant) traceAgentFail(agentNode types.Node, err error) { + if agentNode == nil { + return + } + + agentNode.Fail(err) +} + +// traceLLMRetryRequest adds a LLM retry trace node to the trace +func (ast *Assistant) traceLLMRetryRequest(ctx *context.Context, connID string, completionMessages []context.Message, completionOptions *context.CompletionOptions) { + trace, _ := ctx.Trace() + if trace == nil { + return + } + + trace.Add( + map[string]any{"messages": completionMessages, "options": completionOptions}, + types.TraceNodeOption{ + Label: fmt.Sprintf("LLM %s (Tool Retry)", connID), + Type: "llm_retry", + Icon: "refresh", + Description: fmt.Sprintf("LLM %s is retrying with tool call error feedback", connID), + }, + ) +} diff --git a/agent/assistant/types.go b/agent/assistant/types.go index 8fb757c9..3768cd46 100644 --- a/agent/assistant/types.go +++ b/agent/assistant/types.go @@ -5,6 +5,7 @@ import ( "io" "github.com/gin-gonic/gin" + jsoniter "github.com/json-iterator/go" "github.com/yaoapp/yao/agent/assistant/hook" chatctx "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/message" @@ -164,3 +165,99 @@ type FileResponse struct { ContentType string Extension string } + +// MCPTool represents a simplified MCP tool for building LLM requests +// This is an internal representation used when collecting tools from MCP servers +// and preparing them for the LLM's tool calling interface +type MCPTool struct { + Name string // Formatted tool name with server prefix (e.g., "server_id__tool_name") + Description string // Tool description from MCP server + Parameters interface{} // JSON Schema for tool parameters (from MCP InputSchema) +} + +// ToolCallResult represents the result of a tool call execution +// Used to track the outcome of MCP tool invocations during agent execution +type ToolCallResult struct { + ToolCallID string // Tool call ID from the LLM (matches the ID in the LLM's tool_calls response) + Name string // Tool name (formatted with server prefix, e.g., "server_id__tool_name") + Content string // Result content (JSON string of the tool's output or error message) + Error error // Error if the call failed (nil if successful) + IsRetryableError bool // Whether the error should be sent to LLM for retry + // true: parameter/validation errors that LLM can fix (e.g., "missing required field") + // false: MCP internal errors that LLM cannot fix (e.g., "network error", "service unavailable") +} + +// Server extracts the MCP server ID from the formatted tool name +// Example: "echo__ping" -> "echo" +func (r *ToolCallResult) Server() string { + serverID, _, _ := ParseMCPToolName(r.Name) + return serverID +} + +// Tool extracts the original tool name without server prefix +// Example: "echo__ping" -> "ping" +func (r *ToolCallResult) Tool() string { + _, toolName, _ := ParseMCPToolName(r.Name) + return toolName +} + +// ParsedContent extracts the actual tool return value from MCP ToolContent array +// According to MCP protocol: +// - Content is []ToolContent array +// - For "text" type, the actual value is in Text field (usually JSON string) +// - For "image" type, returns the Data field +// - For "resource" type, returns the Resource object +// If there are multiple content items, returns an array of parsed values +func (r *ToolCallResult) ParsedContent() (interface{}, error) { + if r.Content == "" { + return nil, nil + } + + // Parse Content as []ToolContent + var toolContents []map[string]interface{} + if err := jsoniter.UnmarshalFromString(r.Content, &toolContents); err != nil { + // If parsing fails, return the string content directly (error message) + return r.Content, nil + } + + // Extract actual values from ToolContent items + var results []interface{} + for _, tc := range toolContents { + contentType, _ := tc["type"].(string) + + switch contentType { + case "text": + // For text type, parse the Text field (usually JSON) + if textStr, ok := tc["text"].(string); ok { + // Try to parse as JSON + var parsed interface{} + if err := jsoniter.UnmarshalFromString(textStr, &parsed); err == nil { + results = append(results, parsed) + } else { + // If not JSON, return as plain string + results = append(results, textStr) + } + } + case "image": + // For image type, return the data and mimeType + results = append(results, map[string]interface{}{ + "type": "image", + "data": tc["data"], + "mimeType": tc["mimeType"], + }) + case "resource": + // For resource type, return the resource object + results = append(results, tc["resource"]) + default: + // Unknown type, return as-is + results = append(results, tc) + } + } + + // If only one result, return it directly (not as array) + if len(results) == 1 { + return results[0], nil + } + + return results, nil +} diff --git a/agent/i18n/builtin.go b/agent/i18n/builtin.go index 04597ec4..b5bf69bb 100644 --- a/agent/i18n/builtin.go +++ b/agent/i18n/builtin.go @@ -12,8 +12,8 @@ func init() { Locale: "en", Messages: map[string]any{ // Assistant: agent.go Stream() function - "assistant.agent.stream.label": "Assistant {{name}}", - "assistant.agent.stream.description": "Assistant {{name}} is processing the request", + "assistant.agent.stream.label": "{{name}}", + "assistant.agent.stream.description": "{{name}} is processing the request", "assistant.agent.stream.history": "Get Chat History", "assistant.agent.stream.capabilities": "Get Connector Capabilities", "assistant.agent.stream.create_hook": "Call Create Hook", @@ -100,8 +100,8 @@ func init() { Locale: "zh-cn", Messages: map[string]any{ // Assistant: agent.go Stream() function - "assistant.agent.stream.label": "助手 {{name}}", - "assistant.agent.stream.description": "助手 {{name}} 正在处理请求", + "assistant.agent.stream.label": "{{name}}", + "assistant.agent.stream.description": "{{name}} 正在处理请求", "assistant.agent.stream.history": "获取聊天历史", "assistant.agent.stream.capabilities": "获取连接器能力", "assistant.agent.stream.create_hook": "调用 Create Hook", @@ -160,8 +160,8 @@ func init() { Locale: "zh", Messages: map[string]any{ // Assistant: agent.go Stream() function - "assistant.agent.stream.label": "助手 {{name}}", - "assistant.agent.stream.description": "助手 {{name}} 正在处理请求", + "assistant.agent.stream.label": "{{name}}", + "assistant.agent.stream.description": "{{name}} 正在处理请求", "assistant.agent.stream.history": "获取聊天历史", "assistant.agent.stream.capabilities": "获取连接器能力", "assistant.agent.stream.create_hook": "调用 Create Hook", diff --git a/agent/i18n/i18n_test.go b/agent/i18n/i18n_test.go index 16aba5a8..3ce700d9 100644 --- a/agent/i18n/i18n_test.go +++ b/agent/i18n/i18n_test.go @@ -961,14 +961,15 @@ func TestBuiltinMessages(t *testing.T) { t.Run("English built-in messages", func(t *testing.T) { // Test assistant messages + // Updated: label now only shows {{name}} without "Assistant" prefix result := TranslateGlobal("en", "{{assistant.agent.stream.label}}") - expected := "Assistant {{name}}" + expected := "{{name}}" if result != expected { t.Errorf("Expected '%s', got '%v'", expected, result) } result = TranslateGlobal("en", "{{assistant.agent.stream.description}}") - expected = "Assistant {{name}} is processing the request" + expected = "{{name}} is processing the request" if result != expected { t.Errorf("Expected '%s', got '%v'", expected, result) } @@ -1002,14 +1003,15 @@ func TestBuiltinMessages(t *testing.T) { t.Run("Chinese (zh-cn) built-in messages", func(t *testing.T) { // Test assistant messages + // Updated: label now only shows {{name}} without "助手" prefix result := TranslateGlobal("zh-cn", "{{assistant.agent.stream.label}}") - expected := "助手 {{name}}" + expected := "{{name}}" if result != expected { t.Errorf("Expected '%s', got '%v'", expected, result) } result = TranslateGlobal("zh-cn", "{{assistant.agent.stream.description}}") - expected = "助手 {{name}} 正在处理请求" + expected = "{{name}} 正在处理请求" if result != expected { t.Errorf("Expected '%s', got '%v'", expected, result) } @@ -1036,8 +1038,9 @@ func TestBuiltinMessages(t *testing.T) { }) t.Run("Chinese (zh) short code", func(t *testing.T) { + // Updated: label now only shows {{name}} without "助手" prefix result := TranslateGlobal("zh", "{{assistant.agent.stream.label}}") - expected := "助手 {{name}}" + expected := "{{name}}" if result != expected { t.Errorf("Expected '%s', got '%v'", expected, result) } @@ -1244,7 +1247,8 @@ func TestTAlias(t *testing.T) { t.Errorf("T and TranslateGlobal should return same result. T: %v, TranslateGlobal: %v", resultT, resultGlobal) } - expected := "Assistant {{name}}" + // Updated: label now only shows {{name}} without "Assistant" prefix + expected := "{{name}}" if resultT != expected { t.Errorf("Expected '%s', got '%v'", expected, resultT) } @@ -1267,17 +1271,18 @@ func TestTAlias(t *testing.T) { }) t.Run("T with nested template (template in template value)", func(t *testing.T) { - // assistant.agent.stream.label = "Assistant {{name}}" (contains {{name}} template) + // assistant.agent.stream.label = "{{name}}" (contains {{name}} template) + // Updated: label now only shows {{name}} without prefix // This tests if we can get the template string itself result := T("en", "{{assistant.agent.stream.label}}") - expected := "Assistant {{name}}" + expected := "{{name}}" if result != expected { t.Errorf("Expected '%s', got '%v'", expected, result) } // Verify Chinese version too resultZh := T("zh-cn", "{{assistant.agent.stream.label}}") - expectedZh := "助手 {{name}}" + expectedZh := "{{name}}" if resultZh != expectedZh { t.Errorf("Expected '%s', got '%v'", expectedZh, resultZh) } diff --git a/agent/llm/providers/openai/openai.go b/agent/llm/providers/openai/openai.go index f387b23c..265e165d 100644 --- a/agent/llm/providers/openai/openai.go +++ b/agent/llm/providers/openai/openai.go @@ -213,7 +213,6 @@ func (p *Provider) Stream(ctx *context.Context, messages []context.Message, opti } maxRetries := 3 - maxValidationRetries := 3 var lastErr error // Get Go context for cancellation support @@ -305,43 +304,14 @@ func (p *Provider) Stream(ctx *context.Context, messages []context.Message, opti }) } - // Check if error is tool call validation failure + // Note: Tool call validation errors should not reach here anymore + // because we now pass through validation failures to Agent layer + // This check is kept for safety but should not trigger if isToolCallValidationError(err) { - // Handle tool call validation retry with feedback to LLM - validationRetryMessages := currentMessages - for validationAttempt := 0; validationAttempt < maxValidationRetries; validationAttempt++ { - if trace != nil { - trace.Warn("Tool call validation failed", map[string]any{ - "attempt": validationAttempt + 1, - "max_retries": maxValidationRetries, - "error": err.Error(), - }) - } - - // Add error feedback to conversation history - validationRetryMessages = append(validationRetryMessages, context.Message{ - Role: context.RoleSystem, - Content: fmt.Sprintf("Tool call validation error: %v. Please correct the tool call arguments to match the required schema.", err), + if trace != nil { + trace.Debug("Tool call validation error (unexpected, should be handled differently)", map[string]any{ + "error": err.Error(), }) - - // Retry with feedback - response, err = p.streamWithRetry(ctx, validationRetryMessages, options, handler) - if err == nil { - return response, nil - } - - // Check if still validation error - if !isToolCallValidationError(err) { - // Different error type, break out of validation retry loop - lastErr = err - break - } - lastErr = err - } - - // If we exhausted validation retries, return the error - if isToolCallValidationError(lastErr) { - return nil, fmt.Errorf("tool call validation failed after %d retries: %w", maxValidationRetries, lastErr) } } @@ -803,12 +773,21 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess response.ToolCalls = toolCalls // Validate tool call results if schema is provided + // Note: If validation fails, we log the error but DO NOT return error + // Instead, we let the response through so Agent layer can handle it + // Agent layer will re-validate and provide better error feedback to LLM if err := p.validateToolCallResults(options, toolCalls); err != nil { + // Log validation error + if trace, _ := ctx.Trace(); trace != nil { + trace.Warn("Tool call validation failed at LLM layer, passing to Agent layer for handling", map[string]any{ + "error": err.Error(), + }) + } // End current message messageTracker.endMessage(handler) - // Tool call validation failed, need to retry with error feedback - return nil, fmt.Errorf("tool call validation failed: %w", err) + // Continue and return response (don't return error) + // Agent layer will handle validation and retry } } @@ -829,7 +808,6 @@ func (p *Provider) Post(ctx *context.Context, messages []context.Message, option } maxRetries := 3 - maxValidationRetries := 3 var lastErr error // Get Go context for cancellation support @@ -880,43 +858,14 @@ func (p *Provider) Post(ctx *context.Context, messages []context.Message, option } lastErr = err - // Check if error is tool call validation failure + // Note: Tool call validation errors should not reach here anymore + // because we now pass through validation failures to Agent layer + // This check is kept for safety but should not trigger if isToolCallValidationError(err) { - // Handle tool call validation retry with feedback to LLM - validationRetryMessages := currentMessages - for validationAttempt := 0; validationAttempt < maxValidationRetries; validationAttempt++ { - if trace != nil { - trace.Warn("Tool call validation failed", map[string]any{ - "attempt": validationAttempt + 1, - "max_retries": maxValidationRetries, - "error": err.Error(), - }) - } - - // Add error feedback to conversation history - validationRetryMessages = append(validationRetryMessages, context.Message{ - Role: context.RoleSystem, - Content: fmt.Sprintf("Tool call validation error: %v. Please correct the tool call arguments to match the required schema.", err), + if trace != nil { + trace.Debug("Tool call validation error in Post (unexpected, should be handled differently)", map[string]any{ + "error": err.Error(), }) - - // Retry with feedback - response, err = p.postWithRetry(ctx, validationRetryMessages, options) - if err == nil { - return response, nil - } - - // Check if still validation error - if !isToolCallValidationError(err) { - // Different error type, break out of validation retry loop - lastErr = err - break - } - lastErr = err - } - - // If we exhausted validation retries, return the error - if isToolCallValidationError(lastErr) { - return nil, fmt.Errorf("tool call validation failed after %d retries: %w", maxValidationRetries, lastErr) } } diff --git a/agent/store/types/convert.go b/agent/store/types/convert.go index 43b50260..d8ddc59f 100644 --- a/agent/store/types/convert.go +++ b/agent/store/types/convert.go @@ -63,17 +63,9 @@ func ToMCPServers(v interface{}) (*MCPServers, error) { case MCPServers: return &mcp, nil - case []string: - return &MCPServers{Servers: mcp}, nil - - case []interface{}: - var servers []string - for _, item := range mcp { - servers = append(servers, cast.ToString(item)) - } - return &MCPServers{Servers: servers}, nil - default: + // For any type (including []string, []interface{}, map[string]interface{}), + // marshal and unmarshal to MCPServers using custom UnmarshalJSON raw, err := jsoniter.Marshal(mcp) if err != nil { return nil, fmt.Errorf("mcp format error: %s", err.Error()) diff --git a/agent/store/types/convert_test.go b/agent/store/types/convert_test.go index e21e8759..ad17acbe 100644 --- a/agent/store/types/convert_test.go +++ b/agent/store/types/convert_test.go @@ -118,7 +118,7 @@ func TestToMCPServers(t *testing.T) { }) t.Run("MCPServersPointer", func(t *testing.T) { - mcp := &MCPServers{Servers: []string{"server1", "server2"}} + mcp := &MCPServers{Servers: []MCPServerConfig{{ServerID: "server1"}, {ServerID: "server2"}}} result, err := ToMCPServers(mcp) if err != nil { t.Errorf("Expected no error, got: %v", err) @@ -129,7 +129,7 @@ func TestToMCPServers(t *testing.T) { }) t.Run("MCPServersValue", func(t *testing.T) { - mcp := MCPServers{Servers: []string{"server1", "server2"}} + mcp := MCPServers{Servers: []MCPServerConfig{{ServerID: "server1"}, {ServerID: "server2"}}} result, err := ToMCPServers(mcp) if err != nil { t.Errorf("Expected no error, got: %v", err) @@ -139,37 +139,9 @@ func TestToMCPServers(t *testing.T) { } }) - t.Run("StringSlice", func(t *testing.T) { - servers := []string{"server1", "server2", "server3"} - result, err := ToMCPServers(servers) - if err != nil { - t.Errorf("Expected no error, got: %v", err) - } - if len(result.Servers) != 3 { - t.Errorf("Expected 3 servers, got %d", len(result.Servers)) - } - if result.Servers[0] != "server1" { - t.Errorf("Expected 'server1', got '%s'", result.Servers[0]) - } - }) - - t.Run("InterfaceSlice", func(t *testing.T) { - servers := []interface{}{"server1", "server2", 456} - result, err := ToMCPServers(servers) - if err != nil { - t.Errorf("Expected no error, got: %v", err) - } - if len(result.Servers) != 3 { - t.Errorf("Expected 3 servers, got %d", len(result.Servers)) - } - if result.Servers[2] != "456" { - t.Errorf("Expected '456', got '%s'", result.Servers[2]) - } - }) - t.Run("MapInput", func(t *testing.T) { data := map[string]interface{}{ - "servers": []string{"server1", "server2"}, + "servers": []interface{}{"server1", "server2"}, } result, err := ToMCPServers(data) if err != nil { @@ -178,6 +150,9 @@ func TestToMCPServers(t *testing.T) { if len(result.Servers) != 2 { t.Errorf("Expected 2 servers, got %d", len(result.Servers)) } + if result.Servers[0].ServerID != "server1" { + t.Errorf("Expected 'server1', got '%s'", result.Servers[0].ServerID) + } }) t.Run("InvalidInput", func(t *testing.T) { diff --git a/agent/store/types/mcp_test.go b/agent/store/types/mcp_test.go new file mode 100644 index 00000000..454dbbc5 --- /dev/null +++ b/agent/store/types/mcp_test.go @@ -0,0 +1,302 @@ +package types + +import ( + "encoding/json" + "testing" +) + +func TestMCPServerConfig_UnmarshalJSON(t *testing.T) { + tests := []struct { + name string + input string + want MCPServerConfig + wantErr bool + }{ + { + name: "Simple string", + input: `"server1"`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: nil, + Tools: nil, + }, + wantErr: false, + }, + { + name: "Tools array only", + input: `{"server1": ["tool1", "tool2"]}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: nil, + Tools: []string{"tool1", "tool2"}, + }, + wantErr: false, + }, + { + name: "Full config with resources and tools", + input: `{"server1": {"resources": ["res1", "res2"], "tools": ["tool1", "tool2"]}}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: []string{"res1", "res2"}, + Tools: []string{"tool1", "tool2"}, + }, + wantErr: false, + }, + { + name: "Only resources", + input: `{"server1": {"resources": ["res1"]}}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: []string{"res1"}, + Tools: nil, + }, + wantErr: false, + }, + { + name: "Only tools", + input: `{"server1": {"tools": ["tool1"]}}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: nil, + Tools: []string{"tool1"}, + }, + wantErr: false, + }, + { + name: "Standard object format", + input: `{"server_id": "server1", "resources": ["res1"], "tools": ["tool1"]}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: []string{"res1"}, + Tools: []string{"tool1"}, + }, + wantErr: false, + }, + { + name: "Standard object format - no resources/tools", + input: `{"server_id": "server1"}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: nil, + Tools: nil, + }, + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got MCPServerConfig + err := json.Unmarshal([]byte(tt.input), &got) + if (err != nil) != tt.wantErr { + t.Errorf("UnmarshalJSON() error = %v, wantErr %v", err, tt.wantErr) + return + } + + if !tt.wantErr { + if got.ServerID != tt.want.ServerID { + t.Errorf("ServerID = %v, want %v", got.ServerID, tt.want.ServerID) + } + if !stringSlicesEqual(got.Resources, tt.want.Resources) { + t.Errorf("Resources = %v, want %v", got.Resources, tt.want.Resources) + } + if !stringSlicesEqual(got.Tools, tt.want.Tools) { + t.Errorf("Tools = %v, want %v", got.Tools, tt.want.Tools) + } + } + }) + } +} + +func TestMCPServers_UnmarshalJSON(t *testing.T) { + tests := []struct { + name string + input string + want []MCPServerConfig + wantErr bool + }{ + { + name: "Simple string array", + input: `{"servers": ["server1", "server2", "server3"]}`, + want: []MCPServerConfig{ + {ServerID: "server1"}, + {ServerID: "server2"}, + {ServerID: "server3"}, + }, + wantErr: false, + }, + { + name: "Mixed formats", + input: `{"servers": ["server1", {"server2": ["tool1", "tool2"]}, {"server3": {"resources": ["res1"], "tools": ["tool3"]}}]}`, + want: []MCPServerConfig{ + {ServerID: "server1"}, + {ServerID: "server2", Tools: []string{"tool1", "tool2"}}, + {ServerID: "server3", Resources: []string{"res1"}, Tools: []string{"tool3"}}, + }, + wantErr: false, + }, + { + name: "Empty servers", + input: `{"servers": []}`, + want: []MCPServerConfig{}, + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got MCPServers + err := json.Unmarshal([]byte(tt.input), &got) + if (err != nil) != tt.wantErr { + t.Errorf("UnmarshalJSON() error = %v, wantErr %v", err, tt.wantErr) + return + } + + if !tt.wantErr { + if len(got.Servers) != len(tt.want) { + t.Errorf("got %d servers, want %d", len(got.Servers), len(tt.want)) + return + } + + for i := range got.Servers { + if got.Servers[i].ServerID != tt.want[i].ServerID { + t.Errorf("Server[%d].ServerID = %v, want %v", i, got.Servers[i].ServerID, tt.want[i].ServerID) + } + if !stringSlicesEqual(got.Servers[i].Resources, tt.want[i].Resources) { + t.Errorf("Server[%d].Resources = %v, want %v", i, got.Servers[i].Resources, tt.want[i].Resources) + } + if !stringSlicesEqual(got.Servers[i].Tools, tt.want[i].Tools) { + t.Errorf("Server[%d].Tools = %v, want %v", i, got.Servers[i].Tools, tt.want[i].Tools) + } + } + } + }) + } +} + +func TestMCPServerConfig_MarshalJSON(t *testing.T) { + tests := []struct { + name string + config MCPServerConfig + want string + }{ + { + name: "Only ServerID - should be simple string", + config: MCPServerConfig{ + ServerID: "server1", + }, + want: `"server1"`, + }, + { + name: "With Tools - should be object", + config: MCPServerConfig{ + ServerID: "server1", + Tools: []string{"tool1", "tool2"}, + }, + want: `{"server_id":"server1","tools":["tool1","tool2"]}`, + }, + { + name: "With Resources - should be object", + config: MCPServerConfig{ + ServerID: "server1", + Resources: []string{"res1"}, + }, + want: `{"server_id":"server1","resources":["res1"]}`, + }, + { + name: "With Both - should be object", + config: MCPServerConfig{ + ServerID: "server1", + Resources: []string{"res1"}, + Tools: []string{"tool1"}, + }, + want: `{"server_id":"server1","resources":["res1"],"tools":["tool1"]}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := json.Marshal(tt.config) + if err != nil { + t.Errorf("MarshalJSON() error = %v", err) + return + } + if string(got) != tt.want { + t.Errorf("MarshalJSON() = %s, want %s", string(got), tt.want) + } + }) + } +} + +func TestMCPServerConfig_RoundTrip(t *testing.T) { + tests := []struct { + name string + config MCPServerConfig + }{ + { + name: "Simple ServerID", + config: MCPServerConfig{ + ServerID: "server1", + }, + }, + { + name: "With Tools", + config: MCPServerConfig{ + ServerID: "server2", + Tools: []string{"tool1", "tool2"}, + }, + }, + { + name: "With Resources and Tools", + config: MCPServerConfig{ + ServerID: "server3", + Resources: []string{"res1", "res2"}, + Tools: []string{"tool3", "tool4"}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Marshal + data, err := json.Marshal(tt.config) + if err != nil { + t.Fatalf("Marshal error = %v", err) + } + + // Unmarshal + var got MCPServerConfig + err = json.Unmarshal(data, &got) + if err != nil { + t.Fatalf("Unmarshal error = %v", err) + } + + // Compare + if got.ServerID != tt.config.ServerID { + t.Errorf("ServerID = %v, want %v", got.ServerID, tt.config.ServerID) + } + if !stringSlicesEqual(got.Resources, tt.config.Resources) { + t.Errorf("Resources = %v, want %v", got.Resources, tt.config.Resources) + } + if !stringSlicesEqual(got.Tools, tt.config.Tools) { + t.Errorf("Tools = %v, want %v", got.Tools, tt.config.Tools) + } + }) + } +} + +// Helper function to compare string slices (nil-safe) +func stringSlicesEqual(a, b []string) bool { + if len(a) == 0 && len(b) == 0 { + return true + } + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} diff --git a/agent/store/types/types.go b/agent/store/types/types.go index 3f3f7867..c145c65f 100644 --- a/agent/store/types/types.go +++ b/agent/store/types/types.go @@ -1,6 +1,9 @@ package types import ( + "encoding/json" + "fmt" + "github.com/yaoapp/xun/dbal/query" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" @@ -99,11 +102,98 @@ type KnowledgeBase struct { } // MCPServers the MCP servers configuration +// Supports multiple formats in the servers array: +// - Simple string: "server_id" +// - With tools: {"server_id": ["tool1", "tool2"]} +// - With resources and tools: {"server_id": {"resources": [...], "tools": [...]}} type MCPServers struct { - Servers []string `json:"servers,omitempty"` // MCP server IDs + Servers []MCPServerConfig `json:"servers,omitempty"` // MCP server configurations Options map[string]interface{} `json:"options,omitempty"` // Additional options for MCP servers } +// MCPServerConfig represents a single MCP server configuration +type MCPServerConfig struct { + ServerID string `json:"server_id,omitempty"` // MCP server ID + Resources []string `json:"resources,omitempty"` // Resources to use (optional) + Tools []string `json:"tools,omitempty"` // Tools to use (optional) +} + +// UnmarshalJSON implements custom JSON unmarshaling for MCPServerConfig +// Supports multiple input formats: +// 1. Simple string: "server_id" +// 2. Standard object: {"server_id": "server1", "resources": [...], "tools": [...]} +// 3. Tools array: {"server_id": ["tool1", "tool2"]} +// 4. Full config: {"server_id": {"resources": [...], "tools": [...]}} +func (m *MCPServerConfig) UnmarshalJSON(data []byte) error { + // Try to unmarshal as string first + var str string + if err := json.Unmarshal(data, &str); err == nil { + m.ServerID = str + return nil + } + + // Try to unmarshal as standard object with server_id field + type Alias MCPServerConfig + var stdObj Alias + if err := json.Unmarshal(data, &stdObj); err == nil && stdObj.ServerID != "" { + *m = MCPServerConfig(stdObj) + return nil + } + + // Try to unmarshal as object with single key (alternative formats) + var obj map[string]json.RawMessage + if err := json.Unmarshal(data, &obj); err != nil { + return err + } + + // Should have exactly one key (the server ID) + if len(obj) != 1 { + return fmt.Errorf("MCPServerConfig object must have exactly one key or server_id field") + } + + // Get the server ID (the only key) + for serverID, value := range obj { + m.ServerID = serverID + + // Try to unmarshal value as array of strings (format c: tools only) + var tools []string + if err := json.Unmarshal(value, &tools); err == nil { + m.Tools = tools + return nil + } + + // Try to unmarshal as object with resources and tools (format b) + var detail struct { + Resources []string `json:"resources,omitempty"` + Tools []string `json:"tools,omitempty"` + } + if err := json.Unmarshal(value, &detail); err == nil { + m.Resources = detail.Resources + m.Tools = detail.Tools + return nil + } + + return fmt.Errorf("invalid format for server '%s'", serverID) + } + + return nil +} + +// MarshalJSON implements custom JSON marshaling for MCPServerConfig +// Serializes to different formats based on content: +// 1. If only ServerID: "server_id" +// 2. If has Resources or Tools: {"server_id": "...", "resources": [...], "tools": [...]} +func (m MCPServerConfig) MarshalJSON() ([]byte, error) { + // If only ServerID, serialize as simple string + if len(m.Resources) == 0 && len(m.Tools) == 0 { + return json.Marshal(m.ServerID) + } + + // Otherwise, use standard object format + type Alias MCPServerConfig + return json.Marshal(Alias(m)) +} + // Workflow the workflow configuration type Workflow struct { Workflows []string `json:"workflows,omitempty"` // Workflow IDs diff --git a/agent/store/xun/assistant_test.go b/agent/store/xun/assistant_test.go index ce50c548..4ba8e4a2 100644 --- a/agent/store/xun/assistant_test.go +++ b/agent/store/xun/assistant_test.go @@ -2,6 +2,7 @@ package xun import ( "fmt" + "os" "strings" "testing" "time" @@ -16,10 +17,12 @@ import ( func TestMain(m *testing.M) { // Setup will be done in each test via test.Prepare - // Run tests and exit test.Prepare(nil, config.Conf) defer test.Clean() - m.Run() + + // Run tests and exit with appropriate exit code + code := m.Run() + os.Exit(code) } // TestSaveAssistant tests creating and updating assistants @@ -199,6 +202,137 @@ func TestSaveAssistant(t *testing.T) { } }) + t.Run("SaveWithMCPServers", func(t *testing.T) { + // Test creating assistant with MCP servers directly + // This will test that: + // - server1 (no tools/resources) serializes as "server1" + // - server2 (with tools) serializes as {"server_id":"server2","tools":[...]} + // - server3 (with both) serializes as {"server_id":"server3","resources":[...],"tools":[...]} + assistant := &types.AssistantModel{ + Name: "MCP Save Test", + Type: "assistant", + Connector: "openai", + Share: "private", + MCP: &types.MCPServers{ + Servers: []types.MCPServerConfig{ + {ServerID: "server1"}, + { + ServerID: "server2", + Tools: []string{"tool1", "tool2"}, + }, + { + ServerID: "server3", + Resources: []string{"res1", "res2"}, + Tools: []string{"tool3", "tool4"}, + }, + }, + Options: map[string]interface{}{ + "timeout": 30, + }, + }, + } + + id, err := store.SaveAssistant(assistant) + if err != nil { + t.Fatalf("Failed to save assistant with MCP: %v", err) + } + + // Retrieve and verify MCP configuration + retrieved, err := store.GetAssistant(id) + if err != nil { + t.Fatalf("Failed to retrieve assistant: %v", err) + } + + if retrieved.MCP == nil { + t.Fatal("Expected MCP to be set") + } + + if len(retrieved.MCP.Servers) != 3 { + t.Errorf("Expected 3 MCP servers, got %d", len(retrieved.MCP.Servers)) + } + + // Verify server1 (simple format) + if retrieved.MCP.Servers[0].ServerID != "server1" { + t.Errorf("Expected server1, got '%s'", retrieved.MCP.Servers[0].ServerID) + } + + // Verify server2 (with tools) + if retrieved.MCP.Servers[1].ServerID != "server2" { + t.Errorf("Expected server2, got '%s'", retrieved.MCP.Servers[1].ServerID) + } + if len(retrieved.MCP.Servers[1].Tools) != 2 { + t.Errorf("Expected 2 tools for server2, got %d", len(retrieved.MCP.Servers[1].Tools)) + } + + // Verify server3 (with resources and tools) + if retrieved.MCP.Servers[2].ServerID != "server3" { + t.Errorf("Expected server3, got '%s'", retrieved.MCP.Servers[2].ServerID) + } + if len(retrieved.MCP.Servers[2].Resources) != 2 { + t.Errorf("Expected 2 resources for server3, got %d", len(retrieved.MCP.Servers[2].Resources)) + } + if len(retrieved.MCP.Servers[2].Tools) != 2 { + t.Errorf("Expected 2 tools for server3, got %d", len(retrieved.MCP.Servers[2].Tools)) + } + + // Verify options + if retrieved.MCP.Options == nil { + t.Error("Expected MCP options to be set") + } + if timeout, ok := retrieved.MCP.Options["timeout"].(float64); !ok || timeout != 30 { + t.Errorf("Expected timeout 30, got %v", retrieved.MCP.Options["timeout"]) + } + + t.Logf("Successfully verified MCP configuration for assistant %s", id) + }) + + t.Run("UpdateWithMCPServers", func(t *testing.T) { + // Create assistant without MCP + assistant := &types.AssistantModel{ + Name: "MCP Update Test", + Type: "assistant", + Connector: "openai", + Share: "private", + } + + id, err := store.SaveAssistant(assistant) + if err != nil { + t.Fatalf("Failed to create assistant: %v", err) + } + + // Update assistant with MCP + assistant.MCP = &types.MCPServers{ + Servers: []types.MCPServerConfig{ + {ServerID: "new-server1"}, + { + ServerID: "new-server2", + Tools: []string{"newtool1"}, + }, + }, + } + + _, err = store.SaveAssistant(assistant) + if err != nil { + t.Fatalf("Failed to update assistant with MCP: %v", err) + } + + // Retrieve and verify + retrieved, err := store.GetAssistant(id) + if err != nil { + t.Fatalf("Failed to retrieve assistant: %v", err) + } + + if retrieved.MCP == nil || len(retrieved.MCP.Servers) != 2 { + t.Errorf("Expected 2 MCP servers, got %v", retrieved.MCP) + } + + if retrieved.MCP.Servers[0].ServerID != "new-server1" { + t.Errorf("Expected new-server1, got '%s'", retrieved.MCP.Servers[0].ServerID) + } + + t.Logf("Successfully updated and verified MCP for assistant %s", id) + }) + t.Run("UsesConfiguration", func(t *testing.T) { // Test assistant with Uses configuration assistant := &types.AssistantModel{ @@ -2071,6 +2205,95 @@ func TestUpdateAssistant(t *testing.T) { if retrieved.MCP == nil || len(retrieved.MCP.Servers) != 2 { t.Errorf("Expected 2 MCP servers, got %v", retrieved.MCP) } + if retrieved.MCP.Servers[0].ServerID != "server1" { + t.Errorf("Expected first server 'server1', got '%s'", retrieved.MCP.Servers[0].ServerID) + } + }) + + t.Run("UpdateMCPWithToolsAndResources", func(t *testing.T) { + // Create assistant + assistant := &types.AssistantModel{ + Name: "MCP Advanced Test", + Type: "assistant", + Connector: "openai", + Share: "private", + } + + id, err := store.SaveAssistant(assistant) + if err != nil { + t.Fatalf("Failed to create assistant: %v", err) + } + + // Update with MCP servers using advanced configuration + updates := map[string]interface{}{ + "mcp": map[string]interface{}{ + "servers": []interface{}{ + "server1", // Simple format + map[string]interface{}{ + "server2": []string{"tool1", "tool2"}, // Tools only + }, + map[string]interface{}{ + "server3": map[string]interface{}{ + "resources": []string{"res1", "res2"}, + "tools": []string{"tool3", "tool4"}, + }, + }, + }, + }, + } + + err = store.UpdateAssistant(id, updates) + if err != nil { + t.Fatalf("Failed to update MCP: %v", err) + } + + // Verify updates + retrieved, err := store.GetAssistant(id) + if err != nil { + t.Fatalf("Failed to retrieve assistant: %v", err) + } + + if retrieved.MCP == nil || len(retrieved.MCP.Servers) != 3 { + t.Fatalf("Expected 3 MCP servers, got %d", len(retrieved.MCP.Servers)) + } + + // Verify server1 (simple format) + if retrieved.MCP.Servers[0].ServerID != "server1" { + t.Errorf("Expected server1, got '%s'", retrieved.MCP.Servers[0].ServerID) + } + if len(retrieved.MCP.Servers[0].Tools) != 0 { + t.Errorf("Expected no tools for server1, got %v", retrieved.MCP.Servers[0].Tools) + } + + // Verify server2 (tools only) + if retrieved.MCP.Servers[1].ServerID != "server2" { + t.Errorf("Expected server2, got '%s'", retrieved.MCP.Servers[1].ServerID) + } + if len(retrieved.MCP.Servers[1].Tools) != 2 { + t.Errorf("Expected 2 tools for server2, got %d", len(retrieved.MCP.Servers[1].Tools)) + } + if retrieved.MCP.Servers[1].Tools[0] != "tool1" { + t.Errorf("Expected tool1, got '%s'", retrieved.MCP.Servers[1].Tools[0]) + } + + // Verify server3 (full config) + if retrieved.MCP.Servers[2].ServerID != "server3" { + t.Errorf("Expected server3, got '%s'", retrieved.MCP.Servers[2].ServerID) + } + if len(retrieved.MCP.Servers[2].Resources) != 2 { + t.Errorf("Expected 2 resources for server3, got %d", len(retrieved.MCP.Servers[2].Resources)) + } + if len(retrieved.MCP.Servers[2].Tools) != 2 { + t.Errorf("Expected 2 tools for server3, got %d", len(retrieved.MCP.Servers[2].Tools)) + } + if retrieved.MCP.Servers[2].Resources[0] != "res1" { + t.Errorf("Expected res1, got '%s'", retrieved.MCP.Servers[2].Resources[0]) + } + if retrieved.MCP.Servers[2].Tools[0] != "tool3" { + t.Errorf("Expected tool3, got '%s'", retrieved.MCP.Servers[2].Tools[0]) + } + + t.Logf("Successfully verified MCP advanced configuration for assistant %s", id) }) t.Run("UpdateUses", func(t *testing.T) {