projax/mcp/mcp_test.go

package mcp

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// rpcJSON builds a JSON-RPC request body.
func rpcJSON(t *testing.T, id any, method string, params any) []byte {
	t.Helper()
	out, err := json.Marshal(map[string]any{
		"jsonrpc": "2.0",
		"id":      id,
		"method":  method,
		"params":  params,
	})
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	return out
}

func doRPC(t *testing.T, srv *Server, body []byte, token string) (*http.Response, []byte) {
	t.Helper()
	mux := http.NewServeMux()
	srv.Routes(mux)
	s := httptest.NewServer(http.StripPrefix("", mux))
	t.Cleanup(s.Close)
	req, _ := http.NewRequest(http.MethodPost, s.URL+"/rpc", bytes.NewReader(body))
	if token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("do: %v", err)
	}
	t.Cleanup(func() { resp.Body.Close() })
	raw, _ := io.ReadAll(resp.Body)
	return resp, raw
}

func TestInitializeAndToolsList(t *testing.T) {
	srv := New("projax-test", "0.0.1", "", nil)
	srv.Register(Tool{
		Name:        "echo",
		Description: "echoes input.message",
		InputSchema: json.RawMessage(`{"type":"object","required":["message"]}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (any, error) {
			var in struct {
				Message string `json:"message"`
			}
			if err := json.Unmarshal(raw, &in); err != nil {
				return nil, err
			}
			return map[string]any{"echo": in.Message}, nil
		},
	})

	_, body := doRPC(t, srv, rpcJSON(t, 1, "initialize", map[string]any{}), "")
	if !strings.Contains(string(body), `"protocolVersion":"2024-11-05"`) {
		t.Fatalf("initialize body missing protocolVersion: %s", body)
	}
	if !strings.Contains(string(body), `"name":"projax-test"`) {
		t.Errorf("initialize missing server name: %s", body)
	}

	_, body = doRPC(t, srv, rpcJSON(t, 2, "tools/list", map[string]any{}), "")
	if !strings.Contains(string(body), `"name":"echo"`) {
		t.Fatalf("tools/list missing echo tool: %s", body)
	}
}

func TestToolsCallSuccessAndError(t *testing.T) {
	srv := New("p", "1", "", nil)
	srv.Register(Tool{
		Name: "echo",
		Handler: func(ctx context.Context, raw json.RawMessage) (any, error) {
			return map[string]any{"got": string(raw)}, nil
		},
	})
	srv.Register(Tool{
		Name: "boom",
		Handler: func(ctx context.Context, raw json.RawMessage) (any, error) {
			return nil, &ToolError{Msg: "kaboom"}
		},
	})

	_, body := doRPC(t, srv, rpcJSON(t, 3, "tools/call", map[string]any{
		"name":      "echo",
		"arguments": map[string]any{"x": 1},
	}), "")
	if !strings.Contains(string(body), `"isError":false`) {
		t.Fatalf("success response missing isError:false: %s", body)
	}
	// Tool results land inside a content[].text block as a JSON-stringified
	// payload, so the inner double-quote bytes show up double-escaped. We
	// just check the original key is present somewhere in the response.
	if !strings.Contains(string(body), "got") {
		t.Errorf("echo did not include 'got' key: %s", body)
	}

	_, body = doRPC(t, srv, rpcJSON(t, 4, "tools/call", map[string]any{
		"name":      "boom",
		"arguments": map[string]any{},
	}), "")
	if !strings.Contains(string(body), `"isError":true`) {
		t.Fatalf("error response missing isError:true: %s", body)
	}
	if !strings.Contains(string(body), "kaboom") {
		t.Errorf("error response missing message: %s", body)
	}
}

func TestAuthBearerRequired(t *testing.T) {
	srv := New("p", "1", "s3cr3t", nil)

	resp, _ := doRPC(t, srv, rpcJSON(t, 1, "initialize", map[string]any{}), "")
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("missing token: expected 401, got %d", resp.StatusCode)
	}

	resp, _ = doRPC(t, srv, rpcJSON(t, 1, "initialize", map[string]any{}), "wrong")
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("wrong token: expected 401, got %d", resp.StatusCode)
	}

	resp, body := doRPC(t, srv, rpcJSON(t, 1, "initialize", map[string]any{}), "s3cr3t")
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("good token: expected 200, got %d (%s)", resp.StatusCode, body)
	}
}

func TestUnknownMethod(t *testing.T) {
	srv := New("p", "1", "", nil)
	_, body := doRPC(t, srv, rpcJSON(t, 1, "fly/me/to/the/moon", map[string]any{}), "")
	if !strings.Contains(string(body), `"code":-32601`) {
		t.Fatalf("expected method-not-found error, got %s", body)
	}
}

func TestNotificationsInitializedNoResponse(t *testing.T) {
	srv := New("p", "1", "", nil)
	// Notification — no id field.
	body, _ := json.Marshal(map[string]any{
		"jsonrpc": "2.0",
		"method":  "notifications/initialized",
	})
	resp, raw := doRPC(t, srv, body, "")
	if resp.StatusCode != http.StatusNoContent {
		t.Fatalf("notif: expected 204, got %d (%s)", resp.StatusCode, raw)
	}
	if len(raw) != 0 {
		t.Errorf("notif: expected empty body, got %s", raw)
	}
}