refactor(mcp): server.py 1617→377 LOC via ToolDispatcher + per-domain handlers

better_memory/mcp/_best_effort.py

@@ -0,0 +1,55 @@
+"""Best-effort runner used by handlers that must never fail their caller.
+
+Extracted from ``better_memory.mcp.server`` so handler modules can import
+it without pulling the server module (which itself imports handlers via
+``all_handlers``) and creating an import cycle.
+
+The canonical re-export remains ``better_memory.mcp.server._run_best_effort``
+for the test suite (``tests/mcp/test_best_effort_logging.py``).
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Callable
+from typing import Any
+
+from better_memory import _diag
+
+logger = logging.getLogger(__name__)
+
+
+def _run_best_effort(
+    operation: str,
+    fn: Callable[[], Any],
+    *,
+    diag_cid: str | None = None,
+) -> None:
+    """Run ``fn`` swallowing any ``Exception`` but logging it via the module logger.
+
+    Used by best-effort hooks inside ``memory.retrieve`` (spool drain,
+    retention scheduler) where a failure must NEVER block the call but
+    must still produce a discoverable diagnostic. The previous behaviour
+    silently dropped the exception, so a broken background path could
+    fail invisibly for weeks.
+
+    When ``diag_cid`` is provided and ``BETTER_MEMORY_EMBED_LOG=1`` is on,
+    emits a ``[bm-retrieve step=<operation> cid=... ms=N status=ok|error]``
+    line through the shared diagnostic logger so callers can localize
+    which step is slow.
+    """
+    t0 = time.monotonic()
+    status = "ok"
+    try:
+        fn()
+    except Exception:  # noqa: BLE001 — best-effort wrapper
+        status = "error"
+        logger.exception("best-effort %s failed", operation)
+    finally:
+        if diag_cid is not None and _diag.enabled():
+            ms = int((time.monotonic() - t0) * 1000)
+            _diag.log(
+                f"[bm-retrieve step={operation} cid={diag_cid} "
+                f"ms={ms} status={status}]"
+            )

better_memory/mcp/_session.py

@@ -0,0 +1,28 @@
+"""Resolve the Claude Code session id for MCP tool calls.
+
+Moved out of ``better_memory.mcp.server`` so handlers in
+``better_memory.mcp.handlers.*`` can import it without a circular
+dependency on the server module.
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from better_memory.runtime.session_marker import read_session_id
+
+
+def resolve_session_id(home: Path) -> str | None:
+    """Resolve the current Claude Code session id.
+
+    Order: ``CLAUDE_SESSION_ID`` env, ``CLAUDE_CODE_SESSION_ID`` env, then
+    the marker file written by the SessionStart hook (see
+    :mod:`better_memory.runtime.session_marker`). Claude Code does not
+    propagate the session id into the spawned stdio MCP server's env, so
+    the marker file is the fallback for every rating call.
+    """
+    return (
+        os.environ.get("CLAUDE_SESSION_ID")
+        or os.environ.get("CLAUDE_CODE_SESSION_ID")
+        or read_session_id(home)
+    )

better_memory/mcp/container.py

@@ -0,0 +1,41 @@
+"""Container bundling every long-lived service the MCP dispatcher uses.
+
+Constructed once at startup by ``create_server``; passed by reference to
+every tool handler. Frozen so handlers can never accidentally rebind a
+service mid-call.
+"""
+from __future__ import annotations
+
+import sqlite3
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from better_memory.config import Config
+    from better_memory.services.episode import EpisodeService
+    from better_memory.services.knowledge import KnowledgeService
+    from better_memory.services.memory_rating import MemoryRatingService
+    from better_memory.services.observation import ObservationService
+    from better_memory.services.reflection import ReflectionSynthesisService
+    from better_memory.services.retention import RetentionService
+    from better_memory.services.semantic import SemanticMemoryService
+    from better_memory.services.session_bootstrap import SessionBootstrapService
+    from better_memory.services.spool import SpoolService
+    from better_memory.storage import StorageBackend
+
+
+@dataclass(frozen=True)
+class ServiceContainer:
+    """All long-lived services + connections, built once in create_server."""
+    config: Config
+    memory_conn: sqlite3.Connection
+    backend: StorageBackend
+    episodes: EpisodeService
+    observations: ObservationService
+    reflections: ReflectionSynthesisService
+    retention: RetentionService
+    memory_rating: MemoryRatingService
+    knowledge: KnowledgeService
+    spool: SpoolService
+    semantic: SemanticMemoryService
+    session_bootstrap: SessionBootstrapService

better_memory/mcp/dispatcher.py

@@ -0,0 +1,60 @@
+"""ToolDispatcher: O(1) name to handler routing for MCP tool calls.
+
+Replaces the 470-LOC ``if name == "..."`` chain in the legacy
+``_call_tool`` closure. Handlers are registered as ``Handler`` dataclass
+instances; the dispatcher owns the lookup, the capability gate, and the
+"unknown tool" error contract.
+"""
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from typing import Any
+
+from mcp.types import TextContent, Tool
+
+from better_memory.mcp.container import ServiceContainer
+
+HandlerFn = Callable[[ServiceContainer, dict[str, Any]], Awaitable[list[TextContent]]]
+
+
+@dataclass(frozen=True)
+class Handler:
+    """One MCP tool: name, JSON-Schema for inputs, async callable, capability flag."""
+    name: str
+    schema: dict[str, Any]
+    call: HandlerFn
+    description: str = ""
+    requires_synthesis: bool = False
+
+
+class ToolDispatcher:
+    """Owns the {name: Handler} table plus the capability gate."""
+
+    def __init__(
+        self, services: ServiceContainer, handlers: list[Handler],
+    ) -> None:
+        self._services = services
+        self._handlers: dict[str, Handler] = {h.name: h for h in handlers}
+
+    def tool_definitions(self) -> list[Tool]:
+        supports = self._services.backend.supports_synthesis
+        return [
+            Tool(name=h.name, description=h.description, inputSchema=h.schema)
+            for h in self._handlers.values()
+            if supports or not h.requires_synthesis
+        ]
+
+    async def call(
+        self, name: str, args: dict[str, Any],
+    ) -> list[TextContent]:
+        handler = self._handlers.get(name)
+        if handler is None:
+            raise ValueError(f"Unknown tool: {name}")
+        if (
+            handler.requires_synthesis
+            and not self._services.backend.supports_synthesis
+        ):
+            # Same error shape as today's fallthrough — clients depend on it.
+            raise ValueError(f"Unknown tool: {name}")
+        return await handler.call(self._services, args)

better_memory/mcp/handlers/__init__.py

@@ -0,0 +1,35 @@
+"""MCP tool handlers, organised by domain.
+
+Each domain module exposes one ``*Handlers`` class with a ``handlers()``
+method returning ``list[Handler]`` for registration with the
+``ToolDispatcher``. The :func:`all_handlers` helper assembles every
+domain's contribution.
+"""
+from __future__ import annotations
+
+from better_memory.mcp.dispatcher import Handler
+from better_memory.mcp.handlers.episodes import EpisodeHandlers
+from better_memory.mcp.handlers.knowledge import KnowledgeHandlers
+from better_memory.mcp.handlers.observations import ObservationHandlers
+from better_memory.mcp.handlers.ratings import RatingHandlers
+from better_memory.mcp.handlers.reflections import ReflectionHandlers
+from better_memory.mcp.handlers.retention import RetentionHandlers
+from better_memory.mcp.handlers.semantics import SemanticHandlers
+from better_memory.mcp.handlers.session import SessionHandlers
+
+
+def all_handlers() -> list[Handler]:
+    """Return the union of every domain's registered handlers.
+
+    Filled in as handler modules land (Tasks 7-14).
+    """
+    return [
+        *ObservationHandlers().handlers(),
+        *SemanticHandlers().handlers(),
+        *EpisodeHandlers().handlers(),
+        *ReflectionHandlers().handlers(),
+        *RetentionHandlers().handlers(),
+        *KnowledgeHandlers().handlers(),
+        *RatingHandlers().handlers(),
+        *SessionHandlers().handlers(),
+    ]

better_memory/mcp/handlers/_audit.py

@@ -0,0 +1,94 @@
+"""Audit-log helpers for the synthesize_* tool handlers.
+
+Moved verbatim from ``better_memory.mcp.server``. Byte-shape of the JSONL
+rows is unchanged — ``tests/mcp/test_synth_audit_log.py`` continues to
+pass without modification beyond the import path.
+
+The synthesize drain loop is driven by the IDE LLM across many
+round-trips (one per pending episode). When that loop appears to
+freeze, server-side timing is the only evidence we have — the LLM
+side is opaque. Each call writes two JSONL rows to
+``{config.home}/logs/synthesize.jsonl``:
+
+    {"phase": "start",    "call_id": "...", "tool": "...", ...}
+    {"phase": "complete", "call_id": "...", "tool": "...",
+     "latency_ms": N, "result_kind": "...", ...}
+
+Paired by call_id. A start row with no matching complete row points
+to the call that hung. Best-effort: any IO error is swallowed so the
+audit log can never block or fail the synthesize tool itself.
+"""
+from __future__ import annotations
+
+import contextlib
+import json
+import logging
+import time
+import uuid
+from collections.abc import Iterator
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def _append_synth_audit(home: Path, payload: dict[str, Any]) -> None:
+    """Append one JSONL row to ``{home}/logs/synthesize.jsonl``."""
+    try:
+        log_dir = home / "logs"
+        log_dir.mkdir(parents=True, exist_ok=True)
+        log_path = log_dir / "synthesize.jsonl"
+        line = json.dumps(payload, separators=(",", ":"))
+        with log_path.open("a", encoding="utf-8") as fh:
+            fh.write(line + "\n")
+    except Exception:  # noqa: BLE001 — best-effort audit
+        logger.exception("synth audit write failed")
+
+
+@contextlib.contextmanager
+def _audit_synth_call(
+    home: Path,
+    *,
+    tool: str,
+    project: str,
+    episode_id: str | None,
+) -> Iterator[dict[str, Any]]:
+    """Bracket a synthesize tool call with start + complete audit rows.
+
+    Yields a mutable ``state`` dict the caller fills in (``result_kind``,
+    ``error``, ``counts``, ``obs_count``, ``refl_count``, and may
+    overwrite ``episode_id`` once known). The complete row is written
+    on both normal exit and exception. Exceptions still propagate.
+    """
+    call_id = uuid.uuid4().hex[:12]
+    t0 = time.perf_counter()
+    _append_synth_audit(home, {
+        "phase": "start",
+        "call_id": call_id,
+        "tool": tool,
+        "ts": datetime.now(UTC).isoformat(),
+        "project": project,
+        "episode_id": episode_id,
+    })
+    state: dict[str, Any] = {
+        "phase": "complete",
+        "call_id": call_id,
+        "tool": tool,
+        "project": project,
+        "episode_id": episode_id,
+        "result_kind": None,
+    }
+    try:
+        yield state
+    except BaseException as exc:
+        if state.get("result_kind") is None:
+            state["result_kind"] = "exception"
+        state.setdefault("error", f"{type(exc).__name__}: {exc}")
+        state["ts"] = datetime.now(UTC).isoformat()
+        state["latency_ms"] = int((time.perf_counter() - t0) * 1000)
+        _append_synth_audit(home, state)
+        raise
+    state["ts"] = datetime.now(UTC).isoformat()
+    state["latency_ms"] = int((time.perf_counter() - t0) * 1000)
+    _append_synth_audit(home, state)