fusion/packages/engine/src/__tests__/session-token-usage.test.ts

import { describe, it, expect, vi, beforeEach } from "vitest";
import type { Task, TaskStore } from "@fusion/core";
import { accumulateSessionTokenUsage, captureSessionTokenBaseline, computeCacheHitRatio } from "../session-token-usage.js";
import { enforceTaskTokenBudgetForPersist } from "../token-budget-enforcer.js";
import { TaskExecutor } from "../executor.js";

const { notificationService } = vi.hoisted(() => ({ notificationService: { dispatch: vi.fn() } }));
vi.mock("../notifier.js", () => ({ getActiveNotificationService: () => notificationService }));

interface MockSessionStats {
  tokens?: { input?: number; output?: number; cacheRead?: number; cacheWrite?: number; total?: number };
}

function createSession(
  stats: MockSessionStats | undefined,
  model?: { provider?: string; id?: string },
) {
  return { getSessionStats: vi.fn(() => stats), ...(model ? { model } : {}) } as unknown as Parameters<typeof accumulateSessionTokenUsage>[2];
}

function createStore(initial: Task["tokenUsage"], budget?: { soft?: number; hard?: number }): TaskStore & { _task: Task; updateTask: ReturnType<typeof vi.fn>; pauseTask: ReturnType<typeof vi.fn> } {
  const task = { id: "FN-1", title: "Budget task", tokenUsage: initial } as Task;
  const updateTask = vi.fn(async (_id: string, updates: Partial<Task>) => {
    Object.assign(task, updates);
    return task;
  });
  const pauseTask = vi.fn(async () => task);
  // Model TaskStore's atomic transaction boundary: each updater observes the latest row.
  let atomicQueue: Promise<void> = Promise.resolve();
  const updateTaskAtomic = vi.fn((_id: string, updater: (current: Task) => Partial<Task> | null) => {
    const result = atomicQueue.then(async () => {
      const patch = await updater(task);
      if (patch) Object.assign(task, patch);
      return task;
    });
    atomicQueue = result.then(() => undefined, () => undefined);
    return result;
  });
  const store = {
    _task: task,
    getTask: vi.fn(async () => task),
    getSettingsByScope: vi.fn(async () => ({ project: budget ? { taskTokenBudget: budget } : {}, global: {} })),
    updateTask,
    updateTaskAtomic,
    pauseTask,
  } as unknown as TaskStore & { _task: Task; updateTask: ReturnType<typeof vi.fn>; pauseTask: ReturnType<typeof vi.fn> };
  return store;
}

describe("accumulateSessionTokenUsage", () => {
  beforeEach(() => {
    vi.useRealTimers();
    vi.restoreAllMocks();
    notificationService.dispatch.mockReset();
  });

  it("writes initial token usage and emits cache metrics log", async () => {
    const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
    const store = createStore(undefined);
    const session = createSession({ tokens: { input: 100, output: 30, cacheRead: 5, cacheWrite: 2 } });

    await accumulateSessionTokenUsage(store, "FN-1", session, { agentId: "agent-1", role: "reviewer" });

    expect(store.updateTask).toHaveBeenCalledTimes(1);
    const call = store.updateTask.mock.calls[0]![1] as { tokenUsage: Task["tokenUsage"] };
    expect(call.tokenUsage).toMatchObject({
      inputTokens: 100,
      outputTokens: 30,
      cachedTokens: 5,
      cacheWriteTokens: 2,
      totalTokens: 137,
    });
    const cacheLogCall = errorSpy.mock.calls.find((entry) => String(entry[0]).includes("[token-cache-metrics]"));
    expect(cacheLogCall).toBeTruthy();
    const payload = JSON.parse(String(cacheLogCall?.[0] ?? "").replace(/^.*\[token-cache-metrics\]\s*/, ""));
    expect(payload).toMatchObject({
      taskId: "FN-1",
      agentId: "agent-1",
      role: "reviewer",
      inputTokens: 100,
      cachedTokens: 5,
      cacheWriteTokens: 2,
      hitRatio: computeCacheHitRatio(100, 5),
    });
  });

  it("uses an explicit task-start baseline to exclude resumed-session lifetime tokens", async () => {
    const store = createStore(undefined);
    const session = createSession({ tokens: { input: 1_000, output: 400, cacheRead: 50, cacheWrite: 10 } });

    captureSessionTokenBaseline(session);
    await accumulateSessionTokenUsage(store, "FN-1", session);
    expect(store.updateTask).not.toHaveBeenCalled();

    session.getSessionStats.mockReturnValue({ tokens: { input: 1_025, output: 410, cacheRead: 55, cacheWrite: 12 } });
    await accumulateSessionTokenUsage(store, "FN-1", session);

    expect(store._task.tokenUsage).toMatchObject({
      inputTokens: 25,
      outputTokens: 10,
      cachedTokens: 5,
      cacheWriteTokens: 2,
      totalTokens: 42,
    });
    expect(store._task.tokenUsage?.perModel?.reduce((sum, bucket) => sum + bucket.totalTokens, 0)).toBe(42);
  });

  it("handles undefined task-start stats by retaining a zero snapshot baseline", async () => {
    const store = createStore(undefined);
    const session = createSession(undefined);

    captureSessionTokenBaseline(session);
    await accumulateSessionTokenUsage(store, "FN-1", session);
    expect(store.updateTask).not.toHaveBeenCalled();

    session.getSessionStats.mockReturnValue({ tokens: { input: 4, output: 0, cacheRead: 0, cacheWrite: 0 } });
    await accumulateSessionTokenUsage(store, "FN-1", session);
    expect(store._task.tokenUsage).toMatchObject({ inputTokens: 4, totalTokens: 4 });
  });

  it("re-baselines a reused session between task IDs without changing fresh-session defaults", async () => {
    const taskAStore = createStore(undefined);
    const taskBStore = createStore(undefined);
    const session = createSession({ tokens: { input: 100, output: 40, cacheRead: 0, cacheWrite: 0 } });

    await accumulateSessionTokenUsage(taskAStore, "FN-1", session);
    expect(taskAStore._task.tokenUsage).toMatchObject({ inputTokens: 100, outputTokens: 40, totalTokens: 140 });

    captureSessionTokenBaseline(session);
    await accumulateSessionTokenUsage(taskBStore, "FN-2", session);
    expect(taskBStore.updateTask).not.toHaveBeenCalled();

    session.getSessionStats.mockReturnValue({ tokens: { input: 112, output: 45, cacheRead: 3, cacheWrite: 0 } });
    await accumulateSessionTokenUsage(taskBStore, "FN-2", session);
    expect(taskBStore._task.tokenUsage).toMatchObject({ inputTokens: 12, outputTokens: 5, cachedTokens: 3, totalTokens: 20 });
  });

  it("persists the actually-used session model snapshot with token usage", async () => {
    const store = createStore(undefined);
    const session = createSession(
      { tokens: { input: 20, output: 10, cacheRead: 0, cacheWrite: 0 } },
      { provider: "anthropic", id: "claude-sonnet-4-5" },
    );

    await accumulateSessionTokenUsage(store, "FN-1", session);

    const call = store.updateTask.mock.calls[0]![1] as { tokenUsage: Task["tokenUsage"] };
    expect(call.tokenUsage).toMatchObject({
      modelProvider: "anthropic",
      modelId: "claude-sonnet-4-5",
    });
  });

  it("creates separate per-model buckets for sequential sessions", async () => {
    const store = createStore(undefined);
    const executorSession = createSession(
      { tokens: { input: 70, output: 30, cacheRead: 0, cacheWrite: 0 } },
      { provider: "anthropic", id: "claude-sonnet-4-5" },
    );
    const validatorSession = createSession(
      { tokens: { input: 25, output: 15, cacheRead: 0, cacheWrite: 0 } },
      { provider: "openai", id: "gpt-5" },
    );

    await accumulateSessionTokenUsage(store, "FN-1", executorSession, { role: "executor" });
    await accumulateSessionTokenUsage(store, "FN-1", validatorSession, { role: "reviewer" });

    const usage = store._task.tokenUsage;
    expect(usage).toMatchObject({ inputTokens: 95, outputTokens: 45, totalTokens: 140 });
    expect(usage?.perModel).toEqual([
      expect.objectContaining({ modelProvider: "anthropic", modelId: "claude-sonnet-4-5", inputTokens: 70, outputTokens: 30, totalTokens: 100 }),
      expect.objectContaining({ modelProvider: "openai", modelId: "gpt-5", inputTokens: 25, outputTokens: 15, totalTokens: 40 }),
    ]);
    expect(usage?.perModel?.reduce((sum, bucket) => sum + bucket.totalTokens, 0)).toBe(usage?.totalTokens);
  });

  it("preserves an existing model snapshot when the session has no model", async () => {
    const store = createStore({
      inputTokens: 50,
      outputTokens: 20,
      cachedTokens: 0,
      cacheWriteTokens: 0,
      totalTokens: 70,
      firstUsedAt: "2024-01-01T00:00:00.000Z",
      lastUsedAt: "2024-01-01T00:00:00.000Z",
      modelProvider: "openai",
      modelId: "gpt-5",
    });
    const session = createSession({ tokens: { input: 55, output: 25, cacheRead: 0, cacheWrite: 0 } });

    await accumulateSessionTokenUsage(store, "FN-1", session);

    const call = store.updateTask.mock.calls[0]![1] as { tokenUsage: Task["tokenUsage"] };
    expect(call.tokenUsage).toMatchObject({
      inputTokens: 105,
      outputTokens: 45,
      modelProvider: "openai",
      modelId: "gpt-5",
    });
  });

  it("does nothing when delta is zero (no write, no metrics log)", async () => {
    const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
    const store = createStore({
      inputTokens: 50,
      outputTokens: 20,
      cachedTokens: 0,
      cacheWriteTokens: 0,
      totalTokens: 70,
      firstUsedAt: "2024-01-01T00:00:00.000Z",
      lastUsedAt: "2024-01-01T00:00:00.000Z",
    });
    const session = createSession({ tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 } });

    await accumulateSessionTokenUsage(store, "FN-1", session);
    await accumulateSessionTokenUsage(store, "FN-1", session);

    expect(store.updateTask).not.toHaveBeenCalled();
    expect(errorSpy.mock.calls.find((entry) => String(entry[0]).includes("[token-cache-metrics]"))).toBeUndefined();
  });

  it("emits token-cache-metrics log when executor persists non-zero delta", async () => {
    const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
    const store = createStore(undefined);
    const executor = Object.create(TaskExecutor.prototype) as any;
    executor.store = store;
    executor.tokenUsageBaselines = new Map();
    executor.activeSessions = new Map();
    executor.currentRunContexts = new Map();

    await executor.persistTokenUsage("FN-1", {
      getSessionStats: () => ({ tokens: { input: 3, output: 2, cacheRead: 1, cacheWrite: 0, total: 6 } }),
      model: { provider: "mock", id: "scripted" },
    });

    const cacheLogCall = errorSpy.mock.calls.find((entry) => String(entry[0]).includes("[token-cache-metrics]"));
    expect(cacheLogCall).toBeTruthy();
    const call = store.updateTask.mock.calls[0]![1] as { tokenUsage: Task["tokenUsage"] };
    expect(call.tokenUsage).toMatchObject({ modelProvider: "mock", modelId: "scripted" });
  });

  it("enforces soft and hard budgets through the real persist helper exactly once", async () => {
    const store = createStore(undefined, { soft: 10, hard: 20 });
    const session = createSession({ tokens: { input: 12, output: 0, cacheRead: 0, cacheWrite: 0 } });

    await accumulateSessionTokenUsage(store, "FN-1", session);
    session.getSessionStats.mockReturnValue({ tokens: { input: 25, output: 0, cacheRead: 0, cacheWrite: 0 } });
    await accumulateSessionTokenUsage(store, "FN-1", session);
    session.getSessionStats.mockReturnValue({ tokens: { input: 30, output: 0, cacheRead: 0, cacheWrite: 0 } });
    await accumulateSessionTokenUsage(store, "FN-1", session);

    expect(store._task.tokenBudgetSoftAlertedAt).toBeTruthy();
    expect(store._task.tokenBudgetHardAlertedAt).toBeTruthy();
    expect(store.pauseTask).toHaveBeenCalledOnce();
    expect(store.pauseTask).toHaveBeenCalledWith("FN-1", true, undefined, { pausedReason: "token_budget_exceeded" });
    expect(notificationService.dispatch).toHaveBeenCalledTimes(2);
    expect(notificationService.dispatch).toHaveBeenNthCalledWith(1, "token-budget", expect.objectContaining({ metadata: expect.objectContaining({ kind: "soft" }) }));
    expect(notificationService.dispatch).toHaveBeenNthCalledWith(2, "token-budget", expect.objectContaining({ metadata: expect.objectContaining({ kind: "hard" }) }));
  });

  it("atomically claims concurrent soft and hard enforcement once", async () => {
    const store = createStore({ inputTokens: 25, outputTokens: 0, cacheWriteTokens: 0, totalTokens: 25 }, { soft: 10, hard: 20 });

    await Promise.all(Array.from({ length: 8 }, () => enforceTaskTokenBudgetForPersist(store, "FN-1")));

    expect(store.pauseTask).toHaveBeenCalledOnce();
    expect(notificationService.dispatch).toHaveBeenCalledTimes(2);
    expect(store._task.tokenBudgetSoftAlertedAt).toBeTruthy();
    expect(store._task.tokenBudgetHardAlertedAt).toBeTruthy();
  });

  it("enforces direct executor token persistence through its shared seam", async () => {
    const store = createStore(undefined, { hard: 10 });
    const executor = Object.create(TaskExecutor.prototype) as any;
    executor.store = store;
    executor.currentRunContexts = new Map();

    await executor.persistTaskTokenUsage("FN-1", { inputTokens: 20, outputTokens: 0, cachedTokens: 0, cacheWriteTokens: 0, totalTokens: 20 });

    expect(store.pauseTask).toHaveBeenCalledOnce();
    expect(notificationService.dispatch).toHaveBeenCalledWith("token-budget", expect.objectContaining({ metadata: expect.objectContaining({ kind: "hard" }) }));
  });

  it("retries a hard pause after a failed persisted enforcement attempt", async () => {
    const store = createStore({ inputTokens: 20, outputTokens: 0, cacheWriteTokens: 0, totalTokens: 20 }, { hard: 10 });
    store.pauseTask.mockRejectedValueOnce(new Error("temporary pause failure"));

    await expect(enforceTaskTokenBudgetForPersist(store, "FN-1")).resolves.toBeUndefined();
    expect(store._task.tokenBudgetHardAlertedAt).toBeNull();
    await enforceTaskTokenBudgetForPersist(store, "FN-1");

    expect(store.pauseTask).toHaveBeenCalledTimes(2);
    expect(store._task.tokenBudgetHardAlertedAt).toBeTruthy();
    expect(notificationService.dispatch).toHaveBeenCalledOnce();
  });

  it("keeps a successful hard pause durable when notification dispatch fails", async () => {
    notificationService.dispatch.mockRejectedValueOnce(new Error("notification unavailable"));
    const store = createStore(undefined, { hard: 10 });

    await expect(accumulateSessionTokenUsage(store, "FN-1", createSession({ tokens: { input: 20, output: 0, cacheRead: 0, cacheWrite: 0 } }))).resolves.toBeUndefined();

    expect(store.pauseTask).toHaveBeenCalledOnce();
    expect(store._task.tokenBudgetHardAlertedAt).toBeTruthy();
  });

  it("does not enforce when no budget is configured", async () => {
    const store = createStore(undefined);
    await accumulateSessionTokenUsage(store, "FN-1", createSession({ tokens: { input: 100, output: 0, cacheRead: 0, cacheWrite: 0 } }));
    expect(store.pauseTask).not.toHaveBeenCalled();
    expect(notificationService.dispatch).not.toHaveBeenCalled();
  });

  it("swallows store errors instead of throwing", async () => {
    const store = createStore(undefined);
    (store.updateTask as ReturnType<typeof vi.fn>).mockRejectedValueOnce(new Error("db down"));
    const session = createSession({ tokens: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 } });

    await expect(accumulateSessionTokenUsage(store, "FN-1", session)).resolves.toBeUndefined();
  });

  it.each([
    { input: 10, output: 5, cacheRead: 0, cacheWrite: 0, total: 15 },
    { input: 1000, output: 500, cacheRead: 800, cacheWrite: 200, total: 2500 },
  ])("FN-4389 canonical semantic parity for stats %#", async (tokens) => {
    const heartbeatStore = createStore(undefined);
    const heartbeatSession = createSession({ tokens });
    await accumulateSessionTokenUsage(heartbeatStore, "FN-1", heartbeatSession);
    const heartbeatUsage = (heartbeatStore.updateTask.mock.calls[0]?.[1] as { tokenUsage?: Task["tokenUsage"] })?.tokenUsage;

    const executor = Object.create(TaskExecutor.prototype) as TaskExecutor;
    const extract = (executor as unknown as {
      extractSessionTokenUsage: (session: unknown) => Promise<{ inputTokens: number; outputTokens: number; cachedTokens: number; cacheWriteTokens: number; totalTokens: number } | undefined>;
      accumulateTokenUsage: (existing: Task["tokenUsage"], delta: { inputTokens: number; outputTokens: number; cachedTokens: number; cacheWriteTokens: number; totalTokens: number }) => Task["tokenUsage"];
    });
    const delta = await extract.extractSessionTokenUsage({ getSessionStats: () => ({ tokens }) });
    const executorUsage = delta ? extract.accumulateTokenUsage(undefined, delta) : undefined;

    expect(heartbeatUsage).toMatchObject({
      inputTokens: tokens.input,
      outputTokens: tokens.output,
      cachedTokens: tokens.cacheRead,
      cacheWriteTokens: tokens.cacheWrite,
      totalTokens: tokens.total,
    });
    expect(executorUsage).toMatchObject({
      inputTokens: tokens.input,
      outputTokens: tokens.output,
      cachedTokens: tokens.cacheRead,
      cacheWriteTokens: tokens.cacheWrite,
      totalTokens: tokens.total,
    });
  });
});