krow_backend/go-api/internal/mcpserver/mcpserver_test.go

package mcpserver

import (
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/krow/krow-backend/go-api/internal/authctx"
	"github.com/krow/krow-backend/go-api/internal/runtime"
	"github.com/krow/krow-backend/go-api/internal/tools"
)

// The tool registry these tests run against is the REAL one — the same
// construction the service uses — with a nil database and a nil retriever.
//
// That is sound because nothing here dispatches: Phase 1's tools/call refuses
// before reaching a handler, and every other assertion is about metadata, which
// is fixed at registration. Testing against a hand-built fixture registry would
// be testing a copy of the thing under test: the whole claim of this package is
// "what MCP publishes is what the registry holds", and a fixture would let that
// claim pass while being false of the real set.
// Both dependencies are nil: the handlers capture them in closures and nothing
// here reaches a handler, so nothing dereferences them. If a future test does
// dispatch, this will nil-panic loudly rather than quietly reading a database
// it should not have.
func testRegistry(t *testing.T) *tools.Registry {
	t.Helper()
	return runtime.DefaultTools(nil, nil)
}

// newServer wires the real registry to a test-only token authenticator.
//
// PHASE 2 NOTE: tools/list and tools/call now require a bearer token, so this
// fixture authenticates. Not one assertion in this file changed — only the
// fixture gained a credential. The behaviour change is the point of Phase 2 and
// is asserted directly in auth_test.go (TestBearerRejection), rather than being
// papered over here.
func newServer(t *testing.T) *Server {
	t.Helper()
	return New(testRegistry(t), &fakeTokens{byToken: map[string]authctx.Identity{
		validToken: identityFor("org-test", "admin"),
	}}, slog.New(slog.NewTextHandler(io.Discard, nil)))
}

// post sends one raw body to the MCP endpoint and returns the recorder.
func post(t *testing.T, s *Server, body string) *httptest.ResponseRecorder {
	t.Helper()
	req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+validToken)
	rec := httptest.NewRecorder()
	s.Handler().ServeHTTP(rec, req)
	return rec
}

// decode reads a JSON-RPC response, failing the test if the envelope is wrong.
func decode(t *testing.T, rec *httptest.ResponseRecorder) response {
	t.Helper()
	var resp response
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("response was not JSON: %v\nbody: %s", err, rec.Body.String())
	}
	if resp.JSONRPC != jsonRPCVersion {
		t.Fatalf("jsonrpc = %q, want %q", resp.JSONRPC, jsonRPCVersion)
	}
	return resp
}

// resultInto re-decodes a successful result into dst.
func resultInto(t *testing.T, rec *httptest.ResponseRecorder, dst any) {
	t.Helper()
	var envelope struct {
		Result json.RawMessage `json:"result"`
		Error  *rpcError       `json:"error"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &envelope); err != nil {
		t.Fatalf("response was not JSON: %v", err)
	}
	if envelope.Error != nil {
		t.Fatalf("expected a result, got error %d: %s", envelope.Error.Code, envelope.Error.Message)
	}
	if err := json.Unmarshal(envelope.Result, dst); err != nil {
		t.Fatalf("result did not decode: %v\nresult: %s", err, envelope.Result)
	}
}

/* ── 1. initialize ──────────────────────────────────────────────────────── */

func TestInitializeSucceeds(t *testing.T) {
	s := newServer(t)
	rec := post(t, s, `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{
		"protocolVersion":"2025-06-18",
		"capabilities":{},
		"clientInfo":{"name":"test-client","version":"1.0"}}}`)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200", rec.Code)
	}
	if ct := rec.Header().Get("Content-Type"); !strings.HasPrefix(ct, "application/json") {
		t.Errorf("Content-Type = %q, want application/json", ct)
	}

	var out initializeResult
	resultInto(t, rec, &out)

	if out.ProtocolVersion != ProtocolVersion {
		t.Errorf("protocolVersion = %q, want %q", out.ProtocolVersion, ProtocolVersion)
	}
	if out.ServerInfo["name"] != ServerName {
		t.Errorf("serverInfo.name = %v, want %q", out.ServerInfo["name"], ServerName)
	}
	// Exactly one capability, because exactly one is implemented. Advertising
	// resources or prompts would promise methods that answer method-not-found.
	if _, ok := out.Capabilities["tools"]; !ok {
		t.Error("capabilities.tools is missing")
	}
	for _, unimplemented := range []string{"resources", "prompts", "sampling", "logging"} {
		if _, ok := out.Capabilities[unimplemented]; ok {
			t.Errorf("capabilities advertises %q, which is not implemented", unimplemented)
		}
	}
}

func TestInitializeEchoesTheRequestID(t *testing.T) {
	s := newServer(t)
	// A string id, to prove ids are echoed verbatim rather than coerced.
	rec := post(t, s, `{"jsonrpc":"2.0","id":"abc-123","method":"initialize","params":{}}`)
	resp := decode(t, rec)
	if string(resp.ID) != `"abc-123"` {
		t.Errorf("id = %s, want \"abc-123\"", resp.ID)
	}
}

/* ── 2 & 3. tools/list and exact exposure ───────────────────────────────── */

// wantExposed is the Phase 0 read-only MVP set, written out in full.
//
// Deliberately a literal rather than a filter over the registry: a test that
// derived the expectation the same way the code does would pass no matter what
// the rule became. This is the list a human agreed to, and it is the thing that
// should fail when the rule changes.
var wantExposed = []string{
	"activity_breakdown",
	"activity_signals",
	"available_workers",
	"candidates_awaiting",
	"candidates_quality",
	"hires_performance",
	"hires_recent",
	"open_positions",
	"operations_risk",
	"positions_risk",
	"talent_pool",
	"workforce_attendance",
	"workforce_coverage",
	"workforce_overtime",
	"workforce_training",
	"workspace_summary",
}

func TestToolsListExposesExactlyTheReadOnlyMVPSet(t *testing.T) {
	s := newServer(t)
	rec := post(t, s, `{"jsonrpc":"2.0","id":2,"method":"tools/list"}`)

	var out toolsListResult
	resultInto(t, rec, &out)

	got := make([]string, 0, len(out.Tools))
	for _, tool := range out.Tools {
		got = append(got, tool.Name)
	}

	if len(got) != len(wantExposed) {
		t.Fatalf("exposed %d tools, want %d\n got: %v\nwant: %v",
			len(got), len(wantExposed), got, wantExposed)
	}
	for i := range wantExposed {
		if got[i] != wantExposed[i] {
			t.Errorf("tool[%d] = %q, want %q", i, got[i], wantExposed[i])
		}
	}
}

func TestToolsListWorksWithoutParams(t *testing.T) {
	s := newServer(t)
	// No params key at all — every exposed tool's arguments are optional, so a
	// bare list call must work.
	rec := post(t, s, `{"jsonrpc":"2.0","id":3,"method":"tools/list"}`)
	var out toolsListResult
	resultInto(t, rec, &out)
	if len(out.Tools) == 0 {
		t.Fatal("tools/list returned nothing")
	}
}

/* ── 4 & 5. write tools and knowledge_search are not exposed ────────────── */

func TestNoWriteToolIsExposed(t *testing.T) {
	s := newServer(t)
	rec := post(t, s, `{"jsonrpc":"2.0","id":4,"method":"tools/list"}`)
	var out toolsListResult
	resultInto(t, rec, &out)

	reg := testRegistry(t)
	for _, published := range out.Tools {
		tool, ok := reg.Get(published.Name)
		if !ok {
			t.Fatalf("tools/list published %q, which is not in the registry", published.Name)
		}
		if tool.Effect != tools.EffectRead {
			t.Errorf("%q is exposed with effect %q; only read tools may be exposed",
				tool.Name, tool.Effect)
		}
		if tool.RequiresConfirmation {
			t.Errorf("%q is exposed and requires confirmation, which this surface cannot obtain",
				tool.Name)
		}
		// The annotation must agree with the registry, or a client shows a
		// person the wrong thing before approving.
		if published.Annotations == nil || !published.Annotations.ReadOnlyHint {
			t.Errorf("%q is not annotated readOnlyHint", tool.Name)
		}
		if published.Annotations != nil && published.Annotations.DestructiveHint {
			t.Errorf("%q is annotated destructiveHint", tool.Name)
		}
	}
}

func TestNamedExclusionsAreNotExposed(t *testing.T) {
	s := newServer(t)
	rec := post(t, s, `{"jsonrpc":"2.0","id":5,"method":"tools/list"}`)
	var out toolsListResult
	resultInto(t, rec, &out)

	published := map[string]bool{}
	for _, tool := range out.Tools {
		published[tool.Name] = true
	}
	// The two writes and the one deferred read. Named explicitly, because these
	// three are the ones a future change is most likely to let through.
	for _, forbidden := range []string{"assign_worker", "move_application", "knowledge_search"} {
		if published[forbidden] {
			t.Errorf("%q must not be exposed", forbidden)
		}
	}
}

// The registry must still hold the excluded tools — they are withheld from this
// surface, not removed from KROW. A test that only checked absence would pass
// if somebody deleted them.
func TestExcludedToolsStillExistInTheRegistry(t *testing.T) {
	reg := testRegistry(t)
	for _, name := range []string{"assign_worker", "move_application", "knowledge_search"} {
		if _, ok := reg.Get(name); !ok {
			t.Errorf("%q is missing from the registry entirely", name)
		}
	}
}

/* ── 6. schemas come from the real registry ─────────────────────────────── */

func TestToolSchemaIsSourcedFromTheRegistry(t *testing.T) {
	s := newServer(t)
	rec := post(t, s, `{"jsonrpc":"2.0","id":6,"method":"tools/list"}`)
	var out toolsListResult
	resultInto(t, rec, &out)

	reg := testRegistry(t)
	for _, published := range out.Tools {
		tool, _ := reg.Get(published.Name)

		if published.Description != tool.Description {
			t.Errorf("%q description differs from the registry's", tool.Name)
		}
		if published.InputSchema == nil {
			t.Fatalf("%q published a nil inputSchema", tool.Name)
		}
		// Every exposed tool declares a schema, so the published one must be
		// the registry's own map and not the empty-schema fallback.
		if tool.InputSchema == nil {
			t.Errorf("%q has no InputSchema in the registry; expected every exposed tool to declare one", tool.Name)
			continue
		}
		wantJSON, _ := json.Marshal(tool.InputSchema)
		gotJSON, _ := json.Marshal(published.InputSchema)
		if string(wantJSON) != string(gotJSON) {
			t.Errorf("%q schema differs from the registry's\n got: %s\nwant: %s",
				tool.Name, gotJSON, wantJSON)
		}
		// A published schema must be a JSON Schema object, or a client may
		// refuse the whole list.
		if published.InputSchema["type"] != "object" {
			t.Errorf("%q schema type = %v, want \"object\"", tool.Name, published.InputSchema["type"])
		}
	}
}

// No exposed tool may take a tenant or principal identifier as an argument.
// An argument is something a client chooses; identity is not the client's to
// choose. Enforced here rather than by review, because the cost of missing it
// once is cross-tenant access.
func TestNoToolAcceptsATenantOrPrincipalArgument(t *testing.T) {
	s := newServer(t)
	rec := post(t, s, `{"jsonrpc":"2.0","id":7,"method":"tools/list"}`)
	var out toolsListResult
	resultInto(t, rec, &out)

	forbidden := []string{
		"org_id", "orgid", "organization_id", "organisation_id",
		"tenant_id", "tenantid", "tenant",
		"user_id", "userid", "principal", "principal_id",
		"account_id", "caller", "caller_id", "on_behalf_of", "impersonate",
	}

	for _, tool := range out.Tools {
		props, ok := tool.InputSchema["properties"].(map[string]any)
		if !ok {
			continue
		}
		for field := range props {
			lower := strings.ToLower(field)
			for _, bad := range forbidden {
				if lower == bad {
					t.Errorf("%q accepts %q as an argument; identity must come from the "+
						"authenticated principal, never from the request", tool.Name, field)
				}
			}
		}
	}
}

/* ── 7 & 8. unknown method, unknown tool ────────────────────────────────── */

func TestUnknownMethodReturnsMethodNotFound(t *testing.T) {
	s := newServer(t)
	rec := post(t, s, `{"jsonrpc":"2.0","id":8,"method":"resources/list"}`)

	if rec.Code != http.StatusOK {
		t.Errorf("status = %d, want 200 — a JSON-RPC error is still a successful HTTP exchange", rec.Code)
	}
	resp := decode(t, rec)
	if resp.Error == nil {
		t.Fatal("expected an error")
	}
	if resp.Error.Code != codeMethodNotFound {
		t.Errorf("code = %d, want %d", resp.Error.Code, codeMethodNotFound)
	}
}

func TestUnknownToolIsReportedAsAToolError(t *testing.T) {
	s := newServer(t)
	rec := post(t, s, `{"jsonrpc":"2.0","id":9,"method":"tools/call",
		"params":{"name":"no_such_tool","arguments":{}}}`)

	var out toolsCallResult
	resultInto(t, rec, &out)
	if !out.IsError {
		t.Error("expected isError on an unknown tool")
	}
	if len(out.Content) == 0 || !strings.Contains(out.Content[0].Text, "mcp.unknown_tool") {
		t.Errorf("expected mcp.unknown_tool, got %+v", out.Content)
	}
}

// An unexposed tool must be indistinguishable from a non-existent one, or the
// error messages become an inventory of what this surface is withholding.
func TestUnexposedToolIsIndistinguishableFromUnknown(t *testing.T) {
	s := newServer(t)

	unknown := post(t, s, `{"jsonrpc":"2.0","id":10,"method":"tools/call","params":{"name":"definitely_not_a_tool"}}`)
	withheld := post(t, s, `{"jsonrpc":"2.0","id":10,"method":"tools/call","params":{"name":"assign_worker"}}`)

	var a, b toolsCallResult
	resultInto(t, unknown, &a)
	resultInto(t, withheld, &b)

	codeOf := func(r toolsCallResult) string {
		var payload struct {
			Error struct {
				Code string `json:"code"`
			} `json:"error"`
		}
		_ = json.Unmarshal([]byte(r.Content[0].Text), &payload)
		return payload.Error.Code
	}
	if codeOf(a) != codeOf(b) {
		t.Errorf("an unexposed tool is distinguishable from an unknown one: %q vs %q",
			codeOf(a), codeOf(b))
	}
	if a.IsError != b.IsError {
		t.Error("isError differs between unknown and unexposed tools")
	}
}

/* ── 9 & 10. malformed JSON and params ──────────────────────────────────── */

func TestMalformedJSONReturnsParseError(t *testing.T) {
	s := newServer(t)
	rec := post(t, s, `{"jsonrpc":"2.0","id":1,"method":`)

	resp := decode(t, rec)
	if resp.Error == nil || resp.Error.Code != codeParseError {
		t.Fatalf("want parse error %d, got %+v", codeParseError, resp.Error)
	}
}

func TestEnvelopeValidation(t *testing.T) {
	for name, tc := range map[string]struct {
		body string
		want int
	}{
		"missing jsonrpc":   {`{"id":1,"method":"tools/list"}`, codeInvalidRequest},
		"wrong jsonrpc":     {`{"jsonrpc":"1.0","id":1,"method":"tools/list"}`, codeInvalidRequest},
		"missing method":    {`{"jsonrpc":"2.0","id":1}`, codeInvalidRequest},
		"empty body":        {``, codeInvalidRequest},
		"object id":         {`{"jsonrpc":"2.0","id":{"a":1},"method":"tools/list"}`, codeInvalidRequest},
		"not an object":     {`"a string"`, codeParseError},
		"params as array":   {`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":[1,2]}`, codeInvalidParams},
		"params wrong type": {`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":42}}`, codeInvalidParams},
		"missing tool name": {`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{}}`, codeInvalidParams},
	} {
		t.Run(name, func(t *testing.T) {
			rec := post(t, newServer(t), tc.body)
			resp := decode(t, rec)
			if resp.Error == nil {
				t.Fatalf("expected an error, got %s", rec.Body.String())
			}
			if resp.Error.Code != tc.want {
				t.Errorf("code = %d, want %d (%s)", resp.Error.Code, tc.want, resp.Error.Message)
			}
		})
	}
}

/* ── 11. batching is explicitly rejected ────────────────────────────────── */

func TestBatchRequestsAreExplicitlyRejected(t *testing.T) {
	s := newServer(t)
	rec := post(t, s, `[{"jsonrpc":"2.0","id":1,"method":"tools/list"},
	                    {"jsonrpc":"2.0","id":2,"method":"tools/list"}]`)

	resp := decode(t, rec)
	if resp.Error == nil {
		t.Fatal("a batch must be refused")
	}
	if resp.Error.Code != codeInvalidRequest {
		t.Errorf("code = %d, want %d", resp.Error.Code, codeInvalidRequest)
	}
	// The refusal must say WHY, so a client can fall back to sending singly
	// rather than retrying the same batch forever.
	detail, _ := resp.Error.Data.(string)
	if !strings.Contains(detail, "batch") {
		t.Errorf("the refusal does not mention batching: %q", detail)
	}
}

/* ── 12. no panic on malformed input ────────────────────────────────────── */

func TestNoPanicOnHostileInput(t *testing.T) {
	// Each of these has crashed a hand-written JSON-RPC server somewhere.
	bodies := []string{
		``,
		` `,
		`null`,
		`[]`,
		`[[[[[]]]]]`,
		`{`,
		`}`,
		`{"jsonrpc":"2.0"}`,
		`{"jsonrpc":null,"id":null,"method":null}`,
		`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":null}`,
		`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":""}}`,
		`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"activity_breakdown","arguments":"not-an-object"}}`,
		`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"activity_breakdown","arguments":[]}}`,
		`{"jsonrpc":"2.0","id":[1,2,3],"method":"tools/list"}`,
		`{"jsonrpc":"2.0","id":1,"method":"` + strings.Repeat("A", 10_000) + `"}`,
		`{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"cursor":` + strings.Repeat("[", 500) + `}}`,
		"\x00\x01\x02",
	}

	s := newServer(t)
	for i, body := range bodies {
		// A panic escaping here fails the test by crashing it, which is the
		// assertion: the handler must answer every one of these.
		rec := post(t, s, body)
		if rec.Code < 200 || rec.Code >= 600 {
			t.Errorf("body %d produced status %d", i, rec.Code)
		}
		if rec.Body.Len() > 0 && rec.Code != http.StatusAccepted {
			var resp response
			if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
				t.Errorf("body %d produced a non-JSON response: %s", i, rec.Body.String())
			}
		}
	}
}

/* ── Transport ──────────────────────────────────────────────────────────── */

func TestOnlyPOSTIsAccepted(t *testing.T) {
	s := newServer(t)
	for _, method := range []string{http.MethodGet, http.MethodPut, http.MethodDelete, http.MethodPatch} {
		req := httptest.NewRequest(method, "/mcp", nil)
		rec := httptest.NewRecorder()
		s.Handler().ServeHTTP(rec, req)

		if rec.Code != http.StatusMethodNotAllowed {
			t.Errorf("%s: status = %d, want 405", method, rec.Code)
		}
		if allow := rec.Header().Get("Allow"); allow != http.MethodPost {
			t.Errorf("%s: Allow = %q, want POST", method, allow)
		}
	}
}

func TestNonJSONContentTypeIsRefused(t *testing.T) {
	s := newServer(t)
	req := httptest.NewRequest(http.MethodPost, "/mcp",
		strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	rec := httptest.NewRecorder()
	s.Handler().ServeHTTP(rec, req)

	if rec.Code != http.StatusUnsupportedMediaType {
		t.Errorf("status = %d, want 415", rec.Code)
	}
}

func TestOversizedBodyIsRefused(t *testing.T) {
	s := newServer(t)
	huge := `{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"cursor":"` +
		strings.Repeat("x", MaxRequestBytes+1) + `"}}`
	rec := post(t, s, huge)

	if rec.Code != http.StatusRequestEntityTooLarge {
		t.Errorf("status = %d, want 413", rec.Code)
	}
}

func TestNotificationGetsNoBody(t *testing.T) {
	s := newServer(t)
	// No id — a notification. The spec forbids a response.
	rec := post(t, s, `{"jsonrpc":"2.0","method":"notifications/initialized"}`)

	if rec.Code != http.StatusAccepted {
		t.Errorf("status = %d, want 202", rec.Code)
	}
	if rec.Body.Len() != 0 {
		t.Errorf("a notification was answered with a body: %s", rec.Body.String())
	}
}

/* ── Phase 1 boundary ───────────────────────────────────────────────────── */

// Handle must refuse a nil identity even when called directly.
//
// The transport answers 401 before this point in the ordinary case, so this is
// the SECOND gate: a future caller of Handle that forgets to authenticate must
// fail closed rather than dispatch as nobody. It exists to fail loudly if
// somebody later supplies a default identity to "make it work".
func TestHandleRefusesANilIdentity(t *testing.T) {
	s := newServer(t)
	req := request{
		JSONRPC: jsonRPCVersion,
		ID:      json.RawMessage(`11`),
		Method:  "tools/call",
		Params:  json.RawMessage(`{"name":"activity_breakdown","arguments":{}}`),
	}
	result, rpcErr := s.Handle(context.Background(), nil, req)
	if rpcErr != nil {
		t.Fatalf("unexpected rpc error: %v", rpcErr)
	}
	out, ok := result.(toolsCallResult)
	if !ok {
		t.Fatalf("unexpected result type %T", result)
	}
	if !out.IsError || !strings.Contains(out.Content[0].Text, "mcp.unauthenticated") {
		t.Errorf("expected mcp.unauthenticated, got: %+v", out.Content)
	}
}

func TestPingIsCheap(t *testing.T) {
	s := newServer(t)
	rec := post(t, s, `{"jsonrpc":"2.0","id":12,"method":"ping"}`)
	var out map[string]any
	resultInto(t, rec, &out)
	if len(out) != 0 {
		t.Errorf("ping returned %v, want an empty result", out)
	}
}