diff --git a/src/agents/pi-embedded-runner/tool-result-char-estimator.test.ts b/src/agents/pi-embedded-runner/tool-result-char-estimator.test.ts new file mode 100644 index 00000000000..6e0ab2adad3 --- /dev/null +++ b/src/agents/pi-embedded-runner/tool-result-char-estimator.test.ts @@ -0,0 +1,82 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; +import { describe, expect, it } from "vitest"; +import { castAgentMessage } from "../test-helpers/agent-message-fixtures.js"; +import { + createMessageCharEstimateCache, + estimateContextChars, + estimateMessageCharsCached, +} from "./tool-result-char-estimator.js"; + +function makeUser(text: string): AgentMessage { + return castAgentMessage({ + role: "user", + content: text, + timestamp: Date.now(), + }); +} + +describe("estimateMessageCharsCached", () => { + it("returns a positive estimate for a valid message", () => { + const cache = createMessageCharEstimateCache(); + const msg = makeUser("hello world"); + expect(estimateMessageCharsCached(msg, cache)).toBeGreaterThan(0); + }); + + it("returns 0 for null", () => { + const cache = createMessageCharEstimateCache(); + expect(estimateMessageCharsCached(null as unknown as AgentMessage, cache)).toBe(0); + }); + + it("returns 0 for undefined", () => { + const cache = createMessageCharEstimateCache(); + expect(estimateMessageCharsCached(undefined as unknown as AgentMessage, cache)).toBe(0); + }); + + it("returns 0 for a non-object primitive", () => { + const cache = createMessageCharEstimateCache(); + expect(estimateMessageCharsCached(42 as unknown as AgentMessage, cache)).toBe(0); + }); + + it("caches the estimate on second call", () => { + const cache = createMessageCharEstimateCache(); + const msg = makeUser("cached test"); + const first = estimateMessageCharsCached(msg, cache); + const second = estimateMessageCharsCached(msg, cache); + expect(first).toBe(second); + expect(first).toBeGreaterThan(0); + }); +}); + +describe("estimateContextChars", () => { + it("sums estimates for valid messages", () => { + const cache = createMessageCharEstimateCache(); + const messages = [makeUser("one"), makeUser("two")]; + const total = estimateContextChars(messages, cache); + expect(total).toBeGreaterThan(0); + }); + + it("skips null entries without crashing", () => { + const cache = createMessageCharEstimateCache(); + const messages = [makeUser("valid"), null as unknown as AgentMessage, makeUser("also valid")]; + const total = estimateContextChars(messages, cache); + expect(total).toBeGreaterThan(0); + }); + + it("skips undefined entries without crashing", () => { + const cache = createMessageCharEstimateCache(); + const messages = [undefined as unknown as AgentMessage, makeUser("valid")]; + const total = estimateContextChars(messages, cache); + expect(total).toBeGreaterThan(0); + }); + + it("handles an entirely null/undefined array", () => { + const cache = createMessageCharEstimateCache(); + const messages = [null as unknown as AgentMessage, undefined as unknown as AgentMessage]; + expect(estimateContextChars(messages, cache)).toBe(0); + }); + + it("handles an empty array", () => { + const cache = createMessageCharEstimateCache(); + expect(estimateContextChars([], cache)).toBe(0); + }); +}); diff --git a/src/agents/pi-embedded-runner/tool-result-char-estimator.ts b/src/agents/pi-embedded-runner/tool-result-char-estimator.ts index 6d022d62289..e6e886116f3 100644 --- a/src/agents/pi-embedded-runner/tool-result-char-estimator.ts +++ b/src/agents/pi-embedded-runner/tool-result-char-estimator.ts @@ -141,6 +141,9 @@ export function estimateMessageCharsCached( msg: AgentMessage, cache: MessageCharEstimateCache, ): number { + if (msg == null || typeof msg !== "object") { + return 0; + } const hit = cache.get(msg); if (hit !== undefined) { return hit; @@ -154,7 +157,9 @@ export function estimateContextChars( messages: AgentMessage[], cache: MessageCharEstimateCache, ): number { - return messages.reduce((sum, msg) => sum + estimateMessageCharsCached(msg, cache), 0); + return messages + .filter((m): m is AgentMessage => m != null) + .reduce((sum, msg) => sum + estimateMessageCharsCached(msg, cache), 0); } export function invalidateMessageCharsCacheEntry(