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

package mcpserver

import (
	"context"
	"encoding/json"
	"log/slog"
	"time"

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

// ProtocolVersion is the MCP revision this server implements.
//
// Echoed back from initialize. A client asking for a different revision is not
// refused: the spec's negotiation is that the server states what it speaks and
// the client decides whether it can work with that. Refusing would turn a
// version skew into an outage where it is usually a compatible difference.
const ProtocolVersion = "2025-06-18"

// ServerName and ServerVersion identify this implementation to a client.
const (
	ServerName    = "krow-mcp"
	ServerVersion = "0.1.0"
)

// maxToolArgumentBytes bounds one tool call's arguments.
//
// The same reasoning as runs.go's maxRunRequestBytes: arguments are a handful
// of scalars against a schema that sets additionalProperties:false, so anything
// large is either a mistake or an attempt to push text past the tool layer into
// a prompt. The transport bounds the whole body too; this bounds the part that
// reaches a handler.
const maxToolArgumentBytes = 64 << 10

// Server answers MCP methods against the existing tool registry.
//
// It holds a *tools.Registry and nothing else that matters. There is no second
// registry, no adapter table and no per-tool code in this package: what MCP
// publishes is what the registry holds, filtered by the rule in tools.go.
type Server struct {
	reg *tools.Registry
	log *slog.Logger

	// tokens resolves a bearer credential into an identity. Nil means this
	// surface cannot authenticate anyone, and every call is refused — see
	// ErrNoAuthenticator. Failing closed is the only safe default for a field
	// whose absence would otherwise mean "let everyone in".
	tokens TokenAuthenticator

	// resourceMetadataURL is where a 401 points a client so it can begin
	// discovery. Empty means the challenge carries no pointer, which is a
	// valid but less useful 401: a client then has nowhere to look.
	resourceMetadataURL string

	// orgLimiter bounds tool calls per organisation.
	//
	// HERE rather than in the HTTP middleware, and that placement is the whole
	// point: an organisation is not knowable until the bearer token has been
	// resolved to a user and that user's row read. A middleware running before
	// authentication could only key by something the CLIENT supplied, which is
	// precisely the identity this surface refuses to trust.
	//
	// Nil means no per-organisation ceiling, which is the correct default for
	// a deployment that has not configured one.
	orgLimiter OrgLimiter
}

// OrgLimiter bounds how much one organisation may ask for.
//
// Takes an org id that the caller has already established from an authenticated
// identity. It cannot be handed anything from a request, because the only
// caller is dispatch, which has an authctx.Identity and nothing else.
type OrgLimiter interface {
	// AllowOrg reports whether this organisation may make another call, and
	// how long until its window rolls over.
	AllowOrg(ctx context.Context, orgID string) (allowed bool, retryAfter time.Duration, err error)
}

// WithOrgLimiter installs the per-organisation ceiling.
func (s *Server) WithOrgLimiter(l OrgLimiter) *Server {
	s.orgLimiter = l
	return s
}

// WithResourceMetadataURL sets the RFC 9728 document a 401 points at.
//
// Supplied by the caller rather than derived here, because this package does
// not know its own deployment's URLs and must not invent them. A hardcoded
// hostname would be one deployment's identity baked into every other one.
func (s *Server) WithResourceMetadataURL(u string) *Server {
	s.resourceMetadataURL = u
	return s
}

// challenge builds the WWW-Authenticate header for a 401.
//
// RFC 9728 section 5.1: the client reads `resource_metadata` from here to find
// the protected-resource document, and from there the authorization server.
// Without the parameter a compliant client has a 401 and nowhere to go, which
// is why this is the difference between "authentication failed" and "here is
// how to authenticate".
func (s *Server) challenge() string {
	c := `Bearer realm="` + ServerName + `"`
	if s.resourceMetadataURL != "" {
		c += `, resource_metadata="` + s.resourceMetadataURL + `"`
	}
	return c
}

// New builds a server over an existing registry.
//
// The registry is the one the rest of the service already built — the caller
// passes runtime.DefaultTools(...)'s result, the same value the HTTP server
// uses for its author catalogue. Taking it as a parameter rather than building
// one here is what guarantees there is only ever one.
func New(reg *tools.Registry, tokens TokenAuthenticator, log *slog.Logger) *Server {
	if log == nil {
		log = slog.Default()
	}
	return &Server{reg: reg, tokens: tokens, log: log}
}

// Handle dispatches one parsed JSON-RPC request on behalf of an identity.
//
// The identity is a PARAMETER, not something read from the context, and that is
// the security property rather than a style choice. If this function resolved
// the caller from ctx, then mounting the endpoint behind the cookie middleware
// would make a browser session sufficient to call MCP tools — the middleware
// puts an Identity in the context, and this code would find it. Taking it as an
// argument means only the MCP transport's own bearer authentication can supply
// one. See auth.go.
//
// A nil identity means unauthenticated. The three handshake methods are allowed
// without one; tools/call is not.
//
// Returns a result or an error, never both. Notifications are handled by the
// transport, which discards whatever comes back.
func (s *Server) Handle(ctx context.Context, ident *authctx.Identity, req request) (any, *rpcError) {
	switch req.Method {
	case "initialize":
		return s.handleInitialize(req.Params)
	case "notifications/initialized":
		// The client telling us it is ready. Nothing to do, and answering is
		// not required — it arrives as a notification.
		return map[string]any{}, nil
	case "ping":
		// Cheap liveness, defined by the spec as an empty result. Costs nothing
		// and saves a client from using tools/list as a heartbeat.
		return map[string]any{}, nil
	case "tools/list":
		return s.handleToolsList(req.Params)
	case "tools/call":
		return s.handleToolsCall(ctx, ident, req.Params)
	default:
		return nil, errMethodNotFound(req.Method)
	}
}

/* ── initialize ─────────────────────────────────────────────────────────── */

type initializeParams struct {
	ProtocolVersion string          `json:"protocolVersion"`
	Capabilities    json.RawMessage `json:"capabilities"`
	ClientInfo      struct {
		Name    string `json:"name"`
		Version string `json:"version"`
	} `json:"clientInfo"`
}

type initializeResult struct {
	ProtocolVersion string         `json:"protocolVersion"`
	Capabilities    map[string]any `json:"capabilities"`
	ServerInfo      map[string]any `json:"serverInfo"`
	Instructions    string         `json:"instructions,omitempty"`
}

// handleInitialize answers the opening handshake.
//
// Declares exactly one capability, because exactly one is implemented. A server
// that advertised resources or prompts here would be promising methods that
// answer method-not-found, and a client would reasonably call them.
//
// listChanged is false: the tool set is fixed at process start by the registry,
// so there is no change to notify anyone about.
func (s *Server) handleInitialize(raw json.RawMessage) (any, *rpcError) {
	var p initializeParams
	if err := decodeParams(raw, &p); err != nil {
		return nil, err
	}

	s.log.Info("mcp initialize",
		"client_name", p.ClientInfo.Name,
		"client_version", p.ClientInfo.Version,
		"client_protocol", p.ProtocolVersion,
		"server_protocol", ProtocolVersion)

	return initializeResult{
		ProtocolVersion: ProtocolVersion,
		Capabilities: map[string]any{
			"tools": map[string]any{"listChanged": false},
		},
		ServerInfo: map[string]any{
			"name":    ServerName,
			"version": ServerVersion,
		},
		Instructions: "Read-only access to KROW workforce and hiring data. " +
			"Every call is scoped to the authenticated user's organisation and role; " +
			"results are structured data for you to summarise, not prose.",
	}, nil
}

/* ── tools/list ─────────────────────────────────────────────────────────── */

type toolsListResult struct {
	Tools []mcpTool `json:"tools"`
}

// handleToolsList publishes the exposed tools, straight from the registry.
func (s *Server) handleToolsList(raw json.RawMessage) (any, *rpcError) {
	// Params are optional here (cursor, for pagination this server does not
	// need), but a malformed object is still worth refusing rather than
	// ignoring — silently accepting nonsense trains a client to send it.
	var p struct {
		Cursor string `json:"cursor,omitempty"`
	}
	if err := decodeParams(raw, &p); err != nil {
		return nil, err
	}

	infos := exposed(s.reg)
	out := make([]mcpTool, 0, len(infos))
	for _, info := range infos {
		out = append(out, toMCPTool(info))
	}
	return toolsListResult{Tools: out}, nil
}

/* ── tools/call ─────────────────────────────────────────────────────────── */

type toolsCallParams struct {
	Name      string          `json:"name"`
	Arguments json.RawMessage `json:"arguments,omitempty"`
}

// toolsCallResult is MCP's shape for a tool's output.
//
// IsError is part of the RESULT, not a JSON-RPC error: a tool that refused is
// not a protocol fault, and reporting it as one would deny the model the chance
// to read the refusal and do something sensible. It is the same distinction
// tools.Result already draws, and runs.go draws for terminations.
type toolsCallResult struct {
	Content []contentBlock `json:"content"`
	IsError bool           `json:"isError,omitempty"`
}

type contentBlock struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

func textResult(payload any, isError bool) (toolsCallResult, *rpcError) {
	encoded, err := json.MarshalIndent(payload, "", "  ")
	if err != nil {
		return toolsCallResult{}, errInternal()
	}
	return toolsCallResult{
		Content: []contentBlock{{Type: "text", Text: string(encoded)}},
		IsError: isError,
	}, nil
}

// handleToolsCall validates a call and dispatches it through the registry.
//
// The order is: exposure, then bounds, then identity, then dispatch.
//
// Exposure is checked BEFORE identity on purpose. "There is no such tool here"
// does not depend on who is asking, and answering it first means the surface's
// tool inventory is not something an attacker can probe by comparing an
// authenticated 404 against an unauthenticated 401.
//
// Authentication is the transport's job and has already happened by the time
// this runs; `ident` is nil only when it failed or was never attempted. This
// function does not read the ambient context for a caller — see Handle.
func (s *Server) handleToolsCall(ctx context.Context, ident *authctx.Identity, raw json.RawMessage) (any, *rpcError) {
	var p toolsCallParams
	if err := decodeParams(raw, &p); err != nil {
		return nil, err
	}
	if p.Name == "" {
		return nil, errInvalidParams("name is required")
	}
	if len(p.Arguments) > maxToolArgumentBytes {
		return nil, errInvalidParams("arguments are too large")
	}

	// Unexposed and unknown are the SAME answer, deliberately. See
	// isExposedName — distinguishing them inventories what this surface is
	// hiding.
	if !isExposedName(s.reg, p.Name) {
		s.log.Warn("mcp tool call refused", "tool", p.Name, "reason", "not_exposed")
		return textResult(map[string]any{
			"error": map[string]any{
				"code":    "mcp.unknown_tool",
				"message": "there is no tool called " + p.Name + " on this surface",
			},
		}, true)
	}

	// Unauthenticated calls never reach a handler. The transport answers 401
	// before this point in the ordinary case; this is the second gate, so that
	// a future caller of Handle that forgets to authenticate fails closed
	// rather than dispatching as nobody.
	if ident == nil {
		s.log.Warn("mcp tool call refused", "tool", p.Name, "reason", "no_identity")
		return textResult(map[string]any{
			"error": map[string]any{
				"code":    "mcp.unauthenticated",
				"message": "this call is not authenticated",
			},
		}, true)
	}

	return s.dispatch(ctx, *ident, p)
}

// dispatch runs the tool through the existing registry.
//
// This is the only place this package touches the tool layer, and it is four
// lines on purpose. Everything that decides what comes back — the policy table,
// the org pre-filter, the row scopes, the opaque denial, the truncation — is
// inside Dispatch and the handler beneath it, unchanged and unreachable from
// here.
//
// The principal is the caller's, from the context. It is never read from
// params: an MCP client that could name its own principal could read anything,
// which is the bug I1 exists to prevent.
func (s *Server) dispatch(ctx context.Context, identity authctx.Identity, p toolsCallParams) (any, *rpcError) {
	// The per-organisation ceiling, checked after authentication and before
	// any work. The org comes from `identity`, which came from the token —
	// there is no path by which a request can name a different bucket, because
	// this function is never given anything from the request except the tool
	// name and its arguments.
	if s.orgLimiter != nil {
		allowed, retryAfter, err := s.orgLimiter.AllowOrg(ctx, identity.OrgID)
		if err != nil {
			// The limiter has already decided whether a failure permits the
			// call. Logged without the org's usage, which is not the caller's
			// business.
			s.log.Error("org rate limiter unavailable", "error", err)
		}
		if !allowed {
			s.log.Warn("mcp org rate limit exceeded",
				"org_id", identity.OrgID, "tool", p.Name)
			return nil, errRateLimited(retryAfter)
		}
	}

	args := p.Arguments
	if len(skipSpace(args)) == 0 {
		args = json.RawMessage(`{}`)
	}

	tc := tools.Context{
		Principal: identity,
		// No RunID: an MCP call is not an agent run and writes no trajectory.
		// No KnowledgeSources: there is no spec, which is why knowledge_search
		// is deferred rather than published — see tools.go.
	}

	res := s.reg.Dispatch(ctx, tc, p.Name, args)

	s.log.Info("mcp tool call",
		"tool", p.Name,
		"user_id", identity.UserID,
		"org_id", identity.OrgID,
		"ok", res.Error == nil,
		"truncated", res.Truncated)

	if res.Error != nil {
		return textResult(map[string]any{"error": res.Error}, true)
	}
	return textResult(map[string]any{
		"data":      res.Data,
		"truncated": res.Truncated,
	}, false)
}