chore-mcp-server-god-module-handler-extraction: extract MCP tool handlers

README.md

@@ -181,7 +181,7 @@ despite the shared name.
 
 ## MCP tools
 
-The server registers 22 tools, grouped below. Full schemas are in [`website/mcp-tools.md`](website/mcp-tools.md) and live in `better_memory/mcp/server.py`.
+The server registers 22 tools, grouped below. Full schemas are in [`website/mcp-tools.md`](website/mcp-tools.md) and live in `better_memory/mcp/tools.py`.
 
 **Episodic memory** — observations the AI writes during a session.
 

better_memory/mcp/_util.py

@@ -0,0 +1,71 @@
+"""Small shared helpers for the MCP server and its tool handlers.
+
+Lives below both :mod:`better_memory.mcp.server` and the handler modules
+under :mod:`better_memory.mcp.handlers` so either side can import it
+without creating a cycle.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+from better_memory import _diag
+from better_memory.runtime.session_marker import read_session_id
+
+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}]"
+            )
+
+
+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/handlers/__init__.py

@@ -0,0 +1,59 @@
+"""Per-domain MCP tool handler classes.
+
+Each module in this package owns one tool domain and exposes a handler
+class whose ``tools()`` method returns a ``{tool_name: coroutine}``
+mapping. ``create_server`` (in :mod:`better_memory.mcp.server`)
+constructs the services once, instantiates each handler class with the
+services it needs, and merges the mappings into a single dispatch
+registry — keeping ``_call_tool`` a pure lookup.
+
+Handler classes hold no state beyond the injected services, so the
+concurrency contract documented in ``create_server`` (one in-flight
+tool call at a time over the shared sqlite connection) is unchanged.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from mcp.types import TextContent
+
+from better_memory.mcp.handlers.episodes import EpisodeToolHandlers
+from better_memory.mcp.handlers.knowledge import KnowledgeToolHandlers
+from better_memory.mcp.handlers.observations import ObservationToolHandlers
+from better_memory.mcp.handlers.reflections import ReflectionToolHandlers
+from better_memory.mcp.handlers.semantics import SemanticToolHandlers
+from better_memory.mcp.handlers.sessions import SessionToolHandlers
+
+# One MCP tool invocation: JSON-decoded arguments in, TextContent out.
+ToolHandler = Callable[[dict[str, Any]], Awaitable[list[TextContent]]]
+
+
+def build_registry(
+    *groups: Any,
+) -> dict[str, ToolHandler]:
+    """Merge handler groups into one ``{tool_name: handler}`` registry.
+
+    Raises ``RuntimeError`` on a duplicate tool name so a wiring mistake
+    fails at server construction, not at first dispatch.
+    """
+    registry: dict[str, ToolHandler] = {}
+    for group in groups:
+        for tool_name, handler in group.tools().items():
+            if tool_name in registry:
+                raise RuntimeError(f"duplicate tool handler: {tool_name}")
+            registry[tool_name] = handler
+    return registry
+
+
+__all__ = [
+    "EpisodeToolHandlers",
+    "KnowledgeToolHandlers",
+    "ObservationToolHandlers",
+    "ReflectionToolHandlers",
+    "SemanticToolHandlers",
+    "SessionToolHandlers",
+    "ToolHandler",
+    "build_registry",
+]

better_memory/mcp/handlers/episodes.py

@@ -0,0 +1,154 @@
+"""Handlers for the episode-lifecycle tools.
+
+Tools: ``memory.start_episode``, ``memory.close_episode``,
+``memory.reconcile_episodes``, ``memory.list_episodes``.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from mcp.types import TextContent
+
+from better_memory.config import project_name
+from better_memory.mcp.serializers import serialize_queue
+from better_memory.services.episode import EpisodeService
+from better_memory.services.observation import ObservationService
+from better_memory.services.reflection import ReflectionSynthesisService
+from better_memory.storage import StorageBackend
+
+
+class EpisodeToolHandlers:
+    """Episode open / close / reconcile / list."""
+
+    def __init__(
+        self,
+        *,
+        episodes: EpisodeService,
+        observations: ObservationService,
+        reflections: ReflectionSynthesisService,
+        backend: StorageBackend,
+    ) -> None:
+        self._episodes = episodes
+        self._observations = observations
+        self._reflections = reflections
+        self._backend = backend
+
+    def tools(self) -> dict[str, Any]:
+        return {
+            "memory.start_episode": self.start_episode,
+            "memory.close_episode": self.close_episode,
+            "memory.reconcile_episodes": self.reconcile_episodes,
+            "memory.list_episodes": self.list_episodes,
+        }
+
+    async def start_episode(self, args: dict[str, Any]) -> list[TextContent]:
+        project = project_name()
+        episode_id = self._episodes.start_foreground(
+            session_id=self._observations.session_id,
+            project=project,
+            goal=args["goal"],
+            tech=args.get("tech"),
+        )
+        # Reflection synthesis is now interactive — driven by the IDE-LLM
+        # via memory.synthesize_next_get_context / _apply. We surface the
+        # current pending count so the LLM knows whether it should run
+        # synthesis (typically before treating retrieved reflections as
+        # canonical).
+        queue = self._reflections._read_queue_counts(project=project)
+        buckets = self._backend.retrieve(
+            project=project, tech=args.get("tech"),
+        )
+        return [
+            TextContent(
+                type="text",
+                text=json.dumps(
+                    {
+                        "episode_id": episode_id,
+                        "reflections": buckets,
+                        "pending_synthesis": serialize_queue(queue),
+                    }
+                ),
+            )
+        ]
+
+    async def close_episode(self, args: dict[str, Any]) -> list[TextContent]:
+        outcome = args["outcome"]
+        # Default close_reason: match outcome for the common paths.
+        default_reasons = {
+            "success": "goal_complete",
+            "partial": "plan_complete",
+            "abandoned": "abandoned",
+            "no_outcome": "session_end_reconciled",
+        }
+        close_reason = args.get("close_reason") or default_reasons[outcome]
+        try:
+            closed_id = self._episodes.close_active(
+                session_id=self._observations.session_id,
+                outcome=outcome,
+                close_reason=close_reason,
+                summary=args.get("summary"),
+            )
+        except ValueError:
+            # No active episode — already closed (e.g. by a prior commit-
+            # trailer drain) or never opened. Matches the "safe no-op"
+            # contract documented in the CLAUDE snippet's plan-complete
+            # section.
+            return [
+                TextContent(
+                    type="text",
+                    text=json.dumps(
+                        {"closed_episode_id": None, "already_closed": True}
+                    ),
+                )
+            ]
+        return [
+            TextContent(
+                type="text",
+                text=json.dumps(
+                    {"closed_episode_id": closed_id, "already_closed": False}
+                ),
+            )
+        ]
+
+    async def reconcile_episodes(
+        self, args: dict[str, Any]
+    ) -> list[TextContent]:
+        open_episodes = self._episodes.unclosed_episodes(
+            exclude_session_ids={self._observations.session_id}
+        )
+        payload = [
+            {
+                "episode_id": e.id,
+                "project": e.project,
+                "tech": e.tech,
+                "goal": e.goal,
+                "started_at": e.started_at,
+            }
+            for e in open_episodes
+        ]
+        return [TextContent(type="text", text=json.dumps(payload))]
+
+    async def list_episodes(self, args: dict[str, Any]) -> list[TextContent]:
+        rows = self._episodes.list_episodes(
+            project=args.get("project"),
+            outcome=args.get("outcome"),
+            only_open=args.get("only_open", False),
+        )
+        payload = [
+            {
+                "episode_id": e.id,
+                "project": e.project,
+                "tech": e.tech,
+                "goal": e.goal,
+                "started_at": e.started_at,
+                "hardened_at": e.hardened_at,
+                "ended_at": e.ended_at,
+                "close_reason": e.close_reason,
+                "outcome": e.outcome,
+                "summary": e.summary,
+            }
+            for e in rows
+        ]
+        return [TextContent(type="text", text=json.dumps(payload))]

better_memory/mcp/handlers/knowledge.py

@@ -0,0 +1,55 @@
+"""Handlers for the knowledge-base introspection tools.
+
+Tools: ``knowledge.search``, ``knowledge.list``.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from mcp.types import TextContent
+
+from better_memory.mcp.serializers import (
+    serialize_knowledge_doc,
+    serialize_knowledge_search,
+)
+from better_memory.services.knowledge import KnowledgeService
+
+
+class KnowledgeToolHandlers:
+    """BM25 search + listing over the knowledge-base corpus."""
+
+    def __init__(self, *, knowledge: KnowledgeService) -> None:
+        self._knowledge = knowledge
+
+    def tools(self) -> dict[str, Any]:
+        return {
+            "knowledge.search": self.search,
+            "knowledge.list": self.list_documents,
+        }
+
+    async def search(self, args: dict[str, Any]) -> list[TextContent]:
+        results = self._knowledge.search(
+            args["query"],
+            project=args.get("project"),
+        )
+        return [
+            TextContent(
+                type="text",
+                text=json.dumps(
+                    [serialize_knowledge_search(r) for r in results]
+                ),
+            )
+        ]
+
+    async def list_documents(self, args: dict[str, Any]) -> list[TextContent]:
+        docs = self._knowledge.list_documents(project=args.get("project"))
+        return [
+            TextContent(
+                type="text",
+                text=json.dumps(
+                    [serialize_knowledge_doc(d) for d in docs]
+                ),
+            )
+        ]