diff --git a/src/hooks/anthropic-context-window-limit-recovery/deduplication-recovery.ts b/src/hooks/anthropic-context-window-limit-recovery/deduplication-recovery.ts new file mode 100644 index 00000000..d7cb0314 --- /dev/null +++ b/src/hooks/anthropic-context-window-limit-recovery/deduplication-recovery.ts @@ -0,0 +1,71 @@ +import type { ParsedTokenLimitError } from "./types" +import type { ExperimentalConfig } from "../../config" +import type { DeduplicationConfig } from "./pruning-deduplication" +import type { PruningState } from "./pruning-types" +import { executeDeduplication } from "./pruning-deduplication" +import { truncateToolOutputsByCallId } from "./pruning-tool-output-truncation" +import { log } from "../../shared/logger" + +function createPruningState(): PruningState { + return { + toolIdsToPrune: new Set(), + currentTurn: 0, + fileOperations: new Map(), + toolSignatures: new Map(), + erroredTools: new Map(), + } +} + +function isPromptTooLongError(parsed: ParsedTokenLimitError): boolean { + return !parsed.errorType.toLowerCase().includes("non-empty content") +} + +function getDeduplicationPlan( + experimental?: ExperimentalConfig, +): { config: DeduplicationConfig; protectedTools: Set } | null { + const pruningConfig = experimental?.dynamic_context_pruning + if (!pruningConfig?.enabled) return null + + const deduplicationEnabled = pruningConfig.strategies?.deduplication?.enabled + if (deduplicationEnabled === false) return null + + const protectedTools = new Set(pruningConfig.protected_tools ?? []) + return { + config: { + enabled: true, + protectedTools: pruningConfig.protected_tools ?? [], + }, + protectedTools, + } +} + +export async function attemptDeduplicationRecovery( + sessionID: string, + parsed: ParsedTokenLimitError, + experimental: ExperimentalConfig | undefined, +): Promise { + if (!isPromptTooLongError(parsed)) return + + const plan = getDeduplicationPlan(experimental) + if (!plan) return + + const pruningState = createPruningState() + const prunedCount = executeDeduplication( + sessionID, + pruningState, + plan.config, + plan.protectedTools, + ) + const { truncatedCount } = truncateToolOutputsByCallId( + sessionID, + pruningState.toolIdsToPrune, + ) + + if (prunedCount > 0 || truncatedCount > 0) { + log("[auto-compact] deduplication recovery applied", { + sessionID, + prunedCount, + truncatedCount, + }) + } +} diff --git a/src/hooks/anthropic-context-window-limit-recovery/index.ts b/src/hooks/anthropic-context-window-limit-recovery/index.ts index cd8d1246..205170a9 100644 --- a/src/hooks/anthropic-context-window-limit-recovery/index.ts +++ b/src/hooks/anthropic-context-window-limit-recovery/index.ts @@ -1,151 +1,5 @@ -import type { PluginInput } from "@opencode-ai/plugin" -import type { AutoCompactState, ParsedTokenLimitError } from "./types" -import type { ExperimentalConfig } from "../../config" -import { parseAnthropicTokenLimitError } from "./parser" -import { executeCompact, getLastAssistant } from "./executor" -import { log } from "../../shared/logger" - -export interface AnthropicContextWindowLimitRecoveryOptions { - experimental?: ExperimentalConfig -} - -function createRecoveryState(): AutoCompactState { - return { - pendingCompact: new Set(), - errorDataBySession: new Map(), - retryStateBySession: new Map(), - truncateStateBySession: new Map(), - emptyContentAttemptBySession: new Map(), - compactionInProgress: new Set(), - } -} - -export function createAnthropicContextWindowLimitRecoveryHook(ctx: PluginInput, options?: AnthropicContextWindowLimitRecoveryOptions) { - const autoCompactState = createRecoveryState() - const experimental = options?.experimental - - const eventHandler = async ({ event }: { event: { type: string; properties?: unknown } }) => { - const props = event.properties as Record | undefined - - if (event.type === "session.deleted") { - const sessionInfo = props?.info as { id?: string } | undefined - if (sessionInfo?.id) { - autoCompactState.pendingCompact.delete(sessionInfo.id) - autoCompactState.errorDataBySession.delete(sessionInfo.id) - autoCompactState.retryStateBySession.delete(sessionInfo.id) - autoCompactState.truncateStateBySession.delete(sessionInfo.id) - autoCompactState.emptyContentAttemptBySession.delete(sessionInfo.id) - autoCompactState.compactionInProgress.delete(sessionInfo.id) - } - return - } - - if (event.type === "session.error") { - const sessionID = props?.sessionID as string | undefined - log("[auto-compact] session.error received", { sessionID, error: props?.error }) - if (!sessionID) return - - const parsed = parseAnthropicTokenLimitError(props?.error) - log("[auto-compact] parsed result", { parsed, hasError: !!props?.error }) - if (parsed) { - autoCompactState.pendingCompact.add(sessionID) - autoCompactState.errorDataBySession.set(sessionID, parsed) - - if (autoCompactState.compactionInProgress.has(sessionID)) { - return - } - - const lastAssistant = await getLastAssistant(sessionID, ctx.client, ctx.directory) - const providerID = parsed.providerID ?? (lastAssistant?.providerID as string | undefined) - const modelID = parsed.modelID ?? (lastAssistant?.modelID as string | undefined) - - await ctx.client.tui - .showToast({ - body: { - title: "Context Limit Hit", - message: "Truncating large tool outputs and recovering...", - variant: "warning" as const, - duration: 3000, - }, - }) - .catch(() => {}) - - setTimeout(() => { - executeCompact( - sessionID, - { providerID, modelID }, - autoCompactState, - ctx.client, - ctx.directory, - experimental - ) - }, 300) - } - return - } - - if (event.type === "message.updated") { - const info = props?.info as Record | undefined - const sessionID = info?.sessionID as string | undefined - - if (sessionID && info?.role === "assistant" && info.error) { - log("[auto-compact] message.updated with error", { sessionID, error: info.error }) - const parsed = parseAnthropicTokenLimitError(info.error) - log("[auto-compact] message.updated parsed result", { parsed }) - if (parsed) { - parsed.providerID = info.providerID as string | undefined - parsed.modelID = info.modelID as string | undefined - autoCompactState.pendingCompact.add(sessionID) - autoCompactState.errorDataBySession.set(sessionID, parsed) - } - } - return - } - - if (event.type === "session.idle") { - const sessionID = props?.sessionID as string | undefined - if (!sessionID) return - - if (!autoCompactState.pendingCompact.has(sessionID)) return - - const errorData = autoCompactState.errorDataBySession.get(sessionID) - const lastAssistant = await getLastAssistant(sessionID, ctx.client, ctx.directory) - - if (lastAssistant?.summary === true) { - autoCompactState.pendingCompact.delete(sessionID) - return - } - - const providerID = errorData?.providerID ?? (lastAssistant?.providerID as string | undefined) - const modelID = errorData?.modelID ?? (lastAssistant?.modelID as string | undefined) - - await ctx.client.tui - .showToast({ - body: { - title: "Auto Compact", - message: "Token limit exceeded. Attempting recovery...", - variant: "warning" as const, - duration: 3000, - }, - }) - .catch(() => {}) - - await executeCompact( - sessionID, - { providerID, modelID }, - autoCompactState, - ctx.client, - ctx.directory, - experimental - ) - } - } - - return { - event: eventHandler, - } -} - +export { createAnthropicContextWindowLimitRecoveryHook } from "./recovery-hook" +export type { AnthropicContextWindowLimitRecoveryOptions } from "./recovery-hook" export type { AutoCompactState, ParsedTokenLimitError, TruncateState } from "./types" export { parseAnthropicTokenLimitError } from "./parser" export { executeCompact, getLastAssistant } from "./executor" diff --git a/src/hooks/anthropic-context-window-limit-recovery/pruning-tool-output-truncation.ts b/src/hooks/anthropic-context-window-limit-recovery/pruning-tool-output-truncation.ts new file mode 100644 index 00000000..0481e94c --- /dev/null +++ b/src/hooks/anthropic-context-window-limit-recovery/pruning-tool-output-truncation.ts @@ -0,0 +1,97 @@ +import { existsSync, readdirSync, readFileSync } from "node:fs" +import { join } from "node:path" +import { getOpenCodeStorageDir } from "../../shared/data-path" +import { truncateToolResult } from "./storage" +import { log } from "../../shared/logger" + +interface StoredToolPart { + type?: string + callID?: string + truncated?: boolean + state?: { + output?: string + } +} + +function getMessageStorage(): string { + return join(getOpenCodeStorageDir(), "message") +} + +function getPartStorage(): string { + return join(getOpenCodeStorageDir(), "part") +} + +function getMessageDir(sessionID: string): string | null { + const messageStorage = getMessageStorage() + if (!existsSync(messageStorage)) return null + + const directPath = join(messageStorage, sessionID) + if (existsSync(directPath)) return directPath + + for (const dir of readdirSync(messageStorage)) { + const sessionPath = join(messageStorage, dir, sessionID) + if (existsSync(sessionPath)) return sessionPath + } + + return null +} + +function getMessageIds(sessionID: string): string[] { + const messageDir = getMessageDir(sessionID) + if (!messageDir) return [] + + const messageIds: string[] = [] + for (const file of readdirSync(messageDir)) { + if (!file.endsWith(".json")) continue + messageIds.push(file.replace(".json", "")) + } + + return messageIds +} + +export function truncateToolOutputsByCallId( + sessionID: string, + callIds: Set, +): { truncatedCount: number } { + if (callIds.size === 0) return { truncatedCount: 0 } + + const messageIds = getMessageIds(sessionID) + if (messageIds.length === 0) return { truncatedCount: 0 } + + let truncatedCount = 0 + + for (const messageID of messageIds) { + const partDir = join(getPartStorage(), messageID) + if (!existsSync(partDir)) continue + + for (const file of readdirSync(partDir)) { + if (!file.endsWith(".json")) continue + const partPath = join(partDir, file) + + try { + const content = readFileSync(partPath, "utf-8") + const part = JSON.parse(content) as StoredToolPart + + if (part.type !== "tool" || !part.callID) continue + if (!callIds.has(part.callID)) continue + if (!part.state?.output || part.truncated) continue + + const result = truncateToolResult(partPath) + if (result.success) { + truncatedCount++ + } + } catch { + continue + } + } + } + + if (truncatedCount > 0) { + log("[auto-compact] pruned duplicate tool outputs", { + sessionID, + truncatedCount, + }) + } + + return { truncatedCount } +} diff --git a/src/hooks/anthropic-context-window-limit-recovery/recovery-deduplication.test.ts b/src/hooks/anthropic-context-window-limit-recovery/recovery-deduplication.test.ts new file mode 100644 index 00000000..2e877277 --- /dev/null +++ b/src/hooks/anthropic-context-window-limit-recovery/recovery-deduplication.test.ts @@ -0,0 +1,122 @@ +import { describe, test, expect, mock, beforeEach } from "bun:test" +import type { PluginInput } from "@opencode-ai/plugin" +import type { ExperimentalConfig } from "../../config" + +const attemptDeduplicationRecoveryMock = mock(async () => {}) + +mock.module("./deduplication-recovery", () => ({ + attemptDeduplicationRecovery: attemptDeduplicationRecoveryMock, +})) + +function createImmediateTimeouts(): () => void { + const originalSetTimeout = globalThis.setTimeout + const originalClearTimeout = globalThis.clearTimeout + + globalThis.setTimeout = ((callback: (...args: unknown[]) => void, _delay?: number, ...args: unknown[]) => { + callback(...args) + return 0 as unknown as ReturnType + }) as typeof setTimeout + + globalThis.clearTimeout = ((_: ReturnType) => {}) as typeof clearTimeout + + return () => { + globalThis.setTimeout = originalSetTimeout + globalThis.clearTimeout = originalClearTimeout + } +} + +describe("createAnthropicContextWindowLimitRecoveryHook", () => { + beforeEach(() => { + attemptDeduplicationRecoveryMock.mockClear() + }) + + test("calls deduplication recovery when compaction is already in progress", async () => { + //#given + const restoreTimeouts = createImmediateTimeouts() + + const experimental = { + dynamic_context_pruning: { + enabled: true, + strategies: { + deduplication: { enabled: true }, + }, + }, + } satisfies ExperimentalConfig + + let resolveSummarize: (() => void) | null = null + const summarizePromise = new Promise((resolve) => { + resolveSummarize = resolve + }) + + const mockClient = { + session: { + messages: mock(() => Promise.resolve({ data: [] })), + summarize: mock(() => summarizePromise), + revert: mock(() => Promise.resolve()), + prompt_async: mock(() => Promise.resolve()), + }, + tui: { + showToast: mock(() => Promise.resolve()), + }, + } + + try { + const { createAnthropicContextWindowLimitRecoveryHook } = await import("./recovery-hook") + const ctx = { client: mockClient, directory: "/tmp" } as PluginInput + const hook = createAnthropicContextWindowLimitRecoveryHook(ctx, { experimental }) + + // first error triggers compaction (setTimeout runs immediately due to mock) + await hook.event({ + event: { + type: "session.error", + properties: { sessionID: "session-96", error: "prompt is too long" }, + }, + }) + + //#when - second error while compaction is in progress + await hook.event({ + event: { + type: "session.error", + properties: { sessionID: "session-96", error: "prompt is too long" }, + }, + }) + + //#then - deduplication recovery was called for the second error + expect(attemptDeduplicationRecoveryMock).toHaveBeenCalledTimes(1) + expect(attemptDeduplicationRecoveryMock.mock.calls[0]![0]).toBe("session-96") + } finally { + if (resolveSummarize) resolveSummarize() + restoreTimeouts() + } + }) + + test("does not call deduplication when compaction is not in progress", async () => { + //#given + const mockClient = { + session: { + messages: mock(() => Promise.resolve({ data: [] })), + summarize: mock(() => Promise.resolve()), + revert: mock(() => Promise.resolve()), + prompt_async: mock(() => Promise.resolve()), + }, + tui: { + showToast: mock(() => Promise.resolve()), + }, + } + + const { createAnthropicContextWindowLimitRecoveryHook } = await import("./recovery-hook") + const ctx = { client: mockClient, directory: "/tmp" } as PluginInput + const hook = createAnthropicContextWindowLimitRecoveryHook(ctx) + + //#when - single error (no compaction in progress) + await hook.event({ + event: { + type: "session.error", + properties: { sessionID: "session-no-dedup", error: "some other error" }, + }, + }) + + //#then + expect(attemptDeduplicationRecoveryMock).not.toHaveBeenCalled() + }) +}) diff --git a/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.ts b/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.ts new file mode 100644 index 00000000..f4bcb0f2 --- /dev/null +++ b/src/hooks/anthropic-context-window-limit-recovery/recovery-hook.ts @@ -0,0 +1,153 @@ +import type { PluginInput } from "@opencode-ai/plugin" +import type { AutoCompactState, ParsedTokenLimitError } from "./types" +import type { ExperimentalConfig } from "../../config" +import { parseAnthropicTokenLimitError } from "./parser" +import { executeCompact, getLastAssistant } from "./executor" +import { attemptDeduplicationRecovery } from "./deduplication-recovery" +import { log } from "../../shared/logger" + +export interface AnthropicContextWindowLimitRecoveryOptions { + experimental?: ExperimentalConfig +} + +function createRecoveryState(): AutoCompactState { + return { + pendingCompact: new Set(), + errorDataBySession: new Map(), + retryStateBySession: new Map(), + truncateStateBySession: new Map(), + emptyContentAttemptBySession: new Map(), + compactionInProgress: new Set(), + } +} + + +export function createAnthropicContextWindowLimitRecoveryHook( + ctx: PluginInput, + options?: AnthropicContextWindowLimitRecoveryOptions, +) { + const autoCompactState = createRecoveryState() + const experimental = options?.experimental + + const eventHandler = async ({ event }: { event: { type: string; properties?: unknown } }) => { + const props = event.properties as Record | undefined + + if (event.type === "session.deleted") { + const sessionInfo = props?.info as { id?: string } | undefined + if (sessionInfo?.id) { + autoCompactState.pendingCompact.delete(sessionInfo.id) + autoCompactState.errorDataBySession.delete(sessionInfo.id) + autoCompactState.retryStateBySession.delete(sessionInfo.id) + autoCompactState.truncateStateBySession.delete(sessionInfo.id) + autoCompactState.emptyContentAttemptBySession.delete(sessionInfo.id) + autoCompactState.compactionInProgress.delete(sessionInfo.id) + } + return + } + + if (event.type === "session.error") { + const sessionID = props?.sessionID as string | undefined + log("[auto-compact] session.error received", { sessionID, error: props?.error }) + if (!sessionID) return + + const parsed = parseAnthropicTokenLimitError(props?.error) + log("[auto-compact] parsed result", { parsed, hasError: !!props?.error }) + if (parsed) { + autoCompactState.pendingCompact.add(sessionID) + autoCompactState.errorDataBySession.set(sessionID, parsed) + + if (autoCompactState.compactionInProgress.has(sessionID)) { + await attemptDeduplicationRecovery(sessionID, parsed, experimental) + return + } + + const lastAssistant = await getLastAssistant(sessionID, ctx.client, ctx.directory) + const providerID = parsed.providerID ?? (lastAssistant?.providerID as string | undefined) + const modelID = parsed.modelID ?? (lastAssistant?.modelID as string | undefined) + + await ctx.client.tui + .showToast({ + body: { + title: "Context Limit Hit", + message: "Truncating large tool outputs and recovering...", + variant: "warning" as const, + duration: 3000, + }, + }) + .catch(() => {}) + + setTimeout(() => { + executeCompact( + sessionID, + { providerID, modelID }, + autoCompactState, + ctx.client, + ctx.directory, + experimental, + ) + }, 300) + } + return + } + + if (event.type === "message.updated") { + const info = props?.info as Record | undefined + const sessionID = info?.sessionID as string | undefined + + if (sessionID && info?.role === "assistant" && info.error) { + log("[auto-compact] message.updated with error", { sessionID, error: info.error }) + const parsed = parseAnthropicTokenLimitError(info.error) + log("[auto-compact] message.updated parsed result", { parsed }) + if (parsed) { + parsed.providerID = info.providerID as string | undefined + parsed.modelID = info.modelID as string | undefined + autoCompactState.pendingCompact.add(sessionID) + autoCompactState.errorDataBySession.set(sessionID, parsed) + } + } + return + } + + if (event.type === "session.idle") { + const sessionID = props?.sessionID as string | undefined + if (!sessionID) return + + if (!autoCompactState.pendingCompact.has(sessionID)) return + + const errorData = autoCompactState.errorDataBySession.get(sessionID) + const lastAssistant = await getLastAssistant(sessionID, ctx.client, ctx.directory) + + if (lastAssistant?.summary === true) { + autoCompactState.pendingCompact.delete(sessionID) + return + } + + const providerID = errorData?.providerID ?? (lastAssistant?.providerID as string | undefined) + const modelID = errorData?.modelID ?? (lastAssistant?.modelID as string | undefined) + + await ctx.client.tui + .showToast({ + body: { + title: "Auto Compact", + message: "Token limit exceeded. Attempting recovery...", + variant: "warning" as const, + duration: 3000, + }, + }) + .catch(() => {}) + + await executeCompact( + sessionID, + { providerID, modelID }, + autoCompactState, + ctx.client, + ctx.directory, + experimental, + ) + } + } + + return { + event: eventHandler, + } +}