diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index cc7a5e839..ade252623 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,7 +1,6 @@ # Changelog ## [Unreleased] - ### Added - Added `hindsight.mentalModelsEnabled`, `hindsight.mentalModelAutoSeed`, `hindsight.mentalModelRefreshIntervalMs`, and `hindsight.mentalModelMaxRenderChars` settings to control curated Hindsight mental-model activation, seeding, refresh cadence, and prompt render budget @@ -12,10 +11,14 @@ ### Changed +- Changed `/memory clear` and `/memory enqueue` to apply only to the current agent session’s Hindsight cache instead of all live Hindsight sessions - Changed the prompt assembly order so `` blocks are appended before `` recall blocks in developer instructions ### Fixed +- Fixed Hindsight memory prompt injection and recall/retain tool execution to resolve against the active session state, preventing context from an unrelated session from being used +- Fixed subagent `/task` sessions to persist memories into the parent agent’s Hindsight bank by explicit parent state wiring +- Fixed per-session memory retention behavior when switching or resuming sessions by rekeying Hindsight state and resetting conversation-tracking counters so first-turn recall and nth-turn retain cadence no longer leak across conversations - Fixed the first-turn startup race so `` appears in the opening system prompt when mental-model loading is enabled - Fixed retention hygiene by stripping `` blocks from retained content to prevent curated summaries from feeding back into future memory writes - Fixed `` rendering to honor the configured character budget and truncate with an explicit truncation marker when the snapshot exceeds limits diff --git a/packages/coding-agent/src/hindsight/backend.ts b/packages/coding-agent/src/hindsight/backend.ts index 470766de5..2d00dee27 100644 --- a/packages/coding-agent/src/hindsight/backend.ts +++ b/packages/coding-agent/src/hindsight/backend.ts @@ -2,10 +2,9 @@ * Hindsight memory backend. * * Wires the per-session lifecycle (recall on first turn, retain every Nth - * agent_end, etc.) on top of the AgentSession event stream. State for each - * live session lives in a module-level Map keyed by session id; the tool - * factories read from this map at execute time so they can fail closed when - * the backend isn't started for a given session. + * agent_end, etc.) on top of the AgentSession event stream. Hindsight runtime + * state is owned by the AgentSession so lifetime follows the actual domain + * owner instead of a parallel session-id registry. */ import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; @@ -13,74 +12,11 @@ import { logger } from "@oh-my-pi/pi-utils"; import type { Settings } from "../config/settings"; import type { MemoryBackend, MemoryBackendStartOptions } from "../memory-backend/types"; import type { AgentSession } from "../session/agent-session"; -import { type BankScope, computeBankScope, ensureBankMission } from "./bank"; -import { createHindsightClient, type HindsightApi } from "./client"; -import { type HindsightConfig, isHindsightConfigured, loadHindsightConfig } from "./config"; -import { - composeRecallQuery, - formatCurrentTime, - formatMemories, - type HindsightMessage, - prepareRetentionTranscript, - sliceLastTurnsByUserBoundary, - truncateRecallQuery, -} from "./content"; -import { - ensureMentalModels, - loadMentalModelsBlock, - MENTAL_MODEL_FIRST_TURN_DEADLINE_MS, - resolveSeedsForScope, -} from "./mental-models"; -import { clearRetainQueueForTest, flushAllRetainQueues, flushSessionQueue } from "./retain-queue"; -import { extractMessages } from "./transcript"; - -/** - * Per-session runtime state. One entry per live session id. - * - * `lastRetainedTurn` tracks the user-turn count at which we last retained, so - * `agent_end` only fires `retain` every `retainEveryNTurns` turns. - * - * `lastRecallSnippet` is the most-recent recall block; `buildDeveloperInstructions` - * folds it into the system prompt so the LLM sees memories injected without - * paying a recall round-trip on every prompt rebuild. - */ -export interface HindsightSessionState { - client: HindsightApi; - bankId: string; - /** Tags applied to every retain — non-empty in per-project-tagged mode. */ - retainTags?: string[]; - /** Tag filter applied to every recall/reflect — non-empty in per-project-tagged mode. */ - recallTags?: string[]; - recallTagsMatch?: "any" | "all" | "any_strict" | "all_strict"; - config: HindsightConfig; - session: AgentSession; - missionsSet: Set; - lastRetainedTurn: number; - hasRecalledForFirstTurn: boolean; - lastRecallSnippet?: string; - /** Cached `` block injected into developer instructions. */ - mentalModelsSnippet?: string; - /** When the cached snippet was last refreshed; gates the agent_end re-list. */ - mentalModelsLoadedAt?: number; - /** - * In-flight ensure+load promise. `beforeAgentStartPrompt` awaits this on - * the first turn so the MM block lands in the system prompt before the - * LLM generates, even though `start()` returns before the load completes. - */ - mentalModelsLoadPromise?: Promise; - unsubscribe?: () => void; - /** - * When set, this entry is a subagent alias that reuses the parent's bank, - * scope, config, client, and missionsSet. Aliases skip auto-recall and - * auto-retain — those run on the parent only — but the recall/retain/reflect - * tools resolve via the alias so they persist to the same bank as the - * parent. Iteration sites (`enqueue`, `buildDeveloperInstructions`) skip - * aliases to avoid double-counting the shared state. - */ - aliasOf?: HindsightSessionState; -} - -const STATE_BY_SESSION_ID = new Map(); +import { computeBankScope } from "./bank"; +import { createHindsightClient } from "./client"; +import { isHindsightConfigured, loadHindsightConfig } from "./config"; +import type { HindsightMessage } from "./content"; +import { HindsightSessionState } from "./state"; const STATIC_INSTRUCTIONS = [ "# Memory", @@ -93,245 +29,12 @@ const STATIC_INSTRUCTIONS = [ "", ].join("\n"); -/** Public accessor for session-scoped Hindsight state (used by tools). */ -export function getHindsightSessionState(sessionId: string): HindsightSessionState | undefined { - return STATE_BY_SESSION_ID.get(sessionId); +/** Reload the active session's mental-model cache and prompt. */ +export async function reloadMentalModelsForSession(session: AgentSession): Promise { + const state = session.getHindsightSessionState(); + if (!state) return false; + return await state.reloadMentalModels(); } - -/** Test-only: register a synthetic session state. Pair with `clearHindsightSessionStateForTest`. */ -export function setHindsightSessionStateForTest(sessionId: string, state: HindsightSessionState): void { - STATE_BY_SESSION_ID.set(sessionId, state); -} - -/** Test-only: drop every registered session state and release subscribed listeners. */ -export function clearHindsightSessionStateForTest(): void { - for (const state of STATE_BY_SESSION_ID.values()) state.unsubscribe?.(); - STATE_BY_SESSION_ID.clear(); - clearRetainQueueForTest(); -} - -/** - * Pick a top-level (non-alias) state. Subagent aliases reuse the parent's - * state, so when wiring a new subagent we need the originating primary entry - * to copy bank/scope/config/missionsSet from. Returns the most recently - * registered primary; with one top-level session per process this is the - * correct one. Returns undefined when no primary state has been registered. - */ -function pickPrimaryState(): HindsightSessionState | undefined { - let result: HindsightSessionState | undefined; - for (const state of STATE_BY_SESSION_ID.values()) { - if (state.aliasOf) continue; - result = state; - } - return result; -} - -interface RecallOutcome { - context: string | null; - ok: boolean; -} - -async function recallForContext( - state: HindsightSessionState, - query: string, - signal?: AbortSignal, -): Promise { - const { client, bankId, recallTags, recallTagsMatch, config } = state; - try { - const response = await client.recall(bankId, query, { - budget: config.recallBudget, - maxTokens: config.recallMaxTokens, - types: config.recallTypes.length > 0 ? config.recallTypes : undefined, - tags: recallTags, - tagsMatch: recallTagsMatch, - }); - if (signal?.aborted) return { context: null, ok: false }; - const results = response.results ?? []; - if (results.length === 0) return { context: null, ok: true }; - const formatted = formatMemories(results); - const block = `\n${config.recallPromptPreamble}\nCurrent time: ${formatCurrentTime()} UTC\n\n${formatted}\n`; - return { context: block, ok: true }; - } catch (err) { - if (config.debug) { - logger.debug("Hindsight: recall failed", { bankId, error: String(err) }); - } - return { context: null, ok: false }; - } -} - -async function retainSession( - state: HindsightSessionState, - sessionId: string, - messages: HindsightMessage[], -): Promise { - const { client, bankId, retainTags, config, missionsSet } = state; - const retainFullWindow = config.retainMode === "full-session"; - - let target: HindsightMessage[]; - let documentId: string; - - if (retainFullWindow) { - target = messages; - documentId = sessionId; - } else { - const windowTurns = config.retainEveryNTurns + config.retainOverlapTurns; - target = sliceLastTurnsByUserBoundary(messages, windowTurns); - documentId = `${sessionId}-${Date.now()}`; - } - - const { transcript } = prepareRetentionTranscript(target, true); - if (!transcript) return; - - await ensureBankMission(client, bankId, config, missionsSet); - await client.retain(bankId, transcript, { - documentId, - context: config.retainContext, - metadata: { session_id: sessionId }, - tags: retainTags, - async: true, - }); -} - -async function maybeRetainOnAgentEnd(state: HindsightSessionState): Promise { - if (!state.config.autoRetain) return; - const messages = extractMessages(state.session.sessionManager); - if (messages.length === 0) return; - const userTurns = messages.filter(m => m.role === "user").length; - if (userTurns - state.lastRetainedTurn < state.config.retainEveryNTurns) return; - - const sessionId = state.session.sessionId; - if (!sessionId) return; - - try { - await retainSession(state, sessionId, messages); - state.lastRetainedTurn = userTurns; - if (state.config.debug) { - logger.debug("Hindsight: auto-retain succeeded", { - sessionId, - bankId: state.bankId, - userTurns, - messages: messages.length, - }); - } - } catch (err) { - logger.warn("Hindsight: auto-retain failed", { - sessionId, - bankId: state.bankId, - error: String(err), - }); - } -} - -async function maybeRecallOnAgentStart(state: HindsightSessionState): Promise { - if (!state.config.autoRecall || state.hasRecalledForFirstTurn) return; - const messages = extractMessages(state.session.sessionManager); - const lastUser = [...messages].reverse().find(m => m.role === "user"); - if (!lastUser) return; - - const query = composeRecallQuery(lastUser.content, messages, state.config.recallContextTurns); - const truncated = truncateRecallQuery(query, lastUser.content, state.config.recallMaxQueryChars); - const { context, ok } = await recallForContext(state, truncated); - if (!ok) return; - - state.hasRecalledForFirstTurn = true; - if (!context) return; - - state.lastRecallSnippet = context; - try { - await state.session.refreshBaseSystemPrompt(); - } catch (err) { - logger.debug("Hindsight: refreshBaseSystemPrompt after recall failed", { error: String(err) }); - } -} - -/** - * Run the seed-and-load pipeline once for a top-level session, populating - * `state.mentalModelsSnippet` and refreshing the base system prompt so the - * `` block lands in the next prompt build. - * - * The first turn races: `start()` returns before this resolves. The race is - * covered in `beforeAgentStartPrompt` by awaiting `mentalModelsLoadPromise` - * with a hard deadline. - */ -async function runMentalModelLoad(state: HindsightSessionState, scope: BankScope): Promise { - if (!state.config.mentalModelsEnabled) return; - - // Seeding is opt-in (`hindsight.mentalModelAutoSeed`). Default behaviour is - // read-only: we surface whatever models the operator has curated on the - // bank, but we do NOT POST to create new ones unless they explicitly - // asked. `/memory mm seed` remains the explicit-write entry point. - if (state.config.mentalModelAutoSeed) { - const seeds = resolveSeedsForScope(scope, state.config.scoping); - if (seeds.length > 0) { - await ensureMentalModels(state.client, state.bankId, seeds, state.config.debug); - } - } - - await refreshMentalModelsSnippet(state); - try { - await state.session.refreshBaseSystemPrompt(); - } catch (err) { - logger.debug("Hindsight: refreshBaseSystemPrompt after MM load failed", { error: String(err) }); - } -} - -async function refreshMentalModelsSnippet(state: HindsightSessionState): Promise { - const snippet = await loadMentalModelsBlock(state.client, state.bankId, state.config.mentalModelMaxRenderChars); - state.mentalModelsSnippet = snippet; - state.mentalModelsLoadedAt = Date.now(); -} - -/** - * Public hook for `/memory mm reload` and the `agent_end` cache TTL. Re-pulls - * the list and updates the cached snippet; safe to call concurrently (the - * promise is not memoised — each call is a discrete refresh). - */ -export async function reloadMentalModelsForSession(sessionId: string): Promise { - const state = STATE_BY_SESSION_ID.get(sessionId); - if (!state || state.aliasOf) return false; - if (!state.config.mentalModelsEnabled) return false; - await refreshMentalModelsSnippet(state); - try { - await state.session.refreshBaseSystemPrompt(); - } catch (err) { - logger.debug("Hindsight: refreshBaseSystemPrompt after MM reload failed", { error: String(err) }); - } - return true; -} - -function attachSessionListeners(state: HindsightSessionState): void { - const sessionId = state.session.sessionId; - const unsubscribe = state.session.subscribe(event => { - if (event.type === "agent_start") { - void maybeRecallOnAgentStart(state); - } else if (event.type === "agent_end") { - void maybeRetainOnAgentEnd(state); - // Drain any queued tool-initiated retain calls now that the turn - // is settled. The queue is also debounced/size-bounded, but - // flushing here keeps the bank fresh between turns. - if (sessionId) void flushSessionQueue(sessionId); - // MM TTL refresh: re-list once we're past the cache deadline. List - // is cheap (no reflect call); the LLM doesn't see this happen. - if ( - state.config.mentalModelsEnabled && - state.mentalModelsLoadedAt !== undefined && - Date.now() - state.mentalModelsLoadedAt >= state.config.mentalModelRefreshIntervalMs - ) { - void refreshMentalModelsSnippet(state).then(async () => { - try { - await state.session.refreshBaseSystemPrompt(); - } catch (err) { - logger.debug("Hindsight: refreshBaseSystemPrompt after MM TTL reload failed", { - error: String(err), - }); - } - }); - } - } - }); - state.unsubscribe = unsubscribe; -} - export const hindsightBackend: MemoryBackend = { id: "hindsight", @@ -345,23 +48,25 @@ export const hindsightBackend: MemoryBackend = { // with the parent — running them per subagent would double-recall and // pollute the bank with internal exploration transcripts. if (options.taskDepth > 0) { - const parent = pickPrimaryState(); + const parent = options.parentHindsightSessionState; if (!parent) return; - const previous = STATE_BY_SESSION_ID.get(sessionId); - previous?.unsubscribe?.(); - STATE_BY_SESSION_ID.set(sessionId, { - client: parent.client, - bankId: parent.bankId, - retainTags: parent.retainTags, - recallTags: parent.recallTags, - recallTagsMatch: parent.recallTagsMatch, - config: parent.config, - session, - missionsSet: parent.missionsSet, - lastRetainedTurn: 0, - hasRecalledForFirstTurn: true, - aliasOf: parent, - }); + const previous = session.setHindsightSessionState( + new HindsightSessionState({ + sessionId, + client: parent.client, + bankId: parent.bankId, + retainTags: parent.retainTags, + recallTags: parent.recallTags, + recallTagsMatch: parent.recallTagsMatch, + config: parent.config, + session, + missionsSet: parent.missionsSet, + lastRetainedTurn: 0, + hasRecalledForFirstTurn: true, + aliasOf: parent, + }), + ); + previous?.dispose(); return; } @@ -374,7 +79,8 @@ export const hindsightBackend: MemoryBackend = { const client = createHindsightClient(config); const scope = computeBankScope(config, session.sessionManager.getCwd()); - const state: HindsightSessionState = { + const state = new HindsightSessionState({ + sessionId, client, bankId: scope.bankId, retainTags: scope.retainTags, @@ -385,42 +91,33 @@ export const hindsightBackend: MemoryBackend = { missionsSet: new Set(), lastRetainedTurn: 0, hasRecalledForFirstTurn: false, - }; + }); - // Cleanup any stale state for this session id (defensive — prevents leaks + // Cleanup any stale state for this session (defensive — prevents leaks // when a session is reused without going through dispose). - const previous = STATE_BY_SESSION_ID.get(sessionId); - previous?.unsubscribe?.(); - - STATE_BY_SESSION_ID.set(sessionId, state); - attachSessionListeners(state); + const previous = session.setHindsightSessionState(state); + previous?.dispose(); + state.attachSessionListeners(); // Kick off mental-model bootstrap. Resolves asynchronously; the first // turn races and is covered in `beforeAgentStartPrompt` via // `mentalModelsLoadPromise`. Subsequent turns see the populated cache // because `runMentalModelLoad` calls `refreshBaseSystemPrompt`. if (config.mentalModelsEnabled) { - state.mentalModelsLoadPromise = runMentalModelLoad(state, scope).catch(err => { + state.mentalModelsLoadPromise = state.runMentalModelLoad(scope).catch(err => { logger.debug("Hindsight: mental-model bootstrap failed", { bankId: state.bankId, error: String(err) }); }); } }, - async buildDeveloperInstructions(_agentDir, settings): Promise { + async buildDeveloperInstructions(_agentDir, settings, session): Promise { const config = loadHindsightConfig(settings); if (!isHindsightConfigured(config)) return undefined; - // Pick the active session-scoped snippets, if any. We can't know the - // caller's session id here (the local backend has the same - // limitation), but with a single top-level session per process the - // freshest snippet across all states is the correct one. - let recallSnippet: string | undefined; - let mentalModelsSnippet: string | undefined; - for (const state of STATE_BY_SESSION_ID.values()) { - if (state.aliasOf) continue; - if (state.lastRecallSnippet) recallSnippet = state.lastRecallSnippet; - if (state.mentalModelsSnippet) mentalModelsSnippet = state.mentalModelsSnippet; - } + const state = session?.getHindsightSessionState(); + const primary = state?.aliasOf ?? state; + const recallSnippet = primary?.lastRecallSnippet; + const mentalModelsSnippet = primary?.mentalModelsSnippet; // Order: static instructions → mental models (stable, curated) → recall // (volatile per turn). Stable context first so the LLM's prior is @@ -432,104 +129,48 @@ export const hindsightBackend: MemoryBackend = { }, async beforeAgentStartPrompt(session: AgentSession, promptText: string): Promise { - const sessionId = session.sessionId; - if (!sessionId) return undefined; - const state = STATE_BY_SESSION_ID.get(sessionId); + const state = session.getHindsightSessionState(); if (!state) return undefined; - // Race-cover the first-turn mental-model bootstrap. `start()` returns - // before MMs are seeded/loaded; without this await the very first - // system prompt (built from `buildDeveloperInstructions` at sdk.ts) is - // already locked in by the time MMs land, so the LLM misses the - // `` block on turn one. Awaiting here gives the load a - // hard deadline; on completion `runMentalModelLoad` has already called - // `refreshBaseSystemPrompt`, so the rebuilt base prompt picked up by - // `#buildSystemPromptForAgentStart` (which reads `#baseSystemPrompt` - // AFTER this hook returns) contains the MM block. - if ( - state.config.mentalModelsEnabled && - state.mentalModelsLoadPromise && - state.mentalModelsLoadedAt === undefined - ) { - await Promise.race([state.mentalModelsLoadPromise, Bun.sleep(MENTAL_MODEL_FIRST_TURN_DEADLINE_MS)]); - } - - if (!state.config.autoRecall || state.hasRecalledForFirstTurn) return undefined; - - const latestPrompt = promptText.trim(); - if (!latestPrompt) return undefined; - - const history = extractMessages(session.sessionManager); - const queryMessages = [...history, { role: "user", content: latestPrompt }]; - const query = composeRecallQuery(latestPrompt, queryMessages, state.config.recallContextTurns); - const truncated = truncateRecallQuery(query, latestPrompt, state.config.recallMaxQueryChars); - const { context, ok } = await recallForContext(state, truncated); - if (!ok) return undefined; - - state.hasRecalledForFirstTurn = true; - if (!context) return undefined; - - state.lastRecallSnippet = context; - return context; + return await state.beforeAgentStartPrompt(promptText); }, - async clear(_agentDir, _cwd): Promise { - // Hindsight memory is server-side. The local cache (per-session WeakMap- - // equivalent) is what we can wipe — operators who want to delete the - // upstream bank should use the Hindsight UI / `deleteBank` directly. - // Drain pending tool-initiated retains first so we don't lose them. - await flushAllRetainQueues(); - for (const state of STATE_BY_SESSION_ID.values()) { - state.unsubscribe?.(); - } - STATE_BY_SESSION_ID.clear(); + async clear(_agentDir, _cwd, session): Promise { + // Hindsight memory is server-side. The local cache is what we can wipe — + // operators who want to delete the upstream bank should use the Hindsight + // UI / `deleteBank` directly. Drain pending tool-initiated retains first + // so we don't lose them. + const state = session?.getHindsightSessionState(); + if (state) await state.flushRetainQueue(); + const previous = session?.setHindsightSessionState(undefined); + previous?.dispose(); logger.warn( "Hindsight memory is server-side; only the local recall cache was cleared. " + "Delete the Hindsight bank from the UI to wipe upstream state.", ); }, - async enqueue(_agentDir, _cwd): Promise { - // Force an immediate retain across every active session, including - // the queued tool-initiated retains that haven't flushed yet. - await flushAllRetainQueues(); - for (const state of STATE_BY_SESSION_ID.values()) { - if (state.aliasOf) continue; - const sessionId = state.session.sessionId; - if (!sessionId) continue; - const messages = extractMessages(state.session.sessionManager); - if (messages.length === 0) continue; - try { - await retainSession(state, sessionId, messages); - state.lastRetainedTurn = messages.filter(m => m.role === "user").length; - } catch (err) { - logger.warn("Hindsight: forced retain failed", { - sessionId, - bankId: state.bankId, - error: String(err), - }); - } - } + async enqueue(_agentDir, _cwd, session): Promise { + const state = session?.getHindsightSessionState(); + const primary = state?.aliasOf ? undefined : state; + if (!primary) return; + await primary.flushRetainQueue(); + await primary.forceRetainCurrentSession(); }, - async preCompactionContext(messages: AgentMessage[], settings: Settings): Promise { + async preCompactionContext( + messages: AgentMessage[], + settings: Settings, + session?: AgentSession, + ): Promise { const config = loadHindsightConfig(settings); if (!isHindsightConfigured(config)) return undefined; - // Find the most recent state — we don't have a session id here either, so - // pick the freshest registered session. - let state: HindsightSessionState | undefined; - for (const candidate of STATE_BY_SESSION_ID.values()) state = candidate; + const state = session?.getHindsightSessionState(); if (!state) return undefined; const flat = flattenMessagesForRecall(messages); - const lastUser = [...flat].reverse().find(m => m.role === "user"); - if (!lastUser) return undefined; - - const query = composeRecallQuery(lastUser.content, flat, state.config.recallContextTurns); - const truncated = truncateRecallQuery(query, lastUser.content, state.config.recallMaxQueryChars); - const { context } = await recallForContext(state, truncated); - return context ?? undefined; + return await state.recallForCompaction(flat); }, }; diff --git a/packages/coding-agent/src/hindsight/index.ts b/packages/coding-agent/src/hindsight/index.ts index 483ead369..d3a99546b 100644 --- a/packages/coding-agent/src/hindsight/index.ts +++ b/packages/coding-agent/src/hindsight/index.ts @@ -4,4 +4,5 @@ export * from "./client"; export * from "./config"; export * from "./content"; export * from "./mental-models"; +export * from "./state"; export * from "./transcript"; diff --git a/packages/coding-agent/src/hindsight/retain-queue.ts b/packages/coding-agent/src/hindsight/retain-queue.ts deleted file mode 100644 index 779cc4795..000000000 --- a/packages/coding-agent/src/hindsight/retain-queue.ts +++ /dev/null @@ -1,166 +0,0 @@ -/** - * Global, debounced batch queue for tool-initiated `retain` calls. - * - * The `retain` tool used to block on a single-item HTTP round trip per - * invocation. Now it pushes onto a per-session queue and returns immediately; - * a flush fires when: - * 1. the queue reaches `FLUSH_BATCH_SIZE`, or - * 2. `FLUSH_INTERVAL_MS` elapses since the queue first became non-empty. - * - * On batch failure we surface a UI-only notice via `session.emitNotice` — - * a single yellow "Hindsight: memory retention failed …" line in the TUI. - * The LLM is NOT told; the agent already received "Memory queued" and has - * moved on. This is purely so the user knows their facts didn't persist. - * - * Auto-retain (`retainSession` in backend.ts) is intentionally NOT routed - * through this queue — it submits a full transcript as one large item and - * already runs `async: true` server-side. - */ - -import { logger } from "@oh-my-pi/pi-utils"; -import { getHindsightSessionState, type HindsightSessionState } from "./backend"; -import { ensureBankMission } from "./bank"; -import type { MemoryItemInput } from "./client"; - -const FLUSH_BATCH_SIZE = 16; -const FLUSH_INTERVAL_MS = 5_000; - -interface PendingItem { - content: string; - context?: string; -} - -interface SessionQueue { - items: PendingItem[]; - timer?: NodeJS.Timeout; - /** Currently in-flight flush; subsequent flushes await it before running. */ - flushing?: Promise; -} - -const QUEUES = new Map(); - -/** Push a memory item onto the session's retain queue. Returns immediately. */ -export function enqueueRetain(sessionId: string, content: string, context?: string): void { - const queue = QUEUES.get(sessionId) ?? createQueue(sessionId); - queue.items.push({ content, context }); - - if (queue.items.length >= FLUSH_BATCH_SIZE) { - void flushSessionQueue(sessionId); - return; - } - if (!queue.timer) { - queue.timer = setTimeout(() => { - void flushSessionQueue(sessionId); - }, FLUSH_INTERVAL_MS); - // Don't pin the event loop alive just for a pending retain flush. - queue.timer.unref?.(); - } -} - -/** Flush a single session's queue. Safe to call when empty or already in flight. */ -export async function flushSessionQueue(sessionId: string): Promise { - const queue = QUEUES.get(sessionId); - if (!queue) return; - - if (queue.timer) { - clearTimeout(queue.timer); - queue.timer = undefined; - } - - if (queue.flushing) { - // Coalesce: wait for the in-flight flush, then drain anything that - // landed after it started so we don't strand items. - await queue.flushing; - if (queue.items.length > 0) { - await flushSessionQueue(sessionId); - } - return; - } - - if (queue.items.length === 0) { - QUEUES.delete(sessionId); - return; - } - - const items = queue.items.splice(0); - const flushPromise = doFlush(sessionId, items); - queue.flushing = flushPromise; - try { - await flushPromise; - } finally { - queue.flushing = undefined; - if (queue.items.length === 0 && !queue.timer) { - QUEUES.delete(sessionId); - } - } -} - -/** Flush every pending session queue. Called from `clear`/`enqueue` backend hooks. */ -export async function flushAllRetainQueues(): Promise { - const ids = [...QUEUES.keys()]; - await Promise.all(ids.map(id => flushSessionQueue(id))); -} - -/** Test helper: clear timers and pending items without triggering flushes. */ -export function clearRetainQueueForTest(): void { - for (const queue of QUEUES.values()) { - if (queue.timer) clearTimeout(queue.timer); - } - QUEUES.clear(); -} - -/** Test helper: peek at queued count for a session. */ -export function getRetainQueueDepthForTest(sessionId: string): number { - return QUEUES.get(sessionId)?.items.length ?? 0; -} - -async function doFlush(sessionId: string, items: PendingItem[]): Promise { - const state = getHindsightSessionState(sessionId); - if (!state) { - // Session went away before we could flush. We can't notify anyone, so - // log and drop — these are best-effort facts, not transactional writes. - logger.warn("Hindsight retain queue: session vanished, dropping batch", { - sessionId, - items: items.length, - }); - return; - } - - try { - await ensureBankMission(state.client, state.bankId, state.config, state.missionsSet); - const batch: MemoryItemInput[] = items.map(item => ({ - content: item.content, - context: item.context ?? state.config.retainContext, - metadata: { session_id: sessionId }, - tags: state.retainTags, - })); - await state.client.retainBatch(state.bankId, batch, { async: true }); - if (state.config.debug) { - logger.debug("Hindsight retain queue: batch flushed", { - sessionId, - bankId: state.bankId, - items: items.length, - }); - } - } catch (err) { - const errorText = err instanceof Error ? err.message : String(err); - logger.warn("Hindsight retain queue: batch flush failed", { - sessionId, - bankId: state.bankId, - items: items.length, - error: errorText, - }); - notifyRetainFailure(state, items.length, errorText); - } -} - -function notifyRetainFailure(state: HindsightSessionState, count: number, errorText: string): void { - const noun = count === 1 ? "memory" : "memories"; - state.session.emitNotice("warning", `Memory retention failed for ${count} ${noun}: ${errorText}`, "Hindsight"); -} - -function createQueue(sessionId: string): SessionQueue { - const queue: SessionQueue = { items: [] }; - QUEUES.set(sessionId, queue); - return queue; -} diff --git a/packages/coding-agent/src/hindsight/state.ts b/packages/coding-agent/src/hindsight/state.ts new file mode 100644 index 000000000..afa3311c6 --- /dev/null +++ b/packages/coding-agent/src/hindsight/state.ts @@ -0,0 +1,469 @@ +import { logger } from "@oh-my-pi/pi-utils"; +import type { AgentSession } from "../session/agent-session"; +import { type BankScope, ensureBankMission } from "./bank"; +import type { HindsightApi, MemoryItemInput } from "./client"; +import type { HindsightConfig } from "./config"; +import { + composeRecallQuery, + formatCurrentTime, + formatMemories, + type HindsightMessage, + prepareRetentionTranscript, + sliceLastTurnsByUserBoundary, + truncateRecallQuery, +} from "./content"; +import { + ensureMentalModels, + loadMentalModelsBlock, + MENTAL_MODEL_FIRST_TURN_DEADLINE_MS, + resolveSeedsForScope, +} from "./mental-models"; +import { extractMessages } from "./transcript"; + +const RETAIN_FLUSH_BATCH_SIZE = 16; +const RETAIN_FLUSH_INTERVAL_MS = 5_000; + +interface PendingRetainItem { + content: string; + context?: string; +} + +interface RecallOutcome { + context: string | null; + ok: boolean; +} + +export interface HindsightSessionStateOptions { + /** Session id used for retain-queue metadata. */ + sessionId: string; + client: HindsightApi; + bankId: string; + /** Tags applied to every retain — non-empty in per-project-tagged mode. */ + retainTags?: string[]; + /** Tag filter applied to every recall/reflect — non-empty in per-project-tagged mode. */ + recallTags?: string[]; + recallTagsMatch?: "any" | "all" | "any_strict" | "all_strict"; + config: HindsightConfig; + session: AgentSession; + missionsSet: Set; + lastRetainedTurn?: number; + hasRecalledForFirstTurn?: boolean; + /** + * When set, this entry is a subagent alias that reuses the parent's bank, + * scope, config, client, and missionsSet. Aliases skip auto-recall and + * auto-retain — those run on the parent only — but the recall/retain/reflect + * tools resolve via the alias so they persist to the same bank as the parent. + */ + aliasOf?: HindsightSessionState; +} + +/** + * Debounced batch queue for tool-initiated `retain` calls owned by one + * Hindsight session state instance. + * + * Auto-retain (`HindsightSessionState.retainSession`) is intentionally not + * routed through this queue — it submits a full transcript as one large item + * and already runs `async: true` server-side. + */ +export class HindsightRetainQueue { + readonly #state: HindsightSessionState; + #items: PendingRetainItem[] = []; + #timer?: NodeJS.Timeout; + #flushing?: Promise; + #closed = false; + + constructor(state: HindsightSessionState) { + this.#state = state; + } + + get depth(): number { + return this.#items.length; + } + + enqueue(content: string, context?: string): void { + if (this.#closed) { + throw new Error("Hindsight retain queue is closed."); + } + this.#items.push({ content, context }); + + if (this.#items.length >= RETAIN_FLUSH_BATCH_SIZE) { + void this.flush(); + return; + } + if (!this.#timer) { + this.#timer = setTimeout(() => { + void this.flush(); + }, RETAIN_FLUSH_INTERVAL_MS); + // Don't pin the event loop alive just for a pending retain flush. + this.#timer.unref?.(); + } + } + + async flush(): Promise { + if (this.#timer) { + clearTimeout(this.#timer); + this.#timer = undefined; + } + + if (this.#flushing) { + // Coalesce: wait for the in-flight flush, then drain anything that + // landed after it started so we don't strand items. + await this.#flushing; + if (this.#items.length > 0) await this.flush(); + return; + } + + if (this.#items.length === 0) return; + + const items = this.#items.splice(0); + const flushPromise = this.#doFlush(items); + this.#flushing = flushPromise; + try { + await flushPromise; + } finally { + this.#flushing = undefined; + } + } + + dispose(): void { + this.#closed = true; + if (this.#timer) { + clearTimeout(this.#timer); + this.#timer = undefined; + } + this.#items = []; + } + + async #doFlush(items: PendingRetainItem[]): Promise { + const state = this.#state; + const sessionId = state.sessionId; + if (state.session.getHindsightSessionState() !== state) { + // Session went away before we could flush. We can't notify anyone, so + // log and drop — these are best-effort facts, not transactional writes. + logger.warn("Hindsight retain queue: session vanished, dropping batch", { + sessionId, + items: items.length, + }); + return; + } + + try { + await ensureBankMission(state.client, state.bankId, state.config, state.missionsSet); + const batch: MemoryItemInput[] = items.map(item => ({ + content: item.content, + context: item.context ?? state.config.retainContext, + metadata: { session_id: sessionId }, + tags: state.retainTags, + })); + await state.client.retainBatch(state.bankId, batch, { async: true }); + if (state.config.debug) { + logger.debug("Hindsight retain queue: batch flushed", { + sessionId, + bankId: state.bankId, + items: items.length, + }); + } + } catch (err) { + const errorText = err instanceof Error ? err.message : String(err); + logger.warn("Hindsight retain queue: batch flush failed", { + sessionId, + bankId: state.bankId, + items: items.length, + error: errorText, + }); + this.#notifyRetainFailure(items.length, errorText); + } + } + + #notifyRetainFailure(count: number, errorText: string): void { + const noun = count === 1 ? "memory" : "memories"; + this.#state.session.emitNotice( + "warning", + `Memory retention failed for ${count} ${noun}: ${errorText}`, + "Hindsight", + ); + } +} + +/** Per-session Hindsight runtime state owned by its AgentSession. */ +export class HindsightSessionState { + /** Session id used for retain-queue metadata. */ + sessionId: string; + client: HindsightApi; + bankId: string; + /** Tags applied to every retain — non-empty in per-project-tagged mode. */ + retainTags?: string[]; + /** Tag filter applied to every recall/reflect — non-empty in per-project-tagged mode. */ + recallTags?: string[]; + recallTagsMatch?: "any" | "all" | "any_strict" | "all_strict"; + config: HindsightConfig; + session: AgentSession; + missionsSet: Set; + lastRetainedTurn: number; + hasRecalledForFirstTurn: boolean; + lastRecallSnippet?: string; + /** Cached `` block injected into developer instructions. */ + mentalModelsSnippet?: string; + /** When the cached snippet was last refreshed; gates the agent_end re-list. */ + mentalModelsLoadedAt?: number; + /** + * In-flight ensure+load promise. `beforeAgentStartPrompt` awaits this on + * the first turn so the MM block lands in the system prompt before the + * LLM generates, even though `start()` returns before the load completes. + */ + mentalModelsLoadPromise?: Promise; + unsubscribe?: () => void; + /** Alias states delegate persistence config to a primary parent state. */ + aliasOf?: HindsightSessionState; + readonly retainQueue: HindsightRetainQueue; + + constructor(options: HindsightSessionStateOptions) { + this.sessionId = options.sessionId; + this.client = options.client; + this.bankId = options.bankId; + this.retainTags = options.retainTags; + this.recallTags = options.recallTags; + this.recallTagsMatch = options.recallTagsMatch; + this.config = options.config; + this.session = options.session; + this.missionsSet = options.missionsSet; + this.lastRetainedTurn = options.lastRetainedTurn ?? 0; + this.hasRecalledForFirstTurn = options.hasRecalledForFirstTurn ?? false; + this.aliasOf = options.aliasOf; + this.retainQueue = new HindsightRetainQueue(this); + } + + setSessionId(sessionId: string): void { + this.sessionId = sessionId; + } + + resetConversationTracking(): void { + this.lastRetainedTurn = 0; + this.hasRecalledForFirstTurn = false; + this.lastRecallSnippet = undefined; + } + + enqueueRetain(content: string, context?: string): void { + this.retainQueue.enqueue(content, context); + } + + async flushRetainQueue(): Promise { + await this.retainQueue.flush(); + } + + async recallForContext(query: string, signal?: AbortSignal): Promise { + try { + const response = await this.client.recall(this.bankId, query, { + budget: this.config.recallBudget, + maxTokens: this.config.recallMaxTokens, + types: this.config.recallTypes.length > 0 ? this.config.recallTypes : undefined, + tags: this.recallTags, + tagsMatch: this.recallTagsMatch, + }); + if (signal?.aborted) return { context: null, ok: false }; + const results = response.results ?? []; + if (results.length === 0) return { context: null, ok: true }; + const formatted = formatMemories(results); + const block = `\n${this.config.recallPromptPreamble}\nCurrent time: ${formatCurrentTime()} UTC\n\n${formatted}\n`; + return { context: block, ok: true }; + } catch (err) { + if (this.config.debug) { + logger.debug("Hindsight: recall failed", { bankId: this.bankId, error: String(err) }); + } + return { context: null, ok: false }; + } + } + + async retainSession(messages: HindsightMessage[]): Promise { + const retainFullWindow = this.config.retainMode === "full-session"; + let target: HindsightMessage[]; + let documentId: string; + + if (retainFullWindow) { + target = messages; + documentId = this.sessionId; + } else { + const windowTurns = this.config.retainEveryNTurns + this.config.retainOverlapTurns; + target = sliceLastTurnsByUserBoundary(messages, windowTurns); + documentId = `${this.sessionId}-${Date.now()}`; + } + + const { transcript } = prepareRetentionTranscript(target, true); + if (!transcript) return; + + await ensureBankMission(this.client, this.bankId, this.config, this.missionsSet); + await this.client.retain(this.bankId, transcript, { + documentId, + context: this.config.retainContext, + metadata: { session_id: this.sessionId }, + tags: this.retainTags, + async: true, + }); + } + + async maybeRetainOnAgentEnd(): Promise { + if (!this.config.autoRetain) return; + const messages = extractMessages(this.session.sessionManager); + if (messages.length === 0) return; + const userTurns = messages.filter(m => m.role === "user").length; + if (userTurns - this.lastRetainedTurn < this.config.retainEveryNTurns) return; + + try { + await this.retainSession(messages); + this.lastRetainedTurn = userTurns; + if (this.config.debug) { + logger.debug("Hindsight: auto-retain succeeded", { + sessionId: this.sessionId, + bankId: this.bankId, + userTurns, + messages: messages.length, + }); + } + } catch (err) { + logger.warn("Hindsight: auto-retain failed", { + sessionId: this.sessionId, + bankId: this.bankId, + error: String(err), + }); + } + } + + async forceRetainCurrentSession(): Promise { + const messages = extractMessages(this.session.sessionManager); + if (messages.length === 0) return; + try { + await this.retainSession(messages); + this.lastRetainedTurn = messages.filter(m => m.role === "user").length; + } catch (err) { + logger.warn("Hindsight: forced retain failed", { + sessionId: this.sessionId, + bankId: this.bankId, + error: String(err), + }); + } + } + + async maybeRecallOnAgentStart(): Promise { + if (!this.config.autoRecall || this.hasRecalledForFirstTurn) return; + const messages = extractMessages(this.session.sessionManager); + const lastUser = [...messages].reverse().find(m => m.role === "user"); + if (!lastUser) return; + + const query = composeRecallQuery(lastUser.content, messages, this.config.recallContextTurns); + const truncated = truncateRecallQuery(query, lastUser.content, this.config.recallMaxQueryChars); + const { context, ok } = await this.recallForContext(truncated); + if (!ok) return; + + this.hasRecalledForFirstTurn = true; + if (!context) return; + + this.lastRecallSnippet = context; + await this.#refreshBaseSystemPromptAfter("recall"); + } + + async beforeAgentStartPrompt(promptText: string): Promise { + if (this.config.mentalModelsEnabled && this.mentalModelsLoadPromise && this.mentalModelsLoadedAt === undefined) { + await Promise.race([this.mentalModelsLoadPromise, Bun.sleep(MENTAL_MODEL_FIRST_TURN_DEADLINE_MS)]); + } + + if (!this.config.autoRecall || this.hasRecalledForFirstTurn) return undefined; + + const latestPrompt = promptText.trim(); + if (!latestPrompt) return undefined; + + const history = extractMessages(this.session.sessionManager); + const queryMessages = [...history, { role: "user" as const, content: latestPrompt }]; + const query = composeRecallQuery(latestPrompt, queryMessages, this.config.recallContextTurns); + const truncated = truncateRecallQuery(query, latestPrompt, this.config.recallMaxQueryChars); + const { context, ok } = await this.recallForContext(truncated); + if (!ok) return undefined; + + this.hasRecalledForFirstTurn = true; + if (!context) return undefined; + + this.lastRecallSnippet = context; + return context; + } + + async recallForCompaction(messages: HindsightMessage[]): Promise { + const lastUser = [...messages].reverse().find(m => m.role === "user"); + if (!lastUser) return undefined; + + const query = composeRecallQuery(lastUser.content, messages, this.config.recallContextTurns); + const truncated = truncateRecallQuery(query, lastUser.content, this.config.recallMaxQueryChars); + const { context } = await this.recallForContext(truncated); + return context ?? undefined; + } + + async runMentalModelLoad(scope: BankScope): Promise { + if (!this.config.mentalModelsEnabled) return; + + // Seeding is opt-in (`hindsight.mentalModelAutoSeed`). Default behaviour is + // read-only: we surface whatever models the operator has curated on the + // bank, but we do NOT POST to create new ones unless they explicitly + // asked. `/memory mm seed` remains the explicit-write entry point. + if (this.config.mentalModelAutoSeed) { + const seeds = resolveSeedsForScope(scope, this.config.scoping); + if (seeds.length > 0) { + await ensureMentalModels(this.client, this.bankId, seeds, this.config.debug); + } + } + + await this.refreshMentalModelsSnippet(); + await this.#refreshBaseSystemPromptAfter("MM load"); + } + + async refreshMentalModelsSnippet(): Promise { + const snippet = await loadMentalModelsBlock(this.client, this.bankId, this.config.mentalModelMaxRenderChars); + this.mentalModelsSnippet = snippet; + this.mentalModelsLoadedAt = Date.now(); + } + + async reloadMentalModels(): Promise { + if (this.aliasOf) return false; + if (!this.config.mentalModelsEnabled) return false; + await this.refreshMentalModelsSnippet(); + await this.#refreshBaseSystemPromptAfter("MM reload"); + return true; + } + + attachSessionListeners(): void { + this.unsubscribe?.(); + this.unsubscribe = this.session.subscribe(event => { + if (event.type === "agent_start") { + void this.maybeRecallOnAgentStart(); + } else if (event.type === "agent_end") { + void this.maybeRetainOnAgentEnd(); + // Drain any queued tool-initiated retain calls now that the turn + // is settled. The queue is also debounced/size-bounded, but + // flushing here keeps the bank fresh between turns. + void this.flushRetainQueue(); + // MM TTL refresh: re-list once we're past the cache deadline. List + // is cheap (no reflect call); the LLM doesn't see this happen. + if ( + this.config.mentalModelsEnabled && + this.mentalModelsLoadedAt !== undefined && + Date.now() - this.mentalModelsLoadedAt >= this.config.mentalModelRefreshIntervalMs + ) { + void this.refreshMentalModelsSnippet().then(async () => { + await this.#refreshBaseSystemPromptAfter("MM TTL reload"); + }); + } + } + }); + } + + dispose(): void { + this.unsubscribe?.(); + this.unsubscribe = undefined; + this.retainQueue.dispose(); + } + + async #refreshBaseSystemPromptAfter(reason: "recall" | "MM load" | "MM reload" | "MM TTL reload"): Promise { + try { + await this.session.refreshBaseSystemPrompt(); + } catch (err) { + logger.debug(`Hindsight: refreshBaseSystemPrompt after ${reason} failed`, { error: String(err) }); + } + } +} diff --git a/packages/coding-agent/src/memory-backend/types.ts b/packages/coding-agent/src/memory-backend/types.ts index 67e83b082..848d7e341 100644 --- a/packages/coding-agent/src/memory-backend/types.ts +++ b/packages/coding-agent/src/memory-backend/types.ts @@ -9,6 +9,7 @@ import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; import type { ModelRegistry } from "../config/model-registry"; import type { Settings } from "../config/settings"; +import type { HindsightSessionState } from "../hindsight/state"; import type { AgentSession } from "../session/agent-session"; export type MemoryBackendId = "off" | "local" | "hindsight"; @@ -19,6 +20,7 @@ export interface MemoryBackendStartOptions { modelRegistry: ModelRegistry; agentDir: string; taskDepth: number; + parentHindsightSessionState?: HindsightSessionState; } export interface MemoryBackend { @@ -37,13 +39,17 @@ export interface MemoryBackend { * Markdown injected as the system-prompt append section. * Returned on every prompt rebuild via `refreshBaseSystemPrompt()`. */ - buildDeveloperInstructions(agentDir: string, settings: Settings): Promise; + buildDeveloperInstructions( + agentDir: string, + settings: Settings, + session?: AgentSession, + ): Promise; /** Wipe all persisted state for this backend (slash `/memory clear`). */ - clear(agentDir: string, cwd: string): Promise; + clear(agentDir: string, cwd: string, session?: AgentSession): Promise; /** Force consolidation/retain to happen now (slash `/memory enqueue`). */ - enqueue(agentDir: string, cwd: string): Promise; + enqueue(agentDir: string, cwd: string, session?: AgentSession): Promise; /** * Optional hook to inject a backend-specific block into the current turn's @@ -65,5 +71,9 @@ export interface MemoryBackend { * to inject nothing — the local backend takes this branch because its * summary is already part of the system prompt. */ - preCompactionContext?(messages: AgentMessage[], settings: Settings): Promise; + preCompactionContext?( + messages: AgentMessage[], + settings: Settings, + session?: AgentSession, + ): Promise; } diff --git a/packages/coding-agent/src/modes/controllers/command-controller.ts b/packages/coding-agent/src/modes/controllers/command-controller.ts index 492fa1fab..fb860310e 100644 --- a/packages/coding-agent/src/modes/controllers/command-controller.ts +++ b/packages/coding-agent/src/modes/controllers/command-controller.ts @@ -19,7 +19,6 @@ import { loadCustomShare } from "../../export/custom-share"; import type { CompactOptions } from "../../extensibility/extensions/types"; import { diffMentalModelContent, - getHindsightSessionState, type HindsightApi, type HindsightSessionState, loadHindsightConfig, @@ -583,7 +582,7 @@ export class CommandController { const backend = resolveMemoryBackend(this.ctx.settings); if (action === "view") { - const payload = await backend.buildDeveloperInstructions(agentDir, this.ctx.settings); + const payload = await backend.buildDeveloperInstructions(agentDir, this.ctx.settings, this.ctx.session); if (!payload) { this.ctx.showWarning("Memory payload is empty (memory backend off, disabled, or no memory available)."); return; @@ -600,7 +599,7 @@ export class CommandController { if (action === "reset" || action === "clear") { try { - await backend.clear(agentDir, this.ctx.sessionManager.getCwd()); + await backend.clear(agentDir, this.ctx.sessionManager.getCwd(), this.ctx.session); await this.ctx.session.refreshBaseSystemPrompt(); this.ctx.showStatus("Memory data cleared and system prompt refreshed."); } catch (error) { @@ -611,7 +610,7 @@ export class CommandController { if (action === "enqueue" || action === "rebuild") { try { - await backend.enqueue(agentDir, this.ctx.sessionManager.getCwd()); + await backend.enqueue(agentDir, this.ctx.sessionManager.getCwd(), this.ctx.session); this.ctx.showStatus("Memory consolidation enqueued."); } catch (error) { this.ctx.showError(`Memory enqueue failed: ${error instanceof Error ? error.message : String(error)}`); @@ -633,12 +632,7 @@ export class CommandController { const verb = parts[0]?.toLowerCase() ?? "list"; const arg = parts[1]; - const sessionId = this.ctx.session.sessionId; - if (!sessionId) { - this.ctx.showError("No active session."); - return; - } - const state = getHindsightSessionState(sessionId); + const state = this.ctx.session.getHindsightSessionState(); const primary = state && !state.aliasOf ? state : undefined; if (!primary) { this.ctx.showError("Hindsight backend is not active for this session."); @@ -668,7 +662,7 @@ export class CommandController { await this.#mmSeed(primary); return; case "reload": - await this.#mmReload(sessionId); + await this.#mmReload(primary); return; case "delete": case "remove": @@ -768,7 +762,7 @@ export class CommandController { // Reload the cache after a brief grace so the new content (if the refresh // completes synchronously on the server) flows into the system prompt. await Bun.sleep(500); - await reloadMentalModelsForSession(state.session.sessionId ?? ""); + await reloadMentalModelsForSession(state.session); } catch (error) { this.ctx.showError(`mm refresh failed: ${error instanceof Error ? error.message : String(error)}`); } @@ -852,8 +846,8 @@ export class CommandController { } } - async #mmReload(sessionId: string): Promise { - const ok = await reloadMentalModelsForSession(sessionId); + async #mmReload(state: HindsightSessionState): Promise { + const ok = await reloadMentalModelsForSession(state.session); if (ok) { this.ctx.showStatus("Mental-model cache reloaded."); } else { @@ -870,7 +864,7 @@ export class CommandController { } // Drop the cached snippet so the closing tag does not silently keep // stale content in the system prompt until the next agent_end TTL. - await reloadMentalModelsForSession(state.session.sessionId ?? ""); + await reloadMentalModelsForSession(state.session); this.ctx.showStatus(`Deleted mental model ${id} from bank ${state.bankId}.`); } catch (error) { this.ctx.showError(`mm delete failed: ${error instanceof Error ? error.message : String(error)}`); diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index 22cc2963d..67784737a 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -61,6 +61,7 @@ import { } from "./extensibility/extensions"; import { loadSkills as loadSkillsInternal, type Skill, type SkillWarning } from "./extensibility/skills"; import { type FileSlashCommand, loadSlashCommands as loadSlashCommandsInternal } from "./extensibility/slash-commands"; +import type { HindsightSessionState } from "./hindsight/state"; import { AgentProtocolHandler, ArtifactProtocolHandler, @@ -216,6 +217,8 @@ export interface CreateAgentSessionOptions { requireYieldTool?: boolean; /** Task recursion depth (for subagent sessions). Default: 0 */ taskDepth?: number; + /** Parent Hindsight state to alias for subagent memory tools. */ + parentHindsightSessionState?: HindsightSessionState; /** Pre-allocated agent identity for IRC routing. Default: "0-Main" for top-level, parentTaskPrefix-derived for sub. */ agentId?: string; /** Display name for the agent in IRC. Default: "main" or "sub". */ @@ -968,6 +971,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} trackEvalExecution: (execution, abortController) => session ? session.trackEvalExecution(execution, abortController) : execution, getSessionId: () => sessionManager.getSessionId?.() ?? null, + getHindsightSessionState: () => session?.getHindsightSessionState(), getAgentId: () => resolvedAgentId, getToolByName: name => session?.getToolByName(name), agentRegistry, @@ -1335,7 +1339,11 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const promptTools = buildSystemPromptToolMetadata(tools, { search_tool_bm25: { description: renderSearchToolBm25Description(discoverableMCPTools) }, }); - const memoryInstructions = await resolveMemoryBackend(settings).buildDeveloperInstructions(agentDir, settings); + const memoryInstructions = await resolveMemoryBackend(settings).buildDeveloperInstructions( + agentDir, + settings, + session, + ); // Build combined append prompt: memory instructions + MCP server instructions const serverInstructions = mcpManager?.getServerInstructions(); @@ -1755,6 +1763,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} modelRegistry, agentDir, taskDepth, + parentHindsightSessionState: options.parentHindsightSessionState, }), ), ); diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index f19a08cff..6324a9e6d 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -102,6 +102,7 @@ import { ExtensionToolWrapper } from "../extensibility/extensions/wrapper"; import type { HookCommandContext } from "../extensibility/hooks/types"; import type { Skill, SkillWarning } from "../extensibility/skills"; import { expandSlashCommand, type FileSlashCommand } from "../extensibility/slash-commands"; +import type { HindsightSessionState } from "../hindsight/state"; import { type LocalProtocolOptions, resolveLocalUrlToPath } from "../internal-urls"; import { buildDiscoverableMCPSearchIndex, @@ -563,6 +564,7 @@ export class AgentSession { #lastSuccessfulYieldToolCallId: string | undefined = undefined; #promptGeneration = 0; #providerSessionState = new Map(); + #hindsightSessionState: HindsightSessionState | undefined = undefined; #startPowerAssertion(): void { if (process.platform !== "darwin") { @@ -702,6 +704,16 @@ export class AgentSession { return this.#providerSessionState; } + getHindsightSessionState(): HindsightSessionState | undefined { + return this.#hindsightSessionState; + } + + setHindsightSessionState(state: HindsightSessionState | undefined): HindsightSessionState | undefined { + const previous = this.#hindsightSessionState; + this.#hindsightSessionState = state; + return previous; + } + /** TTSR manager for time-traveling stream rules */ get ttsrManager(): TtsrManager | undefined { return this.#ttsrManager; @@ -1964,6 +1976,22 @@ export class AgentSession { this.#unsubscribeAgent = this.agent.subscribe(this.#handleAgentEvent); } + /** Keep Hindsight metadata aligned when the underlying agent session id changes. */ + #rekeyHindsightMemoryForCurrentSessionId(): void { + if (resolveMemoryBackend(this.settings).id !== "hindsight") return; + const sid = this.agent.sessionId; + if (!sid) return; + this.getHindsightSessionState()?.setSessionId(sid); + } + + /** New session file: reset auto-recall / retain-threshold counters for the new transcript. */ + #resetHindsightConversationTrackingIfHindsight(): void { + if (resolveMemoryBackend(this.settings).id !== "hindsight") return; + const state = this.getHindsightSessionState(); + if (!state || state.aliasOf) return; + state.resetConversationTracking(); + } + /** * Remove all listeners, flush pending writes, and disconnect from agent. * Call this when completely done with the session. @@ -1994,6 +2022,9 @@ export class AgentSession { this.#stopPowerAssertion(); await this.sessionManager.close(); this.#closeAllProviderSessions("dispose"); + const hindsightState = this.setHindsightSessionState(undefined); + await hindsightState?.flushRetainQueue(); + hindsightState?.dispose(); this.#disconnectFromAgent(); this.#eventListeners = []; } @@ -3626,6 +3657,8 @@ export class AgentSession { await this.sessionManager.newSession(options); this.setTodoPhases([]); this.agent.sessionId = this.sessionManager.getSessionId(); + this.#rekeyHindsightMemoryForCurrentSessionId(); + this.#resetHindsightConversationTrackingIfHindsight(); this.#steeringMessages = []; this.#followUpMessages = []; this.#pendingNextTurnMessages = []; @@ -3719,6 +3752,7 @@ export class AgentSession { // Update agent session ID this.agent.sessionId = this.sessionManager.getSessionId(); + this.#rekeyHindsightMemoryForCurrentSessionId(); // Emit session_switch event with reason "fork" to hooks if (this.#extensionRunner) { @@ -4261,7 +4295,7 @@ export class AgentSession { if (!backend.preCompactionContext) return undefined; const messages = preparation.messagesToSummarize.concat(preparation.turnPrefixMessages); try { - return await backend.preCompactionContext(messages, this.settings); + return await backend.preCompactionContext(messages, this.settings, this); } catch (err) { logger.debug("Memory backend preCompactionContext failed", { backend: backend.id, @@ -4425,6 +4459,8 @@ export class AgentSession { await this.sessionManager.newSession(previousSessionFile ? { parentSession: previousSessionFile } : undefined); this.agent.reset(); this.agent.sessionId = this.sessionManager.getSessionId(); + this.#rekeyHindsightMemoryForCurrentSessionId(); + this.#resetHindsightConversationTrackingIfHindsight(); this.#steeringMessages = []; this.#followUpMessages = []; this.#pendingNextTurnMessages = []; @@ -6571,6 +6607,7 @@ export class AgentSession { try { await this.sessionManager.setSessionFile(sessionPath); this.agent.sessionId = this.sessionManager.getSessionId(); + this.#rekeyHindsightMemoryForCurrentSessionId(); const sessionContext = this.buildDisplaySessionContext(); const didReloadConversationChange = @@ -6640,11 +6677,15 @@ export class AgentSession { ? undefined : configuredServiceTier; + if (switchingToDifferentSession) { + this.#resetHindsightConversationTrackingIfHindsight(); + } this.#reconnectToAgent(); return true; } catch (error) { this.sessionManager.restoreState(previousSessionState); this.agent.sessionId = previousSessionState.sessionId; + this.#rekeyHindsightMemoryForCurrentSessionId(); let restoreMcpError: unknown; try { await this.#restoreMCPSelectionsForSessionContext(previousSessionContext, { @@ -6736,6 +6777,8 @@ export class AgentSession { } this.#syncTodoPhasesFromBranch(); this.agent.sessionId = this.sessionManager.getSessionId(); + this.#rekeyHindsightMemoryForCurrentSessionId(); + this.#resetHindsightConversationTrackingIfHindsight(); // Reload messages from entries (works for both file and in-memory mode) const sessionContext = this.buildDisplaySessionContext(); diff --git a/packages/coding-agent/src/task/executor.ts b/packages/coding-agent/src/task/executor.ts index 8bdccb7f9..0a0674587 100644 --- a/packages/coding-agent/src/task/executor.ts +++ b/packages/coding-agent/src/task/executor.ts @@ -17,6 +17,7 @@ import { SETTINGS_SCHEMA, type SettingPath } from "../config/settings-schema"; import type { CustomTool } from "../extensibility/custom-tools/types"; import { runExtensionCompact, runExtensionSetModel } from "../extensibility/extensions/compact-handler"; import type { Skill } from "../extensibility/skills"; +import type { HindsightSessionState } from "../hindsight/state"; import type { LocalProtocolOptions } from "../internal-urls"; import { callTool } from "../mcp/client"; import type { MCPManager } from "../mcp/manager"; @@ -163,6 +164,7 @@ export interface ExecutorOptions { settings?: Settings; /** Override local:// protocol options so subagent shares parent's local:// root */ localProtocolOptions?: LocalProtocolOptions; + parentHindsightSessionState?: HindsightSessionState; } function parseStringifiedJson(value: unknown): unknown { @@ -979,6 +981,7 @@ export async function runSubprocess(options: ExecutorOptions): Promise { skills: availableSkills, promptTemplates, localProtocolOptions, + parentHindsightSessionState: this.session.getHindsightSessionState?.(), }); } @@ -918,6 +919,7 @@ export class TaskTool implements AgentTool { skills: availableSkills, promptTemplates, localProtocolOptions, + parentHindsightSessionState: this.session.getHindsightSessionState?.(), }); if (mergeMode === "branch" && result.exitCode === 0) { try { diff --git a/packages/coding-agent/src/tools/hindsight-recall.ts b/packages/coding-agent/src/tools/hindsight-recall.ts index 5b9f4120b..5f5a60dbb 100644 --- a/packages/coding-agent/src/tools/hindsight-recall.ts +++ b/packages/coding-agent/src/tools/hindsight-recall.ts @@ -1,7 +1,6 @@ import type { AgentTool, AgentToolResult } from "@oh-my-pi/pi-agent-core"; import { logger, untilAborted } from "@oh-my-pi/pi-utils"; import { type Static, Type } from "@sinclair/typebox"; -import { getHindsightSessionState } from "../hindsight/backend"; import { formatCurrentTime, formatMemories } from "../hindsight/content"; import recallDescription from "../prompts/tools/recall.md" with { type: "text" }; import type { ToolSession } from "."; @@ -30,8 +29,7 @@ export class HindsightRecallTool implements AgentTool { return untilAborted(signal, async () => { - const sessionId = this.session.getSessionId?.(); - const state = sessionId ? getHindsightSessionState(sessionId) : undefined; + const state = this.session.getHindsightSessionState?.(); if (!state) { throw new Error("Hindsight backend is not initialised for this session."); } diff --git a/packages/coding-agent/src/tools/hindsight-reflect.ts b/packages/coding-agent/src/tools/hindsight-reflect.ts index 42581f7ba..732a0f29a 100644 --- a/packages/coding-agent/src/tools/hindsight-reflect.ts +++ b/packages/coding-agent/src/tools/hindsight-reflect.ts @@ -1,7 +1,6 @@ import type { AgentTool, AgentToolResult } from "@oh-my-pi/pi-agent-core"; import { logger, untilAborted } from "@oh-my-pi/pi-utils"; import { type Static, Type } from "@sinclair/typebox"; -import { getHindsightSessionState } from "../hindsight/backend"; import { ensureBankMission } from "../hindsight/bank"; import reflectDescription from "../prompts/tools/reflect.md" with { type: "text" }; import type { ToolSession } from "."; @@ -29,8 +28,7 @@ export class HindsightReflectTool implements AgentTool { return untilAborted(signal, async () => { - const sessionId = this.session.getSessionId?.(); - const state = sessionId ? getHindsightSessionState(sessionId) : undefined; + const state = this.session.getHindsightSessionState?.(); if (!state) { throw new Error("Hindsight backend is not initialised for this session."); } diff --git a/packages/coding-agent/src/tools/hindsight-retain.ts b/packages/coding-agent/src/tools/hindsight-retain.ts index 63b7c2de9..2ea7c55b8 100644 --- a/packages/coding-agent/src/tools/hindsight-retain.ts +++ b/packages/coding-agent/src/tools/hindsight-retain.ts @@ -1,7 +1,5 @@ import type { AgentTool, AgentToolResult } from "@oh-my-pi/pi-agent-core"; import { type Static, Type } from "@sinclair/typebox"; -import { getHindsightSessionState } from "../hindsight/backend"; -import { enqueueRetain } from "../hindsight/retain-queue"; import retainDescription from "../prompts/tools/retain.md" with { type: "text" }; import type { ToolSession } from "."; @@ -39,18 +37,17 @@ export class HindsightRetainTool implements AgentTool { - const sessionId = this.session.getSessionId?.(); - const state = sessionId ? getHindsightSessionState(sessionId) : undefined; - if (!state || !sessionId) { + const state = this.session.getHindsightSessionState?.(); + if (!state) { throw new Error("Hindsight backend is not initialised for this session."); } - // Push every item onto the global queue and return immediately. The - // queue flushes either when it reaches its batch threshold or when its - // debounce timer fires. If the eventual batch fails, the queue + // Push every item onto the session-owned queue and return immediately. + // The queue flushes either when it reaches its batch threshold or when + // its debounce timer fires. If the eventual batch fails, the queue // surfaces a UI-only warning notice — the LLM is not informed. for (const item of params.items) { - enqueueRetain(sessionId, item.content, item.context); + state.enqueueRetain(item.content, item.context); } const count = params.items.length; diff --git a/packages/coding-agent/src/tools/index.ts b/packages/coding-agent/src/tools/index.ts index 851c217e1..9b71260d5 100644 --- a/packages/coding-agent/src/tools/index.ts +++ b/packages/coding-agent/src/tools/index.ts @@ -7,6 +7,7 @@ import type { Settings } from "../config/settings"; import { EditTool } from "../edit"; import { checkPythonKernelAvailability } from "../eval/py/kernel"; import type { Skill } from "../extensibility/skills"; +import type { HindsightSessionState } from "../hindsight/state"; import type { InternalUrlRouter } from "../internal-urls"; import { LspTool } from "../lsp"; import type { DiscoverableMCPSearchIndex, DiscoverableMCPTool } from "../mcp/discoverable-tool-metadata"; @@ -141,6 +142,8 @@ export interface ToolSession { trackEvalExecution?(execution: Promise, abortController: AbortController): Promise; /** Get session ID */ getSessionId?: () => string | null; + /** Get Hindsight runtime state for this agent session. */ + getHindsightSessionState?: () => HindsightSessionState | undefined; /** Agent identity used for IRC routing. Returns the registry id (e.g. "0-Main", "0-AuthLoader"). */ getAgentId?: () => string | null; /** Look up a registered tool by name (used by the eval js backend's tool bridge). */ diff --git a/packages/coding-agent/test/hindsight-backend.test.ts b/packages/coding-agent/test/hindsight-backend.test.ts index 284ee5e3f..a9174b62b 100644 --- a/packages/coding-agent/test/hindsight-backend.test.ts +++ b/packages/coding-agent/test/hindsight-backend.test.ts @@ -10,13 +10,9 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; import { _resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; -import { - clearHindsightSessionStateForTest, - getHindsightSessionState, - hindsightBackend, - reloadMentalModelsForSession, -} from "@oh-my-pi/pi-coding-agent/hindsight/backend"; +import { hindsightBackend, reloadMentalModelsForSession } from "@oh-my-pi/pi-coding-agent/hindsight/backend"; import { HindsightApi } from "@oh-my-pi/pi-coding-agent/hindsight/client"; +import type { HindsightSessionState } from "@oh-my-pi/pi-coding-agent/hindsight/state"; import type { AgentSessionEventListener } from "@oh-my-pi/pi-coding-agent/session/agent-session"; interface FakeSessionDeps { @@ -28,6 +24,7 @@ interface FakeSessionDeps { function makeFakeSession(deps: FakeSessionDeps) { const listeners = new Set(); const entries = deps.entries ?? []; + let hindsightState: HindsightSessionState | undefined; const session = { sessionId: deps.sessionId, settings: Settings.isolated(), @@ -64,6 +61,12 @@ function makeFakeSession(deps: FakeSessionDeps) { return () => listeners.delete(listener); }, refreshBaseSystemPrompt: vi.fn().mockResolvedValue(undefined), + getHindsightSessionState: () => hindsightState, + setHindsightSessionState(state: HindsightSessionState | undefined) { + const previous = hindsightState; + hindsightState = state; + return previous; + }, emit(event: Parameters[0]) { for (const l of [...listeners]) l(event); }, @@ -74,12 +77,10 @@ function makeFakeSession(deps: FakeSessionDeps) { describe("hindsightBackend.start", () => { beforeEach(() => { _resetSettingsForTest(); - clearHindsightSessionStateForTest(); }); afterEach(() => { vi.restoreAllMocks(); - clearHindsightSessionStateForTest(); }); it("does nothing when memory.backend is hindsight but apiUrl is empty", async () => { @@ -94,7 +95,7 @@ describe("hindsightBackend.start", () => { taskDepth: 0, }); - expect(getHindsightSessionState("s1")).toBeUndefined(); + expect(session.getHindsightSessionState()).toBeUndefined(); }); it("registers per-session state and subscribes to agent events when configured", async () => { @@ -112,8 +113,30 @@ describe("hindsightBackend.start", () => { taskDepth: 0, }); - expect(getHindsightSessionState("s2")).toBeDefined(); - expect(getHindsightSessionState("s2")?.bankId).toBeTruthy(); + expect(session.getHindsightSessionState()).toBeDefined(); + expect(session.getHindsightSessionState()?.bankId).toBeTruthy(); + }); + + it("rekeys state when the same AgentSession gets a new session id (resume/switch)", async () => { + const settings = Settings.isolated({ + "memory.backend": "hindsight", + "hindsight.apiUrl": "http://localhost:8888", + }); + const session = makeFakeSession({ sessionId: "s-before" }); + + await hindsightBackend.start({ + session: session as never, + settings, + modelRegistry: {} as never, + agentDir: "/tmp", + taskDepth: 0, + }); + + expect(session.getHindsightSessionState()).toBeDefined(); + (session as { sessionId: string | null }).sessionId = "s-after"; + session.getHindsightSessionState()?.setSessionId("s-after"); + expect(session.getHindsightSessionState()?.sessionId).toBe("s-after"); + expect(session.getHindsightSessionState()?.bankId).toBeTruthy(); }); it("retains every Nth user turn on agent_end and skips intermediate turns", async () => { @@ -166,7 +189,7 @@ describe("hindsightBackend.start", () => { agentDir: "/tmp", taskDepth: 0, }); - const parentState = getHindsightSessionState("parent"); + const parentState = parentSession.getHindsightSessionState(); expect(parentState).toBeDefined(); // Subagent runs with taskDepth > 0 should alias the parent. @@ -177,8 +200,9 @@ describe("hindsightBackend.start", () => { modelRegistry: {} as never, agentDir: "/tmp", taskDepth: 1, + parentHindsightSessionState: parentState, }); - const subState = getHindsightSessionState("sub"); + const subState = subSession.getHindsightSessionState(); expect(subState).toBeDefined(); expect(subState?.aliasOf).toBe(parentState); expect(subState?.bankId).toBe(parentState?.bankId); @@ -205,19 +229,17 @@ describe("hindsightBackend.start", () => { taskDepth: 1, }); - expect(getHindsightSessionState("orphan-sub")).toBeUndefined(); + expect(session.getHindsightSessionState()).toBeUndefined(); }); }); describe("hindsightBackend.preCompactionContext", () => { beforeEach(() => { _resetSettingsForTest(); - clearHindsightSessionStateForTest(); }); afterEach(() => { vi.restoreAllMocks(); - clearHindsightSessionStateForTest(); }); it("returns undefined when no apiUrl is configured", async () => { @@ -246,7 +268,7 @@ describe("hindsightBackend.preCompactionContext", () => { } as never); const messages: AgentMessage[] = [{ role: "user", content: "What did we decide?", timestamp: 0 } as never]; - const ctx = await hindsightBackend.preCompactionContext?.(messages, settings); + const ctx = await hindsightBackend.preCompactionContext?.(messages, settings, session as never); expect(ctx).toBeDefined(); expect(ctx).toContain(""); expect(ctx).toContain("remembered fact"); @@ -268,7 +290,7 @@ describe("hindsightBackend.preCompactionContext", () => { vi.spyOn(HindsightApi.prototype, "recall").mockResolvedValue({ results: [] } as never); const messages: AgentMessage[] = [{ role: "user", content: "anything", timestamp: 0 } as never]; - const ctx = await hindsightBackend.preCompactionContext?.(messages, settings); + const ctx = await hindsightBackend.preCompactionContext?.(messages, settings, session as never); expect(ctx).toBeUndefined(); }); }); @@ -276,12 +298,10 @@ describe("hindsightBackend.preCompactionContext", () => { describe("hindsightBackend first-turn injection", () => { beforeEach(() => { _resetSettingsForTest(); - clearHindsightSessionStateForTest(); }); afterEach(() => { vi.restoreAllMocks(); - clearHindsightSessionStateForTest(); }); it("returns a tagged block for the current first turn before agent_start", async () => { @@ -311,8 +331,8 @@ describe("hindsightBackend first-turn injection", () => { ); expect(block).toContain(""); expect(block).toContain("Can prefers concise communication"); - expect(getHindsightSessionState("s8")?.hasRecalledForFirstTurn).toBe(true); - expect(getHindsightSessionState("s8")?.lastRecallSnippet).toBe(block); + expect(session.getHindsightSessionState()?.hasRecalledForFirstTurn).toBe(true); + expect(session.getHindsightSessionState()?.lastRecallSnippet).toBe(block); }); it("keeps the wrapper in buildDeveloperInstructions", async () => { @@ -329,11 +349,11 @@ describe("hindsightBackend first-turn injection", () => { taskDepth: 0, }); - const state = getHindsightSessionState("s9"); + const state = session.getHindsightSessionState(); expect(state).toBeDefined(); state!.lastRecallSnippet = "\nremembered fact\n"; - const prompt = await hindsightBackend.buildDeveloperInstructions("/tmp", settings); + const prompt = await hindsightBackend.buildDeveloperInstructions("/tmp", settings, session as never); expect(prompt).toContain(""); expect(prompt).toContain(""); expect(prompt).toContain("remembered fact"); @@ -356,12 +376,12 @@ describe("hindsightBackend first-turn injection", () => { agentDir: "/tmp", taskDepth: 0, }); - const state = getHindsightSessionState("s-order"); + const state = session.getHindsightSessionState(); expect(state).toBeDefined(); state!.mentalModelsSnippet = "\n# User Preferences\nprefers tabs\n"; state!.lastRecallSnippet = "\nrecalled fact\n"; - const prompt = await hindsightBackend.buildDeveloperInstructions("/tmp", settings); + const prompt = await hindsightBackend.buildDeveloperInstructions("/tmp", settings, session as never); expect(prompt).toBeDefined(); // `` and `` are mentioned in STATIC_INSTRUCTIONS // bullets too. Match the actual injected block opener (tag + newline) @@ -395,13 +415,13 @@ describe("hindsightBackend first-turn injection", () => { taskDepth: 0, }); // Wait for the kicked-off load to settle. - await getHindsightSessionState("s-ttl")?.mentalModelsLoadPromise; - const state = getHindsightSessionState("s-ttl"); + await session.getHindsightSessionState()?.mentalModelsLoadPromise; + const state = session.getHindsightSessionState(); expect(state).toBeDefined(); expect(state!.mentalModelsSnippet).toBeUndefined(); expect(state!.mentalModelsLoadedAt).toBeDefined(); const initialLoadedAt = state!.mentalModelsLoadedAt!; - const refreshSpy = session.refreshBaseSystemPrompt as ReturnType; + const refreshSpy = session.refreshBaseSystemPrompt; const callsBefore = refreshSpy.mock.calls.length; // Now publish content and trigger a reload. @@ -418,7 +438,7 @@ describe("hindsightBackend first-turn injection", () => { // Force the loadedAt timestamp to differ so the next assertion is meaningful. state!.mentalModelsLoadedAt = initialLoadedAt - 1000; - const ok = await reloadMentalModelsForSession("s-ttl"); + const ok = await reloadMentalModelsForSession(session as never); expect(ok).toBe(true); expect(state!.mentalModelsSnippet).toBeDefined(); expect(state!.mentalModelsSnippet).toContain("# User Preferences"); @@ -451,8 +471,9 @@ describe("hindsightBackend first-turn injection", () => { modelRegistry: {} as never, agentDir: "/tmp", taskDepth: 1, + parentHindsightSessionState: parent.getHindsightSessionState(), }); - const ok = await reloadMentalModelsForSession("alias-child"); + const ok = await reloadMentalModelsForSession(child as never); expect(ok).toBe(false); }); }); @@ -460,12 +481,10 @@ describe("hindsightBackend first-turn injection", () => { describe("hindsightBackend.clear", () => { beforeEach(() => { _resetSettingsForTest(); - clearHindsightSessionStateForTest(); }); afterEach(() => { vi.restoreAllMocks(); - clearHindsightSessionStateForTest(); }); it("drops every registered session state", async () => { @@ -481,10 +500,10 @@ describe("hindsightBackend.clear", () => { agentDir: "/tmp", taskDepth: 0, }); - expect(getHindsightSessionState("s7")).toBeDefined(); + expect(session.getHindsightSessionState()).toBeDefined(); - await hindsightBackend.clear("/tmp", "/tmp"); - expect(getHindsightSessionState("s7")).toBeUndefined(); + await hindsightBackend.clear("/tmp", "/tmp", session as never); + expect(session.getHindsightSessionState()).toBeUndefined(); }); it("does not delete server-side mental models on /memory clear (server-side state is sacred)", async () => { @@ -509,7 +528,7 @@ describe("hindsightBackend.clear", () => { taskDepth: 0, }); - await hindsightBackend.clear("/tmp", "/tmp"); + await hindsightBackend.clear("/tmp", "/tmp", session as never); expect(deleteSpy).not.toHaveBeenCalled(); }); }); diff --git a/packages/coding-agent/test/hindsight-tools.test.ts b/packages/coding-agent/test/hindsight-tools.test.ts index fb1f9e18c..2f90f3092 100644 --- a/packages/coding-agent/test/hindsight-tools.test.ts +++ b/packages/coding-agent/test/hindsight-tools.test.ts @@ -3,30 +3,22 @@ * * These exercise the public tool surface (factory gating + execute path) by * spying on `HindsightApi.prototype.{retain, recall, reflect}` and stubbing - * a per-session state via `setHindsightSessionStateForTest`. We deliberately - * do not boot a real session — these tools only need a populated state - * accessor and a Settings instance. + * Hindsight state on the fake ToolSession. We deliberately do not boot a real + * session — these tools only need a populated state accessor and Settings. */ import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import { _resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; -import { - clearHindsightSessionStateForTest, - setHindsightSessionStateForTest, -} from "@oh-my-pi/pi-coding-agent/hindsight/backend"; import { HindsightApi } from "@oh-my-pi/pi-coding-agent/hindsight/client"; import type { HindsightConfig } from "@oh-my-pi/pi-coding-agent/hindsight/config"; -import { - clearRetainQueueForTest, - flushSessionQueue, - getRetainQueueDepthForTest, -} from "@oh-my-pi/pi-coding-agent/hindsight/retain-queue"; +import { HindsightSessionState } from "@oh-my-pi/pi-coding-agent/hindsight/state"; import { HindsightRecallTool } from "@oh-my-pi/pi-coding-agent/tools/hindsight-recall"; import { HindsightReflectTool } from "@oh-my-pi/pi-coding-agent/tools/hindsight-reflect"; import { HindsightRetainTool } from "@oh-my-pi/pi-coding-agent/tools/hindsight-retain"; import type { ToolSession } from "@oh-my-pi/pi-coding-agent/tools/index"; const TEST_SESSION_ID = "test-session-id"; +let registeredState: HindsightSessionState | undefined; function makeConfig(overrides: Partial = {}): HindsightConfig { return { @@ -66,6 +58,7 @@ function makeSession(settings: Settings, sessionId: string | null = TEST_SESSION getSessionFile: () => null, getSessionId: () => sessionId, getSessionSpawns: () => null, + getHindsightSessionState: () => (sessionId === TEST_SESSION_ID ? registeredState : undefined), } as unknown as ToolSession; } @@ -77,7 +70,8 @@ interface RegisterStateOptions { } function registerState(client: HindsightApi, settings?: Settings, opts: RegisterStateOptions = {}) { - setHindsightSessionStateForTest(TEST_SESSION_ID, { + registeredState = new HindsightSessionState({ + sessionId: TEST_SESSION_ID, client, bankId: "test-bank", retainTags: opts.retainTags, @@ -88,6 +82,7 @@ function registerState(client: HindsightApi, settings?: Settings, opts: Register sessionId: TEST_SESSION_ID, sessionManager: { getEntries: () => [] } as never, emitNotice: () => {}, + getHindsightSessionState: () => registeredState, ...opts.sessionOverrides, } as never, missionsSet: new Set(), @@ -100,12 +95,12 @@ function registerState(client: HindsightApi, settings?: Settings, opts: Register describe("Hindsight tool factories", () => { beforeEach(() => { _resetSettingsForTest(); - clearHindsightSessionStateForTest(); + registeredState = undefined; }); afterEach(() => { vi.restoreAllMocks(); - clearHindsightSessionStateForTest(); + registeredState = undefined; }); it("retain/recall/reflect factories return null when memory.backend !== hindsight", () => { @@ -128,14 +123,12 @@ describe("Hindsight tool factories", () => { describe("retain.execute", () => { beforeEach(() => { _resetSettingsForTest(); - clearHindsightSessionStateForTest(); - clearRetainQueueForTest(); + registeredState = undefined; }); afterEach(() => { vi.restoreAllMocks(); - clearHindsightSessionStateForTest(); - clearRetainQueueForTest(); + registeredState = undefined; }); it("queues the memory and reports success without calling the API", async () => { @@ -152,7 +145,7 @@ describe("retain.execute", () => { // Tool returns before any HTTP work happens. expect(retainBatchSpy).not.toHaveBeenCalled(); expect(retainSpy).not.toHaveBeenCalled(); - expect(getRetainQueueDepthForTest(TEST_SESSION_ID)).toBe(1); + expect(registeredState?.retainQueue.depth).toBe(1); }); it("flushes a multi-item tool call as a single retainBatch call with per-item context", async () => { @@ -167,7 +160,7 @@ describe("retain.execute", () => { }); expect(result.content[0]).toEqual({ type: "text", text: "2 memories queued." }); - await flushSessionQueue(TEST_SESSION_ID); + await registeredState?.flushRetainQueue(); expect(retainBatchSpy).toHaveBeenCalledTimes(1); const [bankId, items, options] = retainBatchSpy.mock.calls[0]; @@ -186,7 +179,7 @@ describe("retain.execute", () => { tags: ["project:pi"], }), ]); - expect(getRetainQueueDepthForTest(TEST_SESSION_ID)).toBe(0); + expect(registeredState?.retainQueue.depth).toBe(0); }); it("emits a UI-only warning notice when the batch flush fails", async () => { @@ -198,7 +191,7 @@ describe("retain.execute", () => { const tool = HindsightRetainTool.createIf(makeSession(settings))!; await tool.execute("call-x", { items: [{ content: "doomed fact" }] }); - await flushSessionQueue(TEST_SESSION_ID); + await registeredState?.flushRetainQueue(); expect(noticeSpy).toHaveBeenCalledTimes(1); const [level, message, source] = noticeSpy.mock.calls[0]; @@ -218,12 +211,12 @@ describe("retain.execute", () => { describe("recall.execute", () => { beforeEach(() => { _resetSettingsForTest(); - clearHindsightSessionStateForTest(); + registeredState = undefined; }); afterEach(() => { vi.restoreAllMocks(); - clearHindsightSessionStateForTest(); + registeredState = undefined; }); it("returns the no-results sentinel when recall yields empty", async () => { @@ -286,12 +279,12 @@ describe("recall.execute", () => { describe("reflect.execute", () => { beforeEach(() => { _resetSettingsForTest(); - clearHindsightSessionStateForTest(); + registeredState = undefined; }); afterEach(() => { vi.restoreAllMocks(); - clearHindsightSessionStateForTest(); + registeredState = undefined; }); it("returns the reflect text and forwards context", async () => {