feat(hooks): add aggressive truncation strategy tests for context-window recovery

src/hooks/anthropic-context-window-limit-recovery/aggressive-truncation-strategy.test.ts

@@ -0,0 +1,190 @@
+/// <reference types="bun-types" />
+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<string, boolean>
+    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<string, string>()
+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<unknown> }
+  tui: { showToast: (input: unknown) => Promise<unknown> }
+}
+
+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<string>(),
+    errorDataBySession: new Map(),
+    retryStateBySession: new Map(),
+    retryTimerBySession: new Map(),
+    truncateStateBySession: new Map(),
+    emptyContentAttemptBySession: new Map(),
+    compactionInProgress: new Set<string>(),
+  }
+}
+
+async function flushDeferredPrompt(): Promise<void> {
+  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)
+  })
+})

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 },