Sidekick/mcp_server_test.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"])
	}
}