clanker-agent/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts

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<void>;
      }
    )._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<void>;
        },
        "_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<void>;
      }
    )._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.",
    });
  });
});