AgentCompany/tests/unit/edge-case-budget.test.ts

import { join } from "node:path";
import { tmpdir } from "node:os";
import { mkdtemp, rm } from "node:fs/promises";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";

vi.mock("electron", () => ({
  safeStorage: {
    isEncryptionAvailable: () => false,
    getSelectedStorageBackend: () => "basic_text",
    encryptString: (value: string) => Buffer.from(value, "utf8"),
    decryptString: (value: Buffer) => value.toString("utf8"),
  },
}));

import { AppDatabase } from "../../src/main/database";
import { createBudgetEnforcementService } from "../../src/main/budget-enforcement-service";

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

function scaffold(db: AppDatabase) {
  const companyId = db.saveCompany({
    name: "Budget Edge Case Co",
    description: "Budget exhaustion mid-run edge case tests (D-05.2)",
    status: "active",
  });
  db.saveConnector({
    id: "gemini-cli",
    command: "gemini",
    model: "gemini-pro",
    envBindingText: "",
    notes: "Test connector",
  });
  const workspaceId = db.saveWorkspace({
    companyId,
    name: "Primary",
    localPath: "/tmp/budget-edge-test",
    repoUrl: "",
    repoRef: "",
    isPrimary: true,
  });
  const agentId = db.saveAgent({
    companyId,
    name: "Worker Agent",
    role: "developer",
    title: "Developer",
    status: "active",
    reportsTo: null,
    connectorId: "gemini-cli",
    workspaceId,
    model: null,
    capabilities: "Executes tasks",
    budgetMonthlyUsd: 100,
  });
  const goalId = db.saveGoal({
    companyId,
    title: "Ship feature",
    description: "Build budget edge case test",
    ownerAgentId: agentId,
    status: "active",
  });
  const projectId = db.saveProject({
    companyId,
    goalId,
    name: "Budget Project",
    description: "Testing budget exhaustion",
    leadAgentId: agentId,
    status: "in_progress",
    targetDate: null,
  });
  const taskId = db.saveTask({
    companyId,
    projectId,
    goalId,
    title: "Implement feature",
    description: "Build something to test budget edge cases.",
    assigneeAgentId: agentId,
    workspaceId,
    priority: "high",
    status: "todo",
  });
  return { companyId, workspaceId, agentId, goalId, projectId, taskId };
}

// ---------------------------------------------------------------------------
// Tests — Edge Case: Budget Exhaustion Mid-Run (D-05.2)
// ---------------------------------------------------------------------------

describe("Edge Case: Budget Exhaustion Mid-Run (D-05.2)", () => {
  let tempDir: string;
  let db: AppDatabase;
  const mockLogger = { info: vi.fn(), warn: vi.fn(), error: vi.fn() };
  const mockNotify = vi.fn();
  const mockPublish = vi.fn();

  beforeEach(async () => {
    tempDir = await mkdtemp(join(tmpdir(), "agentcompany-budget-edge-"));
    db = new AppDatabase(tempDir);
    db.init();
    mockLogger.info.mockClear();
    mockLogger.warn.mockClear();
    mockLogger.error.mockClear();
    mockNotify.mockClear();
    mockPublish.mockClear();
  });

  afterEach(async () => {
    db.close();
    await rm(tempDir, { recursive: true, force: true });
  });

  it("in-flight run completes normally when budget is exhausted during execution", () => {
    const { companyId, workspaceId, agentId, projectId, taskId } = scaffold(db);
    const budgetService = createBudgetEnforcementService({
      db,
      logger: mockLogger,
      notify: mockNotify,
      publishDomainChanged: mockPublish,
    });

    // Create company budget policy: 500 cents ($5), hard stop enabled
    db.saveBudgetPolicy({
      companyId,
      scopeType: "company",
      scopeId: companyId,
      amount: 500,
      hardStopEnabled: true,
    });

    // Start Run-1 (the in-flight run)
    const run1 = db.createRun({
      companyId,
      taskId,
      agentId,
      workspaceId,
      connectorId: "gemini-cli",
    });
    db.checkoutTask(taskId, run1.id);
    db.markRunStatus(run1.id, "running", "");

    // Create a second task and run that exhausts budget WHILE Run-1 is in-flight
    const taskId2 = db.saveTask({
      companyId,
      title: "Other Task",
      description: "Another task that uses up budget",
      assigneeAgentId: agentId,
      workspaceId,
      priority: "medium",
      status: "todo",
    });
    const run0 = db.createRun({
      companyId,
      taskId: taskId2,
      agentId,
      workspaceId,
      connectorId: "gemini-cli",
    });
    db.checkoutTask(taskId2, run0.id);
    db.markRunStatus(run0.id, "running", "");
    // Complete Run-0 with $6 cost (exceeds $5 budget)
    db.recordCostEntry({ runId: run0.id, companyId, agentId, taskId: taskId2, connectorId: "gemini-cli", amountUsd: 6.0, unattributed: false });
    db.finishRun({
      runId: run0.id,
      status: "succeeded",
      summary: "Expensive run",
      errorMessage: null,
      exitCode: 0,
      signal: null,
      model: "gemini-pro",
      sessionDisplayId: null,
      unattributedCost: false,
    });

    // Now complete Run-1 (the in-flight run) -- it should succeed normally
    // Per D-04/D-05: in-flight runs are unaffected by budget exhaustion
    db.recordCostEntry({ runId: run1.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 0.5, unattributed: false });
    db.finishRun({
      runId: run1.id,
      status: "succeeded",
      summary: "In-flight run completed",
      errorMessage: null,
      exitCode: 0,
      signal: null,
      model: "gemini-pro",
      sessionDisplayId: null,
      unattributedCost: false,
    });

    // Verify Run-1 completed successfully
    const finishedRun = db.getRun(run1.id);
    expect(finishedRun.status).toBe("succeeded");

    // Verify Run-1's cost was recorded (total spend now $6.50)
    const gate = budgetService.evaluateBudgetGate(companyId, agentId, projectId);
    expect(gate.allowed).toBe(false);
    // Spent should be at least 650 cents ($6.50 = $6 + $0.50)
    expect(gate.blockedBy[0].spent).toBeGreaterThanOrEqual(650);
  });

  it("budget gate blocks next run after exhaustion", () => {
    const { companyId, workspaceId, agentId, projectId, taskId } = scaffold(db);
    const budgetService = createBudgetEnforcementService({
      db,
      logger: mockLogger,
      notify: mockNotify,
      publishDomainChanged: mockPublish,
    });

    // Create company budget policy: 500 cents ($5)
    db.saveBudgetPolicy({
      companyId,
      scopeType: "company",
      scopeId: companyId,
      amount: 500,
      hardStopEnabled: true,
    });

    // Complete a run that exceeds budget
    const run = db.createRun({
      companyId,
      taskId,
      agentId,
      workspaceId,
      connectorId: "gemini-cli",
    });
    db.checkoutTask(taskId, run.id);
    db.markRunStatus(run.id, "running", "");
    db.recordCostEntry({ runId: run.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 6.0, unattributed: false });
    db.finishRun({
      runId: run.id,
      status: "succeeded",
      summary: "Done",
      errorMessage: null,
      exitCode: 0,
      signal: null,
      model: "gemini-pro",
      sessionDisplayId: null,
      unattributedCost: false,
    });

    // Budget gate should now block
    const result = budgetService.evaluateBudgetGate(companyId, agentId, projectId);
    expect(result.allowed).toBe(false);
    expect(result.blockedBy.length).toBeGreaterThanOrEqual(1);
    expect(result.blockedBy[0].spent).toBeGreaterThanOrEqual(500);
    expect(result.blockedBy[0].limit).toBe(500);
  });

  it("budget gate returns detailed blockedBy information", () => {
    const { companyId, workspaceId, agentId, projectId, taskId } = scaffold(db);
    const budgetService = createBudgetEnforcementService({
      db,
      logger: mockLogger,
      notify: mockNotify,
      publishDomainChanged: mockPublish,
    });

    // Create company budget policy with known ID
    const policyId = db.saveBudgetPolicy({
      companyId,
      scopeType: "company",
      scopeId: companyId,
      amount: 500,
      hardStopEnabled: true,
    });

    // Exhaust budget
    const run = db.createRun({
      companyId,
      taskId,
      agentId,
      workspaceId,
      connectorId: "gemini-cli",
    });
    db.checkoutTask(taskId, run.id);
    db.markRunStatus(run.id, "running", "");
    db.recordCostEntry({ runId: run.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 6.0, unattributed: false });
    db.finishRun({
      runId: run.id,
      status: "succeeded",
      summary: "Done",
      errorMessage: null,
      exitCode: 0,
      signal: null,
      model: "gemini-pro",
      sessionDisplayId: null,
      unattributedCost: false,
    });

    const result = budgetService.evaluateBudgetGate(companyId, agentId, projectId);
    expect(result.allowed).toBe(false);
    // Verify blockedBy contains correct details
    expect(result.blockedBy[0].policyId).toBe(policyId);
    expect(result.blockedBy[0].scopeType).toBe("company");
    expect(result.blockedBy[0].scopeId).toBe(companyId);
  });

  it("agent-scope budget exhaustion blocks even when company budget has room", () => {
    const { companyId, workspaceId, agentId, projectId, taskId } = scaffold(db);
    const budgetService = createBudgetEnforcementService({
      db,
      logger: mockLogger,
      notify: mockNotify,
      publishDomainChanged: mockPublish,
    });

    // Company-level: plenty of room (10000 cents = $100)
    db.saveBudgetPolicy({
      companyId,
      scopeType: "company",
      scopeId: companyId,
      amount: 10000,
      hardStopEnabled: true,
    });

    // Agent-level: tight budget (100 cents = $1)
    db.saveBudgetPolicy({
      companyId,
      scopeType: "agent",
      scopeId: agentId,
      amount: 100,
      hardStopEnabled: true,
    });

    // Complete a run costing $1.50 (150 cents > 100 agent limit)
    const run = db.createRun({
      companyId,
      taskId,
      agentId,
      workspaceId,
      connectorId: "gemini-cli",
    });
    db.checkoutTask(taskId, run.id);
    db.markRunStatus(run.id, "running", "");
    db.recordCostEntry({ runId: run.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 1.5, unattributed: false });
    db.finishRun({
      runId: run.id,
      status: "succeeded",
      summary: "Done",
      errorMessage: null,
      exitCode: 0,
      signal: null,
      model: "gemini-pro",
      sessionDisplayId: null,
      unattributedCost: false,
    });

    const result = budgetService.evaluateBudgetGate(companyId, agentId, projectId);
    expect(result.allowed).toBe(false);
    // Should be blocked by agent scope, not company scope
    expect(result.blockedBy.length).toBe(1);
    expect(result.blockedBy[0].scopeType).toBe("agent");
    expect(result.blockedBy[0].scopeId).toBe(agentId);
  });

  it("budget gate transitions from allowed to blocked after spending", () => {
    const { companyId, workspaceId, agentId, projectId, taskId } = scaffold(db);
    const budgetService = createBudgetEnforcementService({
      db,
      logger: mockLogger,
      notify: mockNotify,
      publishDomainChanged: mockPublish,
    });

    // Create policy with 300 cents ($3) budget
    db.saveBudgetPolicy({
      companyId,
      scopeType: "company",
      scopeId: companyId,
      amount: 300,
      hardStopEnabled: true,
    });

    // Budget gate initially allows (nothing spent yet)
    const before = budgetService.evaluateBudgetGate(companyId, agentId, projectId);
    expect(before.allowed).toBe(true);

    // Complete a run costing $4 (400 cents > 300 limit)
    const run = db.createRun({
      companyId,
      taskId,
      agentId,
      workspaceId,
      connectorId: "gemini-cli",
    });
    db.checkoutTask(taskId, run.id);
    db.markRunStatus(run.id, "running", "");
    db.recordCostEntry({ runId: run.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 4.0, unattributed: false });
    db.finishRun({
      runId: run.id,
      status: "succeeded",
      summary: "Done",
      errorMessage: null,
      exitCode: 0,
      signal: null,
      model: "gemini-pro",
      sessionDisplayId: null,
      unattributedCost: false,
    });

    // Budget gate now blocks
    const after = budgetService.evaluateBudgetGate(companyId, agentId, projectId);
    expect(after.allowed).toBe(false);
    expect(after.blockedBy.length).toBeGreaterThanOrEqual(1);
  });

  it("in-flight run cost entry is recorded even when budget would be exceeded", () => {
    const { companyId, workspaceId, agentId, projectId, taskId } = scaffold(db);
    const budgetService = createBudgetEnforcementService({
      db,
      logger: mockLogger,
      notify: mockNotify,
      publishDomainChanged: mockPublish,
    });

    // Create tight budget
    db.saveBudgetPolicy({
      companyId,
      scopeType: "company",
      scopeId: companyId,
      amount: 100,
      hardStopEnabled: true,
    });

    // Run that will exceed budget on completion
    const run = db.createRun({
      companyId,
      taskId,
      agentId,
      workspaceId,
      connectorId: "gemini-cli",
    });
    db.checkoutTask(taskId, run.id);
    db.markRunStatus(run.id, "running", "");

    // Complete with cost that exceeds budget -- cost entry is still recorded
    db.recordCostEntry({ runId: run.id, companyId, agentId, taskId, connectorId: "gemini-cli", amountUsd: 5.0, unattributed: false });
    db.finishRun({
      runId: run.id,
      status: "succeeded",
      summary: "Done",
      errorMessage: null,
      exitCode: 0,
      signal: null,
      model: "gemini-pro",
      sessionDisplayId: null,
      unattributedCost: false,
    });

    // Verify cost was recorded by checking budget gate sees the spend
    const result = budgetService.evaluateBudgetGate(companyId, agentId, projectId);
    expect(result.allowed).toBe(false);
    expect(result.blockedBy[0].spent).toBe(500); // 500 cents = $5
    expect(result.blockedBy[0].limit).toBe(100);
  });
});