From cfde900ada9cc9fab1fa6f94a7ec7a92286fbd0b Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 29 Nov 2025 10:04:19 +0800 Subject: [PATCH] Refactor MCP server configuration and enhance Assistant functionality - Updated MCP server configuration to support advanced formats, including tools and resources. - Refactored the Assistant's Stream method for improved readability and maintainability, adding clear section comments. - Enhanced error handling and logging for better traceability during streaming operations. - Removed deprecated methods and streamlined MCP server handling in tests to ensure robust functionality and validation. --- agent/assistant/agent.go | 38 ++-- agent/assistant/assistant.go | 18 +- agent/assistant/history.go | 31 +++ agent/assistant/llm.go | 1 - agent/store/types/convert.go | 12 +- agent/store/types/convert_test.go | 37 +--- agent/store/types/mcp_test.go | 302 ++++++++++++++++++++++++++++++ agent/store/types/types.go | 92 ++++++++- agent/store/xun/assistant_test.go | 227 +++++++++++++++++++++- 9 files changed, 692 insertions(+), 66 deletions(-) create mode 100644 agent/assistant/history.go create mode 100644 agent/store/types/mcp_test.go diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index f6385097..9868a4f1 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -32,6 +32,10 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa }) } + // ================================================ + // Initialize + // ================================================ + // Initialize stack and auto-handle completion/failure/restore _, _, done := context.EnterStack(ctx, ast.ID, ctx.Referer) defer done() @@ -54,18 +58,18 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa // Initialize agent trace node agentNode := ast.initAgentTraceNode(ctx, inputMessages) - // Full input messages with chat history - fullMessages, err := ast.WithHistory(ctx, inputMessages) + // ================================================ + // Get Full Messages with chat history + // ================================================ + fullMessages, err := ast.WithHistory(ctx, inputMessages, agentNode) if err != nil { - ast.traceAgentFail(agentNode, err) - // Send error stream_end for root stack ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) return nil, err } - // Log the chat history - ast.traceAgentHistory(ctx, agentNode, fullMessages) - + // ================================================ + // Execute Create Hook + // ================================================ // Request Create hook ( Optional ) var createResponse *context.HookCreateResponse if ast.Script != nil { @@ -82,6 +86,9 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa ast.traceCreateHook(agentNode, createResponse) } + // ================================================ + // Execute LLM Call Stream + // ================================================ // LLM Call Stream ( Optional ) var completionResponse *context.CompletionResponse if ast.Prompts != nil || ast.MCP != nil { @@ -103,19 +110,17 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa } } - // Request MCP hook ( Optional ) - var mcpResponse *context.ResponseHookMCP - if ast.MCP != nil { - _ = mcpResponse // mcpResponse is available for further processing - - // MCP Execution Loop + // ================================================ + // Execute tool calls + // ================================================ + if completionResponse != nil && completionResponse.ToolCalls != nil { } // 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 { // Send error stream_end for root stack ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) @@ -271,11 +276,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 { diff --git a/agent/assistant/assistant.go b/agent/assistant/assistant.go index 2053a818..b2a7b8ef 100644 --- a/agent/assistant/assistant.go +++ b/agent/assistant/assistant.go @@ -152,8 +152,22 @@ func (ast *Assistant) Clone() *Assistant { if ast.MCP != nil { clone.MCP = &store.MCPServers{} if ast.MCP.Servers != nil { - clone.MCP.Servers = make([]string, len(ast.MCP.Servers)) - copy(clone.MCP.Servers, ast.MCP.Servers) + clone.MCP.Servers = make([]store.MCPServerConfig, len(ast.MCP.Servers)) + for i, server := range ast.MCP.Servers { + clone.MCP.Servers[i] = store.MCPServerConfig{ + ServerID: server.ServerID, + } + // Deep copy Resources slice + if server.Resources != nil { + clone.MCP.Servers[i].Resources = make([]string, len(server.Resources)) + copy(clone.MCP.Servers[i].Resources, server.Resources) + } + // Deep copy Tools slice + if server.Tools != nil { + clone.MCP.Servers[i].Tools = make([]string, len(server.Tools)) + copy(clone.MCP.Servers[i].Tools, server.Tools) + } + } } if ast.MCP.Options != nil { clone.MCP.Options = make(map[string]interface{}) diff --git a/agent/assistant/history.go b/agent/assistant/history.go new file mode 100644 index 00000000..bbecb73c --- /dev/null +++ b/agent/assistant/history.go @@ -0,0 +1,31 @@ +package assistant + +import ( + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/trace/types" +) + +// WithHistory merges the input messages with chat history and traces it +// This method can be overridden or extended to implement actual history loading +func (ast *Assistant) WithHistory( + ctx *context.Context, + inputMessages []context.Message, + agentNode types.Node, +) ([]context.Message, error) { + + // TODO: Implement actual history loading logic here + // For now, just simulate a check and return the input messages as is + + // Simulate error check (this is where actual history loading would happen) + // if some_condition { + // ast.traceAgentFail(agentNode, err) + // return nil, err + // } + + fullMessages := inputMessages + + // Log the chat history + ast.traceAgentHistory(ctx, agentNode, fullMessages) + + return fullMessages, nil +} diff --git a/agent/assistant/llm.go b/agent/assistant/llm.go index 125d4ff1..aba9246e 100644 --- a/agent/assistant/llm.go +++ b/agent/assistant/llm.go @@ -56,4 +56,3 @@ func (ast *Assistant) executeLLMStream( return completionResponse, nil } - diff --git a/agent/store/types/convert.go b/agent/store/types/convert.go index 43b50260..d8ddc59f 100644 --- a/agent/store/types/convert.go +++ b/agent/store/types/convert.go @@ -63,17 +63,9 @@ func ToMCPServers(v interface{}) (*MCPServers, error) { case MCPServers: return &mcp, nil - case []string: - return &MCPServers{Servers: mcp}, nil - - case []interface{}: - var servers []string - for _, item := range mcp { - servers = append(servers, cast.ToString(item)) - } - return &MCPServers{Servers: servers}, nil - default: + // For any type (including []string, []interface{}, map[string]interface{}), + // marshal and unmarshal to MCPServers using custom UnmarshalJSON raw, err := jsoniter.Marshal(mcp) if err != nil { return nil, fmt.Errorf("mcp format error: %s", err.Error()) diff --git a/agent/store/types/convert_test.go b/agent/store/types/convert_test.go index e21e8759..ad17acbe 100644 --- a/agent/store/types/convert_test.go +++ b/agent/store/types/convert_test.go @@ -118,7 +118,7 @@ func TestToMCPServers(t *testing.T) { }) t.Run("MCPServersPointer", func(t *testing.T) { - mcp := &MCPServers{Servers: []string{"server1", "server2"}} + mcp := &MCPServers{Servers: []MCPServerConfig{{ServerID: "server1"}, {ServerID: "server2"}}} result, err := ToMCPServers(mcp) if err != nil { t.Errorf("Expected no error, got: %v", err) @@ -129,7 +129,7 @@ func TestToMCPServers(t *testing.T) { }) t.Run("MCPServersValue", func(t *testing.T) { - mcp := MCPServers{Servers: []string{"server1", "server2"}} + mcp := MCPServers{Servers: []MCPServerConfig{{ServerID: "server1"}, {ServerID: "server2"}}} result, err := ToMCPServers(mcp) if err != nil { t.Errorf("Expected no error, got: %v", err) @@ -139,37 +139,9 @@ func TestToMCPServers(t *testing.T) { } }) - t.Run("StringSlice", func(t *testing.T) { - servers := []string{"server1", "server2", "server3"} - result, err := ToMCPServers(servers) - if err != nil { - t.Errorf("Expected no error, got: %v", err) - } - if len(result.Servers) != 3 { - t.Errorf("Expected 3 servers, got %d", len(result.Servers)) - } - if result.Servers[0] != "server1" { - t.Errorf("Expected 'server1', got '%s'", result.Servers[0]) - } - }) - - t.Run("InterfaceSlice", func(t *testing.T) { - servers := []interface{}{"server1", "server2", 456} - result, err := ToMCPServers(servers) - if err != nil { - t.Errorf("Expected no error, got: %v", err) - } - if len(result.Servers) != 3 { - t.Errorf("Expected 3 servers, got %d", len(result.Servers)) - } - if result.Servers[2] != "456" { - t.Errorf("Expected '456', got '%s'", result.Servers[2]) - } - }) - t.Run("MapInput", func(t *testing.T) { data := map[string]interface{}{ - "servers": []string{"server1", "server2"}, + "servers": []interface{}{"server1", "server2"}, } result, err := ToMCPServers(data) if err != nil { @@ -178,6 +150,9 @@ func TestToMCPServers(t *testing.T) { if len(result.Servers) != 2 { t.Errorf("Expected 2 servers, got %d", len(result.Servers)) } + if result.Servers[0].ServerID != "server1" { + t.Errorf("Expected 'server1', got '%s'", result.Servers[0].ServerID) + } }) t.Run("InvalidInput", func(t *testing.T) { diff --git a/agent/store/types/mcp_test.go b/agent/store/types/mcp_test.go new file mode 100644 index 00000000..454dbbc5 --- /dev/null +++ b/agent/store/types/mcp_test.go @@ -0,0 +1,302 @@ +package types + +import ( + "encoding/json" + "testing" +) + +func TestMCPServerConfig_UnmarshalJSON(t *testing.T) { + tests := []struct { + name string + input string + want MCPServerConfig + wantErr bool + }{ + { + name: "Simple string", + input: `"server1"`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: nil, + Tools: nil, + }, + wantErr: false, + }, + { + name: "Tools array only", + input: `{"server1": ["tool1", "tool2"]}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: nil, + Tools: []string{"tool1", "tool2"}, + }, + wantErr: false, + }, + { + name: "Full config with resources and tools", + input: `{"server1": {"resources": ["res1", "res2"], "tools": ["tool1", "tool2"]}}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: []string{"res1", "res2"}, + Tools: []string{"tool1", "tool2"}, + }, + wantErr: false, + }, + { + name: "Only resources", + input: `{"server1": {"resources": ["res1"]}}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: []string{"res1"}, + Tools: nil, + }, + wantErr: false, + }, + { + name: "Only tools", + input: `{"server1": {"tools": ["tool1"]}}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: nil, + Tools: []string{"tool1"}, + }, + wantErr: false, + }, + { + name: "Standard object format", + input: `{"server_id": "server1", "resources": ["res1"], "tools": ["tool1"]}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: []string{"res1"}, + Tools: []string{"tool1"}, + }, + wantErr: false, + }, + { + name: "Standard object format - no resources/tools", + input: `{"server_id": "server1"}`, + want: MCPServerConfig{ + ServerID: "server1", + Resources: nil, + Tools: nil, + }, + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got MCPServerConfig + err := json.Unmarshal([]byte(tt.input), &got) + if (err != nil) != tt.wantErr { + t.Errorf("UnmarshalJSON() error = %v, wantErr %v", err, tt.wantErr) + return + } + + if !tt.wantErr { + if got.ServerID != tt.want.ServerID { + t.Errorf("ServerID = %v, want %v", got.ServerID, tt.want.ServerID) + } + if !stringSlicesEqual(got.Resources, tt.want.Resources) { + t.Errorf("Resources = %v, want %v", got.Resources, tt.want.Resources) + } + if !stringSlicesEqual(got.Tools, tt.want.Tools) { + t.Errorf("Tools = %v, want %v", got.Tools, tt.want.Tools) + } + } + }) + } +} + +func TestMCPServers_UnmarshalJSON(t *testing.T) { + tests := []struct { + name string + input string + want []MCPServerConfig + wantErr bool + }{ + { + name: "Simple string array", + input: `{"servers": ["server1", "server2", "server3"]}`, + want: []MCPServerConfig{ + {ServerID: "server1"}, + {ServerID: "server2"}, + {ServerID: "server3"}, + }, + wantErr: false, + }, + { + name: "Mixed formats", + input: `{"servers": ["server1", {"server2": ["tool1", "tool2"]}, {"server3": {"resources": ["res1"], "tools": ["tool3"]}}]}`, + want: []MCPServerConfig{ + {ServerID: "server1"}, + {ServerID: "server2", Tools: []string{"tool1", "tool2"}}, + {ServerID: "server3", Resources: []string{"res1"}, Tools: []string{"tool3"}}, + }, + wantErr: false, + }, + { + name: "Empty servers", + input: `{"servers": []}`, + want: []MCPServerConfig{}, + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got MCPServers + err := json.Unmarshal([]byte(tt.input), &got) + if (err != nil) != tt.wantErr { + t.Errorf("UnmarshalJSON() error = %v, wantErr %v", err, tt.wantErr) + return + } + + if !tt.wantErr { + if len(got.Servers) != len(tt.want) { + t.Errorf("got %d servers, want %d", len(got.Servers), len(tt.want)) + return + } + + for i := range got.Servers { + if got.Servers[i].ServerID != tt.want[i].ServerID { + t.Errorf("Server[%d].ServerID = %v, want %v", i, got.Servers[i].ServerID, tt.want[i].ServerID) + } + if !stringSlicesEqual(got.Servers[i].Resources, tt.want[i].Resources) { + t.Errorf("Server[%d].Resources = %v, want %v", i, got.Servers[i].Resources, tt.want[i].Resources) + } + if !stringSlicesEqual(got.Servers[i].Tools, tt.want[i].Tools) { + t.Errorf("Server[%d].Tools = %v, want %v", i, got.Servers[i].Tools, tt.want[i].Tools) + } + } + } + }) + } +} + +func TestMCPServerConfig_MarshalJSON(t *testing.T) { + tests := []struct { + name string + config MCPServerConfig + want string + }{ + { + name: "Only ServerID - should be simple string", + config: MCPServerConfig{ + ServerID: "server1", + }, + want: `"server1"`, + }, + { + name: "With Tools - should be object", + config: MCPServerConfig{ + ServerID: "server1", + Tools: []string{"tool1", "tool2"}, + }, + want: `{"server_id":"server1","tools":["tool1","tool2"]}`, + }, + { + name: "With Resources - should be object", + config: MCPServerConfig{ + ServerID: "server1", + Resources: []string{"res1"}, + }, + want: `{"server_id":"server1","resources":["res1"]}`, + }, + { + name: "With Both - should be object", + config: MCPServerConfig{ + ServerID: "server1", + Resources: []string{"res1"}, + Tools: []string{"tool1"}, + }, + want: `{"server_id":"server1","resources":["res1"],"tools":["tool1"]}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := json.Marshal(tt.config) + if err != nil { + t.Errorf("MarshalJSON() error = %v", err) + return + } + if string(got) != tt.want { + t.Errorf("MarshalJSON() = %s, want %s", string(got), tt.want) + } + }) + } +} + +func TestMCPServerConfig_RoundTrip(t *testing.T) { + tests := []struct { + name string + config MCPServerConfig + }{ + { + name: "Simple ServerID", + config: MCPServerConfig{ + ServerID: "server1", + }, + }, + { + name: "With Tools", + config: MCPServerConfig{ + ServerID: "server2", + Tools: []string{"tool1", "tool2"}, + }, + }, + { + name: "With Resources and Tools", + config: MCPServerConfig{ + ServerID: "server3", + Resources: []string{"res1", "res2"}, + Tools: []string{"tool3", "tool4"}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Marshal + data, err := json.Marshal(tt.config) + if err != nil { + t.Fatalf("Marshal error = %v", err) + } + + // Unmarshal + var got MCPServerConfig + err = json.Unmarshal(data, &got) + if err != nil { + t.Fatalf("Unmarshal error = %v", err) + } + + // Compare + if got.ServerID != tt.config.ServerID { + t.Errorf("ServerID = %v, want %v", got.ServerID, tt.config.ServerID) + } + if !stringSlicesEqual(got.Resources, tt.config.Resources) { + t.Errorf("Resources = %v, want %v", got.Resources, tt.config.Resources) + } + if !stringSlicesEqual(got.Tools, tt.config.Tools) { + t.Errorf("Tools = %v, want %v", got.Tools, tt.config.Tools) + } + }) + } +} + +// Helper function to compare string slices (nil-safe) +func stringSlicesEqual(a, b []string) bool { + if len(a) == 0 && len(b) == 0 { + return true + } + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} diff --git a/agent/store/types/types.go b/agent/store/types/types.go index 3f3f7867..c145c65f 100644 --- a/agent/store/types/types.go +++ b/agent/store/types/types.go @@ -1,6 +1,9 @@ package types import ( + "encoding/json" + "fmt" + "github.com/yaoapp/xun/dbal/query" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" @@ -99,11 +102,98 @@ type KnowledgeBase struct { } // MCPServers the MCP servers configuration +// Supports multiple formats in the servers array: +// - Simple string: "server_id" +// - With tools: {"server_id": ["tool1", "tool2"]} +// - With resources and tools: {"server_id": {"resources": [...], "tools": [...]}} type MCPServers struct { - Servers []string `json:"servers,omitempty"` // MCP server IDs + Servers []MCPServerConfig `json:"servers,omitempty"` // MCP server configurations Options map[string]interface{} `json:"options,omitempty"` // Additional options for MCP servers } +// MCPServerConfig represents a single MCP server configuration +type MCPServerConfig struct { + ServerID string `json:"server_id,omitempty"` // MCP server ID + Resources []string `json:"resources,omitempty"` // Resources to use (optional) + Tools []string `json:"tools,omitempty"` // Tools to use (optional) +} + +// UnmarshalJSON implements custom JSON unmarshaling for MCPServerConfig +// Supports multiple input formats: +// 1. Simple string: "server_id" +// 2. Standard object: {"server_id": "server1", "resources": [...], "tools": [...]} +// 3. Tools array: {"server_id": ["tool1", "tool2"]} +// 4. Full config: {"server_id": {"resources": [...], "tools": [...]}} +func (m *MCPServerConfig) UnmarshalJSON(data []byte) error { + // Try to unmarshal as string first + var str string + if err := json.Unmarshal(data, &str); err == nil { + m.ServerID = str + return nil + } + + // Try to unmarshal as standard object with server_id field + type Alias MCPServerConfig + var stdObj Alias + if err := json.Unmarshal(data, &stdObj); err == nil && stdObj.ServerID != "" { + *m = MCPServerConfig(stdObj) + return nil + } + + // Try to unmarshal as object with single key (alternative formats) + var obj map[string]json.RawMessage + if err := json.Unmarshal(data, &obj); err != nil { + return err + } + + // Should have exactly one key (the server ID) + if len(obj) != 1 { + return fmt.Errorf("MCPServerConfig object must have exactly one key or server_id field") + } + + // Get the server ID (the only key) + for serverID, value := range obj { + m.ServerID = serverID + + // Try to unmarshal value as array of strings (format c: tools only) + var tools []string + if err := json.Unmarshal(value, &tools); err == nil { + m.Tools = tools + return nil + } + + // Try to unmarshal as object with resources and tools (format b) + var detail struct { + Resources []string `json:"resources,omitempty"` + Tools []string `json:"tools,omitempty"` + } + if err := json.Unmarshal(value, &detail); err == nil { + m.Resources = detail.Resources + m.Tools = detail.Tools + return nil + } + + return fmt.Errorf("invalid format for server '%s'", serverID) + } + + return nil +} + +// MarshalJSON implements custom JSON marshaling for MCPServerConfig +// Serializes to different formats based on content: +// 1. If only ServerID: "server_id" +// 2. If has Resources or Tools: {"server_id": "...", "resources": [...], "tools": [...]} +func (m MCPServerConfig) MarshalJSON() ([]byte, error) { + // If only ServerID, serialize as simple string + if len(m.Resources) == 0 && len(m.Tools) == 0 { + return json.Marshal(m.ServerID) + } + + // Otherwise, use standard object format + type Alias MCPServerConfig + return json.Marshal(Alias(m)) +} + // Workflow the workflow configuration type Workflow struct { Workflows []string `json:"workflows,omitempty"` // Workflow IDs diff --git a/agent/store/xun/assistant_test.go b/agent/store/xun/assistant_test.go index ce50c548..4ba8e4a2 100644 --- a/agent/store/xun/assistant_test.go +++ b/agent/store/xun/assistant_test.go @@ -2,6 +2,7 @@ package xun import ( "fmt" + "os" "strings" "testing" "time" @@ -16,10 +17,12 @@ import ( func TestMain(m *testing.M) { // Setup will be done in each test via test.Prepare - // Run tests and exit test.Prepare(nil, config.Conf) defer test.Clean() - m.Run() + + // Run tests and exit with appropriate exit code + code := m.Run() + os.Exit(code) } // TestSaveAssistant tests creating and updating assistants @@ -199,6 +202,137 @@ func TestSaveAssistant(t *testing.T) { } }) + t.Run("SaveWithMCPServers", func(t *testing.T) { + // Test creating assistant with MCP servers directly + // This will test that: + // - server1 (no tools/resources) serializes as "server1" + // - server2 (with tools) serializes as {"server_id":"server2","tools":[...]} + // - server3 (with both) serializes as {"server_id":"server3","resources":[...],"tools":[...]} + assistant := &types.AssistantModel{ + Name: "MCP Save Test", + Type: "assistant", + Connector: "openai", + Share: "private", + MCP: &types.MCPServers{ + Servers: []types.MCPServerConfig{ + {ServerID: "server1"}, + { + ServerID: "server2", + Tools: []string{"tool1", "tool2"}, + }, + { + ServerID: "server3", + Resources: []string{"res1", "res2"}, + Tools: []string{"tool3", "tool4"}, + }, + }, + Options: map[string]interface{}{ + "timeout": 30, + }, + }, + } + + id, err := store.SaveAssistant(assistant) + if err != nil { + t.Fatalf("Failed to save assistant with MCP: %v", err) + } + + // Retrieve and verify MCP configuration + retrieved, err := store.GetAssistant(id) + if err != nil { + t.Fatalf("Failed to retrieve assistant: %v", err) + } + + if retrieved.MCP == nil { + t.Fatal("Expected MCP to be set") + } + + if len(retrieved.MCP.Servers) != 3 { + t.Errorf("Expected 3 MCP servers, got %d", len(retrieved.MCP.Servers)) + } + + // Verify server1 (simple format) + if retrieved.MCP.Servers[0].ServerID != "server1" { + t.Errorf("Expected server1, got '%s'", retrieved.MCP.Servers[0].ServerID) + } + + // Verify server2 (with tools) + if retrieved.MCP.Servers[1].ServerID != "server2" { + t.Errorf("Expected server2, got '%s'", retrieved.MCP.Servers[1].ServerID) + } + if len(retrieved.MCP.Servers[1].Tools) != 2 { + t.Errorf("Expected 2 tools for server2, got %d", len(retrieved.MCP.Servers[1].Tools)) + } + + // Verify server3 (with resources and tools) + if retrieved.MCP.Servers[2].ServerID != "server3" { + t.Errorf("Expected server3, got '%s'", retrieved.MCP.Servers[2].ServerID) + } + if len(retrieved.MCP.Servers[2].Resources) != 2 { + t.Errorf("Expected 2 resources for server3, got %d", len(retrieved.MCP.Servers[2].Resources)) + } + if len(retrieved.MCP.Servers[2].Tools) != 2 { + t.Errorf("Expected 2 tools for server3, got %d", len(retrieved.MCP.Servers[2].Tools)) + } + + // Verify options + if retrieved.MCP.Options == nil { + t.Error("Expected MCP options to be set") + } + if timeout, ok := retrieved.MCP.Options["timeout"].(float64); !ok || timeout != 30 { + t.Errorf("Expected timeout 30, got %v", retrieved.MCP.Options["timeout"]) + } + + t.Logf("Successfully verified MCP configuration for assistant %s", id) + }) + + t.Run("UpdateWithMCPServers", func(t *testing.T) { + // Create assistant without MCP + assistant := &types.AssistantModel{ + Name: "MCP Update Test", + Type: "assistant", + Connector: "openai", + Share: "private", + } + + id, err := store.SaveAssistant(assistant) + if err != nil { + t.Fatalf("Failed to create assistant: %v", err) + } + + // Update assistant with MCP + assistant.MCP = &types.MCPServers{ + Servers: []types.MCPServerConfig{ + {ServerID: "new-server1"}, + { + ServerID: "new-server2", + Tools: []string{"newtool1"}, + }, + }, + } + + _, err = store.SaveAssistant(assistant) + if err != nil { + t.Fatalf("Failed to update assistant with MCP: %v", err) + } + + // Retrieve and verify + retrieved, err := store.GetAssistant(id) + if err != nil { + t.Fatalf("Failed to retrieve assistant: %v", err) + } + + if retrieved.MCP == nil || len(retrieved.MCP.Servers) != 2 { + t.Errorf("Expected 2 MCP servers, got %v", retrieved.MCP) + } + + if retrieved.MCP.Servers[0].ServerID != "new-server1" { + t.Errorf("Expected new-server1, got '%s'", retrieved.MCP.Servers[0].ServerID) + } + + t.Logf("Successfully updated and verified MCP for assistant %s", id) + }) + t.Run("UsesConfiguration", func(t *testing.T) { // Test assistant with Uses configuration assistant := &types.AssistantModel{ @@ -2071,6 +2205,95 @@ func TestUpdateAssistant(t *testing.T) { if retrieved.MCP == nil || len(retrieved.MCP.Servers) != 2 { t.Errorf("Expected 2 MCP servers, got %v", retrieved.MCP) } + if retrieved.MCP.Servers[0].ServerID != "server1" { + t.Errorf("Expected first server 'server1', got '%s'", retrieved.MCP.Servers[0].ServerID) + } + }) + + t.Run("UpdateMCPWithToolsAndResources", func(t *testing.T) { + // Create assistant + assistant := &types.AssistantModel{ + Name: "MCP Advanced Test", + Type: "assistant", + Connector: "openai", + Share: "private", + } + + id, err := store.SaveAssistant(assistant) + if err != nil { + t.Fatalf("Failed to create assistant: %v", err) + } + + // Update with MCP servers using advanced configuration + updates := map[string]interface{}{ + "mcp": map[string]interface{}{ + "servers": []interface{}{ + "server1", // Simple format + map[string]interface{}{ + "server2": []string{"tool1", "tool2"}, // Tools only + }, + map[string]interface{}{ + "server3": map[string]interface{}{ + "resources": []string{"res1", "res2"}, + "tools": []string{"tool3", "tool4"}, + }, + }, + }, + }, + } + + err = store.UpdateAssistant(id, updates) + if err != nil { + t.Fatalf("Failed to update MCP: %v", err) + } + + // Verify updates + retrieved, err := store.GetAssistant(id) + if err != nil { + t.Fatalf("Failed to retrieve assistant: %v", err) + } + + if retrieved.MCP == nil || len(retrieved.MCP.Servers) != 3 { + t.Fatalf("Expected 3 MCP servers, got %d", len(retrieved.MCP.Servers)) + } + + // Verify server1 (simple format) + if retrieved.MCP.Servers[0].ServerID != "server1" { + t.Errorf("Expected server1, got '%s'", retrieved.MCP.Servers[0].ServerID) + } + if len(retrieved.MCP.Servers[0].Tools) != 0 { + t.Errorf("Expected no tools for server1, got %v", retrieved.MCP.Servers[0].Tools) + } + + // Verify server2 (tools only) + if retrieved.MCP.Servers[1].ServerID != "server2" { + t.Errorf("Expected server2, got '%s'", retrieved.MCP.Servers[1].ServerID) + } + if len(retrieved.MCP.Servers[1].Tools) != 2 { + t.Errorf("Expected 2 tools for server2, got %d", len(retrieved.MCP.Servers[1].Tools)) + } + if retrieved.MCP.Servers[1].Tools[0] != "tool1" { + t.Errorf("Expected tool1, got '%s'", retrieved.MCP.Servers[1].Tools[0]) + } + + // Verify server3 (full config) + if retrieved.MCP.Servers[2].ServerID != "server3" { + t.Errorf("Expected server3, got '%s'", retrieved.MCP.Servers[2].ServerID) + } + if len(retrieved.MCP.Servers[2].Resources) != 2 { + t.Errorf("Expected 2 resources for server3, got %d", len(retrieved.MCP.Servers[2].Resources)) + } + if len(retrieved.MCP.Servers[2].Tools) != 2 { + t.Errorf("Expected 2 tools for server3, got %d", len(retrieved.MCP.Servers[2].Tools)) + } + if retrieved.MCP.Servers[2].Resources[0] != "res1" { + t.Errorf("Expected res1, got '%s'", retrieved.MCP.Servers[2].Resources[0]) + } + if retrieved.MCP.Servers[2].Tools[0] != "tool3" { + t.Errorf("Expected tool3, got '%s'", retrieved.MCP.Servers[2].Tools[0]) + } + + t.Logf("Successfully verified MCP advanced configuration for assistant %s", id) }) t.Run("UpdateUses", func(t *testing.T) {