diff --git a/README.md b/README.md index 439713b..4e7796f 100644 --- a/README.md +++ b/README.md @@ -181,7 +181,7 @@ despite the shared name. ## MCP tools -The server registers 22 tools, grouped below. Full schemas are in [`website/mcp-tools.md`](website/mcp-tools.md) and live in `better_memory/mcp/server.py`. +The server registers 22 tools, grouped below. Full schemas are in [`website/mcp-tools.md`](website/mcp-tools.md) and live in `better_memory/mcp/tools.py`. **Episodic memory** — observations the AI writes during a session. diff --git a/better_memory/mcp/_util.py b/better_memory/mcp/_util.py new file mode 100644 index 0000000..3e5d9d3 --- /dev/null +++ b/better_memory/mcp/_util.py @@ -0,0 +1,71 @@ +"""Small shared helpers for the MCP server and its tool handlers. + +Lives below both :mod:`better_memory.mcp.server` and the handler modules +under :mod:`better_memory.mcp.handlers` so either side can import it +without creating a cycle. +""" + +from __future__ import annotations + +import logging +import os +import time +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from better_memory import _diag +from better_memory.runtime.session_marker import read_session_id + +logger = logging.getLogger(__name__) + + +def run_best_effort( + operation: str, + fn: Callable[[], Any], + *, + diag_cid: str | None = None, +) -> None: + """Run ``fn`` swallowing any ``Exception`` but logging it via the module logger. + + Used by best-effort hooks inside ``memory.retrieve`` (spool drain, + retention scheduler) where a failure must NEVER block the call but + must still produce a discoverable diagnostic. The previous behaviour + silently dropped the exception, so a broken background path could + fail invisibly for weeks. + + When ``diag_cid`` is provided and ``BETTER_MEMORY_EMBED_LOG=1`` is on, + emits a ``[bm-retrieve step= cid=... ms=N status=ok|error]`` + line through the shared diagnostic logger so callers can localize + which step is slow. + """ + t0 = time.monotonic() + status = "ok" + try: + fn() + except Exception: # noqa: BLE001 — best-effort wrapper + status = "error" + logger.exception("best-effort %s failed", operation) + finally: + if diag_cid is not None and _diag.enabled(): + ms = int((time.monotonic() - t0) * 1000) + _diag.log( + f"[bm-retrieve step={operation} cid={diag_cid} " + f"ms={ms} status={status}]" + ) + + +def resolve_session_id(home: Path) -> str | None: + """Resolve the current Claude Code session id. + + Order: ``CLAUDE_SESSION_ID`` env, ``CLAUDE_CODE_SESSION_ID`` env, then + the marker file written by the SessionStart hook (see + :mod:`better_memory.runtime.session_marker`). Claude Code does not + propagate the session id into the spawned stdio MCP server's env, so + the marker file is the fallback for every rating call. + """ + return ( + os.environ.get("CLAUDE_SESSION_ID") + or os.environ.get("CLAUDE_CODE_SESSION_ID") + or read_session_id(home) + ) diff --git a/better_memory/mcp/handlers/__init__.py b/better_memory/mcp/handlers/__init__.py new file mode 100644 index 0000000..d311171 --- /dev/null +++ b/better_memory/mcp/handlers/__init__.py @@ -0,0 +1,59 @@ +"""Per-domain MCP tool handler classes. + +Each module in this package owns one tool domain and exposes a handler +class whose ``tools()`` method returns a ``{tool_name: coroutine}`` +mapping. ``create_server`` (in :mod:`better_memory.mcp.server`) +constructs the services once, instantiates each handler class with the +services it needs, and merges the mappings into a single dispatch +registry — keeping ``_call_tool`` a pure lookup. + +Handler classes hold no state beyond the injected services, so the +concurrency contract documented in ``create_server`` (one in-flight +tool call at a time over the shared sqlite connection) is unchanged. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import Any + +from mcp.types import TextContent + +from better_memory.mcp.handlers.episodes import EpisodeToolHandlers +from better_memory.mcp.handlers.knowledge import KnowledgeToolHandlers +from better_memory.mcp.handlers.observations import ObservationToolHandlers +from better_memory.mcp.handlers.reflections import ReflectionToolHandlers +from better_memory.mcp.handlers.semantics import SemanticToolHandlers +from better_memory.mcp.handlers.sessions import SessionToolHandlers + +# One MCP tool invocation: JSON-decoded arguments in, TextContent out. +ToolHandler = Callable[[dict[str, Any]], Awaitable[list[TextContent]]] + + +def build_registry( + *groups: Any, +) -> dict[str, ToolHandler]: + """Merge handler groups into one ``{tool_name: handler}`` registry. + + Raises ``RuntimeError`` on a duplicate tool name so a wiring mistake + fails at server construction, not at first dispatch. + """ + registry: dict[str, ToolHandler] = {} + for group in groups: + for tool_name, handler in group.tools().items(): + if tool_name in registry: + raise RuntimeError(f"duplicate tool handler: {tool_name}") + registry[tool_name] = handler + return registry + + +__all__ = [ + "EpisodeToolHandlers", + "KnowledgeToolHandlers", + "ObservationToolHandlers", + "ReflectionToolHandlers", + "SemanticToolHandlers", + "SessionToolHandlers", + "ToolHandler", + "build_registry", +] diff --git a/better_memory/mcp/handlers/episodes.py b/better_memory/mcp/handlers/episodes.py new file mode 100644 index 0000000..eb57a58 --- /dev/null +++ b/better_memory/mcp/handlers/episodes.py @@ -0,0 +1,154 @@ +"""Handlers for the episode-lifecycle tools. + +Tools: ``memory.start_episode``, ``memory.close_episode``, +``memory.reconcile_episodes``, ``memory.list_episodes``. +""" + +from __future__ import annotations + +import json +from typing import Any + +from mcp.types import TextContent + +from better_memory.config import project_name +from better_memory.mcp.serializers import serialize_queue +from better_memory.services.episode import EpisodeService +from better_memory.services.observation import ObservationService +from better_memory.services.reflection import ReflectionSynthesisService +from better_memory.storage import StorageBackend + + +class EpisodeToolHandlers: + """Episode open / close / reconcile / list.""" + + def __init__( + self, + *, + episodes: EpisodeService, + observations: ObservationService, + reflections: ReflectionSynthesisService, + backend: StorageBackend, + ) -> None: + self._episodes = episodes + self._observations = observations + self._reflections = reflections + self._backend = backend + + def tools(self) -> dict[str, Any]: + return { + "memory.start_episode": self.start_episode, + "memory.close_episode": self.close_episode, + "memory.reconcile_episodes": self.reconcile_episodes, + "memory.list_episodes": self.list_episodes, + } + + async def start_episode(self, args: dict[str, Any]) -> list[TextContent]: + project = project_name() + episode_id = self._episodes.start_foreground( + session_id=self._observations.session_id, + project=project, + goal=args["goal"], + tech=args.get("tech"), + ) + # Reflection synthesis is now interactive — driven by the IDE-LLM + # via memory.synthesize_next_get_context / _apply. We surface the + # current pending count so the LLM knows whether it should run + # synthesis (typically before treating retrieved reflections as + # canonical). + queue = self._reflections._read_queue_counts(project=project) + buckets = self._backend.retrieve( + project=project, tech=args.get("tech"), + ) + return [ + TextContent( + type="text", + text=json.dumps( + { + "episode_id": episode_id, + "reflections": buckets, + "pending_synthesis": serialize_queue(queue), + } + ), + ) + ] + + async def close_episode(self, args: dict[str, Any]) -> list[TextContent]: + outcome = args["outcome"] + # Default close_reason: match outcome for the common paths. + default_reasons = { + "success": "goal_complete", + "partial": "plan_complete", + "abandoned": "abandoned", + "no_outcome": "session_end_reconciled", + } + close_reason = args.get("close_reason") or default_reasons[outcome] + try: + closed_id = self._episodes.close_active( + session_id=self._observations.session_id, + outcome=outcome, + close_reason=close_reason, + summary=args.get("summary"), + ) + except ValueError: + # No active episode — already closed (e.g. by a prior commit- + # trailer drain) or never opened. Matches the "safe no-op" + # contract documented in the CLAUDE snippet's plan-complete + # section. + return [ + TextContent( + type="text", + text=json.dumps( + {"closed_episode_id": None, "already_closed": True} + ), + ) + ] + return [ + TextContent( + type="text", + text=json.dumps( + {"closed_episode_id": closed_id, "already_closed": False} + ), + ) + ] + + async def reconcile_episodes( + self, args: dict[str, Any] + ) -> list[TextContent]: + open_episodes = self._episodes.unclosed_episodes( + exclude_session_ids={self._observations.session_id} + ) + payload = [ + { + "episode_id": e.id, + "project": e.project, + "tech": e.tech, + "goal": e.goal, + "started_at": e.started_at, + } + for e in open_episodes + ] + return [TextContent(type="text", text=json.dumps(payload))] + + async def list_episodes(self, args: dict[str, Any]) -> list[TextContent]: + rows = self._episodes.list_episodes( + project=args.get("project"), + outcome=args.get("outcome"), + only_open=args.get("only_open", False), + ) + payload = [ + { + "episode_id": e.id, + "project": e.project, + "tech": e.tech, + "goal": e.goal, + "started_at": e.started_at, + "hardened_at": e.hardened_at, + "ended_at": e.ended_at, + "close_reason": e.close_reason, + "outcome": e.outcome, + "summary": e.summary, + } + for e in rows + ] + return [TextContent(type="text", text=json.dumps(payload))] diff --git a/better_memory/mcp/handlers/knowledge.py b/better_memory/mcp/handlers/knowledge.py new file mode 100644 index 0000000..0a2bd5d --- /dev/null +++ b/better_memory/mcp/handlers/knowledge.py @@ -0,0 +1,55 @@ +"""Handlers for the knowledge-base introspection tools. + +Tools: ``knowledge.search``, ``knowledge.list``. +""" + +from __future__ import annotations + +import json +from typing import Any + +from mcp.types import TextContent + +from better_memory.mcp.serializers import ( + serialize_knowledge_doc, + serialize_knowledge_search, +) +from better_memory.services.knowledge import KnowledgeService + + +class KnowledgeToolHandlers: + """BM25 search + listing over the knowledge-base corpus.""" + + def __init__(self, *, knowledge: KnowledgeService) -> None: + self._knowledge = knowledge + + def tools(self) -> dict[str, Any]: + return { + "knowledge.search": self.search, + "knowledge.list": self.list_documents, + } + + async def search(self, args: dict[str, Any]) -> list[TextContent]: + results = self._knowledge.search( + args["query"], + project=args.get("project"), + ) + return [ + TextContent( + type="text", + text=json.dumps( + [serialize_knowledge_search(r) for r in results] + ), + ) + ] + + async def list_documents(self, args: dict[str, Any]) -> list[TextContent]: + docs = self._knowledge.list_documents(project=args.get("project")) + return [ + TextContent( + type="text", + text=json.dumps( + [serialize_knowledge_doc(d) for d in docs] + ), + ) + ] diff --git a/better_memory/mcp/handlers/observations.py b/better_memory/mcp/handlers/observations.py new file mode 100644 index 0000000..e6c5dea --- /dev/null +++ b/better_memory/mcp/handlers/observations.py @@ -0,0 +1,107 @@ +"""Handlers for the observation-lifecycle tools. + +Tools: ``memory.observe``, ``memory.retrieve_observations``, +``memory.record_use``, ``memory.run_retention``. +""" + +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.services.observation import ObservationService +from better_memory.services.retention import RetentionService + + +class ObservationToolHandlers: + """Observation create / drill-down / reinforcement / retention.""" + + def __init__( + self, + *, + observations: ObservationService, + retention: RetentionService, + ) -> None: + self._observations = observations + self._retention = retention + + def tools(self) -> dict[str, Any]: + return { + "memory.observe": self.observe, + "memory.retrieve_observations": self.retrieve_observations, + "memory.record_use": self.record_use, + "memory.run_retention": self.run_retention, + } + + async def observe(self, 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 self._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_observations( + self, args: dict[str, Any] + ) -> list[TextContent]: + project = args.get("project") or project_name() + results = await self._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, args: dict[str, Any]) -> list[TextContent]: + self._observations.record_use( + args["id"], + outcome=args.get("outcome"), + ) + return [TextContent(type="text", text=json.dumps({"ok": True}))] + + async def run_retention(self, args: dict[str, Any]) -> list[TextContent]: + report = self._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, + }), + ) + ] diff --git a/better_memory/mcp/handlers/reflections.py b/better_memory/mcp/handlers/reflections.py new file mode 100644 index 0000000..2549edd --- /dev/null +++ b/better_memory/mcp/handlers/reflections.py @@ -0,0 +1,204 @@ +"""Handlers for reflection retrieval and IDE-driven synthesis. + +Tools: ``memory.retrieve``, ``memory.synthesize_next_get_context``, +``memory.synthesize_next_apply``. +""" + +from __future__ import annotations + +import json +import sqlite3 +import time +import uuid +from pathlib import Path +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._util import run_best_effort +from better_memory.mcp.serializers import ( + serialize_synth_apply_ok, + serialize_synth_apply_state_error, + serialize_synth_apply_validation_error, + serialize_synth_get_context, +) +from better_memory.mcp.synth_audit import audit_synth_call +from better_memory.services.reflection import ( + ReflectionSynthesisService, + SynthesisResponseError, +) +from better_memory.services.retention_scheduler import RetentionScheduler +from better_memory.services.spool import SpoolService +from better_memory.storage import StorageBackend + + +class ReflectionToolHandlers: + """Reflection retrieval (with its best-effort pre-hooks) + synthesis.""" + + def __init__( + self, + *, + backend: StorageBackend, + reflections: ReflectionSynthesisService, + spool: SpoolService, + memory_conn: sqlite3.Connection, + home: Path, + ) -> None: + self._backend = backend + self._reflections = reflections + self._spool = spool + self._memory_conn = memory_conn + self._home = home + + def tools(self) -> dict[str, Any]: + return { + "memory.retrieve": self.retrieve, + "memory.synthesize_next_get_context": self.synthesize_next_get_context, + "memory.synthesize_next_apply": self.synthesize_next_apply, + } + + async def retrieve(self, 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", self._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( + self._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 = self._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 synthesize_next_get_context( + self, args: dict[str, Any] + ) -> list[TextContent]: + project = args.get("project") or project_name() + with audit_synth_call( + self._home, + tool="get_context", + project=project, + episode_id=None, + ) as audit: + ctx = self._reflections.get_next_pending_context(project=project) + queue = self._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 synthesize_next_apply( + self, 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( + self._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 = self._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 = self._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)), + ) + ] diff --git a/better_memory/mcp/handlers/semantics.py b/better_memory/mcp/handlers/semantics.py new file mode 100644 index 0000000..13157cd --- /dev/null +++ b/better_memory/mcp/handlers/semantics.py @@ -0,0 +1,67 @@ +"""Handlers for the semantic-memory CRUD tools. + +Tools: ``memory.semantic_observe``, ``memory.semantic_retrieve``, +``memory.semantic_update``, ``memory.semantic_delete``. +""" + +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.services.semantic import SemanticMemoryService + + +class SemanticToolHandlers: + """User-stated facts/preferences CRUD.""" + + def __init__(self, *, semantic: SemanticMemoryService) -> None: + self._semantic = semantic + + def tools(self) -> dict[str, Any]: + return { + "memory.semantic_observe": self.semantic_observe, + "memory.semantic_retrieve": self.semantic_retrieve, + "memory.semantic_update": self.semantic_update, + "memory.semantic_delete": self.semantic_delete, + } + + async def semantic_observe(self, 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. + # Same fix as PR #25's BugBot finding on memory.observe. + memory_id = self._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 semantic_retrieve(self, args: dict[str, Any]) -> list[TextContent]: + project = args.get("project") or project_name() + memories = self._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 semantic_update(self, args: dict[str, Any]) -> list[TextContent]: + self._semantic.update_text(id=args["id"], content=args["content"]) + return [TextContent(type="text", text=json.dumps({"ok": True}))] + + async def semantic_delete(self, args: dict[str, Any]) -> list[TextContent]: + self._semantic.delete(id=args["id"]) + return [TextContent(type="text", text=json.dumps({"ok": True}))] diff --git a/better_memory/mcp/handlers/sessions.py b/better_memory/mcp/handlers/sessions.py new file mode 100644 index 0000000..1560c4d --- /dev/null +++ b/better_memory/mcp/handlers/sessions.py @@ -0,0 +1,116 @@ +"""Handlers for session-scoped tools: bootstrap, rating, credit, UI. + +Tools: ``memory.session_bootstrap``, ``memory.list_session_exposures``, +``memory.apply_session_ratings``, ``memory.credit``, ``memory.start_ui``. +""" + +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._util import resolve_session_id +from better_memory.services import ui_launcher +from better_memory.services.memory_rating import MemoryRatingService +from better_memory.services.session_bootstrap import SessionBootstrapService + + +class SessionToolHandlers: + """Session bootstrap + exposure rating + management-UI launch.""" + + def __init__( + self, + *, + session_bootstrap: SessionBootstrapService, + memory_rating: MemoryRatingService, + home: Path, + ) -> None: + self._session_bootstrap = session_bootstrap + self._memory_rating = memory_rating + self._home = home + + def tools(self) -> dict[str, Any]: + return { + "memory.session_bootstrap": self.session_bootstrap, + "memory.list_session_exposures": self.list_session_exposures, + "memory.apply_session_ratings": self.apply_session_ratings, + "memory.credit": self.credit, + "memory.start_ui": self.start_ui, + } + + async def session_bootstrap(self, 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 = self._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 list_session_exposures( + self, args: dict[str, Any] + ) -> list[TextContent]: + sid = resolve_session_id(self._home) or "" + payload = self._session_bootstrap.list_session_exposures( + session_id=sid, + ) + return [TextContent(type="text", text=json.dumps(payload))] + + async def apply_session_ratings( + self, args: dict[str, Any] + ) -> list[TextContent]: + sid = resolve_session_id(self._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 = self._memory_rating.apply_session_ratings( + session_id=sid, + ratings=args["ratings"], + ) + return [TextContent(type="text", text=json.dumps(payload))] + + async def credit(self, args: dict[str, Any]) -> list[TextContent]: + sid = resolve_session_id(self._home) + if not sid: + payload = {"applied": None, "skipped": "no_session"} + else: + payload = self._memory_rating.credit_one( + session_id=sid, + kind=args["kind"], + id=args["id"], + classification=args["class"], + ) + return [TextContent(type="text", text=json.dumps(payload))] + + async def start_ui(self, args: dict[str, Any]) -> list[TextContent]: + result = ui_launcher.start_ui() + return [ + TextContent(type="text", text=json.dumps(result)) + ] diff --git a/better_memory/mcp/serializers.py b/better_memory/mcp/serializers.py new file mode 100644 index 0000000..b9778a0 --- /dev/null +++ b/better_memory/mcp/serializers.py @@ -0,0 +1,141 @@ +"""JSON payload builders shared by the MCP tool handlers. + +Pure functions: domain read-models in, plain JSON-ready dicts out. The +handler modules call these so the wire format is defined in exactly one +place and unit tests can assert on it without driving a live server. +""" + +from __future__ import annotations + +import json +from typing import Any + +from better_memory.services.knowledge import ( + KnowledgeDocument, + KnowledgeSearchResult, +) +from better_memory.services.reflection import ( + EpisodeContext, + EpisodeQueueCounts, + SynthesisStep, +) + + +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, + } diff --git a/better_memory/mcp/server.py b/better_memory/mcp/server.py index d1422aa..13dbaf7 100644 --- a/better_memory/mcp/server.py +++ b/better_memory/mcp/server.py @@ -4,6 +4,15 @@ 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. +Module layout +------------- +This module owns wiring only: connection + service construction, the +handler registry, and the stdio run loop. The per-tool behaviour lives in +:mod:`better_memory.mcp.handlers` (one module per tool domain), the tool +name/schema declarations in :mod:`better_memory.mcp.tools`, the JSON +payload builders in :mod:`better_memory.mcp.serializers`, and the +synthesize audit log in :mod:`better_memory.mcp.synth_audit`. + Tools ----- * ``memory.observe`` — create an observation. @@ -36,19 +45,12 @@ from __future__ import annotations import asyncio -import contextlib -import json -import logging -import os import sqlite3 import sys -import time import urllib.error import urllib.request -import uuid -from collections.abc import Callable, Coroutine, Iterator +from collections.abc import Callable, Coroutine from dataclasses import dataclass -from datetime import UTC, datetime from pathlib import Path from typing import Any, cast @@ -56,29 +58,29 @@ from mcp.server.stdio import stdio_server from mcp.types import TextContent, Tool -from better_memory import _diag from better_memory.config import get_config, project_name from better_memory.db.connection import connect from better_memory.db.schema import apply_migrations from better_memory.embeddings.ollama import OllamaEmbedder -from better_memory.runtime.session_marker import read_session_id -from better_memory.services import ui_launcher -from better_memory.services.episode import EpisodeService -from better_memory.services.knowledge import ( - KnowledgeDocument, - KnowledgeSearchResult, - KnowledgeService, -) -from better_memory.services.observation import ObservationService -from better_memory.services.reflection import ( - EpisodeContext, - EpisodeQueueCounts, - ReflectionSynthesisService, - SynthesisStep, +from better_memory.mcp._util import resolve_session_id as _resolve_session_id +from better_memory.mcp.handlers import ( + EpisodeToolHandlers, + KnowledgeToolHandlers, + ObservationToolHandlers, + ReflectionToolHandlers, + SemanticToolHandlers, + SessionToolHandlers, + build_registry, ) +from better_memory.mcp.tools import tool_definitions as _tool_definitions +from better_memory.services.episode import EpisodeService +from better_memory.services.knowledge import KnowledgeService from better_memory.services.memory_rating import MemoryRatingService +from better_memory.services.observation import ObservationService +from better_memory.services.reflection import ReflectionSynthesisService from better_memory.services.retention import RetentionService -from better_memory.services.retention_scheduler import RetentionScheduler +from better_memory.services.semantic import SemanticMemoryService +from better_memory.services.session_bootstrap import SessionBootstrapService from better_memory.services.spool import SpoolService from better_memory.storage import StorageBackend, build_backend @@ -89,137 +91,6 @@ _OLLAMA_PROBE_TIMEOUT_SEC = 2.0 -logger = logging.getLogger(__name__) - - -def _run_best_effort( - operation: str, - fn: Callable[[], Any], - *, - diag_cid: str | None = None, -) -> None: - """Run ``fn`` swallowing any ``Exception`` but logging it via the module logger. - - Used by best-effort hooks inside ``memory.retrieve`` (spool drain, - retention scheduler) where a failure must NEVER block the call but - must still produce a discoverable diagnostic. The previous behaviour - silently dropped the exception, so a broken background path could - fail invisibly for weeks. - - When ``diag_cid`` is provided and ``BETTER_MEMORY_EMBED_LOG=1`` is on, - emits a ``[bm-retrieve step= cid=... ms=N status=ok|error]`` - line through the shared diagnostic logger so callers can localize - which step is slow. - """ - t0 = time.monotonic() - status = "ok" - try: - fn() - except Exception: # noqa: BLE001 — best-effort wrapper - status = "error" - logger.exception("best-effort %s failed", operation) - finally: - if diag_cid is not None and _diag.enabled(): - ms = int((time.monotonic() - t0) * 1000) - _diag.log( - f"[bm-retrieve step={operation} cid={diag_cid} " - f"ms={ms} status={status}]" - ) - - -# --------------------------------------------------------- synthesize audit log -# -# The synthesize drain loop is driven by the IDE LLM across many -# round-trips (one per pending episode). When that loop appears to -# freeze, server-side timing is the only evidence we have — the LLM -# side is opaque. Each call writes two JSONL rows to -# ``{config.home}/logs/synthesize.jsonl``: -# -# {"phase": "start", "call_id": "...", "tool": "...", ...} -# {"phase": "complete", "call_id": "...", "tool": "...", -# "latency_ms": N, "result_kind": "...", ...} -# -# Paired by call_id. A start row with no matching complete row points -# to the call that hung. Best-effort: any IO error is swallowed so the -# audit log can never block or fail the synthesize tool itself. - - -def _resolve_session_id(home: Path) -> str | None: - """Resolve the current Claude Code session id. - - Order: ``CLAUDE_SESSION_ID`` env, ``CLAUDE_CODE_SESSION_ID`` env, then - the marker file written by the SessionStart hook (see - :mod:`better_memory.runtime.session_marker`). Claude Code does not - propagate the session id into the spawned stdio MCP server's env, so - the marker file is the fallback for every rating call. - """ - return ( - os.environ.get("CLAUDE_SESSION_ID") - or os.environ.get("CLAUDE_CODE_SESSION_ID") - or read_session_id(home) - ) - - -def _append_synth_audit(home: Path, payload: dict[str, Any]) -> None: - """Append one JSONL row to ``{home}/logs/synthesize.jsonl``.""" - try: - log_dir = home / "logs" - log_dir.mkdir(parents=True, exist_ok=True) - log_path = log_dir / "synthesize.jsonl" - line = json.dumps(payload, separators=(",", ":")) - with log_path.open("a", encoding="utf-8") as fh: - fh.write(line + "\n") - except Exception: # noqa: BLE001 — best-effort audit - logger.exception("synth audit write failed") - - -@contextlib.contextmanager -def _audit_synth_call( - home: Path, - *, - tool: str, - project: str, - episode_id: str | None, -) -> Iterator[dict[str, Any]]: - """Bracket a synthesize tool call with start + complete audit rows. - - Yields a mutable ``state`` dict the caller fills in (``result_kind``, - ``error``, ``counts``, ``obs_count``, ``refl_count``, and may - overwrite ``episode_id`` once known). The complete row is written - on both normal exit and exception. Exceptions still propagate. - """ - call_id = uuid.uuid4().hex[:12] - t0 = time.perf_counter() - _append_synth_audit(home, { - "phase": "start", - "call_id": call_id, - "tool": tool, - "ts": datetime.now(UTC).isoformat(), - "project": project, - "episode_id": episode_id, - }) - state: dict[str, Any] = { - "phase": "complete", - "call_id": call_id, - "tool": tool, - "project": project, - "episode_id": episode_id, - "result_kind": None, - } - try: - yield state - except BaseException as exc: - if state.get("result_kind") is None: - state["result_kind"] = "exception" - state.setdefault("error", f"{type(exc).__name__}: {exc}") - state["ts"] = datetime.now(UTC).isoformat() - state["latency_ms"] = int((time.perf_counter() - t0) * 1000) - _append_synth_audit(home, state) - raise - state["ts"] = datetime.now(UTC).isoformat() - state["latency_ms"] = int((time.perf_counter() - t0) * 1000) - _append_synth_audit(home, state) - def _probe_ollama(host: str) -> None: """Log a clear stderr warning if Ollama isn't reachable. Never raise. @@ -253,682 +124,6 @@ def _probe_ollama(host: str) -> None: ) -# --------------------------------------------------------------------------- tools - - -def _tool_definitions(*, supports_synthesis: bool = True) -> list[Tool]: - """Return the list of tools exposed over MCP. - - 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. - """ - 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"]}, - }, - }, - ), - ] - - 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 @@ -985,17 +180,18 @@ 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 + # Concurrency invariant: the 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 + # RetentionService, SpoolService, SemanticMemoryService, + # SessionBootstrapService, MemoryRatingService) all share + # ``memory_conn``. Each service's docstring documents that it owns the + # connection and the caller must not share it with another service that + # has an open transaction. That contract holds here only because the MCP + # stdio transport serialises requests — _call_tool runs one tool + # invocation at a time, so at most one service's SAVEPOINT is ever in + # flight. If any future change introduces async fan-out, a background + # task, or a worker thread that writes to ``memory_conn``, this + # assumption breaks and the services must be reworked to use per-task # connections (or a connection pool with explicit checkout). episodes = EpisodeService(memory_conn) observations = ObservationService(memory_conn, embedder=embedder, episodes=episodes) @@ -1026,6 +222,8 @@ def create_server() -> tuple[ reflections = ReflectionSynthesisService(memory_conn) retention = RetentionService(conn=memory_conn) memory_rating = MemoryRatingService(memory_conn) + semantic = SemanticMemoryService(memory_conn) + session_bootstrap = SessionBootstrapService(memory_conn) knowledge = KnowledgeService( knowledge_conn, @@ -1043,6 +241,35 @@ def create_server() -> tuple[ except Exception: # noqa: BLE001 — best-effort startup hook pass + # All tool behaviour lives in the per-domain handler classes; this + # registry is the single dispatch surface for _call_tool. The + # synthesize handlers stay registered even when the backend does not + # support synthesis — only the *advertised* tool list is gated (see + # _list_tools), matching the pre-extraction dispatcher. + registry = build_registry( + ObservationToolHandlers(observations=observations, retention=retention), + ReflectionToolHandlers( + backend=backend, + reflections=reflections, + spool=spool, + memory_conn=memory_conn, + home=config.home, + ), + EpisodeToolHandlers( + episodes=episodes, + observations=observations, + reflections=reflections, + backend=backend, + ), + SemanticToolHandlers(semantic=semantic), + KnowledgeToolHandlers(knowledge=knowledge), + SessionToolHandlers( + session_bootstrap=session_bootstrap, + memory_rating=memory_rating, + home=config.home, + ), + ) + server: Server = Server(name="better-memory") @server.list_tools() @@ -1053,476 +280,10 @@ async def _list_tools() -> list[Tool]: async def _call_tool( 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}") + handler = registry.get(name) + if handler is None: + raise ValueError(f"Unknown tool: {name}") + return await handler(arguments or {}) cleaned = False diff --git a/better_memory/mcp/synth_audit.py b/better_memory/mcp/synth_audit.py new file mode 100644 index 0000000..fa7320d --- /dev/null +++ b/better_memory/mcp/synth_audit.py @@ -0,0 +1,91 @@ +"""JSONL audit log for the synthesize drain loop. + +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/tools.py b/better_memory/mcp/tools.py new file mode 100644 index 0000000..b613195 --- /dev/null +++ b/better_memory/mcp/tools.py @@ -0,0 +1,560 @@ +"""MCP tool definitions (names, descriptions, input schemas). + +Pure data: nothing here touches a database or a service. The server's +``list_tools`` handler calls :func:`tool_definitions`; the actual +behaviour behind each name lives in :mod:`better_memory.mcp.handlers`. +""" + +from __future__ import annotations + +from mcp.types import Tool + + +def tool_definitions(*, supports_synthesis: bool = True) -> list[Tool]: + """Return the list of tools exposed over MCP. + + 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. + """ + 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"]}, + }, + }, + ), + ] + + 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 diff --git a/tests/mcp/test_best_effort_logging.py b/tests/mcp/test_best_effort_logging.py index c0f9c33..c6e1015 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._util import run_best_effort as _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._util"): _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._util"): _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._util" and r.levelno == logging.ERROR and "spool.drain" in r.getMessage() ] diff --git a/tests/mcp/test_synth_audit_log.py b/tests/mcp/test_synth_audit_log.py index c9760a2..7a2e4dc 100644 --- a/tests/mcp/test_synth_audit_log.py +++ b/tests/mcp/test_synth_audit_log.py @@ -15,7 +15,12 @@ import pytest -from better_memory.mcp.server import _append_synth_audit, _audit_synth_call +from better_memory.mcp.synth_audit import ( + append_synth_audit as _append_synth_audit, +) +from better_memory.mcp.synth_audit import ( + audit_synth_call as _audit_synth_call, +) def _read_jsonl(home: Path) -> list[dict]: @@ -53,7 +58,7 @@ def test_swallows_io_error( not_a_dir = tmp_path / "blocker" not_a_dir.write_text("blocker", encoding="utf-8") - with caplog.at_level(logging.ERROR, logger="better_memory.mcp.server"): + with caplog.at_level(logging.ERROR, logger="better_memory.mcp.synth_audit"): _append_synth_audit(not_a_dir, {"phase": "start"}) # must not raise assert any( @@ -152,7 +157,7 @@ def test_audit_io_failure_does_not_prevent_handler_logic( not_a_dir.write_text("blocker", encoding="utf-8") ran: list[int] = [] - with caplog.at_level(logging.ERROR, logger="better_memory.mcp.server"): + with caplog.at_level(logging.ERROR, logger="better_memory.mcp.synth_audit"): with _audit_synth_call( not_a_dir, tool="apply", project="p1", episode_id="ep-1", ) as state: diff --git a/tests/mcp/test_synthesize_tools.py b/tests/mcp/test_synthesize_tools.py index 96d2024..895c054 100644 --- a/tests/mcp/test_synthesize_tools.py +++ b/tests/mcp/test_synthesize_tools.py @@ -32,13 +32,19 @@ from better_memory.db.connection import connect from better_memory.db.schema import apply_migrations -from better_memory.mcp.server import ( - _serialize_synth_apply_ok, - _serialize_synth_apply_state_error, - _serialize_synth_apply_validation_error, - _serialize_synth_get_context, - _tool_definitions, +from better_memory.mcp.serializers import ( + serialize_synth_apply_ok as _serialize_synth_apply_ok, ) +from better_memory.mcp.serializers import ( + serialize_synth_apply_state_error as _serialize_synth_apply_state_error, +) +from better_memory.mcp.serializers import ( + serialize_synth_apply_validation_error as _serialize_synth_apply_validation_error, +) +from better_memory.mcp.serializers import ( + serialize_synth_get_context as _serialize_synth_get_context, +) +from better_memory.mcp.server import _tool_definitions from better_memory.services.reflection import ( EpisodeContext, EpisodeForPrompt,