Merge pull request #1357 from trheyi/main

Streamlining and Enhancing Assistant Functionality

agent/assistant/agent.go

@@ -7,12 +7,11 @@ import (
 	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/kun/log"
+	"github.com/yaoapp/kun/utils"
 	"github.com/yaoapp/yao/agent/assistant/handlers"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/i18n"
-	"github.com/yaoapp/yao/agent/llm"
 	"github.com/yaoapp/yao/agent/output/message"
-	"github.com/yaoapp/yao/trace/types"
 )
 
 // Stream stream the agent
@@ -34,174 +33,198 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 		})
 	}
 
+	// ================================================
+	// Initialize
+	// ================================================
+
 	// Initialize stack and auto-handle completion/failure/restore
-	_, traceID, done := context.EnterStack(ctx, ast.ID, ctx.Referer)
+	_, _, done := context.EnterStack(ctx, ast.ID, ctx.Referer)
 	defer done()
 
-	_ = traceID // traceID is available for trace logging
-
-	// Get connector and capabilities early (before sending stream_start)
-	// so that output adapters can use them when converting stream_start event
-	if ast.Prompts != nil || ast.MCP != nil {
-		_, capabilities, err := ast.GetConnector(ctx)
-		if err != nil {
-			streamHandler := ast.getStreamHandler(ctx, handler...)
-			ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
-			return nil, err
-		}
-
-		// Set capabilities in context for output adapters to use
-		if capabilities != nil {
-			ctx.Capabilities = capabilities
-		}
-	}
-
 	// Determine stream handler
 	streamHandler := ast.getStreamHandler(ctx, handler...)
 
+	// Get connector and capabilities early (before sending stream_start)
+	// so that output adapters can use them when converting stream_start event
+	err = ast.initializeCapabilities(ctx)
+	if err != nil {
+		ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
+		return nil, err
+	}
+
 	// Send ChunkStreamStart only for root stack (agent-level stream start)
 	// Now ctx.Capabilities is set, so output adapters can use it
 	ast.sendAgentStreamStart(ctx, streamHandler, streamStartTime)
 
-	// Trace Add
+	// Initialize agent trace node
-	trace, _ := ctx.Trace()
+	agentNode := ast.initAgentTraceNode(ctx, inputMessages)
-	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"
-		})
-	}
 
-	// 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 {
+		ast.traceAgentFail(agentNode, err)
-			agentNode.Fail(err)
-		}
-		// Send error stream_end for root stack
 		ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
 		return nil, err
 	}
 
-	// Log the chat history
+	// ================================================
-	if agentNode != nil {
+	//  Execute Create Hook
-		agentNode.Info(i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.history"), map[string]any{"messages": fullMessages}) // "Get Chat History"
+	// ================================================
-	}
-
 	// 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 {
+			ast.traceAgentFail(agentNode, err)
-				agentNode.Fail(err)
-			}
 			// Send error stream_end for root stack
 			ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
 			return nil, err
 		}
 
 		// Log the create response
-		if agentNode != nil {
+		ast.traceCreateHook(agentNode, createResponse)
-			agentNode.Debug("Call Create Hook", map[string]any{"response": 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 {
+			ast.traceAgentFail(agentNode, err)
-				agentNode.Fail(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)
+		// Execute the LLM streaming call
-		conn, capabilities, err := ast.GetConnector(ctx)
+		completionResponse, err = ast.executeLLMStream(ctx, completionMessages, completionOptions, agentNode, streamHandler)
 		if err != nil {
-			if agentNode != nil {
+			ast.traceAgentFail(agentNode, err)
-				agentNode.Fail(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
+		// === Debug Completion Response ===
-		if completionOptions.Capabilities == nil && capabilities != nil {
+		fmt.Println("--- Debug Completion Response ----------------------")
-			completionOptions.Capabilities = capabilities
+		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 {
+	// Execute tool calls with retry
-			agentNode.Debug("Get Connector Capabilities", map[string]any{"capabilities": capabilities})
+	// ================================================
-		}
+	if completionResponse != nil && completionResponse.ToolCalls != nil {
 
-		// Trace Add
+		// === Debug Tool Calls ===
-		if trace != nil {
+		fmt.Println("--- Debug Tool Calls --------------------------------")
-			trace.Add(
+		utils.Dump(completionResponse.ToolCalls)
-				map[string]any{"messages": completionMessages, "options": completionOptions},
-				types.TraceNodeOption{
-					Label:       fmt.Sprintf(i18n.Tr(ast.ID, ctx.Locale, "llm.openai.stream.label"), conn.ID()), // "LLM %s"
-					Type:        "llm",
-					Icon:        "psychology",
-					Description: fmt.Sprintf(i18n.Tr(ast.ID, ctx.Locale, "llm.openai.stream.description"), conn.ID()), // "LLM %s is processing the request"
-				},
-			)
-		}
 
-		// Create LLM instance with connector and options
+		// === End Debug Tool Calls ===
-		llmInstance, err := llm.New(conn, completionOptions)
+
+		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)
+			// If it's the last attempt, return error
-		log.Trace("[AGENT] Calling LLM Stream: assistant=%s", ast.ID)
+			if attempt == maxToolRetries-1 {
-		completionResponse, err = llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler)
+				err := fmt.Errorf("tool calls failed after %d attempts", maxToolRetries)
-		log.Trace("[AGENT] LLM Stream returned: assistant=%s, err=%v", ast.ID, err)
+				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.Error("[AGENT] LLM retry failed: %v", err)
-			log.Trace("[AGENT] Calling sendStreamEndOnError")
+				ast.traceAgentFail(agentNode, err)
 				ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
-			log.Trace("[AGENT] sendStreamEndOnError returned")
 				return nil, err
 			}
 
-		// Mark LLM Request Complete
+			// If LLM didn't return tool calls, it might have given up
-		if trace != nil {
+			if currentResponse.ToolCalls == nil {
-			trace.Complete(completionResponse)
+				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 )
+			// Update messages for next iteration
-	var mcpResponse *context.ResponseHookMCP
+			currentMessages = retryMessages
-	if ast.MCP != nil {
+		}
-		_ = mcpResponse // mcpResponse is available for further processing
 
-		// 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 {
+	ast.traceAgentOutput(agentNode, createResponse, doneResponse, completionResponse)
-		agentNode.SetOutput(context.Response{Create: createResponse, Done: doneResponse, Completion: 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
+}

agent/assistant/assistant.go

@@ -152,8 +152,22 @@ func (ast *Assistant) Clone() *Assistant {
 	if ast.MCP != nil {
 		clone.MCP = &store.MCPServers{}
 		if ast.MCP.Servers != nil {
-			clone.MCP.Servers = make([]string, len(ast.MCP.Servers))
+			clone.MCP.Servers = make([]store.MCPServerConfig, len(ast.MCP.Servers))
-			copy(clone.MCP.Servers, 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{})

agent/assistant/build.go

@@ -9,14 +9,14 @@ import (
 
 // BuildRequest build the LLM request
 func (ast *Assistant) BuildRequest(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse) ([]context.Message, *context.CompletionOptions, error) {
-	// Build final messages with proper priority
+	// Build completion options from createResponse and ctx (includes MCP tools)
-	finalMessages, err := ast.buildMessages(ctx, messages, createResponse)
+	options, mcpSamplesPrompt, err := ast.buildCompletionOptions(ctx, createResponse)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	// Build completion options from createResponse and ctx
+	// Build final messages with proper priority (includes MCP samples if available)
-	options, err := ast.buildCompletionOptions(ctx, createResponse)
+	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
+}

agent/assistant/build_mcp_test.go

@@ -0,0 +1,222 @@
+package assistant_test
+
+import (
+	"testing"
+
+	"github.com/yaoapp/yao/agent/assistant"
+	"github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/agent/testutils"
+)
+
+// TestBuildRequest_MCP tests MCP tool integration in BuildRequest
+func TestBuildRequest_MCP(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	agent, err := assistant.Get("tests.mcptest")
+	if err != nil {
+		t.Fatalf("Failed to get tests.mcptest assistant: %s", err.Error())
+	}
+
+	ctx := newTestContext("chat-test-mcp", "tests.mcptest")
+
+	t.Run("MCPToolsLoaded", func(t *testing.T) {
+		inputMessages := []context.Message{{Role: context.RoleUser, Content: "test mcp tools"}}
+
+		// Build LLM request
+		_, options, err := agent.BuildRequest(ctx, inputMessages, nil)
+		if err != nil {
+			t.Fatalf("Failed to build LLM request: %s", err.Error())
+		}
+
+		// Verify that tools are loaded
+		if options.Tools == nil {
+			t.Fatal("Expected tools to be loaded, got nil")
+		}
+
+		if len(options.Tools) == 0 {
+			t.Fatal("Expected at least some MCP tools, got empty list")
+		}
+
+		// Count MCP tools (should be filtered to only ping and echo)
+		mcpToolCount := 0
+		var toolNames []string
+		for _, toolMap := range options.Tools {
+			fn, ok := toolMap["function"].(map[string]interface{})
+			if !ok {
+				continue
+			}
+			name, ok := fn["name"].(string)
+			if ok {
+				toolNames = append(toolNames, name)
+				mcpToolCount++
+			}
+		}
+
+		t.Logf("Found %d MCP tools: %v", mcpToolCount, toolNames)
+
+		// Verify tool count (should be exactly 2: ping and echo)
+		if mcpToolCount != 2 {
+			t.Errorf("Expected 2 MCP tools (ping, echo), got %d: %v", mcpToolCount, toolNames)
+		}
+
+		// Verify specific tools exist
+		hasEchoPing := false
+		hasEchoEcho := false
+		for _, name := range toolNames {
+			if name == "echo__ping" {
+				hasEchoPing = true
+			}
+			if name == "echo__echo" {
+				hasEchoEcho = true
+			}
+		}
+
+		if !hasEchoPing {
+			t.Error("Expected 'echo__ping' tool to be present")
+		}
+		if !hasEchoEcho {
+			t.Error("Expected 'echo__echo' tool to be present")
+		}
+
+		// Verify that 'status' tool is NOT included (filtered out)
+		for _, name := range toolNames {
+			if name == "echo__status" {
+				t.Error("Tool 'echo__status' should be filtered out but was found")
+			}
+		}
+
+		t.Log("✓ MCP tools loaded and filtered correctly")
+	})
+
+	t.Run("MCPSamplesPrompt", func(t *testing.T) {
+		inputMessages := []context.Message{{Role: context.RoleUser, Content: "test mcp samples"}}
+
+		// Build LLM request
+		finalMessages, _, err := agent.BuildRequest(ctx, inputMessages, nil)
+		if err != nil {
+			t.Fatalf("Failed to build LLM request: %s", err.Error())
+		}
+
+		// Check if messages contain MCP samples prompt
+		// The samples prompt should be added as a system message
+		hasMCPSamples := false
+		for _, msg := range finalMessages {
+			if msg.Role == context.RoleSystem {
+				if content, ok := msg.Content.(string); ok {
+					if len(content) > 50 &&
+						(contains(content, "MCP Tool Usage Examples") ||
+							contains(content, "echo.ping") ||
+							contains(content, "echo.echo")) {
+						hasMCPSamples = true
+						t.Logf("Found MCP samples prompt (length: %d chars)", len(content))
+						break
+					}
+				}
+			}
+		}
+
+		// Note: samples may not exist for echo tools, so this is informational
+		if hasMCPSamples {
+			t.Log("✓ MCP samples prompt included in messages")
+		} else {
+			t.Log("ℹ No MCP samples prompt found (may not have sample files)")
+		}
+	})
+
+	t.Run("MCPToolNameFormat", func(t *testing.T) {
+		inputMessages := []context.Message{{Role: context.RoleUser, Content: "test tool format"}}
+
+		// Build LLM request
+		_, options, err := agent.BuildRequest(ctx, inputMessages, nil)
+		if err != nil {
+			t.Fatalf("Failed to build LLM request: %s", err.Error())
+		}
+
+		// Verify tool name format: server_id.tool_name
+		for _, toolMap := range options.Tools {
+			fn, ok := toolMap["function"].(map[string]interface{})
+			if !ok {
+				continue
+			}
+			name, ok := fn["name"].(string)
+			if ok {
+				// Parse tool name
+				serverID, toolName, ok := assistant.ParseMCPToolName(name)
+				if !ok {
+					t.Errorf("Tool name '%s' is not in correct format (server_id.tool_name)", name)
+					continue
+				}
+
+				// Verify server ID
+				if serverID != "echo" {
+					t.Errorf("Expected server_id 'echo', got '%s' for tool '%s'", serverID, name)
+				}
+
+				// Verify tool name is either ping or echo
+				if toolName != "ping" && toolName != "echo" {
+					t.Errorf("Expected tool name 'ping' or 'echo', got '%s'", toolName)
+				}
+
+				t.Logf("✓ Tool name format correct: %s → (%s, %s)", name, serverID, toolName)
+			}
+		}
+	})
+
+	t.Run("MCPToolSchema", func(t *testing.T) {
+		inputMessages := []context.Message{{Role: context.RoleUser, Content: "test tool schema"}}
+
+		// Build LLM request
+		_, options, err := agent.BuildRequest(ctx, inputMessages, nil)
+		if err != nil {
+			t.Fatalf("Failed to build LLM request: %s", err.Error())
+		}
+
+		// Verify tool schema structure
+		for _, toolMap := range options.Tools {
+			// Verify type field
+			if toolType, ok := toolMap["type"].(string); !ok || toolType != "function" {
+				t.Errorf("Expected tool type 'function', got: %v", toolMap["type"])
+			}
+
+			// Verify function field exists
+			fn, ok := toolMap["function"].(map[string]interface{})
+			if !ok {
+				t.Error("Tool missing 'function' field or wrong type")
+				continue
+			}
+
+			// Verify required fields
+			if _, hasName := fn["name"]; !hasName {
+				t.Error("Tool function missing 'name' field")
+			}
+			if _, hasDesc := fn["description"]; !hasDesc {
+				t.Error("Tool function missing 'description' field")
+			}
+			if _, hasParams := fn["parameters"]; !hasParams {
+				t.Error("Tool function missing 'parameters' field")
+			}
+
+			t.Logf("✓ Tool schema valid: %v", fn["name"])
+		}
+	})
+}
+
+// Helper function to check if string contains substring
+func contains(s, substr string) bool {
+	return len(s) >= len(substr) &&
+		(s == substr ||
+			len(s) > len(substr) &&
+				(s[:len(substr)] == substr ||
+					s[len(s)-len(substr):] == substr ||
+					findSubstring(s, substr)))
+}
+
+func findSubstring(s, substr string) bool {
+	for i := 0; i <= len(s)-len(substr); i++ {
+		if s[i:i+len(substr)] == substr {
+			return true
+		}
+	}
+	return false
+}

agent/assistant/history.go

@@ -0,0 +1,31 @@
+package assistant
+
+import (
+	"github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/trace/types"
+)
+
+// WithHistory merges the input messages with chat history and traces it
+// This method can be overridden or extended to implement actual history loading
+func (ast *Assistant) WithHistory(
+	ctx *context.Context,
+	inputMessages []context.Message,
+	agentNode types.Node,
+) ([]context.Message, error) {
+
+	// TODO: Implement actual history loading logic here
+	// For now, just simulate a check and return the input messages as is
+
+	// Simulate error check (this is where actual history loading would happen)
+	// if some_condition {
+	//     ast.traceAgentFail(agentNode, err)
+	//     return nil, err
+	// }
+
+	fullMessages := inputMessages
+
+	// Log the chat history
+	ast.traceAgentHistory(ctx, agentNode, fullMessages)
+
+	return fullMessages, nil
+}

agent/assistant/hook/mcp.go

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

agent/assistant/llm.go

@@ -0,0 +1,126 @@
+package assistant
+
+import (
+	"fmt"
+
+	"github.com/yaoapp/kun/log"
+	"github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/agent/llm"
+	"github.com/yaoapp/yao/agent/output/message"
+	"github.com/yaoapp/yao/trace/types"
+)
+
+// executeLLMStream executes the LLM streaming call with pre-built request
+// Returns completionResponse and error
+func (ast *Assistant) executeLLMStream(
+	ctx *context.Context,
+	completionMessages []context.Message,
+	completionOptions *context.CompletionOptions,
+	agentNode types.Node,
+	streamHandler message.StreamFunc,
+) (*context.CompletionResponse, error) {
+
+	// === Debug LLM Stream Start ===
+	fmt.Println(">>> executeLLMStream: STARTING")
+	fmt.Printf(">>> Messages count: %d\n", len(completionMessages))
+	fmt.Printf(">>> Tools count: %d\n", len(completionOptions.Tools))
+	// === End Debug ===
+
+	// Get connector object (capabilities were already set above, before stream_start)
+	conn, capabilities, err := ast.GetConnector(ctx)
+	if err != nil {
+		ast.traceAgentFail(agentNode, err)
+		return nil, err
+	}
+
+	// Set capabilities in options if not already set
+	if completionOptions.Capabilities == nil && capabilities != nil {
+		completionOptions.Capabilities = capabilities
+	}
+
+	// Log the capabilities
+	ast.traceConnectorCapabilities(agentNode, capabilities)
+
+	// Trace Add LLM request
+	ast.traceLLMRequest(ctx, conn.ID(), completionMessages, completionOptions)
+
+	// Create LLM instance with connector and options
+	llmInstance, err := llm.New(conn, completionOptions)
+	if err != nil {
+		return nil, err
+	}
+
+	// Call the LLM Completion Stream (streamHandler was set earlier)
+	log.Trace("[AGENT] Calling LLM Stream: assistant=%s", ast.ID)
+
+	// === Debug LLM Stream Call ===
+	fmt.Println(">>> executeLLMStream: CALLING llmInstance.Stream()")
+	// === End Debug ===
+
+	completionResponse, err := llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler)
+
+	// === Debug LLM Stream Return ===
+	fmt.Println(">>> executeLLMStream: llmInstance.Stream() RETURNED")
+	fmt.Printf(">>> err: %v\n", err)
+	if completionResponse != nil {
+		fmt.Printf(">>> ToolCalls: %d\n", len(completionResponse.ToolCalls))
+	}
+	// === End Debug ===
+
+	log.Trace("[AGENT] LLM Stream returned: assistant=%s, err=%v", ast.ID, err)
+	if err != nil {
+		log.Trace("[AGENT] Calling sendStreamEndOnError")
+		return nil, err
+	}
+
+	// Mark LLM Request Complete
+	ast.traceLLMComplete(ctx, completionResponse)
+
+	return completionResponse, nil
+}
+
+// executeLLMForToolRetry executes LLM call for tool retry with streaming output
+// This is used when retrying tool calls - we still want to show LLM's response to users
+// Returns completionResponse and error
+func (ast *Assistant) executeLLMForToolRetry(
+	ctx *context.Context,
+	completionMessages []context.Message,
+	completionOptions *context.CompletionOptions,
+	agentNode types.Node,
+	streamHandler message.StreamFunc,
+) (*context.CompletionResponse, error) {
+
+	// Get connector object
+	conn, capabilities, err := ast.GetConnector(ctx)
+	if err != nil {
+		ast.traceAgentFail(agentNode, err)
+		return nil, err
+	}
+
+	// Set capabilities in options if not already set
+	if completionOptions.Capabilities == nil && capabilities != nil {
+		completionOptions.Capabilities = capabilities
+	}
+
+	// Trace Add LLM retry request
+	ast.traceLLMRetryRequest(ctx, conn.ID(), completionMessages, completionOptions)
+
+	// Create LLM instance with connector and options
+	llmInstance, err := llm.New(conn, completionOptions)
+	if err != nil {
+		return nil, err
+	}
+
+	// Call the LLM Completion Stream (still streaming for tool retry)
+	log.Trace("[AGENT] Calling LLM Stream for tool retry: assistant=%s", ast.ID)
+	completionResponse, err := llmInstance.Stream(ctx, completionMessages, completionOptions, streamHandler)
+	log.Trace("[AGENT] LLM tool retry stream returned: assistant=%s, err=%v", ast.ID, err)
+	if err != nil {
+		return nil, err
+	}
+
+	// Mark LLM Request Complete
+	ast.traceLLMComplete(ctx, completionResponse)
+
+	return completionResponse, nil
+}

agent/assistant/mcp.go

@@ -0,0 +1,834 @@
+package assistant
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	jsoniter "github.com/json-iterator/go"
+	gouJson "github.com/yaoapp/gou/json"
+	"github.com/yaoapp/gou/mcp"
+	mcpTypes "github.com/yaoapp/gou/mcp/types"
+	"github.com/yaoapp/kun/log"
+	agentContext "github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/trace/types"
+)
+
+const (
+	// MaxMCPTools maximum number of MCP tools to include (to avoid overwhelming the LLM)
+	MaxMCPTools = 20
+)
+
+// MCPToolName formats a tool name with MCP server prefix
+// Format: server_id__tool_name (double underscore separator)
+// Dots in server_id are replaced with single underscores
+// Examples:
+//   - ("echo", "ping") → "echo__ping"
+//   - ("github.enterprise", "search") → "github_enterprise__search"
+//
+// Naming constraint: MCP server_id MUST NOT contain underscores (_)
+// Only dots (.), letters, numbers, and hyphens (-) are allowed in server_id
+func MCPToolName(serverID, toolName string) string {
+	if serverID == "" || toolName == "" {
+		return ""
+	}
+	// Replace dots with single underscores in server_id
+	cleanServerID := strings.ReplaceAll(serverID, ".", "_")
+	// Use double underscore as separator
+	return fmt.Sprintf("%s__%s", cleanServerID, toolName)
+}
+
+// ParseMCPToolName parses a formatted MCP tool name into server ID and tool name
+// Splits by double underscore (__), then restores dots in server_id
+// Examples:
+//   - "echo__ping" → ("echo", "ping")
+//   - "github_enterprise__search" → ("github.enterprise", "search")
+//
+// Returns (serverID, toolName, true) if valid format, ("", "", false) otherwise
+func ParseMCPToolName(formattedName string) (string, string, bool) {
+	if formattedName == "" {
+		return "", "", false
+	}
+
+	// Split by double underscore
+	parts := strings.Split(formattedName, "__")
+	if len(parts) != 2 {
+		return "", "", false
+	}
+
+	cleanServerID := parts[0]
+	toolName := parts[1]
+
+	// Validate that both parts are non-empty
+	if cleanServerID == "" || toolName == "" {
+		return "", "", false
+	}
+
+	// Restore dots in server_id (replace single underscores back to dots)
+	serverID := strings.ReplaceAll(cleanServerID, "_", ".")
+
+	return serverID, toolName, true
+}
+
+// buildMCPTools builds tool definitions and samples system prompt from MCP servers
+// Returns (tools, samplesPrompt, error)
+func (ast *Assistant) buildMCPTools(ctx *agentContext.Context) ([]MCPTool, string, error) {
+	if ast.MCP == nil || len(ast.MCP.Servers) == 0 {
+		return nil, "", nil
+	}
+
+	// Use the agent context for cancellation and timeout control
+	mcpCtx := ctx.Context
+	if mcpCtx == nil {
+		mcpCtx = context.Background()
+	}
+
+	allTools := make([]MCPTool, 0)
+	samplesBuilder := strings.Builder{}
+	hasSamples := false
+
+	// Process each MCP server in order
+	for _, serverConfig := range ast.MCP.Servers {
+		if len(allTools) >= MaxMCPTools {
+			log.Warn("[Assistant MCP] Reached maximum tool limit (%d), skipping remaining servers", MaxMCPTools)
+			break
+		}
+
+		// Get MCP client
+		client, err := mcp.Select(serverConfig.ServerID)
+		if err != nil {
+			log.Warn("[Assistant MCP] Failed to select MCP client '%s': %v", serverConfig.ServerID, err)
+			continue
+		}
+
+		// Get tools list (filter by serverConfig.Tools if specified)
+		toolsResponse, err := client.ListTools(mcpCtx, "")
+		if err != nil {
+			log.Warn("[Assistant MCP] Failed to list tools for '%s': %v", serverConfig.ServerID, err)
+			continue
+		}
+
+		// Build tool filter map if specified
+		toolFilter := make(map[string]bool)
+		if len(serverConfig.Tools) > 0 {
+			for _, toolName := range serverConfig.Tools {
+				toolFilter[toolName] = true
+			}
+		}
+
+		// Process each tool
+		for _, tool := range toolsResponse.Tools {
+			// Check tool limit
+			if len(allTools) >= MaxMCPTools {
+				break
+			}
+
+			// Apply tool filter if specified
+			if len(toolFilter) > 0 && !toolFilter[tool.Name] {
+				continue
+			}
+
+			// Format tool name with server prefix
+			formattedName := MCPToolName(serverConfig.ServerID, tool.Name)
+
+			// Convert MCP tool to MCPTool format
+			mcpTool := MCPTool{
+				Name:        formattedName,
+				Description: tool.Description,
+				Parameters:  tool.InputSchema,
+			}
+
+			allTools = append(allTools, mcpTool)
+
+			// Try to get samples for this tool
+			samples, err := client.ListSamples(mcpCtx, mcpTypes.SampleTool, tool.Name)
+			if err == nil && len(samples.Samples) > 0 {
+				if !hasSamples {
+					samplesBuilder.WriteString("\n\n## MCP Tool Usage Examples\n\n")
+					samplesBuilder.WriteString("The following examples demonstrate how to use MCP tools correctly:\n\n")
+					hasSamples = true
+				}
+
+				samplesBuilder.WriteString(fmt.Sprintf("### %s\n\n", formattedName))
+				if tool.Description != "" {
+					samplesBuilder.WriteString(fmt.Sprintf("**Description**: %s\n\n", tool.Description))
+				}
+
+				for i, sample := range samples.Samples {
+					if i >= 3 { // Limit to 3 examples per tool
+						break
+					}
+
+					samplesBuilder.WriteString(fmt.Sprintf("**Example %d", i+1))
+					if sample.Name != "" {
+						samplesBuilder.WriteString(fmt.Sprintf(" - %s", sample.Name))
+					}
+					samplesBuilder.WriteString("**:\n")
+
+					// Check metadata for description
+					if sample.Metadata != nil {
+						if desc, ok := sample.Metadata["description"].(string); ok && desc != "" {
+							samplesBuilder.WriteString(fmt.Sprintf("- Description: %s\n", desc))
+						}
+					}
+
+					if sample.Input != nil {
+						samplesBuilder.WriteString(fmt.Sprintf("- Input: `%v`\n", sample.Input))
+					}
+
+					if sample.Output != nil {
+						samplesBuilder.WriteString(fmt.Sprintf("- Output: `%v`\n", sample.Output))
+					}
+
+					samplesBuilder.WriteString("\n")
+				}
+			}
+		}
+
+		log.Trace("[Assistant MCP] Loaded %d tools from server '%s'", len(toolsResponse.Tools), serverConfig.ServerID)
+	}
+
+	samplesPrompt := ""
+	if hasSamples {
+		samplesPrompt = samplesBuilder.String()
+	}
+
+	log.Trace("[Assistant MCP] Total MCP tools loaded: %d", len(allTools))
+	return allTools, samplesPrompt, nil
+}
+
+// ToolCallResult represents the result of a tool call execution
+// executeToolCalls executes tool calls with intelligent strategy and trace logging:
+// - Single tool: use CallTool, single trace node
+// - Multiple tools: use CallToolsParallel with parallel trace nodes, fallback to sequential on certain errors
+// Returns (results, hasErrors)
+func (ast *Assistant) executeToolCalls(ctx *agentContext.Context, toolCalls []agentContext.ToolCall, attempt int) ([]ToolCallResult, bool) {
+	if len(toolCalls) == 0 {
+		return nil, false
+	}
+
+	// === Debug ===
+	fmt.Printf(">>> executeToolCalls: START (attempt %d, toolCalls count: %d)\n", attempt, len(toolCalls))
+	// === End Debug ===
+
+	log.Trace("[Assistant MCP] Executing %d tool calls (attempt %d)", len(toolCalls), attempt)
+
+	// Single tool call
+	if len(toolCalls) == 1 {
+		fmt.Println(">>> executeToolCalls: Calling executeSingleToolCall")
+		results, hasErrors := ast.executeSingleToolCall(ctx, toolCalls[0])
+		fmt.Printf(">>> executeToolCalls: executeSingleToolCall RETURNED (hasErrors: %v)\n", hasErrors)
+		return results, hasErrors
+	}
+
+	// Multiple tool calls - try parallel first
+	fmt.Println(">>> executeToolCalls: Calling executeMultipleToolCallsParallel")
+	results, hasErrors := ast.executeMultipleToolCallsParallel(ctx, toolCalls)
+	fmt.Printf(">>> executeToolCalls: executeMultipleToolCallsParallel RETURNED (hasErrors: %v)\n", hasErrors)
+	return results, hasErrors
+}
+
+// executeSingleToolCall executes a single tool call with trace logging
+func (ast *Assistant) executeSingleToolCall(ctx *agentContext.Context, toolCall agentContext.ToolCall) ([]ToolCallResult, bool) {
+	// === Debug ===
+	fmt.Printf(">>> executeSingleToolCall: START (tool: %s)\n", toolCall.Function.Name)
+	// === End Debug ===
+
+	trace, _ := ctx.Trace()
+
+	// Use the agent context for cancellation and timeout control
+	mcpCtx := ctx.Context
+	if mcpCtx == nil {
+		mcpCtx = context.Background()
+	}
+
+	result := ToolCallResult{
+		ToolCallID: toolCall.ID,
+		Name:       toolCall.Function.Name,
+	}
+
+	// Parse tool name
+	fmt.Println(">>> executeSingleToolCall: Parsing tool name")
+	serverID, toolName, ok := ParseMCPToolName(toolCall.Function.Name)
+	if !ok {
+		result.Error = fmt.Errorf("invalid MCP tool name format: %s", toolCall.Function.Name)
+		result.Content = result.Error.Error()
+		log.Error("[Assistant MCP] %v", result.Error)
+		return []ToolCallResult{result}, true
+	}
+
+	// Get MCP client
+	client, err := mcp.Select(serverID)
+	if err != nil {
+		result.Error = fmt.Errorf("failed to select MCP client '%s': %w", serverID, err)
+		result.Content = result.Error.Error()
+		result.IsRetryableError = false // MCP client selection error is not retryable
+		log.Error("[Assistant MCP] %v", result.Error)
+		return []ToolCallResult{result}, true
+	}
+
+	// Get tool info for description and schema
+	toolsResponse, err := client.ListTools(mcpCtx, "")
+	var toolDescription string
+	var toolSchema interface{}
+	if err == nil {
+		for _, t := range toolsResponse.Tools {
+			if t.Name == toolName {
+				toolDescription = t.Description
+				toolSchema = t.InputSchema
+				break
+			}
+		}
+	}
+	if toolDescription == "" {
+		toolDescription = fmt.Sprintf("MCP tool '%s'", toolName)
+	}
+
+	// Add trace node for this tool call
+	var toolNode types.Node
+	if trace != nil {
+		toolNode, _ = trace.Add(
+			map[string]any{
+				"tool_call_id": toolCall.ID,
+				"server":       serverID,
+				"tool":         toolName,
+				"arguments":    toolCall.Function.Arguments,
+			},
+			types.TraceNodeOption{
+				Label:       toolDescription,
+				Type:        "mcp_tool",
+				Icon:        "build",
+				Description: fmt.Sprintf("Calling '%s' on server '%s'", toolName, serverID),
+			},
+		)
+	}
+
+	// Parse arguments with repair support for better tolerance
+	var args map[string]interface{}
+	if toolCall.Function.Arguments != "" {
+		parsed, err := gouJson.Parse(toolCall.Function.Arguments)
+		if err != nil {
+			result.Error = fmt.Errorf("failed to parse arguments: %w", err)
+			result.Content = result.Error.Error()
+			result.IsRetryableError = true // Argument parsing error is retryable by LLM
+			log.Error("[Assistant MCP] %v", result.Error)
+			if toolNode != nil {
+				toolNode.Fail(result.Error)
+			}
+			return []ToolCallResult{result}, true
+		}
+
+		// Convert to map
+		if argsMap, ok := parsed.(map[string]interface{}); ok {
+			args = argsMap
+		} else {
+			result.Error = fmt.Errorf("arguments must be an object, got %T", parsed)
+			result.Content = result.Error.Error()
+			result.IsRetryableError = true // Type error is retryable by LLM
+			log.Error("[Assistant MCP] %v", result.Error)
+			if toolNode != nil {
+				toolNode.Fail(result.Error)
+			}
+			return []ToolCallResult{result}, true
+		}
+
+		// Validate arguments against tool schema if available
+		if toolSchema != nil {
+			fmt.Println(">>> executeSingleToolCall: Validating arguments against schema")
+			if err := gouJson.Validate(args, toolSchema); err != nil {
+				fmt.Printf(">>> executeSingleToolCall: Validation FAILED: %v\n", err)
+				result.Error = fmt.Errorf("argument validation failed: %w", err)
+				result.Content = result.Error.Error()
+				result.IsRetryableError = true // Validation error is retryable by LLM
+				log.Error("[Assistant MCP] %v", result.Error)
+				if toolNode != nil {
+					fmt.Println(">>> executeSingleToolCall: Failing toolNode due to validation error")
+					toolNode.Fail(result.Error)
+					fmt.Println(">>> executeSingleToolCall: toolNode.Fail() finished")
+				}
+				fmt.Println(">>> executeSingleToolCall: RETURNING with validation error")
+				return []ToolCallResult{result}, true
+			}
+			fmt.Println(">>> executeSingleToolCall: Validation PASSED")
+		}
+	}
+
+	// Call the tool
+	log.Trace("[Assistant MCP] Calling tool: %s (server: %s)", toolName, serverID)
+	fmt.Printf(">>> executeSingleToolCall: CALLING client.CallTool (tool: %s, server: %s)\n", toolName, serverID)
+	callResult, err := client.CallTool(mcpCtx, toolName, args)
+	fmt.Printf(">>> executeSingleToolCall: client.CallTool RETURNED (err: %v)\n", err)
+	if err != nil {
+		result.Error = fmt.Errorf("tool call failed: %w", err)
+		result.Content = result.Error.Error()
+		// Check if error is retryable (parameter/validation errors)
+		result.IsRetryableError = isRetryableToolError(err)
+		log.Error("[Assistant MCP] Tool call failed: %v (retryable: %v)", err, result.IsRetryableError)
+		if toolNode != nil {
+			toolNode.Fail(result.Error)
+		}
+		fmt.Println(">>> executeSingleToolCall: RETURNING with error")
+		return []ToolCallResult{result}, true
+	}
+
+	// Check if result is an error
+	if callResult.IsError {
+		result.Error = fmt.Errorf("MCP tool error")
+		result.IsRetryableError = false // MCP internal error is not retryable
+	}
+
+	// Serialize the Content field only ([]ToolContent)
+	contentBytes, err := jsoniter.Marshal(callResult.Content)
+	if err != nil {
+		result.Error = fmt.Errorf("failed to serialize result: %w", err)
+		result.Content = result.Error.Error()
+		result.IsRetryableError = false
+		log.Error("[Assistant MCP] %v", result.Error)
+		if toolNode != nil {
+			toolNode.Fail(result.Error)
+		}
+		return []ToolCallResult{result}, true
+	}
+
+	result.Content = string(contentBytes)
+	log.Trace("[Assistant MCP] Tool call succeeded: %s", toolName)
+
+	if toolNode != nil {
+		fmt.Println(">>> executeSingleToolCall: Completing toolNode")
+		toolNode.Complete(map[string]any{
+			"result": callResult,
+		})
+		fmt.Println(">>> executeSingleToolCall: toolNode.Complete() finished")
+	}
+
+	fmt.Println(">>> executeSingleToolCall: RETURNING success")
+	return []ToolCallResult{result}, false
+}
+
+// executeMultipleToolCallsParallel executes multiple tool calls in parallel with trace logging
+func (ast *Assistant) executeMultipleToolCallsParallel(ctx *agentContext.Context, toolCalls []agentContext.ToolCall) ([]ToolCallResult, bool) {
+	trace, _ := ctx.Trace()
+
+	// Use the agent context for cancellation and timeout control
+	mcpCtx := ctx.Context
+	if mcpCtx == nil {
+		mcpCtx = context.Background()
+	}
+
+	// Group tool calls by server
+	serverGroups := make(map[string][]agentContext.ToolCall)
+	for _, tc := range toolCalls {
+		serverID, _, ok := ParseMCPToolName(tc.Function.Name)
+		if !ok {
+			log.Warn("[Assistant MCP] Invalid tool name format: %s", tc.Function.Name)
+			continue
+		}
+		serverGroups[serverID] = append(serverGroups[serverID], tc)
+	}
+
+	results := make([]ToolCallResult, 0, len(toolCalls))
+	hasErrors := false
+
+	// Process each server's tools
+	for serverID, calls := range serverGroups {
+		client, err := mcp.Select(serverID)
+		if err != nil {
+			log.Error("[Assistant MCP] Failed to select MCP client '%s': %v", serverID, err)
+			// Add error results for all calls to this server
+			for _, tc := range calls {
+				results = append(results, ToolCallResult{
+					ToolCallID: tc.ID,
+					Name:       tc.Function.Name,
+					Content:    fmt.Sprintf("Failed to select MCP client: %v", err),
+					Error:      err,
+				})
+			}
+			hasErrors = true
+			continue
+		}
+
+		// Try parallel execution
+		serverResults, serverHasErrors := ast.executeServerToolsParallelWithTrace(
+			mcpCtx, trace, client, serverID, calls,
+		)
+
+		// If parallel execution failed with retryable error, try sequential
+		if serverHasErrors && ast.shouldRetrySequential(serverResults) {
+			log.Warn("[Assistant MCP] Parallel execution had parameter errors for server '%s', retrying sequentially", serverID)
+			serverResults, serverHasErrors = ast.executeServerToolsSequentialWithTrace(
+				mcpCtx, trace, client, serverID, calls,
+			)
+		}
+
+		results = append(results, serverResults...)
+		if serverHasErrors {
+			hasErrors = true
+		}
+	}
+
+	return results, hasErrors
+}
+
+// isRetryableToolError checks if an error is retryable by LLM (parameter/validation errors)
+// Returns true for errors that LLM can potentially fix by adjusting parameters
+// Returns false for MCP internal errors (network, auth, service unavailable, etc.)
+func isRetryableToolError(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	errMsg := strings.ToLower(err.Error())
+
+	// These are NOT retryable (MCP internal issues)
+	nonRetryablePatterns := []string{
+		"network",
+		"timeout",
+		"connection",
+		"unauthorized",
+		"forbidden",
+		"unavailable",
+		"failed to select",
+		"context canceled",
+		"context deadline",
+		"server error",
+		"internal error",
+	}
+
+	for _, pattern := range nonRetryablePatterns {
+		if strings.Contains(errMsg, pattern) {
+			return false
+		}
+	}
+
+	// These ARE retryable (parameter/validation issues LLM can fix)
+	retryablePatterns := []string{
+		"invalid",
+		"required",
+		"missing",
+		"validation",
+		"schema",
+		"type",
+		"format",
+		"parse",
+		"argument",
+		"parameter",
+	}
+
+	for _, pattern := range retryablePatterns {
+		if strings.Contains(errMsg, pattern) {
+			return true
+		}
+	}
+
+	// Default: assume it's retryable unless proven otherwise
+	// This allows LLM to attempt fixes for unknown error types
+	return true
+}
+
+// shouldRetrySequential checks if errors are retryable (parameter issues, not network/service issues)
+func (ast *Assistant) shouldRetrySequential(results []ToolCallResult) bool {
+	// Check if any result has a retryable error
+	hasRetryable := false
+	for _, result := range results {
+		if result.Error != nil && result.IsRetryableError {
+			hasRetryable = true
+			break
+		}
+	}
+	return hasRetryable
+}
+
+// executeServerToolsParallelWithTrace executes tools for a single server in parallel with trace
+func (ast *Assistant) executeServerToolsParallelWithTrace(mcpCtx context.Context, trace types.Manager, client mcp.Client, serverID string, toolCalls []agentContext.ToolCall) ([]ToolCallResult, bool) {
+	// Prepare parallel trace inputs
+	var parallelInputs []types.TraceParallelInput
+	mcpCalls := make([]mcpTypes.ToolCall, 0, len(toolCalls))
+	callMap := make(map[string]agentContext.ToolCall)
+
+	for _, tc := range toolCalls {
+		_, toolName, ok := ParseMCPToolName(tc.Function.Name)
+		if !ok {
+			continue
+		}
+
+		var args map[string]interface{}
+		if tc.Function.Arguments != "" {
+			if err := jsoniter.UnmarshalFromString(tc.Function.Arguments, &args); err != nil {
+				log.Error("[Assistant MCP] Failed to parse arguments for %s: %v", toolName, err)
+				continue
+			}
+		}
+
+		mcpCalls = append(mcpCalls, mcpTypes.ToolCall{
+			Name:      toolName,
+			Arguments: args,
+		})
+		callMap[toolName] = tc
+
+		// Add trace input for this tool
+		parallelInputs = append(parallelInputs, types.TraceParallelInput{
+			Input: map[string]any{
+				"tool_call_id": tc.ID,
+				"server":       serverID,
+				"tool":         toolName,
+				"arguments":    tc.Function.Arguments,
+			},
+			Option: types.TraceNodeOption{
+				Label:       fmt.Sprintf("Tool: %s", toolName),
+				Type:        "mcp_tool",
+				Icon:        "build",
+				Description: fmt.Sprintf("Calling MCP tool '%s' on server '%s'", toolName, serverID),
+			},
+		})
+	}
+
+	// Create parallel trace nodes
+	var toolNodes []types.Node
+	if trace != nil && len(parallelInputs) > 0 {
+		fmt.Printf(">>> executeServerToolsParallelWithTrace: Creating %d parallel trace nodes\n", len(parallelInputs))
+		var err error
+		toolNodes, err = trace.Parallel(parallelInputs)
+		if err != nil {
+			fmt.Printf(">>> executeServerToolsParallelWithTrace: trace.Parallel() FAILED: %v\n", err)
+		} else {
+			fmt.Printf(">>> executeServerToolsParallelWithTrace: Created %d trace nodes\n", len(toolNodes))
+		}
+	} else {
+		fmt.Printf(">>> executeServerToolsParallelWithTrace: NOT creating trace nodes (trace: %v, inputs: %d)\n", trace != nil, len(parallelInputs))
+	}
+
+	// Call tools in parallel
+	log.Trace("[Assistant MCP] Calling %d tools in parallel on server '%s'", len(mcpCalls), serverID)
+	mcpResponse, err := client.CallToolsParallel(mcpCtx, mcpCalls)
+	if err != nil {
+		log.Error("[Assistant MCP] Parallel call failed: %v", err)
+		// Mark all trace nodes as failed
+		for _, node := range toolNodes {
+			if node != nil {
+				node.Fail(err)
+			}
+		}
+		return nil, true
+	}
+
+	// Process results
+	results := make([]ToolCallResult, 0, len(mcpResponse.Results))
+	hasErrors := false
+
+	for i, mcpResult := range mcpResponse.Results {
+		toolName := mcpCalls[i].Name
+		originalCall := callMap[toolName]
+		var toolNode types.Node
+		if i < len(toolNodes) {
+			toolNode = toolNodes[i]
+		}
+
+		result := ToolCallResult{
+			ToolCallID: originalCall.ID,
+			Name:       originalCall.Function.Name,
+		}
+
+		// Serialize content
+		contentBytes, err := jsoniter.Marshal(mcpResult.Content)
+		if err != nil {
+			result.Error = fmt.Errorf("failed to serialize result: %w", err)
+			result.Content = result.Error.Error()
+			result.IsRetryableError = false // Serialization error is not retryable
+			hasErrors = true
+			if toolNode != nil {
+				toolNode.Fail(result.Error)
+			}
+		} else {
+			result.Content = string(contentBytes)
+
+			// Check if it's an error result
+			if mcpResult.IsError {
+				result.Error = fmt.Errorf("tool call error: %s", result.Content)
+				result.IsRetryableError = isRetryableToolError(result.Error)
+				hasErrors = true
+				log.Error("[Assistant MCP] Tool call failed: %s - %s (retryable: %v)", toolName, result.Content, result.IsRetryableError)
+				if toolNode != nil {
+					toolNode.Fail(result.Error)
+				}
+			} else {
+				// Success
+				if toolNode != nil {
+					fmt.Printf(">>> executeServerToolsParallelWithTrace: Completing toolNode %d\n", i)
+					toolNode.Complete(map[string]any{
+						"result": mcpResult.Content,
+					})
+					fmt.Printf(">>> executeServerToolsParallelWithTrace: toolNode %d completed\n", i)
+				} else {
+					fmt.Printf(">>> executeServerToolsParallelWithTrace: toolNode %d is nil!\n", i)
+				}
+			}
+		}
+
+		results = append(results, result)
+	}
+
+	return results, hasErrors
+}
+
+// executeServerToolsSequentialWithTrace executes tools for a single server sequentially with trace
+func (ast *Assistant) executeServerToolsSequentialWithTrace(mcpCtx context.Context, trace types.Manager, client mcp.Client, serverID string, toolCalls []agentContext.ToolCall) ([]ToolCallResult, bool) {
+	results := make([]ToolCallResult, 0, len(toolCalls))
+	hasErrors := false
+
+	log.Trace("[Assistant MCP] Calling %d tools sequentially on server '%s'", len(toolCalls), serverID)
+
+	for _, tc := range toolCalls {
+		_, toolName, ok := ParseMCPToolName(tc.Function.Name)
+		if !ok {
+			results = append(results, ToolCallResult{
+				ToolCallID: tc.ID,
+				Name:       tc.Function.Name,
+				Content:    fmt.Sprintf("Invalid tool name format: %s", tc.Function.Name),
+				Error:      fmt.Errorf("invalid tool name format"),
+			})
+			hasErrors = true
+			continue
+		}
+
+		// Get tool schema for validation
+		toolsResponse, err := client.ListTools(mcpCtx, "")
+		var toolSchema interface{}
+		if err == nil {
+			for _, t := range toolsResponse.Tools {
+				if t.Name == toolName {
+					toolSchema = t.InputSchema
+					break
+				}
+			}
+		}
+
+		// Add trace node for this tool call
+		var toolNode types.Node
+		if trace != nil {
+			toolNode, _ = trace.Add(
+				map[string]any{
+					"tool_call_id": tc.ID,
+					"server":       serverID,
+					"tool":         toolName,
+					"arguments":    tc.Function.Arguments,
+				},
+				types.TraceNodeOption{
+					Label:       fmt.Sprintf("Tool: %s (sequential retry)", toolName),
+					Type:        "mcp_tool",
+					Icon:        "build",
+					Description: fmt.Sprintf("Retrying MCP tool '%s' on server '%s' sequentially", toolName, serverID),
+				},
+			)
+		}
+
+		// Parse arguments with repair support
+		var args map[string]interface{}
+		if tc.Function.Arguments != "" {
+			parsed, err := gouJson.Parse(tc.Function.Arguments)
+			if err != nil {
+				result := ToolCallResult{
+					ToolCallID:       tc.ID,
+					Name:             tc.Function.Name,
+					Content:          fmt.Sprintf("Failed to parse arguments: %v", err),
+					Error:            err,
+					IsRetryableError: true, // Parsing error is retryable
+				}
+				results = append(results, result)
+				hasErrors = true
+				if toolNode != nil {
+					toolNode.Fail(err)
+				}
+				continue
+			}
+
+			// Convert to map
+			if argsMap, ok := parsed.(map[string]interface{}); ok {
+				args = argsMap
+			} else {
+				err := fmt.Errorf("arguments must be an object, got %T", parsed)
+				result := ToolCallResult{
+					ToolCallID:       tc.ID,
+					Name:             tc.Function.Name,
+					Content:          err.Error(),
+					Error:            err,
+					IsRetryableError: true, // Type error is retryable
+				}
+				results = append(results, result)
+				hasErrors = true
+				if toolNode != nil {
+					toolNode.Fail(err)
+				}
+				continue
+			}
+
+			// Validate arguments against tool schema if available
+			if toolSchema != nil {
+				if err := gouJson.Validate(args, toolSchema); err != nil {
+					result := ToolCallResult{
+						ToolCallID:       tc.ID,
+						Name:             tc.Function.Name,
+						Content:          fmt.Sprintf("Argument validation failed: %v", err),
+						Error:            err,
+						IsRetryableError: true, // Validation error is retryable
+					}
+					results = append(results, result)
+					hasErrors = true
+					if toolNode != nil {
+						toolNode.Fail(err)
+					}
+					continue
+				}
+			}
+		}
+
+		// Call single tool
+		log.Trace("[Assistant MCP] Calling tool: %s", toolName)
+		mcpResult, err := client.CallTool(mcpCtx, toolName, args)
+
+		result := ToolCallResult{
+			ToolCallID: tc.ID,
+			Name:       tc.Function.Name,
+		}
+
+		if err != nil {
+			result.Error = err
+			result.Content = fmt.Sprintf("Tool call failed: %v", err)
+			result.IsRetryableError = isRetryableToolError(err)
+			hasErrors = true
+			log.Error("[Assistant MCP] Tool call failed: %s - %v (retryable: %v)", toolName, err, result.IsRetryableError)
+			if toolNode != nil {
+				toolNode.Fail(err)
+			}
+		} else {
+			// Check if result is an error
+			if mcpResult.IsError {
+				result.Error = fmt.Errorf("MCP tool error")
+				result.IsRetryableError = false // MCP internal error is not retryable
+				hasErrors = true
+			}
+
+			// Serialize the Content field only ([]ToolContent)
+			contentBytes, err := jsoniter.Marshal(mcpResult.Content)
+			if err != nil {
+				result.Error = err
+				result.Content = fmt.Sprintf("Failed to serialize result: %v", err)
+				result.IsRetryableError = false // Serialization error is not retryable
+				hasErrors = true
+				if toolNode != nil {
+					toolNode.Fail(err)
+				}
+			} else {
+				result.Content = string(contentBytes)
+				if toolNode != nil {
+					toolNode.Complete(map[string]any{
+						"result": mcpResult.Content,
+					})
+				}
+			}
+		}
+
+		results = append(results, result)
+	}
+
+	return results, hasErrors
+}

agent/assistant/mcp_test.go

@@ -0,0 +1,236 @@
+package assistant_test
+
+import (
+	"testing"
+
+	"github.com/yaoapp/yao/agent/assistant"
+	"github.com/yaoapp/yao/agent/testutils"
+)
+
+func TestMCPToolName(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	tests := []struct {
+		name       string
+		serverID   string
+		toolName   string
+		wantResult string
+	}{
+		{
+			name:       "Simple tool name",
+			serverID:   "github",
+			toolName:   "search_repos",
+			wantResult: "github__search_repos",
+		},
+		{
+			name:       "Server with dots",
+			serverID:   "github.enterprise",
+			toolName:   "search_repos",
+			wantResult: "github_enterprise__search_repos",
+		},
+		{
+			name:       "Tool with underscores",
+			serverID:   "customer-db",
+			toolName:   "create_customer",
+			wantResult: "customer-db__create_customer",
+		},
+		{
+			name:       "Complex server with multiple dots",
+			serverID:   "com.example.mcp",
+			toolName:   "tool_name",
+			wantResult: "com_example_mcp__tool_name",
+		},
+		{
+			name:       "Empty server ID",
+			serverID:   "",
+			toolName:   "tool",
+			wantResult: "",
+		},
+		{
+			name:       "Empty tool name",
+			serverID:   "server",
+			toolName:   "",
+			wantResult: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := assistant.MCPToolName(tt.serverID, tt.toolName)
+			if result != tt.wantResult {
+				t.Errorf("MCPToolName() = %v, want %v", result, tt.wantResult)
+			}
+		})
+	}
+}
+
+func TestParseMCPToolName(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	tests := []struct {
+		name          string
+		formattedName string
+		wantServerID  string
+		wantToolName  string
+		wantOK        bool
+	}{
+		{
+			name:          "Valid simple format",
+			formattedName: "github__search_repos",
+			wantServerID:  "github",
+			wantToolName:  "search_repos",
+			wantOK:        true,
+		},
+		{
+			name:          "Server with dots restored",
+			formattedName: "github_enterprise__search_repos",
+			wantServerID:  "github.enterprise",
+			wantToolName:  "search_repos",
+			wantOK:        true,
+		},
+		{
+			name:          "Complex server ID with multiple dots",
+			formattedName: "com_example_mcp_server__tool_name",
+			wantServerID:  "com.example.mcp.server",
+			wantToolName:  "tool_name",
+			wantOK:        true,
+		},
+		{
+			name:          "Tool name with underscores",
+			formattedName: "server__create_new_user",
+			wantServerID:  "server",
+			wantToolName:  "create_new_user",
+			wantOK:        true,
+		},
+		{
+			name:          "Server with hyphens",
+			formattedName: "mcp-server__tool",
+			wantServerID:  "mcp-server",
+			wantToolName:  "tool",
+			wantOK:        true,
+		},
+		{
+			name:          "Invalid format - no double underscore",
+			formattedName: "invalid",
+			wantServerID:  "",
+			wantToolName:  "",
+			wantOK:        false,
+		},
+		{
+			name:          "Invalid format - empty string",
+			formattedName: "",
+			wantServerID:  "",
+			wantToolName:  "",
+			wantOK:        false,
+		},
+		{
+			name:          "Invalid format - only double underscore",
+			formattedName: "__",
+			wantServerID:  "",
+			wantToolName:  "",
+			wantOK:        false,
+		},
+		{
+			name:          "Invalid format - ends with double underscore",
+			formattedName: "server__",
+			wantServerID:  "",
+			wantToolName:  "",
+			wantOK:        false,
+		},
+		{
+			name:          "Invalid format - starts with double underscore",
+			formattedName: "__tool",
+			wantServerID:  "",
+			wantToolName:  "",
+			wantOK:        false,
+		},
+		{
+			name:          "Invalid format - multiple double underscores",
+			formattedName: "server__middle__tool",
+			wantServerID:  "",
+			wantToolName:  "",
+			wantOK:        false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			serverID, toolName, ok := assistant.ParseMCPToolName(tt.formattedName)
+			if serverID != tt.wantServerID {
+				t.Errorf("ParseMCPToolName() serverID = %v, want %v", serverID, tt.wantServerID)
+			}
+			if toolName != tt.wantToolName {
+				t.Errorf("ParseMCPToolName() toolName = %v, want %v", toolName, tt.wantToolName)
+			}
+			if ok != tt.wantOK {
+				t.Errorf("ParseMCPToolName() ok = %v, want %v", ok, tt.wantOK)
+			}
+		})
+	}
+}
+
+func TestMCPToolName_RoundTrip(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	tests := []struct {
+		name     string
+		serverID string
+		toolName string
+	}{
+		{
+			name:     "Simple IDs",
+			serverID: "github",
+			toolName: "search_repos",
+		},
+		{
+			name:     "Server with dots",
+			serverID: "github.enterprise",
+			toolName: "search",
+		},
+		{
+			name:     "Complex server ID",
+			serverID: "com.example.mcp.server",
+			toolName: "tool_name",
+		},
+		{
+			name:     "Server with dashes",
+			serverID: "mcp-server-123",
+			toolName: "tool_with_underscores",
+		},
+		{
+			name:     "Mixed dots and dashes",
+			serverID: "github.enterprise-prod",
+			toolName: "api_call",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Format
+			formatted := assistant.MCPToolName(tt.serverID, tt.toolName)
+			if formatted == "" {
+				t.Fatal("MCPToolName() returned empty string")
+			}
+
+			// Parse
+			serverID, toolName, ok := assistant.ParseMCPToolName(formatted)
+
+			// Verify round-trip
+			if !ok {
+				t.Fatal("ParseMCPToolName() failed")
+			}
+			if serverID != tt.serverID {
+				t.Errorf("Round-trip failed: serverID = %v, want %v", serverID, tt.serverID)
+			}
+			if toolName != tt.toolName {
+				t.Errorf("Round-trip failed: toolName = %v, want %v", toolName, tt.toolName)
+			}
+
+			t.Logf("✓ Round-trip successful: (%s, %s) → %s → (%s, %s)",
+				tt.serverID, tt.toolName, formatted, serverID, toolName)
+		})
+	}
+}

agent/assistant/trace.go

@@ -0,0 +1,126 @@
+package assistant
+
+import (
+	"fmt"
+
+	"github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/agent/i18n"
+	"github.com/yaoapp/yao/trace/types"
+)
+
+// initAgentTraceNode creates and returns the agent trace node
+func (ast *Assistant) initAgentTraceNode(ctx *context.Context, inputMessages []context.Message) types.Node {
+	trace, _ := ctx.Trace()
+	if trace == nil {
+		return nil
+	}
+
+	agentNode, _ := trace.Add(inputMessages, types.TraceNodeOption{
+		Label:       i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.label"), // "Assistant {{name}}"
+		Type:        "agent",
+		Icon:        "assistant",
+		Description: i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.description"), // "Assistant {{name}} is processing the request"
+	})
+
+	return agentNode
+}
+
+// traceAgentHistory logs the chat history to the agent trace node
+func (ast *Assistant) traceAgentHistory(ctx *context.Context, agentNode types.Node, fullMessages []context.Message) {
+	if agentNode == nil {
+		return
+	}
+
+	agentNode.Info(
+		i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.stream.history"), // "Get Chat History"
+		map[string]any{"messages": fullMessages},
+	)
+}
+
+// traceCreateHook logs the create hook response to the agent trace node
+func (ast *Assistant) traceCreateHook(agentNode types.Node, createResponse *context.HookCreateResponse) {
+	if agentNode == nil {
+		return
+	}
+
+	agentNode.Debug("Call Create Hook", map[string]any{"response": createResponse})
+}
+
+// traceConnectorCapabilities logs the connector capabilities to the agent trace node
+func (ast *Assistant) traceConnectorCapabilities(agentNode types.Node, capabilities *context.ModelCapabilities) {
+	if agentNode == nil {
+		return
+	}
+
+	agentNode.Debug("Get Connector Capabilities", map[string]any{"capabilities": capabilities})
+}
+
+// traceLLMRequest adds a LLM trace node to the trace
+func (ast *Assistant) traceLLMRequest(ctx *context.Context, connID string, completionMessages []context.Message, completionOptions *context.CompletionOptions) {
+	trace, _ := ctx.Trace()
+	if trace == nil {
+		return
+	}
+
+	trace.Add(
+		map[string]any{"messages": completionMessages, "options": completionOptions},
+		types.TraceNodeOption{
+			Label:       fmt.Sprintf(i18n.Tr(ast.ID, ctx.Locale, "llm.openai.stream.label"), connID), // "LLM %s"
+			Type:        "llm",
+			Icon:        "psychology",
+			Description: fmt.Sprintf(i18n.Tr(ast.ID, ctx.Locale, "llm.openai.stream.description"), connID), // "LLM %s is processing the request"
+		},
+	)
+}
+
+// traceLLMComplete marks the LLM request as complete in the trace
+func (ast *Assistant) traceLLMComplete(ctx *context.Context, completionResponse *context.CompletionResponse) {
+	trace, _ := ctx.Trace()
+	if trace == nil {
+		return
+	}
+
+	trace.Complete(completionResponse)
+}
+
+// traceAgentOutput sets the output of the agent trace node
+func (ast *Assistant) traceAgentOutput(agentNode types.Node, createResponse *context.HookCreateResponse, doneResponse *context.ResponseHookDone, completionResponse *context.CompletionResponse) {
+	if agentNode == nil {
+		return
+	}
+
+	output := context.Response{
+		Create:     createResponse,
+		Done:       doneResponse,
+		Completion: completionResponse,
+	}
+
+	agentNode.Complete(output)
+}
+
+// traceAgentFail marks the agent trace node as failed
+func (ast *Assistant) traceAgentFail(agentNode types.Node, err error) {
+	if agentNode == nil {
+		return
+	}
+
+	agentNode.Fail(err)
+}
+
+// traceLLMRetryRequest adds a LLM retry trace node to the trace
+func (ast *Assistant) traceLLMRetryRequest(ctx *context.Context, connID string, completionMessages []context.Message, completionOptions *context.CompletionOptions) {
+	trace, _ := ctx.Trace()
+	if trace == nil {
+		return
+	}
+
+	trace.Add(
+		map[string]any{"messages": completionMessages, "options": completionOptions},
+		types.TraceNodeOption{
+			Label:       fmt.Sprintf("LLM %s (Tool Retry)", connID),
+			Type:        "llm_retry",
+			Icon:        "refresh",
+			Description: fmt.Sprintf("LLM %s is retrying with tool call error feedback", connID),
+		},
+	)
+}

agent/assistant/types.go

@@ -5,6 +5,7 @@ import (
 	"io"
 
 	"github.com/gin-gonic/gin"
+	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/yao/agent/assistant/hook"
 	chatctx "github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/message"
@@ -164,3 +165,99 @@ type FileResponse struct {
 	ContentType string
 	Extension   string
 }
+
+// MCPTool represents a simplified MCP tool for building LLM requests
+// This is an internal representation used when collecting tools from MCP servers
+// and preparing them for the LLM's tool calling interface
+type MCPTool struct {
+	Name        string      // Formatted tool name with server prefix (e.g., "server_id__tool_name")
+	Description string      // Tool description from MCP server
+	Parameters  interface{} // JSON Schema for tool parameters (from MCP InputSchema)
+}
+
+// ToolCallResult represents the result of a tool call execution
+// Used to track the outcome of MCP tool invocations during agent execution
+type ToolCallResult struct {
+	ToolCallID       string // Tool call ID from the LLM (matches the ID in the LLM's tool_calls response)
+	Name             string // Tool name (formatted with server prefix, e.g., "server_id__tool_name")
+	Content          string // Result content (JSON string of the tool's output or error message)
+	Error            error  // Error if the call failed (nil if successful)
+	IsRetryableError bool   // Whether the error should be sent to LLM for retry
+	// true: parameter/validation errors that LLM can fix (e.g., "missing required field")
+	// false: MCP internal errors that LLM cannot fix (e.g., "network error", "service unavailable")
+}
+
+// Server extracts the MCP server ID from the formatted tool name
+// Example: "echo__ping" -> "echo"
+func (r *ToolCallResult) Server() string {
+	serverID, _, _ := ParseMCPToolName(r.Name)
+	return serverID
+}
+
+// Tool extracts the original tool name without server prefix
+// Example: "echo__ping" -> "ping"
+func (r *ToolCallResult) Tool() string {
+	_, toolName, _ := ParseMCPToolName(r.Name)
+	return toolName
+}
+
+// ParsedContent extracts the actual tool return value from MCP ToolContent array
+// According to MCP protocol:
+// - Content is []ToolContent array
+// - For "text" type, the actual value is in Text field (usually JSON string)
+// - For "image" type, returns the Data field
+// - For "resource" type, returns the Resource object
+// If there are multiple content items, returns an array of parsed values
+func (r *ToolCallResult) ParsedContent() (interface{}, error) {
+	if r.Content == "" {
+		return nil, nil
+	}
+
+	// Parse Content as []ToolContent
+	var toolContents []map[string]interface{}
+	if err := jsoniter.UnmarshalFromString(r.Content, &toolContents); err != nil {
+		// If parsing fails, return the string content directly (error message)
+		return r.Content, nil
+	}
+
+	// Extract actual values from ToolContent items
+	var results []interface{}
+	for _, tc := range toolContents {
+		contentType, _ := tc["type"].(string)
+
+		switch contentType {
+		case "text":
+			// For text type, parse the Text field (usually JSON)
+			if textStr, ok := tc["text"].(string); ok {
+				// Try to parse as JSON
+				var parsed interface{}
+				if err := jsoniter.UnmarshalFromString(textStr, &parsed); err == nil {
+					results = append(results, parsed)
+				} else {
+					// If not JSON, return as plain string
+					results = append(results, textStr)
+				}
+			}
+		case "image":
+			// For image type, return the data and mimeType
+			results = append(results, map[string]interface{}{
+				"type":     "image",
+				"data":     tc["data"],
+				"mimeType": tc["mimeType"],
+			})
+		case "resource":
+			// For resource type, return the resource object
+			results = append(results, tc["resource"])
+		default:
+			// Unknown type, return as-is
+			results = append(results, tc)
+		}
+	}
+
+	// If only one result, return it directly (not as array)
+	if len(results) == 1 {
+		return results[0], nil
+	}
+
+	return results, nil
+}

agent/i18n/builtin.go

@@ -12,8 +12,8 @@ func init() {
 		Locale: "en",
 		Messages: map[string]any{
 			// Assistant: agent.go Stream() function
-			"assistant.agent.stream.label":        "Assistant {{name}}",
+			"assistant.agent.stream.label":        "{{name}}",
-			"assistant.agent.stream.description":  "Assistant {{name}} is processing the request",
+			"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.label":        "{{name}}",
-			"assistant.agent.stream.description":  "助手 {{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.label":        "{{name}}",
-			"assistant.agent.stream.description":  "助手 {{name}} 正在处理请求",
+			"assistant.agent.stream.description":  "{{name}} 正在处理请求",
 			"assistant.agent.stream.history":      "获取聊天历史",
 			"assistant.agent.stream.capabilities": "获取连接器能力",
 			"assistant.agent.stream.create_hook":  "调用 Create Hook",

agent/i18n/i18n_test.go

@@ -961,14 +961,15 @@ func TestBuiltinMessages(t *testing.T) {
 
 	t.Run("English built-in messages", func(t *testing.T) {
 		// Test assistant messages
+		// Updated: label now only shows {{name}} without "Assistant" prefix
 		result := TranslateGlobal("en", "{{assistant.agent.stream.label}}")
-		expected := "Assistant {{name}}"
+		expected := "{{name}}"
 		if result != expected {
 			t.Errorf("Expected '%s', got '%v'", expected, result)
 		}
 
 		result = TranslateGlobal("en", "{{assistant.agent.stream.description}}")
-		expected = "Assistant {{name}} is processing the request"
+		expected = "{{name}} is processing the request"
 		if result != expected {
 			t.Errorf("Expected '%s', got '%v'", expected, result)
 		}
@@ -1002,14 +1003,15 @@ func TestBuiltinMessages(t *testing.T) {
 
 	t.Run("Chinese (zh-cn) built-in messages", func(t *testing.T) {
 		// Test assistant messages
+		// Updated: label now only shows {{name}} without "助手" prefix
 		result := TranslateGlobal("zh-cn", "{{assistant.agent.stream.label}}")
-		expected := "助手 {{name}}"
+		expected := "{{name}}"
 		if result != expected {
 			t.Errorf("Expected '%s', got '%v'", expected, result)
 		}
 
 		result = TranslateGlobal("zh-cn", "{{assistant.agent.stream.description}}")
-		expected = "助手 {{name}} 正在处理请求"
+		expected = "{{name}} 正在处理请求"
 		if result != expected {
 			t.Errorf("Expected '%s', got '%v'", expected, result)
 		}
@@ -1036,8 +1038,9 @@ func TestBuiltinMessages(t *testing.T) {
 	})
 
 	t.Run("Chinese (zh) short code", func(t *testing.T) {
+		// Updated: label now only shows {{name}} without "助手" prefix
 		result := TranslateGlobal("zh", "{{assistant.agent.stream.label}}")
-		expected := "助手 {{name}}"
+		expected := "{{name}}"
 		if result != expected {
 			t.Errorf("Expected '%s', got '%v'", expected, result)
 		}
@@ -1244,7 +1247,8 @@ func TestTAlias(t *testing.T) {
 			t.Errorf("T and TranslateGlobal should return same result. T: %v, TranslateGlobal: %v", resultT, resultGlobal)
 		}
 
-		expected := "Assistant {{name}}"
+		// Updated: label now only shows {{name}} without "Assistant" prefix
+		expected := "{{name}}"
 		if resultT != expected {
 			t.Errorf("Expected '%s', got '%v'", expected, resultT)
 		}
@@ -1267,17 +1271,18 @@ func TestTAlias(t *testing.T) {
 	})
 
 	t.Run("T with nested template (template in template value)", func(t *testing.T) {
-		// assistant.agent.stream.label = "Assistant {{name}}" (contains {{name}} template)
+		// assistant.agent.stream.label = "{{name}}" (contains {{name}} template)
+		// Updated: label now only shows {{name}} without prefix
 		// This tests if we can get the template string itself
 		result := T("en", "{{assistant.agent.stream.label}}")
-		expected := "Assistant {{name}}"
+		expected := "{{name}}"
 		if result != expected {
 			t.Errorf("Expected '%s', got '%v'", expected, result)
 		}
 
 		// Verify Chinese version too
 		resultZh := T("zh-cn", "{{assistant.agent.stream.label}}")
-		expectedZh := "助手 {{name}}"
+		expectedZh := "{{name}}"
 		if resultZh != expectedZh {
 			t.Errorf("Expected '%s', got '%v'", expectedZh, resultZh)
 		}

agent/llm/providers/openai/openai.go

@@ -213,7 +213,6 @@ func (p *Provider) Stream(ctx *context.Context, messages []context.Message, opti
 	}
 
 	maxRetries := 3
-	maxValidationRetries := 3
 	var lastErr error
 
 	// Get Go context for cancellation support
@@ -305,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{
+				trace.Debug("Tool call validation error (unexpected, should be handled differently)", map[string]any{
-						"attempt":     validationAttempt + 1,
-						"max_retries": maxValidationRetries,
 					"error": err.Error(),
 				})
 			}
-
-				// Add error feedback to conversation history
-				validationRetryMessages = append(validationRetryMessages, context.Message{
-					Role:    context.RoleSystem,
-					Content: fmt.Sprintf("Tool call validation error: %v. Please correct the tool call arguments to match the required schema.", err),
-				})
-
-				// 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
+			// Continue and return response (don't return error)
-			return nil, fmt.Errorf("tool call validation failed: %w", err)
+			// 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{
+				trace.Debug("Tool call validation error in Post (unexpected, should be handled differently)", map[string]any{
-						"attempt":     validationAttempt + 1,
-						"max_retries": maxValidationRetries,
 					"error": err.Error(),
 				})
 			}
-
-				// Add error feedback to conversation history
-				validationRetryMessages = append(validationRetryMessages, context.Message{
-					Role:    context.RoleSystem,
-					Content: fmt.Sprintf("Tool call validation error: %v. Please correct the tool call arguments to match the required schema.", err),
-				})
-
-				// 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.)

agent/store/types/convert.go

@@ -63,17 +63,9 @@ func ToMCPServers(v interface{}) (*MCPServers, error) {
 	case MCPServers:
 		return &mcp, nil
 
-	case []string:
-		return &MCPServers{Servers: mcp}, nil
-
-	case []interface{}:
-		var servers []string
-		for _, item := range mcp {
-			servers = append(servers, cast.ToString(item))
-		}
-		return &MCPServers{Servers: servers}, nil
-
 	default:
+		// For any type (including []string, []interface{}, map[string]interface{}),
+		// marshal and unmarshal to MCPServers using custom UnmarshalJSON
 		raw, err := jsoniter.Marshal(mcp)
 		if err != nil {
 			return nil, fmt.Errorf("mcp format error: %s", err.Error())

agent/store/types/convert_test.go

@@ -118,7 +118,7 @@ func TestToMCPServers(t *testing.T) {
 	})
 
 	t.Run("MCPServersPointer", func(t *testing.T) {
-		mcp := &MCPServers{Servers: []string{"server1", "server2"}}
+		mcp := &MCPServers{Servers: []MCPServerConfig{{ServerID: "server1"}, {ServerID: "server2"}}}
 		result, err := ToMCPServers(mcp)
 		if err != nil {
 			t.Errorf("Expected no error, got: %v", err)
@@ -129,7 +129,7 @@ func TestToMCPServers(t *testing.T) {
 	})
 
 	t.Run("MCPServersValue", func(t *testing.T) {
-		mcp := MCPServers{Servers: []string{"server1", "server2"}}
+		mcp := MCPServers{Servers: []MCPServerConfig{{ServerID: "server1"}, {ServerID: "server2"}}}
 		result, err := ToMCPServers(mcp)
 		if err != nil {
 			t.Errorf("Expected no error, got: %v", err)
@@ -139,37 +139,9 @@ func TestToMCPServers(t *testing.T) {
 		}
 	})
 
-	t.Run("StringSlice", func(t *testing.T) {
-		servers := []string{"server1", "server2", "server3"}
-		result, err := ToMCPServers(servers)
-		if err != nil {
-			t.Errorf("Expected no error, got: %v", err)
-		}
-		if len(result.Servers) != 3 {
-			t.Errorf("Expected 3 servers, got %d", len(result.Servers))
-		}
-		if result.Servers[0] != "server1" {
-			t.Errorf("Expected 'server1', got '%s'", result.Servers[0])
-		}
-	})
-
-	t.Run("InterfaceSlice", func(t *testing.T) {
-		servers := []interface{}{"server1", "server2", 456}
-		result, err := ToMCPServers(servers)
-		if err != nil {
-			t.Errorf("Expected no error, got: %v", err)
-		}
-		if len(result.Servers) != 3 {
-			t.Errorf("Expected 3 servers, got %d", len(result.Servers))
-		}
-		if result.Servers[2] != "456" {
-			t.Errorf("Expected '456', got '%s'", result.Servers[2])
-		}
-	})
-
 	t.Run("MapInput", func(t *testing.T) {
 		data := map[string]interface{}{
-			"servers": []string{"server1", "server2"},
+			"servers": []interface{}{"server1", "server2"},
 		}
 		result, err := ToMCPServers(data)
 		if err != nil {
@@ -178,6 +150,9 @@ func TestToMCPServers(t *testing.T) {
 		if len(result.Servers) != 2 {
 			t.Errorf("Expected 2 servers, got %d", len(result.Servers))
 		}
+		if result.Servers[0].ServerID != "server1" {
+			t.Errorf("Expected 'server1', got '%s'", result.Servers[0].ServerID)
+		}
 	})
 
 	t.Run("InvalidInput", func(t *testing.T) {

agent/store/types/mcp_test.go

@@ -0,0 +1,302 @@
+package types
+
+import (
+	"encoding/json"
+	"testing"
+)
+
+func TestMCPServerConfig_UnmarshalJSON(t *testing.T) {
+	tests := []struct {
+		name    string
+		input   string
+		want    MCPServerConfig
+		wantErr bool
+	}{
+		{
+			name:  "Simple string",
+			input: `"server1"`,
+			want: MCPServerConfig{
+				ServerID:  "server1",
+				Resources: nil,
+				Tools:     nil,
+			},
+			wantErr: false,
+		},
+		{
+			name:  "Tools array only",
+			input: `{"server1": ["tool1", "tool2"]}`,
+			want: MCPServerConfig{
+				ServerID:  "server1",
+				Resources: nil,
+				Tools:     []string{"tool1", "tool2"},
+			},
+			wantErr: false,
+		},
+		{
+			name:  "Full config with resources and tools",
+			input: `{"server1": {"resources": ["res1", "res2"], "tools": ["tool1", "tool2"]}}`,
+			want: MCPServerConfig{
+				ServerID:  "server1",
+				Resources: []string{"res1", "res2"},
+				Tools:     []string{"tool1", "tool2"},
+			},
+			wantErr: false,
+		},
+		{
+			name:  "Only resources",
+			input: `{"server1": {"resources": ["res1"]}}`,
+			want: MCPServerConfig{
+				ServerID:  "server1",
+				Resources: []string{"res1"},
+				Tools:     nil,
+			},
+			wantErr: false,
+		},
+		{
+			name:  "Only tools",
+			input: `{"server1": {"tools": ["tool1"]}}`,
+			want: MCPServerConfig{
+				ServerID:  "server1",
+				Resources: nil,
+				Tools:     []string{"tool1"},
+			},
+			wantErr: false,
+		},
+		{
+			name:  "Standard object format",
+			input: `{"server_id": "server1", "resources": ["res1"], "tools": ["tool1"]}`,
+			want: MCPServerConfig{
+				ServerID:  "server1",
+				Resources: []string{"res1"},
+				Tools:     []string{"tool1"},
+			},
+			wantErr: false,
+		},
+		{
+			name:  "Standard object format - no resources/tools",
+			input: `{"server_id": "server1"}`,
+			want: MCPServerConfig{
+				ServerID:  "server1",
+				Resources: nil,
+				Tools:     nil,
+			},
+			wantErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var got MCPServerConfig
+			err := json.Unmarshal([]byte(tt.input), &got)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("UnmarshalJSON() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+
+			if !tt.wantErr {
+				if got.ServerID != tt.want.ServerID {
+					t.Errorf("ServerID = %v, want %v", got.ServerID, tt.want.ServerID)
+				}
+				if !stringSlicesEqual(got.Resources, tt.want.Resources) {
+					t.Errorf("Resources = %v, want %v", got.Resources, tt.want.Resources)
+				}
+				if !stringSlicesEqual(got.Tools, tt.want.Tools) {
+					t.Errorf("Tools = %v, want %v", got.Tools, tt.want.Tools)
+				}
+			}
+		})
+	}
+}
+
+func TestMCPServers_UnmarshalJSON(t *testing.T) {
+	tests := []struct {
+		name    string
+		input   string
+		want    []MCPServerConfig
+		wantErr bool
+	}{
+		{
+			name:  "Simple string array",
+			input: `{"servers": ["server1", "server2", "server3"]}`,
+			want: []MCPServerConfig{
+				{ServerID: "server1"},
+				{ServerID: "server2"},
+				{ServerID: "server3"},
+			},
+			wantErr: false,
+		},
+		{
+			name:  "Mixed formats",
+			input: `{"servers": ["server1", {"server2": ["tool1", "tool2"]}, {"server3": {"resources": ["res1"], "tools": ["tool3"]}}]}`,
+			want: []MCPServerConfig{
+				{ServerID: "server1"},
+				{ServerID: "server2", Tools: []string{"tool1", "tool2"}},
+				{ServerID: "server3", Resources: []string{"res1"}, Tools: []string{"tool3"}},
+			},
+			wantErr: false,
+		},
+		{
+			name:    "Empty servers",
+			input:   `{"servers": []}`,
+			want:    []MCPServerConfig{},
+			wantErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var got MCPServers
+			err := json.Unmarshal([]byte(tt.input), &got)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("UnmarshalJSON() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+
+			if !tt.wantErr {
+				if len(got.Servers) != len(tt.want) {
+					t.Errorf("got %d servers, want %d", len(got.Servers), len(tt.want))
+					return
+				}
+
+				for i := range got.Servers {
+					if got.Servers[i].ServerID != tt.want[i].ServerID {
+						t.Errorf("Server[%d].ServerID = %v, want %v", i, got.Servers[i].ServerID, tt.want[i].ServerID)
+					}
+					if !stringSlicesEqual(got.Servers[i].Resources, tt.want[i].Resources) {
+						t.Errorf("Server[%d].Resources = %v, want %v", i, got.Servers[i].Resources, tt.want[i].Resources)
+					}
+					if !stringSlicesEqual(got.Servers[i].Tools, tt.want[i].Tools) {
+						t.Errorf("Server[%d].Tools = %v, want %v", i, got.Servers[i].Tools, tt.want[i].Tools)
+					}
+				}
+			}
+		})
+	}
+}
+
+func TestMCPServerConfig_MarshalJSON(t *testing.T) {
+	tests := []struct {
+		name   string
+		config MCPServerConfig
+		want   string
+	}{
+		{
+			name: "Only ServerID - should be simple string",
+			config: MCPServerConfig{
+				ServerID: "server1",
+			},
+			want: `"server1"`,
+		},
+		{
+			name: "With Tools - should be object",
+			config: MCPServerConfig{
+				ServerID: "server1",
+				Tools:    []string{"tool1", "tool2"},
+			},
+			want: `{"server_id":"server1","tools":["tool1","tool2"]}`,
+		},
+		{
+			name: "With Resources - should be object",
+			config: MCPServerConfig{
+				ServerID:  "server1",
+				Resources: []string{"res1"},
+			},
+			want: `{"server_id":"server1","resources":["res1"]}`,
+		},
+		{
+			name: "With Both - should be object",
+			config: MCPServerConfig{
+				ServerID:  "server1",
+				Resources: []string{"res1"},
+				Tools:     []string{"tool1"},
+			},
+			want: `{"server_id":"server1","resources":["res1"],"tools":["tool1"]}`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := json.Marshal(tt.config)
+			if err != nil {
+				t.Errorf("MarshalJSON() error = %v", err)
+				return
+			}
+			if string(got) != tt.want {
+				t.Errorf("MarshalJSON() = %s, want %s", string(got), tt.want)
+			}
+		})
+	}
+}
+
+func TestMCPServerConfig_RoundTrip(t *testing.T) {
+	tests := []struct {
+		name   string
+		config MCPServerConfig
+	}{
+		{
+			name: "Simple ServerID",
+			config: MCPServerConfig{
+				ServerID: "server1",
+			},
+		},
+		{
+			name: "With Tools",
+			config: MCPServerConfig{
+				ServerID: "server2",
+				Tools:    []string{"tool1", "tool2"},
+			},
+		},
+		{
+			name: "With Resources and Tools",
+			config: MCPServerConfig{
+				ServerID:  "server3",
+				Resources: []string{"res1", "res2"},
+				Tools:     []string{"tool3", "tool4"},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Marshal
+			data, err := json.Marshal(tt.config)
+			if err != nil {
+				t.Fatalf("Marshal error = %v", err)
+			}
+
+			// Unmarshal
+			var got MCPServerConfig
+			err = json.Unmarshal(data, &got)
+			if err != nil {
+				t.Fatalf("Unmarshal error = %v", err)
+			}
+
+			// Compare
+			if got.ServerID != tt.config.ServerID {
+				t.Errorf("ServerID = %v, want %v", got.ServerID, tt.config.ServerID)
+			}
+			if !stringSlicesEqual(got.Resources, tt.config.Resources) {
+				t.Errorf("Resources = %v, want %v", got.Resources, tt.config.Resources)
+			}
+			if !stringSlicesEqual(got.Tools, tt.config.Tools) {
+				t.Errorf("Tools = %v, want %v", got.Tools, tt.config.Tools)
+			}
+		})
+	}
+}
+
+// Helper function to compare string slices (nil-safe)
+func stringSlicesEqual(a, b []string) bool {
+	if len(a) == 0 && len(b) == 0 {
+		return true
+	}
+	if len(a) != len(b) {
+		return false
+	}
+	for i := range a {
+		if a[i] != b[i] {
+			return false
+		}
+	}
+	return true
+}

agent/store/types/types.go

@@ -1,6 +1,9 @@
 package types
 
 import (
+	"encoding/json"
+	"fmt"
+
 	"github.com/yaoapp/xun/dbal/query"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/i18n"
@@ -99,11 +102,98 @@ type KnowledgeBase struct {
 }
 
 // MCPServers the MCP servers configuration
+// Supports multiple formats in the servers array:
+// - Simple string: "server_id"
+// - With tools: {"server_id": ["tool1", "tool2"]}
+// - With resources and tools: {"server_id": {"resources": [...], "tools": [...]}}
 type MCPServers struct {
-	Servers []string               `json:"servers,omitempty"` // MCP server IDs
+	Servers []MCPServerConfig      `json:"servers,omitempty"` // MCP server configurations
 	Options map[string]interface{} `json:"options,omitempty"` // Additional options for MCP servers
 }
 
+// MCPServerConfig represents a single MCP server configuration
+type MCPServerConfig struct {
+	ServerID  string   `json:"server_id,omitempty"` // MCP server ID
+	Resources []string `json:"resources,omitempty"` // Resources to use (optional)
+	Tools     []string `json:"tools,omitempty"`     // Tools to use (optional)
+}
+
+// UnmarshalJSON implements custom JSON unmarshaling for MCPServerConfig
+// Supports multiple input formats:
+// 1. Simple string: "server_id"
+// 2. Standard object: {"server_id": "server1", "resources": [...], "tools": [...]}
+// 3. Tools array: {"server_id": ["tool1", "tool2"]}
+// 4. Full config: {"server_id": {"resources": [...], "tools": [...]}}
+func (m *MCPServerConfig) UnmarshalJSON(data []byte) error {
+	// Try to unmarshal as string first
+	var str string
+	if err := json.Unmarshal(data, &str); err == nil {
+		m.ServerID = str
+		return nil
+	}
+
+	// Try to unmarshal as standard object with server_id field
+	type Alias MCPServerConfig
+	var stdObj Alias
+	if err := json.Unmarshal(data, &stdObj); err == nil && stdObj.ServerID != "" {
+		*m = MCPServerConfig(stdObj)
+		return nil
+	}
+
+	// Try to unmarshal as object with single key (alternative formats)
+	var obj map[string]json.RawMessage
+	if err := json.Unmarshal(data, &obj); err != nil {
+		return err
+	}
+
+	// Should have exactly one key (the server ID)
+	if len(obj) != 1 {
+		return fmt.Errorf("MCPServerConfig object must have exactly one key or server_id field")
+	}
+
+	// Get the server ID (the only key)
+	for serverID, value := range obj {
+		m.ServerID = serverID
+
+		// Try to unmarshal value as array of strings (format c: tools only)
+		var tools []string
+		if err := json.Unmarshal(value, &tools); err == nil {
+			m.Tools = tools
+			return nil
+		}
+
+		// Try to unmarshal as object with resources and tools (format b)
+		var detail struct {
+			Resources []string `json:"resources,omitempty"`
+			Tools     []string `json:"tools,omitempty"`
+		}
+		if err := json.Unmarshal(value, &detail); err == nil {
+			m.Resources = detail.Resources
+			m.Tools = detail.Tools
+			return nil
+		}
+
+		return fmt.Errorf("invalid format for server '%s'", serverID)
+	}
+
+	return nil
+}
+
+// MarshalJSON implements custom JSON marshaling for MCPServerConfig
+// Serializes to different formats based on content:
+// 1. If only ServerID: "server_id"
+// 2. If has Resources or Tools: {"server_id": "...", "resources": [...], "tools": [...]}
+func (m MCPServerConfig) MarshalJSON() ([]byte, error) {
+	// If only ServerID, serialize as simple string
+	if len(m.Resources) == 0 && len(m.Tools) == 0 {
+		return json.Marshal(m.ServerID)
+	}
+
+	// Otherwise, use standard object format
+	type Alias MCPServerConfig
+	return json.Marshal(Alias(m))
+}
+
 // Workflow the workflow configuration
 type Workflow struct {
 	Workflows []string               `json:"workflows,omitempty"` // Workflow IDs

agent/store/xun/assistant_test.go

@@ -2,6 +2,7 @@ package xun
 
 import (
 	"fmt"
+	"os"
 	"strings"
 	"testing"
 	"time"
@@ -16,10 +17,12 @@ import (
 
 func TestMain(m *testing.M) {
 	// Setup will be done in each test via test.Prepare
-	// Run tests and exit
 	test.Prepare(nil, config.Conf)
 	defer test.Clean()
-	m.Run()
+
+	// Run tests and exit with appropriate exit code
+	code := m.Run()
+	os.Exit(code)
 }
 
 // TestSaveAssistant tests creating and updating assistants
@@ -199,6 +202,137 @@ func TestSaveAssistant(t *testing.T) {
 		}
 	})
 
+	t.Run("SaveWithMCPServers", func(t *testing.T) {
+		// Test creating assistant with MCP servers directly
+		// This will test that:
+		// - server1 (no tools/resources) serializes as "server1"
+		// - server2 (with tools) serializes as {"server_id":"server2","tools":[...]}
+		// - server3 (with both) serializes as {"server_id":"server3","resources":[...],"tools":[...]}
+		assistant := &types.AssistantModel{
+			Name:      "MCP Save Test",
+			Type:      "assistant",
+			Connector: "openai",
+			Share:     "private",
+			MCP: &types.MCPServers{
+				Servers: []types.MCPServerConfig{
+					{ServerID: "server1"},
+					{
+						ServerID: "server2",
+						Tools:    []string{"tool1", "tool2"},
+					},
+					{
+						ServerID:  "server3",
+						Resources: []string{"res1", "res2"},
+						Tools:     []string{"tool3", "tool4"},
+					},
+				},
+				Options: map[string]interface{}{
+					"timeout": 30,
+				},
+			},
+		}
+
+		id, err := store.SaveAssistant(assistant)
+		if err != nil {
+			t.Fatalf("Failed to save assistant with MCP: %v", err)
+		}
+
+		// Retrieve and verify MCP configuration
+		retrieved, err := store.GetAssistant(id)
+		if err != nil {
+			t.Fatalf("Failed to retrieve assistant: %v", err)
+		}
+
+		if retrieved.MCP == nil {
+			t.Fatal("Expected MCP to be set")
+		}
+
+		if len(retrieved.MCP.Servers) != 3 {
+			t.Errorf("Expected 3 MCP servers, got %d", len(retrieved.MCP.Servers))
+		}
+
+		// Verify server1 (simple format)
+		if retrieved.MCP.Servers[0].ServerID != "server1" {
+			t.Errorf("Expected server1, got '%s'", retrieved.MCP.Servers[0].ServerID)
+		}
+
+		// Verify server2 (with tools)
+		if retrieved.MCP.Servers[1].ServerID != "server2" {
+			t.Errorf("Expected server2, got '%s'", retrieved.MCP.Servers[1].ServerID)
+		}
+		if len(retrieved.MCP.Servers[1].Tools) != 2 {
+			t.Errorf("Expected 2 tools for server2, got %d", len(retrieved.MCP.Servers[1].Tools))
+		}
+
+		// Verify server3 (with resources and tools)
+		if retrieved.MCP.Servers[2].ServerID != "server3" {
+			t.Errorf("Expected server3, got '%s'", retrieved.MCP.Servers[2].ServerID)
+		}
+		if len(retrieved.MCP.Servers[2].Resources) != 2 {
+			t.Errorf("Expected 2 resources for server3, got %d", len(retrieved.MCP.Servers[2].Resources))
+		}
+		if len(retrieved.MCP.Servers[2].Tools) != 2 {
+			t.Errorf("Expected 2 tools for server3, got %d", len(retrieved.MCP.Servers[2].Tools))
+		}
+
+		// Verify options
+		if retrieved.MCP.Options == nil {
+			t.Error("Expected MCP options to be set")
+		}
+		if timeout, ok := retrieved.MCP.Options["timeout"].(float64); !ok || timeout != 30 {
+			t.Errorf("Expected timeout 30, got %v", retrieved.MCP.Options["timeout"])
+		}
+
+		t.Logf("Successfully verified MCP configuration for assistant %s", id)
+	})
+
+	t.Run("UpdateWithMCPServers", func(t *testing.T) {
+		// Create assistant without MCP
+		assistant := &types.AssistantModel{
+			Name:      "MCP Update Test",
+			Type:      "assistant",
+			Connector: "openai",
+			Share:     "private",
+		}
+
+		id, err := store.SaveAssistant(assistant)
+		if err != nil {
+			t.Fatalf("Failed to create assistant: %v", err)
+		}
+
+		// Update assistant with MCP
+		assistant.MCP = &types.MCPServers{
+			Servers: []types.MCPServerConfig{
+				{ServerID: "new-server1"},
+				{
+					ServerID: "new-server2",
+					Tools:    []string{"newtool1"},
+				},
+			},
+		}
+
+		_, err = store.SaveAssistant(assistant)
+		if err != nil {
+			t.Fatalf("Failed to update assistant with MCP: %v", err)
+		}
+
+		// Retrieve and verify
+		retrieved, err := store.GetAssistant(id)
+		if err != nil {
+			t.Fatalf("Failed to retrieve assistant: %v", err)
+		}
+
+		if retrieved.MCP == nil || len(retrieved.MCP.Servers) != 2 {
+			t.Errorf("Expected 2 MCP servers, got %v", retrieved.MCP)
+		}
+
+		if retrieved.MCP.Servers[0].ServerID != "new-server1" {
+			t.Errorf("Expected new-server1, got '%s'", retrieved.MCP.Servers[0].ServerID)
+		}
+
+		t.Logf("Successfully updated and verified MCP for assistant %s", id)
+	})
+
 	t.Run("UsesConfiguration", func(t *testing.T) {
 		// Test assistant with Uses configuration
 		assistant := &types.AssistantModel{
@@ -2071,6 +2205,95 @@ func TestUpdateAssistant(t *testing.T) {
 		if retrieved.MCP == nil || len(retrieved.MCP.Servers) != 2 {
 			t.Errorf("Expected 2 MCP servers, got %v", retrieved.MCP)
 		}
+		if retrieved.MCP.Servers[0].ServerID != "server1" {
+			t.Errorf("Expected first server 'server1', got '%s'", retrieved.MCP.Servers[0].ServerID)
+		}
+	})
+
+	t.Run("UpdateMCPWithToolsAndResources", func(t *testing.T) {
+		// Create assistant
+		assistant := &types.AssistantModel{
+			Name:      "MCP Advanced Test",
+			Type:      "assistant",
+			Connector: "openai",
+			Share:     "private",
+		}
+
+		id, err := store.SaveAssistant(assistant)
+		if err != nil {
+			t.Fatalf("Failed to create assistant: %v", err)
+		}
+
+		// Update with MCP servers using advanced configuration
+		updates := map[string]interface{}{
+			"mcp": map[string]interface{}{
+				"servers": []interface{}{
+					"server1", // Simple format
+					map[string]interface{}{
+						"server2": []string{"tool1", "tool2"}, // Tools only
+					},
+					map[string]interface{}{
+						"server3": map[string]interface{}{
+							"resources": []string{"res1", "res2"},
+							"tools":     []string{"tool3", "tool4"},
+						},
+					},
+				},
+			},
+		}
+
+		err = store.UpdateAssistant(id, updates)
+		if err != nil {
+			t.Fatalf("Failed to update MCP: %v", err)
+		}
+
+		// Verify updates
+		retrieved, err := store.GetAssistant(id)
+		if err != nil {
+			t.Fatalf("Failed to retrieve assistant: %v", err)
+		}
+
+		if retrieved.MCP == nil || len(retrieved.MCP.Servers) != 3 {
+			t.Fatalf("Expected 3 MCP servers, got %d", len(retrieved.MCP.Servers))
+		}
+
+		// Verify server1 (simple format)
+		if retrieved.MCP.Servers[0].ServerID != "server1" {
+			t.Errorf("Expected server1, got '%s'", retrieved.MCP.Servers[0].ServerID)
+		}
+		if len(retrieved.MCP.Servers[0].Tools) != 0 {
+			t.Errorf("Expected no tools for server1, got %v", retrieved.MCP.Servers[0].Tools)
+		}
+
+		// Verify server2 (tools only)
+		if retrieved.MCP.Servers[1].ServerID != "server2" {
+			t.Errorf("Expected server2, got '%s'", retrieved.MCP.Servers[1].ServerID)
+		}
+		if len(retrieved.MCP.Servers[1].Tools) != 2 {
+			t.Errorf("Expected 2 tools for server2, got %d", len(retrieved.MCP.Servers[1].Tools))
+		}
+		if retrieved.MCP.Servers[1].Tools[0] != "tool1" {
+			t.Errorf("Expected tool1, got '%s'", retrieved.MCP.Servers[1].Tools[0])
+		}
+
+		// Verify server3 (full config)
+		if retrieved.MCP.Servers[2].ServerID != "server3" {
+			t.Errorf("Expected server3, got '%s'", retrieved.MCP.Servers[2].ServerID)
+		}
+		if len(retrieved.MCP.Servers[2].Resources) != 2 {
+			t.Errorf("Expected 2 resources for server3, got %d", len(retrieved.MCP.Servers[2].Resources))
+		}
+		if len(retrieved.MCP.Servers[2].Tools) != 2 {
+			t.Errorf("Expected 2 tools for server3, got %d", len(retrieved.MCP.Servers[2].Tools))
+		}
+		if retrieved.MCP.Servers[2].Resources[0] != "res1" {
+			t.Errorf("Expected res1, got '%s'", retrieved.MCP.Servers[2].Resources[0])
+		}
+		if retrieved.MCP.Servers[2].Tools[0] != "tool3" {
+			t.Errorf("Expected tool3, got '%s'", retrieved.MCP.Servers[2].Tools[0])
+		}
+
+		t.Logf("Successfully verified MCP advanced configuration for assistant %s", id)
 	})
 
 	t.Run("UpdateUses", func(t *testing.T) {