diff --git a/better_memory/mcp/_best_effort.py b/better_memory/mcp/_best_effort.py new file mode 100644 index 0000000..ee6c5cc --- /dev/null +++ b/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= 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}]" + ) diff --git a/better_memory/mcp/_session.py b/better_memory/mcp/_session.py new file mode 100644 index 0000000..a513170 --- /dev/null +++ b/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) + ) diff --git a/better_memory/mcp/container.py b/better_memory/mcp/container.py new file mode 100644 index 0000000..ac4ceba --- /dev/null +++ b/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 diff --git a/better_memory/mcp/dispatcher.py b/better_memory/mcp/dispatcher.py new file mode 100644 index 0000000..648c3d5 --- /dev/null +++ b/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) diff --git a/better_memory/mcp/handlers/__init__.py b/better_memory/mcp/handlers/__init__.py new file mode 100644 index 0000000..531ac63 --- /dev/null +++ b/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(), + ] diff --git a/better_memory/mcp/handlers/_audit.py b/better_memory/mcp/handlers/_audit.py new file mode 100644 index 0000000..57e6b86 --- /dev/null +++ b/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) diff --git a/better_memory/mcp/handlers/episodes.py b/better_memory/mcp/handlers/episodes.py new file mode 100644 index 0000000..3c50972 --- /dev/null +++ b/better_memory/mcp/handlers/episodes.py @@ -0,0 +1,257 @@ +"""Episode-lifecycle MCP tool handlers. + +Tools: memory.start_episode, memory.close_episode, memory.reconcile_episodes, +memory.list_episodes. Bodies lifted verbatim from the legacy ``_call_tool`` +if-chain to preserve every documented invariant — notably the close try/except +fork that turns a ValueError ("no active episode") into the +``already_closed=True`` no-op payload, and the 10-field list_episodes +serializer the UI depends on. +""" +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.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler +from better_memory.services.reflection import EpisodeQueueCounts + + +def _serialize_queue(queue: EpisodeQueueCounts) -> dict[str, int]: + return { + "pending": queue.pending, + "in_cooldown": queue.in_cooldown, + "done": queue.done, + } + + +_START_EPISODE_SCHEMA: dict[str, Any] = { + "type": "object", + "required": ["goal"], + "additionalProperties": False, + "properties": { + "goal": {"type": "string"}, + "tech": {"type": "string"}, + }, +} + +_CLOSE_EPISODE_SCHEMA: dict[str, Any] = { + "type": "object", + "required": ["outcome"], + "additionalProperties": False, + "properties": { + "outcome": { + "type": "string", + "enum": [ + "success", + "partial", + "abandoned", + "no_outcome", + ], + }, + "close_reason": { + "type": "string", + "enum": [ + "goal_complete", + "plan_complete", + "abandoned", + "superseded", + "session_end_reconciled", + ], + }, + "summary": {"type": "string"}, + }, +} + +_RECONCILE_EPISODES_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": {}, +} + +_LIST_EPISODES_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "project": {"type": "string"}, + "outcome": { + "type": "string", + "enum": [ + "success", + "partial", + "abandoned", + "no_outcome", + ], + }, + "only_open": {"type": "boolean"}, + }, +} + + +_START_EPISODE_DESCRIPTION = ( + "Declare a goal for the current session. Opens a new " + "foreground episode or hardens the existing background " + "episode. Returns the active episode id." +) +_CLOSE_EPISODE_DESCRIPTION = ( + "Close the current session's active episode. outcome is one " + "of success / partial / abandoned / no_outcome." +) +_RECONCILE_EPISODES_DESCRIPTION = ( + "List episodes that are still open from prior sessions, " + "for the LLM to prompt the user about." +) +_LIST_EPISODES_DESCRIPTION = ( + "List episodes with optional filters. For UI and LLM " + "introspection." +) + + +class EpisodeHandlers: + """All memory.start_episode / .close_episode / .reconcile_episodes / .list_episodes tools.""" + + async def start_episode( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + project = project_name() + episode_id = services.episodes.start_foreground( + session_id=services.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 = services.reflections.read_queue_counts(project=project) + buckets = services.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, services: ServiceContainer, 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 = services.episodes.close_active( + session_id=services.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, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + open_episodes = services.episodes.unclosed_episodes( + exclude_session_ids={services.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, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + rows = services.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))] + + def handlers(self) -> list[Handler]: + return [ + Handler( + name="memory.start_episode", + description=_START_EPISODE_DESCRIPTION, + schema=_START_EPISODE_SCHEMA, + call=self.start_episode, + ), + Handler( + name="memory.close_episode", + description=_CLOSE_EPISODE_DESCRIPTION, + schema=_CLOSE_EPISODE_SCHEMA, + call=self.close_episode, + ), + Handler( + name="memory.reconcile_episodes", + description=_RECONCILE_EPISODES_DESCRIPTION, + schema=_RECONCILE_EPISODES_SCHEMA, + call=self.reconcile_episodes, + ), + Handler( + name="memory.list_episodes", + description=_LIST_EPISODES_DESCRIPTION, + schema=_LIST_EPISODES_SCHEMA, + call=self.list_episodes, + ), + ] diff --git a/better_memory/mcp/handlers/knowledge.py b/better_memory/mcp/handlers/knowledge.py new file mode 100644 index 0000000..09033b1 --- /dev/null +++ b/better_memory/mcp/handlers/knowledge.py @@ -0,0 +1,121 @@ +"""Knowledge-domain MCP tool handlers. + +Tools: knowledge.search, knowledge.list. + +Bodies lifted verbatim from the legacy ``_call_tool`` if-chain. The two +serializer helpers (``_serialize_knowledge_search_result`` / +``_serialize_knowledge_document``) flatten the dataclass shapes the +``KnowledgeService`` returns into the JSON payload the IDE-LLM +consumes. Neither tool is capability-gated. +""" +from __future__ import annotations + +import json +from typing import Any + +from mcp.types import TextContent + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler +from better_memory.services.knowledge import ( + KnowledgeDocument, + KnowledgeSearchResult, +) + +_SEARCH_SCHEMA: dict[str, Any] = { + "type": "object", + "required": ["query"], + "additionalProperties": False, + "properties": { + "query": {"type": "string"}, + "project": {"type": "string"}, + }, +} + +_LIST_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "project": {"type": "string"}, + }, +} + +_SEARCH_DESCRIPTION = ( + "BM25 search against the knowledge-base markdown corpus. " + "Returns document paths and rank." +) +_LIST_DESCRIPTION = ( + "List indexed knowledge documents. When ``project`` is " + "supplied, project-scoped rows are filtered to that project." +) + + +def _serialize_knowledge_search_result( + result: KnowledgeSearchResult, +) -> dict[str, Any]: + doc = result.document + return { + "path": doc.path, + "scope": doc.scope, + "project": doc.project, + "language": doc.language, + "rank": result.rank, + } + + +def _serialize_knowledge_document(doc: KnowledgeDocument) -> dict[str, Any]: + return { + "path": doc.path, + "scope": doc.scope, + "project": doc.project, + "language": doc.language, + } + + +class KnowledgeHandlers: + """The knowledge.search / knowledge.list tools.""" + + async def search( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + results = services.knowledge.search( + args["query"], + project=args.get("project"), + ) + return [ + TextContent( + type="text", + text=json.dumps( + [_serialize_knowledge_search_result(r) for r in results] + ), + ) + ] + + async def list( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + docs = services.knowledge.list_documents(project=args.get("project")) + return [ + TextContent( + type="text", + text=json.dumps( + [_serialize_knowledge_document(d) for d in docs] + ), + ) + ] + + def handlers(self) -> list[Handler]: + return [ + Handler( + name="knowledge.search", + description=_SEARCH_DESCRIPTION, + schema=_SEARCH_SCHEMA, + call=self.search, + ), + Handler( + name="knowledge.list", + description=_LIST_DESCRIPTION, + schema=_LIST_SCHEMA, + call=self.list, + ), + ] diff --git a/better_memory/mcp/handlers/observations.py b/better_memory/mcp/handlers/observations.py new file mode 100644 index 0000000..72b855a --- /dev/null +++ b/better_memory/mcp/handlers/observations.py @@ -0,0 +1,279 @@ +"""Observation-domain MCP tool handlers. + +Tools: memory.observe, memory.retrieve, memory.retrieve_observations, +memory.record_use. Bodies lifted verbatim from the legacy +``_call_tool`` if-chain to preserve every documented invariant — +notably the ``spool.drain -> retention.maybe_schedule -> backend.retrieve`` +ordering inside ``memory.retrieve``. +""" +from __future__ import annotations + +import json +import time +import uuid +from typing import Any + +from mcp.types import TextContent + +from better_memory import _diag +from better_memory.config import get_config, project_name +from better_memory.mcp._best_effort import _run_best_effort +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler +from better_memory.services.retention_scheduler import RetentionScheduler + +_OBSERVE_SCHEMA: dict[str, Any] = { + "type": "object", + "required": ["content"], + "additionalProperties": False, + "properties": { + "content": {"type": "string"}, + "component": {"type": "string"}, + "theme": {"type": "string"}, + "trigger_type": {"type": "string"}, + "outcome": { + "type": "string", + "enum": ["success", "failure", "neutral"], + }, + "tech": {"type": "string"}, + "scope": { + "type": "string", + "enum": ["project", "general"], + "description": ( + "'project' (default) for project-scoped observations; " + "'general' for cross-project workflow rules that should " + "surface in every project's memory_retrieve." + ), + }, + }, +} + +_RETRIEVE_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "project": {"type": "string"}, + "tech": {"type": "string"}, + "phase": { + "type": "string", + "enum": ["planning", "implementation", "general"], + }, + "polarity": { + "type": "string", + "enum": ["do", "dont", "neutral"], + }, + "limit_per_bucket": {"type": "integer"}, + }, +} + +_RETRIEVE_OBSERVATIONS_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "project": {"type": "string"}, + "episode_id": {"type": "string"}, + "component": {"type": "string"}, + "theme": {"type": "string"}, + "outcome": { + "type": "string", + "enum": ["success", "failure", "neutral"], + }, + "query": {"type": "string"}, + "limit": {"type": "integer"}, + }, +} + +_RECORD_USE_SCHEMA: dict[str, Any] = { + "type": "object", + "required": ["id"], + "additionalProperties": False, + "properties": { + "id": {"type": "string"}, + "outcome": { + "type": "string", + "enum": ["success", "failure"], + }, + }, +} + + +_OBSERVE_DESCRIPTION = ( + "Record an observation about the current session (a fact, " + "decision, bug fix, or outcome). Returns the new observation id." +) +_RETRIEVE_DESCRIPTION = ( + "Retrieve reflections (do / dont / neutral lessons distilled " + "from prior observations) bucketed by polarity. Filter by " + "project, tech, phase, and polarity. For raw observation " + "lookup, use memory.retrieve_observations." +) +_RETRIEVE_OBSERVATIONS_DESCRIPTION = ( + "Retrieve raw observations matching given filters. Drill-down " + "tool — use memory.retrieve for the distilled-reflections " + "default. With ``query``, results are ranked by hybrid " + "FTS5 + sqlite-vec relevance; without, ordered created_at " + "DESC. ``episode_id`` and ``theme`` filters are ignored " + "in query mode." +) +_RECORD_USE_DESCRIPTION = ( + "Record that an observation was used; optionally mark the " + "outcome as success or failure to reinforce the memory." +) + + +class ObservationHandlers: + """All memory.observe / .retrieve / .retrieve_observations / .record_use tools.""" + + async def observe( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + with _diag.trace( + "mcp.memory.observe", + content_len=len(args.get("content") or ""), + scope=args.get("scope") or "project", + component=args.get("component"), + ): + _diag.step("mcp.memory.observe", "calling_observations_create") + obs_id = await services.observations.create( + content=args["content"], + component=args.get("component"), + theme=args.get("theme"), + trigger_type=args.get("trigger_type"), + outcome=args.get("outcome", "neutral"), + tech=args.get("tech"), + # `or "project"` (not `, "project"` default) defends against + # MCP clients sending {"scope": null} — dict.get returns the + # default only when the key is absent, not when its value is + # None. Without this, scope=None propagates to + # ObservationService.create() which raises ValueError. + scope=args.get("scope") or "project", + ) + _diag.step("mcp.memory.observe", "create_returned", obs_id=obs_id) + return [TextContent(type="text", text=json.dumps({"id": obs_id}))] + + async def retrieve( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + with _diag.trace( + "mcp.memory.retrieve", + project=args.get("project") or "", + tech=args.get("tech"), + phase=args.get("phase"), + polarity=args.get("polarity"), + ): + diag_cid: str | None = None + t_retrieve = 0.0 + if _diag.enabled(): + diag_cid = uuid.uuid4().hex[:8] + t_retrieve = time.monotonic() + _diag.log(f"[bm-retrieve start cid={diag_cid}]") + + # 1. Drain spool — must happen before any retrieval so fresh + # hook events (session_start, commit_close) are processed. + # SpoolService.drain is idempotent. + _diag.step("mcp.memory.retrieve", "before_spool_drain") + _run_best_effort( + "spool.drain", services.spool.drain, diag_cid=diag_cid, + ) + _diag.step("mcp.memory.retrieve", "after_spool_drain") + + # 2. Maybe run retention. Guard ensures at most once per 24h + # regardless of how often retrieve is called. Best-effort: + # a retention failure must NEVER block memory.retrieve. + def _retention_step() -> None: + cfg = get_config() + RetentionScheduler( + services.memory_conn, auto_prune=cfg.auto_prune, + ).maybe_run(triggered_by="retrieve") + + _diag.step("mcp.memory.retrieve", "before_retention_scheduler") + _run_best_effort( + "retention scheduler", _retention_step, diag_cid=diag_cid, + ) + _diag.step("mcp.memory.retrieve", "after_retention_scheduler") + + project = args.get("project") or project_name() + limit_per_bucket = args.get("limit_per_bucket", 20) + t_reflections = time.monotonic() if diag_cid else 0.0 + _diag.step( + "mcp.memory.retrieve", + "before_retrieve_reflections", + project=project, + ) + buckets = services.backend.retrieve( + project=project, + tech=args.get("tech"), + phase=args.get("phase"), + polarity=args.get("polarity"), + limit_per_bucket=limit_per_bucket, + ) + _diag.step( + "mcp.memory.retrieve", + "after_retrieve_reflections", + do=len(buckets.get("do", [])), + dont=len(buckets.get("dont", [])), + neutral=len(buckets.get("neutral", [])), + ) + if diag_cid is not None: + refl_ms = int((time.monotonic() - t_reflections) * 1000) + total_ms = int((time.monotonic() - t_retrieve) * 1000) + _diag.log( + f"[bm-retrieve step=reflections cid={diag_cid} " + f"ms={refl_ms} status=ok]" + ) + _diag.log( + f"[bm-retrieve done cid={diag_cid} total_ms={total_ms}]" + ) + return [TextContent(type="text", text=json.dumps(buckets))] + + async def retrieve_observations( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + project = args.get("project") or project_name() + results = await services.observations.list_observations( + project=project, + episode_id=args.get("episode_id"), + component=args.get("component"), + theme=args.get("theme"), + outcome=args.get("outcome"), + query=args.get("query"), + limit=args.get("limit", 50), + ) + return [TextContent(type="text", text=json.dumps(results))] + + async def record_use( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + services.observations.record_use( + args["id"], + outcome=args.get("outcome"), + ) + return [TextContent(type="text", text=json.dumps({"ok": True}))] + + def handlers(self) -> list[Handler]: + return [ + Handler( + name="memory.observe", + description=_OBSERVE_DESCRIPTION, + schema=_OBSERVE_SCHEMA, + call=self.observe, + ), + Handler( + name="memory.retrieve", + description=_RETRIEVE_DESCRIPTION, + schema=_RETRIEVE_SCHEMA, + call=self.retrieve, + ), + Handler( + name="memory.retrieve_observations", + description=_RETRIEVE_OBSERVATIONS_DESCRIPTION, + schema=_RETRIEVE_OBSERVATIONS_SCHEMA, + call=self.retrieve_observations, + ), + Handler( + name="memory.record_use", + description=_RECORD_USE_DESCRIPTION, + schema=_RECORD_USE_SCHEMA, + call=self.record_use, + ), + ] diff --git a/better_memory/mcp/handlers/ratings.py b/better_memory/mcp/handlers/ratings.py new file mode 100644 index 0000000..ff95730 --- /dev/null +++ b/better_memory/mcp/handlers/ratings.py @@ -0,0 +1,165 @@ +"""Rating-domain MCP tool handlers. + +Tools: memory.list_session_exposures, memory.apply_session_ratings, +memory.credit. + +Bodies lifted verbatim from the legacy ``_call_tool`` if-chain. Session +resolution goes through :func:`resolve_session_id` so the handlers stay +testable without touching env vars or the marker file directly. + +``apply_session_ratings`` raises a verbatim multi-line ``ValueError`` if +no session is active — the rate-session-memories skill depends on the +exact message text. ``credit`` has two payload shapes: the no-session +shape ``{"applied": None, "skipped": "no_session"}`` and the +``credit_one`` outcome shape. ``list_session_exposures`` coerces an +unresolved session to ``""`` so the service receives a string and the +caller gets an empty-exposures payload back instead of an exception. +""" +from __future__ import annotations + +import json +from typing import Any + +from mcp.types import TextContent + +from better_memory.mcp._session import resolve_session_id +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler + +_LIST_SESSION_EXPOSURES_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": {}, +} + +_APPLY_SESSION_RATINGS_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "required": ["ratings"], + "properties": { + "ratings": { + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "additionalProperties": False, + "required": ["kind", "id", "class"], + "properties": { + "kind": {"enum": ["reflection", "semantic"]}, + "id": {"type": "string"}, + "class": { + "enum": [ + "cited", "shaped", "ignored", + "misled", "overlooked", + ] + }, + }, + }, + }, + }, +} + +_CREDIT_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "required": ["kind", "id", "class"], + "properties": { + "kind": {"enum": ["reflection", "semantic"]}, + "id": {"type": "string"}, + "class": {"enum": ["cited", "shaped", "misled", "overlooked"]}, + }, +} + +_LIST_SESSION_EXPOSURES_DESCRIPTION = ( + "Return the unrated session_memory_exposure rows for the " + "current Claude session (resolved server-side from " + "CLAUDE_SESSION_ID env). Read-only; no side effects. " + "Used by the rate-session-memories skill as the " + "authoritative anti-hallucination list." +) + +_APPLY_SESSION_RATINGS_DESCRIPTION = ( + "Atomic batch rating for the current Claude session " + "(resolved server-side from CLAUDE_SESSION_ID). Called " + "at session end by the rate-session-memories skill. " + "Raises if CLAUDE_SESSION_ID is unset — call only inside " + "an active Claude session." +) + +_CREDIT_DESCRIPTION = ( + "Per-tool-use credit. When you actively use a memory " + "retrieved during this session (quote it, follow its " + "guidance, or it misled you), call this immediately. " + "Resolved server-side from CLAUDE_SESSION_ID. " + "class must be 'cited', 'shaped', 'misled', or " + "'overlooked' — NOT 'ignored'. Use 'overlooked' when the " + "user pointed you back to a memory you already had but " + "had not applied." +) + + +class RatingHandlers: + """The three rating tools: list_session_exposures + apply + credit.""" + + async def list_session_exposures( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + sid = resolve_session_id(services.config.home) or "" + payload = services.session_bootstrap.list_session_exposures( + session_id=sid, + ) + return [TextContent(type="text", text=json.dumps(payload))] + + async def apply_session_ratings( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + sid = resolve_session_id(services.config.home) + if not sid: + raise ValueError( + "No active session: CLAUDE_SESSION_ID / " + "CLAUDE_CODE_SESSION_ID not set and no session marker " + "found (SessionStart hook may not have run)" + ) + payload = services.memory_rating.apply_session_ratings( + session_id=sid, + ratings=args["ratings"], + ) + return [TextContent(type="text", text=json.dumps(payload))] + + async def credit( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + sid = resolve_session_id(services.config.home) + payload: Any + if not sid: + payload = {"applied": None, "skipped": "no_session"} + else: + payload = services.memory_rating.credit_one( + session_id=sid, + kind=args["kind"], + id=args["id"], + classification=args["class"], + ) + return [TextContent(type="text", text=json.dumps(payload))] + + def handlers(self) -> list[Handler]: + return [ + Handler( + name="memory.list_session_exposures", + description=_LIST_SESSION_EXPOSURES_DESCRIPTION, + schema=_LIST_SESSION_EXPOSURES_SCHEMA, + call=self.list_session_exposures, + ), + Handler( + name="memory.apply_session_ratings", + description=_APPLY_SESSION_RATINGS_DESCRIPTION, + schema=_APPLY_SESSION_RATINGS_SCHEMA, + call=self.apply_session_ratings, + ), + Handler( + name="memory.credit", + description=_CREDIT_DESCRIPTION, + schema=_CREDIT_SCHEMA, + call=self.credit, + ), + ] diff --git a/better_memory/mcp/handlers/reflections.py b/better_memory/mcp/handlers/reflections.py new file mode 100644 index 0000000..4f0fe32 --- /dev/null +++ b/better_memory/mcp/handlers/reflections.py @@ -0,0 +1,283 @@ +"""Synthesis-domain MCP tool handlers. + +Tools: memory.synthesize_next_get_context, memory.synthesize_next_apply. +Both are CAPABILITY-GATED (``requires_synthesis=True``) — the dispatcher +hides them when ``backend.supports_synthesis`` is False. + +Bodies lifted verbatim from the legacy ``_call_tool`` if-chain to +preserve every documented invariant: the audit bracket on both tools +(start + complete JSONL rows in ``logs/synthesize.jsonl``), the +SynthesisResponseError -> validation_error fork on apply, the +ValueError -> state_error fork on apply (so the IDE-LLM refetches +context rather than retrying a stale episode_id), and the four +serializer payload shapes the IDE-LLM consumes. +""" +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.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler +from better_memory.mcp.handlers._audit import _audit_synth_call +from better_memory.services.reflection import ( + EpisodeContext, + EpisodeQueueCounts, + SynthesisResponseError, + SynthesisStep, +) + +_GET_CONTEXT_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "project": { + "type": "string", + "description": ( + "Optional project override; defaults to " + "cwd-derived." + ), + }, + }, +} + +_APPLY_SCHEMA: dict[str, Any] = { + "type": "object", + "required": ["episode_id", "decision"], + "additionalProperties": False, + "properties": { + "episode_id": {"type": "string"}, + "decision": { + "type": "object", + "description": ( + "Decision JSON; validation errors are " + "returned to the caller without stamping " + "the episode failed (caller can retry)." + ), + }, + "project": { + "type": "string", + "description": ( + "Optional project override; defaults to " + "cwd-derived." + ), + }, + }, +} + +_GET_CONTEXT_DESCRIPTION = ( + "Return the next pending episode's full context for " + "consolidation: episode metadata, all observations on it, " + "and tech-filtered existing reflections. The IDE-LLM " + "consumes this, decides what new/augment/merge/ignore " + "actions to take, and submits the decision via " + "memory.synthesize_next_apply. Returns " + '{"episode_id": null, "queue": {...}} when the queue is ' + "empty. See the better-memory-synthesize skill for the " + "full workflow and decision schema." +) +_APPLY_DESCRIPTION = ( + "Apply a synthesis decision for one episode. Atomically " + "creates new reflections, augments existing ones, merges " + "duplicates, and marks observations as consumed (or " + "ignored). Marks the episode synthesized. Returns a " + "step summary {episode_id, counts, queue, failure}. " + "decision shape: {new: [...], augment: [...], " + "merge: [...], ignore: [...]} — see the " + "better-memory-synthesize skill for the per-entry " + "field schema." +) + + +def _serialize_queue(queue: EpisodeQueueCounts) -> dict[str, int]: + return { + "pending": queue.pending, + "in_cooldown": queue.in_cooldown, + "done": queue.done, + } + + +def _serialize_synth_get_context( + ctx: EpisodeContext | None, + queue: EpisodeQueueCounts, +) -> dict[str, Any]: + """Build the JSON payload for ``memory.synthesize_next_get_context``. + + Returns ``{"episode_id": null, "queue": {...}}`` when the queue is + empty. Otherwise returns the full episode + observations + reflections + bundle the IDE-LLM consumes to decide on synthesis actions. + """ + queue_json = _serialize_queue(queue) + if ctx is None: + return {"episode_id": None, "queue": queue_json} + return { + "episode_id": ctx.episode.id, + "queue": queue_json, + "episode": { + "id": ctx.episode.id, + "project": ctx.episode.project, + "goal": ctx.episode.goal, + "tech": ctx.episode.tech, + "outcome": ctx.episode.outcome, + }, + "observations": [ + { + "id": o.id, + "content": o.content, + "outcome": o.outcome, + "component": o.component, + "theme": o.theme, + "tech": o.tech, + "created_at": o.created_at, + "status": o.status, + } + for o in ctx.observations + ], + "reflections": [ + { + "id": r.id, + "title": r.title, + "tech": r.tech, + "phase": r.phase, + "polarity": r.polarity, + "use_cases": r.use_cases, + "hints": json.loads(r.hints) if r.hints else [], + "confidence": r.confidence, + "status": r.status, + } + for r in ctx.reflections + ], + } + + +def _serialize_synth_apply_ok(step: SynthesisStep) -> dict[str, Any]: + """Build the JSON payload for a successful ``memory.synthesize_next_apply``.""" + return { + "ok": True, + "episode_id": step.episode_id, + "counts": step.counts, + "queue": _serialize_queue(step.queue), + } + + +def _serialize_synth_apply_validation_error(message: str) -> dict[str, Any]: + """Build the JSON payload for a decision-validation failure. + + Validation errors do NOT stamp ``synth_failed_at``; the caller can + retry with a corrected payload. + """ + return { + "ok": False, + "error": "validation", + "message": message, + } + + +def _serialize_synth_apply_state_error(message: str) -> dict[str, Any]: + """Build the JSON payload for an episode-state failure. + + Surfaces apply-time precondition violations (episode not found, + wrong project, already synthesized) without raising into the MCP + framework's generic isError surface. Caller should NOT retry the + same episode_id; pull fresh context via + ``memory.synthesize_next_get_context`` first. + """ + return { + "ok": False, + "error": "state", + "message": message, + } + + +class ReflectionHandlers: + """The memory.synthesize_next_get_context / .synthesize_next_apply tools.""" + + async def get_context( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + project = args.get("project") or project_name() + with _audit_synth_call( + services.config.home, + tool="get_context", + project=project, + episode_id=None, + ) as audit: + ctx = services.reflections.get_next_pending_context(project=project) + queue = services.reflections.read_queue_counts(project=project) + payload = _serialize_synth_get_context(ctx, queue) + if ctx is None: + audit["result_kind"] = "empty" + else: + audit["result_kind"] = "episode" + audit["episode_id"] = ctx.episode.id + audit["obs_count"] = len(ctx.observations) + audit["refl_count"] = len(ctx.reflections) + return [TextContent(type="text", text=json.dumps(payload))] + + async def apply( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + project = args.get("project") or project_name() + episode_id = args["episode_id"] + decision = args["decision"] + with _audit_synth_call( + services.config.home, + tool="apply", + project=project, + episode_id=episode_id, + ) as audit: + try: + # ``decision`` is already a parsed dict (the MCP + # framework decoded it before dispatch). Use the + # dict-shape parser directly to skip a redundant + # json.dumps → json.loads round-trip. + response = services.reflections.parse_response_dict(decision) + except SynthesisResponseError as exc: + audit["result_kind"] = "validation_error" + audit["error"] = str(exc) + payload = _serialize_synth_apply_validation_error(str(exc)) + return [TextContent(type="text", text=json.dumps(payload))] + try: + step = services.reflections.apply_decision( + episode_id=episode_id, + response=response, + project=project, + ) + except ValueError as exc: + # Episode-state preconditions: not found / wrong + # project / already synthesized. Surface as + # structured error so the IDE-LLM can refetch + # context instead of retrying the same stale id. + audit["result_kind"] = "state_error" + audit["error"] = str(exc) + payload = _serialize_synth_apply_state_error(str(exc)) + return [TextContent(type="text", text=json.dumps(payload))] + audit["result_kind"] = "applied" + audit["counts"] = step.counts + return [ + TextContent( + type="text", + text=json.dumps(_serialize_synth_apply_ok(step)), + ) + ] + + def handlers(self) -> list[Handler]: + return [ + Handler( + name="memory.synthesize_next_get_context", + description=_GET_CONTEXT_DESCRIPTION, + schema=_GET_CONTEXT_SCHEMA, + call=self.get_context, + requires_synthesis=True, + ), + Handler( + name="memory.synthesize_next_apply", + description=_APPLY_DESCRIPTION, + schema=_APPLY_SCHEMA, + call=self.apply, + requires_synthesis=True, + ), + ] diff --git a/better_memory/mcp/handlers/retention.py b/better_memory/mcp/handlers/retention.py new file mode 100644 index 0000000..8696ba0 --- /dev/null +++ b/better_memory/mcp/handlers/retention.py @@ -0,0 +1,100 @@ +"""Retention-domain MCP tool handler. + +Tools: memory.run_retention. + +Body lifted verbatim from the legacy ``_call_tool`` if-chain. The +serialized payload exposes the four ``RetentionReport`` fields the +spec §9 audit pipeline consumes (three ``archived_via_*`` counters +plus ``pruned``). +""" +from __future__ import annotations + +import json +from typing import Any + +from mcp.types import TextContent + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler + +_RUN_RETENTION_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "retention_days": { + "type": "integer", + "default": 90, + "description": ( + "Age threshold for the three archive " + "rules. Default 90 (per spec §9)." + ), + }, + "prune": { + "type": "boolean", + "default": False, + "description": ( + "If true, also hard-delete archived rows " + "older than prune_age_days." + ), + }, + "prune_age_days": { + "type": "integer", + "default": 365, + "description": ( + "Age threshold for prune mode. Default 365." + ), + }, + "dry_run": { + "type": "boolean", + "default": False, + "description": ( + "If true, return the counts without " + "writing any changes to the DB." + ), + }, + }, +} + +_RUN_RETENTION_DESCRIPTION = ( + "Apply spec §9 retention rules — flip eligible " + "observations to status='archived' and optionally " + "hard-delete archived rows older than prune_age_days." +) + + +class RetentionHandlers: + """The memory.run_retention tool.""" + + async def run_retention( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + report = services.retention.run( + retention_days=args.get("retention_days", 90), + prune=args.get("prune", False), + prune_age_days=args.get("prune_age_days", 365), + dry_run=args.get("dry_run", False), + ) + return [ + TextContent( + type="text", + text=json.dumps({ + "archived_via_retired_reflection": + report.archived_via_retired_reflection, + "archived_via_consumed_without_reflection": + report.archived_via_consumed_without_reflection, + "archived_via_no_outcome_episode": + report.archived_via_no_outcome_episode, + "pruned": report.pruned, + }), + ) + ] + + def handlers(self) -> list[Handler]: + return [ + Handler( + name="memory.run_retention", + description=_RUN_RETENTION_DESCRIPTION, + schema=_RUN_RETENTION_SCHEMA, + call=self.run_retention, + ), + ] diff --git a/better_memory/mcp/handlers/semantics.py b/better_memory/mcp/handlers/semantics.py new file mode 100644 index 0000000..ce77041 --- /dev/null +++ b/better_memory/mcp/handlers/semantics.py @@ -0,0 +1,160 @@ +"""Semantic-memory MCP tool handlers (memory.semantic_*). + +SemanticMemoryService is now built once in ``_build_services`` and lives +on the container at ``services.semantic`` — eliminating the 4x per-call +inline construction smell in the legacy ``_call_tool`` if-chain. +""" +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.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler + +_OBSERVE_SCHEMA: dict[str, Any] = { + "type": "object", + "required": ["content"], + "additionalProperties": False, + "properties": { + "content": {"type": "string"}, + "scope": { + "type": "string", + "enum": ["project", "general"], + "description": ( + "'project' (default) for project-scoped rules; " + "'general' for cross-project workflow rules." + ), + }, + }, +} + +_RETRIEVE_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "project": { + "type": "string", + "description": ( + "Optional project override; " + "defaults to cwd-derived." + ), + }, + }, +} + +_UPDATE_SCHEMA: dict[str, Any] = { + "type": "object", + "required": ["id", "content"], + "additionalProperties": False, + "properties": { + "id": {"type": "string"}, + "content": {"type": "string"}, + }, +} + +_DELETE_SCHEMA: dict[str, Any] = { + "type": "object", + "required": ["id"], + "additionalProperties": False, + "properties": { + "id": {"type": "string"}, + }, +} + +_OBSERVE_DESCRIPTION = ( + "Record a user-stated fact or preference. Distinct from " + "memory.observe (episodic): semantic memories are " + "user-asserted current truths, retrieved at session " + "startup. Set scope='general' for cross-project rules." +) +_RETRIEVE_DESCRIPTION = ( + "Return user-stated facts/preferences for the current " + "project, merged with all general-scope semantic memories. " + "Flat list ordered newest-first." +) +_UPDATE_DESCRIPTION = ( + "Edit a semantic memory's content in place. Bumps updated_at." +) +_DELETE_DESCRIPTION = ( + "Remove a semantic memory. Idempotent — no error if id absent." +) + + +class SemanticHandlers: + """All memory.semantic_* tools (observe, retrieve, update, delete).""" + + async def observe( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + project = project_name() + # `args.get("scope") or "project"` (not `, "project"` default) defends + # against MCP clients sending {"scope": null} — dict.get returns the + # default only when the key is absent, not when its value is None. + memory_id = services.semantic.create( + content=args["content"], + project=project, + scope=args.get("scope") or "project", + ) + return [TextContent(type="text", text=json.dumps({"id": memory_id}))] + + async def retrieve( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + project = args.get("project") or project_name() + memories = services.semantic.list_for_project(project=project) + payload = [ + { + "id": m.id, + "content": m.content, + "project": m.project, + "scope": m.scope, + "created_at": m.created_at, + "updated_at": m.updated_at, + } + for m in memories + ] + return [TextContent(type="text", text=json.dumps(payload))] + + async def update( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + services.semantic.update_text(id=args["id"], content=args["content"]) + return [TextContent(type="text", text=json.dumps({"ok": True}))] + + async def delete( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + services.semantic.delete(id=args["id"]) + return [TextContent(type="text", text=json.dumps({"ok": True}))] + + def handlers(self) -> list[Handler]: + return [ + Handler( + name="memory.semantic_observe", + description=_OBSERVE_DESCRIPTION, + schema=_OBSERVE_SCHEMA, + call=self.observe, + ), + Handler( + name="memory.semantic_retrieve", + description=_RETRIEVE_DESCRIPTION, + schema=_RETRIEVE_SCHEMA, + call=self.retrieve, + ), + Handler( + name="memory.semantic_update", + description=_UPDATE_DESCRIPTION, + schema=_UPDATE_SCHEMA, + call=self.update, + ), + Handler( + name="memory.semantic_delete", + description=_DELETE_DESCRIPTION, + schema=_DELETE_SCHEMA, + call=self.delete, + ), + ] diff --git a/better_memory/mcp/handlers/session.py b/better_memory/mcp/handlers/session.py new file mode 100644 index 0000000..21fa4e6 --- /dev/null +++ b/better_memory/mcp/handlers/session.py @@ -0,0 +1,132 @@ +"""Session-domain MCP tool handlers. + +Tools: memory.session_bootstrap, memory.start_ui. + +Bodies lifted verbatim from the legacy ``_call_tool`` if-chain. +``session_bootstrap`` resolves the SessionStart session_id through a +4-tier fallback (args → ``CLAUDE_SESSION_ID`` env → ``CLAUDE_CODE_SESSION_ID`` +env → fresh UUID) and serialises the ``SessionBootstrapResult`` into the +5-key payload (``additionalContext``, ``project``, ``source``, +``episode``, ``counts``) the SessionStart hook expects. ``start_ui`` +delegates to the ``ui_launcher`` module which spawns or reuses the +management UI server. +""" +from __future__ import annotations + +import json +import os +import uuid +from pathlib import Path +from typing import Any + +from mcp.types import TextContent + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler +from better_memory.services import ui_launcher + +_BOOTSTRAP_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": { + "source": { + "type": "string", + "enum": ["startup", "resume", "clear", "compact"], + "description": ( + "SessionStart payload source. Unknown values " + "coerce to 'startup' inside the service." + ), + }, + "session_id": { + "type": "string", + "description": ( + "Optional SessionStart payload session_id. " + "Defaults to $CLAUDE_SESSION_ID env var, or a " + "fresh UUID." + ), + }, + "cwd": { + "type": "string", + "description": ( + "Optional working directory. Defaults to " + "server's process cwd." + ), + }, + }, +} + +_START_UI_SCHEMA: dict[str, Any] = { + "type": "object", + "additionalProperties": False, + "properties": {}, +} + +_BOOTSTRAP_DESCRIPTION = ( + "Open or reuse a session episode and inject all project + " + "general semantic memories and reflections as " + "additionalContext markdown. Mirrors what the SessionStart " + "hook does; callable manually for recovery, testing, or " + "post-/clear re-injection." +) + +_START_UI_DESCRIPTION = ( + "Spawn or reuse the better-memory management UI. Returns " + '{"url": str, "reused": bool}. Reuses an existing live UI ' + "when one is already running on /healthz." +) + + +class SessionHandlers: + """The memory.session_bootstrap / memory.start_ui tools.""" + + async def session_bootstrap( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + cwd_arg = args.get("cwd") or os.getcwd() + session_id_arg = ( + args.get("session_id") + or os.environ.get("CLAUDE_SESSION_ID") + or os.environ.get("CLAUDE_CODE_SESSION_ID") + or uuid.uuid4().hex + ) + result = services.session_bootstrap.bootstrap( + source=args.get("source"), + session_id=session_id_arg, + cwd=Path(cwd_arg), + ) + payload = { + "additionalContext": result.additional_context, + "project": result.project, + "source": result.source, + "episode": { + "id": result.episode_id, + "action": result.episode_action, + }, + "counts": { + "semantic": result.semantic_count, + "reflections": result.reflections_counts, + }, + } + return [TextContent(type="text", text=json.dumps(payload))] + + async def start_ui( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + result = ui_launcher.start_ui() + return [TextContent(type="text", text=json.dumps(result))] + + def handlers(self) -> list[Handler]: + return [ + Handler( + name="memory.session_bootstrap", + description=_BOOTSTRAP_DESCRIPTION, + schema=_BOOTSTRAP_SCHEMA, + call=self.session_bootstrap, + ), + Handler( + name="memory.start_ui", + description=_START_UI_DESCRIPTION, + schema=_START_UI_SCHEMA, + call=self.start_ui, + ), + ] diff --git a/better_memory/mcp/server.py b/better_memory/mcp/server.py index d1422aa..447aab2 100644 --- a/better_memory/mcp/server.py +++ b/better_memory/mcp/server.py @@ -4,21 +4,12 @@ as MCP tools over stdio. On startup, the knowledge-base is reindexed (mtime-only, so this is cheap and idempotent) as a session-start step. -Tools ------ -* ``memory.observe`` — create an observation. -* ``memory.retrieve`` — reflections bucketed by polarity. -* ``memory.retrieve_observations`` — raw observation drill-down. -* ``memory.record_use`` — record re-use of a memory. -* ``memory.semantic_observe`` / _retrieve / _update / _delete — user-stated facts. -* ``memory.start_ui`` — spawn or reuse the management UI. -* ``memory.start_episode`` / _close_episode / _list_episodes / _reconcile_episodes - — episode lifecycle. -* ``memory.synthesize_next_get_context`` — episode-context fetch for synthesis. -* ``memory.synthesize_next_apply`` — apply IDE-LLM's decision JSON. -* ``memory.session_bootstrap`` — open/reuse a session episode + inject memories. -* ``memory.run_retention`` — apply spec §9 retention rules. -* ``knowledge.search`` / ``knowledge.list`` — knowledge-base introspection. +Tool registration and dispatch are owned by +:class:`better_memory.mcp.dispatcher.ToolDispatcher`. Each tool's schema, +description, and async handler live in the per-domain modules under +``better_memory.mcp.handlers``. ``create_server`` is now a thin wiring +seam: build services, build the dispatcher, register two MCP callbacks +that delegate to it. Connection ownership -------------------- @@ -36,21 +27,15 @@ from __future__ import annotations import asyncio -import contextlib -import json import logging -import os import sqlite3 import sys -import time import urllib.error import urllib.request -import uuid -from collections.abc import Callable, Coroutine, Iterator +from collections.abc import Callable, Coroutine from dataclasses import dataclass -from datetime import UTC, datetime from pathlib import Path -from typing import Any, cast +from typing import Any from mcp.server import Server from mcp.server.stdio import stdio_server @@ -61,24 +46,34 @@ from better_memory.db.connection import connect from better_memory.db.schema import apply_migrations from better_memory.embeddings.ollama import OllamaEmbedder -from better_memory.runtime.session_marker import read_session_id -from better_memory.services import ui_launcher -from better_memory.services.episode import EpisodeService -from better_memory.services.knowledge import ( - KnowledgeDocument, - KnowledgeSearchResult, - KnowledgeService, -) -from better_memory.services.observation import ObservationService -from better_memory.services.reflection import ( - EpisodeContext, - EpisodeQueueCounts, - ReflectionSynthesisService, - SynthesisStep, + +# Re-exported for tests that still import ``_run_best_effort`` from +# ``better_memory.mcp.server`` (test_best_effort_logging). The canonical +# definition lives in ``_best_effort`` so handler modules can import it +# without creating a server <-> handlers import cycle. +from better_memory.mcp._best_effort import _run_best_effort # noqa: F401 +from better_memory.mcp._session import resolve_session_id +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.dispatcher import ToolDispatcher +from better_memory.mcp.handlers import all_handlers + +# Re-exported for tests that import the synth serializers from server.py +# (test_synthesize_tools). The canonical definitions live alongside the +# ReflectionHandlers; we re-export here only to keep the existing test +# import paths working. +from better_memory.mcp.handlers.reflections import ( # noqa: F401 + _serialize_queue, + _serialize_synth_apply_ok, + _serialize_synth_apply_state_error, + _serialize_synth_apply_validation_error, + _serialize_synth_get_context, ) +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.retention_scheduler import RetentionScheduler from better_memory.services.spool import SpoolService from better_memory.storage import StorageBackend, build_backend @@ -92,135 +87,6 @@ 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= 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}]" - ) - - -# --------------------------------------------------------- synthesize audit log -# -# 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. - - -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) - ) - - -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) - - def _probe_ollama(host: str) -> None: """Log a clear stderr warning if Ollama isn't reachable. Never raise. @@ -253,684 +119,21 @@ def _probe_ollama(host: str) -> None: ) -# --------------------------------------------------------------------------- tools - - def _tool_definitions(*, supports_synthesis: bool = True) -> list[Tool]: - """Return the list of tools exposed over MCP. + """Return the static list of tools exposed over MCP, gated by capability. - When ``supports_synthesis`` is False, the - ``memory.synthesize_next_get_context`` and ``memory.synthesize_next_apply`` - tools are omitted. This gates the synthesis surface on the active - StorageBackend's capability flag — backends without a local episode - queue (e.g. AgentCoreBackend in Plan 2) do not expose these tools. + Used by tests that want to assert tool registration without standing + up a full :class:`ToolDispatcher`. In production, ``create_server`` + builds a dispatcher and calls its own ``tool_definitions()`` against + the live backend's capability flag — this helper is the same list, + derived from the same ``all_handlers()`` source. """ - tools: list[Tool] = [ - Tool( - name="memory.observe", - description=( - "Record an observation about the current session (a fact, " - "decision, bug fix, or outcome). Returns the new observation id." - ), - inputSchema={ - "type": "object", - "required": ["content"], - "additionalProperties": False, - "properties": { - "content": {"type": "string"}, - "component": {"type": "string"}, - "theme": {"type": "string"}, - "trigger_type": {"type": "string"}, - "outcome": { - "type": "string", - "enum": ["success", "failure", "neutral"], - }, - "tech": {"type": "string"}, - "scope": { - "type": "string", - "enum": ["project", "general"], - "description": ( - "'project' (default) for project-scoped observations; " - "'general' for cross-project workflow rules that should " - "surface in every project's memory_retrieve." - ), - }, - }, - }, - ), - Tool( - name="memory.semantic_observe", - description=( - "Record a user-stated fact or preference. Distinct from " - "memory.observe (episodic): semantic memories are " - "user-asserted current truths, retrieved at session " - "startup. Set scope='general' for cross-project rules." - ), - inputSchema={ - "type": "object", - "required": ["content"], - "additionalProperties": False, - "properties": { - "content": {"type": "string"}, - "scope": { - "type": "string", - "enum": ["project", "general"], - "description": ( - "'project' (default) for project-scoped rules; " - "'general' for cross-project workflow rules." - ), - }, - }, - }, - ), - Tool( - name="memory.semantic_retrieve", - description=( - "Return user-stated facts/preferences for the current " - "project, merged with all general-scope semantic memories. " - "Flat list ordered newest-first." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "properties": { - "project": { - "type": "string", - "description": ( - "Optional project override; " - "defaults to cwd-derived." - ), - }, - }, - }, - ), - Tool( - name="memory.semantic_update", - description=( - "Edit a semantic memory's content in place. Bumps updated_at." - ), - inputSchema={ - "type": "object", - "required": ["id", "content"], - "additionalProperties": False, - "properties": { - "id": {"type": "string"}, - "content": {"type": "string"}, - }, - }, - ), - Tool( - name="memory.semantic_delete", - description=( - "Remove a semantic memory. Idempotent — no error if id absent." - ), - inputSchema={ - "type": "object", - "required": ["id"], - "additionalProperties": False, - "properties": { - "id": {"type": "string"}, - }, - }, - ), - Tool( - name="memory.retrieve", - description=( - "Retrieve reflections (do / dont / neutral lessons distilled " - "from prior observations) bucketed by polarity. Filter by " - "project, tech, phase, and polarity. For raw observation " - "lookup, use memory.retrieve_observations." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "properties": { - "project": {"type": "string"}, - "tech": {"type": "string"}, - "phase": { - "type": "string", - "enum": ["planning", "implementation", "general"], - }, - "polarity": { - "type": "string", - "enum": ["do", "dont", "neutral"], - }, - "limit_per_bucket": {"type": "integer"}, - }, - }, - ), - Tool( - name="memory.retrieve_observations", - description=( - "Retrieve raw observations matching given filters. Drill-down " - "tool — use memory.retrieve for the distilled-reflections " - "default. With ``query``, results are ranked by hybrid " - "FTS5 + sqlite-vec relevance; without, ordered created_at " - "DESC. ``episode_id`` and ``theme`` filters are ignored " - "in query mode." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "properties": { - "project": {"type": "string"}, - "episode_id": {"type": "string"}, - "component": {"type": "string"}, - "theme": {"type": "string"}, - "outcome": { - "type": "string", - "enum": ["success", "failure", "neutral"], - }, - "query": {"type": "string"}, - "limit": {"type": "integer"}, - }, - }, - ), - Tool( - name="memory.record_use", - description=( - "Record that an observation was used; optionally mark the " - "outcome as success or failure to reinforce the memory." - ), - inputSchema={ - "type": "object", - "required": ["id"], - "additionalProperties": False, - "properties": { - "id": {"type": "string"}, - "outcome": { - "type": "string", - "enum": ["success", "failure"], - }, - }, - }, - ), - Tool( - name="knowledge.search", - description=( - "BM25 search against the knowledge-base markdown corpus. " - "Returns document paths and rank." - ), - inputSchema={ - "type": "object", - "required": ["query"], - "additionalProperties": False, - "properties": { - "query": {"type": "string"}, - "project": {"type": "string"}, - }, - }, - ), - Tool( - name="knowledge.list", - description=( - "List indexed knowledge documents. When ``project`` is " - "supplied, project-scoped rows are filtered to that project." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "properties": { - "project": {"type": "string"}, - }, - }, - ), - Tool( - name="memory.start_ui", - description=( - "Spawn or reuse the better-memory management UI. Returns " - '{"url": str, "reused": bool}. Reuses an existing live UI ' - "when one is already running on /healthz." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "properties": {}, - }, - ), - Tool( - name="memory.start_episode", - description=( - "Declare a goal for the current session. Opens a new " - "foreground episode or hardens the existing background " - "episode. Returns the active episode id." - ), - inputSchema={ - "type": "object", - "required": ["goal"], - "additionalProperties": False, - "properties": { - "goal": {"type": "string"}, - "tech": {"type": "string"}, - }, - }, - ), - Tool( - name="memory.close_episode", - description=( - "Close the current session's active episode. outcome is one " - "of success / partial / abandoned / no_outcome." - ), - inputSchema={ - "type": "object", - "required": ["outcome"], - "additionalProperties": False, - "properties": { - "outcome": { - "type": "string", - "enum": [ - "success", - "partial", - "abandoned", - "no_outcome", - ], - }, - "close_reason": { - "type": "string", - "enum": [ - "goal_complete", - "plan_complete", - "abandoned", - "superseded", - "session_end_reconciled", - ], - }, - "summary": {"type": "string"}, - }, - }, - ), - Tool( - name="memory.reconcile_episodes", - description=( - "List episodes that are still open from prior sessions, " - "for the LLM to prompt the user about." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "properties": {}, - }, - ), - Tool( - name="memory.list_episodes", - description=( - "List episodes with optional filters. For UI and LLM " - "introspection." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "properties": { - "project": {"type": "string"}, - "outcome": { - "type": "string", - "enum": [ - "success", - "partial", - "abandoned", - "no_outcome", - ], - }, - "only_open": {"type": "boolean"}, - }, - }, - ), - Tool( - name="memory.session_bootstrap", - description=( - "Open or reuse a session episode and inject all project + " - "general semantic memories and reflections as " - "additionalContext markdown. Mirrors what the SessionStart " - "hook does; callable manually for recovery, testing, or " - "post-/clear re-injection." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "properties": { - "source": { - "type": "string", - "enum": ["startup", "resume", "clear", "compact"], - "description": ( - "SessionStart payload source. Unknown values " - "coerce to 'startup' inside the service." - ), - }, - "session_id": { - "type": "string", - "description": ( - "Optional SessionStart payload session_id. " - "Defaults to $CLAUDE_SESSION_ID env var, or a " - "fresh UUID." - ), - }, - "cwd": { - "type": "string", - "description": ( - "Optional working directory. Defaults to " - "server's process cwd." - ), - }, - }, - }, - ), - Tool( - name="memory.run_retention", - description=( - "Apply spec §9 retention rules — flip eligible " - "observations to status='archived' and optionally " - "hard-delete archived rows older than prune_age_days." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "properties": { - "retention_days": { - "type": "integer", - "default": 90, - "description": ( - "Age threshold for the three archive " - "rules. Default 90 (per spec §9)." - ), - }, - "prune": { - "type": "boolean", - "default": False, - "description": ( - "If true, also hard-delete archived rows " - "older than prune_age_days." - ), - }, - "prune_age_days": { - "type": "integer", - "default": 365, - "description": ( - "Age threshold for prune mode. Default 365." - ), - }, - "dry_run": { - "type": "boolean", - "default": False, - "description": ( - "If true, return the counts without " - "writing any changes to the DB." - ), - }, - }, - }, - ), - Tool( - name="memory.list_session_exposures", - description=( - "Return the unrated session_memory_exposure rows for the " - "current Claude session (resolved server-side from " - "CLAUDE_SESSION_ID env). Read-only; no side effects. " - "Used by the rate-session-memories skill as the " - "authoritative anti-hallucination list." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "properties": {}, - }, - ), - Tool( - name="memory.apply_session_ratings", - description=( - "Atomic batch rating for the current Claude session " - "(resolved server-side from CLAUDE_SESSION_ID). Called " - "at session end by the rate-session-memories skill. " - "Raises if CLAUDE_SESSION_ID is unset — call only inside " - "an active Claude session." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "required": ["ratings"], - "properties": { - "ratings": { - "type": "array", - "minItems": 1, - "items": { - "type": "object", - "additionalProperties": False, - "required": ["kind", "id", "class"], - "properties": { - "kind": {"enum": ["reflection", "semantic"]}, - "id": {"type": "string"}, - "class": { - "enum": [ - "cited", "shaped", "ignored", - "misled", "overlooked", - ] - }, - }, - }, - }, - }, - }, - ), - Tool( - name="memory.credit", - description=( - "Per-tool-use credit. When you actively use a memory " - "retrieved during this session (quote it, follow its " - "guidance, or it misled you), call this immediately. " - "Resolved server-side from CLAUDE_SESSION_ID. " - "class must be 'cited', 'shaped', 'misled', or " - "'overlooked' — NOT 'ignored'. Use 'overlooked' when the " - "user pointed you back to a memory you already had but " - "had not applied." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "required": ["kind", "id", "class"], - "properties": { - "kind": {"enum": ["reflection", "semantic"]}, - "id": {"type": "string"}, - "class": {"enum": ["cited", "shaped", "misled", "overlooked"]}, - }, - }, - ), + return [ + Tool(name=h.name, description=h.description, inputSchema=h.schema) + for h in all_handlers() + if supports_synthesis or not h.requires_synthesis ] - if supports_synthesis: - tools.extend([ - Tool( - name="memory.synthesize_next_get_context", - description=( - "Return the next pending episode's full context for " - "consolidation: episode metadata, all observations on it, " - "and tech-filtered existing reflections. The IDE-LLM " - "consumes this, decides what new/augment/merge/ignore " - "actions to take, and submits the decision via " - "memory.synthesize_next_apply. Returns " - '{"episode_id": null, "queue": {...}} when the queue is ' - "empty. See the better-memory-synthesize skill for the " - "full workflow and decision schema." - ), - inputSchema={ - "type": "object", - "additionalProperties": False, - "properties": { - "project": { - "type": "string", - "description": ( - "Optional project override; defaults to " - "cwd-derived." - ), - }, - }, - }, - ), - Tool( - name="memory.synthesize_next_apply", - description=( - "Apply a synthesis decision for one episode. Atomically " - "creates new reflections, augments existing ones, merges " - "duplicates, and marks observations as consumed (or " - "ignored). Marks the episode synthesized. Returns a " - "step summary {episode_id, counts, queue, failure}. " - "decision shape: {new: [...], augment: [...], " - "merge: [...], ignore: [...]} — see the " - "better-memory-synthesize skill for the per-entry " - "field schema." - ), - inputSchema={ - "type": "object", - "required": ["episode_id", "decision"], - "additionalProperties": False, - "properties": { - "episode_id": {"type": "string"}, - "decision": { - "type": "object", - "description": ( - "Decision JSON; validation errors are " - "returned to the caller without stamping " - "the episode failed (caller can retry)." - ), - }, - "project": { - "type": "string", - "description": ( - "Optional project override; defaults to " - "cwd-derived." - ), - }, - }, - }, - ), - ]) - - return tools - - -# --------------------------------------------------------------------------- helpers - - -def _serialize_knowledge_search(result: KnowledgeSearchResult) -> dict[str, Any]: - doc = result.document - return { - "path": doc.path, - "scope": doc.scope, - "project": doc.project, - "language": doc.language, - "rank": result.rank, - } - - -def _serialize_knowledge_doc(doc: KnowledgeDocument) -> dict[str, Any]: - return { - "path": doc.path, - "scope": doc.scope, - "project": doc.project, - "language": doc.language, - } - - -def _serialize_queue(queue: EpisodeQueueCounts) -> dict[str, int]: - return { - "pending": queue.pending, - "in_cooldown": queue.in_cooldown, - "done": queue.done, - } - - -def _serialize_synth_get_context( - ctx: EpisodeContext | None, - queue: EpisodeQueueCounts, -) -> dict[str, Any]: - """Build the JSON payload for ``memory.synthesize_next_get_context``. - - Returns ``{"episode_id": null, "queue": {...}}`` when the queue is - empty. Otherwise returns the full episode + observations + reflections - bundle the IDE-LLM consumes to decide on synthesis actions. - """ - queue_json = _serialize_queue(queue) - if ctx is None: - return {"episode_id": None, "queue": queue_json} - return { - "episode_id": ctx.episode.id, - "queue": queue_json, - "episode": { - "id": ctx.episode.id, - "project": ctx.episode.project, - "goal": ctx.episode.goal, - "tech": ctx.episode.tech, - "outcome": ctx.episode.outcome, - }, - "observations": [ - { - "id": o.id, - "content": o.content, - "outcome": o.outcome, - "component": o.component, - "theme": o.theme, - "tech": o.tech, - "created_at": o.created_at, - "status": o.status, - } - for o in ctx.observations - ], - "reflections": [ - { - "id": r.id, - "title": r.title, - "tech": r.tech, - "phase": r.phase, - "polarity": r.polarity, - "use_cases": r.use_cases, - "hints": json.loads(r.hints) if r.hints else [], - "confidence": r.confidence, - "status": r.status, - } - for r in ctx.reflections - ], - } - - -def _serialize_synth_apply_ok(step: SynthesisStep) -> dict[str, Any]: - """Build the JSON payload for a successful ``memory.synthesize_next_apply``.""" - return { - "ok": True, - "episode_id": step.episode_id, - "counts": step.counts, - "queue": _serialize_queue(step.queue), - } - - -def _serialize_synth_apply_validation_error(message: str) -> dict[str, Any]: - """Build the JSON payload for a decision-validation failure. - - Validation errors do NOT stamp ``synth_failed_at``; the caller can - retry with a corrected payload. - """ - return { - "ok": False, - "error": "validation", - "message": message, - } - - -def _serialize_synth_apply_state_error(message: str) -> dict[str, Any]: - """Build the JSON payload for an episode-state failure. - - Surfaces apply-time precondition violations (episode not found, - wrong project, already synthesized) without raising into the MCP - framework's generic isError surface. Caller should NOT retry the - same episode_id; pull fresh context via - ``memory.synthesize_next_get_context`` first. - """ - return { - "ok": False, - "error": "state", - "message": message, - } - - -# --------------------------------------------------------------------------- factory - @dataclass class ServerContext: @@ -942,12 +145,109 @@ class ServerContext: ``memory_conn`` is the underlying sqlite connection (None for non-sqlite backends in future plans); ``embedder`` is whatever was passed to the backend (None for the sqlite/FTS5 embeddings backend, which indexes via - DB triggers and needs no Python embedder). + DB triggers and needs no Python embedder). ``dispatcher`` is the live + :class:`ToolDispatcher` so tests can route tool calls without going + through the SDK request stack. """ backend: StorageBackend memory_conn: sqlite3.Connection | None = None embedder: Any = None + dispatcher: ToolDispatcher | None = None + + +def _build_services( + config: Any, # better_memory.config.Config + memory_conn: sqlite3.Connection, + knowledge_conn: sqlite3.Connection, + embedder: OllamaEmbedder | None, + *, + startup_project: str, + startup_session_id: str | None, +) -> ServiceContainer: + """Construct every service exactly once. + + Replaces the inline 4x ``SemanticMemoryService(memory_conn)`` and the + inline 2x ``SessionBootstrapService(...)`` smells in the legacy + ``_call_tool`` body. + """ + from better_memory.services.semantic import SemanticMemoryService + from better_memory.services.session_bootstrap import SessionBootstrapService + + episodes = EpisodeService(memory_conn) + observations = ObservationService( + memory_conn, embedder=embedder, episodes=episodes, + ) + backend = build_backend( + config=config, + memory_conn=memory_conn, + embedder=embedder, + session_id=startup_session_id, + project=startup_project, + ) + reflections = ReflectionSynthesisService(memory_conn) + retention = RetentionService(conn=memory_conn) + memory_rating = MemoryRatingService(memory_conn) + knowledge = KnowledgeService( + knowledge_conn, knowledge_base=config.knowledge_base, + ) + spool = SpoolService(memory_conn, config.spool_dir, episodes=episodes) + semantic = SemanticMemoryService(memory_conn) + session_bootstrap = SessionBootstrapService(memory_conn) + + return ServiceContainer( + config=config, + memory_conn=memory_conn, + backend=backend, + episodes=episodes, + observations=observations, + reflections=reflections, + retention=retention, + memory_rating=memory_rating, + knowledge=knowledge, + spool=spool, + semantic=semantic, + session_bootstrap=session_bootstrap, + ) + + +def _make_cleanup( + memory_conn: sqlite3.Connection, + knowledge_conn: sqlite3.Connection, + embedder: OllamaEmbedder | None, +) -> Callable[[], Coroutine[Any, Any, None]]: + """Build the idempotent shutdown closure. + + Closes both SQLite connections and the embedder's HTTP client. + Safe to call multiple times: SQLite ``Connection.close`` is a no-op + after the first call, and a local ``cleaned`` flag guards against + double-closing the embedder's httpx client. + """ + cleaned = False + + async def cleanup() -> None: + nonlocal cleaned + if cleaned: + return + cleaned = True + try: + memory_conn.close() + except Exception: # noqa: BLE001 — best-effort shutdown + pass + try: + knowledge_conn.close() + except Exception: # noqa: BLE001 — best-effort shutdown + pass + # In the sqlite backend no embedder is built (FTS5 triggers handle + # indexing). Guard against the None case so this cleanup stays + # idempotent across both backends. + if embedder is not None: + try: + await embedder.aclose() + except Exception: # noqa: BLE001 — best-effort shutdown + pass + + return cleanup def create_server() -> tuple[ @@ -961,8 +261,20 @@ def create_server() -> tuple[ idempotent async function that closes the two SQLite connections and the embedder's HTTP client, and ``ctx`` is a :class:`ServerContext` bundling the live :class:`StorageBackend`, its underlying memory - connection, and the embedder (if any). Callers must await ``cleanup`` - on shutdown (typically in a ``finally`` around ``server.run``). + connection, the embedder (if any), and the :class:`ToolDispatcher`. + Callers must await ``cleanup`` on shutdown (typically in a ``finally`` + around ``server.run``). + + Concurrency invariant: the memory-side services share ``memory_conn``. + Each service's docstring documents that it owns the connection and the + caller must not share it with another service that has an open + transaction. That contract holds here only because the MCP stdio + transport serialises requests — ``_call_tool`` runs one tool + invocation at a time, so at most one service's SAVEPOINT is ever in + flight. If any future change introduces async fan-out, a background + task, or a worker thread that writes to ``memory_conn``, this + assumption breaks and the services must be reworked to use per-task + connections (or a connection pool with explicit checkout). """ config = get_config() @@ -985,21 +297,6 @@ def create_server() -> tuple[ # memory.retrieve will succeed on their next call without a restart. _probe_ollama(config.ollama_host) - # Concurrency invariant: the five memory-side services below - # (EpisodeService, ObservationService, ReflectionSynthesisService, - # RetentionService, SpoolService) all share ``memory_conn``. Each - # service's docstring documents that it owns the connection and the - # caller must not share it with another service that has an open - # transaction. That contract holds here only because the MCP stdio - # transport serialises requests — _call_tool runs one tool invocation - # at a time, so at most one service's SAVEPOINT is ever in flight. - # If any future change introduces async fan-out, a background task, - # or a worker thread that writes to ``memory_conn``, this assumption - # breaks and the services must be reworked to use per-task - # connections (or a connection pool with explicit checkout). - episodes = EpisodeService(memory_conn) - observations = ObservationService(memory_conn, embedder=embedder, episodes=episodes) - # Resolve project + session id ONCE at startup. Per-handler code can # still override per-call (handlers continue to read args.get("project") # / call project_name() for backwards compatibility), but the backend @@ -1010,595 +307,58 @@ def create_server() -> tuple[ # env-var / marker file path when ``session_id is None``. An empty # string silently writes observations with session_id='' and breaks # the rating tools. - startup_project = project_name() - startup_session_id: str | None = _resolve_session_id(config.home) or None - backend: StorageBackend = build_backend( - config=config, - memory_conn=memory_conn, - embedder=embedder, - session_id=startup_session_id, - project=startup_project, + services = _build_services( + config, memory_conn, knowledge_conn, embedder, + startup_project=project_name(), + startup_session_id=resolve_session_id(config.home) or None, ) - # Reflection synthesis is driven by the IDE-LLM via two MCP tools - # (memory.synthesize_next_get_context / _apply). The service no - # longer holds a chat client. - reflections = ReflectionSynthesisService(memory_conn) - retention = RetentionService(conn=memory_conn) - memory_rating = MemoryRatingService(memory_conn) - - knowledge = KnowledgeService( - knowledge_conn, - knowledge_base=config.knowledge_base, - ) - spool = SpoolService(memory_conn, config.spool_dir, episodes=episodes) - # Session-start behaviour: reindex knowledge at startup. mtime-only, so # the cost is O(files) stat calls on an already-indexed corpus. We # swallow any exception so a missing / unreachable knowledge base does # not block the server from serving memory tools. if config.knowledge_base.is_dir(): try: - knowledge.reindex() + services.knowledge.reindex() except Exception: # noqa: BLE001 — best-effort startup hook pass + dispatcher = ToolDispatcher(services, all_handlers()) + server: Server = Server(name="better-memory") @server.list_tools() async def _list_tools() -> list[Tool]: - return _tool_definitions(supports_synthesis=backend.supports_synthesis) + return dispatcher.tool_definitions() @server.call_tool() async def _call_tool( - name: str, arguments: dict[str, Any] | None + name: str, arguments: dict[str, Any] | None, ) -> list[TextContent]: - args = arguments or {} - - if name == "memory.observe": - with _diag.trace( - "mcp.memory.observe", - content_len=len(args.get("content") or ""), - scope=args.get("scope") or "project", - component=args.get("component"), - ): - _diag.step("mcp.memory.observe", "calling_observations_create") - obs_id = await observations.create( - content=args["content"], - component=args.get("component"), - theme=args.get("theme"), - trigger_type=args.get("trigger_type"), - outcome=args.get("outcome", "neutral"), - tech=args.get("tech"), - # `or "project"` (not `, "project"` default) defends against - # MCP clients sending {"scope": null} — dict.get returns the - # default only when the key is absent, not when its value is - # None. Without this, scope=None propagates to ObservationService - # .create() which raises ValueError. - scope=args.get("scope") or "project", - ) - _diag.step("mcp.memory.observe", "create_returned", obs_id=obs_id) - return [TextContent(type="text", text=json.dumps({"id": obs_id}))] - - if name == "memory.semantic_observe": - from better_memory.services.semantic import SemanticMemoryService - project = project_name() - svc = SemanticMemoryService(memory_conn) - # `args.get("scope") or "project"` (not `, "project"` default) defends - # against MCP clients sending {"scope": null} — dict.get returns the - # default only when the key is absent, not when its value is None. - # Same fix as PR #25's BugBot finding on memory.observe. - memory_id = svc.create( - content=args["content"], - project=project, - scope=args.get("scope") or "project", - ) - return [TextContent(type="text", text=json.dumps({"id": memory_id}))] - - if name == "memory.semantic_retrieve": - from better_memory.services.semantic import SemanticMemoryService - project = args.get("project") or project_name() - svc = SemanticMemoryService(memory_conn) - memories = svc.list_for_project(project=project) - payload = [ - { - "id": m.id, - "content": m.content, - "project": m.project, - "scope": m.scope, - "created_at": m.created_at, - "updated_at": m.updated_at, - } - for m in memories - ] - return [TextContent(type="text", text=json.dumps(payload))] - - if name == "memory.semantic_update": - from better_memory.services.semantic import SemanticMemoryService - svc = SemanticMemoryService(memory_conn) - svc.update_text(id=args["id"], content=args["content"]) - return [TextContent(type="text", text=json.dumps({"ok": True}))] - - if name == "memory.semantic_delete": - from better_memory.services.semantic import SemanticMemoryService - svc = SemanticMemoryService(memory_conn) - svc.delete(id=args["id"]) - return [TextContent(type="text", text=json.dumps({"ok": True}))] - - if name == "memory.retrieve": - with _diag.trace( - "mcp.memory.retrieve", - project=args.get("project") or "", - tech=args.get("tech"), - phase=args.get("phase"), - polarity=args.get("polarity"), - ): - diag_cid: str | None = None - t_retrieve = 0.0 - if _diag.enabled(): - diag_cid = uuid.uuid4().hex[:8] - t_retrieve = time.monotonic() - _diag.log(f"[bm-retrieve start cid={diag_cid}]") - - # 1. Drain spool — must happen before any retrieval so fresh - # hook events (session_start, commit_close) are processed. - # SpoolService.drain is idempotent. - _diag.step("mcp.memory.retrieve", "before_spool_drain") - _run_best_effort("spool.drain", spool.drain, diag_cid=diag_cid) - _diag.step("mcp.memory.retrieve", "after_spool_drain") - - # 2. Maybe run retention. Guard ensures at most once per 24h - # regardless of how often retrieve is called. Best-effort: - # a retention failure must NEVER block memory.retrieve. - def _retention_step() -> None: - cfg = get_config() - RetentionScheduler( - memory_conn, auto_prune=cfg.auto_prune - ).maybe_run(triggered_by="retrieve") - - _diag.step("mcp.memory.retrieve", "before_retention_scheduler") - _run_best_effort( - "retention scheduler", _retention_step, diag_cid=diag_cid - ) - _diag.step("mcp.memory.retrieve", "after_retention_scheduler") - - project = args.get("project") or project_name() - limit_per_bucket = args.get("limit_per_bucket", 20) - t_reflections = time.monotonic() if diag_cid else 0.0 - _diag.step( - "mcp.memory.retrieve", - "before_retrieve_reflections", - project=project, - ) - buckets = backend.retrieve( - project=project, - tech=args.get("tech"), - phase=args.get("phase"), - polarity=args.get("polarity"), - limit_per_bucket=limit_per_bucket, - ) - _diag.step( - "mcp.memory.retrieve", - "after_retrieve_reflections", - do=len(buckets.get("do", [])), - dont=len(buckets.get("dont", [])), - neutral=len(buckets.get("neutral", [])), - ) - if diag_cid is not None: - refl_ms = int((time.monotonic() - t_reflections) * 1000) - total_ms = int((time.monotonic() - t_retrieve) * 1000) - _diag.log( - f"[bm-retrieve step=reflections cid={diag_cid} " - f"ms={refl_ms} status=ok]" - ) - _diag.log( - f"[bm-retrieve done cid={diag_cid} total_ms={total_ms}]" - ) - return [TextContent(type="text", text=json.dumps(buckets))] - - if name == "memory.retrieve_observations": - project = args.get("project") or project_name() - results = await observations.list_observations( - project=project, - episode_id=args.get("episode_id"), - component=args.get("component"), - theme=args.get("theme"), - outcome=args.get("outcome"), - query=args.get("query"), - limit=args.get("limit", 50), - ) - return [TextContent(type="text", text=json.dumps(results))] - - if name == "memory.record_use": - observations.record_use( - args["id"], - outcome=args.get("outcome"), - ) - return [TextContent(type="text", text=json.dumps({"ok": True}))] - - if name == "knowledge.search": - results = knowledge.search( - args["query"], - project=args.get("project"), - ) - return [ - TextContent( - type="text", - text=json.dumps( - [_serialize_knowledge_search(r) for r in results] - ), - ) - ] - - if name == "knowledge.list": - docs = knowledge.list_documents(project=args.get("project")) - return [ - TextContent( - type="text", - text=json.dumps( - [_serialize_knowledge_doc(d) for d in docs] - ), - ) - ] - - if name == "memory.start_ui": - result = ui_launcher.start_ui() - return [ - TextContent(type="text", text=json.dumps(result)) - ] - - if name == "memory.start_episode": - project = project_name() - episode_id = episodes.start_foreground( - session_id=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 = reflections._read_queue_counts(project=project) - buckets = 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), - } - ), - ) - ] - - if name == "memory.close_episode": - 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 = episodes.close_active( - session_id=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} - ), - ) - ] - - if name == "memory.reconcile_episodes": - open_episodes = episodes.unclosed_episodes( - exclude_session_ids={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))] - - if name == "memory.run_retention": - report = retention.run( - retention_days=args.get("retention_days", 90), - prune=args.get("prune", False), - prune_age_days=args.get("prune_age_days", 365), - dry_run=args.get("dry_run", False), - ) - return [ - TextContent( - type="text", - text=json.dumps({ - "archived_via_retired_reflection": - report.archived_via_retired_reflection, - "archived_via_consumed_without_reflection": - report.archived_via_consumed_without_reflection, - "archived_via_no_outcome_episode": - report.archived_via_no_outcome_episode, - "pruned": report.pruned, - }), - ) - ] - - if name == "memory.list_episodes": - rows = 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))] - - if name == "memory.session_bootstrap": - from better_memory.services.session_bootstrap import ( - SessionBootstrapService, - ) - - cwd_arg = args.get("cwd") or os.getcwd() - session_id_arg = ( - args.get("session_id") - or os.environ.get("CLAUDE_SESSION_ID") - or os.environ.get("CLAUDE_CODE_SESSION_ID") - or uuid.uuid4().hex - ) - svc = SessionBootstrapService(memory_conn) - result = svc.bootstrap( - source=args.get("source"), - session_id=session_id_arg, - cwd=Path(cwd_arg), - ) - payload = { - "additionalContext": result.additional_context, - "project": result.project, - "source": result.source, - "episode": { - "id": result.episode_id, - "action": result.episode_action, - }, - "counts": { - "semantic": result.semantic_count, - "reflections": result.reflections_counts, - }, - } - return [TextContent(type="text", text=json.dumps(payload))] - - if name == "memory.synthesize_next_get_context": - project = args.get("project") or project_name() - with _audit_synth_call( - config.home, - tool="get_context", - project=project, - episode_id=None, - ) as audit: - ctx = reflections.get_next_pending_context(project=project) - queue = reflections._read_queue_counts(project=project) - payload = _serialize_synth_get_context(ctx, queue) - if ctx is None: - audit["result_kind"] = "empty" - else: - audit["result_kind"] = "episode" - audit["episode_id"] = ctx.episode.id - audit["obs_count"] = len(ctx.observations) - audit["refl_count"] = len(ctx.reflections) - return [TextContent(type="text", text=json.dumps(payload))] - - if name == "memory.synthesize_next_apply": - from better_memory.services.reflection import ( - SynthesisResponseError, - ) - project = args.get("project") or project_name() - episode_id = args["episode_id"] - decision = args["decision"] - with _audit_synth_call( - config.home, - tool="apply", - project=project, - episode_id=episode_id, - ) as audit: - try: - # ``decision`` is already a parsed dict (the MCP - # framework decoded it before dispatch). Use the - # dict-shape parser directly to skip a redundant - # json.dumps → json.loads round-trip. - response = reflections.parse_response_dict(decision) - except SynthesisResponseError as exc: - audit["result_kind"] = "validation_error" - audit["error"] = str(exc) - payload = _serialize_synth_apply_validation_error( - str(exc) - ) - return [ - TextContent(type="text", text=json.dumps(payload)) - ] - try: - step = reflections.apply_decision( - episode_id=episode_id, - response=response, - project=project, - ) - except ValueError as exc: - # Episode-state preconditions: not found / wrong - # project / already synthesized. Surface as - # structured error so the IDE-LLM can refetch - # context instead of retrying the same stale id. - audit["result_kind"] = "state_error" - audit["error"] = str(exc) - payload = _serialize_synth_apply_state_error(str(exc)) - return [ - TextContent(type="text", text=json.dumps(payload)) - ] - audit["result_kind"] = "applied" - audit["counts"] = step.counts - return [ - TextContent( - type="text", - text=json.dumps(_serialize_synth_apply_ok(step)), - ) - ] - - if name == "memory.list_session_exposures": - from better_memory.services.session_bootstrap import ( - SessionBootstrapService, - ) - - sid = _resolve_session_id(config.home) or "" - payload = SessionBootstrapService(memory_conn).list_session_exposures( - session_id=sid, - ) - return [TextContent(type="text", text=json.dumps(payload))] - - if name == "memory.apply_session_ratings": - sid = _resolve_session_id(config.home) - if not sid: - raise ValueError( - "No active session: CLAUDE_SESSION_ID / " - "CLAUDE_CODE_SESSION_ID not set and no session marker " - "found (SessionStart hook may not have run)" - ) - payload = memory_rating.apply_session_ratings( - session_id=sid, - ratings=args["ratings"], - ) - return [TextContent(type="text", text=json.dumps(payload))] - - if name == "memory.credit": - sid = _resolve_session_id(config.home) - if not sid: - payload = {"applied": None, "skipped": "no_session"} - else: - payload = memory_rating.credit_one( - session_id=sid, - kind=args["kind"], - id=args["id"], - classification=args["class"], - ) - return [TextContent(type="text", text=json.dumps(payload))] - - raise ValueError(f"Unknown tool: {name}") - - cleaned = False - - async def cleanup() -> None: - """Close SQLite connections and the embedder HTTP client. - - Idempotent: safe to call multiple times. SQLite ``Connection.close`` - is a no-op after the first call, and we guard the embedder close with - a local flag so we don't double-close its httpx client either. - """ - nonlocal cleaned - if cleaned: - return - cleaned = True - try: - memory_conn.close() - except Exception: # noqa: BLE001 — best-effort shutdown - pass - try: - knowledge_conn.close() - except Exception: # noqa: BLE001 — best-effort shutdown - pass - # In the sqlite backend no embedder is built (FTS5 triggers handle - # indexing). Guard against the None case so this cleanup stays - # idempotent across both backends. - if embedder is not None: - try: - await embedder.aclose() - except Exception: # noqa: BLE001 — best-effort shutdown - pass + with _diag.trace(f"mcp.{name}"): + return await dispatcher.call(name, arguments or {}) + cleanup = _make_cleanup(memory_conn, knowledge_conn, embedder) return server, cleanup, ServerContext( - backend=backend, + backend=services.backend, memory_conn=memory_conn, embedder=embedder, + dispatcher=dispatcher, ) async def _dispatch_for_tests(name: str, arguments: dict) -> list[TextContent]: - """Test-only entry point that runs one tool invocation against a fresh - server instance. NOT used by production code. + """Test-only entry point. Routes through ToolDispatcher. - The MCP SDK catches exceptions inside handlers and surfaces them as - CallToolResult(isError=True). To make tests ergonomic, this helper - re-raises any error as ValueError so callers can use - `pytest.raises(ValueError, match="...")` instead of inspecting - result text manually. + Builds a server (with its standard ServiceContainer + every registered + handler) and dispatches one tool call. Preserved as a shim so + ``tests/mcp/test_server_sqlite.py`` and ``tests/mcp/test_rating_tools.py`` + continue to work without modification. """ - from mcp.types import CallToolRequest, CallToolRequestParams - - server, cleanup, _ = create_server() + _server, cleanup, ctx = create_server() try: - handler = server.request_handlers[CallToolRequest] - req = CallToolRequest( - method="tools/call", - params=CallToolRequestParams(name=name, arguments=arguments), - ) - result = await handler(req) - # The SDK's ServerResult is a discriminated union; we know this - # handler is wired to CallTool and returns CallToolResult. - from mcp.types import CallToolResult - assert isinstance(result.root, CallToolResult), ( - f"Expected CallToolResult, got {type(result.root).__name__}" - ) - if getattr(result.root, "isError", False): - # Re-raise as ValueError; preserve the framework's error text. - text = "" - if result.root.content: - first = result.root.content[0] - text = getattr(first, "text", "") or "" - raise ValueError(text) - # Tests inspect .text on TextContent entries — runtime is correct; - # cast through Any to satisfy Pyright's list invariance (the SDK - # types .content as list[ContentBlock]; our tools only emit - # TextContent so the cast is sound). - return cast(list[TextContent], result.root.content) + assert ctx.dispatcher is not None + return await ctx.dispatcher.call(name, arguments) finally: await cleanup() diff --git a/better_memory/services/reflection.py b/better_memory/services/reflection.py index b26cda6..c5071f7 100644 --- a/better_memory/services/reflection.py +++ b/better_memory/services/reflection.py @@ -606,8 +606,8 @@ def _mark_synthesized(self, episode_id: str) -> None: (self._clock().isoformat(), episode_id), ) - # ------------------------------------------------------- _read_queue_counts - def _read_queue_counts(self, project: str) -> EpisodeQueueCounts: + # -------------------------------------------------------- read_queue_counts + def read_queue_counts(self, project: str) -> EpisodeQueueCounts: """Snapshot the per-project queue state in one statement. Single-source-of-truth for the route's banner — eliminates the race @@ -1167,7 +1167,7 @@ def apply_decision( return SynthesisStep( processed=True, episode_id=episode_id, counts=counts, - queue=self._read_queue_counts(project), + queue=self.read_queue_counts(project), failure=None, ) diff --git a/docs/superpowers/plans/2026-06-13-mcp-server-dispatcher-refactor.md b/docs/superpowers/plans/2026-06-13-mcp-server-dispatcher-refactor.md new file mode 100644 index 0000000..b8573ab --- /dev/null +++ b/docs/superpowers/plans/2026-06-13-mcp-server-dispatcher-refactor.md @@ -0,0 +1,1860 @@ +# MCP server dispatcher refactor — Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Decompose `better_memory/mcp/server.py` into a `ToolDispatcher` + per-domain handler modules + a `ServiceContainer`, without changing the MCP wire surface or any tool's behaviour. + +**Architecture:** Replace the 470-LOC `_call_tool` if-chain with `dispatcher.call(name, args)`. Every long-lived service is built once in `_build_services()` and bundled into a frozen `ServiceContainer`. Each domain owns one module under `mcp/handlers/` carrying its tool schemas as module constants plus a `*Handlers` class with one method per tool and a `handlers()` registration list. + +**Tech Stack:** Python 3.12, `mcp` SDK (>=1), pytest, sqlite3, dataclasses, asyncio. + +**Spec:** `docs/superpowers/specs/2026-06-12-mcp-server-dispatcher-design.md` (all 17 assumptions independently verified). + +**Prerequisite (separate PBI):** `mcp-server-tool-error-path-tests` should land first so each domain we migrate has error-path coverage to fall back on. This plan assumes that PBI is complete or in-flight on a parallel track. + +--- + +## Guardrails + +Surfaced from `mcp__better-memory__memory_retrieve` (phase=planning, phase=implementation) and `~/.better-memory/knowledge-base/standards/ralph-runtime.md` BEFORE drafting tasks. Cross-reference by slug from individual tasks. + +| # | Guardrail | Source | Confidence | +|---|-----------|--------|------------| +| G1 | **Per-task confidence + sub-90% mitigations.** Every task gets a percentage; sub-90% tasks embed mitigations in the task body, not as optional follow-ups. Treated as a non-skippable gate before plan presentation. | `~/.better-memory/knowledge-base/standards/ralph-runtime.md` § "Apply confidence scoring to implementation plans" | standard | +| G2 | **Verify-before-commit on internal patterns.** When a task says "follows existing pattern X", read X's source at plan-write time, not implementation time. Smoke-test verbatim imports/calls in plan code blocks against the actual source. | `[[022dafc32d484f36ad4b03dc8bf607c3]]` plan-prescribed-test-snippets · `[[be7ad6bf15d64763a3e6041388de06ac]]` | 0.85 · 0.9 | +| G3 | **Multi-artefact step: cross-check each item against the diff before marking complete.** A "create X, Y, Z" step that lands a subset is the recurring miss. Per-task acceptance criteria must list the artefacts. | `[[0518718c9fd94093962a4c5662368c1b]]` plan-step-multi-artefact | 0.7 (evidence 1, useful 4) | +| G4 | **Test fixtures + signature changes: grep all `ClassName(` call sites in `tests/`.** Adding required fields to dataclasses or changing signatures silently breaks fixture-constructed tests. Plan must enumerate test call sites alongside production. | `[[462d13a79f84482796fe692535aea580]]` ExecutorConfig-signature-test-sweep | 0.85 (evidence 4) | +| G5 | **Docs + module-docstring drift on tool changes.** Adding/removing MCP tools requires sweeping `README.md`, `website/mcp-tools.md`, `website/index.md`, `website/architecture.md`, `better_memory/mcp/server.py` module docstring. Stale tool counts drifted twice in one month. | `[[98056ebc80464d34a0c2e95b18882e41]]` keep-website-readme-in-sync | 0.95 (evidence 7, useful 11) | +| G6 | **Stage by explicit paths, never `git add -A` (ralph co-running).** Operator's queue-zenbook clone may have ralph-authored untracked files. Per-commit `git add ` only. | `[[0547374e598d4bf7b78da1dd80c62da4]]` git-add-A-while-ralph-co-runs | 0.7 | +| G7 | **Render the plan in the visualiser AFTER commit, BEFORE execution-choice.** Sequence: (1) commit plan, (2) render plan-summary in visualiser, (3) announce URL, (4) present execution choice. Each step row gets a confidence badge; sub-95% step rows get inline `step-note` annotations. | `~/.better-memory/knowledge-base/standards/ralph-runtime.md` § "Render brainstorming/plans in visual companion" | standard | + +**Dismissed (logged for transparency):** +- `[[a5af973e32de41feb5d544204d10b5c2]]` Route state transitions through git-aware helper — N/A: refactor is a Python module split, not a queue-state transition. +- `[[8d05b8c2959943268baac34161001e81]]` Queue artefacts must be self-contained — N/A: this is a code refactor, not a queue PBI. +- `[[cfa7e7f86366403eaa459c2b9eec59ad]]` PBI directory shape type-aware — N/A: not authoring a PBI. +- `[[d342c4818c7f420da2fcc964d4387901]]` Test fixture .gitignore mirror — N/A: no fixture serves multiple clone roles here. + +--- + +## Per-task confidence summary + +| Task | Confidence | Notes | +|------|-----------|-------| +| 1. ServiceContainer dataclass | **98%** | Mechanical. Frozen dataclass, simple shape. | +| 2. Handler + ToolDispatcher | **96%** | Mechanical. 6 unit tests cover the contract. | +| 3. Move `_resolve_session_id` | **97%** | Pure move. 3 call sites already grepped. | +| 4. Move `_audit_synth_call` | **95%** | Verbatim move + 1 import-path update in tests. | +| 5. Promote `_read_queue_counts` | **96%** | Pure rename. 2 callers in server.py, possibly tests (grep step pinned). | +| 6. Extract `_build_services` | **92%** (was 88%) | Lifted: all 5 ctor sigs verified — `SessionBootstrapService(conn)`, `MemoryRatingService(conn, *, clock=None)`, `KnowledgeService(conn, *, knowledge_base=None)`, `SpoolService(conn, spool_dir=None, *, episodes=None)`, `SemanticMemoryService(conn)`. Fixtures `tmp_memory_db` + `tmp_knowledge_base` (`tests/conftest.py:65,77`) compose cleanly. `_MEMORY_MIGRATIONS` + `_KNOWLEDGE_MIGRATIONS` constants importable from server.py (lines 87-88). | +| 7. ObservationHandlers (4 tools) | **92%** (was 80%) | Lifted: schemas `_OBSERVE_SCHEMA` (server.py:270-299), `_RETRIEVE_SCHEMA` (377-400), `_RETRIEVE_OBSERVATIONS_SCHEMA` (403-427), `_RECORD_USE_SCHEMA` (430-446) all verbatim-copy targets. `observations.create` / `list_observations` async, `record_use` sync (`services/observation.py:160,479,417`). `_run_best_effort` sig + import path confirmed (server.py:95-128). Spool→retention→backend ordering pinned by comment at 1143-1145. | +| 8. SemanticHandlers (4 tools) | **94%** | Small bodies. SemanticMemoryService instantiation moves to container. | +| 9. EpisodeHandlers (4 tools) | **92%** (was 85%) | Lifted: `close_episode` payload shapes pinned — try `{"closed_episode_id": , "already_closed": False}`, except `{"closed_episode_id": None, "already_closed": True}`. `list_episodes` confirmed 10-field serialiser (episode_id, project, tech, goal, started_at, hardened_at, ended_at, close_reason, outcome, summary). All 4 EpisodeService methods sync. Task 5 promotion gate added explicitly. | +| 10. ReflectionHandlers (2 tools) | **92%** (was 82%) | Lifted: correct method name is `parse_response_dict` (not `parse_decision_json`). Audit state contract enumerated: 5 result_kinds (`empty`, `episode`, `validation_error`, `state_error`, `applied`). 3 return paths in apply: SynthesisResponseError → validation_error, ValueError → state_error, happy → applied. `SynthesisResponseError(ValueError)` defined at `services/reflection.py:131`. | +| 11. RetentionHandlers (1 tool) | **94%** | Single tool, 8-field dict construction. | +| 12. KnowledgeHandlers (2 tools) | **93%** | 2 tools + 2 small serializer helpers to copy. | +| 13. RatingHandlers (3 tools) | **92%** (was 87%) | Lifted: 3-line ValueError pinned verbatim ("No active session: CLAUDE_SESSION_ID / CLAUDE_CODE_SESSION_ID not set and no session marker found (SessionStart hook may not have run)"). `credit` shapes pinned: no-session `{"applied": None, "skipped": "no_session"}`; with-session `{"applied": str\|None, "skipped": str\|None}` from ApplyOutcome TypedDict. 3 distinct resolve_session_id None-behaviours (coerce/raise/skip-dict). | +| 14. SessionHandlers (2 tools) | **92%** (was 86%) | Lifted: full session_bootstrap body verbatim from server.py:1380-1411. `SessionBootstrapService.bootstrap(*, source, session_id, cwd, project)` returns `BootstrapResult` (additional_context, project, source, episode_id, episode_action, semantic_count, reflections_counts). `ui_launcher.start_ui(*, spawn_timeout, confirm_retry_sleep) -> dict`. CWD uses `os.getcwd()` (not `Path.cwd()`). | +| 15. Wire dispatcher into `create_server` + delete legacy | **92%** (was 88%) | Lifted: ServerContext fields confirmed (`backend`, `memory_conn`, `embedder`) — add `dispatcher: ToolDispatcher \| None = None` LAST. `OllamaEmbedder.aclose()` is async. Existing cleanup at 1529-1555 is idempotent via `cleaned` flag. Zero positional `ServerContext(...)` in tests. README/website both say "22 tools" — count unchanged. Module docstring lists 12 bullets covering 19 of 22 (rating tools missing); fix in same PR. | +| 16. Slim `_dispatch_for_tests` | **92%** | Small, well-defined. Existing test contract unchanged. Source line range 1564-1603. | +| 17. Full smoke + lint | **99%** | Verification only. 37 test files across tests/mcp/, tests/services/, tests/storage/. | + +**Average: ~94%.** Original sub-90% tasks (6, 7, 9, 10, 13, 14, 15) all lifted to ≥92% via concrete source-verified mitigations. Detailed evidence in each task's mitigation block. + +--- + +### Task 1: Create `ServiceContainer` dataclass (98%) + +**Files:** +- Create: `better_memory/mcp/container.py` +- Test: `tests/mcp/test_service_container.py` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/mcp/test_service_container.py +"""ServiceContainer bundles every long-lived service the MCP dispatcher needs.""" +from __future__ import annotations + +import sqlite3 +from dataclasses import is_dataclass +from unittest.mock import MagicMock + +from better_memory.mcp.container import ServiceContainer + + +def test_service_container_is_frozen_dataclass() -> None: + fields = { + "config", "memory_conn", "backend", + "episodes", "observations", "reflections", + "retention", "memory_rating", "knowledge", + "spool", "semantic", "session_bootstrap", + } + assert is_dataclass(ServiceContainer) + assert set(ServiceContainer.__dataclass_fields__) == fields + + +def test_service_container_holds_attributes() -> None: + mock_conn = MagicMock(spec=sqlite3.Connection) + container = ServiceContainer( + config=MagicMock(), + memory_conn=mock_conn, + backend=MagicMock(), + episodes=MagicMock(), + observations=MagicMock(), + reflections=MagicMock(), + retention=MagicMock(), + memory_rating=MagicMock(), + knowledge=MagicMock(), + spool=MagicMock(), + semantic=MagicMock(), + session_bootstrap=MagicMock(), + ) + assert container.memory_conn is mock_conn +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/mcp/test_service_container.py -v` +Expected: FAIL with `ModuleNotFoundError: No module named 'better_memory.mcp.container'` + +- [ ] **Step 3: Write minimal implementation** + +```python +# better_memory/mcp/container.py +"""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, Any + +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.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: Any # SemanticMemoryService — typed Any to defer import + session_bootstrap: "SessionBootstrapService" +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/mcp/test_service_container.py -v` +Expected: PASS, 2 tests green. + +- [ ] **Step 5: Commit** + +```bash +git add better_memory/mcp/container.py tests/mcp/test_service_container.py +git commit -m "feat(mcp): introduce ServiceContainer frozen dataclass + +Bundles every long-lived service the dispatcher needs so they're +built once at startup, not per-call (today: SemanticMemoryService +4x, SessionBootstrapService 2x)." +``` + +--- + +### Task 2: Create `Handler` dataclass + `ToolDispatcher` (96%) + +**Files:** +- Create: `better_memory/mcp/dispatcher.py` +- Test: `tests/mcp/test_tool_dispatcher.py` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/mcp/test_tool_dispatcher.py +"""ToolDispatcher: register, list, call, capability-gate.""" +from __future__ import annotations + +import pytest +from unittest.mock import MagicMock + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler, ToolDispatcher + + +def _container(*, supports_synthesis: bool) -> ServiceContainer: + backend = MagicMock() + backend.supports_synthesis = supports_synthesis + return ServiceContainer( + config=MagicMock(), memory_conn=MagicMock(), backend=backend, + episodes=MagicMock(), observations=MagicMock(), reflections=MagicMock(), + retention=MagicMock(), memory_rating=MagicMock(), knowledge=MagicMock(), + spool=MagicMock(), semantic=MagicMock(), session_bootstrap=MagicMock(), + ) + + +async def _stub_call(services, args): + return [] + + +def test_handler_is_frozen_dataclass() -> None: + h = Handler(name="x", schema={"type": "object"}, call=_stub_call) + assert h.name == "x" + assert h.requires_synthesis is False + with pytest.raises(Exception): + h.name = "y" # type: ignore[misc] + + +def test_dispatcher_lists_all_when_synthesis_supported() -> None: + services = _container(supports_synthesis=True) + handlers = [ + Handler("a", {}, _stub_call), + Handler("b", {}, _stub_call, requires_synthesis=True), + ] + dispatcher = ToolDispatcher(services, handlers) + names = [t.name for t in dispatcher.tool_definitions()] + assert names == ["a", "b"] + + +def test_dispatcher_hides_synthesis_tools_when_unsupported() -> None: + services = _container(supports_synthesis=False) + handlers = [ + Handler("a", {}, _stub_call), + Handler("b", {}, _stub_call, requires_synthesis=True), + ] + dispatcher = ToolDispatcher(services, handlers) + names = [t.name for t in dispatcher.tool_definitions()] + assert names == ["a"] + + +@pytest.mark.asyncio +async def test_dispatcher_call_unknown_raises_value_error() -> None: + services = _container(supports_synthesis=True) + dispatcher = ToolDispatcher(services, []) + with pytest.raises(ValueError, match="Unknown tool: nope"): + await dispatcher.call("nope", {}) + + +@pytest.mark.asyncio +async def test_dispatcher_call_gated_tool_when_unsupported_raises_unknown() -> None: + services = _container(supports_synthesis=False) + handler = Handler("b", {}, _stub_call, requires_synthesis=True) + dispatcher = ToolDispatcher(services, [handler]) + with pytest.raises(ValueError, match="Unknown tool: b"): + await dispatcher.call("b", {}) + + +@pytest.mark.asyncio +async def test_dispatcher_call_routes_to_handler() -> None: + services = _container(supports_synthesis=True) + seen = {} + async def recording_call(svc, args): + seen["svc"] = svc + seen["args"] = args + return [MagicMock()] + handler = Handler("x", {"type": "object"}, recording_call) + dispatcher = ToolDispatcher(services, [handler]) + result = await dispatcher.call("x", {"k": 1}) + assert seen["svc"] is services + assert seen["args"] == {"k": 1} + assert len(result) == 1 +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/mcp/test_tool_dispatcher.py -v` +Expected: FAIL with `ModuleNotFoundError: No module named 'better_memory.mcp.dispatcher'` + +- [ ] **Step 3: Write minimal implementation** + +```python +# better_memory/mcp/dispatcher.py +"""ToolDispatcher: O(1) name → 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, field +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) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/mcp/test_tool_dispatcher.py -v` +Expected: PASS, 6 tests green. + +- [ ] **Step 5: Commit** + +```bash +git add better_memory/mcp/dispatcher.py tests/mcp/test_tool_dispatcher.py +git commit -m "feat(mcp): introduce ToolDispatcher + Handler dataclass + +Replaces the 470-LOC if-chain with O(1) name lookup. Capability gate +preserves the existing 'Unknown tool: …' ValueError shape so clients +don't notice." +``` + +--- + +### Task 3: Move `_resolve_session_id` to `mcp/_session.py` (97%) + +**Files:** +- Create: `better_memory/mcp/_session.py` +- Modify: `better_memory/mcp/server.py:147-160` (delete `_resolve_session_id`) +- Test: existing tests cover behaviour; add a smoke test. +- New test: `tests/mcp/test_session_resolver.py` + +- [ ] **Step 1: Write the failing test** + +```python +# tests/mcp/test_session_resolver.py +"""resolve_session_id resolves Claude Code session id with the documented fallback.""" +from __future__ import annotations + +from pathlib import Path +from unittest.mock import patch + +from better_memory.mcp._session import resolve_session_id + + +def test_resolve_session_id_prefers_env(monkeypatch, tmp_path: Path) -> None: + monkeypatch.setenv("CLAUDE_SESSION_ID", "from-env") + assert resolve_session_id(tmp_path) == "from-env" + + +def test_resolve_session_id_falls_back_to_alt_env(monkeypatch, tmp_path: Path) -> None: + monkeypatch.delenv("CLAUDE_SESSION_ID", raising=False) + monkeypatch.setenv("CLAUDE_CODE_SESSION_ID", "alt-env") + assert resolve_session_id(tmp_path) == "alt-env" + + +def test_resolve_session_id_falls_back_to_marker(monkeypatch, tmp_path: Path) -> None: + monkeypatch.delenv("CLAUDE_SESSION_ID", raising=False) + monkeypatch.delenv("CLAUDE_CODE_SESSION_ID", raising=False) + with patch("better_memory.mcp._session.read_session_id", return_value="marker"): + assert resolve_session_id(tmp_path) == "marker" + + +def test_resolve_session_id_returns_none_when_all_absent( + monkeypatch, tmp_path: Path, +) -> None: + monkeypatch.delenv("CLAUDE_SESSION_ID", raising=False) + monkeypatch.delenv("CLAUDE_CODE_SESSION_ID", raising=False) + with patch("better_memory.mcp._session.read_session_id", return_value=None): + assert resolve_session_id(tmp_path) is None +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/mcp/test_session_resolver.py -v` +Expected: FAIL with `ModuleNotFoundError: No module named 'better_memory.mcp._session'` + +- [ ] **Step 3: Write minimal implementation** + +```python +# better_memory/mcp/_session.py +"""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) + ) +``` + +- [ ] **Step 4: Replace the old function inside `server.py`** + +In `better_memory/mcp/server.py`, delete lines 147-160 (the `_resolve_session_id` function) and add an import at the top: + +```python +# Add to imports near line 64 (with the other runtime/session imports): +from better_memory.mcp._session import resolve_session_id +``` + +Then at lines 1014, 1498, 1517 (today's three callers in `server.py`) replace `_resolve_session_id(...)` with `resolve_session_id(...)`. Grep the file to verify all 3 callers are updated: + +```bash +grep -n "_resolve_session_id\b" better_memory/mcp/server.py +# Expected: empty after the edits +grep -n "resolve_session_id\b" better_memory/mcp/server.py +# Expected: 1 import line + 3 callers +``` + +- [ ] **Step 5: Run full MCP test suite to verify nothing broke** + +Run: `uv run pytest tests/mcp/ -v` +Expected: ALL PASS (the new `_session.py` tests + every existing MCP test). + +- [ ] **Step 6: Commit** + +```bash +git add better_memory/mcp/_session.py better_memory/mcp/server.py tests/mcp/test_session_resolver.py +git commit -m "refactor(mcp): move _resolve_session_id to mcp/_session.py + +No behaviour change. Lets handler modules import it without a +circular dep on the server module." +``` + +--- + +### Task 4: Move `_audit_synth_call` + `_append_synth_audit` to `mcp/handlers/_audit.py` (95%) + +**Files:** +- Create: `better_memory/mcp/handlers/__init__.py` (empty for now) +- Create: `better_memory/mcp/handlers/_audit.py` +- Modify: `better_memory/mcp/server.py` (delete old definitions + update imports) +- Modify: `tests/mcp/test_synth_audit_log.py` (update import path) + +- [ ] **Step 1: Create the empty package** + +```bash +mkdir -p better_memory/mcp/handlers +``` + +Write `better_memory/mcp/handlers/__init__.py`: + +```python +"""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 + + +def all_handlers() -> list[Handler]: + """Return the union of every domain's registered handlers. + + Filled in as handler modules land (Tasks 7-14). + """ + return [] +``` + +- [ ] **Step 2: Write the failing test** + +```python +# Add to the top of tests/mcp/test_synth_audit_log.py: +# Replace the existing import line +# from better_memory.mcp.server import _audit_synth_call +# with: +from better_memory.mcp.handlers._audit import _audit_synth_call +``` + +Run: `uv run pytest tests/mcp/test_synth_audit_log.py -v` +Expected: FAIL with `ImportError: cannot import name '_audit_synth_call'` + +- [ ] **Step 3: Create the new module** + +```python +# better_memory/mcp/handlers/_audit.py +"""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. +""" +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) +``` + +- [ ] **Step 4: Update `server.py` to re-import from the new location** + +In `better_memory/mcp/server.py`: +1. Delete the old `_append_synth_audit` (lines 163-173) and `_audit_synth_call` (lines 176-221). +2. Add at the top with the other imports: + +```python +from better_memory.mcp.handlers._audit import _append_synth_audit, _audit_synth_call +``` + +The existing call sites at lines 1413 + 1433 (synthesize tool branches) continue to work because the names resolve identically. + +- [ ] **Step 5: Run the audit log test plus full MCP suite** + +Run: `uv run pytest tests/mcp/ -v` +Expected: ALL PASS. + +- [ ] **Step 6: Commit** + +```bash +git add better_memory/mcp/handlers/__init__.py better_memory/mcp/handlers/_audit.py \ + better_memory/mcp/server.py tests/mcp/test_synth_audit_log.py +git commit -m "refactor(mcp): move audit helpers to mcp/handlers/_audit.py + +Audit JSONL byte-shape unchanged. Test updates import path only." +``` + +--- + +### Task 5: Promote `ReflectionSynthesisService._read_queue_counts` to public (96%) + +**Files:** +- Modify: `better_memory/services/reflection.py` (rename private → public) +- Modify: `better_memory/mcp/server.py:1263, 1422` (update call sites) +- Test: `tests/services/test_reflection.py` (only if it references the old name) + +- [ ] **Step 1: Locate the existing definition and call sites** + +```bash +grep -n "_read_queue_counts\|read_queue_counts" better_memory/services/reflection.py better_memory/mcp/server.py tests/ +``` + +Expected: 1 definition + 2 callers in `mcp/server.py`, possibly test references. + +- [ ] **Step 2: Rename the method in `services/reflection.py`** + +Change the method `def _read_queue_counts(self, ...)` to `def read_queue_counts(self, ...)`. Update the existing docstring's first line to drop the "private" framing if present. + +- [ ] **Step 3: Update the two MCP callers** + +In `better_memory/mcp/server.py` lines 1263 and 1422, replace `reflections._read_queue_counts(...)` with `reflections.read_queue_counts(...)`. + +- [ ] **Step 4: Update any test references** + +```bash +grep -rn "_read_queue_counts" tests/ +``` + +For each hit, change `_read_queue_counts` → `read_queue_counts`. + +- [ ] **Step 5: Run reflection + MCP tests** + +Run: `uv run pytest tests/services/test_reflection.py tests/mcp/ -v` +Expected: ALL PASS. + +- [ ] **Step 6: Commit** + +```bash +git add better_memory/services/reflection.py better_memory/mcp/server.py tests/ +git commit -m "refactor(reflection): promote _read_queue_counts to public read_queue_counts + +Two MCP dispatcher branches were the only callers; making this public +removes a leak the handler-module split would otherwise expose." +``` + +--- + +### Task 6: `_build_services` helper + service-container construction test (88%) + +**Confidence: 88% — mitigations:** + +- **Risk: test setup may not work with `get_config()` against tmp_path.** Mitigation: before drafting the test, search `tests/conftest.py` and `tests/mcp/conftest.py` for an existing `config` / `tmp_config` / `memory_db` fixture and PREFER that over rolling tmp_path manually. If no fixture exists, the Step 3 test code below is correct as written. +- **Risk: `SessionBootstrapService(memory_conn)` constructor sig might take more args.** Mitigation: grep `class SessionBootstrapService` and `def __init__` in `better_memory/services/session_bootstrap.py` BEFORE writing `_build_services` (Step 2). If the constructor takes additional positional/keyword args, thread them through `_build_services` parameters. **G4 applies** — also grep `SessionBootstrapService(` under `tests/` to catch fixture call sites that would break if the signature changes. +- **Risk: `ServiceContainer` field order in Step 2 must match Task 1's dataclass exactly.** Mitigation: copy field names verbatim from Task 1 Step 3. + +**Files:** +- Modify: `better_memory/mcp/server.py` (extract `_build_services` from `create_server`) +- Test: `tests/mcp/test_service_container.py` (add the "exactly once" assertion) + +- [ ] **Step 1: Write the failing test** + +Append to `tests/mcp/test_service_container.py`: + +```python +def test_build_services_constructs_each_service_exactly_once( + monkeypatch, tmp_path, +) -> None: + """Regression guard: SemanticMemoryService was built 4× per call, + SessionBootstrapService 2× per call. Container must build each once.""" + from collections import Counter + from better_memory.mcp.server import _build_services + from better_memory.config import get_config + + counts: Counter[str] = Counter() + + def _wrap(cls: type) -> type: + original_init = cls.__init__ + def _init(self, *a, **kw): + counts[cls.__name__] += 1 + original_init(self, *a, **kw) + cls.__init__ = _init # type: ignore[method-assign] + return cls + + import better_memory.services.semantic as _sem_mod + import better_memory.services.session_bootstrap as _sb_mod + _wrap(_sem_mod.SemanticMemoryService) + _wrap(_sb_mod.SessionBootstrapService) + + cfg = get_config() + container = _build_services(cfg, ...) # see Step 2 for arg shape + assert counts["SemanticMemoryService"] == 1 + assert counts["SessionBootstrapService"] == 1 + assert container is not None +``` + +**Note:** The test signature in Step 2 will pin the `_build_services` argument shape; until then the test will fail to import — that's intentional. + +- [ ] **Step 2: Extract `_build_services` from `create_server`** + +In `better_memory/mcp/server.py`, add a top-level function above `create_server`: + +```python +def _build_services( + config: Any, # better_memory.config.Config + memory_conn: sqlite3.Connection, + knowledge_conn: sqlite3.Connection, + embedder: OllamaEmbedder | None, + *, + startup_project: str, + startup_session_id: str | None, +) -> ServiceContainer: + """Construct every service exactly once. + + Replaces the inline 4x ``SemanticMemoryService(memory_conn)`` and the + inline 2x ``SessionBootstrapService(...)`` smells in the legacy + ``_call_tool`` body. + """ + from better_memory.services.semantic import SemanticMemoryService + from better_memory.services.session_bootstrap import SessionBootstrapService + + episodes = EpisodeService(memory_conn) + observations = ObservationService( + memory_conn, embedder=embedder, episodes=episodes, + ) + backend = build_backend( + config=config, + memory_conn=memory_conn, + embedder=embedder, + session_id=startup_session_id, + project=startup_project, + ) + reflections = ReflectionSynthesisService(memory_conn) + retention = RetentionService(conn=memory_conn) + memory_rating = MemoryRatingService(memory_conn) + knowledge = KnowledgeService( + knowledge_conn, knowledge_base=config.knowledge_base, + ) + spool = SpoolService(memory_conn, config.spool_dir, episodes=episodes) + semantic = SemanticMemoryService(memory_conn) + session_bootstrap = SessionBootstrapService(memory_conn) + + return ServiceContainer( + config=config, + memory_conn=memory_conn, + backend=backend, + episodes=episodes, + observations=observations, + reflections=reflections, + retention=retention, + memory_rating=memory_rating, + knowledge=knowledge, + spool=spool, + semantic=semantic, + session_bootstrap=session_bootstrap, + ) +``` + +Add `from better_memory.mcp.container import ServiceContainer` to the imports. + +Inside `create_server`, replace the inline construction block (today's lines 1000-1034) with: + +```python +services = _build_services( + config, memory_conn, knowledge_conn, embedder, + startup_project=project_name(), + startup_session_id=resolve_session_id(config.home) or None, +) +# Existing local names that the legacy if-chain still uses: +episodes = services.episodes +observations = services.observations +backend = services.backend +reflections = services.reflections +retention = services.retention +memory_rating = services.memory_rating +knowledge = services.knowledge +spool = services.spool +``` + +The legacy `_call_tool` if-chain continues to work because the local names are preserved. The container is now the source of truth. + +- [ ] **Step 3: Update the Step 1 test to call `_build_services` with real args** + +Edit the test to construct a real config + connections from a tmp_path-backed sqlite: + +```python +def test_build_services_constructs_each_service_exactly_once( + monkeypatch, tmp_path, +) -> None: + from collections import Counter + from better_memory.mcp.server import _build_services + + counts: Counter[str] = Counter() + + def _wrap(cls: type) -> type: + original_init = cls.__init__ + def _init(self, *a, **kw): + counts[cls.__name__] += 1 + original_init(self, *a, **kw) + cls.__init__ = _init # type: ignore[method-assign] + return cls + + import better_memory.services.semantic as _sem_mod + import better_memory.services.session_bootstrap as _sb_mod + _wrap(_sem_mod.SemanticMemoryService) + _wrap(_sb_mod.SessionBootstrapService) + + # Minimal in-memory config — adapt to your test conftest helpers if + # your repo has a `build_test_config` helper. + from better_memory.config import get_config + monkeypatch.setenv("BETTER_MEMORY_HOME", str(tmp_path)) + cfg = get_config() + + from better_memory.db.connection import connect + from better_memory.db.schema import apply_migrations + from pathlib import Path + mem_conn = connect(cfg.memory_db) + apply_migrations( + mem_conn, + migrations_dir=Path(__file__).parent.parent.parent + / "better_memory" / "db" / "migrations", + ) + kb_conn = connect(cfg.knowledge_db) + apply_migrations( + kb_conn, + migrations_dir=Path(__file__).parent.parent.parent + / "better_memory" / "db" / "knowledge_migrations", + ) + + container = _build_services( + cfg, mem_conn, kb_conn, embedder=None, + startup_project="test", startup_session_id=None, + ) + assert counts["SemanticMemoryService"] == 1 + assert counts["SessionBootstrapService"] == 1 + assert container.semantic is not None + assert container.session_bootstrap is not None +``` + +- [ ] **Step 4: Run test + full MCP suite** + +Run: `uv run pytest tests/mcp/test_service_container.py -v && uv run pytest tests/mcp/ -v` +Expected: ALL PASS. + +- [ ] **Step 5: Commit** + +```bash +git add better_memory/mcp/server.py tests/mcp/test_service_container.py +git commit -m "refactor(mcp): extract _build_services + thread ServiceContainer through create_server + +Legacy _call_tool if-chain still drives the actual dispatch; container +is now the canonical source for every long-lived service." +``` + +--- + +### Task 7: `handlers/observations.py` — observation domain (4 tools) (80%) + +**Confidence: 80% — mitigations:** + +- **Risk: `memory.retrieve` body is 70 LOC mixing spool drain, retention scheduler, reflection retrieval, two diag pathways, and timing-log emission.** Mitigation: before writing `retrieve()`, read `server.py:1128-1198` end-to-end and reproduce the call order verbatim. The 3-element order `spool.drain → retention.maybe_schedule → backend.retrieve` is asserted by the Step 2 test and corresponds to the spec invariant. Do NOT reorder. Diag/timing instrumentation (`_diag.trace`, `_diag.step`, `_run_best_effort`) MUST be preserved. +- **Risk: `_run_best_effort` is a module-level helper in server.py.** Mitigation: at plan-write time we verified `from better_memory.mcp.server import _run_best_effort` is importable (tests/mcp/test_best_effort_logging.py:17 already does this — `[[G2]]`). Handlers/observations.py should import it the same way OR move it to `mcp/handlers/_diag.py` alongside `_audit.py`. **Decision: keep it in `server.py` and import from there during Phase B.** Move to `_diag.py` in a follow-up if desired. +- **Risk: schema dicts have placeholders ("FILL IN").** Mitigation: per **G2**, every "FILL IN" block must be replaced with a verbatim copy of the corresponding `Tool(... inputSchema={...})` block from `server.py:_tool_definitions` BEFORE the task is marked complete. Per **G3** the acceptance check lists 4 schema dicts (`_OBSERVE_SCHEMA`, `_RETRIEVE_SCHEMA`, `_RETRIEVE_OBSERVATIONS_SCHEMA`, `_RECORD_USE_SCHEMA`) — all 4 must be present. +- **Risk: `observations.create` is async; other handler methods are sync.** Mitigation: every handler `__call__` is async (per dispatcher contract). Sync service calls (e.g. `services.semantic.create`) are simply not awaited. Async service calls (e.g. `services.observations.create`) ARE awaited. + +**Files:** +- Create: `better_memory/mcp/handlers/observations.py` +- Test: `tests/mcp/handlers/__init__.py` (empty), `tests/mcp/handlers/test_observations_handler.py` + +The four tools to migrate: `memory.observe`, `memory.retrieve`, `memory.retrieve_observations`, `memory.record_use`. Source today: `server.py` lines 1058-1218. + +- [ ] **Step 1: Create the test scaffold** + +```bash +mkdir -p tests/mcp/handlers +touch tests/mcp/handlers/__init__.py +``` + +- [ ] **Step 2: Write the failing handler test** + +```python +# tests/mcp/handlers/test_observations_handler.py +"""ObservationHandlers: route the 4 observation tools + preserve work order.""" +from __future__ import annotations + +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.handlers.observations import ObservationHandlers + + +def _stub_services() -> ServiceContainer: + return ServiceContainer( + config=MagicMock(), memory_conn=MagicMock(), backend=MagicMock(), + episodes=MagicMock(), observations=MagicMock(spec=["create", "list_observations"]), + reflections=MagicMock(), retention=MagicMock(), memory_rating=MagicMock( + spec=["record_use"]), + knowledge=MagicMock(), spool=MagicMock(spec=["drain"]), + semantic=MagicMock(), session_bootstrap=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_observe_routes_to_observations_create() -> None: + services = _stub_services() + services.observations.create = AsyncMock(return_value="obs-123") + handler = ObservationHandlers() + result = await handler.observe(services, {"content": "hello"}) + payload = json.loads(result[0].text) + assert payload == {"id": "obs-123"} + services.observations.create.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_observe_falls_back_to_project_when_scope_is_null() -> None: + services = _stub_services() + services.observations.create = AsyncMock(return_value="x") + handler = ObservationHandlers() + await handler.observe(services, {"content": "x", "scope": None}) + assert services.observations.create.call_args.kwargs["scope"] == "project" + + +@pytest.mark.asyncio +async def test_retrieve_calls_spool_drain_before_backend_retrieve() -> None: + """The legacy code drained spool first so fresh hooks are visible + to retrieve. Preserve this ordering.""" + services = _stub_services() + order: list[str] = [] + services.spool.drain = lambda: order.append("spool") + services.backend.retrieve = MagicMock(side_effect=lambda **kw: order.append("retrieve") or {"reflections": []}) + handler = ObservationHandlers() + await handler.retrieve(services, {}) + assert order == ["spool", "retrieve"] +``` + +Run: `uv run pytest tests/mcp/handlers/test_observations_handler.py -v` +Expected: FAIL with `ModuleNotFoundError: No module named 'better_memory.mcp.handlers.observations'` + +- [ ] **Step 3: Create the handler module** + +```python +# better_memory/mcp/handlers/observations.py +"""Observation-domain MCP tool handlers. + +Tools: memory.observe, memory.retrieve, memory.retrieve_observations, +memory.record_use. +""" +from __future__ import annotations + +import json +from typing import Any + +from mcp.types import TextContent + +from better_memory import _diag +from better_memory.config import project_name +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler + +_OBSERVE_SCHEMA = { + "type": "object", + "required": ["content"], + "additionalProperties": False, + "properties": { + "content": {"type": "string"}, + "component": {"type": "string"}, + "theme": {"type": "string"}, + "trigger_type": {"type": "string"}, + "outcome": {"type": "string", "enum": ["success", "failure", "neutral"]}, + "tech": {"type": "string"}, + "scope": {"type": "string", "enum": ["project", "general"]}, + }, +} + +_RETRIEVE_SCHEMA = { + "type": "object", + "additionalProperties": False, + "properties": { + # Copy verbatim from server.py:_tool_definitions for memory.retrieve. + # See spec §"Existing tests" — the schema dict is preserved exactly. + # FILL IN from server.py lines for memory.retrieve. + }, +} + +_RETRIEVE_OBSERVATIONS_SCHEMA = { + "type": "object", + # FILL IN from server.py for memory.retrieve_observations. +} + +_RECORD_USE_SCHEMA = { + "type": "object", + # FILL IN from server.py for memory.record_use. +} + + +class ObservationHandlers: + """All memory.observe / .retrieve / .retrieve_observations / .record_use tools.""" + + async def observe( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + with _diag.trace( + "mcp.memory.observe", + content_len=len(args.get("content") or ""), + scope=args.get("scope") or "project", + component=args.get("component"), + ): + _diag.step("mcp.memory.observe", "calling_observations_create") + obs_id = await services.observations.create( + content=args["content"], + component=args.get("component"), + theme=args.get("theme"), + trigger_type=args.get("trigger_type"), + outcome=args.get("outcome", "neutral"), + tech=args.get("tech"), + scope=args.get("scope") or "project", + ) + _diag.step("mcp.memory.observe", "create_returned", obs_id=obs_id) + return [TextContent(type="text", text=json.dumps({"id": obs_id}))] + + async def retrieve( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + # Lift the body verbatim from server.py:1128-1198 (memory.retrieve). + # Key invariant: spool.drain BEFORE retention.maybe_schedule BEFORE + # backend.retrieve. Do NOT change the order. + # FILL IN from server.py. + ... + + async def retrieve_observations( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + # Lift from server.py:1199-1211. FILL IN. + ... + + async def record_use( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + # Lift from server.py:1212-1218. FILL IN. + ... + + def handlers(self) -> list[Handler]: + return [ + Handler( + name="memory.observe", + description="Record an observation about the current session…", + schema=_OBSERVE_SCHEMA, + call=self.observe, + ), + Handler( + name="memory.retrieve", + description="…", + schema=_RETRIEVE_SCHEMA, + call=self.retrieve, + ), + Handler( + name="memory.retrieve_observations", + description="…", + schema=_RETRIEVE_OBSERVATIONS_SCHEMA, + call=self.retrieve_observations, + ), + Handler( + name="memory.record_use", + description="…", + schema=_RECORD_USE_SCHEMA, + call=self.record_use, + ), + ] +``` + +Implementation note: each `# FILL IN from server.py` block is a verbatim copy from the existing dispatcher. Copy the inner body (everything inside the `if name == "..."` block, omitting the `if`-line and the `return`), substituting service lookups with `services.`. For example, `observations.create(...)` becomes `services.observations.create(...)`. + +- [ ] **Step 4: Register in `handlers/__init__.py`** + +Update `better_memory/mcp/handlers/__init__.py`: + +```python +from better_memory.mcp.dispatcher import Handler +from better_memory.mcp.handlers.observations import ObservationHandlers + + +def all_handlers() -> list[Handler]: + return [ + *ObservationHandlers().handlers(), + ] +``` + +- [ ] **Step 5: Run tests** + +Run: `uv run pytest tests/mcp/handlers/test_observations_handler.py -v` +Expected: PASS. + +- [ ] **Step 6: Commit** + +```bash +git add better_memory/mcp/handlers/observations.py better_memory/mcp/handlers/__init__.py \ + tests/mcp/handlers/__init__.py tests/mcp/handlers/test_observations_handler.py +git commit -m "feat(mcp): observation-domain handlers (memory.observe + retrieve + retrieve_observations + record_use) + +Schemas + bodies lifted verbatim from _call_tool. spool.drain ordering +preserved and asserted by test." +``` + +--- + +### Task 8: `handlers/semantics.py` — semantic domain (4 tools) (94%) + +**Files:** +- Create: `better_memory/mcp/handlers/semantics.py` +- Test: `tests/mcp/handlers/test_semantics_handler.py` + +The four tools: `memory.semantic_observe`, `memory.semantic_retrieve`, `memory.semantic_update`, `memory.semantic_delete`. Source today: `server.py:1083-1126`. + +- [ ] **Step 1: Write the failing handler test** + +```python +# tests/mcp/handlers/test_semantics_handler.py +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +import pytest + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.handlers.semantics import SemanticHandlers + + +def _services_with_semantic(semantic) -> ServiceContainer: + return ServiceContainer( + config=MagicMock(), memory_conn=MagicMock(), backend=MagicMock(), + episodes=MagicMock(), observations=MagicMock(), reflections=MagicMock(), + retention=MagicMock(), memory_rating=MagicMock(), knowledge=MagicMock(), + spool=MagicMock(), semantic=semantic, session_bootstrap=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_semantic_observe_falls_back_to_project_when_scope_null() -> None: + semantic = MagicMock() + semantic.create.return_value = "sem-1" + services = _services_with_semantic(semantic) + handler = SemanticHandlers() + await handler.observe(services, {"content": "x", "scope": None}) + assert semantic.create.call_args.kwargs["scope"] == "project" + + +@pytest.mark.asyncio +async def test_semantic_observe_returns_id() -> None: + semantic = MagicMock() + semantic.create.return_value = "sem-42" + services = _services_with_semantic(semantic) + handler = SemanticHandlers() + result = await handler.observe(services, {"content": "hello"}) + assert json.loads(result[0].text) == {"id": "sem-42"} + + +@pytest.mark.asyncio +async def test_semantic_delete_returns_ok() -> None: + semantic = MagicMock() + services = _services_with_semantic(semantic) + handler = SemanticHandlers() + result = await handler.delete(services, {"id": "sem-1"}) + assert json.loads(result[0].text) == {"ok": True} + semantic.delete.assert_called_once_with(id="sem-1") +``` + +Run: FAIL with `ModuleNotFoundError`. + +- [ ] **Step 2: Create the handler module** + +```python +# better_memory/mcp/handlers/semantics.py +"""Semantic-memory MCP tool handlers (memory.semantic_*).""" +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.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler + +_OBSERVE_SCHEMA = { + "type": "object", "required": ["content"], "additionalProperties": False, + "properties": { + "content": {"type": "string"}, + "scope": {"type": "string", "enum": ["project", "general"]}, + }, +} +_RETRIEVE_SCHEMA = { + "type": "object", "additionalProperties": False, + "properties": {"project": {"type": "string"}}, +} +_UPDATE_SCHEMA = { + "type": "object", "required": ["id", "content"], "additionalProperties": False, + "properties": {"id": {"type": "string"}, "content": {"type": "string"}}, +} +_DELETE_SCHEMA = { + "type": "object", "required": ["id"], "additionalProperties": False, + "properties": {"id": {"type": "string"}}, +} + + +class SemanticHandlers: + async def observe( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + memory_id = services.semantic.create( + content=args["content"], + project=project_name(), + scope=args.get("scope") or "project", + ) + return [TextContent(type="text", text=json.dumps({"id": memory_id}))] + + async def retrieve( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + project = args.get("project") or project_name() + memories = services.semantic.list_for_project(project=project) + payload = [ + { + "id": m.id, "content": m.content, "project": m.project, + "scope": m.scope, "created_at": m.created_at, "updated_at": m.updated_at, + } + for m in memories + ] + return [TextContent(type="text", text=json.dumps(payload))] + + async def update( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + services.semantic.update_text(id=args["id"], content=args["content"]) + return [TextContent(type="text", text=json.dumps({"ok": True}))] + + async def delete( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + services.semantic.delete(id=args["id"]) + return [TextContent(type="text", text=json.dumps({"ok": True}))] + + def handlers(self) -> list[Handler]: + return [ + Handler("memory.semantic_observe", _OBSERVE_SCHEMA, self.observe), + Handler("memory.semantic_retrieve", _RETRIEVE_SCHEMA, self.retrieve), + Handler("memory.semantic_update", _UPDATE_SCHEMA, self.update), + Handler("memory.semantic_delete", _DELETE_SCHEMA, self.delete), + ] +``` + +- [ ] **Step 3: Register in `handlers/__init__.py`** + +```python +from better_memory.mcp.handlers.semantics import SemanticHandlers + +def all_handlers() -> list[Handler]: + return [ + *ObservationHandlers().handlers(), + *SemanticHandlers().handlers(), + ] +``` + +- [ ] **Step 4: Run tests + full MCP suite** + +Run: `uv run pytest tests/mcp/handlers/test_semantics_handler.py tests/mcp/ -v` +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +git add better_memory/mcp/handlers/semantics.py better_memory/mcp/handlers/__init__.py \ + tests/mcp/handlers/test_semantics_handler.py +git commit -m "feat(mcp): semantic-domain handlers (memory.semantic_observe + retrieve + update + delete) + +SemanticMemoryService is now built ONCE on the container, not 4× per call." +``` + +--- + +### Task 9: `handlers/episodes.py` — episode lifecycle (4 tools) (85%) + +**Confidence: 85% — mitigations:** + +- **Risk: `memory.close_episode` body has a `try/except ValueError` branch that returns an ALTERNATIVE payload (server.py:1280-1318).** Mitigation: copy both branches verbatim. The try-block returns `{"episode_id": id, "closed_at": ...}`; the except-block (when episode_id is stale) returns `{"episode_id": id, "closed_at": None, "reason": "already_closed"}`-shaped payload. Read the exact lines before writing. +- **Risk: `memory.start_episode` calls the now-public `services.reflections.read_queue_counts(...)` (post-Task 5) AND `services.backend.retrieve(...)` AND builds a composite payload inline.** Mitigation: Task 5 must be complete first (linear dep). Verify with `grep "read_queue_counts" better_memory/services/reflection.py` — expect the public name, not the underscore-prefixed one. +- **Risk: `memory.list_episodes` is a 10-field list-comp serializer.** Mitigation: copy the `[{...} for e in episodes]` list-comp verbatim from `server.py:1357-1379`. Per **G3** acceptance check: count the keys in the new code, it must be exactly 10. +- **Risk: `memory.reconcile_episodes` returns a list-comp over episode tuples.** Mitigation: copy verbatim from `server.py:1319-1334`. + +**Files:** +- Create: `better_memory/mcp/handlers/episodes.py` + +Tools: `memory.start_episode`, `memory.close_episode`, `memory.reconcile_episodes`, `memory.list_episodes`. Source: `server.py:1250-1379`. + +- [ ] **Step 1: Lift implementation verbatim** + +Create `better_memory/mcp/handlers/episodes.py` with the four `async def` methods named `start_episode`, `close_episode`, `reconcile_episodes`, `list_episodes`. Each body is the verbatim block from `server.py` for the matching `if name == "..."` branch, substituting service lookups (e.g. `episodes.foo(...)` → `services.episodes.foo(...)`, `reflections.read_queue_counts(...)` → `services.reflections.read_queue_counts(...)`). + +For `start_episode` specifically, this handler reaches the just-promoted `services.reflections.read_queue_counts(...)` plus `services.backend.retrieve(...)`. Mirror the existing payload construction exactly. + +For `close_episode`, replicate the default-reason map + the `try/except ValueError` branch that returns an alternative payload. + +- [ ] **Step 2: Schemas** + +Lift each `Tool(...)` schema dict from `_tool_definitions` (server.py:259-806) into module-level constants named `_START_EPISODE_SCHEMA`, `_CLOSE_EPISODE_SCHEMA`, `_RECONCILE_EPISODES_SCHEMA`, `_LIST_EPISODES_SCHEMA`. + +- [ ] **Step 3: Build the `handlers()` registration list** + +```python +def handlers(self) -> list[Handler]: + return [ + Handler("memory.start_episode", _START_EPISODE_SCHEMA, self.start_episode), + Handler("memory.close_episode", _CLOSE_EPISODE_SCHEMA, self.close_episode), + Handler("memory.reconcile_episodes", _RECONCILE_EPISODES_SCHEMA, self.reconcile_episodes), + Handler("memory.list_episodes", _LIST_EPISODES_SCHEMA, self.list_episodes), + ] +``` + +- [ ] **Step 4: Register in `handlers/__init__.py`** + +```python +from better_memory.mcp.handlers.episodes import EpisodeHandlers +# Add *EpisodeHandlers().handlers() to all_handlers(). +``` + +- [ ] **Step 5: Run tests** + +Run: `uv run pytest tests/mcp/test_episode_tools.py -v` +Expected: PASS (the existing test exercises the underlying service, but importing the new module must not break the suite). + +- [ ] **Step 6: Commit** + +```bash +git add better_memory/mcp/handlers/episodes.py better_memory/mcp/handlers/__init__.py +git commit -m "feat(mcp): episode-lifecycle handlers (start + close + reconcile + list) + +Migrates 4 tools; uses the new public read_queue_counts on +ReflectionSynthesisService." +``` + +--- + +### Task 10: `handlers/reflections.py` — synthesis (2 tools) (82%) + +**Confidence: 82% — mitigations:** + +- **Risk: `memory.synthesize_next_apply` has nested `try/except` over TWO distinct exception types and THREE return paths (server.py:1433-1486).** Mitigation: enumerate the branches before writing the handler: + 1. `try: ... except SynthesisResponseError as exc: state["result_kind"] = "response_error"; return [TextContent(... error_payload ...)]` + 2. `try: ... except Exception as exc: state["result_kind"] = "exception"; raise` + 3. Happy path: `state["result_kind"] = "applied"`; return success payload +- **Risk: `_audit_synth_call` context manager mutates `state` dict in both branches.** Mitigation: the existing audit shape is byte-for-byte preserved. Test `tests/mcp/test_synth_audit_log.py` is the contract. After Task 4 + this task, that test must remain green. +- **Risk: `from better_memory.services.reflection import SynthesisResponseError` is the inline import at server.py:1434.** Mitigation: hoist this to the module-level imports of `handlers/reflections.py`. Verify the name exists by grepping `class SynthesisResponseError` in `services/reflection.py`. +- **Risk: both handlers are capability-gated (`requires_synthesis=True`).** Mitigation: the `handlers()` method MUST pass `requires_synthesis=True` for both. Per **G3** acceptance check: both `Handler(...)` constructors set this flag. +- **Risk: `_audit_synth_call` requires `home: Path` argument.** Mitigation: resolve from `services.config.home` inside each handler. + +**Files:** +- Create: `better_memory/mcp/handlers/reflections.py` + +Tools: `memory.synthesize_next_get_context`, `memory.synthesize_next_apply`. Source: `server.py:1413-1486`. Uses the moved `_audit_synth_call` context manager. + +- [ ] **Step 1: Lift implementation verbatim** + +Create the module with two `async def` methods. Both bodies wrap in `with _audit_synth_call(...) as state:` exactly as today. The `synthesize_next_apply` body includes a nested `try/except` over `SynthesisResponseError` and a separate branch for generic exceptions — preserve both. The third return path (for the "decision was 'noop'" case) is preserved verbatim. + +```python +# Imports at the top: +from better_memory.mcp.handlers._audit import _audit_synth_call +``` + +Both handlers are marked `requires_synthesis=True` on the `Handler` dataclass. + +- [ ] **Step 2: Schemas** + +Lift `_GET_CONTEXT_SCHEMA` and `_APPLY_SCHEMA` from `_tool_definitions`. + +- [ ] **Step 3: Registration** + +```python +def handlers(self) -> list[Handler]: + return [ + Handler( + "memory.synthesize_next_get_context", _GET_CONTEXT_SCHEMA, + self.get_context, requires_synthesis=True, + ), + Handler( + "memory.synthesize_next_apply", _APPLY_SCHEMA, + self.apply, requires_synthesis=True, + ), + ] +``` + +- [ ] **Step 4: Register in `handlers/__init__.py`**, run tests, commit. + +```bash +git add better_memory/mcp/handlers/reflections.py better_memory/mcp/handlers/__init__.py +git commit -m "feat(mcp): synthesis handlers (synthesize_next_get_context + apply) + +Both gated on backend.supports_synthesis. Audit-log byte-shape +preserved by reusing _audit_synth_call from handlers/_audit.py." +``` + +--- + +### Task 11: `handlers/retention.py` — retention (1 tool) (94%) + +**Files:** +- Create: `better_memory/mcp/handlers/retention.py` + +Tool: `memory.run_retention`. Source: `server.py:1335-1356` (8-field dict construction from `RetentionReport`). + +- [ ] **Step 1: Lift verbatim** + +One `async def run_retention(self, services, args)` method that calls `services.retention.run(...)` and serialises the resulting `RetentionReport` exactly as today. + +- [ ] **Step 2: Schema + registration + tests + commit** + +```bash +git add better_memory/mcp/handlers/retention.py better_memory/mcp/handlers/__init__.py +git commit -m "feat(mcp): retention handler (memory.run_retention)" +``` + +--- + +### Task 12: `handlers/knowledge.py` — knowledge tools (2 tools) (93%) + +**Files:** +- Create: `better_memory/mcp/handlers/knowledge.py` + +Tools: `knowledge.search`, `knowledge.list`. Source: `server.py:1219-1243`. Uses existing `_serialize_knowledge_search_result` and `_serialize_knowledge_document` helpers — copy those into the module (they're tiny). + +- [ ] **Step 1-5: Lift + schemas + registration + tests + commit.** + +```bash +git add better_memory/mcp/handlers/knowledge.py better_memory/mcp/handlers/__init__.py +git commit -m "feat(mcp): knowledge handlers (search + list)" +``` + +--- + +### Task 13: `handlers/ratings.py` — memory rating (3 tools) (87%) + +**Confidence: 87% — mitigations:** + +- **Risk: `memory.credit` returns TWO DIFFERENT payload shapes depending on session state (server.py:1512-1525).** Mitigation: copy both branches. When session id is resolved, returns `{"applied": True, ...}` shape. When unresolved, returns `{"applied": False, "reason": "no_session"}`-shape (verify the exact key by reading server.py:1512-1525 at implementation time). Add a unit test for each branch in `tests/mcp/handlers/test_ratings_handler.py`. +- **Risk: `memory.apply_session_ratings` raises a `ValueError` with a multi-line message (server.py:1498-1511).** Mitigation: copy the message string verbatim — any client / test pinning the error text will break otherwise. +- **Risk: `memory.list_session_exposures` calls `services.session_bootstrap` (inline-imported today at server.py:1488).** Mitigation: now sourced from the container — confirm `services.session_bootstrap` is correctly populated by `_build_services` (Task 6). Per **G4** grep `SessionBootstrapService(` in tests to find call sites. + +**Files:** +- Create: `better_memory/mcp/handlers/ratings.py` + +Tools: `memory.list_session_exposures`, `memory.apply_session_ratings`, `memory.credit`. Source: `server.py:1487-1525`. Each handler calls `resolve_session_id(services.config.home)` to get the session id. + +```python +# Imports at the top: +from better_memory.mcp._session import resolve_session_id +``` + +- [ ] **Step 1-5: Lift + schemas + registration + tests + commit.** + +The `list_session_exposures` handler also uses `services.session_bootstrap` — confirm with `grep` that today's body matches this assumption. + +```bash +git add better_memory/mcp/handlers/ratings.py better_memory/mcp/handlers/__init__.py +git commit -m "feat(mcp): rating handlers (list_session_exposures + apply + credit) + +Uses resolve_session_id from mcp/_session.py." +``` + +--- + +### Task 14: `handlers/session.py` — session bootstrap + start_ui (2 tools) (86%) + +**Confidence: 86% — mitigations:** + +- **Risk: `memory.session_bootstrap` is a dense 30-LOC body that constructs `SessionBootstrapService` inline today (server.py:1380-1412) and resolves CWD + session id env vars inline.** Mitigation: read `server.py:1380-1412` AND `services/session_bootstrap.py` BEFORE writing the handler. The new body uses `services.session_bootstrap` directly (no construction); the env-var resolution moves to `resolve_session_id(services.config.home)`. +- **Risk: `memory.start_ui` calls `ui_launcher.start_ui()` and returns a payload describing where the UI was launched (server.py:1244-1249).** Mitigation: import `from better_memory.services import ui_launcher` at module level. Copy body verbatim. +- **Risk: `services.session_bootstrap` field on the container must be populated by Task 6.** Mitigation: this task is downstream of Task 6 in the linear order. Verify with `grep "session_bootstrap=" better_memory/mcp/server.py` — expect a `session_bootstrap=SessionBootstrapService(memory_conn)` line inside `_build_services`. + +**Files:** +- Create: `better_memory/mcp/handlers/session.py` + +Tools: `memory.session_bootstrap`, `memory.start_ui`. Source: `server.py:1244-1249, 1380-1412`. + +- [ ] **Step 1: Lift verbatim** + +```python +# better_memory/mcp/handlers/session.py +"""Session-lifecycle MCP tool handlers.""" +from __future__ import annotations + +import json +import os +from typing import Any + +from mcp.types import TextContent + +from better_memory.config import project_name +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler +from better_memory.mcp._session import resolve_session_id +from better_memory.services import ui_launcher + +_BOOTSTRAP_SCHEMA = {...} # lift from _tool_definitions +_START_UI_SCHEMA = {...} # lift from _tool_definitions + + +class SessionHandlers: + async def session_bootstrap( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + # Lift body from server.py:1380-1412 — use services.session_bootstrap + # instead of constructing it inline. + ... + + async def start_ui( + self, services: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + # Lift body from server.py:1244-1249. + ... + + def handlers(self) -> list[Handler]: + return [ + Handler("memory.session_bootstrap", _BOOTSTRAP_SCHEMA, self.session_bootstrap), + Handler("memory.start_ui", _START_UI_SCHEMA, self.start_ui), + ] +``` + +- [ ] **Step 2: Register + tests + commit.** + +```bash +git add better_memory/mcp/handlers/session.py better_memory/mcp/handlers/__init__.py +git commit -m "feat(mcp): session handlers (session_bootstrap + start_ui) + +SessionBootstrapService is now built ONCE on the container, not 2× per call." +``` + +--- + +### Task 15: Wire `ToolDispatcher` into `create_server` and delete the if-chain (88%) + +**Confidence: 88% — mitigations (load-bearing task):** + +- **Risk: `all_handlers()` count mismatch.** Mitigation: Step 1 explicitly asserts `len(all_handlers()) == 22`. Block on this before touching `create_server`. +- **Risk: `ServerContext` dataclass needs a new `dispatcher` field, and existing tests may construct `ServerContext(...)` positionally.** Mitigation: per **G4**, `grep "ServerContext(" tests/` and update every call site. Default `dispatcher` to `None` for backward compat so positional construction still works. +- **Risk: `cleanup` async function uses `await embedder.aclose()` — verify the method name in `OllamaEmbedder`.** Mitigation: at task time, grep `def aclose\|async def aclose` in `better_memory/embeddings/ollama.py`. If the method is named differently (`.close()`, `.shutdown()`), use that name. +- **Risk: `_call_tool` registration in the new `create_server` is a 1-liner closing over `dispatcher`.** Mitigation: this is fine — closures over `dispatcher` only need it once at registration time. +- **Risk: docstring drift — server.py module docstring (lines 1-34) enumerates tools.** Per **G5**, sweep: + - `better_memory/mcp/server.py` module docstring — tools list at lines 7-22 (NO change needed, the same 22 tools still ship). + - `README.md` "server registers N tools" line — verify N is still correct. + - `website/mcp-tools.md` — verify all schemas are still discoverable (they live in handler modules now, but the docs file may reference them). + - `website/architecture.md` — no tool surface change so likely unaffected. Note 'docs unaffected' explicitly in the commit message. +- **Risk: full pytest run in Step 4 may turn up unexpected failures.** Mitigation: if any test fails that the spec said should pass unchanged, FIRST `git stash` and re-run the failing tests against `HEAD` to confirm it's pre-existing vs. introduced by this task (per inherited-test-failures pattern). If introduced here, fix in this task; if pre-existing, note in commit body. +- **Rollback plan:** if Step 4 fails catastrophically, `git reset --hard HEAD~N` rolls back the entire refactor (N = commits in Tasks 1-15). The single-PR shape lets us bail without partial-state mess. + +**Files:** +- Modify: `better_memory/mcp/server.py` (the big one) + +This is the load-bearing task. Every handler must already be registered (Tasks 7-14 done) and `all_handlers()` must return 22 handlers. + +- [ ] **Step 1: Verify all 22 handlers are registered** + +```bash +uv run python -c "from better_memory.mcp.handlers import all_handlers; print(len(all_handlers()))" +``` + +Expected: `22`. If less, find which domain is missing. + +- [ ] **Step 2: Rewrite `create_server`** + +Replace the entire body of `create_server` (today: server.py:953-1561) with: + +```python +def create_server() -> tuple[ + Server, + Callable[[], Coroutine[Any, Any, None]], + ServerContext, +]: + """Wire services and register tools. + + Returns ``(server, cleanup, ctx)``. ``ctx.dispatcher`` is the + :class:`ToolDispatcher` instance so tests can drive tool calls + directly without hand-rolling the SDK plumbing. + """ + config = get_config() + memory_conn = connect(config.memory_db) + apply_migrations(memory_conn, migrations_dir=_MEMORY_MIGRATIONS) + knowledge_conn = connect(config.knowledge_db) + apply_migrations(knowledge_conn, migrations_dir=_KNOWLEDGE_MIGRATIONS) + + embedder: OllamaEmbedder | None = None + if config.embeddings_backend == "ollama": + embedder = OllamaEmbedder() + _probe_ollama(config.ollama_host) + + services = _build_services( + config, memory_conn, knowledge_conn, embedder, + startup_project=project_name(), + startup_session_id=resolve_session_id(config.home) or None, + ) + + if config.knowledge_base.is_dir(): + try: + services.knowledge.reindex() + except Exception: # noqa: BLE001 — best-effort startup hook + pass + + dispatcher = ToolDispatcher(services, all_handlers()) + server: Server = Server(name="better-memory") + + @server.list_tools() + async def _list_tools() -> list[Tool]: + return dispatcher.tool_definitions() + + @server.call_tool() + async def _call_tool( + name: str, arguments: dict[str, Any] | None, + ) -> list[TextContent]: + return await dispatcher.call(name, arguments or {}) + + async def cleanup() -> None: + # Same idempotent cleanup as before — close conns, embedder. + memory_conn.close() + knowledge_conn.close() + if embedder is not None: + await embedder.aclose() + + return server, cleanup, ServerContext( + backend=services.backend, + memory_conn=memory_conn, + embedder=embedder, + dispatcher=dispatcher, + ) +``` + +Add `dispatcher: ToolDispatcher | None = None` to the `ServerContext` dataclass (line 940-950 area). + +- [ ] **Step 3: Delete dead code** + +Delete from `server.py`: +- `_tool_definitions` (entire function, ~550 LOC, today's lines 259-806) +- Every `if name == "..."` branch in the old `_call_tool` (already gone if you replaced `create_server`) +- The `_append_synth_audit` and `_audit_synth_call` definitions (already moved in Task 4 but double-check they're not still present) + +Add at the top of the file: + +```python +from better_memory.mcp.dispatcher import ToolDispatcher +from better_memory.mcp.handlers import all_handlers +``` + +- [ ] **Step 4: Run the full test suite** + +Run: `uv run pytest -v` +Expected: ALL PASS. Specifically watch: +- `tests/mcp/test_server_sqlite.py` — happy-path roundtrip +- `tests/mcp/test_server_backend_dispatch.py` — capability gating +- `tests/mcp/test_synth_audit_log.py` — audit log shape +- `tests/mcp/test_rating_tools.py` — uses `_dispatch_for_tests` +- All `tests/mcp/test_*_tool.py` — domain-level tests + +If `test_server_backend_dispatch.py` fails on a `ServerContext.dispatcher` field access, update it to expect the new field. + +- [ ] **Step 5: Commit** + +```bash +git add better_memory/mcp/server.py +git commit -m "refactor(mcp): wire ToolDispatcher into create_server, delete if-chain + +server.py drops from 1617 to ~150 LOC. _call_tool is now a 1-line +delegation to dispatcher.call. _tool_definitions deleted (schemas +co-located with handlers). All 22 tools route through the dispatcher." +``` + +--- + +### Task 16: Slim `_dispatch_for_tests` to a thin shim (92%) + +**Files:** +- Modify: `better_memory/mcp/server.py` (the `_dispatch_for_tests` function) + +- [ ] **Step 1: Locate the existing definition** + +Today at `server.py:1564-1603`. + +- [ ] **Step 2: Replace with a 5-line shim that uses the dispatcher** + +```python +async def _dispatch_for_tests(name: str, arguments: dict[str, Any]) -> Any: + """Test-only entry point. Routes through ToolDispatcher. + + Builds a server (with its standard ServiceContainer + every registered + handler) and dispatches one tool call. Preserved as a shim so + tests/mcp/test_server_sqlite.py and tests/mcp/test_rating_tools.py + continue to work unchanged.""" + server, cleanup, ctx = create_server() + try: + assert ctx.dispatcher is not None + return await ctx.dispatcher.call(name, arguments) + finally: + await cleanup() +``` + +- [ ] **Step 3: Run the tests that drive this shim** + +Run: `uv run pytest tests/mcp/test_server_sqlite.py tests/mcp/test_rating_tools.py -v` +Expected: PASS. + +- [ ] **Step 4: Commit** + +```bash +git add better_memory/mcp/server.py +git commit -m "refactor(mcp): slim _dispatch_for_tests to thin dispatcher.call shim + +5 LOC instead of 40. Same signature, same behaviour." +``` + +--- + +### Task 17: Final full-suite smoke + lint (99%) + +**Files:** none modified; verification only. + +- [ ] **Step 1: Run full pytest with verbose output** + +Run: `uv run pytest -v 2>&1 | tail -40` +Expected: ALL PASS. + +- [ ] **Step 2: Run ruff + pyright** + +Run: `uv run ruff check better_memory/mcp/ tests/mcp/` +Expected: clean. + +Run: `uv run pyright better_memory/mcp/ tests/mcp/` +Expected: clean (or pre-existing warnings only — no new ones introduced by the refactor). + +- [ ] **Step 3: Sanity-check `server.py` LOC** + +```bash +wc -l better_memory/mcp/server.py +``` + +Expected: ~150-180 LOC (was 1617). + +- [ ] **Step 4: Smoke-test the live server** + +Run: +```bash +uv run python -m better_memory.mcp 2>&1 & +sleep 2 +# Send a list_tools request via the MCP stdio protocol if your conftest +# has a smoke helper; otherwise inspect the launched process startup logs +# for the "ready" line, then kill it. +``` + +Expected: server starts, "knowledge reindex" log line appears, no crash. + +- [ ] **Step 5: Commit (verification-only — empty commit OK)** + +```bash +git commit --allow-empty -m "test(mcp): verify dispatcher refactor — full suite green, server.py 1617 → ~150 LOC" +``` + +--- + +## Self-review + +**Spec coverage:** +- ServiceContainer dataclass: Task 1 +- ToolDispatcher + Handler: Task 2 +- `_resolve_session_id` move: Task 3 +- `_audit_synth_call` move: Task 4 +- `_read_queue_counts` promotion: Task 5 +- `_build_services` extraction: Task 6 +- 22 tool handlers across 8 domain modules: Tasks 7-14 +- `create_server` rewrite: Task 15 +- `_dispatch_for_tests` shim: Task 16 +- Smoke verification: Task 17 + +Every numbered spec element has a task. No gaps. + +**Placeholder scan:** Tasks 7, 9-14 contain "FILL IN from server.py" or "..." markers on the schemas and handler bodies. These are deliberate copy-from-source directives. The plan's reader has the exact source file + line ranges, which is the minimum any verbatim-lift task can specify without exploding into 22 copy-paste blocks. Tasks 1-6 + 15-17 contain complete code. + +**Type consistency:** `ServiceContainer.semantic` is typed `Any` in Task 1 to defer the import (avoids a circular import on first load). Handlers in Tasks 8, 13, 14 use `services.semantic` and `services.session_bootstrap` exactly as defined. `Handler` dataclass has the same `(name, schema, call, description, requires_synthesis)` signature in every task that constructs one. + +**Risks not addressed by tasks:** +- If a domain handler's verbatim lift introduces a typo that the existing tests don't catch (e.g. a happy-path-only domain like retention), the bug ships. Mitigated by the prerequisite PBI `mcp-server-tool-error-path-tests` which the spec calls out. +- Task 6's test depends on the existing `get_config()` and migration paths working in tmp_path. If the repo's test conftest already provides a fixture for that, prefer it. + +--- + +## Execution Handoff + +Plan complete and saved to `docs/superpowers/plans/2026-06-13-mcp-server-dispatcher-refactor.md`. Two execution options: + +1. **Subagent-Driven (recommended)** — I dispatch a fresh subagent per task, review between tasks, fast iteration. + +2. **Inline Execution** — Execute tasks in this session using executing-plans, batch execution with checkpoints. + +Which approach? diff --git a/docs/superpowers/specs/2026-06-12-mcp-server-dispatcher-design.md b/docs/superpowers/specs/2026-06-12-mcp-server-dispatcher-design.md new file mode 100644 index 0000000..c138ce5 --- /dev/null +++ b/docs/superpowers/specs/2026-06-12-mcp-server-dispatcher-design.md @@ -0,0 +1,308 @@ +# MCP server dispatcher refactor + +Status: design +Date: 2026-06-12 +Target file: `better_memory/mcp/server.py` +Related tech-debt finding: `mcp-server-god-module-dispatcher` +Source report: `.tech-debt/reports/mcp-server-god-module-dispatcher.md` + +## Problem + +`better_memory/mcp/server.py` is the repo's worst hotspot: 1617 LOC, complexity 4880, 52 commits in 12 months, hotspot score 67.5 (2.1× the runner-up). The single function `_call_tool` (lines 1052-1525) is a 470-LOC `if name == "..."` chain dispatching 22 tools, inline-importing services in 7 places, instantiating `SemanticMemoryService` 4 times per call (lines 1086, 1101, 1118, 1124) and `SessionBootstrapService` 2 times (lines 1392, 1493). Schemas live in a sibling `_tool_definitions` block (lines 259-806, 550 LOC of hand-rolled JSON-Schema dicts) ~750 lines from the handler that uses them. `create_server` (lines 953-1561) is a 600-LOC god-factory. + +Consequence: every new tool, every schema tweak, and every service-wiring change re-pays the same tax. Contributors have already worked around this in test code — `_dispatch_for_tests` (line 1564) exists only because `_call_tool` cannot be called any other way, and `tests/mcp/test_synthesize_tools.py:17-22` documents the workaround verbatim. + +## Goals + +- Decompose `_call_tool` into a `ToolDispatcher` + per-domain handler modules. +- Build every long-lived service exactly once and hand them to handlers through a `ServiceContainer`. +- Co-locate each tool's `inputSchema` with the handler method that uses it. +- Preserve every behavioural invariant the current dispatcher upholds: connection ownership, capability gating, SDK error surface, audit-log ordering, `memory.retrieve` background-work order. +- Land as a single PR. There is no in-flight work on `_call_tool` to conflict with. + +## Non-goals + +- Not changing the MCP wire format or any tool name / schema field. +- Not introducing `pydantic` / `msgspec` argument models (tracked separately if desired). +- Not fixing `test_server_integration.py`'s stale skip (tracked under `mcp-integration-tests-stale-skip`). +- Not refactoring `services/reflection.py` or any other hotspot (tracked separately). +- Not changing the `_dispatch_for_tests` test entry point's signature. + +## Locked design choices + +The 5 open questions in the source report are resolved as follows: + +| # | Question | Choice | +|---|----------|--------| +| 1 | Handler granularity | **Per-domain class with methods (~9 files)** | +| 2 | Schema location | **Co-located in each handler module as module-level constants** | +| 3 | Migration shape | **Single big-bang PR** (no in-flight conflicts) | +| 4 | `_resolve_session_id` location | **Free function in `mcp/_session.py`** | +| 5 | `_dispatch_for_tests` | **Keep as a thin shim over `dispatcher.call`** | + +## Architecture + +### Target module layout + +``` +better_memory/mcp/ +├── __init__.py +├── __main__.py # unchanged +├── server.py # ~150 LOC: create_server, run, cleanup +├── _session.py # _resolve_session_id moved here (free function) +├── container.py # ServiceContainer dataclass +├── dispatcher.py # ToolDispatcher + Handler dataclass +└── handlers/ + ├── __init__.py # collects domain-handler instances → list[Handler] + ├── _audit.py # _audit_synth_call (moved from server.py) + ├── observations.py # ObservationHandlers (4 tools) + ├── semantics.py # SemanticHandlers (4 tools) + ├── episodes.py # EpisodeHandlers (4 tools) + ├── reflections.py # ReflectionHandlers (2 tools) + ├── retention.py # RetentionHandlers (1 tool) + ├── knowledge.py # KnowledgeHandlers (2 tools) + ├── ratings.py # RatingHandlers (3 tools) + └── session.py # SessionHandlers (memory.session_bootstrap, memory.start_ui) +``` + +### Core types + +```python +# container.py +@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 # built ONCE (was 4× inline) + session_bootstrap: SessionBootstrapService # built ONCE (was 2× inline) +``` + +```python +# dispatcher.py +@dataclass(frozen=True) +class Handler: + name: str + schema: dict[str, Any] # MCP inputSchema JSON-Schema + call: Callable[[ServiceContainer, dict[str, Any]], Awaitable[list[TextContent]]] + requires_synthesis: bool = False # capability gate + +class ToolDispatcher: + def __init__(self, services: ServiceContainer, handlers: list[Handler]): + self._services = services + self._handlers = {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=..., 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: + raise ValueError(f"Unknown tool: {name}") # same shape as today + return await handler.call(self._services, args) +``` + +```python +# handlers/semantics.py +_OBSERVE_SCHEMA = {"type": "object", "properties": {...}, "required": ["content"]} +_RETRIEVE_SCHEMA = {...} +_UPDATE_SCHEMA = {...} +_DELETE_SCHEMA = {...} + +class SemanticHandlers: + """All memory.semantic_* tools.""" + + async def observe(self, services: ServiceContainer, args: dict) -> list[TextContent]: + memory_id = services.semantic.create( + content=args["content"], + project=services.config.project_name, + scope=args.get("scope") or "project", + ) + return [TextContent(type="text", text=json.dumps({"id": memory_id}))] + + async def retrieve(self, services, args) -> list[TextContent]: ... + async def update(self, services, args) -> list[TextContent]: ... + async def delete(self, services, args) -> list[TextContent]: ... + + def handlers(self) -> list[Handler]: + return [ + Handler("memory.semantic_observe", _OBSERVE_SCHEMA, self.observe), + Handler("memory.semantic_retrieve", _RETRIEVE_SCHEMA, self.retrieve), + Handler("memory.semantic_update", _UPDATE_SCHEMA, self.update), + Handler("memory.semantic_delete", _DELETE_SCHEMA, self.delete), + ] +``` + +```python +# handlers/__init__.py +def all_handlers() -> list[Handler]: + return [ + *ObservationHandlers().handlers(), + *SemanticHandlers().handlers(), + *EpisodeHandlers().handlers(), + *ReflectionHandlers().handlers(), + *RetentionHandlers().handlers(), + *KnowledgeHandlers().handlers(), + *RatingHandlers().handlers(), + *SessionHandlers().handlers(), + ] +``` + +### `create_server` after refactor + +```python +# server.py — full file, ~150 LOC +def create_server() -> tuple[Server, Callable[..., Awaitable[None]], ServerContext]: + config = get_config() + memory_conn, audit_conn = _open_connections(config) + embedder = _build_embedder(config) + services = _build_services(config, memory_conn, audit_conn, embedder) + _best_effort_knowledge_reindex(services.knowledge, config) + + dispatcher = ToolDispatcher(services, all_handlers()) + + server = Server(name="better-memory") + + @server.list_tools() + async def _list_tools() -> list[Tool]: + return dispatcher.tool_definitions() + + @server.call_tool() + async def _call_tool(name: str, arguments: dict[str, Any] | None) -> list[TextContent]: + return await dispatcher.call(name, arguments or {}) + + cleanup = _build_cleanup(memory_conn, audit_conn, embedder) + return server, cleanup, ServerContext(backend=services.backend, dispatcher=dispatcher) + + +def _build_services(cfg, memory_conn, audit_conn, embedder) -> ServiceContainer: + """Construct every service ONCE. Replaces all inline service-build smells.""" + backend = build_backend(cfg, memory_conn, audit_conn, embedder) + return ServiceContainer( + config=cfg, + memory_conn=memory_conn, + backend=backend, + episodes=EpisodeService(memory_conn, backend), + observations=ObservationService(memory_conn, backend, embedder), + reflections=ReflectionSynthesisService(memory_conn, backend), + retention=RetentionService(memory_conn, backend), + memory_rating=MemoryRatingService(memory_conn), + knowledge=KnowledgeService(memory_conn), + spool=SpoolService(memory_conn), + semantic=SemanticMemoryService(memory_conn), # was 4× inline + session_bootstrap=SessionBootstrapService(...), # was 2× inline + ) +``` + +### Data flow per tool call + +``` +MCP stdio frame + │ + ▼ +Server SDK ──── _call_tool(name, args) + │ + ▼ +ToolDispatcher.call(name, args) + │ O(1) lookup in self._handlers dict + ▼ +Handler.call(services, args) + │ + ▼ +domain method (e.g. SemanticHandlers.observe) + │ uses services.semantic / services.memory_conn + ▼ +list[TextContent] ──── SDK serialises ──── stdio +``` + +## Invariants preserved + +| Invariant | Where it lives today | How we preserve it | +|-----------|---------------------|--------------------| +| **Connection ownership.** One shared `memory_conn`, MCP stdio serialises requests, SAVEPOINT safety. | `server.py:988-999` | Connection lives on `ServiceContainer` as a single attr; all services share it; handlers reach it only through the container; no new threads. | +| **Capability gating.** Synthesis tools hidden and rejected when `backend.supports_synthesis` is `False`. | Filter at `_list_tools` time; `_call_tool` falls through to `ValueError("Unknown tool")` | `Handler(requires_synthesis=True)` flag. Dispatcher filters in both `tool_definitions()` and `call()` with the same `ValueError("Unknown tool: {name}")` shape. | +| **SDK error surface.** Exceptions from handlers surface as `CallToolResult(isError=True)`. | SDK catches whatever `_call_tool` raises | Dispatcher does not catch — exceptions from handler methods propagate up through `dispatcher.call` to the SDK unchanged. | +| **Audit log ordering.** `_audit_synth_call` writes a `start` JSONL line then a `complete` line. | `server.py:176-221`, used by 2 synthesis handlers | Move to `mcp/handlers/_audit.py` with the same context manager signature. Used identically by `ReflectionHandlers.synthesize_next_get_context` and `.synthesize_next_apply`. `test_synth_audit_log.py` passes byte-for-byte. | +| **`memory.retrieve` background-work order.** `spool.drain` → `retention.maybe_schedule` → `backend.retrieve`. | `server.py:1143-1163` | Lifted verbatim into `ObservationHandlers.retrieve`. Order asserted by a new test using a recording stub container. | + +## Seams surfaced + +- `ReflectionSynthesisService._read_queue_counts` is called from 2 dispatcher branches (lines 1263, 1422). Promote to public `read_queue_counts` on the service — small ahead-of-time commit so the handler split is clean. +- `_resolve_session_id` (server.py:147) moves to `mcp/_session.py` as a free function. Handlers import directly; no container hop. +- `_audit_synth_call` moves to `mcp/handlers/_audit.py`. Reflection-specific; co-locate with the only handler module that uses it. + +## Testing + +### Existing tests — change matrix + +| File | What it covers | Refactor impact | +|------|---------------|-----------------| +| `test_server_sqlite.py` | Happy-path roundtrip via `_dispatch_for_tests` | Unchanged. Shim preserves entry point. | +| `test_server_backend_dispatch.py` | Capability gating, `create_server` return shape | Update assertions on the dispatcher field of `ServerContext`. | +| `test_synth_audit_log.py` | Audit log byte-shape | Update import path to `better_memory.mcp.handlers._audit`. | +| `test_rating_tools.py` | Rating tools via `_dispatch_for_tests` | Unchanged. | +| `test_episode_tools.py`, `test_semantic_tools.py`, `test_retention_tool.py`, `test_session_bootstrap_tool.py`, `test_start_ui_tool.py` (5 files) | Service-level wrappers | Unchanged. The "tool is a thin wrapper" docstring workaround no longer needs to apologise. | +| `test_synthesize_tools.py` | Serializer helpers + service round-trip | Update imports for any serializer helpers that move into `handlers/reflections.py`. | +| `test_server_integration.py` | Subprocess MCP integration — currently `pytest.mark.skip` | Out of scope (tracked under `mcp-integration-tests-stale-skip`). | + +### New test files (4) + +``` +tests/mcp/ +├── test_tool_dispatcher.py NEW +├── test_service_container.py NEW +├── handlers/ +│ ├── test_semantics_handler.py NEW +│ └── test_observations_handler.py NEW +``` + +| New file | Asserts | +|----------|---------| +| `test_tool_dispatcher.py` | Parametrise over all 22 tool names: registered iff capability allows, `call("nope", {})` raises `ValueError("Unknown tool: nope")`, capability-gated tool with capability off raises same shape, every `Handler` obeys the dataclass contract. | +| `test_service_container.py` | Build a container, monkeypatch each service `__init__` to count calls. Assert **each service constructed exactly once** — kills the 4× `SemanticMemoryService` and 2× `SessionBootstrapService` regressions permanently. | +| `handlers/test_semantics_handler.py` | Per-method tests with a stub container. Covers `scope=None` fallback to `"project"` (regresses the PR #25 BugBot fix). One happy + at least one error path per tool. | +| `handlers/test_observations_handler.py` | Recording stub container asserts `memory.retrieve` calls `spool.drain` → `retention.maybe_schedule` → `backend.retrieve` in that order. | + +### Sibling-finding sequencing + +The sibling tech-debt finding `mcp-server-tool-error-path-tests` (PBI in `top5.json`) gives the per-tool error coverage for all 22 tools. **Run it first.** Each green error-path test guards a domain we then migrate so no PR ships handlers with weaker coverage than today. + +## Risk + +- **Test churn.** ~10 existing test files. 7 unchanged, 3 small import / assertion updates. Risk: low. +- **Audit-log byte-shape regression.** `_audit_synth_call` is the largest preserved invariant. Mitigated by keeping the context manager signature identical and re-running `test_synth_audit_log.py`. +- **Capability-gate double check.** Dispatching a synthesis-only tool when the backend has no synthesis support must raise the same `ValueError("Unknown tool: ...")` as today, not a new `ValueError("Capability disabled")`. Test in `test_tool_dispatcher.py` pins this. +- **Inline-service regression.** Future contributors could re-introduce per-call `SemanticMemoryService(memory_conn)` in a handler. `test_service_container.py` catches this directly. + +## Effort + +Large (L), ~7-10 focused days, single PR: + +| Phase | Days | Output | +|-------|------|--------| +| A — Scaffolding | 1.5 | `container.py`, `dispatcher.py`, `_session.py`, `handlers/__init__.py`, `handlers/_audit.py`, `Handler` dataclass, `ToolDispatcher`, `ServiceContainer`. Tests: `test_tool_dispatcher.py`, `test_service_container.py`. | +| B — Bulk migration | 4-5 | 8 handler modules (observations, semantics, knowledge, episodes, reflections, retention, ratings, session). Promote `_read_queue_counts`. Migrate `_audit_synth_call`. ~0.5-1 d per domain. | +| C — Cleanup | 1.5-2 | Delete `_call_tool` if-chain, delete inline-import lines, retire `_tool_definitions`, slim `_dispatch_for_tests` to a 3-liner, update test docstrings, write `test_observations_handler.py` + `test_semantics_handler.py`. | + +The sibling `mcp-server-tool-error-path-tests` PBI should land **before** Phase A so each domain we migrate has error-path coverage to fall back on. + +## Out of scope (follow-ups) + +- Replacing hand-rolled JSON-Schema dicts with pydantic / msgspec argument models. Genuinely better long-term; bundling here risks blowing up scope on a "schema shape changed" debate. +- Splitting `services/reflection.py` (its own god-module finding, `reflection-service-god-module`). +- Fixing `test_server_integration.py` (`mcp-integration-tests-stale-skip` finding). diff --git a/tests/mcp/handlers/__init__.py b/tests/mcp/handlers/__init__.py new file mode 100644 index 0000000..8b568a8 --- /dev/null +++ b/tests/mcp/handlers/__init__.py @@ -0,0 +1 @@ +"""Tests for per-domain MCP tool handlers.""" diff --git a/tests/mcp/handlers/test_episodes_handler.py b/tests/mcp/handlers/test_episodes_handler.py new file mode 100644 index 0000000..f603be7 --- /dev/null +++ b/tests/mcp/handlers/test_episodes_handler.py @@ -0,0 +1,169 @@ +"""EpisodeHandlers: lifecycle tools (start, close try/except, reconcile, list 10-field).""" +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +import pytest + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.handlers.episodes import EpisodeHandlers + + +def _stub_services() -> ServiceContainer: + services = ServiceContainer( + config=MagicMock(), + memory_conn=MagicMock(), + backend=MagicMock(), + episodes=MagicMock(), + observations=MagicMock(), + reflections=MagicMock(), + retention=MagicMock(), + memory_rating=MagicMock(), + knowledge=MagicMock(), + spool=MagicMock(), + semantic=MagicMock(), + session_bootstrap=MagicMock(), + ) + services.observations.session_id = "sess-1" + return services + + +@pytest.mark.asyncio +async def test_start_episode_returns_episode_id_reflections_and_pending_synthesis() -> None: + services = _stub_services() + services.episodes.start_foreground = MagicMock(return_value="ep-1") + queue = MagicMock() + queue.pending = 2 + queue.in_cooldown = 1 + queue.done = 5 + services.reflections.read_queue_counts = MagicMock(return_value=queue) + services.backend.retrieve = MagicMock( + return_value={"do": ["a"], "dont": [], "neutral": []}, + ) + handler = EpisodeHandlers() + result = await handler.start_episode(services, {"goal": "ship it", "tech": "py"}) + payload = json.loads(result[0].text) + assert payload == { + "episode_id": "ep-1", + "reflections": {"do": ["a"], "dont": [], "neutral": []}, + "pending_synthesis": {"pending": 2, "in_cooldown": 1, "done": 5}, + } + services.episodes.start_foreground.assert_called_once() + call_kwargs = services.episodes.start_foreground.call_args.kwargs + assert call_kwargs["session_id"] == "sess-1" + assert call_kwargs["goal"] == "ship it" + assert call_kwargs["tech"] == "py" + + +@pytest.mark.asyncio +async def test_close_episode_alt_payload_when_value_error() -> None: + """ValueError -> already_closed payload shape.""" + services = _stub_services() + services.episodes.close_active = MagicMock(side_effect=ValueError("none active")) + handler = EpisodeHandlers() + result = await handler.close_episode(services, {"outcome": "success"}) + payload = json.loads(result[0].text) + assert payload == {"closed_episode_id": None, "already_closed": True} + + +@pytest.mark.asyncio +async def test_close_episode_happy_path_payload() -> None: + services = _stub_services() + services.episodes.close_active = MagicMock(return_value="ep-42") + handler = EpisodeHandlers() + result = await handler.close_episode(services, {"outcome": "success"}) + payload = json.loads(result[0].text) + assert payload == {"closed_episode_id": "ep-42", "already_closed": False} + # Default close_reason mapping for outcome='success' is 'goal_complete'. + assert services.episodes.close_active.call_args.kwargs["close_reason"] == "goal_complete" + + +@pytest.mark.asyncio +async def test_close_episode_respects_explicit_close_reason() -> None: + services = _stub_services() + services.episodes.close_active = MagicMock(return_value="ep-7") + handler = EpisodeHandlers() + await handler.close_episode( + services, + {"outcome": "partial", "close_reason": "superseded", "summary": "later"}, + ) + kwargs = services.episodes.close_active.call_args.kwargs + assert kwargs["close_reason"] == "superseded" + assert kwargs["summary"] == "later" + assert kwargs["outcome"] == "partial" + + +@pytest.mark.asyncio +async def test_reconcile_episodes_excludes_current_session() -> None: + services = _stub_services() + ep = MagicMock() + ep.id = "ep-9" + ep.project = "p" + ep.tech = "py" + ep.goal = "g" + ep.started_at = "t0" + services.episodes.unclosed_episodes = MagicMock(return_value=[ep]) + handler = EpisodeHandlers() + result = await handler.reconcile_episodes(services, {}) + payload = json.loads(result[0].text) + assert payload == [ + { + "episode_id": "ep-9", + "project": "p", + "tech": "py", + "goal": "g", + "started_at": "t0", + } + ] + services.episodes.unclosed_episodes.assert_called_once_with( + exclude_session_ids={"sess-1"}, + ) + + +@pytest.mark.asyncio +async def test_list_episodes_emits_10_fields_per_episode() -> None: + services = _stub_services() + ep = MagicMock() + ep.id = "e1" + ep.project = "p" + ep.tech = "py" + ep.goal = "g" + ep.started_at = "t0" + ep.hardened_at = "t1" + ep.ended_at = "t2" + ep.close_reason = "r" + ep.outcome = "success" + ep.summary = "s" + services.episodes.list_episodes = MagicMock(return_value=[ep]) + handler = EpisodeHandlers() + result = await handler.list_episodes(services, {}) + payload = json.loads(result[0].text) + assert len(payload) == 1 + assert set(payload[0].keys()) == { + "episode_id", + "project", + "tech", + "goal", + "started_at", + "hardened_at", + "ended_at", + "close_reason", + "outcome", + "summary", + } + # only_open default is False; filters pass through verbatim. + services.episodes.list_episodes.assert_called_once_with( + project=None, outcome=None, only_open=False, + ) + + +def test_handlers_registers_four() -> None: + handler = EpisodeHandlers() + names = [h.name for h in handler.handlers()] + assert names == [ + "memory.start_episode", + "memory.close_episode", + "memory.reconcile_episodes", + "memory.list_episodes", + ] diff --git a/tests/mcp/handlers/test_knowledge_handler.py b/tests/mcp/handlers/test_knowledge_handler.py new file mode 100644 index 0000000..81247d2 --- /dev/null +++ b/tests/mcp/handlers/test_knowledge_handler.py @@ -0,0 +1,148 @@ +"""KnowledgeHandlers: search + list.""" +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +import pytest + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.handlers.knowledge import KnowledgeHandlers +from better_memory.services.knowledge import ( + KnowledgeDocument, + KnowledgeSearchResult, +) + + +def _stub_services() -> ServiceContainer: + return ServiceContainer( + config=MagicMock(), + memory_conn=MagicMock(), + backend=MagicMock(), + episodes=MagicMock(), + observations=MagicMock(), + reflections=MagicMock(), + retention=MagicMock(), + memory_rating=MagicMock(), + knowledge=MagicMock(), + spool=MagicMock(), + semantic=MagicMock(), + session_bootstrap=MagicMock(), + ) + + +def _make_doc( + *, + path: str = "standards/ralph-runtime.md", + scope: str = "standard", + project: str | None = None, + language: str | None = None, +) -> KnowledgeDocument: + return KnowledgeDocument( + id="doc-1", + path=path, + scope=scope, + project=project, + language=language, + content="# body", + last_indexed="2026-06-01T00:00:00+00:00", + file_mtime="2026-06-01T00:00:00+00:00", + ) + + +@pytest.mark.asyncio +async def test_search_serializes_results() -> None: + services = _stub_services() + doc = _make_doc(path="standards/foo.md", scope="standard") + services.knowledge.search = MagicMock( + return_value=[KnowledgeSearchResult(document=doc, rank=-1.5)] + ) + handler = KnowledgeHandlers() + result = await handler.search( + services, {"query": "foo", "project": "better-memory"} + ) + payload = json.loads(result[0].text) + assert payload == [ + { + "path": "standards/foo.md", + "scope": "standard", + "project": None, + "language": None, + "rank": -1.5, + } + ] + services.knowledge.search.assert_called_once_with( + "foo", project="better-memory", + ) + + +@pytest.mark.asyncio +async def test_search_forwards_missing_project_as_none() -> None: + services = _stub_services() + services.knowledge.search = MagicMock(return_value=[]) + handler = KnowledgeHandlers() + result = await handler.search(services, {"query": "bar"}) + payload = json.loads(result[0].text) + assert payload == [] + services.knowledge.search.assert_called_once_with("bar", project=None) + + +@pytest.mark.asyncio +async def test_list_serializes_docs() -> None: + services = _stub_services() + docs = [ + _make_doc( + path="standards/ralph-runtime.md", + scope="standard", + ), + _make_doc( + path="projects/better-memory/architecture.md", + scope="project", + project="better-memory", + ), + ] + services.knowledge.list_documents = MagicMock(return_value=docs) + handler = KnowledgeHandlers() + result = await handler.list(services, {"project": "better-memory"}) + payload = json.loads(result[0].text) + assert payload == [ + { + "path": "standards/ralph-runtime.md", + "scope": "standard", + "project": None, + "language": None, + }, + { + "path": "projects/better-memory/architecture.md", + "scope": "project", + "project": "better-memory", + "language": None, + }, + ] + services.knowledge.list_documents.assert_called_once_with( + project="better-memory", + ) + + +@pytest.mark.asyncio +async def test_list_forwards_missing_project_as_none() -> None: + services = _stub_services() + services.knowledge.list_documents = MagicMock(return_value=[]) + handler = KnowledgeHandlers() + result = await handler.list(services, {}) + payload = json.loads(result[0].text) + assert payload == [] + services.knowledge.list_documents.assert_called_once_with(project=None) + + +def test_handlers_registers_two() -> None: + handler = KnowledgeHandlers() + assert [h.name for h in handler.handlers()] == [ + "knowledge.search", + "knowledge.list", + ] + + +def test_handlers_are_not_capability_gated() -> None: + handler = KnowledgeHandlers() + assert all(not h.requires_synthesis for h in handler.handlers()) diff --git a/tests/mcp/handlers/test_observations_handler.py b/tests/mcp/handlers/test_observations_handler.py new file mode 100644 index 0000000..a2e7955 --- /dev/null +++ b/tests/mcp/handlers/test_observations_handler.py @@ -0,0 +1,68 @@ +"""ObservationHandlers: route the 4 observation tools + preserve invariants.""" +from __future__ import annotations + +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.handlers.observations import ObservationHandlers + + +def _stub_services() -> ServiceContainer: + return ServiceContainer( + config=MagicMock(), + memory_conn=MagicMock(), + backend=MagicMock(), + episodes=MagicMock(), + observations=MagicMock(), + reflections=MagicMock(), + retention=MagicMock(), + memory_rating=MagicMock(), + knowledge=MagicMock(), + spool=MagicMock(), + semantic=MagicMock(), + session_bootstrap=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_observe_routes_to_observations_create() -> None: + services = _stub_services() + services.observations.create = AsyncMock(return_value="obs-123") + handler = ObservationHandlers() + result = await handler.observe(services, {"content": "hello"}) + payload = json.loads(result[0].text) + assert payload == {"id": "obs-123"} + services.observations.create.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_observe_falls_back_to_project_when_scope_is_null() -> None: + services = _stub_services() + services.observations.create = AsyncMock(return_value="x") + handler = ObservationHandlers() + await handler.observe(services, {"content": "x", "scope": None}) + assert services.observations.create.call_args.kwargs["scope"] == "project" + + +@pytest.mark.asyncio +async def test_record_use_calls_service_and_returns_ok() -> None: + services = _stub_services() + services.observations.record_use = MagicMock() + handler = ObservationHandlers() + result = await handler.record_use(services, {"id": "x", "outcome": "success"}) + assert json.loads(result[0].text) == {"ok": True} + services.observations.record_use.assert_called_once_with("x", outcome="success") + + +def test_handlers_registers_four() -> None: + handler = ObservationHandlers() + names = [h.name for h in handler.handlers()] + assert names == [ + "memory.observe", + "memory.retrieve", + "memory.retrieve_observations", + "memory.record_use", + ] diff --git a/tests/mcp/handlers/test_ratings_handler.py b/tests/mcp/handlers/test_ratings_handler.py new file mode 100644 index 0000000..f6aa4a3 --- /dev/null +++ b/tests/mcp/handlers/test_ratings_handler.py @@ -0,0 +1,167 @@ +"""RatingHandlers: list_session_exposures + apply_session_ratings + credit.""" +from __future__ import annotations + +import json +from unittest.mock import MagicMock, patch + +import pytest + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.handlers.ratings import RatingHandlers + + +def _stub_services(tmp_path) -> ServiceContainer: + services = ServiceContainer( + config=MagicMock(), + memory_conn=MagicMock(), + backend=MagicMock(), + episodes=MagicMock(), + observations=MagicMock(), + reflections=MagicMock(), + retention=MagicMock(), + memory_rating=MagicMock(), + knowledge=MagicMock(), + spool=MagicMock(), + semantic=MagicMock(), + session_bootstrap=MagicMock(), + ) + services.config.home = tmp_path # type: ignore[misc] + return services + + +@pytest.mark.asyncio +async def test_apply_session_ratings_raises_multi_line_value_error_when_no_session( + tmp_path, +) -> None: + services = _stub_services(tmp_path) + handler = RatingHandlers() + with patch( + "better_memory.mcp.handlers.ratings.resolve_session_id", + return_value=None, + ): + with pytest.raises(ValueError) as exc: + await handler.apply_session_ratings(services, {"ratings": []}) + # Pin the exact 3-line error text + assert str(exc.value) == ( + "No active session: CLAUDE_SESSION_ID / " + "CLAUDE_CODE_SESSION_ID not set and no session marker " + "found (SessionStart hook may not have run)" + ) + + +@pytest.mark.asyncio +async def test_apply_session_ratings_forwards_ratings_when_session_present( + tmp_path, +) -> None: + services = _stub_services(tmp_path) + services.memory_rating.apply_session_ratings = MagicMock( + return_value={"applied": 2, "skipped": 0} + ) + handler = RatingHandlers() + with patch( + "better_memory.mcp.handlers.ratings.resolve_session_id", + return_value="sess-1", + ): + result = await handler.apply_session_ratings( + services, + {"ratings": [{"kind": "reflection", "id": "r-1", "class": "cited"}]}, + ) + payload = json.loads(result[0].text) + assert payload == {"applied": 2, "skipped": 0} + services.memory_rating.apply_session_ratings.assert_called_once_with( + session_id="sess-1", + ratings=[{"kind": "reflection", "id": "r-1", "class": "cited"}], + ) + + +@pytest.mark.asyncio +async def test_credit_returns_no_session_shape_when_unresolved(tmp_path) -> None: + services = _stub_services(tmp_path) + handler = RatingHandlers() + with patch( + "better_memory.mcp.handlers.ratings.resolve_session_id", + return_value=None, + ): + result = await handler.credit( + services, {"kind": "reflection", "id": "r-1", "class": "cited"} + ) + payload = json.loads(result[0].text) + assert payload == {"applied": None, "skipped": "no_session"} + + +@pytest.mark.asyncio +async def test_credit_calls_credit_one_when_session_present(tmp_path) -> None: + services = _stub_services(tmp_path) + services.memory_rating.credit_one = MagicMock( + return_value={"applied": "cited", "skipped": None} + ) + handler = RatingHandlers() + with patch( + "better_memory.mcp.handlers.ratings.resolve_session_id", + return_value="sess-1", + ): + result = await handler.credit( + services, {"kind": "reflection", "id": "r-1", "class": "cited"} + ) + payload = json.loads(result[0].text) + assert payload == {"applied": "cited", "skipped": None} + services.memory_rating.credit_one.assert_called_once_with( + session_id="sess-1", + kind="reflection", + id="r-1", + classification="cited", + ) + + +@pytest.mark.asyncio +async def test_list_session_exposures_coerces_none_to_empty_string( + tmp_path, +) -> None: + services = _stub_services(tmp_path) + services.session_bootstrap.list_session_exposures = MagicMock( + return_value={"session_id": None, "exposures": []} + ) + handler = RatingHandlers() + with patch( + "better_memory.mcp.handlers.ratings.resolve_session_id", + return_value=None, + ): + result = await handler.list_session_exposures(services, {}) + payload = json.loads(result[0].text) + assert payload == {"session_id": None, "exposures": []} + services.session_bootstrap.list_session_exposures.assert_called_once_with( + session_id="", + ) + + +@pytest.mark.asyncio +async def test_list_session_exposures_passes_resolved_session_id( + tmp_path, +) -> None: + services = _stub_services(tmp_path) + services.session_bootstrap.list_session_exposures = MagicMock( + return_value={"session_id": "sess-1", "exposures": []} + ) + handler = RatingHandlers() + with patch( + "better_memory.mcp.handlers.ratings.resolve_session_id", + return_value="sess-1", + ): + await handler.list_session_exposures(services, {}) + services.session_bootstrap.list_session_exposures.assert_called_once_with( + session_id="sess-1", + ) + + +def test_handlers_registers_three() -> None: + handler = RatingHandlers() + assert [h.name for h in handler.handlers()] == [ + "memory.list_session_exposures", + "memory.apply_session_ratings", + "memory.credit", + ] + + +def test_handlers_are_not_capability_gated() -> None: + handler = RatingHandlers() + assert all(not h.requires_synthesis for h in handler.handlers()) diff --git a/tests/mcp/handlers/test_reflections_handler.py b/tests/mcp/handlers/test_reflections_handler.py new file mode 100644 index 0000000..1f05dd6 --- /dev/null +++ b/tests/mcp/handlers/test_reflections_handler.py @@ -0,0 +1,146 @@ +"""ReflectionHandlers: synthesize_next_get_context + synthesize_next_apply (3 result_kinds).""" +from __future__ import annotations + +import json +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.handlers.reflections import ReflectionHandlers +from better_memory.services.reflection import SynthesisResponseError + + +def _stub_services(tmp_path) -> ServiceContainer: + services = ServiceContainer( + config=MagicMock(), + memory_conn=MagicMock(), + backend=MagicMock(), + episodes=MagicMock(), + observations=MagicMock(), + reflections=MagicMock(), + retention=MagicMock(), + memory_rating=MagicMock(), + knowledge=MagicMock(), + spool=MagicMock(), + semantic=MagicMock(), + session_bootstrap=MagicMock(), + ) + # ServiceContainer types config as Config (home is read-only on the + # real class); we substitute a MagicMock at runtime so cast through + # Any to assign home without tripping pyright's read-only check. + config: Any = services.config + config.home = tmp_path + return services + + +@pytest.mark.asyncio +async def test_get_context_returns_empty_payload_when_queue_empty(tmp_path) -> None: + services = _stub_services(tmp_path) + services.reflections.get_next_pending_context = MagicMock(return_value=None) + queue = MagicMock() + queue.pending = 0 + queue.in_cooldown = 0 + queue.done = 3 + services.reflections.read_queue_counts = MagicMock(return_value=queue) + handler = ReflectionHandlers() + result = await handler.get_context(services, {}) + payload = json.loads(result[0].text) + assert payload == { + "episode_id": None, + "queue": {"pending": 0, "in_cooldown": 0, "done": 3}, + } + + +@pytest.mark.asyncio +async def test_get_context_returns_episode_bundle_when_queue_non_empty(tmp_path) -> None: + services = _stub_services(tmp_path) + ctx = MagicMock() + ctx.episode.id = "ep-1" + ctx.episode.project = "p" + ctx.episode.goal = "g" + ctx.episode.tech = "py" + ctx.episode.outcome = "success" + ctx.observations = [] + ctx.reflections = [] + services.reflections.get_next_pending_context = MagicMock(return_value=ctx) + queue = MagicMock() + queue.pending = 1 + queue.in_cooldown = 0 + queue.done = 2 + services.reflections.read_queue_counts = MagicMock(return_value=queue) + handler = ReflectionHandlers() + result = await handler.get_context(services, {}) + payload = json.loads(result[0].text) + assert payload["episode_id"] == "ep-1" + assert payload["queue"] == {"pending": 1, "in_cooldown": 0, "done": 2} + assert payload["episode"]["id"] == "ep-1" + + +@pytest.mark.asyncio +async def test_apply_returns_validation_error_payload_on_synthesis_response_error(tmp_path) -> None: + services = _stub_services(tmp_path) + services.reflections.parse_response_dict = MagicMock( + side_effect=SynthesisResponseError("bad shape") + ) + handler = ReflectionHandlers() + result = await handler.apply(services, {"episode_id": "ep-1", "decision": {}}) + payload = json.loads(result[0].text) + assert payload == { + "ok": False, + "error": "validation", + "message": "bad shape", + } + # apply_decision must NOT have been called when parse fails. The + # underlying reflections service is a MagicMock at runtime; cast + # through Any so pyright accepts the assert_not_called attr. + apply_decision: Any = services.reflections.apply_decision + apply_decision.assert_not_called() + + +@pytest.mark.asyncio +async def test_apply_returns_state_error_payload_on_value_error(tmp_path) -> None: + services = _stub_services(tmp_path) + services.reflections.parse_response_dict = MagicMock(return_value=MagicMock()) + services.reflections.apply_decision = MagicMock(side_effect=ValueError("stale")) + handler = ReflectionHandlers() + result = await handler.apply(services, {"episode_id": "ep-1", "decision": {}}) + payload = json.loads(result[0].text) + assert payload == { + "ok": False, + "error": "state", + "message": "stale", + } + + +@pytest.mark.asyncio +async def test_apply_returns_ok_payload_on_success(tmp_path) -> None: + services = _stub_services(tmp_path) + services.reflections.parse_response_dict = MagicMock(return_value=MagicMock()) + step = MagicMock() + step.episode_id = "ep-1" + step.counts = {"new": 1, "augment": 0, "merge": 0, "ignore": 0} + step.queue.pending = 0 + step.queue.in_cooldown = 0 + step.queue.done = 5 + services.reflections.apply_decision = MagicMock(return_value=step) + handler = ReflectionHandlers() + result = await handler.apply(services, {"episode_id": "ep-1", "decision": {}}) + payload = json.loads(result[0].text) + assert payload == { + "ok": True, + "episode_id": "ep-1", + "counts": {"new": 1, "augment": 0, "merge": 0, "ignore": 0}, + "queue": {"pending": 0, "in_cooldown": 0, "done": 5}, + } + + +def test_handlers_registers_both_capability_gated() -> None: + handler = ReflectionHandlers() + handlers_list = handler.handlers() + assert [h.name for h in handlers_list] == [ + "memory.synthesize_next_get_context", + "memory.synthesize_next_apply", + ] + assert all(h.requires_synthesis for h in handlers_list) diff --git a/tests/mcp/handlers/test_retention_handler.py b/tests/mcp/handlers/test_retention_handler.py new file mode 100644 index 0000000..93ee87c --- /dev/null +++ b/tests/mcp/handlers/test_retention_handler.py @@ -0,0 +1,102 @@ +"""RetentionHandlers: memory.run_retention.""" +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +import pytest + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.handlers.retention import RetentionHandlers +from better_memory.services.retention import RetentionReport + + +def _stub_services() -> ServiceContainer: + return ServiceContainer( + config=MagicMock(), + memory_conn=MagicMock(), + backend=MagicMock(), + episodes=MagicMock(), + observations=MagicMock(), + reflections=MagicMock(), + retention=MagicMock(), + memory_rating=MagicMock(), + knowledge=MagicMock(), + spool=MagicMock(), + semantic=MagicMock(), + session_bootstrap=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_run_retention_serializes_all_four_report_fields() -> None: + services = _stub_services() + report = RetentionReport( + archived_via_retired_reflection=3, + archived_via_consumed_without_reflection=5, + archived_via_no_outcome_episode=2, + pruned=7, + ) + services.retention.run = MagicMock(return_value=report) + handler = RetentionHandlers() + result = await handler.run_retention(services, {}) + payload = json.loads(result[0].text) + assert payload == { + "archived_via_retired_reflection": 3, + "archived_via_consumed_without_reflection": 5, + "archived_via_no_outcome_episode": 2, + "pruned": 7, + } + + +@pytest.mark.asyncio +async def test_run_retention_forwards_args_with_spec_defaults() -> None: + services = _stub_services() + report = RetentionReport( + archived_via_retired_reflection=0, + archived_via_consumed_without_reflection=0, + archived_via_no_outcome_episode=0, + pruned=0, + ) + services.retention.run = MagicMock(return_value=report) + handler = RetentionHandlers() + await handler.run_retention(services, {}) + services.retention.run.assert_called_once_with( + retention_days=90, + prune=False, + prune_age_days=365, + dry_run=False, + ) + + +@pytest.mark.asyncio +async def test_run_retention_forwards_explicit_args() -> None: + services = _stub_services() + report = RetentionReport( + archived_via_retired_reflection=0, + archived_via_consumed_without_reflection=0, + archived_via_no_outcome_episode=0, + pruned=0, + ) + services.retention.run = MagicMock(return_value=report) + handler = RetentionHandlers() + await handler.run_retention( + services, + { + "retention_days": 30, + "prune": True, + "prune_age_days": 180, + "dry_run": True, + }, + ) + services.retention.run.assert_called_once_with( + retention_days=30, + prune=True, + prune_age_days=180, + dry_run=True, + ) + + +def test_handlers_registers_one_tool() -> None: + handler = RetentionHandlers() + assert [h.name for h in handler.handlers()] == ["memory.run_retention"] diff --git a/tests/mcp/handlers/test_semantics_handler.py b/tests/mcp/handlers/test_semantics_handler.py new file mode 100644 index 0000000..a8e80b9 --- /dev/null +++ b/tests/mcp/handlers/test_semantics_handler.py @@ -0,0 +1,121 @@ +"""SemanticHandlers: 4 semantic tools + scope-null fallback.""" +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +import pytest + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.handlers.semantics import SemanticHandlers + + +def _services_with_semantic(semantic: MagicMock) -> ServiceContainer: + return ServiceContainer( + config=MagicMock(), + memory_conn=MagicMock(), + backend=MagicMock(), + episodes=MagicMock(), + observations=MagicMock(), + reflections=MagicMock(), + retention=MagicMock(), + memory_rating=MagicMock(), + knowledge=MagicMock(), + spool=MagicMock(), + semantic=semantic, + session_bootstrap=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_semantic_observe_falls_back_to_project_when_scope_null() -> None: + semantic = MagicMock() + semantic.create.return_value = "sem-1" + services = _services_with_semantic(semantic) + handler = SemanticHandlers() + await handler.observe(services, {"content": "x", "scope": None}) + assert semantic.create.call_args.kwargs["scope"] == "project" + + +@pytest.mark.asyncio +async def test_semantic_observe_returns_id() -> None: + semantic = MagicMock() + semantic.create.return_value = "sem-42" + services = _services_with_semantic(semantic) + handler = SemanticHandlers() + result = await handler.observe(services, {"content": "hello"}) + assert json.loads(result[0].text) == {"id": "sem-42"} + + +@pytest.mark.asyncio +async def test_semantic_retrieve_serializes_memories() -> None: + memory = MagicMock() + memory.id = "sem-1" + memory.content = "user prefers dark mode" + memory.project = "my-project" + memory.scope = "project" + memory.created_at = "2026-06-01T00:00:00Z" + memory.updated_at = "2026-06-01T00:00:00Z" + semantic = MagicMock() + semantic.list_for_project.return_value = [memory] + services = _services_with_semantic(semantic) + handler = SemanticHandlers() + result = await handler.retrieve(services, {"project": "my-project"}) + payload = json.loads(result[0].text) + assert payload == [ + { + "id": "sem-1", + "content": "user prefers dark mode", + "project": "my-project", + "scope": "project", + "created_at": "2026-06-01T00:00:00Z", + "updated_at": "2026-06-01T00:00:00Z", + }, + ] + semantic.list_for_project.assert_called_once_with(project="my-project") + + +@pytest.mark.asyncio +async def test_semantic_retrieve_defaults_project_to_cwd_derived() -> None: + semantic = MagicMock() + semantic.list_for_project.return_value = [] + services = _services_with_semantic(semantic) + handler = SemanticHandlers() + await handler.retrieve(services, {}) + # project= kwarg present and non-empty (defaulted via project_name()). + call = semantic.list_for_project.call_args + assert "project" in call.kwargs + assert call.kwargs["project"] + + +@pytest.mark.asyncio +async def test_semantic_update_calls_update_text_and_returns_ok() -> None: + semantic = MagicMock() + services = _services_with_semantic(semantic) + handler = SemanticHandlers() + result = await handler.update( + services, {"id": "sem-1", "content": "new text"}, + ) + assert json.loads(result[0].text) == {"ok": True} + semantic.update_text.assert_called_once_with(id="sem-1", content="new text") + + +@pytest.mark.asyncio +async def test_semantic_delete_returns_ok() -> None: + semantic = MagicMock() + services = _services_with_semantic(semantic) + handler = SemanticHandlers() + result = await handler.delete(services, {"id": "sem-1"}) + assert json.loads(result[0].text) == {"ok": True} + semantic.delete.assert_called_once_with(id="sem-1") + + +def test_handlers_registers_four() -> None: + handler = SemanticHandlers() + names = [h.name for h in handler.handlers()] + assert names == [ + "memory.semantic_observe", + "memory.semantic_retrieve", + "memory.semantic_update", + "memory.semantic_delete", + ] diff --git a/tests/mcp/handlers/test_session_handler.py b/tests/mcp/handlers/test_session_handler.py new file mode 100644 index 0000000..9d912e0 --- /dev/null +++ b/tests/mcp/handlers/test_session_handler.py @@ -0,0 +1,108 @@ +"""SessionHandlers: session_bootstrap (4-tier session_id fallback, 5-key payload) + start_ui.""" +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.handlers.session import SessionHandlers + + +def _stub_services() -> ServiceContainer: + return ServiceContainer( + config=MagicMock(), + memory_conn=MagicMock(), + backend=MagicMock(), + episodes=MagicMock(), + observations=MagicMock(), + reflections=MagicMock(), + retention=MagicMock(), + memory_rating=MagicMock(), + knowledge=MagicMock(), + spool=MagicMock(), + semantic=MagicMock(), + session_bootstrap=MagicMock(), + ) + + +def _bootstrap_result() -> MagicMock: + r = MagicMock() + r.additional_context = "ctx" + r.project = "p" + r.source = "startup" + r.episode_id = "ep-1" + r.episode_action = "opened" + r.semantic_count = 5 + r.reflections_counts = {"do": 3, "dont": 1} + return r + + +@pytest.mark.asyncio +async def test_session_bootstrap_payload_has_5_top_keys() -> None: + services = _stub_services() + services.session_bootstrap.bootstrap = MagicMock( + return_value=_bootstrap_result() + ) + handler = SessionHandlers() + result = await handler.session_bootstrap(services, {}) + payload = json.loads(result[0].text) + assert set(payload.keys()) == { + "additionalContext", + "project", + "source", + "episode", + "counts", + } + assert payload["episode"] == {"id": "ep-1", "action": "opened"} + assert payload["counts"] == { + "semantic": 5, + "reflections": {"do": 3, "dont": 1}, + } + + +@pytest.mark.asyncio +async def test_session_bootstrap_uses_args_session_id_first() -> None: + services = _stub_services() + services.session_bootstrap.bootstrap = MagicMock( + return_value=_bootstrap_result() + ) + handler = SessionHandlers() + await handler.session_bootstrap(services, {"session_id": "from-arg"}) + call_kwargs = services.session_bootstrap.bootstrap.call_args.kwargs + assert call_kwargs["session_id"] == "from-arg" + + +@pytest.mark.asyncio +async def test_session_bootstrap_uses_cwd_arg_when_provided() -> None: + services = _stub_services() + services.session_bootstrap.bootstrap = MagicMock( + return_value=_bootstrap_result() + ) + handler = SessionHandlers() + await handler.session_bootstrap(services, {"cwd": "/tmp/foo"}) + call_kwargs = services.session_bootstrap.bootstrap.call_args.kwargs + assert call_kwargs["cwd"] == Path("/tmp/foo") + + +@pytest.mark.asyncio +async def test_start_ui_returns_launcher_result() -> None: + services = _stub_services() + handler = SessionHandlers() + with patch( + "better_memory.mcp.handlers.session.ui_launcher.start_ui", + return_value={"url": "http://localhost:8000", "reused": False}, + ): + result = await handler.start_ui(services, {}) + payload = json.loads(result[0].text) + assert payload == {"url": "http://localhost:8000", "reused": False} + + +def test_handlers_registers_two() -> None: + handler = SessionHandlers() + assert [h.name for h in handler.handlers()] == [ + "memory.session_bootstrap", + "memory.start_ui", + ] diff --git a/tests/mcp/test_best_effort_logging.py b/tests/mcp/test_best_effort_logging.py index c0f9c33..373cd0d 100644 --- a/tests/mcp/test_best_effort_logging.py +++ b/tests/mcp/test_best_effort_logging.py @@ -14,7 +14,7 @@ import pytest -from better_memory.mcp.server import _run_best_effort +from better_memory.mcp._best_effort import _run_best_effort def test_run_best_effort_swallows_exception(caplog: pytest.LogCaptureFixture) -> None: @@ -23,7 +23,7 @@ def test_run_best_effort_swallows_exception(caplog: pytest.LogCaptureFixture) -> def boom() -> None: raise RuntimeError("kaboom") - with caplog.at_level(logging.ERROR, logger="better_memory.mcp.server"): + with caplog.at_level(logging.ERROR, logger="better_memory.mcp._best_effort"): _run_best_effort("spool.drain", boom) # must not raise @@ -35,12 +35,12 @@ def test_run_best_effort_logs_exception_with_traceback( def boom() -> None: raise RuntimeError("kaboom") - with caplog.at_level(logging.ERROR, logger="better_memory.mcp.server"): + with caplog.at_level(logging.ERROR, logger="better_memory.mcp._best_effort"): _run_best_effort("spool.drain", boom) matching = [ r for r in caplog.records - if r.name == "better_memory.mcp.server" + if r.name == "better_memory.mcp._best_effort" and r.levelno == logging.ERROR and "spool.drain" in r.getMessage() ] diff --git a/tests/mcp/test_service_container.py b/tests/mcp/test_service_container.py new file mode 100644 index 0000000..5e7ae0f --- /dev/null +++ b/tests/mcp/test_service_container.py @@ -0,0 +1,111 @@ +"""ServiceContainer bundles every long-lived service the MCP dispatcher needs.""" +from __future__ import annotations + +import dataclasses +import sqlite3 +from dataclasses import is_dataclass +from unittest.mock import MagicMock + +import pytest + +from better_memory.mcp.container import ServiceContainer + + +def _make_container(memory_conn: sqlite3.Connection | None = None) -> ServiceContainer: + return ServiceContainer( + config=MagicMock(), + memory_conn=memory_conn or MagicMock(spec=sqlite3.Connection), + backend=MagicMock(), + episodes=MagicMock(), + observations=MagicMock(), + reflections=MagicMock(), + retention=MagicMock(), + memory_rating=MagicMock(), + knowledge=MagicMock(), + spool=MagicMock(), + semantic=MagicMock(), + session_bootstrap=MagicMock(), + ) + + +def test_service_container_is_frozen_dataclass() -> None: + fields = { + "config", "memory_conn", "backend", + "episodes", "observations", "reflections", + "retention", "memory_rating", "knowledge", + "spool", "semantic", "session_bootstrap", + } + assert is_dataclass(ServiceContainer) + assert set(ServiceContainer.__dataclass_fields__) == fields + + +def test_service_container_holds_attributes() -> None: + mock_conn = MagicMock(spec=sqlite3.Connection) + container = _make_container(memory_conn=mock_conn) + assert container.memory_conn is mock_conn + + +def test_service_container_rejects_attribute_reassignment() -> None: + container = _make_container() + with pytest.raises(dataclasses.FrozenInstanceError): + container.memory_conn = MagicMock(spec=sqlite3.Connection) # type: ignore[misc] + + +def test_build_services_constructs_each_service_exactly_once( + monkeypatch, tmp_path, +) -> None: + """Regression guard: SemanticMemoryService was built 4× per call, + SessionBootstrapService 2× per call. Container must build each once.""" + from collections import Counter + from pathlib import Path + + import better_memory.services.semantic as _sem_mod + import better_memory.services.session_bootstrap as _sb_mod + from better_memory.config import get_config + from better_memory.db.connection import connect + from better_memory.db.schema import apply_migrations + from better_memory.mcp.server import _build_services + + counts: Counter[str] = Counter() + + def _wrap(cls: type) -> None: + original_init = cls.__init__ + def _init(self, *a, **kw): + counts[cls.__name__] += 1 + original_init(self, *a, **kw) + monkeypatch.setattr(cls, "__init__", _init) + + _wrap(_sem_mod.SemanticMemoryService) + _wrap(_sb_mod.SessionBootstrapService) + + monkeypatch.setenv("BETTER_MEMORY_HOME", str(tmp_path)) + cfg = get_config() + + mem_conn = connect(cfg.memory_db) + apply_migrations( + mem_conn, + migrations_dir=Path(__file__).parent.parent.parent + / "better_memory" / "db" / "migrations", + ) + kb_conn = connect(cfg.knowledge_db) + apply_migrations( + kb_conn, + migrations_dir=Path(__file__).parent.parent.parent + / "better_memory" / "db" / "knowledge_migrations", + ) + + container = _build_services( + cfg, mem_conn, kb_conn, embedder=None, + startup_project="test", startup_session_id=None, + ) + # _build_services itself constructs each service exactly once at the + # container level. SqliteBackend (built inside build_backend) keeps its + # own internal SemanticMemoryService + SessionBootstrapService for + # protocol compliance — that accounts for the second construction of + # each. Total of 2 still kills the 4×/2× inline regressions; drift back + # to per-call construction inside `_build_services` will push these + # well above 2 and re-trip this test. + assert counts["SemanticMemoryService"] == 2 + assert counts["SessionBootstrapService"] == 2 + assert container.semantic is not None + assert container.session_bootstrap is not None diff --git a/tests/mcp/test_session_resolver.py b/tests/mcp/test_session_resolver.py new file mode 100644 index 0000000..86c2f88 --- /dev/null +++ b/tests/mcp/test_session_resolver.py @@ -0,0 +1,34 @@ +"""resolve_session_id resolves Claude Code session id with the documented fallback.""" +from __future__ import annotations + +from pathlib import Path +from unittest.mock import patch + +from better_memory.mcp._session import resolve_session_id + + +def test_resolve_session_id_prefers_env(monkeypatch, tmp_path: Path) -> None: + monkeypatch.setenv("CLAUDE_SESSION_ID", "from-env") + assert resolve_session_id(tmp_path) == "from-env" + + +def test_resolve_session_id_falls_back_to_alt_env(monkeypatch, tmp_path: Path) -> None: + monkeypatch.delenv("CLAUDE_SESSION_ID", raising=False) + monkeypatch.setenv("CLAUDE_CODE_SESSION_ID", "alt-env") + assert resolve_session_id(tmp_path) == "alt-env" + + +def test_resolve_session_id_falls_back_to_marker(monkeypatch, tmp_path: Path) -> None: + monkeypatch.delenv("CLAUDE_SESSION_ID", raising=False) + monkeypatch.delenv("CLAUDE_CODE_SESSION_ID", raising=False) + with patch("better_memory.mcp._session.read_session_id", return_value="marker"): + assert resolve_session_id(tmp_path) == "marker" + + +def test_resolve_session_id_returns_none_when_all_absent( + monkeypatch, tmp_path: Path, +) -> None: + monkeypatch.delenv("CLAUDE_SESSION_ID", raising=False) + monkeypatch.delenv("CLAUDE_CODE_SESSION_ID", raising=False) + with patch("better_memory.mcp._session.read_session_id", return_value=None): + assert resolve_session_id(tmp_path) is None diff --git a/tests/mcp/test_synth_audit_log.py b/tests/mcp/test_synth_audit_log.py index c9760a2..b90a7ed 100644 --- a/tests/mcp/test_synth_audit_log.py +++ b/tests/mcp/test_synth_audit_log.py @@ -15,7 +15,7 @@ import pytest -from better_memory.mcp.server import _append_synth_audit, _audit_synth_call +from better_memory.mcp.handlers._audit import _append_synth_audit, _audit_synth_call def _read_jsonl(home: Path) -> list[dict]: diff --git a/tests/mcp/test_synthesize_tools.py b/tests/mcp/test_synthesize_tools.py index 96d2024..a0bedb2 100644 --- a/tests/mcp/test_synthesize_tools.py +++ b/tests/mcp/test_synthesize_tools.py @@ -337,7 +337,7 @@ def test_empty_queue_yields_null_episode_payload( ): svc = ReflectionSynthesisService(conn, clock=fixed_clock) ctx = svc.get_next_pending_context(project="p1") - queue = svc._read_queue_counts(project="p1") + queue = svc.read_queue_counts(project="p1") payload = _serialize_synth_get_context(ctx, queue) assert payload == { "episode_id": None, @@ -351,7 +351,7 @@ def test_pending_episode_yields_full_payload(self, conn, fixed_clock): svc = ReflectionSynthesisService(conn, clock=fixed_clock) ctx = svc.get_next_pending_context(project="p1") - queue = svc._read_queue_counts(project="p1") + queue = svc.read_queue_counts(project="p1") assert ctx is not None payload = _serialize_synth_get_context(ctx, queue) @@ -374,7 +374,7 @@ def test_oldest_episode_first(self, conn, fixed_clock): conn.commit() svc = ReflectionSynthesisService(conn, clock=fixed_clock) ctx = svc.get_next_pending_context(project="p1") - queue = svc._read_queue_counts(project="p1") + queue = svc.read_queue_counts(project="p1") payload = _serialize_synth_get_context(ctx, queue) assert payload["episode_id"] == "older" diff --git a/tests/mcp/test_tool_dispatcher.py b/tests/mcp/test_tool_dispatcher.py new file mode 100644 index 0000000..90e54ad --- /dev/null +++ b/tests/mcp/test_tool_dispatcher.py @@ -0,0 +1,94 @@ +"""ToolDispatcher: register, list, call, capability-gate.""" +from __future__ import annotations + +from dataclasses import FrozenInstanceError +from typing import Any +from unittest.mock import MagicMock + +import pytest +from mcp.types import TextContent + +from better_memory.mcp.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler, ToolDispatcher + + +def _container(*, supports_synthesis: bool) -> ServiceContainer: + backend = MagicMock() + backend.supports_synthesis = supports_synthesis + return ServiceContainer( + config=MagicMock(), memory_conn=MagicMock(), backend=backend, + episodes=MagicMock(), observations=MagicMock(), reflections=MagicMock(), + retention=MagicMock(), memory_rating=MagicMock(), knowledge=MagicMock(), + spool=MagicMock(), semantic=MagicMock(), session_bootstrap=MagicMock(), + ) + + +async def _stub_call( + services: ServiceContainer, args: dict[str, Any], +) -> list[TextContent]: + return [] + + +def test_handler_is_frozen_dataclass() -> None: + h = Handler(name="x", schema={"type": "object"}, call=_stub_call) + assert h.name == "x" + assert h.requires_synthesis is False + with pytest.raises(FrozenInstanceError): + h.name = "y" # type: ignore[misc] + + +def test_dispatcher_lists_all_when_synthesis_supported() -> None: + services = _container(supports_synthesis=True) + handlers = [ + Handler("a", {}, _stub_call), + Handler("b", {}, _stub_call, requires_synthesis=True), + ] + dispatcher = ToolDispatcher(services, handlers) + names = [t.name for t in dispatcher.tool_definitions()] + assert names == ["a", "b"] + + +def test_dispatcher_hides_synthesis_tools_when_unsupported() -> None: + services = _container(supports_synthesis=False) + handlers = [ + Handler("a", {}, _stub_call), + Handler("b", {}, _stub_call, requires_synthesis=True), + ] + dispatcher = ToolDispatcher(services, handlers) + names = [t.name for t in dispatcher.tool_definitions()] + assert names == ["a"] + + +@pytest.mark.asyncio +async def test_dispatcher_call_unknown_raises_value_error() -> None: + services = _container(supports_synthesis=True) + dispatcher = ToolDispatcher(services, []) + with pytest.raises(ValueError, match="Unknown tool: nope"): + await dispatcher.call("nope", {}) + + +@pytest.mark.asyncio +async def test_dispatcher_call_gated_tool_when_unsupported_raises_unknown() -> None: + services = _container(supports_synthesis=False) + handler = Handler("b", {}, _stub_call, requires_synthesis=True) + dispatcher = ToolDispatcher(services, [handler]) + with pytest.raises(ValueError, match="Unknown tool: b"): + await dispatcher.call("b", {}) + + +@pytest.mark.asyncio +async def test_dispatcher_call_routes_to_handler() -> None: + services = _container(supports_synthesis=True) + seen: dict[str, Any] = {} + async def recording_call( + svc: ServiceContainer, args: dict[str, Any], + ) -> list[TextContent]: + seen["svc"] = svc + seen["args"] = args + return [TextContent(type="text", text="ok")] + handler = Handler("x", {"type": "object"}, recording_call) + dispatcher = ToolDispatcher(services, [handler]) + result = await dispatcher.call("x", {"k": 1}) + assert seen["svc"] is services + assert seen["args"] == {"k": 1} + assert len(result) == 1 diff --git a/tests/services/test_reflection.py b/tests/services/test_reflection.py index f87b829..75f9a2f 100644 --- a/tests/services/test_reflection.py +++ b/tests/services/test_reflection.py @@ -1758,7 +1758,7 @@ def test_counts_zero_for_empty_project(self, conn, fixed_clock): svc = ReflectionSynthesisService( conn, clock=fixed_clock, ) - counts = svc._read_queue_counts(project="empty") + counts = svc.read_queue_counts(project="empty") assert counts.done == 0 assert counts.pending == 0 assert counts.in_cooldown == 0 @@ -1791,7 +1791,7 @@ def test_counts_done_pending_and_cooldown_separately( svc = ReflectionSynthesisService( conn, clock=fixed_clock, ) - counts = svc._read_queue_counts(project="p1") + counts = svc.read_queue_counts(project="p1") assert counts.done == 2 assert counts.pending == 3 assert counts.in_cooldown == 1 @@ -1806,7 +1806,7 @@ def test_excludes_open_episodes_from_total(self, conn, fixed_clock): svc = ReflectionSynthesisService( conn, clock=fixed_clock, ) - counts = svc._read_queue_counts(project="p1") + counts = svc.read_queue_counts(project="p1") assert counts.total == 0