ai-gateway-go/internal/workbench/mcp_client_test.go

package workbench

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

// minimalMCP is a scriptable JSON-RPC MCP server for exercising the client.
type minimalMCP struct {
	t         *testing.T
	server    *httptest.Server
	tools     []MCPTool
	initCount int32
	listCount int32
	callCount int32
	lastArgs  map[string]any
	lastTool  string
	versioned bool
}

func newMinimalMCP(t *testing.T) *minimalMCP {
	m := &minimalMCP{t: t, tools: []MCPTool{
		{Name: "lookup", Description: "look something up", InputSchema: json.RawMessage(`{"type":"object","properties":{"q":{"type":"string"}},"required":["q"]}`)},
		{Name: "add", Description: "add two numbers", InputSchema: json.RawMessage(`{"type":"object","properties":{"a":{"type":"number"},"b":{"type":"number"}},"required":["a","b"]}`)},
	}}
	m.server = httptest.NewServer(http.HandlerFunc(m.serve))
	t.Cleanup(m.server.Close)
	return m
}

func (m *minimalMCP) endpoint() string { return m.server.URL + "/mcp" }

func (m *minimalMCP) serve(w http.ResponseWriter, r *http.Request) {
	var req mcpRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "bad json", http.StatusBadRequest)
		return
	}
	switch req.Method {
	case "initialize":
		atomic.AddInt32(&m.initCount, 1)
		params, _ := req.Params.(map[string]any)
		clientInfo, _ := params["clientInfo"].(map[string]any)
		if clientInfo["name"] == nil || clientInfo["name"] == "" {
			m.t.Error("initialize missing clientInfo.name")
		}
		result := map[string]any{
			"protocolVersion": "2025-06-18",
			"capabilities":    map[string]any{"tools": map[string]any{}},
			"serverInfo":      map[string]any{"name": "test-mcp", "version": "1.0.0"},
		}
		m.write(w, req.ID, result, false)
	case "notifications/initialized":
		w.WriteHeader(http.StatusAccepted)
	case "tools/list":
		atomic.AddInt32(&m.listCount, 1)
		m.write(w, req.ID, map[string]any{"tools": m.tools}, false)
	case "tools/call":
		atomic.AddInt32(&m.callCount, 1)
		params, _ := req.Params.(map[string]any)
		name, _ := params["name"].(string)
		args, _ := params["arguments"].(map[string]any)
		m.lastArgs = args
		m.lastTool = name
		if name == "boom" {
			m.write(w, req.ID, map[string]any{"content": []map[string]any{{"type": "text", "text": "failed intentionally"}}, "isError": true}, false)
			return
		}
		if name == "nope" {
			m.write(w, req.ID, nil, true)
			return
		}
		m.write(w, req.ID, map[string]any{"content": []map[string]any{{"type": "text", "text": "result for " + name}}, "isError": false}, false)
	case "tools/fail":
		m.write(w, req.ID, nil, true)
	default:
		http.Error(w, "unknown method", http.StatusBadRequest)
	}
}

func (m *minimalMCP) write(w http.ResponseWriter, id any, result any, isError bool) {
	w.Header().Set("Content-Type", "application/json")
	if isError {
		json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": id, "error": map[string]any{"code": -32601, "message": "method not found"}})
		return
	}
	json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": id, "result": result})
}

func testServer(code string) MCPServer {
	return MCPServer{ID: "00000000-0000-0000-0000-000000000001", Code: code}
}

func TestMCPDiscoverTools(t *testing.T) {
	mcp := newMinimalMCP(t)
	client := NewMCPClient(true, time.Minute)
	server := testServer("demo")
	server.EndpointURL = mcp.endpoint()

	tools, err := client.DiscoverTools(context.Background(), server, nil)
	if err != nil {
		t.Fatalf("discover failed: %v", err)
	}
	if len(tools) != 2 || tools[0].Name != "lookup" {
		t.Fatalf("expected 2 tools starting with lookup, got %+v", tools)
	}
	if atomic.LoadInt32(&mcp.initCount) != 1 {
		t.Fatalf("expected exactly one initialize handshake, got %d", mcp.initCount)
	}
	// Second call hits the tool cache, no new tools/list.
	if _, err = client.DiscoverTools(context.Background(), server, nil); err != nil {
		t.Fatalf("cached discover failed: %v", err)
	}
	if atomic.LoadInt32(&mcp.listCount) != 1 {
		t.Fatalf("expected tools/list to run once, got %d", mcp.listCount)
	}
}

func TestMCPCacheExpiryRediscover(t *testing.T) {
	mcp := newMinimalMCP(t)
	client := NewMCPClient(true, 50*time.Millisecond)
	server := testServer("expire")
	server.EndpointURL = mcp.endpoint()

	if _, err := client.DiscoverTools(context.Background(), server, nil); err != nil {
		t.Fatalf("discover failed: %v", err)
	}
	time.Sleep(80 * time.Millisecond)
	if _, err := client.DiscoverTools(context.Background(), server, nil); err != nil {
		t.Fatalf("rediscover failed: %v", err)
	}
	if atomic.LoadInt32(&mcp.listCount) != 2 {
		t.Fatalf("expected tools/list to run twice after expiry, got %d", mcp.listCount)
	}
}

func TestMCPCallTool(t *testing.T) {
	mcp := newMinimalMCP(t)
	client := NewMCPClient(true, time.Minute)
	server := testServer("call")
	server.EndpointURL = mcp.endpoint()

	result, err := client.CallTool(context.Background(), server, nil, "lookup", map[string]any{"q": "hello"})
	if err != nil {
		t.Fatalf("call failed: %v", err)
	}
	if result.Content != "result for lookup" {
		t.Fatalf("unexpected content %q", result.Content)
	}
	if mcp.lastArgs["q"] != "hello" {
		t.Fatalf("arguments not forwarded: %v", mcp.lastArgs)
	}
	// The runtime exposes tools under the mcp__{code}__{tool} prefix; CallTool
	// must strip it so the remote server sees the real tool name.
	if mcp.lastTool != "lookup" {
		t.Fatalf("tool name not sent as-is: %q", mcp.lastTool)
	}
	if _, err := client.CallTool(context.Background(), server, nil, "mcp__call__lookup", map[string]any{"q": "x"}); err != nil {
		t.Fatalf("prefixed call failed: %v", err)
	}
	if mcp.lastTool != "lookup" {
		t.Fatalf("prefixed tool name was not stripped: %q", mcp.lastTool)
	}
}

func TestMCPCallToolIsError(t *testing.T) {
	mcp := newMinimalMCP(t)
	client := NewMCPClient(true, time.Minute)
	server := testServer("err")
	server.EndpointURL = mcp.endpoint()

	if _, err := client.CallTool(context.Background(), server, nil, "boom", nil); err == nil {
		t.Fatal("expected error for isError tool result")
	}
}

func TestMCPJSONRPCError(t *testing.T) {
	mcp := newMinimalMCP(t)
	client := NewMCPClient(true, time.Minute)
	server := testServer("rpc")
	server.EndpointURL = mcp.endpoint()

	if _, err := client.CallTool(context.Background(), server, nil, "nope", nil); err == nil || !strings.Contains(err.Error(), "method not found") {
		t.Fatalf("expected JSON-RPC error surfaced, got %v", err)
	}
}

func TestMCPToolNamePrefixRoundTrip(t *testing.T) {
	prefixed := mcpToolName("github", "create-issue")
	if prefixed != "mcp__github__create-issue" {
		t.Fatalf("unexpected prefix %q", prefixed)
	}
	code, tool, ok := resolveMCPTool(prefixed)
	if !ok || code != "github" || tool != "create-issue" {
		t.Fatalf("resolve failed: %q %q %v", code, tool, ok)
	}
	if _, _, ok := resolveMCPTool("plain-name"); ok {
		t.Fatal("non-prefixed name should not resolve as MCP tool")
	}
}

func TestMCPSSESingleFrame(t *testing.T) {
	handler := func(w http.ResponseWriter, r *http.Request) {
		var req mcpRequest
		_ = json.NewDecoder(r.Body).Decode(&req)
		w.Header().Set("Content-Type", "text/event-stream")
		switch req.Method {
		case "initialize":
			payload, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": map[string]any{"protocolVersion": "2025-06-18", "capabilities": map[string]any{}, "serverInfo": map[string]any{"name": "sse", "version": "1"}}})
			w.Write([]byte("event: message\ndata: " + string(payload) + "\n\n"))
		case "notifications/initialized":
			w.WriteHeader(http.StatusAccepted)
		case "tools/list":
			payload, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": map[string]any{"tools": []MCPTool{{Name: "sse-tool", Description: "", InputSchema: json.RawMessage(`{}`)}}}})
			w.Write([]byte("event: message\ndata: " + string(payload) + "\n\n"))
		}
	}
	server := httptest.NewServer(http.HandlerFunc(handler))
	defer server.Close()

	client := NewMCPClient(true, time.Minute)
	svr := testServer("sse")
	svr.EndpointURL = server.URL
	tools, err := client.DiscoverTools(context.Background(), svr, nil)
	if err != nil {
		t.Fatalf("SSE discover failed: %v", err)
	}
	if len(tools) != 1 || tools[0].Name != "sse-tool" {
		t.Fatalf("unexpected SSE tools: %+v", tools)
	}
}