diff --git a/src/hooks/context-window-monitor.test.ts b/src/hooks/context-window-monitor.test.ts index e2252fd0..d0f8de3f 100644 --- a/src/hooks/context-window-monitor.test.ts +++ b/src/hooks/context-window-monitor.test.ts @@ -1,6 +1,28 @@ -import { describe, it, expect, mock, beforeEach } from "bun:test" +/// + +import { describe, it, expect, mock, beforeEach, afterEach } from "bun:test" import { createContextWindowMonitorHook } from "./context-window-monitor" +const ANTHROPIC_CONTEXT_ENV_KEY = "ANTHROPIC_1M_CONTEXT" +const VERTEX_CONTEXT_ENV_KEY = "VERTEX_ANTHROPIC_1M_CONTEXT" + +const originalAnthropicContextEnv = process.env[ANTHROPIC_CONTEXT_ENV_KEY] +const originalVertexContextEnv = process.env[VERTEX_CONTEXT_ENV_KEY] + +function resetContextLimitEnv(): void { + if (originalAnthropicContextEnv === undefined) { + delete process.env[ANTHROPIC_CONTEXT_ENV_KEY] + } else { + process.env[ANTHROPIC_CONTEXT_ENV_KEY] = originalAnthropicContextEnv + } + + if (originalVertexContextEnv === undefined) { + delete process.env[VERTEX_CONTEXT_ENV_KEY] + } else { + process.env[VERTEX_CONTEXT_ENV_KEY] = originalVertexContextEnv + } +} + function createMockCtx() { return { client: { @@ -17,6 +39,12 @@ describe("context-window-monitor", () => { beforeEach(() => { ctx = createMockCtx() + delete process.env[ANTHROPIC_CONTEXT_ENV_KEY] + delete process.env[VERTEX_CONTEXT_ENV_KEY] + }) + + afterEach(() => { + resetContextLimitEnv() }) // #given event caches token info from message.updated @@ -218,4 +246,77 @@ describe("context-window-monitor", () => { ) expect(output.output).toBe("test") }) + + it("should use 1M limit when model cache flag is enabled", async () => { + //#given + const hook = createContextWindowMonitorHook(ctx as never, true) + const sessionID = "ses_1m_flag" + + await hook.event({ + event: { + type: "message.updated", + properties: { + info: { + role: "assistant", + sessionID, + providerID: "anthropic", + finish: true, + tokens: { + input: 300000, + output: 1000, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + }, + }, + }, + }) + + //#when + const output = { title: "", output: "original", metadata: null } + await hook["tool.execute.after"]( + { tool: "bash", sessionID, callID: "call_1" }, + output + ) + + //#then + expect(output.output).toBe("original") + }) + + it("should keep env var fallback when model cache flag is disabled", async () => { + //#given + process.env[ANTHROPIC_CONTEXT_ENV_KEY] = "true" + const hook = createContextWindowMonitorHook(ctx as never, false) + const sessionID = "ses_env_fallback" + + await hook.event({ + event: { + type: "message.updated", + properties: { + info: { + role: "assistant", + sessionID, + providerID: "anthropic", + finish: true, + tokens: { + input: 300000, + output: 1000, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + }, + }, + }, + }) + + //#when + const output = { title: "", output: "original", metadata: null } + await hook["tool.execute.after"]( + { tool: "bash", sessionID, callID: "call_1" }, + output + ) + + //#then + expect(output.output).toBe("original") + }) }) diff --git a/src/hooks/preemptive-compaction.test.ts b/src/hooks/preemptive-compaction.test.ts index c6c38fa5..41e704e9 100644 --- a/src/hooks/preemptive-compaction.test.ts +++ b/src/hooks/preemptive-compaction.test.ts @@ -1,4 +1,26 @@ -import { describe, it, expect, mock, beforeEach } from "bun:test" +/// + +import { describe, it, expect, mock, beforeEach, afterEach } from "bun:test" + +const ANTHROPIC_CONTEXT_ENV_KEY = "ANTHROPIC_1M_CONTEXT" +const VERTEX_CONTEXT_ENV_KEY = "VERTEX_ANTHROPIC_1M_CONTEXT" + +const originalAnthropicContextEnv = process.env[ANTHROPIC_CONTEXT_ENV_KEY] +const originalVertexContextEnv = process.env[VERTEX_CONTEXT_ENV_KEY] + +function resetContextLimitEnv(): void { + if (originalAnthropicContextEnv === undefined) { + delete process.env[ANTHROPIC_CONTEXT_ENV_KEY] + } else { + process.env[ANTHROPIC_CONTEXT_ENV_KEY] = originalAnthropicContextEnv + } + + if (originalVertexContextEnv === undefined) { + delete process.env[VERTEX_CONTEXT_ENV_KEY] + } else { + process.env[VERTEX_CONTEXT_ENV_KEY] = originalVertexContextEnv + } +} const logMock = mock(() => {}) @@ -29,6 +51,12 @@ describe("preemptive-compaction", () => { beforeEach(() => { ctx = createMockCtx() logMock.mockClear() + delete process.env[ANTHROPIC_CONTEXT_ENV_KEY] + delete process.env[VERTEX_CONTEXT_ENV_KEY] + }) + + afterEach(() => { + resetContextLimitEnv() }) // #given event caches token info from message.updated @@ -238,4 +266,77 @@ describe("preemptive-compaction", () => { error: String(summarizeError), }) }) + + it("should use 1M limit when model cache flag is enabled", async () => { + //#given + const hook = createPreemptiveCompactionHook(ctx as never, true) + const sessionID = "ses_1m_flag" + + await hook.event({ + event: { + type: "message.updated", + properties: { + info: { + role: "assistant", + sessionID, + providerID: "anthropic", + modelID: "claude-sonnet-4-5", + finish: true, + tokens: { + input: 300000, + output: 1000, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + }, + }, + }, + }) + + //#when + await hook["tool.execute.after"]( + { tool: "bash", sessionID, callID: "call_1" }, + { title: "", output: "test", metadata: null } + ) + + //#then + expect(ctx.client.session.summarize).not.toHaveBeenCalled() + }) + + it("should keep env var fallback when model cache flag is disabled", async () => { + //#given + process.env[ANTHROPIC_CONTEXT_ENV_KEY] = "true" + const hook = createPreemptiveCompactionHook(ctx as never, false) + const sessionID = "ses_env_fallback" + + await hook.event({ + event: { + type: "message.updated", + properties: { + info: { + role: "assistant", + sessionID, + providerID: "anthropic", + modelID: "claude-sonnet-4-5", + finish: true, + tokens: { + input: 300000, + output: 1000, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + }, + }, + }, + }) + + //#when + await hook["tool.execute.after"]( + { tool: "bash", sessionID, callID: "call_1" }, + { title: "", output: "test", metadata: null } + ) + + //#then + expect(ctx.client.session.summarize).not.toHaveBeenCalled() + }) }) diff --git a/src/shared/dynamic-truncator.test.ts b/src/shared/dynamic-truncator.test.ts new file mode 100644 index 00000000..91105bc7 --- /dev/null +++ b/src/shared/dynamic-truncator.test.ts @@ -0,0 +1,96 @@ +/// + +import { describe, expect, it, afterEach } from "bun:test" + +import { getContextWindowUsage } from "./dynamic-truncator" + +const ANTHROPIC_CONTEXT_ENV_KEY = "ANTHROPIC_1M_CONTEXT" +const VERTEX_CONTEXT_ENV_KEY = "VERTEX_ANTHROPIC_1M_CONTEXT" + +const originalAnthropicContextEnv = process.env[ANTHROPIC_CONTEXT_ENV_KEY] +const originalVertexContextEnv = process.env[VERTEX_CONTEXT_ENV_KEY] + +function resetContextLimitEnv(): void { + if (originalAnthropicContextEnv === undefined) { + delete process.env[ANTHROPIC_CONTEXT_ENV_KEY] + } else { + process.env[ANTHROPIC_CONTEXT_ENV_KEY] = originalAnthropicContextEnv + } + + if (originalVertexContextEnv === undefined) { + delete process.env[VERTEX_CONTEXT_ENV_KEY] + } else { + process.env[VERTEX_CONTEXT_ENV_KEY] = originalVertexContextEnv + } +} + +function createContextUsageMockContext(inputTokens: number) { + return { + client: { + session: { + messages: async () => ({ + data: [ + { + info: { + role: "assistant", + tokens: { + input: inputTokens, + output: 0, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + }, + }, + ], + }), + }, + }, + } +} + +describe("getContextWindowUsage", () => { + afterEach(() => { + resetContextLimitEnv() + }) + + it("uses 1M limit when model cache flag is enabled", async () => { + //#given + delete process.env[ANTHROPIC_CONTEXT_ENV_KEY] + delete process.env[VERTEX_CONTEXT_ENV_KEY] + const ctx = createContextUsageMockContext(300000) + + //#when + const usage = await getContextWindowUsage(ctx as never, "ses_1m_flag", true) + + //#then + expect(usage?.usagePercentage).toBe(0.3) + expect(usage?.remainingTokens).toBe(700000) + }) + + it("uses 200K limit when model cache flag is disabled and env vars are unset", async () => { + //#given + delete process.env[ANTHROPIC_CONTEXT_ENV_KEY] + delete process.env[VERTEX_CONTEXT_ENV_KEY] + const ctx = createContextUsageMockContext(150000) + + //#when + const usage = await getContextWindowUsage(ctx as never, "ses_default", false) + + //#then + expect(usage?.usagePercentage).toBe(0.75) + expect(usage?.remainingTokens).toBe(50000) + }) + + it("keeps env var fallback when model cache flag is disabled", async () => { + //#given + process.env[ANTHROPIC_CONTEXT_ENV_KEY] = "true" + const ctx = createContextUsageMockContext(300000) + + //#when + const usage = await getContextWindowUsage(ctx as never, "ses_env_fallback", false) + + //#then + expect(usage?.usagePercentage).toBe(0.3) + expect(usage?.remainingTokens).toBe(700000) + }) +})