Maven/internal/mcp/mcp_test.go

package mcp

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"
)

// fakePoster answers POSTs from a canned handler, in either JSON or SSE form.
type fakePoster struct {
	mu      sync.Mutex
	handler func(method string, params json.RawMessage) (any, *rpcError)
	sse     bool
	session string
	seen    []map[string]string // headers of each request, for the session test
	calls   []string
}

func (f *fakePoster) Post(_ context.Context, _, _ string, body []byte, hdr map[string]string) (*PostResponse, error) {
	var req struct {
		ID     *int64          `json:"id"`
		Method string          `json:"method"`
		Params json.RawMessage `json:"params"`
	}
	if err := json.Unmarshal(body, &req); err != nil {
		return nil, err
	}
	f.mu.Lock()
	f.seen = append(f.seen, hdr)
	f.calls = append(f.calls, req.Method)
	f.mu.Unlock()

	if req.ID == nil { // notification
		return &PostResponse{Status: 202, Body: []byte(`{}`)}, nil
	}
	result, rerr := f.handler(req.Method, req.Params)
	resp := map[string]any{"jsonrpc": "2.0", "id": *req.ID}
	if rerr != nil {
		resp["error"] = map[string]any{"code": rerr.Code, "message": rerr.Message}
	} else {
		resp["result"] = result
	}
	raw, _ := json.Marshal(resp)
	out := &PostResponse{Status: 200, Body: raw, ContentType: "application/json", Header: map[string]string{}}
	if f.sse {
		out.ContentType = "text/event-stream"
		out.Body = []byte("event: message\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\"}\n\nevent: message\ndata: " + string(raw) + "\n\n")
	}
	if f.session != "" {
		out.Header["Mcp-Session-Id"] = f.session
	}
	return out, nil
}

// echoServer is a handler with two tools, one read-only and one not.
func echoServer() func(string, json.RawMessage) (any, *rpcError) {
	return func(method string, params json.RawMessage) (any, *rpcError) {
		switch method {
		case "initialize":
			return map[string]any{
				"protocolVersion": ProtocolVersion,
				"serverInfo":      map[string]any{"name": "fake", "version": "0.1"},
			}, nil
		case "tools/list":
			return map[string]any{"tools": []any{
				map[string]any{
					"name": "read_thing", "description": "reads",
					"inputSchema": map[string]any{"type": "object"},
					"annotations": map[string]any{"readOnlyHint": true},
				},
				map[string]any{"name": "break_thing", "description": "mutates"},
			}}, nil
		case "tools/call":
			var p struct {
				Name string         `json:"name"`
				Args map[string]any `json:"arguments"`
			}
			_ = json.Unmarshal(params, &p)
			if p.Name == "break_thing" {
				return map[string]any{"isError": true, "content": []any{
					map[string]any{"type": "text", "text": "не вышло"}}}, nil
			}
			return map[string]any{"content": []any{
				map[string]any{"type": "text", "text": fmt.Sprintf("%s:%v", p.Name, p.Args["q"])},
				map[string]any{"type": "image", "text": "ignored"},
			}}, nil
		case "resources/list":
			return map[string]any{"resources": []any{
				map[string]any{"uri": "note://one", "name": "one", "mimeType": "text/plain"},
				map[string]any{"uri": "", "name": "nameless"},
			}}, nil
		case "resources/read":
			return map[string]any{"contents": []any{map[string]any{"text": "тело ресурса"}}}, nil
		}
		return nil, &rpcError{Code: -32601, Message: "method not found"}
	}
}

func dialFake(t *testing.T, p *fakePoster) *Client {
	t.Helper()
	c := newClient("fake", newHTTPTransport(p, "http://example.test/mcp"))
	if err := c.Initialize(context.Background()); err != nil {
		t.Fatalf("initialize: %v", err)
	}
	return c
}

func TestHandshakeAndDiscovery(t *testing.T) {
	for _, sse := range []bool{false, true} {
		name := "json"
		if sse {
			name = "sse"
		}
		t.Run(name, func(t *testing.T) {
			p := &fakePoster{handler: echoServer(), sse: sse}
			c := dialFake(t, p)
			if got := c.Info().Name; got != "fake" {
				t.Fatalf("server name = %q", got)
			}
			if got := c.Info().ProtocolVersion; got != ProtocolVersion {
				t.Fatalf("protocol = %q", got)
			}
			tools, err := c.ListTools(context.Background())
			if err != nil {
				t.Fatalf("list tools: %v", err)
			}
			if len(tools) != 2 {
				t.Fatalf("tools = %+v", tools)
			}
			byName := map[string]Tool{}
			for _, tl := range tools {
				byName[tl.Name] = tl
			}
			if !byName["read_thing"].ReadOnly {
				t.Error("read_thing should be read-only (readOnlyHint true)")
			}
			// The important direction: no annotation ⇒ assume it mutates.
			if byName["break_thing"].ReadOnly {
				t.Error("break_thing has no readOnlyHint, must NOT be treated as read-only")
			}
			if byName["read_thing"].Server != "fake" {
				t.Error("tool should carry its server handle")
			}
		})
	}
}

func TestCallToolTextOnly(t *testing.T) {
	c := dialFake(t, &fakePoster{handler: echoServer()})
	out, err := c.CallTool(context.Background(), "read_thing", map[string]any{"q": "привет"})
	if err != nil {
		t.Fatalf("call: %v", err)
	}
	if out != "read_thing:привет" {
		t.Fatalf("out = %q (non-text content must be dropped)", out)
	}
}

func TestCallToolErrorResult(t *testing.T) {
	c := dialFake(t, &fakePoster{handler: echoServer()})
	out, err := c.CallTool(context.Background(), "break_thing", nil)
	if err == nil {
		t.Fatal("isError result must surface as an error")
	}
	if out != "не вышло" {
		t.Fatalf("text should still come back, got %q", out)
	}
}

func TestResources(t *testing.T) {
	c := dialFake(t, &fakePoster{handler: echoServer()})
	rs, err := c.ListResources(context.Background())
	if err != nil {
		t.Fatalf("list resources: %v", err)
	}
	if len(rs) != 1 || rs[0].URI != "note://one" {
		t.Fatalf("resources = %+v (a uri-less entry must be dropped)", rs)
	}
	body, err := c.ReadResource(context.Background(), "note://one")
	if err != nil {
		t.Fatalf("read: %v", err)
	}
	if body != "тело ресурса" {
		t.Fatalf("body = %q", body)
	}
}

func TestCallBeforeInitializeRefused(t *testing.T) {
	c := newClient("fake", newHTTPTransport(&fakePoster{handler: echoServer()}, "http://example.test/mcp"))
	if _, err := c.CallTool(context.Background(), "read_thing", nil); err != ErrNotInitialized {
		t.Fatalf("err = %v, want ErrNotInitialized", err)
	}
}

func TestSessionIDEchoed(t *testing.T) {
	p := &fakePoster{handler: echoServer(), session: "sess-1"}
	c := dialFake(t, p)
	if _, err := c.ListTools(context.Background()); err != nil {
		t.Fatal(err)
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	last := p.seen[len(p.seen)-1]
	if last["Mcp-Session-Id"] != "sess-1" {
		t.Fatalf("session header not echoed: %+v", last)
	}
	if !strings.Contains(last["Accept"], "text/event-stream") {
		t.Fatalf("Accept must offer both forms: %q", last["Accept"])
	}
}

func TestHandshakeWithoutProtocolVersionRefused(t *testing.T) {
	p := &fakePoster{handler: func(m string, _ json.RawMessage) (any, *rpcError) {
		return map[string]any{"serverInfo": map[string]any{"name": "not-mcp"}}, nil
	}}
	c := newClient("x", newHTTPTransport(p, "http://example.test/mcp"))
	if err := c.Initialize(context.Background()); err == nil {
		t.Fatal("a reply with no protocolVersion is not an MCP server")
	}
}

func TestRPCErrorSurfaces(t *testing.T) {
	c := dialFake(t, &fakePoster{handler: echoServer()})
	if _, err := c.callRaw(context.Background(), "nope/nope"); err == nil {
		t.Fatal("want an rpc error")
	} else if !strings.Contains(err.Error(), "method not found") {
		t.Fatalf("err = %v", err)
	}
}

// callRaw is a test-only shim so the rpc-error path can be exercised without a
// typed wrapper for a method the server does not implement.
func (c *Client) callRaw(ctx context.Context, method string) (any, error) {
	var out any
	err := c.call(ctx, method, map[string]any{}, &out)
	return out, err
}

func TestDecodeFrame(t *testing.T) {
	cases := []struct {
		name, in, want string
		wantErr        bool
	}{
		{name: "plain json", in: `{"id":1,"result":{}}`, want: `{"id":1,"result":{}}`},
		{name: "sse single", in: "event: message\ndata: {\"id\":1,\"result\":1}\n\n", want: `{"id":1,"result":1}`},
		{
			name: "sse picks the response not the notification",
			in:   "data: {\"method\":\"notifications/progress\"}\n\ndata: {\"id\":2,\"result\":2}\n\n",
			want: `{"id":2,"result":2}`,
		},
		{name: "empty", in: "  ", wantErr: true},
		{name: "sse with no response", in: "data: {\"method\":\"x\"}\n\n", wantErr: true},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			got, err := decodeFrame([]byte(tc.in))
			if tc.wantErr {
				if err == nil {
					t.Fatalf("want error, got %q", got)
				}
				return
			}
			if err != nil {
				t.Fatal(err)
			}
			if string(got) != tc.want {
				t.Fatalf("got %q want %q", got, tc.want)
			}
		})
	}
}

func TestValidate(t *testing.T) {
	cases := []struct {
		name    string
		cfg     ServerConfig
		wantErr bool
	}{
		{name: "stdio ok", cfg: ServerConfig{Name: "a", Command: "echo"}},
		{name: "http ok", cfg: ServerConfig{Name: "a", URL: "http://x.test/mcp"}},
		{name: "no name", cfg: ServerConfig{Command: "echo"}, wantErr: true},
		{name: "spacey name", cfg: ServerConfig{Name: "a b", Command: "echo"}, wantErr: true},
		{name: "neither", cfg: ServerConfig{Name: "a"}, wantErr: true},
		{name: "both", cfg: ServerConfig{Name: "a", Command: "echo", URL: "http://x.test"}, wantErr: true},
		{name: "bad scheme", cfg: ServerConfig{Name: "a", URL: "file:///etc/passwd"}, wantErr: true},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			err := Validate([]ServerConfig{tc.cfg})
			if (err != nil) != tc.wantErr {
				t.Fatalf("err = %v, wantErr = %v", err, tc.wantErr)
			}
		})
	}
	if err := Validate([]ServerConfig{{Name: "a", Command: "x"}, {Name: "a", Command: "y"}}); err == nil {
		t.Error("duplicate names must be refused")
	}
}

func TestManagerOffWhenNothingEnabled(t *testing.T) {
	m, err := NewManager(nil, []ServerConfig{{Name: "a", Command: "echo"}}) // Enabled=false
	if err != nil {
		t.Fatal(err)
	}
	if !m.Empty() {
		t.Fatal("a server that is not enabled must not be wired")
	}
	m.Connect(context.Background())
	if got := m.Tools(); len(got) != 0 {
		t.Fatalf("tools = %+v", got)
	}
}

func TestManagerDiscoversAndCalls(t *testing.T) {
	p := &fakePoster{handler: echoServer()}
	m, err := NewManager(func(ServerConfig) (Poster, error) { return p, nil },
		[]ServerConfig{{Name: "fake", URL: "http://example.test/mcp", Enabled: true}})
	if err != nil {
		t.Fatal(err)
	}
	m.Connect(context.Background())
	defer m.Close()

	tools := m.Tools()
	if len(tools) != 2 {
		t.Fatalf("tools = %+v", tools)
	}
	out, err := m.Call(context.Background(), "fake", "read_thing", map[string]any{"q": "да"})
	if err != nil {
		t.Fatalf("call: %v", err)
	}
	if out != "read_thing:да" {
		t.Fatalf("out = %q", out)
	}
	// The discovered set is a second allowlist.
	if _, err := m.Call(context.Background(), "fake", "not_offered", nil); err == nil {
		t.Error("a tool the server does not offer must be refused")
	}
	if _, err := m.Call(context.Background(), "other", "read_thing", nil); err == nil {
		t.Error("an unconfigured server must be refused")
	}
	st := m.Status()
	if len(st) != 1 || !st[0].Connected || st[0].Transport != "http" || st[0].Tools != 2 {
		t.Fatalf("status = %+v", st)
	}
}

func TestManagerAllowToolsAndMaxTools(t *testing.T) {
	p := &fakePoster{handler: echoServer()}
	mk := func(cfg ServerConfig) *Manager {
		cfg.Name, cfg.URL, cfg.Enabled = "fake", "http://example.test/mcp", true
		m, err := NewManager(func(ServerConfig) (Poster, error) { return p, nil }, []ServerConfig{cfg})
		if err != nil {
			t.Fatal(err)
		}
		m.Connect(context.Background())
		return m
	}
	m := mk(ServerConfig{AllowTools: []string{"read_thing"}})
	defer m.Close()
	if got := m.Tools(); len(got) != 1 || got[0].Name != "read_thing" {
		t.Fatalf("allow_tools ignored: %+v", got)
	}
	if _, err := m.Call(context.Background(), "fake", "break_thing", nil); err == nil {
		t.Error("a tool excluded by allow_tools must be unreachable")
	}
	m2 := mk(ServerConfig{MaxTools: 1})
	defer m2.Close()
	if got := m2.Tools(); len(got) != 1 || got[0].Name != "break_thing" {
		t.Fatalf("max_tools should keep the first name-sorted tool: %+v", got)
	}
}

func TestManagerURLServerWithoutHTTPDoor(t *testing.T) {
	m, err := NewManager(nil, []ServerConfig{{Name: "fake", URL: "http://example.test/mcp", Enabled: true}})
	if err != nil {
		t.Fatal(err)
	}
	m.Connect(context.Background())
	st := m.Status()
	if len(st) != 1 || st[0].Connected || st[0].Err == "" {
		t.Fatalf("a url server with no poster must be recorded as failed: %+v", st)
	}
}

func TestManagerReconnectAfterFailure(t *testing.T) {
	var mu sync.Mutex
	fail := true
	m, err := NewManager(func(ServerConfig) (Poster, error) {
		mu.Lock()
		defer mu.Unlock()
		if fail {
			return nil, fmt.Errorf("down")
		}
		return &fakePoster{handler: echoServer()}, nil
	}, []ServerConfig{{Name: "fake", URL: "http://example.test/mcp", Enabled: true}})
	if err != nil {
		t.Fatal(err)
	}
	defer m.Close()
	m.Connect(context.Background())
	if m.Status()[0].Connected {
		t.Fatal("should be down")
	}
	mu.Lock()
	fail = false
	mu.Unlock()
	// Refresh honours the backoff, so pretend the last attempt was long ago.
	m.mu.Lock()
	m.conns["fake"].lastTry = time.Now().Add(-2 * DefaultReconnectEvery)
	m.mu.Unlock()
	m.Refresh(context.Background())
	if !m.Status()[0].Connected {
		t.Fatalf("should have reconnected: %+v", m.Status())
	}
}

func TestLocalNameAndCmd(t *testing.T) {
	cases := [][3]string{
		{"vikunja", "list_tasks", "vikunja_list_tasks"},
		{"Vikunja", "Get Task Details", "vikunja_get_task_details"},
		{"fs", "read-file", "fs_read_file"},
		{"", "search", "search"},
	}
	for _, c := range cases {
		if got := LocalName(c[0], c[1]); got != c[2] {
			t.Errorf("LocalName(%q,%q) = %q want %q", c[0], c[1], got, c[2])
		}
	}
	server, tool, ok := ParseCmd(Cmd("vikunja", "list_tasks"))
	if !ok || server != "vikunja" || tool != "list_tasks" {
		t.Fatalf("ParseCmd round-trip: %q %q %v", server, tool, ok)
	}
	for _, bad := range [][]string{nil, {"systemctl", "restart", "nginx"}, {"mcp", "vikunja"}, {"mcp", "", "x"}} {
		if _, _, ok := ParseCmd(bad); ok {
			t.Errorf("ParseCmd(%v) must not claim an ordinary tool row", bad)
		}
	}
	if Scope("vikunja") != "mcp:vikunja" {
		t.Error("scope")
	}
}

func TestBindPositional(t *testing.T) {
	cases := []struct {
		name, schema string
		args         []string
		mutating     bool
		want         map[string]any
		wantErr      bool
	}{
		{
			name: "no required runs with nothing",
			// A spare tail is fine: "покажи проекты пожалуйста" still lists them.
			schema: `{"type":"object","properties":{},"required":[]}`,
			args:   []string{"пожалуйста"},
			want:   map[string]any{},
		},
		{
			name:   "empty schema",
			schema: ``,
			want:   map[string]any{},
		},
		{
			name:   "one required string gets the tail",
			schema: `{"properties":{"q":{"type":"string"}},"required":["q"]}`,
			args:   []string{"почему", "небо", "синее"},
			want:   map[string]any{"q": "почему небо синее"},
		},
		{
			name:    "one required string with no tail",
			schema:  `{"properties":{"q":{"type":"string"}},"required":["q"]}`,
			wantErr: true,
		},
		{
			name:   "one required integer parses",
			schema: `{"properties":{"task_id":{"type":"integer"}},"required":["task_id"]}`,
			args:   []string{"251"},
			want:   map[string]any{"task_id": float64(251)},
		},
		{
			name:    "one required integer with words",
			schema:  `{"properties":{"task_id":{"type":"integer"}},"required":["task_id"]}`,
			args:    []string{"двести", "пятьдесят", "один"},
			wantErr: true,
		},
		{
			name:    "two required is refused rather than guessed",
			schema:  `{"properties":{"a":{"type":"string"},"b":{"type":"string"}},"required":["a","b"]}`,
			args:    []string{"что-то"},
			wantErr: true,
		},
		{
			name:    "one required object is refused",
			schema:  `{"properties":{"payload":{"type":"object"}},"required":["payload"]}`,
			args:    []string{"что-то"},
			wantErr: true,
		},
		{
			// Learned from Vikunja's update_task: required ["task_id"], every
			// other field optional, so one guessed argument blanks the rest.
			name:     "one required on a mutating tool is refused",
			schema:   `{"properties":{"task_id":{"type":"integer"}},"required":["task_id"]}`,
			args:     []string{"251"},
			mutating: true,
			wantErr:  true,
		},
		{
			// Nothing was guessed, so there is nothing to get wrong. It still
			// goes through the confirm turn upstream.
			name:     "no required on a mutating tool still runs",
			schema:   `{"properties":{},"required":[]}`,
			mutating: true,
			want:     map[string]any{},
		},
		{
			name:    "unreadable schema",
			schema:  `not json`,
			wantErr: true,
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			got, err := bindPositional(json.RawMessage(tc.schema), tc.args, !tc.mutating)
			if tc.wantErr {
				if err == nil {
					t.Fatalf("want an error, got %v", got)
				}
				return
			}
			if err != nil {
				t.Fatal(err)
			}
			if fmt.Sprint(got) != fmt.Sprint(tc.want) {
				t.Fatalf("got %v want %v", got, tc.want)
			}
		})
	}
}

func TestCallPositionalThroughManager(t *testing.T) {
	p := &fakePoster{handler: echoServer()}
	m, err := NewManager(func(ServerConfig) (Poster, error) { return p, nil },
		[]ServerConfig{{Name: "fake", URL: "http://example.test/mcp", Enabled: true}})
	if err != nil {
		t.Fatal(err)
	}
	m.Connect(context.Background())
	defer m.Close()
	// echoServer's tools declare no required properties.
	out, err := m.CallPositional(context.Background(), "fake", "read_thing", []string{"хвост"})
	if err != nil {
		t.Fatalf("call: %v", err)
	}
	if out != "read_thing:<nil>" {
		t.Fatalf("out = %q", out)
	}
	if _, err := m.CallPositional(context.Background(), "fake", "absent", nil); err == nil {
		t.Error("an unknown tool must be refused")
	}
}