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, and yours: recall before acting,
record as you go, keep one copy here rather than in local memory files.

Each reflex is stated in full in the using-scribe skill (if your client reads
Agent Skills) and 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, what_might_apply("what you are about to do") — fifty
  ranked, no bar. search(content_type="rule") reads one you suspect. Silence
  means nothing matched, not none. Rules bind; preferences guide and you keep
  them current; lessons inform.
- MISSED: a rule that missed you is a trigger to fix, not a floor to move
  (retrieval_telemetry).
- JUDGE: you are the judge of record — what a shape is, whether a finding
  holds, whether work is done. Surfacing one for them to rule on is the
  judgment not made. Escalate their acts, not your decisions.
- 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.
- PLAN with an arc: find the existing plan first
  (search(content_type="milestone")) and add steps to it; else
  start_planning(steps=[...]).
- IDS exist only once a create returns them. Records citing 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 from the `placement` a task write returns: where it sits, what
  changed, what needs them, what 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: EVERY
# registered tool must appear in exactly one of _READ_ONLY_TOOLS, _WRITE_TOOLS or
# _DELIBERATELY_WRITE_SCOPED below, so adding one forces a decision instead of
# silently denying it — whatever the tool is called (#3191).
#
# 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",
    "get_lesson", "list_lessons",
    # 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", "stamps_to_review",
    # 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.
    "notes_due_for_verification",
    # Its rule twin and a rule's edit history (milestones 312 and 323). Both
    # pure reads, and both sat unlisted — so a read key was refused them — for
    # the same reason: no read prefix, back when the completeness test only
    # looked at names that had one (#3191). rule_history records a pull the way
    # the getters above do.
    "rules_due_for_verification", "rule_history",
    # What each retrieval surface's floor and budget currently are, and what
    # has been changed about them (#4102). Both pure reads; `tune_retrieval` is
    # the write and is deliberately NOT here. Read access matters more than
    # usual for these two: a session that cannot see the bar in force, or the
    # reason it was last moved, is a session that will move it again blind.
    "retrieval_surfaces", "retrieval_tuning_history",
    # The wide net (#4103) — ranked rule candidates with no bar, for the
    # moment before a consequential act. A pure read, and one a read key needs
    # most: it is the surface the "ask before acting" reflex calls, and a key
    # that could not reach it would be denied exactly the check it should run.
    "what_might_apply",
})

# Every tool that WRITES, by name. Nothing reads this set at runtime — a tool
# absent from _READ_ONLY_TOOLS is already denied to a read key. It exists so the
# classification is total: test_mcp_auth requires every registered tool to sit
# in exactly one of the three sets, which is what makes forgetting impossible
# rather than merely unlikely. Before #3191 the test only asked about tools whose
# names looked like reads, and two reads with other names were denied for weeks.
_WRITE_TOOLS = frozenset({
    # notes, tasks, planning
    "create_note", "update_note", "delete_note",
    "create_lesson", "update_lesson", "delete_lesson",
    "create_task", "update_task", "delete_task", "add_task_log",
    "create_records", "start_planning",
    "create_milestone", "update_milestone", "delete_milestone",
    "mark_note_verified",
    # projects, Systems, repos
    "create_project", "update_project", "delete_project", "decide_project_inception",
    "create_system", "update_system", "delete_system", "map_system_to_canonical",
    "bind_repo", "unbind_repo",
    # snippets, processes, the shape ledger
    "create_snippet", "update_snippet", "delete_snippet", "verify_snippet",
    "merge_snippets", "unmerge_snippet",
    "create_process", "update_process", "delete_process",
    "classify_shapes", "classify_shapes_by_rule", "confirm_shape_proposals",
    "refresh_pattern_coverage",
    # design systems
    "create_design_system", "update_design_system", "delete_design_system",
    "create_design_token", "update_design_token", "delete_design_token",
    "set_project_design_system",
    # rules
    "create_rulebook", "update_rulebook", "delete_rulebook",
    "create_topic", "update_topic", "delete_topic",
    "create_rule", "create_project_rule", "update_rule", "move_rule", "delete_rule",
    "create_preference", "update_preference",
    "relate_rules", "unrelate_rules", "mark_rule_verified",
    # rule_outcome writes only telemetry, which is the case _READ_ONLY_TOOLS
    # above explicitly tolerates for getters that call record_pulled. It is
    # classed as a WRITE anyway, on the difference that matters: those are
    # reads that happen to leave a trace, while this tool's entire effect is
    # the row — and the row carries `detail`, free prose the agent authored.
    # A read-scoped key that can put text into the operator's database is not
    # read-scoped, whatever table it lands in (#4212).
    "rule_outcome",
    # retrieval tuning — a write in both senses: it moves the number the arm
    # reads, and it appends the reason to the audit trail (#4102).
    "tune_retrieval",
    "migrate_retrieval_floor",
    # trash
    "restore", "purge_trash",
})

# 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