From 255eb531109418399700b2a90896c9a38ad993fc Mon Sep 17 00:00:00 2001 From: Fazle Elahee Date: Tue, 28 Apr 2026 00:39:41 +0100 Subject: [PATCH] =?UTF-8?q?feat(memory):=20recall=20=E2=80=94=20extend=20s?= =?UTF-8?q?ession=5Frecall=20+=20add=20session=5Ftimeline=20+=20session=5F?= =?UTF-8?q?event?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR 4 of 5 against feature/memory-claude-mem-parity, stacked on PR 3. Wires the MCP retrieval surface to memory.db while preserving the existing JSON path for backward compatibility. - ContextEngineMCP opens memory.db on startup and seeds an INSERT OR IGNORE sessions row so manual record_decision / record_code_area dual-writes don't fail the FK constraint when the SessionStart hook hasn't fired yet (test envs, future-non-CC clients). - session_recall now folds three new candidate sources on top of the existing JSON sessions / consolidated decisions: - decisions (manual + migrated, last 200 by recency) - code_areas (manual + migrated, last 200 by recency) - turn_summaries (last 200 turns, the layer-1 compact index) Tags include source and session_id so the agent can drill via the new tools. - session_timeline(session_id, limit=20) — layer 2. Returns the session's turn_summaries with rollup + status header. - session_event(event_id) — layer 3. Returns the raw input/output payload for one tool_event, with a dedicated "aged out" message when the payload row was pruned by retention. - record_decision and record_code_area now dual-write to memory.db with source='manual'. Prior JSON write path remains active so a rollback to a previous PR doesn't lose recall coverage. The JSON write side can be retired once parity is confirmed in production. Tests (10 new, full suite 350 passed): - tests/memory/test_mcp_recall.py covers dual-write of decisions and code_areas, session_timeline with seeded summaries, session_event payload roundtrip + aged-out message + invalid id, session_recall surfacing memory.db decisions, and TOOL_NAMES registration. --- src/context_engine/integration/mcp_server.py | 208 ++++++++++++++++++- tests/memory/test_mcp_recall.py | 183 ++++++++++++++++ 2 files changed, 390 insertions(+), 1 deletion(-) create mode 100644 tests/memory/test_mcp_recall.py diff --git a/src/context_engine/integration/mcp_server.py b/src/context_engine/integration/mcp_server.py index 46da147..b1ba42d 100644 --- a/src/context_engine/integration/mcp_server.py +++ b/src/context_engine/integration/mcp_server.py @@ -20,6 +20,7 @@ get_working_state, ) from context_engine.integration.session_capture import SessionCapture +from context_engine.memory import db as memory_db log = logging.getLogger(__name__) @@ -123,6 +124,8 @@ class ContextEngineMCP: "expand_chunk", "related_context", "session_recall", + "session_timeline", + "session_event", "record_decision", "record_code_area", "index_status", @@ -155,10 +158,33 @@ def __init__(self, retriever, backend, compressor, embedder, config) -> None: ) # Session capture — persists decisions and code-area notes across runs. + # Both the legacy JSON path and the new memory.db path are written to + # for record_decision / record_code_area; recall queries both. Once a + # release cycle of dual-write confirms parity, the JSON write side + # can be retired. self._session_capture = SessionCapture( sessions_dir=str(self._storage_base / "sessions") ) self._session_id = self._session_capture.start_session(project_name) + try: + self._memory_conn = memory_db.connect( + memory_db.memory_db_path(self._storage_base) + ) + # Ensure the sessions row exists so dual-writes don't trip the FK. + # The SessionStart hook normally creates this, but the MCP server + # may start in environments without hook coverage (e.g. tests). + import time as _t + _epoch = int(_t.time()) + self._memory_conn.execute( + "INSERT OR IGNORE INTO sessions (id, project, started_at_epoch, " + "started_at, status) VALUES (?, ?, ?, ?, 'active')", + (self._session_id, project_name, _epoch, + _t.strftime("%Y-%m-%dT%H:%M:%S", _t.gmtime(_epoch))), + ) + self._memory_conn.commit() + except Exception as exc: + log.warning("memory.db open failed; recall will fall back to JSON: %s", exc) + self._memory_conn = None # Cheap maintenance on start: if the project has accumulated more than # _PRUNE_THRESHOLD session files, consolidate the oldest decisions # into decisions_log.json and remove the source files. No-op when @@ -314,13 +340,45 @@ async def list_tools(): ), Tool( name="session_recall", - description="Recall past decisions and code-area notes recorded in this or prior sessions", + description=( + "Recall past decisions, prompts, and turn summaries via topic search. " + "Returns compact-index hits across the whole project history." + ), inputSchema={ "type": "object", "properties": {"topic": {"type": "string"}}, "required": ["topic"], }, ), + Tool( + name="session_timeline", + description=( + "List the turn summaries for a session, oldest first. " + "Layer 2 of progressive disclosure — drill into a session_id " + "returned by session_recall." + ), + inputSchema={ + "type": "object", + "properties": { + "session_id": {"type": "string"}, + "limit": {"type": "integer", "default": 20}, + }, + "required": ["session_id"], + }, + ), + Tool( + name="session_event", + description=( + "Return the raw input/output payload for a single tool_event. " + "Layer 3 of progressive disclosure — drill into an event_id " + "from session_timeline." + ), + inputSchema={ + "type": "object", + "properties": {"event_id": {"type": "integer"}}, + "required": ["event_id"], + }, + ), Tool( name="record_decision", description="Record a decision (with reason) for future session_recall", @@ -393,6 +451,10 @@ async def call_tool(name: str, arguments: dict): return await self._handle_related_context(arguments) elif name == "session_recall": return await self._handle_session_recall(arguments) + elif name == "session_timeline": + return self._handle_session_timeline(arguments) + elif name == "session_event": + return self._handle_session_event(arguments) elif name == "record_decision": return self._handle_record_decision(arguments) elif name == "record_code_area": @@ -579,6 +641,21 @@ def _handle_record_decision(self, args): return [TextContent(type="text", text="decision is required.")] self._session_capture.record_decision(self._session_id, decision, reason) self._persist_current_session() + # Dual-write into memory.db so the new recall path (FTS5) sees it too. + if self._memory_conn is not None: + try: + import time as _time + epoch = int(_time.time()) + self._memory_conn.execute( + "INSERT INTO decisions (session_id, decision, reason, source, " + "created_at_epoch, created_at) " + "VALUES (?, ?, ?, 'manual', ?, ?)", + (self._session_id, decision, reason, epoch, + _time.strftime("%Y-%m-%dT%H:%M:%S", _time.gmtime(epoch))), + ) + self._memory_conn.commit() + except Exception: + log.exception("memory.db decision dual-write failed") return [ TextContent( type="text", @@ -595,6 +672,18 @@ def _handle_record_code_area(self, args): self._session_id, file_path, description ) self._persist_current_session() + if self._memory_conn is not None: + try: + import time as _time + epoch = int(_time.time()) + self._memory_conn.execute( + "INSERT INTO code_areas (session_id, file_path, description, " + "source, created_at_epoch) VALUES (?, ?, ?, 'manual', ?)", + (self._session_id, file_path, description, epoch), + ) + self._memory_conn.commit() + except Exception: + log.exception("memory.db code_area dual-write failed") return [ TextContent( type="text", @@ -602,6 +691,81 @@ def _handle_record_code_area(self, args): ) ] + def _handle_session_timeline(self, args): + session_id = (args.get("session_id") or "").strip() + limit = int(args.get("limit") or 20) + if not session_id: + return [TextContent(type="text", text="session_id is required.")] + if self._memory_conn is None: + return [TextContent(type="text", text="Memory store not available.")] + try: + rows = list(self._memory_conn.execute( + "SELECT prompt_number, summary, tier FROM turn_summaries " + "WHERE session_id = ? ORDER BY prompt_number ASC LIMIT ?", + (session_id, limit), + )) + except Exception as exc: + return [TextContent(type="text", text=f"timeline query failed: {exc}")] + if not rows: + return [TextContent( + type="text", + text=f"No turn summaries for session {session_id} yet.", + )] + meta = self._memory_conn.execute( + "SELECT project, started_at, ended_at, status, prompt_count, " + "rollup_summary FROM sessions WHERE id = ?", + (session_id,), + ).fetchone() + header = [] + if meta: + header.append(f"session: {session_id} · {meta['project']} · {meta['status']}") + header.append(f"started: {meta['started_at']} ended: {meta['ended_at'] or '—'}") + if meta["rollup_summary"]: + header.append(f"rollup: {meta['rollup_summary']}") + body = "\n".join( + f" turn {r['prompt_number']:>3} [{r['tier']}] {r['summary']}" + for r in rows + ) + return [TextContent( + type="text", + text="\n".join(header) + ("\n\n" + body if header else body), + )] + + def _handle_session_event(self, args): + try: + event_id = int(args.get("event_id")) + except (TypeError, ValueError): + return [TextContent(type="text", text="event_id must be an integer.")] + if self._memory_conn is None: + return [TextContent(type="text", text="Memory store not available.")] + row = self._memory_conn.execute( + "SELECT te.tool_name, te.session_id, te.prompt_number, te.created_at, " + "p.raw_input, p.raw_output FROM tool_events te " + "LEFT JOIN tool_event_payloads p ON p.id = te.payload_id " + "WHERE te.id = ?", + (event_id,), + ).fetchone() + if row is None: + return [TextContent( + type="text", + text=f"No event with id={event_id}.", + )] + if row["raw_input"] is None and row["raw_output"] is None: + return [TextContent( + type="text", + text=( + f"Event {event_id} ({row['tool_name']}) was retained as a summary " + "only — its raw payload aged out of the retention window." + ), + )] + body = ( + f"event {event_id} · {row['tool_name']} · session {row['session_id']} · " + f"turn {row['prompt_number']} · {row['created_at']}\n\n" + f"input:\n{row['raw_input']}\n\n" + f"output:\n{row['raw_output']}" + ) + return [TextContent(type="text", text=body)] + async def _handle_index_status(self): queries = self._stats["queries"] raw = self._stats["raw_tokens"] @@ -762,6 +926,48 @@ def _search_sessions(self, topic: str) -> list[str]: seen.add(text) candidates.append(text) + # Fold in memory.db rows: decisions/code_areas (manual or migrated), + # plus the index of turn_summaries. Tagged with [layer:..|sid:..] so + # the agent knows where to drill from session_recall results — those + # tags map onto session_timeline / session_event drill-downs. + if self._memory_conn is not None: + try: + for row in self._memory_conn.execute( + "SELECT decision, reason, source, session_id " + "FROM decisions ORDER BY created_at_epoch DESC LIMIT 200" + ): + text = ( + f"[decision src={row['source']}|sid:{row['session_id'] or '-'}] " + f"{row['decision']} — {row['reason']}" + ) + if text not in seen: + seen.add(text) + candidates.append(text) + for row in self._memory_conn.execute( + "SELECT file_path, description, source, session_id " + "FROM code_areas ORDER BY created_at_epoch DESC LIMIT 200" + ): + text = ( + f"[code_area src={row['source']}|sid:{row['session_id'] or '-'}] " + f"{row['file_path']} — {row['description']}" + ) + if text not in seen: + seen.add(text) + candidates.append(text) + for row in self._memory_conn.execute( + "SELECT session_id, prompt_number, summary " + "FROM turn_summaries ORDER BY created_at_epoch DESC LIMIT 200" + ): + text = ( + f"[turn sid:{row['session_id']}|n:{row['prompt_number']}] " + f"{row['summary']}" + ) + if text not in seen: + seen.add(text) + candidates.append(text) + except Exception: + log.exception("memory.db recall query failed; using JSON only") + if not candidates: return [] diff --git a/tests/memory/test_mcp_recall.py b/tests/memory/test_mcp_recall.py new file mode 100644 index 0000000..622ceca --- /dev/null +++ b/tests/memory/test_mcp_recall.py @@ -0,0 +1,183 @@ +"""Tests for PR 4 — extended session_recall + new MCP tools. + +These tests exercise the recall handlers directly without a stdio transport +by calling the private _handle_* methods on a constructed ContextEngineMCP. +""" +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from context_engine.config import Config +from context_engine.integration.mcp_server import ContextEngineMCP +from context_engine.memory import db as memory_db + + +@pytest.fixture +def mcp(tmp_path, monkeypatch): + """A ContextEngineMCP bound to a tmp project + storage. Stub deps.""" + project_dir = tmp_path / "demo" + project_dir.mkdir() + storage_path = tmp_path / "storage" + monkeypatch.chdir(project_dir) + + config = Config( + storage_path=str(storage_path), + embedding_model="BAAI/bge-small-en-v1.5", + ) + + backend = MagicMock() + backend._vector_store.count.return_value = 0 + compressor = MagicMock() + embedder = MagicMock() + # The recall path embeds candidates; return a stable vector so + # cosine ranking is deterministic. + embedder.embed_query = lambda text: [1.0, 0.0] if "KEY" in text else [0.0, 1.0] + retriever = MagicMock() + + server = ContextEngineMCP( + retriever=retriever, backend=backend, compressor=compressor, + embedder=embedder, config=config, + ) + yield server + if server._memory_conn is not None: + server._memory_conn.close() + + +def test_record_decision_dual_writes_to_memory_db(mcp): + out = mcp._handle_record_decision({ + "decision": "Use bge-small for KEY recall", + "reason": "Already loaded for the index", + }) + assert "Decision recorded" in out[0].text + + rows = list(mcp._memory_conn.execute( + "SELECT decision, reason, source FROM decisions" + )) + assert len(rows) == 1 + assert rows[0]["source"] == "manual" + assert "bge-small" in rows[0]["decision"] + + +def test_record_code_area_dual_writes_to_memory_db(mcp): + mcp._handle_record_code_area({ + "file_path": "src/foo.py", + "description": "memory bootstrap", + }) + rows = list(mcp._memory_conn.execute( + "SELECT file_path, description, source FROM code_areas" + )) + assert len(rows) == 1 + assert rows[0]["source"] == "manual" + + +def test_session_timeline_returns_turn_summaries_for_session(mcp): + sid = "tl-test" + mcp._memory_conn.execute( + "INSERT INTO sessions (id, project, started_at_epoch, started_at, " + "status, prompt_count) VALUES (?, 'demo', 1700000000, " + "'2023-11-14T22:13:20', 'completed', 2)", + (sid,), + ) + mcp._memory_conn.execute( + "INSERT INTO turn_summaries (session_id, prompt_number, summary, tier, " + "created_at_epoch) VALUES (?, 1, 'first turn summary', 'extractive', " + "1700000010)", (sid,), + ) + mcp._memory_conn.execute( + "INSERT INTO turn_summaries (session_id, prompt_number, summary, tier, " + "created_at_epoch) VALUES (?, 2, 'second turn summary', 'extractive', " + "1700000020)", (sid,), + ) + mcp._memory_conn.commit() + + out = mcp._handle_session_timeline({"session_id": sid}) + text = out[0].text + assert "first turn summary" in text + assert "second turn summary" in text + assert "turn 1" in text and "turn 2" in text + + +def test_session_timeline_empty_session(mcp): + out = mcp._handle_session_timeline({"session_id": "missing"}) + assert "No turn summaries" in out[0].text + + +def test_session_timeline_requires_session_id(mcp): + out = mcp._handle_session_timeline({}) + assert "required" in out[0].text + + +def test_session_event_returns_raw_payload(mcp): + mcp._memory_conn.execute( + "INSERT INTO sessions (id, project, started_at_epoch, started_at) " + "VALUES ('sx', 'demo', 1700000000, '2023-11-14T22:13:20')" + ) + cur = mcp._memory_conn.execute( + "INSERT INTO tool_event_payloads (raw_input, raw_output, size_bytes) " + "VALUES (?, ?, ?)", + (json.dumps({"file_path": "/tmp/x"}), "x = 1", 5), + ) + payload_id = cur.lastrowid + cur = mcp._memory_conn.execute( + "INSERT INTO tool_events (session_id, prompt_number, tool_name, " + "payload_id, created_at_epoch, created_at) " + "VALUES ('sx', 1, 'Read', ?, 1700000000, '2023-11-14T22:13:20')", + (payload_id,), + ) + event_id = cur.lastrowid + mcp._memory_conn.commit() + + out = mcp._handle_session_event({"event_id": event_id}) + text = out[0].text + assert "Read" in text + assert "/tmp/x" in text + assert "x = 1" in text + + +def test_session_event_returns_summary_only_message_when_payload_pruned(mcp): + mcp._memory_conn.execute( + "INSERT INTO sessions (id, project, started_at_epoch, started_at) " + "VALUES ('sx', 'demo', 1700000000, '2023-11-14T22:13:20')" + ) + cur = mcp._memory_conn.execute( + "INSERT INTO tool_events (session_id, prompt_number, tool_name, " + "payload_id, created_at_epoch, created_at) " + "VALUES ('sx', 1, 'Read', NULL, 1700000000, '2023-11-14T22:13:20')", + ) + event_id = cur.lastrowid + mcp._memory_conn.commit() + + out = mcp._handle_session_event({"event_id": event_id}) + assert "aged out" in out[0].text + + +def test_session_event_invalid_id(mcp): + out = mcp._handle_session_event({"event_id": "abc"}) + assert "must be an integer" in out[0].text + + +def test_session_recall_includes_memory_db_decisions(mcp): + """A decision in memory.db should surface via session_recall.""" + # Seed memory.db with a relevant decision. + mcp._memory_conn.execute( + "INSERT INTO decisions (decision, reason, source, " + "created_at_epoch, created_at) VALUES (?, ?, 'manual', 1700000000, " + "'2023-11-14T22:13:20')", + ("Pick KEY library for X", "KEY rationale here"), + ) + mcp._memory_conn.commit() + + matches = mcp._search_sessions("KEY") + # The candidate text contains "[decision src=manual|sid:-]" prefix + + # decision text. We just need at least one match referencing KEY. + assert any("KEY" in m for m in matches), matches + + +def test_tool_names_includes_new_tools(mcp): + assert "session_timeline" in mcp.TOOL_NAMES + assert "session_event" in mcp.TOOL_NAMES + assert "session_recall" in mcp.TOOL_NAMES