Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 207 additions & 1 deletion src/context_engine/integration/mcp_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down Expand Up @@ -123,6 +124,8 @@ class ContextEngineMCP:
"expand_chunk",
"related_context",
"session_recall",
"session_timeline",
"session_event",
"record_decision",
"record_code_area",
"index_status",
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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":
Expand Down Expand Up @@ -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.

Copilot AI Apr 27, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comment says the dual-write is so the "FTS5" recall path sees it, but session_recall currently ranks candidates via embeddings/cosine similarity and does not query the FTS tables. Please update the comment to reflect the actual mechanism (or switch recall to use the FTS index).

Suggested change
# Dual-write into memory.db so the new recall path (FTS5) sees it too.
# Mirror manual decisions into memory.db as well so they remain
# available to the database-backed memory flows alongside the
# session-capture copy used by session_recall.

Copilot uses AI. Check for mistakes.
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",
Expand All @@ -595,13 +672,100 @@ 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",
text=f"✓ Code area noted: {file_path} — {description}",
)
]

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.")]
Comment on lines +695 to +700

Copilot AI Apr 27, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

limit = int(args.get('limit') or 20) will raise ValueError for non-integer inputs and also allows negative/very large limits, which can crash the handler (when called directly) or produce unbounded output. Consider parsing with try/except and clamping to a sane range (e.g., 1..200), similar to _clamp_top_k.

Copilot uses AI. Check for mistakes.
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()
Comment on lines +714 to +718

Copilot AI Apr 27, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The sessions metadata query (SELECT project, started_at, ... FROM sessions) is executed outside the try/except. If the DB is unavailable/corrupt or schema is mid-migration, this will raise and bubble up as a generic tool failure. Consider wrapping this query in the same error handling as the turn_summaries query and returning a consistent, user-facing error message.

Copilot uses AI. Check for mistakes.
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()
Comment on lines +741 to +747

Copilot AI Apr 27, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_handle_session_event executes the SQL query without a try/except. Any sqlite error (locked DB, corruption, missing table during partial init) will raise and be caught only by the outer call_tool handler, returning a generic "Tool session_event failed" message. Consider catching DB exceptions here and returning a specific "event query failed" response (matching _handle_session_timeline behavior).

Suggested change
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()
try:
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()
except Exception as exc:
return [TextContent(type="text", text=f"event query failed: {exc}")]

Copilot uses AI. Check for mistakes.
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']}"
Comment on lines +761 to +765

Copilot AI Apr 27, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

tool_event_payloads.raw_output is nullable, but this formatter interpolates it directly. When raw_output is NULL, the response will contain the literal string None, which is misleading. Consider normalizing NULL to an empty string or an explicit placeholder (e.g., "") before formatting.

Suggested change
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']}"
raw_input = row["raw_input"] if row["raw_input"] is not None else "<no input>"
raw_output = row["raw_output"] if row["raw_output"] is not None else "<no output>"
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{raw_input}\n\n"
f"output:\n{raw_output}"

Copilot uses AI. Check for mistakes.
)
return [TextContent(type="text", text=body)]

async def _handle_index_status(self):
queries = self._stats["queries"]
raw = self._stats["raw_tokens"]
Expand Down Expand Up @@ -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"
):
Comment on lines +929 to +960

Copilot AI Apr 27, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_search_sessions now unconditionally adds up to 600 additional memory.db candidates (decisions + code_areas + turn_summaries) and then calls embedder.embed_query() for every candidate. This increases worst-case recall latency and CPU substantially for large histories. Consider using the existing FTS5 indexes to prefilter candidates by topic (or otherwise reducing candidate count) before embedding/ranking.

Copilot uses AI. Check for mistakes.
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 []

Expand Down
Loading