diff --git a/tests/unit/edge-case-budget.test.ts b/tests/unit/edge-case-budget.test.ts new file mode 100644 index 00000000..e9a73ac8 --- /dev/null +++ b/tests/unit/edge-case-budget.test.ts @@ -0,0 +1,457 @@ +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_local", + 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_local", + 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_local", + }); + 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_local", + }); + db.checkoutTask(taskId2, run0.id); + db.markRunStatus(run0.id, "running", ""); + // Complete Run-0 with $6 cost (exceeds $5 budget) + db.finishRun({ + runId: run0.id, + status: "succeeded", + summary: "Expensive run", + errorMessage: null, + exitCode: 0, + signal: null, + model: "gemini-pro", + sessionDisplayId: null, + costUsd: 6.0, + 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.finishRun({ + runId: run1.id, + status: "succeeded", + summary: "In-flight run completed", + errorMessage: null, + exitCode: 0, + signal: null, + model: "gemini-pro", + sessionDisplayId: null, + costUsd: 0.50, + 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_local", + }); + db.checkoutTask(taskId, run.id); + db.markRunStatus(run.id, "running", ""); + db.finishRun({ + runId: run.id, + status: "succeeded", + summary: "Done", + errorMessage: null, + exitCode: 0, + signal: null, + model: "gemini-pro", + sessionDisplayId: null, + costUsd: 6.0, + 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_local", + }); + db.checkoutTask(taskId, run.id); + db.markRunStatus(run.id, "running", ""); + db.finishRun({ + runId: run.id, + status: "succeeded", + summary: "Done", + errorMessage: null, + exitCode: 0, + signal: null, + model: "gemini-pro", + sessionDisplayId: null, + costUsd: 6.0, + 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_local", + }); + db.checkoutTask(taskId, run.id); + db.markRunStatus(run.id, "running", ""); + db.finishRun({ + runId: run.id, + status: "succeeded", + summary: "Done", + errorMessage: null, + exitCode: 0, + signal: null, + model: "gemini-pro", + sessionDisplayId: null, + costUsd: 1.5, + 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_local", + }); + db.checkoutTask(taskId, run.id); + db.markRunStatus(run.id, "running", ""); + db.finishRun({ + runId: run.id, + status: "succeeded", + summary: "Done", + errorMessage: null, + exitCode: 0, + signal: null, + model: "gemini-pro", + sessionDisplayId: null, + costUsd: 4.0, + 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_local", + }); + db.checkoutTask(taskId, run.id); + db.markRunStatus(run.id, "running", ""); + + // Complete with cost that exceeds budget -- cost entry is still recorded + db.finishRun({ + runId: run.id, + status: "succeeded", + summary: "Done", + errorMessage: null, + exitCode: 0, + signal: null, + model: "gemini-pro", + sessionDisplayId: null, + costUsd: 5.0, + 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); + }); +});