From 22c3d2606d04652ca7d42ae2cd4342413d0439e6 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 14:01:53 +0100 Subject: [PATCH 01/20] docs(mcp): spec + plan for server.py dispatcher refactor Spec: extract _call_tool if-chain into ToolDispatcher + ServiceContainer + per-domain handler modules. 17 tasks, single big-bang PR, ~94% avg confidence with all sub-90% tasks lifted via parallel investigation. Co-Authored-By: Claude Opus 4.7 --- ...26-06-13-mcp-server-dispatcher-refactor.md | 1860 +++++++++++++++++ ...2026-06-12-mcp-server-dispatcher-design.md | 308 +++ 2 files changed, 2168 insertions(+) create mode 100644 docs/superpowers/plans/2026-06-13-mcp-server-dispatcher-refactor.md create mode 100644 docs/superpowers/specs/2026-06-12-mcp-server-dispatcher-design.md 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). From cb33488ec429d9a7e7111f69a4b43cfbd9b1a252 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 14:07:05 +0100 Subject: [PATCH 02/20] 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). --- better_memory/mcp/container.py | 40 +++++++++++++++++++++++++++++ tests/mcp/test_service_container.py | 38 +++++++++++++++++++++++++++ 2 files changed, 78 insertions(+) create mode 100644 better_memory/mcp/container.py create mode 100644 tests/mcp/test_service_container.py diff --git a/better_memory/mcp/container.py b/better_memory/mcp/container.py new file mode 100644 index 0000000..c3b027c --- /dev/null +++ b/better_memory/mcp/container.py @@ -0,0 +1,40 @@ +"""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" diff --git a/tests/mcp/test_service_container.py b/tests/mcp/test_service_container.py new file mode 100644 index 0000000..0b368e4 --- /dev/null +++ b/tests/mcp/test_service_container.py @@ -0,0 +1,38 @@ +"""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 From 60ea7d62f76da4f5aa1a1d077a8173079f4dcc50 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 14:11:42 +0100 Subject: [PATCH 03/20] fix(mcp): address code-review on ServiceContainer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Strip unnecessary quote-string forward refs (ruff UP037) — the `from __future__ import annotations` import already defers evaluation, so quoting is redundant. - Restore `semantic: SemanticMemoryService` type (was `Any`) by adding the import to the TYPE_CHECKING block alongside the other services. No runtime circular-import risk: container.py is only imported by the MCP dispatcher and services/semantic.py only depends on services/memory_rating. - Strengthen frozen-dataclass test: assert `dataclasses.FrozenInstanceError` is raised on attribute reassignment, in addition to the existing structural checks. Co-Authored-By: Claude Opus 4.7 --- better_memory/mcp/container.py | 25 +++++++++--------- tests/mcp/test_service_container.py | 41 +++++++++++++++++++---------- 2 files changed, 40 insertions(+), 26 deletions(-) diff --git a/better_memory/mcp/container.py b/better_memory/mcp/container.py index c3b027c..ac4ceba 100644 --- a/better_memory/mcp/container.py +++ b/better_memory/mcp/container.py @@ -8,7 +8,7 @@ import sqlite3 from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING if TYPE_CHECKING: from better_memory.config import Config @@ -18,6 +18,7 @@ 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 @@ -26,15 +27,15 @@ @dataclass(frozen=True) class ServiceContainer: """All long-lived services + connections, built once in create_server.""" - config: "Config" + 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" + 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/tests/mcp/test_service_container.py b/tests/mcp/test_service_container.py index 0b368e4..5007475 100644 --- a/tests/mcp/test_service_container.py +++ b/tests/mcp/test_service_container.py @@ -1,13 +1,33 @@ """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", @@ -21,18 +41,11 @@ def test_service_container_is_frozen_dataclass() -> None: 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(), - ) + 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] From e9deac6f732ae9f68f7f3bd2b17b7a2c2be2433a Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 14:15:53 +0100 Subject: [PATCH 04/20] 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. --- better_memory/mcp/dispatcher.py | 60 ++++++++++++++++++++ tests/mcp/test_tool_dispatcher.py | 94 +++++++++++++++++++++++++++++++ 2 files changed, 154 insertions(+) create mode 100644 better_memory/mcp/dispatcher.py create mode 100644 tests/mcp/test_tool_dispatcher.py 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/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 From fe1696243913c252a1fbac62fb43ec13cff37b20 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 14:24:19 +0100 Subject: [PATCH 05/20] 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. --- better_memory/mcp/_session.py | 28 ++++++++++++++++++++++++ better_memory/mcp/server.py | 28 ++++++------------------ tests/mcp/test_session_resolver.py | 34 ++++++++++++++++++++++++++++++ 3 files changed, 68 insertions(+), 22 deletions(-) create mode 100644 better_memory/mcp/_session.py create mode 100644 tests/mcp/test_session_resolver.py 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/server.py b/better_memory/mcp/server.py index d1422aa..b9a81c7 100644 --- a/better_memory/mcp/server.py +++ b/better_memory/mcp/server.py @@ -61,7 +61,7 @@ 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.mcp._session import resolve_session_id from better_memory.services import ui_launcher from better_memory.services.episode import EpisodeService from better_memory.services.knowledge import ( @@ -69,6 +69,7 @@ KnowledgeSearchResult, KnowledgeService, ) +from better_memory.services.memory_rating import MemoryRatingService from better_memory.services.observation import ObservationService from better_memory.services.reflection import ( EpisodeContext, @@ -76,7 +77,6 @@ ReflectionSynthesisService, SynthesisStep, ) -from better_memory.services.memory_rating import MemoryRatingService from better_memory.services.retention import RetentionService from better_memory.services.retention_scheduler import RetentionScheduler from better_memory.services.spool import SpoolService @@ -144,22 +144,6 @@ def _run_best_effort( # 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: @@ -1011,7 +995,7 @@ def create_server() -> tuple[ # 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 + startup_session_id: str | None = resolve_session_id(config.home) or None backend: StorageBackend = build_backend( config=config, memory_conn=memory_conn, @@ -1489,14 +1473,14 @@ def _retention_step() -> None: SessionBootstrapService, ) - sid = _resolve_session_id(config.home) or "" + 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) + sid = resolve_session_id(config.home) if not sid: raise ValueError( "No active session: CLAUDE_SESSION_ID / " @@ -1510,7 +1494,7 @@ def _retention_step() -> None: return [TextContent(type="text", text=json.dumps(payload))] if name == "memory.credit": - sid = _resolve_session_id(config.home) + sid = resolve_session_id(config.home) if not sid: payload = {"applied": None, "skipped": "no_session"} else: 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 From 3e9e9d4e3ba470c5d02a7c5d49adedbc28b06179 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 14:33:06 +0100 Subject: [PATCH 06/20] refactor(mcp): move audit helpers to mcp/handlers/_audit.py Audit JSONL byte-shape unchanged. Test updates import path only. --- better_memory/mcp/handlers/__init__.py | 18 +++++ better_memory/mcp/handlers/_audit.py | 94 ++++++++++++++++++++++++++ better_memory/mcp/server.py | 83 +---------------------- tests/mcp/test_synth_audit_log.py | 2 +- 4 files changed, 115 insertions(+), 82 deletions(-) create mode 100644 better_memory/mcp/handlers/__init__.py create mode 100644 better_memory/mcp/handlers/_audit.py diff --git a/better_memory/mcp/handlers/__init__.py b/better_memory/mcp/handlers/__init__.py new file mode 100644 index 0000000..4e43269 --- /dev/null +++ b/better_memory/mcp/handlers/__init__.py @@ -0,0 +1,18 @@ +"""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 [] 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/server.py b/better_memory/mcp/server.py index b9a81c7..15fd3ca 100644 --- a/better_memory/mcp/server.py +++ b/better_memory/mcp/server.py @@ -36,7 +36,6 @@ from __future__ import annotations import asyncio -import contextlib import json import logging import os @@ -46,9 +45,8 @@ 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 @@ -62,6 +60,7 @@ from better_memory.db.schema import apply_migrations from better_memory.embeddings.ollama import OllamaEmbedder from better_memory.mcp._session import resolve_session_id +from better_memory.mcp.handlers._audit import _audit_synth_call from better_memory.services import ui_launcher from better_memory.services.episode import EpisodeService from better_memory.services.knowledge import ( @@ -127,84 +126,6 @@ def _run_best_effort( ) -# --------------------------------------------------------- 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 _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. 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]: From 0e088e3467874a9830ceee9dbccab6e07d717a66 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 14:38:46 +0100 Subject: [PATCH 07/20] 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. Co-Authored-By: Claude Opus 4.7 --- better_memory/mcp/server.py | 4 ++-- better_memory/services/reflection.py | 6 +++--- tests/mcp/test_synthesize_tools.py | 6 +++--- tests/services/test_reflection.py | 6 +++--- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/better_memory/mcp/server.py b/better_memory/mcp/server.py index 15fd3ca..a82a7c9 100644 --- a/better_memory/mcp/server.py +++ b/better_memory/mcp/server.py @@ -1165,7 +1165,7 @@ def _retention_step() -> None: # 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) + queue = reflections.read_queue_counts(project=project) buckets = backend.retrieve( project=project, tech=args.get("tech"), ) @@ -1324,7 +1324,7 @@ def _retention_step() -> None: episode_id=None, ) as audit: ctx = reflections.get_next_pending_context(project=project) - queue = reflections._read_queue_counts(project=project) + queue = reflections.read_queue_counts(project=project) payload = _serialize_synth_get_context(ctx, queue) if ctx is None: audit["result_kind"] = "empty" 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/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/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 From c79760a8261266b988409a896448c5592ff8c18c Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 14:47:34 +0100 Subject: [PATCH 08/20] 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. Co-Authored-By: Claude Opus 4.7 --- better_memory/mcp/server.py | 94 +++++++++++++++++++++-------- tests/mcp/test_service_container.py | 60 ++++++++++++++++++ 2 files changed, 130 insertions(+), 24 deletions(-) diff --git a/better_memory/mcp/server.py b/better_memory/mcp/server.py index a82a7c9..692134a 100644 --- a/better_memory/mcp/server.py +++ b/better_memory/mcp/server.py @@ -60,6 +60,7 @@ from better_memory.db.schema import apply_migrations from better_memory.embeddings.ollama import OllamaEmbedder from better_memory.mcp._session import resolve_session_id +from better_memory.mcp.container import ServiceContainer from better_memory.mcp.handlers._audit import _audit_synth_call from better_memory.services import ui_launcher from better_memory.services.episode import EpisodeService @@ -855,6 +856,61 @@ class ServerContext: embedder: Any = 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 create_server() -> tuple[ Server, Callable[[], Coroutine[Any, Any, None]], @@ -902,9 +958,7 @@ def create_server() -> tuple[ # 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 @@ -915,28 +969,20 @@ 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) + # 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 # Session-start behaviour: reindex knowledge at startup. mtime-only, so # the cost is O(files) stat calls on an already-indexed corpus. We diff --git a/tests/mcp/test_service_container.py b/tests/mcp/test_service_container.py index 5007475..d25b2ab 100644 --- a/tests/mcp/test_service_container.py +++ b/tests/mcp/test_service_container.py @@ -49,3 +49,63 @@ 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) + cls.__init__ = _init # type: ignore[method-assign] + + _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; any + # future drift back to per-call construction 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 From ded513187e7ff3ae6fe5cd053aa6aa5d5f4b220d Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 14:56:02 +0100 Subject: [PATCH 09/20] test(mcp): use monkeypatch.setattr to revert __init__ wrapper on teardown The _wrap helper inside test_build_services_constructs_each_service_exactly_once was permanently mutating SemanticMemoryService.__init__ and SessionBootstrapService.__init__ via direct cls.__init__ = _init assignment. Any test running after this one continued executing the wrapped __init__, a slow leak that monkeypatch was already in scope to prevent. Switch to monkeypatch.setattr so pytest reverts the class patch on teardown. Also tighten the assertion comment to scope the regression statement to _build_services (the only thing this test exercises). Co-Authored-By: Claude Opus 4.7 --- tests/mcp/test_service_container.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/mcp/test_service_container.py b/tests/mcp/test_service_container.py index d25b2ab..5e7ae0f 100644 --- a/tests/mcp/test_service_container.py +++ b/tests/mcp/test_service_container.py @@ -73,7 +73,7 @@ def _wrap(cls: type) -> None: def _init(self, *a, **kw): counts[cls.__name__] += 1 original_init(self, *a, **kw) - cls.__init__ = _init # type: ignore[method-assign] + monkeypatch.setattr(cls, "__init__", _init) _wrap(_sem_mod.SemanticMemoryService) _wrap(_sb_mod.SessionBootstrapService) @@ -102,9 +102,9 @@ def _init(self, *a, **kw): # 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; any - # future drift back to per-call construction will push these well - # above 2 and re-trip this test. + # 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 From dee7641e8e8504acf43d3a9056065c3ad98b71db Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 15:01:31 +0100 Subject: [PATCH 10/20] feat(mcp): observation-domain handlers (memory.observe + retrieve + retrieve_observations + record_use) Schemas + bodies lifted verbatim from _call_tool. spool->retention->backend ordering preserved. ObservationHandlers exposes one method per tool + handlers() registration list for ToolDispatcher. Co-Authored-By: Claude Opus 4.7 --- better_memory/mcp/handlers/__init__.py | 5 +- better_memory/mcp/handlers/observations.py | 279 ++++++++++++++++++ tests/mcp/handlers/__init__.py | 1 + .../mcp/handlers/test_observations_handler.py | 68 +++++ 4 files changed, 352 insertions(+), 1 deletion(-) create mode 100644 better_memory/mcp/handlers/observations.py create mode 100644 tests/mcp/handlers/__init__.py create mode 100644 tests/mcp/handlers/test_observations_handler.py diff --git a/better_memory/mcp/handlers/__init__.py b/better_memory/mcp/handlers/__init__.py index 4e43269..4a58c12 100644 --- a/better_memory/mcp/handlers/__init__.py +++ b/better_memory/mcp/handlers/__init__.py @@ -8,6 +8,7 @@ from __future__ import annotations from better_memory.mcp.dispatcher import Handler +from better_memory.mcp.handlers.observations import ObservationHandlers def all_handlers() -> list[Handler]: @@ -15,4 +16,6 @@ def all_handlers() -> list[Handler]: Filled in as handler modules land (Tasks 7-14). """ - return [] + return [ + *ObservationHandlers().handlers(), + ] diff --git a/better_memory/mcp/handlers/observations.py b/better_memory/mcp/handlers/observations.py new file mode 100644 index 0000000..0338388 --- /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.container import ServiceContainer +from better_memory.mcp.dispatcher import Handler +from better_memory.mcp.server import _run_best_effort +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/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_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", + ] From c12ea45f47ed2dc974a9e6244ee863d7e884920c Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 15:09:49 +0100 Subject: [PATCH 11/20] feat(mcp): semantic-domain handlers (memory.semantic_observe + retrieve + update + delete) SemanticMemoryService is now built ONCE on the container, not 4x per call. --- better_memory/mcp/handlers/__init__.py | 2 + better_memory/mcp/handlers/semantics.py | 160 +++++++++++++++++++ tests/mcp/handlers/test_semantics_handler.py | 121 ++++++++++++++ 3 files changed, 283 insertions(+) create mode 100644 better_memory/mcp/handlers/semantics.py create mode 100644 tests/mcp/handlers/test_semantics_handler.py diff --git a/better_memory/mcp/handlers/__init__.py b/better_memory/mcp/handlers/__init__.py index 4a58c12..ba608ad 100644 --- a/better_memory/mcp/handlers/__init__.py +++ b/better_memory/mcp/handlers/__init__.py @@ -9,6 +9,7 @@ from better_memory.mcp.dispatcher import Handler from better_memory.mcp.handlers.observations import ObservationHandlers +from better_memory.mcp.handlers.semantics import SemanticHandlers def all_handlers() -> list[Handler]: @@ -18,4 +19,5 @@ def all_handlers() -> list[Handler]: """ return [ *ObservationHandlers().handlers(), + *SemanticHandlers().handlers(), ] 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/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", + ] From 1ceed718c1f56315328e431a4ef9e215379282d5 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 15:19:06 +0100 Subject: [PATCH 12/20] feat(mcp): episode-lifecycle handlers (start + close + reconcile + list) Task 9 of the MCP server dispatcher refactor. Lifts the four episode lifecycle tools out of the legacy _call_tool if-chain into a domain handler module. Bodies are verbatim, including the close_episode try/except fork that converts EpisodeService.close_active's ValueError into the already_closed=True no-op payload, and the 10-field list_episodes serializer the UI depends on. _serialize_queue is duplicated as a small module-local helper to keep the handler module self-contained. After this task, all_handlers() returns 12 (4 obs + 4 sem + 4 ep). Co-Authored-By: Claude Opus 4.7 --- better_memory/mcp/handlers/__init__.py | 2 + better_memory/mcp/handlers/episodes.py | 257 ++++++++++++++++++++ tests/mcp/handlers/test_episodes_handler.py | 169 +++++++++++++ 3 files changed, 428 insertions(+) create mode 100644 better_memory/mcp/handlers/episodes.py create mode 100644 tests/mcp/handlers/test_episodes_handler.py diff --git a/better_memory/mcp/handlers/__init__.py b/better_memory/mcp/handlers/__init__.py index ba608ad..27d449e 100644 --- a/better_memory/mcp/handlers/__init__.py +++ b/better_memory/mcp/handlers/__init__.py @@ -8,6 +8,7 @@ from __future__ import annotations from better_memory.mcp.dispatcher import Handler +from better_memory.mcp.handlers.episodes import EpisodeHandlers from better_memory.mcp.handlers.observations import ObservationHandlers from better_memory.mcp.handlers.semantics import SemanticHandlers @@ -20,4 +21,5 @@ def all_handlers() -> list[Handler]: return [ *ObservationHandlers().handlers(), *SemanticHandlers().handlers(), + *EpisodeHandlers().handlers(), ] 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/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", + ] From 550cdd9c5a789c299d10e3424c0b1cae2e6ef810 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 15:29:20 +0100 Subject: [PATCH 13/20] feat(mcp): synthesis handlers (synthesize_next_get_context + apply) Adds ReflectionHandlers covering the two capability-gated synthesis tools. Bodies lifted verbatim from the legacy _call_tool if-chain to preserve every documented invariant: audit bracketing on both tools, the SynthesisResponseError -> validation_error fork (no synth_failed_at stamp; caller can retry), the ValueError -> state_error fork (caller must refetch context), and the four serializer payload shapes the IDE-LLM consumes. Both handlers are marked requires_synthesis=True so the dispatcher hides them when backend.supports_synthesis is False. all_handlers() now returns 14 tools (4 obs + 4 sem + 4 ep + 2 ref). Task 10 of 17 of the MCP dispatcher refactor. Co-Authored-By: Claude Opus 4.7 --- better_memory/mcp/handlers/__init__.py | 2 + better_memory/mcp/handlers/reflections.py | 283 ++++++++++++++++++ .../mcp/handlers/test_reflections_handler.py | 146 +++++++++ 3 files changed, 431 insertions(+) create mode 100644 better_memory/mcp/handlers/reflections.py create mode 100644 tests/mcp/handlers/test_reflections_handler.py diff --git a/better_memory/mcp/handlers/__init__.py b/better_memory/mcp/handlers/__init__.py index 27d449e..663dc9e 100644 --- a/better_memory/mcp/handlers/__init__.py +++ b/better_memory/mcp/handlers/__init__.py @@ -10,6 +10,7 @@ from better_memory.mcp.dispatcher import Handler from better_memory.mcp.handlers.episodes import EpisodeHandlers from better_memory.mcp.handlers.observations import ObservationHandlers +from better_memory.mcp.handlers.reflections import ReflectionHandlers from better_memory.mcp.handlers.semantics import SemanticHandlers @@ -22,4 +23,5 @@ def all_handlers() -> list[Handler]: *ObservationHandlers().handlers(), *SemanticHandlers().handlers(), *EpisodeHandlers().handlers(), + *ReflectionHandlers().handlers(), ] 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/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) From ef8d299824b77b20c9cbbb21cd9a7570d3447754 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 15:36:16 +0100 Subject: [PATCH 14/20] feat(mcp): retention handler (memory.run_retention) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 11 of MCP server dispatcher refactor. Lifts the memory.run_retention branch out of the _call_tool if-chain into RetentionHandlers, registered via handlers/__init__.all_handlers(). Body and schema are verbatim from server.py; `retention.run(...)` is re-routed through `services.retention.run(...)` on the ServiceContainer. The serialized payload is the 4-field RetentionReport (three archived_via_* counters + pruned) — not 8 as the task description suggested; the dataclass has 4 fields. Tests cover the success serialization, default-arg fan-out, explicit overrides, and the one-tool registration. all_handlers() now returns 15 handlers. Co-Authored-By: Claude Opus 4.7 --- better_memory/mcp/handlers/__init__.py | 2 + better_memory/mcp/handlers/retention.py | 100 ++++++++++++++++++ tests/mcp/handlers/test_retention_handler.py | 102 +++++++++++++++++++ 3 files changed, 204 insertions(+) create mode 100644 better_memory/mcp/handlers/retention.py create mode 100644 tests/mcp/handlers/test_retention_handler.py diff --git a/better_memory/mcp/handlers/__init__.py b/better_memory/mcp/handlers/__init__.py index 663dc9e..9373982 100644 --- a/better_memory/mcp/handlers/__init__.py +++ b/better_memory/mcp/handlers/__init__.py @@ -11,6 +11,7 @@ from better_memory.mcp.handlers.episodes import EpisodeHandlers from better_memory.mcp.handlers.observations import ObservationHandlers from better_memory.mcp.handlers.reflections import ReflectionHandlers +from better_memory.mcp.handlers.retention import RetentionHandlers from better_memory.mcp.handlers.semantics import SemanticHandlers @@ -24,4 +25,5 @@ def all_handlers() -> list[Handler]: *SemanticHandlers().handlers(), *EpisodeHandlers().handlers(), *ReflectionHandlers().handlers(), + *RetentionHandlers().handlers(), ] 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/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"] From deb2d05d49f81156c7646fc56a8906ae0d76ebd0 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 15:42:39 +0100 Subject: [PATCH 15/20] feat(mcp): knowledge handlers (search + list) --- better_memory/mcp/handlers/__init__.py | 2 + better_memory/mcp/handlers/knowledge.py | 121 +++++++++++++++ tests/mcp/handlers/test_knowledge_handler.py | 148 +++++++++++++++++++ 3 files changed, 271 insertions(+) create mode 100644 better_memory/mcp/handlers/knowledge.py create mode 100644 tests/mcp/handlers/test_knowledge_handler.py diff --git a/better_memory/mcp/handlers/__init__.py b/better_memory/mcp/handlers/__init__.py index 9373982..5931e90 100644 --- a/better_memory/mcp/handlers/__init__.py +++ b/better_memory/mcp/handlers/__init__.py @@ -9,6 +9,7 @@ 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.reflections import ReflectionHandlers from better_memory.mcp.handlers.retention import RetentionHandlers @@ -26,4 +27,5 @@ def all_handlers() -> list[Handler]: *EpisodeHandlers().handlers(), *ReflectionHandlers().handlers(), *RetentionHandlers().handlers(), + *KnowledgeHandlers().handlers(), ] 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/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()) From e37ea33b90db8b1b6cd572cadd678d865f582f45 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 15:49:17 +0100 Subject: [PATCH 16/20] feat(mcp): rating handlers (list_session_exposures + apply + credit) Task 13 of the MCP server dispatcher refactor. Adds the three rating-domain handlers as a RatingHandlers class behind the ToolDispatcher / Handler abstraction, lifting the bodies and JSON schemas verbatim from the legacy _call_tool if-chain. - memory.list_session_exposures: coerces an unresolved session id to "" so the bootstrap service can return an empty exposures payload instead of raising. - memory.apply_session_ratings: raises the verbatim 3-line ValueError when no session is active (the rate-session-memories skill matches the exact text). - memory.credit: two payload shapes - {"applied": None, "skipped": "no_session"} when no session, otherwise the ApplyOutcome from MemoryRatingService.credit_one. Note the schema's class enum intentionally omits "ignored" (apply-only). Registered in handlers/__init__.py so all_handlers() now returns 20 handlers across 7 domain modules. Co-Authored-By: Claude Opus 4.7 --- better_memory/mcp/handlers/__init__.py | 2 + better_memory/mcp/handlers/ratings.py | 165 ++++++++++++++++++++ tests/mcp/handlers/test_ratings_handler.py | 167 +++++++++++++++++++++ 3 files changed, 334 insertions(+) create mode 100644 better_memory/mcp/handlers/ratings.py create mode 100644 tests/mcp/handlers/test_ratings_handler.py diff --git a/better_memory/mcp/handlers/__init__.py b/better_memory/mcp/handlers/__init__.py index 5931e90..9d433de 100644 --- a/better_memory/mcp/handlers/__init__.py +++ b/better_memory/mcp/handlers/__init__.py @@ -11,6 +11,7 @@ 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 @@ -28,4 +29,5 @@ def all_handlers() -> list[Handler]: *ReflectionHandlers().handlers(), *RetentionHandlers().handlers(), *KnowledgeHandlers().handlers(), + *RatingHandlers().handlers(), ] 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/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()) From f8b4bb750e98f0bb93c675619ad2f4137cfe879d Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 15:56:28 +0100 Subject: [PATCH 17/20] feat(mcp): session handlers (session_bootstrap + start_ui) Task 14/17 of the MCP dispatcher refactor. Lands handlers/session.py with SessionHandlers exposing memory.session_bootstrap (4-tier session_id fallback: args -> CLAUDE_SESSION_ID -> CLAUDE_CODE_SESSION_ID -> uuid; 5-key payload) and memory.start_ui. Uses the services.session_bootstrap singleton from the container (built once in Task 6) instead of re-instantiating SessionBootstrapService per call. After this task all_handlers() returns 22 -- the full tool set. Task 15 will wire ToolDispatcher into create_server. --- better_memory/mcp/handlers/__init__.py | 2 + better_memory/mcp/handlers/session.py | 132 +++++++++++++++++++++ tests/mcp/handlers/test_session_handler.py | 108 +++++++++++++++++ 3 files changed, 242 insertions(+) create mode 100644 better_memory/mcp/handlers/session.py create mode 100644 tests/mcp/handlers/test_session_handler.py diff --git a/better_memory/mcp/handlers/__init__.py b/better_memory/mcp/handlers/__init__.py index 9d433de..531ac63 100644 --- a/better_memory/mcp/handlers/__init__.py +++ b/better_memory/mcp/handlers/__init__.py @@ -15,6 +15,7 @@ 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]: @@ -30,4 +31,5 @@ def all_handlers() -> list[Handler]: *RetentionHandlers().handlers(), *KnowledgeHandlers().handlers(), *RatingHandlers().handlers(), + *SessionHandlers().handlers(), ] 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/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", + ] From 4d0e99ebb5c57c638b91dc752e40c3c5e54b72e0 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 16:59:12 +0100 Subject: [PATCH 18/20] refactor(mcp): wire ToolDispatcher into create_server, delete if-chain server.py drops from 1568 to 401 LOC. _call_tool is now a 1-line delegation to dispatcher.call(name, arguments or {}) (under _diag.trace). _tool_definitions deleted as the 550-LOC schema function; replaced by a thin shim built from all_handlers() so tests asserting tool registration keep working. All 22 tools route through the dispatcher. To break a server <-> handlers import cycle introduced by importing all_handlers() at module load, extracted _run_best_effort into mcp/_best_effort.py. server.py re-exports it so the existing test import path (better_memory.mcp.server._run_best_effort) still works; test_best_effort_logging now asserts the canonical logger name (better_memory.mcp._best_effort). ServerContext gains a `dispatcher` field LAST so positional construction keeps compiling. Synth serializers (_serialize_synth_*, _serialize_queue) re-exported from handlers.reflections for test backwards compat. _serialize_knowledge_* deleted (no callers). Co-Authored-By: Claude Opus 4.7 --- better_memory/mcp/_best_effort.py | 55 + better_memory/mcp/handlers/observations.py | 2 +- better_memory/mcp/server.py | 1375 ++------------------ tests/mcp/test_best_effort_logging.py | 8 +- 4 files changed, 165 insertions(+), 1275 deletions(-) create mode 100644 better_memory/mcp/_best_effort.py 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/handlers/observations.py b/better_memory/mcp/handlers/observations.py index 0338388..72b855a 100644 --- a/better_memory/mcp/handlers/observations.py +++ b/better_memory/mcp/handlers/observations.py @@ -17,9 +17,9 @@ 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.mcp.server import _run_best_effort from better_memory.services.retention_scheduler import RetentionScheduler _OBSERVE_SCHEMA: dict[str, Any] = { diff --git a/better_memory/mcp/server.py b/better_memory/mcp/server.py index 692134a..1d971b3 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,15 +27,11 @@ from __future__ import annotations import asyncio -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 from dataclasses import dataclass from pathlib import Path @@ -59,26 +46,34 @@ from better_memory.db.connection import connect from better_memory.db.schema import apply_migrations from better_memory.embeddings.ollama import OllamaEmbedder + +# 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.handlers._audit import _audit_synth_call -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.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 ( - EpisodeContext, - EpisodeQueueCounts, - ReflectionSynthesisService, - SynthesisStep, -) +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,41 +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}]" - ) - - def _probe_ollama(host: str) -> None: """Log a clear stderr warning if Ollama isn't reachable. Never raise. @@ -159,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: @@ -848,12 +145,15 @@ 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( @@ -911,6 +211,45 @@ def _build_services( ) +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[ Server, Callable[[], Coroutine[Any, Any, None]], @@ -922,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() @@ -946,19 +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). - # # 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 @@ -974,15 +312,6 @@ def create_server() -> tuple[ 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 # Session-start behaviour: reindex knowledge at startup. mtime-only, so # the cost is O(files) stat calls on an already-indexed corpus. We @@ -990,525 +319,31 @@ def create_server() -> tuple[ # 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, ) 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() ] From 9917350e98998eb8819c3e607b83e1ac806c6369 Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 17:07:43 +0100 Subject: [PATCH 19/20] refactor(mcp): slim _dispatch_for_tests to thin dispatcher.call shim 5 LOC instead of 40. Same signature, same behaviour. --- better_memory/mcp/server.py | 46 ++++++++----------------------------- 1 file changed, 10 insertions(+), 36 deletions(-) diff --git a/better_memory/mcp/server.py b/better_memory/mcp/server.py index 1d971b3..447aab2 100644 --- a/better_memory/mcp/server.py +++ b/better_memory/mcp/server.py @@ -35,7 +35,7 @@ from collections.abc import Callable, Coroutine from dataclasses import dataclass 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 @@ -348,43 +348,17 @@ async def _call_tool( 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. - - 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. - """ - from mcp.types import CallToolRequest, CallToolRequestParams + """Test-only entry point. Routes through ToolDispatcher. - server, cleanup, _ = create_server() + 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. + """ + _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() From ab47920f23a91c382f8e262166b39c1e342484cd Mon Sep 17 00:00:00 2001 From: gethin Date: Sat, 13 Jun 2026 17:46:14 +0100 Subject: [PATCH 20/20] =?UTF-8?q?test(mcp):=20verify=20dispatcher=20refact?= =?UTF-8?q?or=20=E2=80=94=20full=20suite=20green,=20server.py=201617=20?= =?UTF-8?q?=E2=86=92=20377=20LOC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All 22 tools route through ToolDispatcher. Handler bodies + schemas co-located in mcp/handlers/. ServiceContainer built once at startup. Final pytest: 1180 passed, 23 skipped, 0 failed (chunked across 4 junit runs because full-suite stdout buffer was truncated by the harness). Pyright clean on better_memory/mcp/ (0 errors, 0 warnings). Ruff on better_memory/mcp/ + tests/mcp/: 5 pre-existing errors (down from 6 at base 22c3d26); none introduced by this refactor. Co-Authored-By: Claude Opus 4.7