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

Refactor MCP server configuration and enhance Assistant functionality

Browse source 
- 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.
This commit is contained in:
Max 2025-11-29 10:04:19 +08:00
parent a0483bac5d
commit cfde900ada
9 changed files with 692 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
agent/assistant/agent.go
View file

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

18
agent/assistant/assistant.go
View file

@ -152,8 +152,22 @@ func (ast *Assistant) Clone() *Assistant {
if ast.MCP != nil {
clone.MCP = &store.MCPServers{}
if ast.MCP.Servers != nil {
clone.MCP.Servers = make([]string, len(ast.MCP.Servers))
copy(clone.MCP.Servers, ast.MCP.Servers)
clone.MCP.Servers = make([]store.MCPServerConfig, len(ast.MCP.Servers))
for i, server := range ast.MCP.Servers {
clone.MCP.Servers[i] = store.MCPServerConfig{
ServerID: server.ServerID,
}
// Deep copy Resources slice
if server.Resources != nil {
clone.MCP.Servers[i].Resources = make([]string, len(server.Resources))
copy(clone.MCP.Servers[i].Resources, server.Resources)
}
// Deep copy Tools slice
if server.Tools != nil {
clone.MCP.Servers[i].Tools = make([]string, len(server.Tools))
copy(clone.MCP.Servers[i].Tools, server.Tools)
}
}
}
if ast.MCP.Options != nil {
clone.MCP.Options = make(map[string]interface{})

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

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

1
agent/assistant/llm.go
View file

@ -56,4 +56,3 @@ func (ast *Assistant) executeLLMStream(
return completionResponse, nil
}

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

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

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

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

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

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

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

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

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

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