diff --git a/src/agents/pi-extensions/context-pruning.test.ts b/src/agents/pi-extensions/context-pruning.test.ts index b316ef87f..9e92e320e 100644 --- a/src/agents/pi-extensions/context-pruning.test.ts +++ b/src/agents/pi-extensions/context-pruning.test.ts @@ -2,7 +2,7 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent"; import { describe, expect, it } from "vitest"; -import { setContextPruningRuntime } from "./context-pruning/runtime.js"; +import { getContextPruningRuntime, setContextPruningRuntime } from "./context-pruning/runtime.js"; import { computeEffectiveSettings, @@ -306,6 +306,72 @@ describe("context-pruning", () => { expect(toolText(findToolResult(result.messages, "t1"))).toBe("[cleared]"); }); + it("cache-ttl prunes once and resets the ttl window", () => { + const sessionManager = {}; + const lastTouch = Date.now() - DEFAULT_CONTEXT_PRUNING_SETTINGS.ttlMs - 1000; + + setContextPruningRuntime(sessionManager, { + settings: { + ...DEFAULT_CONTEXT_PRUNING_SETTINGS, + keepLastAssistants: 0, + softTrimRatio: 0, + hardClearRatio: 0, + minPrunableToolChars: 0, + hardClear: { enabled: true, placeholder: "[cleared]" }, + softTrim: { maxChars: 10, headChars: 3, tailChars: 3 }, + }, + contextWindowTokens: 1000, + isToolPrunable: () => true, + lastCacheTouchAt: lastTouch, + }); + + const messages: AgentMessage[] = [ + makeUser("u1"), + makeAssistant("a1"), + makeToolResult({ + toolCallId: "t1", + toolName: "exec", + text: "x".repeat(20_000), + }), + ]; + + let handler: + | (( + event: { messages: AgentMessage[] }, + ctx: ExtensionContext, + ) => { messages: AgentMessage[] } | undefined) + | undefined; + + const api = { + on: (name: string, fn: unknown) => { + if (name === "context") { + handler = fn as typeof handler; + } + }, + appendEntry: (_type: string, _data?: unknown) => {}, + } as unknown as ExtensionAPI; + + contextPruningExtension(api); + if (!handler) throw new Error("missing context handler"); + + const first = handler({ messages }, { + model: undefined, + sessionManager, + } as unknown as ExtensionContext); + if (!first) throw new Error("expected first prune"); + expect(toolText(findToolResult(first.messages, "t1"))).toBe("[cleared]"); + + const runtime = getContextPruningRuntime(sessionManager); + if (!runtime?.lastCacheTouchAt) throw new Error("expected lastCacheTouchAt"); + expect(runtime.lastCacheTouchAt).toBeGreaterThan(lastTouch); + + const second = handler({ messages }, { + model: undefined, + sessionManager, + } as unknown as ExtensionContext); + expect(second).toBeUndefined(); + }); + it("respects tools allow/deny (deny wins; wildcards supported)", () => { const messages: AgentMessage[] = [ makeUser("u1"), diff --git a/src/agents/pi-extensions/context-pruning/extension.ts b/src/agents/pi-extensions/context-pruning/extension.ts index 7e48141c4..411cc9a44 100644 --- a/src/agents/pi-extensions/context-pruning/extension.ts +++ b/src/agents/pi-extensions/context-pruning/extension.ts @@ -29,6 +29,11 @@ export default function contextPruningExtension(api: ExtensionAPI): void { }); if (next === event.messages) return undefined; + + if (runtime.settings.mode === "cache-ttl") { + runtime.lastCacheTouchAt = Date.now(); + } + return { messages: next }; }); }