diff --git a/src/features/background-agent/compaction-aware-message-resolver.test.ts b/src/features/background-agent/compaction-aware-message-resolver.test.ts index 17cc07024e..5b9bed5af1 100644 --- a/src/features/background-agent/compaction-aware-message-resolver.test.ts +++ b/src/features/background-agent/compaction-aware-message-resolver.test.ts @@ -2,7 +2,15 @@ import { describe, test, expect, beforeEach, afterEach } from "bun:test" import { mkdtempSync, writeFileSync, rmSync } from "node:fs" import { join } from "node:path" import { tmpdir } from "node:os" -import { isCompactionAgent, findNearestMessageExcludingCompaction } from "./compaction-aware-message-resolver" +import { + isCompactionAgent, + findNearestMessageExcludingCompaction, + resolvePromptContextFromSessionMessages, +} from "./compaction-aware-message-resolver" +import { + clearCompactionAgentConfigCheckpoint, + setCompactionAgentConfigCheckpoint, +} from "../../shared/compaction-agent-config-checkpoint" describe("isCompactionAgent", () => { describe("#given agent name variations", () => { @@ -65,6 +73,7 @@ describe("findNearestMessageExcludingCompaction", () => { afterEach(() => { rmSync(tempDir, { force: true, recursive: true }) + clearCompactionAgentConfigCheckpoint("ses_checkpoint") }) describe("#given directory with messages", () => { @@ -186,5 +195,65 @@ describe("findNearestMessageExcludingCompaction", () => { expect(result).not.toBeNull() expect(result?.agent).toBe("newer") }) + + test("merges partial metadata from multiple recent messages", () => { + // given + writeFileSync( + join(tempDir, "003.json"), + JSON.stringify({ model: { providerID: "anthropic", modelID: "claude-opus-4-1" } }), + ) + writeFileSync(join(tempDir, "002.json"), JSON.stringify({ agent: "atlas" })) + writeFileSync(join(tempDir, "001.json"), JSON.stringify({ tools: { bash: true } })) + + // when + const result = findNearestMessageExcludingCompaction(tempDir) + + // then + expect(result).toEqual({ + agent: "atlas", + model: { providerID: "anthropic", modelID: "claude-opus-4-1" }, + tools: { bash: true }, + }) + }) + + test("fills missing metadata from compaction checkpoint", () => { + // given + setCompactionAgentConfigCheckpoint("ses_checkpoint", { + agent: "sisyphus", + model: { providerID: "openai", modelID: "gpt-5" }, + }) + writeFileSync(join(tempDir, "001.json"), JSON.stringify({ tools: { bash: true } })) + + // when + const result = findNearestMessageExcludingCompaction(tempDir, "ses_checkpoint") + + // then + expect(result).toEqual({ + agent: "sisyphus", + model: { providerID: "openai", modelID: "gpt-5" }, + tools: { bash: true }, + }) + }) + }) +}) + +describe("resolvePromptContextFromSessionMessages", () => { + test("merges partial prompt context from recent SDK messages", () => { + // given + const messages = [ + { info: { agent: "atlas" } }, + { info: { model: { providerID: "anthropic", modelID: "claude-opus-4-1" } } }, + { info: { tools: { bash: true } } }, + ] + + // when + const result = resolvePromptContextFromSessionMessages(messages) + + // then + expect(result).toEqual({ + agent: "atlas", + model: { providerID: "anthropic", modelID: "claude-opus-4-1" }, + tools: { bash: true }, + }) }) }) diff --git a/src/features/background-agent/compaction-aware-message-resolver.ts b/src/features/background-agent/compaction-aware-message-resolver.ts index 1bf94bfdbe..60b3949b3e 100644 --- a/src/features/background-agent/compaction-aware-message-resolver.ts +++ b/src/features/background-agent/compaction-aware-message-resolver.ts @@ -1,6 +1,21 @@ import { readdirSync, readFileSync } from "node:fs" import { join } from "node:path" import type { StoredMessage } from "../hook-message-injector" +import { getCompactionAgentConfigCheckpoint } from "../../shared/compaction-agent-config-checkpoint" + +type SessionMessage = { + info?: { + agent?: string + model?: { + providerID?: string + modelID?: string + variant?: string + } + providerID?: string + modelID?: string + tools?: StoredMessage["tools"] + } +} export function isCompactionAgent(agent: string | undefined): boolean { return agent?.trim().toLowerCase() === "compaction" @@ -16,42 +31,121 @@ function hasFullAgentAndModel(message: StoredMessage): boolean { function hasPartialAgentOrModel(message: StoredMessage): boolean { const hasAgent = !!message.agent && !isCompactionAgent(message.agent) const hasModel = !!message.model?.providerID && !!message.model?.modelID - return hasAgent || hasModel + return hasAgent || hasModel || !!message.tools } -export function findNearestMessageExcludingCompaction(messageDir: string): StoredMessage | null { +function convertSessionMessageToStoredMessage(message: SessionMessage): StoredMessage | null { + const info = message.info + if (!info) { + return null + } + + const providerID = info.model?.providerID ?? info.providerID + const modelID = info.model?.modelID ?? info.modelID + + return { + ...(info.agent ? { agent: info.agent } : {}), + ...(providerID && modelID + ? { + model: { + providerID, + modelID, + ...(info.model?.variant ? { variant: info.model.variant } : {}), + }, + } + : {}), + ...(info.tools ? { tools: info.tools } : {}), + } +} + +function mergeStoredMessages( + messages: Array, + sessionID?: string, +): StoredMessage | null { + const merged: StoredMessage = {} + + for (const message of messages) { + if (!message || isCompactionAgent(message.agent)) { + continue + } + + if (!merged.agent && message.agent) { + merged.agent = message.agent + } + + if (!merged.model?.providerID && message.model?.providerID && message.model.modelID) { + merged.model = { + providerID: message.model.providerID, + modelID: message.model.modelID, + ...(message.model.variant ? { variant: message.model.variant } : {}), + } + } + + if (!merged.tools && message.tools) { + merged.tools = message.tools + } + + if (hasFullAgentAndModel(merged) && merged.tools) { + break + } + } + + const checkpoint = sessionID + ? getCompactionAgentConfigCheckpoint(sessionID) + : undefined + + if (!merged.agent && checkpoint?.agent) { + merged.agent = checkpoint.agent + } + + if (!merged.model && checkpoint?.model) { + merged.model = { + providerID: checkpoint.model.providerID, + modelID: checkpoint.model.modelID, + } + } + + if (!merged.tools && checkpoint?.tools) { + merged.tools = checkpoint.tools + } + + return hasPartialAgentOrModel(merged) ? merged : null +} + +export function resolvePromptContextFromSessionMessages( + messages: SessionMessage[], + sessionID?: string, +): StoredMessage | null { + const convertedMessages = messages + .map(convertSessionMessageToStoredMessage) + .reverse() + + return mergeStoredMessages(convertedMessages, sessionID) +} + +export function findNearestMessageExcludingCompaction( + messageDir: string, + sessionID?: string, +): StoredMessage | null { try { const files = readdirSync(messageDir) - .filter((name) => name.endsWith(".json")) + .filter((name: string) => name.endsWith(".json")) .sort() .reverse() - for (const file of files) { - try { - const content = readFileSync(join(messageDir, file), "utf-8") - const parsed = JSON.parse(content) as StoredMessage - if (hasFullAgentAndModel(parsed)) { - return parsed - } - } catch { - continue - } - } + const messages: Array = [] for (const file of files) { try { const content = readFileSync(join(messageDir, file), "utf-8") - const parsed = JSON.parse(content) as StoredMessage - if (hasPartialAgentOrModel(parsed)) { - return parsed - } + messages.push(JSON.parse(content) as StoredMessage) } catch { continue } } + + return mergeStoredMessages(messages, sessionID) } catch { return null } - - return null } diff --git a/src/features/background-agent/manager.ts b/src/features/background-agent/manager.ts index 26e9525188..6cd60273ea 100644 --- a/src/features/background-agent/manager.ts +++ b/src/features/background-agent/manager.ts @@ -41,7 +41,10 @@ import { } from "./error-classifier" import { tryFallbackRetry } from "./fallback-retry-handler" import { registerManagerForCleanup, unregisterManagerForCleanup } from "./process-cleanup" -import { isCompactionAgent, findNearestMessageExcludingCompaction } from "./compaction-aware-message-resolver" +import { + findNearestMessageExcludingCompaction, + resolvePromptContextFromSessionMessages, +} from "./compaction-aware-message-resolver" import { handleSessionIdleBackgroundEvent } from "./session-idle-event-handler" import { MESSAGE_STORAGE } from "../hook-message-injector" import { join } from "node:path" @@ -1323,20 +1326,20 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea tools?: Record } }>) - for (let i = messages.length - 1; i >= 0; i--) { - const info = messages[i].info - if (isCompactionAgent(info?.agent)) { - continue - } - const normalizedTools = isRecord(info?.tools) - ? normalizePromptTools(info.tools as Record) + const promptContext = resolvePromptContextFromSessionMessages( + messages, + task.parentSessionID, + ) + const normalizedTools = isRecord(promptContext?.tools) + ? normalizePromptTools(promptContext.tools) + : undefined + + if (promptContext?.agent || promptContext?.model || normalizedTools) { + agent = promptContext?.agent ?? task.parentAgent + model = promptContext?.model?.providerID && promptContext.model.modelID + ? { providerID: promptContext.model.providerID, modelID: promptContext.model.modelID } : undefined - if (info?.agent || info?.model || (info?.modelID && info?.providerID) || normalizedTools) { - agent = info?.agent ?? task.parentAgent - model = info?.model ?? (info?.providerID && info?.modelID ? { providerID: info.providerID, modelID: info.modelID } : undefined) - tools = normalizedTools ?? tools - break - } + tools = normalizedTools ?? tools } } catch (error) { if (isAbortedSessionError(error)) { @@ -1346,7 +1349,9 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea }) } const messageDir = join(MESSAGE_STORAGE, task.parentSessionID) - const currentMessage = messageDir ? findNearestMessageExcludingCompaction(messageDir) : null + const currentMessage = messageDir + ? findNearestMessageExcludingCompaction(messageDir, task.parentSessionID) + : null agent = currentMessage?.agent ?? task.parentAgent model = currentMessage?.model?.providerID && currentMessage?.model?.modelID ? { providerID: currentMessage.model.providerID, modelID: currentMessage.model.modelID } diff --git a/src/hooks/compaction-context-injector/compaction-context-prompt.ts b/src/hooks/compaction-context-injector/compaction-context-prompt.ts new file mode 100644 index 0000000000..11bc25747f --- /dev/null +++ b/src/hooks/compaction-context-injector/compaction-context-prompt.ts @@ -0,0 +1,56 @@ +import { + createSystemDirective, + SystemDirectiveTypes, +} from "../../shared/system-directive" + +export const COMPACTION_CONTEXT_PROMPT = `${createSystemDirective(SystemDirectiveTypes.COMPACTION_CONTEXT)} + +When summarizing this session, you MUST include the following sections in your summary: + +## 1. User Requests (As-Is) +- List all original user requests exactly as they were stated +- Preserve the user's exact wording and intent + +## 2. Final Goal +- What the user ultimately wanted to achieve +- The end result or deliverable expected + +## 3. Work Completed +- What has been done so far +- Files created/modified +- Features implemented +- Problems solved + +## 4. Remaining Tasks +- What still needs to be done +- Pending items from the original request +- Follow-up tasks identified during the work + +## 5. Active Working Context (For Seamless Continuation) +- **Files**: Paths of files currently being edited or frequently referenced +- **Code in Progress**: Key code snippets, function signatures, or data structures under active development +- **External References**: Documentation URLs, library APIs, or external resources being consulted +- **State & Variables**: Important variable names, configuration values, or runtime state relevant to ongoing work + +## 6. Explicit Constraints (Verbatim Only) +- Include ONLY constraints explicitly stated by the user or in existing AGENTS.md context +- Quote constraints verbatim (do not paraphrase) +- Do NOT invent, add, or modify constraints +- If no explicit constraints exist, write "None" + +## 7. Agent Verification State (Critical for Reviewers) +- **Current Agent**: What agent is running (momus, oracle, etc.) +- **Verification Progress**: Files already verified/validated +- **Pending Verifications**: Files still needing verification +- **Previous Rejections**: If reviewer agent, what was rejected and why +- **Acceptance Status**: Current state of review process + +This section is CRITICAL for reviewer agents (momus, oracle) to maintain continuity. + +## 8. Delegated Agent Sessions +- List ALL background agent tasks spawned during this session +- For each: agent name, category, status, description, and **session_id** +- **RESUME, DON'T RESTART.** Each listed session retains full context. After compaction, use \`session_id\` to continue existing agent sessions instead of spawning new ones. This saves tokens, preserves learned context, and prevents duplicate work. + +This context is critical for maintaining continuity after compaction. +` diff --git a/src/hooks/compaction-context-injector/constants.ts b/src/hooks/compaction-context-injector/constants.ts new file mode 100644 index 0000000000..b57c245196 --- /dev/null +++ b/src/hooks/compaction-context-injector/constants.ts @@ -0,0 +1,5 @@ +export const HOOK_NAME = "compaction-context-injector" +export const AGENT_RECOVERY_PROMPT = "[restore checkpointed session agent configuration after compaction]" +export const NO_TEXT_TAIL_THRESHOLD = 5 +export const RECOVERY_COOLDOWN_MS = 60_000 +export const RECENT_COMPACTION_WINDOW_MS = 10 * 60 * 1000 diff --git a/src/hooks/compaction-context-injector/hook.ts b/src/hooks/compaction-context-injector/hook.ts index 0f3a912fe5..462dc18a60 100644 --- a/src/hooks/compaction-context-injector/hook.ts +++ b/src/hooks/compaction-context-injector/hook.ts @@ -1,63 +1,60 @@ import type { BackgroundManager } from "../../features/background-agent" import { - createSystemDirective, - SystemDirectiveTypes, -} from "../../shared/system-directive" - -const COMPACTION_CONTEXT_PROMPT = `${createSystemDirective(SystemDirectiveTypes.COMPACTION_CONTEXT)} - -When summarizing this session, you MUST include the following sections in your summary: - -## 1. User Requests (As-Is) -- List all original user requests exactly as they were stated -- Preserve the user's exact wording and intent - -## 2. Final Goal -- What the user ultimately wanted to achieve -- The end result or deliverable expected - -## 3. Work Completed -- What has been done so far -- Files created/modified -- Features implemented -- Problems solved - -## 4. Remaining Tasks -- What still needs to be done -- Pending items from the original request -- Follow-up tasks identified during the work - -## 5. Active Working Context (For Seamless Continuation) -- **Files**: Paths of files currently being edited or frequently referenced -- **Code in Progress**: Key code snippets, function signatures, or data structures under active development -- **External References**: Documentation URLs, library APIs, or external resources being consulted -- **State & Variables**: Important variable names, configuration values, or runtime state relevant to ongoing work - -## 6. Explicit Constraints (Verbatim Only) -- Include ONLY constraints explicitly stated by the user or in existing AGENTS.md context -- Quote constraints verbatim (do not paraphrase) -- Do NOT invent, add, or modify constraints -- If no explicit constraints exist, write "None" - -## 7. Agent Verification State (Critical for Reviewers) -- **Current Agent**: What agent is running (momus, oracle, etc.) -- **Verification Progress**: Files already verified/validated -- **Pending Verifications**: Files still needing verification -- **Previous Rejections**: If reviewer agent, what was rejected and why -- **Acceptance Status**: Current state of review process - -This section is CRITICAL for reviewer agents (momus, oracle) to maintain continuity. - -## 8. Delegated Agent Sessions -- List ALL background agent tasks spawned during this session -- For each: agent name, category, status, description, and **session_id** -- **RESUME, DON'T RESTART.** Each listed session retains full context. After compaction, use \`session_id\` to continue existing agent sessions instead of spawning new ones. This saves tokens, preserves learned context, and prevents duplicate work. - -This context is critical for maintaining continuity after compaction. -` - -export function createCompactionContextInjector(backgroundManager?: BackgroundManager) { - return (sessionID?: string): string => { + clearCompactionAgentConfigCheckpoint, + setCompactionAgentConfigCheckpoint, +} from "../../shared/compaction-agent-config-checkpoint" +import { log } from "../../shared/logger" +import { COMPACTION_CONTEXT_PROMPT } from "./compaction-context-prompt" +import { resolveSessionPromptConfig } from "./session-prompt-config-resolver" +import { finalizeTrackedAssistantMessage, shouldTreatAssistantPartAsOutput, trackAssistantOutput, type TailMonitorState } from "./tail-monitor" +import { resolveSessionID } from "./session-id" +import type { CompactionContextClient, CompactionContextInjector } from "./types" +import { createRecoveryLogic } from "./recovery" + +export function createCompactionContextInjector(options?: { + ctx?: CompactionContextClient + backgroundManager?: BackgroundManager +}): CompactionContextInjector { + const ctx = options?.ctx + const backgroundManager = options?.backgroundManager + const tailStates = new Map() + + const getTailState = (sessionID: string): TailMonitorState => { + const existing = tailStates.get(sessionID) + if (existing) { + return existing + } + + const created: TailMonitorState = { + currentHasOutput: false, + consecutiveNoTextMessages: 0, + } + tailStates.set(sessionID, created) + return created + } + + const { recoverCheckpointedAgentConfig, maybeWarnAboutNoTextTail } = createRecoveryLogic(ctx, getTailState) + + const capture = async (sessionID: string): Promise => { + if (!ctx || !sessionID) { + return + } + + const promptConfig = await resolveSessionPromptConfig(ctx, sessionID) + if (!promptConfig.agent && !promptConfig.model && !promptConfig.tools) { + return + } + + setCompactionAgentConfigCheckpoint(sessionID, promptConfig) + log(`[compaction-context-injector] Captured agent checkpoint before compaction`, { + sessionID, + agent: promptConfig.agent, + model: promptConfig.model, + hasTools: !!promptConfig.tools, + }) + } + + const inject = (sessionID?: string): string => { let prompt = COMPACTION_CONTEXT_PROMPT if (backgroundManager && sessionID) { @@ -69,4 +66,99 @@ export function createCompactionContextInjector(backgroundManager?: BackgroundMa return prompt } + + const event = async ({ event }: { event: { type: string; properties?: unknown } }): Promise => { + const props = event.properties as Record | undefined + + if (event.type === "session.deleted") { + const sessionID = resolveSessionID(props) + if (sessionID) { + clearCompactionAgentConfigCheckpoint(sessionID) + tailStates.delete(sessionID) + } + return + } + + if (event.type === "session.idle") { + const sessionID = resolveSessionID(props) + if (!sessionID) { + return + } + + const noTextCount = finalizeTrackedAssistantMessage(getTailState(sessionID)) + if (noTextCount > 0) { + await maybeWarnAboutNoTextTail(sessionID) + } + return + } + + if (event.type === "session.compacted") { + const sessionID = resolveSessionID(props) + if (!sessionID) { + return + } + + const tailState = getTailState(sessionID) + finalizeTrackedAssistantMessage(tailState) + tailState.lastCompactedAt = Date.now() + await maybeWarnAboutNoTextTail(sessionID) + await recoverCheckpointedAgentConfig(sessionID, "session.compacted") + return + } + + if (event.type === "message.updated") { + const info = props?.info as { + id?: string + role?: string + sessionID?: string + } | undefined + + if (!info?.sessionID || info.role !== "assistant" || !info.id) { + return + } + + const tailState = getTailState(info.sessionID) + if (tailState.currentMessageID && tailState.currentMessageID !== info.id) { + finalizeTrackedAssistantMessage(tailState) + await maybeWarnAboutNoTextTail(info.sessionID) + } + + if (tailState.currentMessageID !== info.id) { + tailState.currentMessageID = info.id + tailState.currentHasOutput = false + } + return + } + + if (event.type === "message.part.delta") { + const sessionID = props?.sessionID as string | undefined + const messageID = props?.messageID as string | undefined + const field = props?.field as string | undefined + const delta = props?.delta as string | undefined + + if (!sessionID || field !== "text" || !delta?.trim()) { + return + } + + trackAssistantOutput(getTailState(sessionID), messageID) + return + } + + if (event.type === "message.part.updated") { + const part = props?.part as { + messageID?: string + sessionID?: string + type?: string + text?: string + } | undefined + + if (!part?.sessionID || !shouldTreatAssistantPartAsOutput(part)) { + return + } + + trackAssistantOutput(getTailState(part.sessionID), part.messageID) + } + } + + return { capture, inject, event } } diff --git a/src/hooks/compaction-context-injector/index.test.ts b/src/hooks/compaction-context-injector/index.test.ts index a2813916fc..9eacd0cddd 100644 --- a/src/hooks/compaction-context-injector/index.test.ts +++ b/src/hooks/compaction-context-injector/index.test.ts @@ -17,6 +17,27 @@ mock.module("../../shared/system-directive", () => ({ import { createCompactionContextInjector } from "./index" import { TaskHistory } from "../../features/background-agent/task-history" +function createMockContext( + messageResponses: Array }>>, + promptAsyncMock = mock(async () => ({})), +) { + let callIndex = 0 + + return { + client: { + session: { + messages: mock(async () => { + const response = messageResponses[Math.min(callIndex, messageResponses.length - 1)] ?? [] + callIndex += 1 + return { data: response } + }), + promptAsync: promptAsyncMock, + }, + }, + directory: "/tmp/test", + } +} + describe("createCompactionContextInjector", () => { describe("Agent Verification State preservation", () => { it("includes Agent Verification State section in compaction prompt", async () => { @@ -24,7 +45,7 @@ describe("createCompactionContextInjector", () => { const injector = createCompactionContextInjector() //#when - const prompt = injector() + const prompt = injector.inject() //#then expect(prompt).toContain("Agent Verification State") @@ -37,7 +58,7 @@ describe("createCompactionContextInjector", () => { const injector = createCompactionContextInjector() //#when - const prompt = injector() + const prompt = injector.inject() //#then expect(prompt).toContain("Previous Rejections") @@ -50,7 +71,7 @@ describe("createCompactionContextInjector", () => { const injector = createCompactionContextInjector() //#when - const prompt = injector() + const prompt = injector.inject() //#then expect(prompt).toContain("Pending Verifications") @@ -63,7 +84,7 @@ describe("createCompactionContextInjector", () => { const injector = createCompactionContextInjector() //#when - const prompt = injector() + const prompt = injector.inject() //#then expect(prompt).toContain("Explicit Constraints (Verbatim Only)") @@ -77,7 +98,7 @@ describe("createCompactionContextInjector", () => { const injector = createCompactionContextInjector() //#when - const prompt = injector() + const prompt = injector.inject() //#then expect(prompt).toContain("Delegated Agent Sessions") @@ -89,10 +110,10 @@ describe("createCompactionContextInjector", () => { //#given const mockManager = { taskHistory: new TaskHistory() } as any mockManager.taskHistory.record("ses_parent", { id: "t1", sessionID: "ses_child", agent: "explore", description: "Find patterns", status: "completed", category: "quick" }) - const injector = createCompactionContextInjector(mockManager) + const injector = createCompactionContextInjector({ backgroundManager: mockManager }) //#when - const prompt = injector("ses_parent") + const prompt = injector.inject("ses_parent") //#then expect(prompt).toContain("Active/Recent Delegated Sessions") @@ -104,13 +125,152 @@ describe("createCompactionContextInjector", () => { it("does not inject task history section when no entries exist", async () => { //#given const mockManager = { taskHistory: new TaskHistory() } as any - const injector = createCompactionContextInjector(mockManager) + const injector = createCompactionContextInjector({ backgroundManager: mockManager }) //#when - const prompt = injector("ses_empty") + const prompt = injector.inject("ses_empty") //#then expect(prompt).not.toContain("Active/Recent Delegated Sessions") }) }) + + describe("agent checkpoint recovery", () => { + it("re-injects checkpointed agent config after compaction when latest agent is lost", async () => { + //#given + const promptAsyncMock = mock(async () => ({})) + const ctx = createMockContext( + [ + [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + tools: { bash: "allow" }, + }, + }, + ], + [ + { + info: { + role: "user", + agent: "compaction", + model: { providerID: "anthropic", modelID: "claude-opus-4-1" }, + }, + }, + ], + [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + }, + }, + ], + ], + promptAsyncMock, + ) + const injector = createCompactionContextInjector({ ctx }) + + //#when + await injector.capture("ses_checkpoint") + await injector.event({ + event: { type: "session.compacted", properties: { sessionID: "ses_checkpoint" } }, + }) + + //#then + expect(promptAsyncMock).toHaveBeenCalledWith({ + path: { id: "ses_checkpoint" }, + body: { + noReply: true, + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + tools: { bash: true }, + parts: [ + { + type: "text", + text: expect.stringContaining("restore checkpointed session agent configuration"), + }, + ], + }, + query: { directory: "/tmp/test" }, + }) + }) + + it("recovers after five consecutive assistant messages with no text", async () => { + //#given + const promptAsyncMock = mock(async () => ({})) + const ctx = createMockContext( + [ + [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + }, + }, + ], + [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + }, + }, + ], + [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + }, + }, + ], + ], + promptAsyncMock, + ) + const injector = createCompactionContextInjector({ ctx }) + + await injector.capture("ses_no_text_tail") + await injector.event({ + event: { type: "session.compacted", properties: { sessionID: "ses_no_text_tail" } }, + }) + + //#when + for (let index = 1; index <= 5; index++) { + await injector.event({ + event: { + type: "message.updated", + properties: { + info: { + id: `msg_${index}`, + role: "assistant", + sessionID: "ses_no_text_tail", + }, + }, + }, + }) + } + await injector.event({ + event: { type: "session.idle", properties: { sessionID: "ses_no_text_tail" } }, + }) + + //#then + expect(promptAsyncMock).toHaveBeenCalledTimes(1) + expect(promptAsyncMock).toHaveBeenCalledWith( + expect.objectContaining({ + path: { id: "ses_no_text_tail" }, + body: expect.objectContaining({ + noReply: true, + agent: "atlas", + }), + }), + ) + }) + }) }) diff --git a/src/hooks/compaction-context-injector/recovery-prompt-config.ts b/src/hooks/compaction-context-injector/recovery-prompt-config.ts new file mode 100644 index 0000000000..f0f8480c81 --- /dev/null +++ b/src/hooks/compaction-context-injector/recovery-prompt-config.ts @@ -0,0 +1,76 @@ +import type { CompactionAgentConfigCheckpoint } from "../../shared/compaction-agent-config-checkpoint" + +export type RecoveryPromptConfig = CompactionAgentConfigCheckpoint & { + agent: string +} + +function isCompactionAgent(agent: string | undefined): boolean { + return agent?.trim().toLowerCase() === "compaction" +} + +function matchesExpectedModel( + actualModel: CompactionAgentConfigCheckpoint["model"], + expectedModel: CompactionAgentConfigCheckpoint["model"], +): boolean { + if (!expectedModel) { + return true + } + + return ( + actualModel?.providerID === expectedModel.providerID && + actualModel.modelID === expectedModel.modelID + ) +} + +function matchesExpectedTools( + actualTools: CompactionAgentConfigCheckpoint["tools"], + expectedTools: CompactionAgentConfigCheckpoint["tools"], +): boolean { + if (!expectedTools) { + return true + } + + if (!actualTools) { + return false + } + + const expectedEntries = Object.entries(expectedTools) + if (expectedEntries.length !== Object.keys(actualTools).length) { + return false + } + + return expectedEntries.every( + ([toolName, isAllowed]) => actualTools[toolName] === isAllowed, + ) +} + +export function createExpectedRecoveryPromptConfig( + checkpoint: Pick & CompactionAgentConfigCheckpoint, + currentPromptConfig: CompactionAgentConfigCheckpoint, +): RecoveryPromptConfig { + const model = checkpoint.model ?? currentPromptConfig.model + const tools = checkpoint.tools ?? currentPromptConfig.tools + + return { + agent: checkpoint.agent, + ...(model ? { model } : {}), + ...(tools ? { tools } : {}), + } +} + +export function isPromptConfigRecovered( + actualPromptConfig: CompactionAgentConfigCheckpoint, + expectedPromptConfig: RecoveryPromptConfig, +): boolean { + const actualAgent = actualPromptConfig.agent + const agentMatches = + typeof actualAgent === "string" && + !isCompactionAgent(actualAgent) && + actualAgent.toLowerCase() === expectedPromptConfig.agent.toLowerCase() + + return ( + agentMatches && + matchesExpectedModel(actualPromptConfig.model, expectedPromptConfig.model) && + matchesExpectedTools(actualPromptConfig.tools, expectedPromptConfig.tools) + ) +} diff --git a/src/hooks/compaction-context-injector/recovery.test.ts b/src/hooks/compaction-context-injector/recovery.test.ts new file mode 100644 index 0000000000..5d2f885049 --- /dev/null +++ b/src/hooks/compaction-context-injector/recovery.test.ts @@ -0,0 +1,294 @@ +/// + +import { describe, expect, it } from "bun:test" +import { createCompactionContextInjector } from "./index" + +type SessionMessageResponse = Array<{ + info?: Record +}> + +type PromptAsyncInput = { + path: { id: string } + body: { + noReply?: boolean + agent?: string + model?: { providerID: string; modelID: string } + tools?: Record + parts: Array<{ type: "text"; text: string }> + } + query?: { directory: string } +} + +function createPromptAsyncRecorder(): { + calls: PromptAsyncInput[] + promptAsync: (input: PromptAsyncInput) => Promise> +} { + const calls: PromptAsyncInput[] = [] + + return { + calls, + promptAsync: async (input: PromptAsyncInput) => { + calls.push(input) + return {} + }, + } +} + +function createMockContext( + messageResponses: SessionMessageResponse[], + promptAsync: (input: PromptAsyncInput) => Promise>, +) { + let callIndex = 0 + + return { + client: { + session: { + messages: async () => { + const response = + messageResponses[Math.min(callIndex, messageResponses.length - 1)] ?? [] + callIndex += 1 + return { data: response } + }, + promptAsync, + }, + }, + directory: "/tmp/test", + } +} + +function createAssistantMessageUpdatedEvent(sessionID: string, messageID: string) { + return { + event: { + type: "message.updated", + properties: { + info: { + id: messageID, + role: "assistant", + sessionID, + }, + }, + }, + } as const +} + +function createMeaningfulPartUpdatedEvent( + sessionID: string, + messageID: string, + type: "reasoning" | "tool_use", +) { + return { + event: { + type: "message.part.updated", + properties: { + part: { + messageID, + sessionID, + type, + ...(type === "reasoning" ? { text: "thinking" } : {}), + }, + }, + }, + } as const +} + +describe("createCompactionContextInjector recovery", () => { + it("re-injects after compaction when agent and model match but tools are missing", async () => { + //#given + const promptAsyncRecorder = createPromptAsyncRecorder() + const ctx = createMockContext( + [ + [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + tools: { bash: true }, + }, + }, + ], + [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + }, + }, + ], + [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + }, + }, + ], + [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + tools: { bash: true }, + }, + }, + ], + ], + promptAsyncRecorder.promptAsync, + ) + const injector = createCompactionContextInjector({ ctx }) + + //#when + await injector.capture("ses_missing_tools") + await injector.event({ + event: { type: "session.compacted", properties: { sessionID: "ses_missing_tools" } }, + }) + + //#then + expect(promptAsyncRecorder.calls.length).toBe(1) + expect(promptAsyncRecorder.calls[0]?.body.agent).toBe("atlas") + expect(promptAsyncRecorder.calls[0]?.body.model).toEqual({ + providerID: "openai", + modelID: "gpt-5", + }) + expect(promptAsyncRecorder.calls[0]?.body.tools).toEqual({ bash: true }) + }) + + it("retries recovery when the recovered prompt config still mismatches expected model or tools", async () => { + //#given + const promptAsyncRecorder = createPromptAsyncRecorder() + const mismatchResponse = [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-4.1" }, + }, + }, + ] + const ctx = createMockContext( + [ + [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + tools: { bash: true }, + }, + }, + ], + mismatchResponse, + mismatchResponse, + mismatchResponse, + mismatchResponse, + mismatchResponse, + mismatchResponse, + ], + promptAsyncRecorder.promptAsync, + ) + const injector = createCompactionContextInjector({ ctx }) + + //#when + await injector.capture("ses_retry_incomplete_recovery") + await injector.event({ + event: { + type: "session.compacted", + properties: { sessionID: "ses_retry_incomplete_recovery" }, + }, + }) + await injector.event({ + event: { + type: "session.compacted", + properties: { sessionID: "ses_retry_incomplete_recovery" }, + }, + }) + + //#then + expect(promptAsyncRecorder.calls.length).toBe(2) + }) + + it("does not treat reasoning-only assistant messages as a no-text tail", async () => { + //#given + const promptAsyncRecorder = createPromptAsyncRecorder() + const matchingPromptConfig = [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + tools: { bash: true }, + }, + }, + ] + const ctx = createMockContext( + [matchingPromptConfig, matchingPromptConfig, matchingPromptConfig], + promptAsyncRecorder.promptAsync, + ) + const injector = createCompactionContextInjector({ ctx }) + const sessionID = "ses_reasoning_tail" + + await injector.capture(sessionID) + await injector.event({ + event: { type: "session.compacted", properties: { sessionID } }, + }) + + //#when + for (let index = 1; index <= 5; index++) { + const messageID = `msg_reasoning_${index}` + await injector.event(createAssistantMessageUpdatedEvent(sessionID, messageID)) + await injector.event( + createMeaningfulPartUpdatedEvent(sessionID, messageID, "reasoning"), + ) + await injector.event({ + event: { type: "session.idle", properties: { sessionID } }, + }) + } + + //#then + expect(promptAsyncRecorder.calls.length).toBe(0) + }) + + it("does not treat tool_use-only assistant messages as a no-text tail", async () => { + //#given + const promptAsyncRecorder = createPromptAsyncRecorder() + const matchingPromptConfig = [ + { + info: { + role: "user", + agent: "atlas", + model: { providerID: "openai", modelID: "gpt-5" }, + tools: { bash: true }, + }, + }, + ] + const ctx = createMockContext( + [matchingPromptConfig, matchingPromptConfig, matchingPromptConfig], + promptAsyncRecorder.promptAsync, + ) + const injector = createCompactionContextInjector({ ctx }) + const sessionID = "ses_tool_use_tail" + + await injector.capture(sessionID) + await injector.event({ + event: { type: "session.compacted", properties: { sessionID } }, + }) + + //#when + for (let index = 1; index <= 5; index++) { + const messageID = `msg_tool_use_${index}` + await injector.event(createAssistantMessageUpdatedEvent(sessionID, messageID)) + await injector.event( + createMeaningfulPartUpdatedEvent(sessionID, messageID, "tool_use"), + ) + await injector.event({ + event: { type: "session.idle", properties: { sessionID } }, + }) + } + + //#then + expect(promptAsyncRecorder.calls.length).toBe(0) + }) +}) diff --git a/src/hooks/compaction-context-injector/recovery.ts b/src/hooks/compaction-context-injector/recovery.ts new file mode 100644 index 0000000000..9542d6e175 --- /dev/null +++ b/src/hooks/compaction-context-injector/recovery.ts @@ -0,0 +1,143 @@ +import { updateSessionAgent } from "../../features/claude-code-session-state" +import { + getCompactionAgentConfigCheckpoint, +} from "../../shared/compaction-agent-config-checkpoint" +import { createInternalAgentTextPart } from "../../shared/internal-initiator-marker" +import { log } from "../../shared/logger" +import { setSessionModel } from "../../shared/session-model-state" +import { setSessionTools } from "../../shared/session-tools-store" +import { + createExpectedRecoveryPromptConfig, + isPromptConfigRecovered, +} from "./recovery-prompt-config" +import { + resolveLatestSessionPromptConfig, + resolveSessionPromptConfig, +} from "./session-prompt-config-resolver" +import { AGENT_RECOVERY_PROMPT, NO_TEXT_TAIL_THRESHOLD, RECOVERY_COOLDOWN_MS, RECENT_COMPACTION_WINDOW_MS } from "./constants" +import type { CompactionContextClient } from "./types" +import type { TailMonitorState } from "./tail-monitor" + +export function createRecoveryLogic( + ctx: CompactionContextClient | undefined, + getTailState: (sessionID: string) => TailMonitorState, +) { + const recoverCheckpointedAgentConfig = async ( + sessionID: string, + reason: "session.compacted" | "no-text-tail", + ): Promise => { + if (!ctx) { + return false + } + + const checkpoint = getCompactionAgentConfigCheckpoint(sessionID) + if (!checkpoint?.agent) { + return false + } + const checkpointWithAgent = { + ...checkpoint, + agent: checkpoint.agent, + } + + const tailState = getTailState(sessionID) + const now = Date.now() + if (tailState.lastRecoveryAt && now - tailState.lastRecoveryAt < RECOVERY_COOLDOWN_MS) { + return false + } + + const currentPromptConfig = await resolveSessionPromptConfig(ctx, sessionID) + const expectedPromptConfig = createExpectedRecoveryPromptConfig( + checkpointWithAgent, + currentPromptConfig, + ) + const model = expectedPromptConfig.model + const tools = expectedPromptConfig.tools + + if (reason === "session.compacted") { + const latestPromptConfig = await resolveLatestSessionPromptConfig(ctx, sessionID) + if (isPromptConfigRecovered(latestPromptConfig, expectedPromptConfig)) { + return false + } + } + + try { + await ctx.client.session.promptAsync({ + path: { id: sessionID }, + body: { + noReply: true, + agent: expectedPromptConfig.agent, + ...(model ? { model } : {}), + ...(tools ? { tools } : {}), + parts: [createInternalAgentTextPart(AGENT_RECOVERY_PROMPT)], + }, + query: { directory: ctx.directory }, + }) + + const recoveredPromptConfig = await resolveLatestSessionPromptConfig(ctx, sessionID) + if (!isPromptConfigRecovered(recoveredPromptConfig, expectedPromptConfig)) { + log(`[compaction-context-injector] Re-injected agent config but recovery is still incomplete`, { + sessionID, + reason, + agent: expectedPromptConfig.agent, + model, + hasTools: !!tools, + recoveredPromptConfig, + }) + return false + } + + updateSessionAgent(sessionID, expectedPromptConfig.agent) + if (model) { + setSessionModel(sessionID, model) + } + if (tools) { + setSessionTools(sessionID, tools) + } + + tailState.lastRecoveryAt = now + tailState.consecutiveNoTextMessages = 0 + + log(`[compaction-context-injector] Re-injected checkpointed agent config`, { + sessionID, + reason, + agent: expectedPromptConfig.agent, + model, + }) + + return true + } catch (error) { + log(`[compaction-context-injector] Failed to re-inject checkpointed agent config`, { + sessionID, + reason, + error: String(error), + }) + return false + } + } + + const maybeWarnAboutNoTextTail = async (sessionID: string): Promise => { + const tailState = getTailState(sessionID) + if (tailState.consecutiveNoTextMessages < NO_TEXT_TAIL_THRESHOLD) { + return + } + + const recentlyCompacted = + tailState.lastCompactedAt !== undefined && + Date.now() - tailState.lastCompactedAt < RECENT_COMPACTION_WINDOW_MS + + log(`[compaction-context-injector] Detected consecutive assistant messages with no text`, { + sessionID, + consecutiveNoTextMessages: tailState.consecutiveNoTextMessages, + recentlyCompacted, + }) + + if (recentlyCompacted) { + await recoverCheckpointedAgentConfig(sessionID, "no-text-tail") + } + } + + return { + recoverCheckpointedAgentConfig, + maybeWarnAboutNoTextTail, + } +} diff --git a/src/hooks/compaction-context-injector/session-id.ts b/src/hooks/compaction-context-injector/session-id.ts new file mode 100644 index 0000000000..29de9c5c48 --- /dev/null +++ b/src/hooks/compaction-context-injector/session-id.ts @@ -0,0 +1,8 @@ +export function isCompactionAgent(agent: string | undefined): boolean { + return agent?.trim().toLowerCase() === "compaction" +} + +export function resolveSessionID(props?: Record): string | undefined { + return (props?.sessionID ?? + (props?.info as { id?: string } | undefined)?.id) as string | undefined +} diff --git a/src/hooks/compaction-context-injector/session-prompt-config-resolver.ts b/src/hooks/compaction-context-injector/session-prompt-config-resolver.ts new file mode 100644 index 0000000000..8d946ef7e6 --- /dev/null +++ b/src/hooks/compaction-context-injector/session-prompt-config-resolver.ts @@ -0,0 +1,131 @@ +import { getSessionAgent } from "../../features/claude-code-session-state" +import type { CompactionAgentConfigCheckpoint } from "../../shared/compaction-agent-config-checkpoint" +import { log } from "../../shared/logger" +import { normalizeSDKResponse } from "../../shared/normalize-sdk-response" +import { normalizePromptTools } from "../../shared/prompt-tools" +import { getSessionModel } from "../../shared/session-model-state" +import { getSessionTools } from "../../shared/session-tools-store" + +type SessionMessage = { + info?: { + agent?: string + model?: { + providerID?: string + modelID?: string + } + providerID?: string + modelID?: string + tools?: Record + } +} + +type ResolverContext = { + client: { + session: { + messages: (input: { path: { id: string } }) => Promise + } + } + directory: string +} + +function isCompactionAgent(agent: string | undefined): boolean { + return agent?.trim().toLowerCase() === "compaction" +} + +function resolveModel( + info: SessionMessage["info"], +): CompactionAgentConfigCheckpoint["model"] | undefined { + const providerID = info?.model?.providerID ?? info?.providerID + const modelID = info?.model?.modelID ?? info?.modelID + + if (!providerID || !modelID) { + return undefined + } + + return { providerID, modelID } +} + +export async function resolveSessionPromptConfig( + ctx: ResolverContext, + sessionID: string, +): Promise { + const promptConfig: CompactionAgentConfigCheckpoint = { + agent: getSessionAgent(sessionID), + model: getSessionModel(sessionID), + tools: getSessionTools(sessionID), + } + + try { + const response = await ctx.client.session.messages({ path: { id: sessionID } }) + const messages = normalizeSDKResponse(response, [] as SessionMessage[], { + preferResponseOnMissingData: true, + }) + + for (let index = messages.length - 1; index >= 0; index--) { + const info = messages[index].info + + if (!promptConfig.agent && info?.agent && !isCompactionAgent(info.agent)) { + promptConfig.agent = info.agent + } + + if (!promptConfig.model) { + const model = resolveModel(info) + if (model) { + promptConfig.model = model + } + } + + if (!promptConfig.tools) { + const tools = normalizePromptTools(info?.tools) + if (tools) { + promptConfig.tools = tools + } + } + + if (promptConfig.agent && promptConfig.model && promptConfig.tools) { + break + } + } + } catch (error) { + log("[compaction-context-injector] Failed to resolve prompt config from messages", { + sessionID, + directory: ctx.directory, + error: String(error), + }) + } + + return promptConfig +} + +export async function resolveLatestSessionPromptConfig( + ctx: ResolverContext, + sessionID: string, +): Promise { + try { + const response = await ctx.client.session.messages({ path: { id: sessionID } }) + const messages = normalizeSDKResponse(response, [] as SessionMessage[], { + preferResponseOnMissingData: true, + }) + const latestInfo = messages.at(-1)?.info + + if (!latestInfo) { + return {} + } + + const model = resolveModel(latestInfo) + const tools = normalizePromptTools(latestInfo.tools) + + return { + ...(latestInfo.agent ? { agent: latestInfo.agent } : {}), + ...(model ? { model } : {}), + ...(tools ? { tools } : {}), + } + } catch (error) { + log("[compaction-context-injector] Failed to resolve latest prompt config", { + sessionID, + directory: ctx.directory, + error: String(error), + }) + return {} + } +} diff --git a/src/hooks/compaction-context-injector/tail-monitor.ts b/src/hooks/compaction-context-injector/tail-monitor.ts new file mode 100644 index 0000000000..c936b94735 --- /dev/null +++ b/src/hooks/compaction-context-injector/tail-monitor.ts @@ -0,0 +1,52 @@ +const MEANINGFUL_ASSISTANT_PART_TYPES = new Set([ + "reasoning", + "tool", + "tool_use", +]) + +export type TailMonitorState = { + currentMessageID?: string + currentHasOutput: boolean + consecutiveNoTextMessages: number + lastCompactedAt?: number + lastRecoveryAt?: number +} + +export function finalizeTrackedAssistantMessage( + state: TailMonitorState, +): number { + if (!state.currentMessageID) { + return state.consecutiveNoTextMessages + } + + state.consecutiveNoTextMessages = state.currentHasOutput + ? 0 + : state.consecutiveNoTextMessages + 1 + state.currentMessageID = undefined + state.currentHasOutput = false + + return state.consecutiveNoTextMessages +} + +export function shouldTreatAssistantPartAsOutput(part: { + type?: string + text?: string +}): boolean { + if (part.type === "text") { + return !!part.text?.trim() + } + + return typeof part.type === "string" && MEANINGFUL_ASSISTANT_PART_TYPES.has(part.type) +} + +export function trackAssistantOutput( + state: TailMonitorState, + messageID?: string, +): void { + if (messageID && !state.currentMessageID) { + state.currentMessageID = messageID + } + + state.currentHasOutput = true + state.consecutiveNoTextMessages = 0 +} diff --git a/src/hooks/compaction-context-injector/types.ts b/src/hooks/compaction-context-injector/types.ts new file mode 100644 index 0000000000..b97c2e6f6e --- /dev/null +++ b/src/hooks/compaction-context-injector/types.ts @@ -0,0 +1,25 @@ +export interface CompactionContextInjector { + capture: (sessionID: string) => Promise + inject: (sessionID?: string) => string + event: (input: { event: { type: string; properties?: unknown } }) => Promise +} + +export type CompactionContextClient = { + client: { + session: { + messages: (input: { path: { id: string } }) => Promise + promptAsync: (input: { + path: { id: string } + body: { + noReply?: boolean + agent?: string + model?: { providerID: string; modelID: string } + tools?: Record + parts: Array<{ type: "text"; text: string }> + } + query?: { directory: string } + }) => Promise + } + } + directory: string +} diff --git a/src/index.test.ts b/src/index.test.ts index c0ca3212f0..0082d1e2c0 100644 --- a/src/index.test.ts +++ b/src/index.test.ts @@ -2,6 +2,10 @@ import { describe, expect, it, mock } from "bun:test" describe("experimental.session.compacting handler", () => { function createCompactingHandler(hooks: { + compactionContextInjector?: { + capture: (sessionID: string) => Promise + inject: (sessionID: string) => string + } compactionTodoPreserver?: { capture: (sessionID: string) => Promise } claudeCodeHooks?: { "experimental.session.compacting"?: ( @@ -9,19 +13,19 @@ describe("experimental.session.compacting handler", () => { output: { context: string[] }, ) => Promise } - compactionContextInjector?: (sessionID: string) => string }) { return async ( _input: { sessionID: string }, output: { context: string[] }, ): Promise => { + await hooks.compactionContextInjector?.capture(_input.sessionID) await hooks.compactionTodoPreserver?.capture(_input.sessionID) await hooks.claudeCodeHooks?.["experimental.session.compacting"]?.( _input, output, ) if (hooks.compactionContextInjector) { - output.context.push(hooks.compactionContextInjector(_input.sessionID)) + output.context.push(hooks.compactionContextInjector.inject(_input.sessionID)) } } } @@ -33,6 +37,15 @@ describe("experimental.session.compacting handler", () => { const callOrder: string[] = [] const handler = createCompactingHandler({ + compactionContextInjector: { + capture: mock(async () => { + callOrder.push("checkpointCapture") + }), + inject: mock((sessionID: string) => { + callOrder.push("contextInjector") + return `context-for-${sessionID}` + }), + }, compactionTodoPreserver: { capture: mock(async () => { callOrder.push("capture") }), }, @@ -41,16 +54,12 @@ describe("experimental.session.compacting handler", () => { callOrder.push("preCompact") }), }, - compactionContextInjector: mock((sessionID: string) => { - callOrder.push("contextInjector") - return `context-for-${sessionID}` - }), }) const output = { context: [] as string[] } await handler({ sessionID: "ses_test" }, output) - expect(callOrder).toEqual(["capture", "preCompact", "contextInjector"]) + expect(callOrder).toEqual(["checkpointCapture", "capture", "preCompact", "contextInjector"]) expect(output.context).toEqual(["context-for-ses_test"]) }) @@ -77,17 +86,22 @@ describe("experimental.session.compacting handler", () => { //#then handler completes without error and other hooks still run it("handles null claudeCodeHooks gracefully", async () => { const captureMock = mock(async () => {}) + const checkpointCaptureMock = mock(async () => {}) const contextMock = mock(() => "injected-context") const handler = createCompactingHandler({ + compactionContextInjector: { + capture: checkpointCaptureMock, + inject: contextMock, + }, compactionTodoPreserver: { capture: captureMock }, claudeCodeHooks: undefined, - compactionContextInjector: contextMock, }) const output = { context: [] as string[] } await handler({ sessionID: "ses_test" }, output) + expect(checkpointCaptureMock).toHaveBeenCalledWith("ses_test") expect(captureMock).toHaveBeenCalledWith("ses_test") expect(contextMock).toHaveBeenCalledWith("ses_test") expect(output.context).toEqual(["injected-context"]) diff --git a/src/index.ts b/src/index.ts index 3ae22411ce..4b6f54dd57 100644 --- a/src/index.ts +++ b/src/index.ts @@ -83,13 +83,14 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { _input: { sessionID: string }, output: { context: string[] }, ): Promise => { + await hooks.compactionContextInjector?.capture(_input.sessionID) await hooks.compactionTodoPreserver?.capture(_input.sessionID) await hooks.claudeCodeHooks?.["experimental.session.compacting"]?.( _input, output, ) if (hooks.compactionContextInjector) { - output.context.push(hooks.compactionContextInjector(_input.sessionID)) + output.context.push(hooks.compactionContextInjector.inject(_input.sessionID)) } }, } diff --git a/src/plugin/event.ts b/src/plugin/event.ts index 0b7c3c2ad0..4d618e83a9 100644 --- a/src/plugin/event.ts +++ b/src/plugin/event.ts @@ -184,6 +184,7 @@ export function createEventHandler(args: { await Promise.resolve(hooks.interactiveBashSession?.event?.(input as EventInput)); await Promise.resolve(hooks.ralphLoop?.event?.(input)); await Promise.resolve(hooks.stopContinuationGuard?.event?.(input)); + await Promise.resolve(hooks.compactionContextInjector?.event?.(input)); await Promise.resolve(hooks.compactionTodoPreserver?.event?.(input)); await Promise.resolve(hooks.writeExistingFileGuard?.event?.(input)); await Promise.resolve(hooks.atlasHook?.handler?.(input)); diff --git a/src/plugin/hooks/create-continuation-hooks.ts b/src/plugin/hooks/create-continuation-hooks.ts index 9cedce2120..092064cbf7 100644 --- a/src/plugin/hooks/create-continuation-hooks.ts +++ b/src/plugin/hooks/create-continuation-hooks.ts @@ -56,7 +56,8 @@ export function createContinuationHooks(args: { : null const compactionContextInjector = isHookEnabled("compaction-context-injector") - ? safeHook("compaction-context-injector", () => createCompactionContextInjector(backgroundManager)) + ? safeHook("compaction-context-injector", () => + createCompactionContextInjector({ ctx, backgroundManager })) : null const compactionTodoPreserver = isHookEnabled("compaction-todo-preserver") diff --git a/src/shared/compaction-agent-config-checkpoint.ts b/src/shared/compaction-agent-config-checkpoint.ts new file mode 100644 index 0000000000..47a3cabb32 --- /dev/null +++ b/src/shared/compaction-agent-config-checkpoint.ts @@ -0,0 +1,42 @@ +export type CompactionAgentConfigCheckpoint = { + agent?: string + model?: { providerID: string; modelID: string } + tools?: Record +} + +const checkpoints = new Map() + +function cloneCheckpoint( + checkpoint: CompactionAgentConfigCheckpoint, +): CompactionAgentConfigCheckpoint { + return { + ...(checkpoint.agent ? { agent: checkpoint.agent } : {}), + ...(checkpoint.model + ? { + model: { + providerID: checkpoint.model.providerID, + modelID: checkpoint.model.modelID, + }, + } + : {}), + ...(checkpoint.tools ? { tools: { ...checkpoint.tools } } : {}), + } +} + +export function setCompactionAgentConfigCheckpoint( + sessionID: string, + checkpoint: CompactionAgentConfigCheckpoint, +): void { + checkpoints.set(sessionID, cloneCheckpoint(checkpoint)) +} + +export function getCompactionAgentConfigCheckpoint( + sessionID: string, +): CompactionAgentConfigCheckpoint | undefined { + const checkpoint = checkpoints.get(sessionID) + return checkpoint ? cloneCheckpoint(checkpoint) : undefined +} + +export function clearCompactionAgentConfigCheckpoint(sessionID: string): void { + checkpoints.delete(sessionID) +}