FabledScribe/src/scribe/mcp/server.py

"""FastMCP instance + Quart mount-point. Tools are registered in mcp/tools/."""
from __future__ import annotations

import difflib

from mcp.server.fastmcp import FastMCP
from mcp.server.transport_security import TransportSecuritySettings
from quart import Quart

## What this block is — read before editing it
#
# ONE OWNER PER PIECE OF GUIDANCE (decision #4027, milestone 410). Scribe's
# guidance to agents lives where it can be delivered, and each topic is stated
# in full exactly once:
#   - Tool docstrings: each tool's contract, delivered with its schema.
#   - In-band tool responses: behaviour prose cannot be trusted to trigger —
#     the duplicate gate, the guessed-id refusal, `systems_hint`,
#     `placement` and `report_back` — at the moment it applies.
#   - The bundled skills (Agent Skills, client-neutral): every reflex in full.
#     `using-scribe` owns the working reflexes; the process skills own arcs.
#   - Client adapters (the Claude Code plugin today): timing and that
#     client's own conventions, never a copy of the above.
# tests/test_guidance_ownership.py holds the topic registry that enforces it.
#
# THIS BLOCK IS THE SERVER'S ORIENTATION, WRITTEN AS AN INDEX. It reaches
# every MCP client, so it names no client, and it points at where each reflex
# is stated rather than restating it. A new topic gets a line here only if it
# is a session-start reflex; its full statement goes to its owner.
#
# BUDGET: at most 2,000 characters (test_instructions_fit_the_fold). Claude
# Code injects only the first ~2,048 characters of a server's instructions and
# cuts the rest mid-word (#2562, observed live), and other clients differ, so
# nothing load-bearing may sit past the fold. The history of what was traded
# for space before the ownership split (milestones 317, 333, 409) is in
# decision #4027 and the notes it supersedes.
_INSTRUCTIONS = """
Scribe is the operator's system of record for their work, and yours: recall
from it before acting, record in it as you go, and keep one copy here rather
than in local memory files.

Every reflex below is stated in full in the using-scribe skill (if your
client reads Agent Skills) and in each tool's description. The index:
- ORIENT: enter_project(id) loads the project, open work, Systems and design
  system. An `inception` key: ask what it inherits, then
  decide_project_inception.
- RULES: nothing preloads; a rule arrives when your work matches it. Before a
  consequential act, search(content_type="rule"). Silence means
  nothing matched, not none. Rules bind; preferences guide.
- RECALL: search before acting, scoped with the active project_id.
- RECORD: create_task; a fix is kind="issue". add_task_log as you go; status
  in_progress on start, done on finish. Tag system_ids as you write.
- PLAN work with an arc: start_planning(steps=[...]). The plan is a milestone
  and each step a task.
- IDS exist only once a create returns them. Records that cite each other go
  through create_records, writing {{ref:N}} for the Nth record.
- REUSE: search snippets before building; create_snippet what you build.
- UI: the project's design system binds; resolve_design_system before
  hand-writing a value.
- REPORT back from the `placement` a task write returns: where the work sits,
  what changed, what needs the operator, what comes next.

Creates are duplicate-gated: a near-match returns the existing id to update.
shared:true records are another user's suggestion, not settled practice.
"""


# Tools a read-only API key may call. Anything not listed is treated as a
# write for read keys (default-deny), so a newly-added tool is locked down
# until explicitly classified here.
#
# The list stays EXPLICIT rather than being derived from the name. A read key is
# what you hand to something you don't fully trust — a dashboard, a CI job, a
# shared integration — and a boundary inferred from a naming convention grants
# access to whatever a future author happens to call `get_*`. Enumerating it is
# the point; staleness is the cost, and test_mcp_auth covers that (a read-shaped
# tool must appear here or in _DELIBERATELY_WRITE_SCOPED below, so adding one
# forces a decision instead of silently denying it).
#
# Membership means "reads the operator's data and mutates none of it". Several
# getters record a retrieval event via record_pulled; that is telemetry about
# the read itself, not a change to what was read, and it must keep working for a
# read key or the corpus's surfaced:pulled ratio silently under-counts whichever
# consumers hold one.
_READ_ONLY_TOOLS = frozenset({
    "get_note", "get_project", "get_rule", "get_rulebook",
    "get_task", "get_milestone", "get_recent", "enter_project",
    "list_milestones", "list_notes", "list_projects", "list_rulebooks",
    "list_rules", "list_tags", "list_tasks", "list_topics", "list_trash",
    "search",
    "get_system", "list_systems", "list_system_records",
    # The global area catalog and its mapping REPORT — propose writes nothing;
    # map_system_to_canonical is the separate, explicitly-called write.
    "list_canonical_systems", "propose_canonical_mappings",
    # Reports on the corpus. Reads only — the merge or supersession each
    # suggests is a separate, explicitly-called write.
    "find_duplicate_snippets", "find_duplicate_records",
    # Snippets and processes are notes with a kind. A key that may read a note
    # but not a snippet inverts the sensitivity ordering: it exposes the
    # free-text records and withholds the structured ones (#2496).
    "get_snippet", "list_snippets",
    "get_process", "list_processes",
    # Design systems: read, resolve (inheritance + mode), render, and compare
    # against recorded snippets. All four compute from stored records and write
    # nothing — the drift report is a report, and applying it is a separate
    # explicit call.
    "get_design_system", "list_design_systems", "resolve_design_system",
    "get_design_system_stylesheet", "list_design_tokens",
    "check_snippets_against_design_system", "list_starter_role_groups",
    # Which repos map to which project. Read-only by nature; bind_repo /
    # unbind_repo are the writes.
    "list_repo_bindings",
    # The shape ledger's todo query (#2789). Reads only — classify_shapes is
    # the write, and it is deliberately NOT here.
    "list_shapes", "shape_history",
    # The retrieval telemetry readout (#2975). Aggregates two log tables and
    # writes nothing. Listed explicitly because its name carries no read
    # prefix, so the completeness test below cannot derive it — the same
    # reason `enter_project` is spelled out above.
    "retrieval_telemetry",
    # The note staleness sweep (milestone 317). A pure read — mark_note_verified
    # is the write, and it is deliberately NOT here. Spelled out for
    # retrieval_telemetry's reason: `notes_due_for_verification` matches none of
    # the prefixes the completeness test derives from, so nothing would have
    # prompted this decision. `rules_due_for_verification` is in the same
    # position and is NOT listed — see #3191.
    "notes_due_for_verification",
})

# Read-SHAPED tools that must NOT be reachable with a read key — a getter that
# creates on miss, a list that has a side effect. Empty today, and deliberately
# kept as a declared escape hatch rather than left implicit: without it, the
# completeness test would push a future `get_or_create_*` into the allow-list
# above, which is exactly the wrong way to make a test pass.
_DELIBERATELY_WRITE_SCOPED: frozenset[str] = frozenset()


async def _buffer_request_body(receive):
    """Drain the ASGI request body and return (body_bytes, replay_receive).

    The MCP sub-app still needs to read the body, so we return a fresh
    `receive` that replays the buffered bytes.
    """
    chunks: list[bytes] = []
    more = True
    while more:
        message = await receive()
        if message["type"] == "http.request":
            chunks.append(message.get("body", b""))
            more = message.get("more_body", False)
        else:  # http.disconnect
            more = False
    body = b"".join(chunks)

    sent = False

    async def replay():
        nonlocal sent
        if not sent:
            sent = True
            return {"type": "http.request", "body": body, "more_body": False}
        return {"type": "http.disconnect"}

    return body, replay


def _body_calls_write_tool(body: bytes) -> bool:
    """True if the JSON-RPC body invokes a tool outside the read all-list."""
    import json
    try:
        payload = json.loads(body)
    except Exception:
        return False
    items = payload if isinstance(payload, list) else [payload]
    for item in items:
        if not isinstance(item, dict):
            continue
        if item.get("method") == "tools/call":
            name = (item.get("params") or {}).get("name", "")
            if name and name not in _READ_ONLY_TOOLS:
                return True
    return False


class StrictArgsFastMCP(FastMCP):
    """A FastMCP that REJECTS tool calls carrying undeclared arguments.

    FastMCP validates arguments with a pydantic model built from the tool
    signature, and pydantic's default extra-field policy is "ignore" — so a
    misnamed argument simply vanishes and the tool runs with that field's
    default. On the create/update tools the default is "", which turns a
    plausible near-miss (`content=` for `body=`, primed by add_task_log's
    `content`) into SILENT DATA LOSS: the call reports success and stores an
    empty body, leaving a record search cannot see (#2709). Two notes were
    persisted body-less that way before anyone noticed.

    An error the caller sees once is strictly better than data half-written
    forever, so the policy is applied to every tool, not just the two that
    bit: nothing here knows tool semantics, only that an argument nobody
    declared cannot have been meant to be dropped.
    """

    async def call_tool(self, name, arguments):
        try:
            tool = self._tool_manager.get_tool(name)
        except Exception:
            tool = None  # unknown tool → let upstream produce its own error
        if tool is not None:
            declared = set((tool.parameters or {}).get("properties", {}))
            unknown = sorted(set(arguments or {}) - declared)
            if unknown:
                hints = []
                for arg in unknown:
                    close = difflib.get_close_matches(arg, sorted(declared), n=1)
                    suggestion = f" (did you mean '{close[0]}'?)" if close else ""
                    hints.append(f"'{arg}'{suggestion}")
                raise ValueError(
                    f"{name} does not accept argument(s) {', '.join(hints)}. "
                    f"It accepts: {', '.join(sorted(declared))}. Nothing was "
                    "created or changed — retry with the declared names."
                )
        return await super().call_tool(name, arguments)


def build_mcp_server() -> FastMCP:
    """Build the FastMCP instance with all tools registered.

    DNS-rebinding protection is disabled: FastMCP's default allow-list is
    just localhost variants, which means any deployment behind a reverse
    proxy (Traefik with a hostname like devassistant.traefik.internal,
    Cloudflare, nginx, etc.) gets 421 Misdirected Request. The threat
    model that protection addresses — a malicious browser page rebinding
    DNS to hit a localhost MCP — doesn't apply here: this is HTTP transport
    behind a reverse proxy with bearer-token auth as the real security
    boundary.
    """
    # stateless_http=True: don't hand the client a persistent Mcp-Session-Id.
    # The stateful default strands Claude Code after a container redeploy —
    # it reconnects with the now-unknown session id, the server returns 404,
    # and the client won't re-initialize on a 404 (Claude Code issue #60949),
    # so the connection stays dead until a manual /mcp retry. Stateless makes
    # every request self-contained (bearer-auth only), so a post-deploy
    # reconnect just works. Trade-off: no server-pushed list_changed stream,
    # which we don't use — tools are re-fetched on reconnect anyway.
    mcp = StrictArgsFastMCP(
        "scribe",
        instructions=_INSTRUCTIONS.strip(),
        stateless_http=True,
        transport_security=TransportSecuritySettings(
            enable_dns_rebinding_protection=False,
        ),
    )
    from scribe.mcp.tools import register_all
    register_all(mcp)
    return mcp


def mount_mcp(app: Quart) -> None:
    """Mount the FastMCP streamable-HTTP ASGI sub-app at /mcp on the Quart app.

    A small ASGI middleware between Quart and the FastMCP sub-app validates the
    Bearer token against the api_keys table. Authenticated requests have their
    user_id attached to the ASGI scope under "scribe_user_id" for tool handlers
    to read.

    FastMCP's streamable_http session manager owns a task group that must be
    running before it can serve requests. In a stand-alone Starlette deployment
    that would happen via the Starlette `lifespan` parameter. Since we're hosted
    inside Quart, we hook the session manager's `run()` async context manager
    into Quart's serving lifecycle (before_serving / after_serving).
    """
    from scribe.mcp.auth import resolve_bearer

    mcp = build_mcp_server()
    mcp_asgi = mcp.streamable_http_app()
    app.mcp_instance = mcp

    @app.before_serving
    async def _start_mcp_session() -> None:
        cm = mcp.session_manager.run()
        await cm.__aenter__()
        app._mcp_session_cm = cm

    @app.after_serving
    async def _stop_mcp_session() -> None:
        cm = getattr(app, "_mcp_session_cm", None)
        if cm is not None:
            await cm.__aexit__(None, None, None)

    async def auth_wrapped(scope, receive, send):
        if scope["type"] != "http":
            return await mcp_asgi(scope, receive, send)
        # ASGI headers are lowercase bytes per spec; lowercase explicitly to be safe.
        headers = {k.decode().lower(): v.decode() for k, v in scope.get("headers", [])}
        resolved = await resolve_bearer(headers.get("authorization"))
        if resolved is None:
            await send({
                "type": "http.response.start",
                "status": 401,
                "headers": [
                    (b"content-type", b"application/json"),
                    (b"www-authenticate", b'Bearer realm="scribe-mcp"'),
                ],
            })
            await send({
                "type": "http.response.body",
                "body": b'{"error":"unauthorized"}',
            })
            return
        user_id, key_scope = resolved

        # Enforce read-only keys: REST blocks non-GET for scope='read', and the
        # MCP surface must match or the read-only guarantee is void. A tool call
        # arrives as a JSON-RPC POST; buffer the body, and if it invokes a tool
        # outside the read all-list, reject before dispatch. (default-deny: any
        # unknown/new tool is treated as a write for read keys.)
        if key_scope == "read" and scope.get("method") == "POST":
            body, receive = await _buffer_request_body(receive)
            if _body_calls_write_tool(body):
                await send({
                    "type": "http.response.start",
                    "status": 403,
                    "headers": [(b"content-type", b"application/json")],
                })
                await send({
                    "type": "http.response.body",
                    "body": b'{"error":"read-only API key cannot call write tools"}',
                })
                return

        scope["scribe_user_id"] = user_id
        from scribe.mcp._context import _user_id_ctx
        token = _user_id_ctx.set(user_id)
        try:
            await mcp_asgi(scope, receive, send)
        finally:
            _user_id_ctx.reset(token)

    original_asgi = app.asgi_app

    async def dispatch(scope, receive, send):
        if scope["type"] == "http":
            path = scope.get("path", "")
            if path == "/mcp" or path.startswith("/mcp/"):
                # Don't rewrite the path: FastMCP's streamable_http_app mounts
                # its handler at /mcp by default. If we strip the prefix to "/",
                # FastMCP's internal routing returns 404 because there's no
                # handler at "/" — only at "/mcp". Pass the scope through
                # untouched and let FastMCP's own routing match.
                return await auth_wrapped(scope, receive, send)
        return await original_asgi(scope, receive, send)

    app.asgi_app = dispatch