opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #1357 from trheyi/main

Browse source 
Streamlining and Enhancing Assistant Functionality
This commit is contained in:
Max 2025-11-29 17:54:51 +08:00 • committed by GitHub
commit b613cdd423
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 2638 additions and 272 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

299
agent/assistant/agent.go
View file

@ -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
// === 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 ===
}
// Log the capabilities
if agentNode != nil {
agentNode.Debug("Get Connector Capabilities", map[string]any{"capabilities": capabilities})
}
// ================================================
// Execute tool calls with retry
// ================================================
if completionResponse != nil && completionResponse.ToolCalls != nil {
// 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"
},
)
}
// === Debug Tool Calls ===
fmt.Println("--- Debug Tool Calls --------------------------------")
utils.Dump(completionResponse.ToolCalls)
// Create LLM instance with connector and options
llmInstance, err := llm.New(conn, completionOptions)
// === 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 {
// Send error stream_end for root stack
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
}
// 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 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 {
// Send error stream_end for root stack
log.Trace("[AGENT] Calling sendStreamEndOnError")
log.Error("[AGENT] LLM retry failed: %v", err)
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
log.Trace("[AGENT] sendStreamEndOnError returned")
return nil, err
}
// Mark LLM Request Complete
if trace != nil {
trace.Complete(completionResponse)
}
// 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
}
// Request MCP hook ( Optional )
var mcpResponse *context.ResponseHookMCP
if ast.MCP != nil {
_ = mcpResponse // mcpResponse is available for further processing
// Update messages for next iteration
currentMessages = retryMessages
}
// MCP Execution Loop
// 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
}

18
agent/assistant/assistant.go
View file

@ -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{})

85
agent/assistant/build.go
View file

@ -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
}

222
agent/assistant/build_mcp_test.go Normal file
View file

@ -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
}

31
agent/assistant/history.go Normal file
View file

@ -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
}

8
agent/assistant/hook/mcp.go
View file

@ -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
}

126
agent/assistant/llm.go Normal file
View file

@ -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
}

834
agent/assistant/mcp.go Normal file
View file

@ -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
}

236
agent/assistant/mcp_test.go Normal file
View file

@ -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)
})
}
}

126
agent/assistant/trace.go Normal file
View file

@ -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),
},
)
}

97
agent/assistant/types.go
View file

@ -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
}

12
agent/i18n/builtin.go
View file

@ -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",

23
agent/i18n/i18n_test.go
View file

@ -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)
}

89
agent/llm/providers/openai/openai.go
View file

@ -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,44 +304,15 @@ 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,
trace.Debug("Tool call validation error (unexpected, should be handled differently)", map[string]any{
"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),
})
// 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)
}
}
// Check if error is retryable (network errors, rate limits, etc.)
@ -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,44 +858,15 @@ 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,
trace.Debug("Tool call validation error in Post (unexpected, should be handled differently)", map[string]any{
"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),
})
// 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)
}
}
// Check if error is retryable (network errors, rate limits, etc.)

12
agent/store/types/convert.go
View file

@ -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())

37
agent/store/types/convert_test.go
View file

@ -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) {

302
agent/store/types/mcp_test.go Normal file
View file

@ -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
}

92
agent/store/types/types.go
View file

@ -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

227
agent/store/xun/assistant_test.go
View file

@ -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) {