153 lines
4.2 KiB
Go
153 lines
4.2 KiB
Go
package sidekick
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/mark3labs/mcp-go/mcp"
|
|
)
|
|
|
|
func TestMCPServerQueryHandler(t *testing.T) {
|
|
// 1. Setup Agent with mock LLM
|
|
ac := AgentConfig{ID: "test-agent", Role: RoleSpecialist, MaxIterations: 2}
|
|
mCfg := ModelConfig{}
|
|
agent := NewSidekick(ac, mCfg, nil)
|
|
|
|
mockLLM := &mockLLMClient{
|
|
responses: []Message{
|
|
{Content: `{"action": "RESPOND", "params": {"response": "I am Sidekick. How can I help?"}}`},
|
|
},
|
|
}
|
|
SetLLMClient(mockLLM)
|
|
|
|
pool := map[string]*Sidekick{"test-agent": agent}
|
|
|
|
// 2. Define the handler (mirrors logic in cmd/sidekick-mcp/main.go)
|
|
handler := func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
msg, err := request.RequireString("message")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var inCtx []Message
|
|
args := request.GetArguments()
|
|
if histInter, ok := args["history"]; ok {
|
|
if histList, ok := histInter.([]interface{}); ok {
|
|
for _, h := range histList {
|
|
if hMap, ok := h.(map[string]interface{}); ok {
|
|
role, _ := hMap["role"].(string)
|
|
content, _ := hMap["content"].(string)
|
|
if role != "" && content != "" {
|
|
inCtx = append(inCtx, Message{
|
|
Role: MessageRole(role),
|
|
Content: content,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
inCtx = append(inCtx, Message{
|
|
Role: MessageRoleUser,
|
|
Content: msg,
|
|
})
|
|
|
|
responseStr, err := agent.Run(ctx, inCtx, pool)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
outHistory := append(inCtx, Message{
|
|
Role: MessageRole("assistant"),
|
|
Content: responseStr,
|
|
})
|
|
|
|
resultObj := map[string]interface{}{
|
|
"response": responseStr,
|
|
"history": outHistory,
|
|
}
|
|
|
|
resultBytes, _ := json.Marshal(resultObj)
|
|
|
|
return &mcp.CallToolResult{
|
|
Content: []mcp.Content{
|
|
mcp.NewTextContent(string(resultBytes)),
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
// 3. Test with simple message
|
|
req := mcp.CallToolRequest{}
|
|
req.Params.Name = "query"
|
|
req.Params.Arguments = map[string]interface{}{
|
|
"message": "Hello",
|
|
}
|
|
|
|
res, err := handler(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("handler failed: %v", err)
|
|
}
|
|
|
|
if len(res.Content) != 1 {
|
|
t.Fatalf("expected 1 content item, got %d", len(res.Content))
|
|
}
|
|
|
|
textContent, ok := res.Content[0].(mcp.TextContent)
|
|
if !ok {
|
|
t.Fatalf("expected TextContent")
|
|
}
|
|
|
|
var result map[string]interface{}
|
|
if err := json.Unmarshal([]byte(textContent.Text), &result); err != nil {
|
|
t.Fatalf("failed to unmarshal result: %v", err)
|
|
}
|
|
|
|
if !strings.Contains(result["response"].(string), "How can I help?") {
|
|
t.Errorf("unexpected response: %v", result["response"])
|
|
}
|
|
|
|
history, ok := result["history"].([]interface{})
|
|
if !ok || len(history) != 2 {
|
|
t.Errorf("expected history of length 2, got %v", len(history))
|
|
}
|
|
|
|
// 4. Test with history
|
|
reqWithHist := mcp.CallToolRequest{}
|
|
reqWithHist.Params.Name = "query"
|
|
reqWithHist.Params.Arguments = map[string]interface{}{
|
|
"message": "What did I say?",
|
|
"history": []interface{}{
|
|
map[string]interface{}{"role": "user", "content": "I like cats"},
|
|
map[string]interface{}{"role": "assistant", "content": "Me too"},
|
|
},
|
|
}
|
|
|
|
// Reset mock for next run
|
|
mockLLM.calls = 0
|
|
mockLLM.responses = []Message{
|
|
{Content: `{"action": "RESPOND", "params": {"response": "You said you like cats."}}`},
|
|
}
|
|
|
|
resHist, err := handler(context.Background(), reqWithHist)
|
|
if err != nil {
|
|
t.Fatalf("handler with history failed: %v", err)
|
|
}
|
|
|
|
var resultHist map[string]interface{}
|
|
textContentHist := resHist.Content[0].(mcp.TextContent)
|
|
json.Unmarshal([]byte(textContentHist.Text), &resultHist)
|
|
|
|
historyFinal, ok := resultHist["history"].([]interface{})
|
|
if !ok || len(historyFinal) != 4 { // 2 old + 1 new user + 1 assistant response
|
|
t.Errorf("expected history of length 4, got %v", len(historyFinal))
|
|
}
|
|
|
|
lastMsg := historyFinal[3].(map[string]interface{})
|
|
if lastMsg["content"] != "You said you like cats." {
|
|
t.Errorf("unexpected last message: %v", lastMsg["content"])
|
|
}
|
|
}
|