import { existsSync, mkdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { Agent } from "@mariozechner/pi-agent-core"; import { type AssistantMessage, getModel } from "@mariozechner/pi-ai"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { AgentSession } from "../src/core/agent-session.js"; import { AuthStorage } from "../src/core/auth-storage.js"; import { ModelRegistry } from "../src/core/model-registry.js"; import { SessionManager } from "../src/core/session-manager.js"; import { SettingsManager } from "../src/core/settings-manager.js"; import { createTestResourceLoader } from "./utilities.js"; vi.mock("../src/core/compaction/index.js", () => ({ calculateContextTokens: () => 0, collectEntriesForBranchSummary: () => ({ entries: [], commonAncestorId: null }), compact: async () => ({ summary: "compacted", firstKeptEntryId: "entry-1", tokensBefore: 100, details: {}, }), estimateContextTokens: () => ({ tokens: 0, usageTokens: 0, trailingTokens: 0, lastUsageIndex: -1 }), generateBranchSummary: async () => ({ summary: "", aborted: false, readFiles: [], modifiedFiles: [] }), prepareCompaction: () => ({ dummy: true }), shouldCompact: () => false, })); describe("AgentSession auto-compaction queue resume", () => { let session: AgentSession; let tempDir: string; beforeEach(() => { tempDir = join(tmpdir(), `pi-auto-compaction-queue-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); vi.useFakeTimers(); const model = getModel("anthropic", "claude-sonnet-4-5")!; const agent = new Agent({ initialState: { model, systemPrompt: "Test", tools: [], }, }); const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); authStorage.setRuntimeApiKey("anthropic", "test-key"); const modelRegistry = new ModelRegistry(authStorage, tempDir); session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader(), }); }); afterEach(() => { session.dispose(); vi.useRealTimers(); vi.restoreAllMocks(); if (tempDir && existsSync(tempDir)) { rmSync(tempDir, { recursive: true }); } }); it("should resume after threshold compaction when only agent-level queued messages exist", async () => { session.agent.followUp({ role: "custom", customType: "test", content: [{ type: "text", text: "Queued custom" }], display: false, timestamp: Date.now(), }); expect(session.pendingMessageCount).toBe(0); expect(session.agent.hasQueuedMessages()).toBe(true); const continueSpy = vi.spyOn(session.agent, "continue").mockResolvedValue(); const runAutoCompaction = ( session as unknown as { _runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise; } )._runAutoCompaction.bind(session); await runAutoCompaction("threshold", false); await vi.advanceTimersByTimeAsync(100); expect(continueSpy).toHaveBeenCalledTimes(1); }); it("should not compact repeatedly after overflow recovery already attempted", async () => { const model = session.model!; const overflowMessage: AssistantMessage = { role: "assistant", content: [{ type: "text", text: "" }], api: model.api, provider: model.provider, model: model.id, usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "error", errorMessage: "prompt is too long", timestamp: Date.now(), }; const runAutoCompactionSpy = vi .spyOn( session as unknown as { _runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise; }, "_runAutoCompaction", ) .mockResolvedValue(); const events: Array<{ type: string; errorMessage?: string }> = []; session.subscribe((event) => { if (event.type === "auto_compaction_end") { events.push({ type: event.type, errorMessage: event.errorMessage }); } }); const checkCompaction = ( session as unknown as { _checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise; } )._checkCompaction.bind(session); await checkCompaction(overflowMessage); await checkCompaction({ ...overflowMessage, timestamp: Date.now() + 1 }); expect(runAutoCompactionSpy).toHaveBeenCalledTimes(1); expect(events).toContainEqual({ type: "auto_compaction_end", errorMessage: "Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.", }); }); });