Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
9.8 kB · 324 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325import type { AssistantMessage } from "@earendil-works/pi-ai";import { describe, expect, it } from "vitest";import { evaluateActiveCacheTtl, extractSessionCacheState, formatCacheTtl, resolvePromptCacheTtl,} from "../src/core/prompt-cache.js";import { buildCompactAgentsViewLayout, formatSessionCacheTtl } from "../src/modes/agents-view/agents-view-mode.js";import type { AgentsViewRow } from "../src/modes/agents-view/agents-view-state.js";
describe("prompt-cache-state", () => { describe("formatCacheTtl", () => { it("formats seconds, minutes, and hours compactly", () => { expect(formatCacheTtl(0)).toBe("0s"); expect(formatCacheTtl(45)).toBe("45s"); expect(formatCacheTtl(60)).toBe("1m"); expect(formatCacheTtl(125)).toBe("2m 5s"); expect(formatCacheTtl(300)).toBe("5m"); expect(formatCacheTtl(1800)).toBe("30m"); expect(formatCacheTtl(3600)).toBe("1h"); expect(formatCacheTtl(3660)).toBe("1h 1m"); expect(formatCacheTtl(43200)).toBe("12h"); });
it("handles negative or non-finite seconds gracefully", () => { expect(formatCacheTtl(-10)).toBe("0s"); expect(formatCacheTtl(Number.NaN)).toBe("0s"); expect(formatCacheTtl(Number.POSITIVE_INFINITY)).toBe("0s"); }); });
describe("resolvePromptCacheTtl", () => { it("resolves Anthropic and Bedrock Claude defaults and long retentions", () => { expect(resolvePromptCacheTtl("anthropic", "claude-3-7-sonnet-20250219")).toEqual({ ttlSeconds: 300, deterministic: true, }); expect(resolvePromptCacheTtl("anthropic", "claude-3-7-sonnet-20250219", "long")).toEqual({ ttlSeconds: 3600, deterministic: true, }); expect(resolvePromptCacheTtl("amazon-bedrock", "anthropic.claude-3-7-sonnet")).toEqual({ ttlSeconds: 300, deterministic: true, }); });
it("resolves OpenAI modern 30m floor and long 24h retention", () => { expect(resolvePromptCacheTtl("openai", "gpt-5.6")).toEqual({ ttlSeconds: 1800, deterministic: true, }); expect(resolvePromptCacheTtl("openai", "gpt-5.6", "long")).toEqual({ ttlSeconds: 86400, deterministic: true, }); expect(resolvePromptCacheTtl("azure-openai-responses", "gpt-5.6")).toEqual({ ttlSeconds: 1800, deterministic: true, }); });
it("resolves OpenAI Codex 30m sliding TTL regardless of long retention", () => { expect(resolvePromptCacheTtl("openai-codex", "gpt-5.5")).toEqual({ ttlSeconds: 1800, deterministic: true, }); expect(resolvePromptCacheTtl("openai-codex", "gpt-5.5", "long")).toEqual({ ttlSeconds: 1800, deterministic: true, }); expect(resolvePromptCacheTtl("openai-codex", "codex-preview")).toEqual({ ttlSeconds: 1800, deterministic: true, }); });
it("resolves conservative TTL estimates for providers without fixed server TTLs", () => { expect(resolvePromptCacheTtl("deepseek", "deepseek-chat")).toEqual({ ttlSeconds: 43200, deterministic: false, }); expect(resolvePromptCacheTtl("google", "gemini-2.5-flash")).toEqual({ ttlSeconds: 300, deterministic: false, }); expect(resolvePromptCacheTtl("groq", "llama-3.3-70b-versatile")).toEqual({ ttlSeconds: 7200, deterministic: false, }); expect(resolvePromptCacheTtl("zai", "glm-5.2")).toEqual({ ttlSeconds: 1800, deterministic: false, }); expect(resolvePromptCacheTtl("mistral", "mistral-large-latest")).toEqual({ ttlSeconds: 1800, deterministic: false, }); expect(resolvePromptCacheTtl("together", "meta-llama/llama-3.3-70b-instruct-turbo")).toEqual({ ttlSeconds: 1800, deterministic: false, }); expect(resolvePromptCacheTtl("fireworks", "accounts/fireworks/models/deepseek-v3")).toEqual({ ttlSeconds: 900, deterministic: false, }); }); });
describe("extractSessionCacheState", () => { const baseUsage = { input: 100, output: 50, cacheRead: 5000, cacheWrite: 0, totalTokens: 5150, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, };
it("extracts cache state from the latest non-aborted assistant message", () => { const messages = [ { role: "assistant", provider: "anthropic", model: "claude-3-7-sonnet", usage: baseUsage, stopReason: "stop", timestamp: 1000000, } as unknown as AssistantMessage, ];
const state = extractSessionCacheState(messages); expect(state).toBeDefined(); expect(state?.provider).toBe("anthropic"); expect(state?.modelId).toBe("claude-3-7-sonnet"); expect(state?.cachedTokens).toBe(5000); expect(state?.ttlSeconds).toBe(300); });
it("respects compaction boundary", () => { const messages = [ { role: "assistant", provider: "anthropic", model: "claude-3-7-sonnet", usage: baseUsage, stopReason: "stop", timestamp: 1000000, } as unknown as AssistantMessage, ];
const postCompactionLimbo = extractSessionCacheState(messages, 1050000); expect(postCompactionLimbo).toBeUndefined();
const messagesAfterTurn = [ ...messages, { role: "assistant", provider: "anthropic", model: "claude-3-7-sonnet", usage: { ...baseUsage, cacheWrite: 8000, cacheRead: 0 }, stopReason: "stop", timestamp: 1100000, } as unknown as AssistantMessage, ];
const freshState = extractSessionCacheState(messagesAfterTurn, 1050000); expect(freshState).toBeDefined(); expect(freshState?.cachedTokens).toBe(8000); expect(freshState?.timestamp).toBe(1100000); }); });
describe("evaluateActiveCacheTtl and model switches", () => { const warmCache = { provider: "anthropic", modelId: "claude-3-7-sonnet", cachedTokens: 12000, timestamp: 1000000, ttlSeconds: 300, deterministic: true, };
it("returns warm status when active model matches and within TTL", () => { const res = evaluateActiveCacheTtl( { provider: "anthropic", id: "claude-3-7-sonnet" }, warmCache, 1000000 + 60000, ); expect(res.status).toBe("warm"); if (res.status === "warm") { expect(res.remainingSeconds).toBe(240); expect(res.label).toBe("4m"); } });
it("returns mismatch when model is switched to a different model", () => { const res = evaluateActiveCacheTtl({ provider: "openai", id: "gpt-5" }, warmCache, 1000000 + 60000); expect(res.status).toBe("mismatch"); if (res.status === "mismatch") { expect(res.cachedModelId).toBe("claude-3-7-sonnet"); expect(res.canRestore).toBe(true); } });
it("resurrects countdown when user cycles back to original model before TTL expires", () => { const switchAway = evaluateActiveCacheTtl({ provider: "openai", id: "gpt-5" }, warmCache, 1000000 + 100000); expect(switchAway.status).toBe("mismatch");
const switchBack = evaluateActiveCacheTtl( { provider: "anthropic", id: "claude-3-7-sonnet" }, warmCache, 1000000 + 120000, ); expect(switchBack.status).toBe("warm"); if (switchBack.status === "warm") { expect(switchBack.remainingSeconds).toBe(180); expect(switchBack.label).toBe("3m"); } });
it("reports mismatch with canRestore false when previous model cache expired", () => { const res = evaluateActiveCacheTtl( { provider: "openai", id: "gpt-5" }, warmCache, 1000000 + 350000, // past the 300s ttl ); expect(res.status).toBe("mismatch"); if (res.status === "mismatch") { expect(res.canRestore).toBe(false); } });
it("returns cold status when TTL has expired", () => { const res = evaluateActiveCacheTtl( { provider: "anthropic", id: "claude-3-7-sonnet" }, warmCache, 1000000 + 301000, ); expect(res.status).toBe("cold"); }); });
describe("agents view layout with cache TTL", () => { it("includes TTL column in layout when any session has active cache", () => { const now = 1000000; const rows: AgentsViewRow[] = [ { identity: "sess-1", kind: "agent", title: "Session 1", subtitle: "", statusLabel: "", runningSubagentCount: 0, depth: 0, section: "idle", selectable: true, pinned: false, descendantCount: 0, recursiveCost: 0.05, summary: { id: "sess-1", sessionId: "sess-1", cwd: "/tmp", lifecycle: "live", activity: "idle", isSessionActive: false, isStreaming: false, isCompacting: false, attachedClients: 0, messageCount: 5, sessionActions: { queuedCount: 0, steering: [], followUps: [] }, cacheExpiresAt: new Date(now + 240000).toISOString(), cachedTokens: 15000, created: new Date(now - 600000).toISOString(), }, }, { identity: "sess-2", kind: "agent", title: "Session 2", subtitle: "", statusLabel: "", runningSubagentCount: 0, depth: 0, section: "idle", selectable: true, pinned: false, descendantCount: 0, recursiveCost: 0.12, summary: { id: "sess-2", sessionId: "sess-2", cwd: "/tmp", lifecycle: "live", activity: "idle", isSessionActive: false, isStreaming: false, isCompacting: false, attachedClients: 0, messageCount: 2, sessionActions: { queuedCount: 0, steering: [], followUps: [] }, created: new Date(now - 300000).toISOString(), }, }, ];
expect(formatSessionCacheTtl(rows[0].summary, now)).toBe("4m"); expect(formatSessionCacheTtl(rows[1].summary, now)).toBe("");
const layout = buildCompactAgentsViewLayout(rows, 120, now); expect(layout.legend).toContain("Cost"); expect(layout.legend).toContain("TTL"); expect(layout.legend).toContain("Age");
const detail1 = layout.details.get("sess-1"); expect(detail1).toContain("4m");
const detail2 = layout.details.get("sess-2"); expect(detail2).toContain("-"); }); });});