diff --git a/src/hooks/anthropic-context-window-limit-recovery/aggressive-truncation-strategy.test.ts b/src/hooks/anthropic-context-window-limit-recovery/aggressive-truncation-strategy.test.ts new file mode 100644 index 000000000..db1810074 --- /dev/null +++ b/src/hooks/anthropic-context-window-limit-recovery/aggressive-truncation-strategy.test.ts @@ -0,0 +1,190 @@ +/// +import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test" + +import type { AutoCompactState } from "./types" + +type PromptAsyncCall = { + path: { id: string } + body: { + auto?: boolean + agent?: string + model?: { providerID: string; modelID: string } + variant?: string + tools?: Record + parts?: unknown + } + query: { directory: string } +} + +const truncateUntilTargetTokensMock = mock(async () => ({ + truncatedCount: 1, + totalBytesRemoved: 1000, + truncatedTools: [{ toolName: "bash" }], + sufficient: true, +})) + +mock.module("./storage", () => ({ + truncateUntilTargetTokens: truncateUntilTargetTokensMock, +})) + +const findNearestMessageWithFieldsFromSDKMock = mock(async () => null) +const findNearestMessageWithFieldsMock = mock(() => null) + +mock.module("../../features/hook-message-injector", () => ({ + findNearestMessageWithFieldsFromSDK: findNearestMessageWithFieldsFromSDKMock, + findNearestMessageWithFields: findNearestMessageWithFieldsMock, +})) + +const sessionAgentMap = new Map() +const resolveRegisteredAgentNameMock = mock((name: string | undefined) => name) + +mock.module("../../features/claude-code-session-state/state", () => ({ + _resetForTesting: () => { sessionAgentMap.clear() }, + setSessionAgent: (sessionID: string, agent: string) => { sessionAgentMap.set(sessionID, agent) }, + getSessionAgent: (sessionID: string) => sessionAgentMap.get(sessionID), + resolveRegisteredAgentName: resolveRegisteredAgentNameMock, + registerAgentName: () => {}, + isAgentRegistered: () => false, + resolveInheritedPromptTools: () => undefined, +})) + +import { runAggressiveTruncationStrategy } from "./aggressive-truncation-strategy" + +type FakeClient = { + session: { promptAsync: (input: PromptAsyncCall) => Promise } + tui: { showToast: (input: unknown) => Promise } +} + +function createRecordingClient(): { client: FakeClient; calls: PromptAsyncCall[] } { + const calls: PromptAsyncCall[] = [] + const client: FakeClient = { + session: { + promptAsync: async (input: PromptAsyncCall) => { + calls.push(input) + return undefined + }, + }, + tui: { + showToast: async () => undefined, + }, + } + return { client, calls } +} + +function createAutoCompactState(): AutoCompactState { + return { + pendingCompact: new Set(), + errorDataBySession: new Map(), + retryStateBySession: new Map(), + retryTimerBySession: new Map(), + truncateStateBySession: new Map(), + emptyContentAttemptBySession: new Map(), + compactionInProgress: new Set(), + } +} + +async function flushDeferredPrompt(): Promise { + await new Promise((resolve) => setTimeout(resolve, 600)) +} + +describe("runAggressiveTruncationStrategy - pins agent/model/variant on recovered promptAsync", () => { + beforeEach(() => { + sessionAgentMap.clear() + truncateUntilTargetTokensMock.mockClear() + findNearestMessageWithFieldsFromSDKMock.mockClear() + findNearestMessageWithFieldsMock.mockClear() + resolveRegisteredAgentNameMock.mockClear() + findNearestMessageWithFieldsFromSDKMock.mockResolvedValue(null) + findNearestMessageWithFieldsMock.mockReturnValue(null) + resolveRegisteredAgentNameMock.mockImplementation((name: string | undefined) => name) + }) + + afterEach(() => { + sessionAgentMap.clear() + }) + + test("includes the session's resolved agent on promptAsync when agent is known", async () => { + // given + const { client, calls } = createRecordingClient() + const sessionID = "session-truncation-agent" + sessionAgentMap.set(sessionID, "sisyphus-junior") + + // when + await runAggressiveTruncationStrategy({ + sessionID, + autoCompactState: createAutoCompactState(), + client: client as never, + directory: "/tmp/test-truncation", + truncateAttempt: 0, + currentTokens: 250_000, + maxTokens: 200_000, + }) + await flushDeferredPrompt() + + // then + expect(calls).toHaveLength(1) + expect(calls[0].path.id).toBe(sessionID) + expect(calls[0].body.agent).toBe("sisyphus-junior") + expect(calls[0].body.auto).toBe(true) + }) + + test("pins provider/model/variant resolved from the nearest prior assistant message", async () => { + // given + const { client, calls } = createRecordingClient() + const sessionID = "session-truncation-model" + findNearestMessageWithFieldsFromSDKMock.mockResolvedValue({ + agent: "atlas", + model: { providerID: "anthropic", modelID: "claude-opus-4-7", variant: "high" }, + tools: undefined, + } as never) + findNearestMessageWithFieldsMock.mockReturnValue({ + agent: "atlas", + model: { providerID: "anthropic", modelID: "claude-opus-4-7", variant: "high" }, + tools: undefined, + } as never) + + // when + await runAggressiveTruncationStrategy({ + sessionID, + autoCompactState: createAutoCompactState(), + client: client as never, + directory: "/tmp/test-truncation", + truncateAttempt: 0, + currentTokens: 250_000, + maxTokens: 200_000, + }) + await flushDeferredPrompt() + + // then + expect(calls).toHaveLength(1) + expect(calls[0].body.agent).toBe("atlas") + expect(calls[0].body.model).toEqual({ providerID: "anthropic", modelID: "claude-opus-4-7" }) + expect(calls[0].body.variant).toBe("high") + expect(calls[0].body.auto).toBe(true) + }) + + test("omits agent/model/variant when the session has nothing resolvable", async () => { + // given + const { client, calls } = createRecordingClient() + const sessionID = "session-truncation-empty" + + // when + await runAggressiveTruncationStrategy({ + sessionID, + autoCompactState: createAutoCompactState(), + client: client as never, + directory: "/tmp/test-truncation", + truncateAttempt: 0, + currentTokens: 250_000, + maxTokens: 200_000, + }) + await flushDeferredPrompt() + + // then + expect(calls).toHaveLength(1) + expect(calls[0].body.agent).toBeUndefined() + expect(calls[0].body.model).toBeUndefined() + expect(calls[0].body.variant).toBeUndefined() + expect(calls[0].body.auto).toBe(true) + }) +}) diff --git a/src/hooks/anthropic-context-window-limit-recovery/aggressive-truncation-strategy.ts b/src/hooks/anthropic-context-window-limit-recovery/aggressive-truncation-strategy.ts index 88f82f1d4..34660e74b 100644 --- a/src/hooks/anthropic-context-window-limit-recovery/aggressive-truncation-strategy.ts +++ b/src/hooks/anthropic-context-window-limit-recovery/aggressive-truncation-strategy.ts @@ -5,7 +5,18 @@ import type { Client } from "./client" import { clearSessionState } from "./state" import { formatBytes } from "./message-builder" import { log } from "../../shared/logger" -import { resolveInheritedPromptTools } from "../../shared" +import { + getMessageDir, + resolveInheritedPromptTools, +} from "../../shared" +import { + getSessionAgent, + resolveRegisteredAgentName, +} from "../../features/claude-code-session-state/state" +import { + findNearestMessageWithFields, + findNearestMessageWithFieldsFromSDK, +} from "../../features/hook-message-injector" export async function runAggressiveTruncationStrategy(params: { sessionID: string @@ -62,11 +73,27 @@ export async function runAggressiveTruncationStrategy(params: { clearSessionState(params.autoCompactState, params.sessionID) setTimeout(async () => { try { - const inheritedTools = resolveInheritedPromptTools(params.sessionID) + const sdkMessage = await findNearestMessageWithFieldsFromSDK(params.client, params.sessionID) + const previousMessage = sdkMessage ?? (() => { + const messageDir = getMessageDir(params.sessionID) + return messageDir ? findNearestMessageWithFields(messageDir) : null + })() + + const agentName = getSessionAgent(params.sessionID) ?? previousMessage?.agent + const launchAgent = resolveRegisteredAgentName(agentName) + const launchModel = previousMessage?.model?.providerID && previousMessage.model.modelID + ? { providerID: previousMessage.model.providerID, modelID: previousMessage.model.modelID } + : undefined + const launchVariant = previousMessage?.model?.variant + const inheritedTools = resolveInheritedPromptTools(params.sessionID, previousMessage?.tools) + await params.client.session.promptAsync({ path: { id: params.sessionID }, body: { auto: true, + ...(launchAgent ? { agent: launchAgent } : {}), + ...(launchModel ? { model: launchModel } : {}), + ...(launchVariant ? { variant: launchVariant } : {}), ...(inheritedTools ? { tools: inheritedTools } : {}), } as never, query: { directory: params.directory },