diff --git a/src/context_engine/integration/mcp_server.py b/src/context_engine/integration/mcp_server.py index d359223..befb549 100644 --- a/src/context_engine/integration/mcp_server.py +++ b/src/context_engine/integration/mcp_server.py @@ -469,6 +469,12 @@ def __init__(self, retriever, backend, compressor, embedder, config) -> None: # Lazy indexing flag — triggers on first context_search if index is empty. self._lazy_indexed = False + # Memory nudge state — tracks tool activity so context_search can + # append reminders when the agent hasn't been recording (#nudges). + # Reset on every record_decision / record_code_area call. + self._searches_since_last_decision = 0 + self._has_recorded_decision = False + self._register_tools() self._register_prompts() @@ -709,6 +715,76 @@ def _apply_output_compression(self, body: str) -> str: ) return out + # ── memory nudges ─────────────────────────────────────────────────── + + # Thresholds for triggering nudges. After the first record_decision, + # the search threshold doubles — the agent has shown it knows how, so + # we nudge less aggressively. + _DECISION_NUDGE_FIRST = 4 + _DECISION_NUDGE_REPEAT = 8 + _CODE_AREA_NUDGE_THRESHOLD = 3 + _CODE_AREA_NUDGE_MAX_FILES = 5 + + def _build_nudge(self) -> str: + """Return a nudge string to append to context_search results, or "". + + Two independent triggers: + 1. Decision nudge: N searches with 0 decisions recorded. + 2. Code area nudge: M files touched (via search/expand) but not + explicitly record_code_area'd this session. + + Both can fire in the same response. The "skip if" clause gives the + agent permission to ignore the nudge when it's just reading. + """ + parts: list[str] = [] + + # ── decision nudge ──────────────────────────────────────────── + threshold = ( + self._DECISION_NUDGE_REPEAT + if self._has_recorded_decision + else self._DECISION_NUDGE_FIRST + ) + if self._searches_since_last_decision >= threshold: + n = self._searches_since_last_decision + snap = self._session_capture.get_session_snapshot(self._session_id) + n_decisions = len(snap["decisions"]) if snap else 0 + parts.append( + f"[CCE] {n} context searches this session, " + f"{n_decisions} decision(s) recorded.\n" + f"If you chose a library, resolved an ambiguity, or " + f"established a convention, call " + f'record_decision(decision="...", reason="...").\n' + f"Skip if this session is just exploration/reading." + ) + + # ── code area nudge ─────────────────────────────────────────── + snap = self._session_capture.get_session_snapshot(self._session_id) + if snap: + touched = set(snap.get("touched_files", {}).keys()) + recorded = { + ca["file_path"] for ca in snap.get("code_areas", []) + } + unrecorded = sorted(touched - recorded) + if len(unrecorded) >= self._CODE_AREA_NUDGE_THRESHOLD: + shown = unrecorded[:self._CODE_AREA_NUDGE_MAX_FILES] + suffix = ( + f" (+{len(unrecorded) - len(shown)} more)" + if len(unrecorded) > len(shown) else "" + ) + file_list = ", ".join(shown) + suffix + parts.append( + f"[CCE] {len(unrecorded)} files explored but not " + f"recorded: {file_list}\n" + f"If you modified or traced a non-obvious flow, call " + f'record_code_area(file_path="...", ' + f'description="...").\n' + f"Skip for files you only glanced at." + ) + + if not parts: + return "" + return "\n\n---\n" + "\n\n".join(parts) + def get_tool_names(self) -> list[str]: return list(self.TOOL_NAMES) @@ -1048,6 +1124,14 @@ async def _handle_context_search(self, args): self._persist_current_session() body = _format_results_with_overflow(inline_chunks, overflow_chunks) + # Memory nudge — appended before output compression so the agent + # sees it as part of the tool result. Fires only when thresholds + # are met (see _build_nudge). Increment BEFORE building the nudge + # so the current search counts toward the threshold. + self._searches_since_last_decision += 1 + nudge = self._build_nudge() + if nudge: + body = body + nudge body = self._apply_output_compression(body) # Only note omissions the retriever actually made (threshold or # marginal stop) — chunk-count heuristics false-positive on every @@ -1246,6 +1330,10 @@ def _handle_record_decision(self, args): reason = memory_db.scrub_pii(reason) self._session_capture.record_decision(self._session_id, decision, reason) self._persist_current_session() + # Reset nudge state — the agent recorded something, give it a + # longer leash before the next nudge. + self._searches_since_last_decision = 0 + self._has_recorded_decision = True # Dual-write into memory.db. `decision` and `reason` are compressed # via the grammar module before INSERT — structured tokens (paths, # versions, identifiers) are preserved byte-for-byte; only prose diff --git a/src/context_engine/memory/hook_installer.py b/src/context_engine/memory/hook_installer.py index 7e657be..eb20ed6 100644 --- a/src/context_engine/memory/hook_installer.py +++ b/src/context_engine/memory/hook_installer.py @@ -97,7 +97,7 @@ def _is_windows() -> bool: bash -c "exec 3<>/dev/tcp/127.0.0.1/${PORT}" 2>/dev/null || exit 0 fi -if [ "${HOOK_NAME}" = "SessionStart" ]; then +if [ "${HOOK_NAME}" = "SessionStart" ] || [ "${HOOK_NAME}" = "Stop" ]; then RESPONSE="$(curl -sf -m 2 -X POST -H "Content-Type: application/json" \\ --data-binary @- "http://127.0.0.1:${PORT}/hooks/${HOOK_NAME}" \\ 2>/dev/null || true)" @@ -154,17 +154,19 @@ def _is_windows() -> bool: powershell -NoProfile -Command "$ErrorActionPreference='Stop'; try { (New-Object Net.Sockets.TcpClient).Connect('127.0.0.1',%PORT%); exit 0 } catch { exit 1 }" >nul 2>&1 if errorlevel 1 exit /b 0 -if /i "%HOOK_NAME%"=="SessionStart" ( - set "TMP_RESP=%TEMP%\\cce_hook_resp_%RANDOM%.txt" - curl -sf -m 2 -X POST -H "Content-Type: application/json" ^ - --data-binary @- "http://127.0.0.1:%PORT%/hooks/%HOOK_NAME%" ^ - > "%TMP_RESP%" 2>nul - if exist "%TMP_RESP%" type "%TMP_RESP%" - if exist "%TMP_RESP%" del "%TMP_RESP%" >nul 2>&1 -) else ( - curl -sf -m 1 -X POST -H "Content-Type: application/json" ^ - --data-binary @- "http://127.0.0.1:%PORT%/hooks/%HOOK_NAME%" >nul 2>&1 -) +if /i "%HOOK_NAME%"=="SessionStart" goto :capture_response +if /i "%HOOK_NAME%"=="Stop" goto :capture_response +curl -sf -m 1 -X POST -H "Content-Type: application/json" ^ + --data-binary @- "http://127.0.0.1:%PORT%/hooks/%HOOK_NAME%" >nul 2>&1 +goto :eof + +:capture_response +set "TMP_RESP=%TEMP%\\cce_hook_resp_%RANDOM%.txt" +curl -sf -m 2 -X POST -H "Content-Type: application/json" ^ + --data-binary @- "http://127.0.0.1:%PORT%/hooks/%HOOK_NAME%" ^ + > "%TMP_RESP%" 2>nul +if exist "%TMP_RESP%" type "%TMP_RESP%" +if exist "%TMP_RESP%" del "%TMP_RESP%" >nul 2>&1 exit /b 0 """ diff --git a/src/context_engine/memory/hooks.py b/src/context_engine/memory/hooks.py index a76bf19..9ac659b 100644 --- a/src/context_engine/memory/hooks.py +++ b/src/context_engine/memory/hooks.py @@ -383,9 +383,75 @@ async def handle_stop(request: web.Request) -> web.Response: except Exception: log.exception("Stop enqueue failed") return web.json_response({"ok": False}, status=202) + + # Memory nudge — summarise unrecorded activity so the agent can fire + # off quick record_decision / record_code_area calls before the session + # ends. The hook script captures Stop stdout (like SessionStart) and + # injects it into the model context. + nudge = _build_stop_nudge(conn, session_id) + if nudge: + return web.Response(text=nudge, content_type="text/plain") return web.json_response({"ok": True}) +def _build_stop_nudge(conn: sqlite3.Connection, session_id: str) -> str: + """Build a session-end nudge summarising unrecorded activity. + + Queries memory.db for tool_events (context_search calls) and decisions + recorded during this session. Returns a short directive if activity + was significant but recording was low. Returns "" if nothing to nudge. + """ + try: + # Count context_search tool calls this session + row = conn.execute( + "SELECT COUNT(*) AS n FROM tool_events " + "WHERE session_id = ? AND tool_name = 'context_search'", + (session_id,), + ).fetchone() + n_searches = int(row["n"]) if row else 0 + + # Count decisions recorded this session + row = conn.execute( + "SELECT COUNT(*) AS n FROM decisions WHERE session_id = ?", + (session_id,), + ).fetchone() + n_decisions = int(row["n"]) if row else 0 + + # Count code_areas recorded this session + row = conn.execute( + "SELECT COUNT(*) AS n FROM code_areas WHERE session_id = ?", + (session_id,), + ).fetchone() + n_code_areas = int(row["n"]) if row else 0 + + # Only nudge if there was real activity (4+ searches) and + # recording was sparse. + if n_searches < 4: + return "" + parts: list[str] = [] + if n_decisions == 0: + parts.append( + "If you chose a library, resolved an ambiguity, or " + "established a convention, call " + 'record_decision(decision="...", reason="...") now.' + ) + if n_code_areas == 0: + parts.append( + "If you modified or traced non-obvious code, call " + 'record_code_area(file_path="...", description="...") now.' + ) + if not parts: + return "" + header = ( + f"[CCE session end] {n_searches} searches, " + f"{n_decisions} decision(s), {n_code_areas} code area(s) recorded." + ) + return header + "\n" + "\n".join(parts) + except Exception: + log.debug("Stop nudge query failed", exc_info=True) + return "" + + async def handle_session_end(request: web.Request) -> web.Response: data = await _read_json(request) session_id = data.get("session_id") diff --git a/tests/integration/test_mcp_server.py b/tests/integration/test_mcp_server.py index d95c800..93234d0 100644 --- a/tests/integration/test_mcp_server.py +++ b/tests/integration/test_mcp_server.py @@ -30,6 +30,9 @@ def _make_server(tmp_path): # _record_bucket already guards on this; tests don't need a real db. server._memory_conn = None server._storage_base = tmp_path + # Memory nudge state (see ContextEngineMCP.__init__) + server._searches_since_last_decision = 0 + server._has_recorded_decision = False return server @@ -176,6 +179,9 @@ async def fake_retrieve(query, top_k=10, confidence_threshold=0.0, server._compressor.compress = AsyncMock(return_value=[stub_chunk]) server._session_capture = MagicMock() server._session_capture.touch_files = MagicMock() + server._session_capture.get_session_snapshot = MagicMock(return_value={ + "decisions": [], "code_areas": [], "touched_files": {}, + }) server._persist_current_session = MagicMock() server._record = MagicMock() server._append_audit_log = MagicMock() @@ -208,3 +214,62 @@ async def test_context_search_no_note_when_nothing_dropped(tmp_path): result = await server._handle_context_search({"query": "find something", "top_k": 5}) text = result[0].text assert "lower-confidence results omitted" not in text + + +@pytest.mark.asyncio +async def test_decision_nudge_fires_after_threshold(tmp_path): + """After 4 context_search calls with no record_decision, the nudge appears.""" + server = _make_search_server(tmp_path, dropped_low_value=0) + # First 3 searches: no nudge + for _ in range(3): + result = await server._handle_context_search({"query": "q", "top_k": 5}) + assert "[CCE]" not in result[0].text + # 4th search: nudge fires + result = await server._handle_context_search({"query": "q", "top_k": 5}) + assert "[CCE] 4 context searches" in result[0].text + assert "record_decision" in result[0].text + + +@pytest.mark.asyncio +async def test_decision_nudge_resets_after_record(tmp_path): + """Recording a decision resets the search counter.""" + server = _make_search_server(tmp_path, dropped_low_value=0) + # Trigger nudge + for _ in range(4): + await server._handle_context_search({"query": "q", "top_k": 5}) + # Simulate record_decision resetting state + server._searches_since_last_decision = 0 + server._has_recorded_decision = True + # Next search: no nudge (counter reset, higher threshold now) + result = await server._handle_context_search({"query": "q", "top_k": 5}) + assert "[CCE]" not in result[0].text + + +@pytest.mark.asyncio +async def test_code_area_nudge_fires_with_unrecorded_files(tmp_path): + """Nudge appears when 3+ files are touched but none recorded as code areas.""" + server = _make_search_server(tmp_path, dropped_low_value=0) + # Simulate 3 touched files with no code_areas recorded + server._session_capture.get_session_snapshot = MagicMock(return_value={ + "decisions": [], "code_areas": [], + "touched_files": {"a.py": 1, "b.py": 2, "c.py": 1}, + }) + result = await server._handle_context_search({"query": "q", "top_k": 5}) + assert "files explored but not recorded" in result[0].text + assert "record_code_area" in result[0].text + + +@pytest.mark.asyncio +async def test_code_area_nudge_excludes_recorded_files(tmp_path): + """Files that were record_code_area'd don't count toward the threshold.""" + server = _make_search_server(tmp_path, dropped_low_value=0) + server._session_capture.get_session_snapshot = MagicMock(return_value={ + "decisions": [], "touched_files": {"a.py": 1, "b.py": 2, "c.py": 1}, + "code_areas": [ + {"file_path": "a.py", "description": "x", "timestamp": 0}, + {"file_path": "b.py", "description": "y", "timestamp": 0}, + ], + }) + result = await server._handle_context_search({"query": "q", "top_k": 5}) + # Only 1 unrecorded file (c.py) — below threshold + assert "[CCE]" not in result[0].text or "files explored" not in result[0].text