summercms/modules/pier/dispatch.go

package pier

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"slices"
	"strconv"
	"strings"
	"time"

	"git.golem15.com/golem15/summercms/modules/backpack"
	"git.golem15.com/golem15/summercms/modules/bouncer"
	"git.golem15.com/golem15/summercms/modules/pact"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// resultTooLarge is the text of the tool error for a route answer above
// mcp.max_result_bytes.
const resultTooLarge = "pier: the route's answer exceeds mcp.max_result_bytes"

// dispatchFailed is the fixed text of the tool error for a guard dispatch
// that failed or contradicted the tool's guard; the cause (which may name a
// database host or user) is logged server-side only (A-WR-08).
const dispatchFailed = "pier: the tool could not be run for this caller"

// tokenInvalid is the tool error text for a provider answering
// pact.ErrMCPTokenInvalid.
const tokenInvalid = "pier: the caller's token is no longer valid"

// callInfo is what a dispatched call copies from the outer MCP request: the
// client address, the Host, the forwarding, real-ip and language headers,
// plus the session the call belongs to, for the audit line.
type callInfo struct {
	remoteAddr     string
	host           string
	forwardedFor   string
	forwardedProto string
	forwardedHost  string
	realIP         string
	acceptLanguage string
	session        string
}

// callInfoFrom reads callInfo off an HTTP request.
func callInfoFrom(r *http.Request) callInfo {
	if r == nil {
		return callInfo{}
	}
	return callInfo{
		remoteAddr:     r.RemoteAddr,
		host:           r.Host,
		forwardedFor:   r.Header.Get("X-Forwarded-For"),
		forwardedProto: r.Header.Get("X-Forwarded-Proto"),
		forwardedHost:  r.Header.Get("X-Forwarded-Host"),
		realIP:         r.Header.Get("X-Real-IP"),
		acceptLanguage: r.Header.Get("Accept-Language"),
	}
}

// dispatcher runs tool calls in-process through the host's compiled route
// chains (D-10, D-11, D-13, D-17).
type dispatcher struct {
	host Host
	app  *backpack.App
	cfg  Config
}

// call validates the arguments against the tool's resolved schema, splits
// them into path values, query string and body, builds the inner request
// without any Authorization or Cookie header, attaches the bouncer dispatch
// marker with the guard's pre-resolved principal, serves the compiled chain
// into a size-capped recorder and maps the answer: 2xx to text plus
// structured content, anything else to a tool error carrying the route's
// own body. Protocol-level failures (a tool the identity may not call, an
// unknown route) are tool errors too, so the agent can read them.
func (d *dispatcher) call(ctx context.Context, id Identity, info callInfo, t *tool, raw json.RawMessage) *mcp.CallToolResult {
	start := time.Now()
	res, status := d.run(ctx, id, info, t, raw)
	logToolCall(logger(), id, t, info.session, status, time.Since(start))
	return res
}

func (d *dispatcher) run(ctx context.Context, id Identity, info callInfo, t *tool, raw json.RawMessage) (*mcp.CallToolResult, int) {
	// tools/call may omit arguments (or send null) for a tool without
	// required input; both mean an empty object.
	if trimmed := bytes.TrimSpace(raw); len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) {
		raw = json.RawMessage("{}")
	}
	var parts map[string]json.RawMessage
	if err := json.Unmarshal(raw, &parts); err != nil {
		return toolError("pier: arguments must be a JSON object: " + err.Error()), 0
	}
	var instance any
	if err := json.Unmarshal(raw, &instance); err != nil {
		return toolError("pier: arguments must be a JSON object: " + err.Error()), 0
	}
	if err := t.resolved.Validate(instance); err != nil {
		return toolError("pier: invalid arguments: " + err.Error()), 0
	}
	pathValues := map[string]string{}
	query := url.Values{}
	for _, f := range t.fields {
		v, ok := parts[f.name]
		if !ok {
			continue
		}
		switch f.kind {
		case argPath:
			s, err := scalarString(v)
			if err != nil {
				return toolError(fmt.Sprintf("pier: invalid arguments: %q must be a string or a number", f.name)), 0
			}
			if s == "" || (strings.Contains(s, "/") && !strings.Contains(t.pattern, "{"+f.key+"...}")) {
				return toolError(fmt.Sprintf("pier: invalid arguments: %q must be one non-empty path segment", f.name)), 0
			}
			// The mux cleans a request path before any route sees it, so a
			// route never receives an empty, . or .. segment (A-WR-01).
			for _, seg := range strings.Split(s, "/") {
				if seg == "" || seg == "." || seg == ".." {
					return toolError(fmt.Sprintf("pier: invalid arguments: %q must not contain empty, . or .. segments", f.name)), 0
				}
			}
			pathValues[f.key] = s
			delete(parts, f.name)
		case argQuery:
			values, err := queryStrings(v)
			if err != nil {
				return toolError(fmt.Sprintf("pier: invalid arguments: %q must be a scalar or a list of scalars", f.name)), 0
			}
			for _, s := range values {
				query.Add(f.key, s)
			}
			delete(parts, f.name)
		}
	}
	// Every {segment} gets a value, as over HTTP (A-WR-02; boot already
	// requires path fields, this guards the dispatch on its own).
	for _, f := range t.fields {
		if _, ok := pathValues[f.key]; f.kind == argPath && !ok {
			return toolError(fmt.Sprintf("pier: invalid arguments: %q is required", f.name)), 0
		}
	}
	path := t.pattern
	for key, value := range pathValues {
		if strings.Contains(path, "{"+key+"...}") {
			path = strings.ReplaceAll(path, "{"+key+"...}", escapePathSegments(value))
		} else {
			path = strings.ReplaceAll(path, "{"+key+"}", url.PathEscape(value))
		}
	}
	target := path
	if enc := query.Encode(); enc != "" {
		target += "?" + enc
	}
	var body *bytes.Reader
	if t.method != http.MethodGet && t.method != http.MethodDelete {
		encoded, err := json.Marshal(parts)
		if err != nil {
			return toolError("pier: arguments could not be encoded: " + err.Error()), 0
		}
		body = bytes.NewReader(encoded)
	}
	subject := id.subjectFor(t.guardName)
	if subject == nil {
		return toolError(fmt.Sprintf("pier: tool %q is not available to this caller", t.name)), 0
	}
	handler, ok := t.compiledHandler(d.host)
	if !ok {
		return toolError(fmt.Sprintf("pier: tool %q is bound to %s %s, which is not compiled", t.name, t.method, t.pattern)), 0
	}
	callCtx, cancel := context.WithTimeout(ctx, d.cfg.CallTimeout)
	defer cancel()
	disp, err := t.guard.MCPDispatch(callCtx, *subject, t.method)
	if err != nil {
		switch {
		case errors.Is(err, pact.ErrMCPTokenInvalid):
			return toolError(tokenInvalid), 0
		case errors.Is(err, errSubjectUnavailable):
			return toolError(errSubjectUnavailable.Error()), 0
		}
		logger().Warn("pier: guard dispatch failed", "tool", t.name, "guard", t.guardName, "error", err)
		return toolError(dispatchFailed), 0
	}
	principal, ok := disp.Principal.(*bouncer.Principal)
	if !ok || principal == nil {
		return toolError(fmt.Sprintf("pier: guard %s did not resolve a principal", t.guardName)), 0
	}
	// The marker's guard list and the principal's kind must agree with the
	// tool's guard: a frontend provider naming the backend guard, or a
	// backend principal on a frontend tool, is a provider bug (A-IN-08).
	backendTool := t.guardName == BackendGuard
	if principal.Backend != backendTool || (!backendTool && slices.Contains(disp.Guards, BackendGuard)) {
		logger().Warn("pier: guard dispatch contradicts the tool's guard", "tool", t.name, "guard", t.guardName, "guards", disp.Guards, "backend_principal", principal.Backend)
		return toolError(dispatchFailed), 0
	}
	credential := disp.Credential
	if credential == nil && t.guardName == BackendGuard {
		credential = id.credential
	}
	callCtx = bouncer.WithDispatch(callCtx, &bouncer.Dispatch{Guards: disp.Guards, Principal: principal, Credential: credential})

	var req *http.Request
	if body != nil {
		req = httptest.NewRequestWithContext(callCtx, t.method, target, body)
		req.Header.Set("Content-Type", "application/json")
	} else {
		req = httptest.NewRequestWithContext(callCtx, t.method, target, nil)
	}
	req.Header.Set("Accept", "application/json")
	if info.remoteAddr != "" {
		req.RemoteAddr = info.remoteAddr
	}
	if info.host != "" {
		req.Host = info.host
	}
	if info.forwardedFor != "" {
		req.Header.Set("X-Forwarded-For", info.forwardedFor)
	}
	if info.forwardedProto != "" {
		req.Header.Set("X-Forwarded-Proto", info.forwardedProto)
	}
	if info.forwardedHost != "" {
		req.Header.Set("X-Forwarded-Host", info.forwardedHost)
	}
	if info.realIP != "" {
		req.Header.Set("X-Real-IP", info.realIP)
	}
	if info.acceptLanguage != "" {
		req.Header.Set("Accept-Language", info.acceptLanguage)
	}
	for _, seg := range t.segments {
		req.SetPathValue(seg, pathValues[seg])
	}
	req.Pattern = t.method + " " + t.pattern

	rec := newCappedRecorder(d.cfg.MaxResultBytes)
	handler.ServeHTTP(rec, req)
	if rec.overflow {
		return toolError(resultTooLarge), rec.rec.Code
	}
	return mapResult(rec.rec.Code, rec.rec.Body.Bytes()), rec.rec.Code
}

// mapResult turns the route's status and body into the tool result (D-13):
// 2xx is text plus structured content (an object as is, any other JSON
// wrapped as {"data": v}, a non-JSON body text only); anything else is a
// tool error whose text is the route's body (or the status text when the
// body is empty) and whose structured content is the body when it is an
// object.
func mapResult(status int, body []byte) *mcp.CallToolResult {
	text := string(body)
	structured := structuredOf(body)
	if status >= 200 && status < 300 {
		res := &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: text}}}
		if structured != nil {
			res.StructuredContent = structured
		} else if v, ok := anyJSON(body); ok {
			res.StructuredContent = map[string]any{"data": v}
		}
		return res
	}
	if strings.TrimSpace(text) == "" {
		text = http.StatusText(status)
		if text == "" {
			text = "status " + strconv.Itoa(status)
		}
	}
	res := toolError(text)
	if structured != nil {
		res.StructuredContent = structured
	}
	return res
}

// toolError is a tool result with IsError set and one text content.
func toolError(text string) *mcp.CallToolResult {
	return &mcp.CallToolResult{IsError: true, Content: []mcp.Content{&mcp.TextContent{Text: text}}}
}

// structuredOf decodes body when it is a JSON object, else nil.
func structuredOf(body []byte) map[string]any {
	trimmed := bytes.TrimSpace(body)
	if len(trimmed) == 0 || trimmed[0] != '{' {
		return nil
	}
	var m map[string]any
	if err := json.Unmarshal(trimmed, &m); err != nil {
		return nil
	}
	return m
}

// anyJSON decodes body as any JSON value.
func anyJSON(body []byte) (any, bool) {
	trimmed := bytes.TrimSpace(body)
	if len(trimmed) == 0 {
		return nil, false
	}
	var v any
	if err := json.Unmarshal(trimmed, &v); err != nil {
		return nil, false
	}
	return v, true
}

// decodeNumbers decodes one JSON value keeping numbers as json.Number, so
// an integer above 2^53 keeps its digits (A-IN-04).
func decodeNumbers(raw json.RawMessage) (any, error) {
	dec := json.NewDecoder(bytes.NewReader(raw))
	dec.UseNumber()
	var v any
	if err := dec.Decode(&v); err != nil {
		return nil, err
	}
	return v, nil
}

// scalarString renders a JSON string, number or boolean as a path value; a
// number is rendered as its literal.
func scalarString(raw json.RawMessage) (string, error) {
	v, err := decodeNumbers(raw)
	if err != nil {
		return "", err
	}
	return scalarAny(v)
}

// queryStrings renders a JSON scalar or list of scalars as query values.
func queryStrings(raw json.RawMessage) ([]string, error) {
	v, err := decodeNumbers(raw)
	if err != nil {
		return nil, err
	}
	switch x := v.(type) {
	case nil:
		return nil, nil
	case []any:
		var out []string
		for _, item := range x {
			s, err := scalarAny(item)
			if err != nil {
				return nil, err
			}
			out = append(out, s)
		}
		return out, nil
	default:
		s, err := scalarAny(x)
		if err != nil {
			return nil, err
		}
		return []string{s}, nil
	}
}

func scalarAny(v any) (string, error) {
	switch x := v.(type) {
	case string:
		return x, nil
	case json.Number:
		return x.String(), nil
	case float64:
		return strconv.FormatFloat(x, 'f', -1, 64), nil
	case bool:
		return strconv.FormatBool(x), nil
	}
	return "", fmt.Errorf("not a scalar")
}

// escapePathSegments escapes a multi-segment value segment by segment.
func escapePathSegments(value string) string {
	parts := strings.Split(value, "/")
	for i, p := range parts {
		parts[i] = url.PathEscape(p)
	}
	return strings.Join(parts, "/")
}

// cappedRecorder records a route answer up to a byte limit; beyond it the
// rest is dropped and overflow is set, so a tool result is bounded.
type cappedRecorder struct {
	rec      *httptest.ResponseRecorder
	limit    int64
	written  int64
	overflow bool
}

func newCappedRecorder(limit int64) *cappedRecorder {
	return &cappedRecorder{rec: httptest.NewRecorder(), limit: limit}
}

func (c *cappedRecorder) Header() http.Header { return c.rec.Header() }

func (c *cappedRecorder) WriteHeader(status int) { c.rec.WriteHeader(status) }

func (c *cappedRecorder) Write(p []byte) (int, error) {
	if c.overflow {
		return len(p), nil
	}
	if c.written+int64(len(p)) > c.limit {
		c.overflow = true
		c.rec.Body.Reset()
		return len(p), nil
	}
	c.written += int64(len(p))
	return c.rec.Write(p)
}

// Flush is a no-op: the answer is read once the handler returns.
func (c *cappedRecorder) Flush() {}