From 285895493d1df7c7321abc8cbfb8109dd73e246a Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sun, 8 Mar 2026 02:23:22 +0900 Subject: [PATCH] fix(compaction): recover checkpointed agent config after compaction Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-opencode) Co-authored-by: Sisyphus --- src/hooks/compaction-context-injector/hook.ts | 323 +++++++++++++++++- .../compaction-context-injector/index.test.ts | 178 +++++++++- 2 files changed, 490 insertions(+), 11 deletions(-) diff --git a/src/hooks/compaction-context-injector/hook.ts b/src/hooks/compaction-context-injector/hook.ts index 0f3a912fe..406846c38 100644 --- a/src/hooks/compaction-context-injector/hook.ts +++ b/src/hooks/compaction-context-injector/hook.ts @@ -1,8 +1,94 @@ import type { BackgroundManager } from "../../features/background-agent" +import { updateSessionAgent } from "../../features/claude-code-session-state" +import { + clearCompactionAgentConfigCheckpoint, + getCompactionAgentConfigCheckpoint, + setCompactionAgentConfigCheckpoint, +} 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 { createSystemDirective, SystemDirectiveTypes, } from "../../shared/system-directive" +import { + resolveLatestSessionPromptConfig, + resolveSessionPromptConfig, +} from "./session-prompt-config-resolver" + +const HOOK_NAME = "compaction-context-injector" +const AGENT_RECOVERY_PROMPT = "[restore checkpointed session agent configuration after compaction]" +const NO_TEXT_TAIL_THRESHOLD = 5 +const RECOVERY_COOLDOWN_MS = 60_000 +const RECENT_COMPACTION_WINDOW_MS = 10 * 60 * 1000 + +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 +} + +type TailMonitorState = { + currentMessageID?: string + currentHasText: boolean + consecutiveNoTextMessages: number + lastCompactedAt?: number + lastRecoveryAt?: number +} + +export interface CompactionContextInjector { + capture: (sessionID: string) => Promise + inject: (sessionID?: string) => string + event: (input: { event: { type: string; properties?: unknown } }) => Promise +} + +function isCompactionAgent(agent: string | undefined): boolean { + return agent?.trim().toLowerCase() === "compaction" +} + +function resolveSessionID(props?: Record): string | undefined { + return (props?.sessionID ?? + (props?.info as { id?: string } | undefined)?.id) as string | undefined +} + +function finalizeTrackedAssistantMessage(state: TailMonitorState): number { + if (!state.currentMessageID) { + return state.consecutiveNoTextMessages + } + + state.consecutiveNoTextMessages = state.currentHasText + ? 0 + : state.consecutiveNoTextMessages + 1 + state.currentMessageID = undefined + state.currentHasText = false + + return state.consecutiveNoTextMessages +} + +function trackAssistantText(state: TailMonitorState, messageID?: string): void { + if (messageID && !state.currentMessageID) { + state.currentMessageID = messageID + } + + state.currentHasText = true + state.consecutiveNoTextMessages = 0 +} const COMPACTION_CONTEXT_PROMPT = `${createSystemDirective(SystemDirectiveTypes.COMPACTION_CONTEXT)} @@ -56,8 +142,146 @@ This section is CRITICAL for reviewer agents (momus, oracle) to maintain continu This context is critical for maintaining continuity after compaction. ` -export function createCompactionContextInjector(backgroundManager?: BackgroundManager) { - return (sessionID?: string): string => { +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 = { + currentHasText: false, + consecutiveNoTextMessages: 0, + } + tailStates.set(sessionID, created) + return created + } + + 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 tailState = getTailState(sessionID) + const now = Date.now() + if (tailState.lastRecoveryAt && now - tailState.lastRecoveryAt < RECOVERY_COOLDOWN_MS) { + return false + } + + if (reason === "session.compacted") { + const latestPromptConfig = await resolveLatestSessionPromptConfig(ctx, sessionID) + const latestAgentMatchesCheckpoint = + typeof latestPromptConfig.agent === "string" && + latestPromptConfig.agent.toLowerCase() === checkpoint.agent.toLowerCase() && + !isCompactionAgent(latestPromptConfig.agent) + + if (latestAgentMatchesCheckpoint && latestPromptConfig.model) { + return false + } + } + + const currentPromptConfig = await resolveSessionPromptConfig(ctx, sessionID) + const model = checkpoint.model ?? currentPromptConfig.model + const tools = checkpoint.tools ?? currentPromptConfig.tools + + try { + await ctx.client.session.promptAsync({ + path: { id: sessionID }, + body: { + noReply: true, + agent: checkpoint.agent, + ...(model ? { model } : {}), + ...(tools ? { tools } : {}), + parts: [createInternalAgentTextPart(AGENT_RECOVERY_PROMPT)], + }, + query: { directory: ctx.directory }, + }) + + updateSessionAgent(sessionID, checkpoint.agent) + if (model) { + setSessionModel(sessionID, model) + } + if (tools) { + setSessionTools(sessionID, tools) + } + + tailState.lastRecoveryAt = now + tailState.consecutiveNoTextMessages = 0 + + log(`[${HOOK_NAME}] Re-injected checkpointed agent config`, { + sessionID, + reason, + agent: checkpoint.agent, + model, + }) + + return true + } catch (error) { + log(`[${HOOK_NAME}] 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(`[${HOOK_NAME}] Detected consecutive assistant messages with no text`, { + sessionID, + consecutiveNoTextMessages: tailState.consecutiveNoTextMessages, + recentlyCompacted, + }) + + if (recentlyCompacted) { + await recoverCheckpointedAgentConfig(sessionID, "no-text-tail") + } + } + + 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(`[${HOOK_NAME}] 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 +293,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.currentHasText = 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 + } + + trackAssistantText(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 || part.type !== "text" || !part.text?.trim()) { + return + } + + trackAssistantText(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 a2813916f..9eacd0cdd 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", + }), + }), + ) + }) + }) })