Files
AgentCompany/tests/integration/database.spec.ts
T
ZacharyZhang-NY 14dc9a4897 refactor(cost): single writer for cost_entries, finishRun reads SUM
Cost was written through three independent paths: worker-runner shipped
`finalUsage.costUsd` in run-finished -> finishRun wrote `runs.cost_usd`,
`tasks.cost_usd += costUsd`, `agents.spent_monthly_usd += costUsd`, and
inserted a per-run aggregate row in cost_entries. CostWriter
additionally inserted per-message rows in cost_entries from main. Any
CostWriter failure created silent drift between `SUM(cost_entries)`
(read by budget enforcement) and `runs.cost_usd` (read by the UI).

Make CostWriter the single writer, finishRun the single reader.

- Add `db.recordCostEntry(...)` as the first-class insert path for
  per-message cost rows. CostWriter calls it; tests call it directly.
- Remove the aggregate INSERT from finishRun. finishRun now reads
  `SELECT COALESCE(SUM(amount_usd), 0) FROM cost_entries WHERE run_id=?
  AND unattributed=0` and writes that as the authoritative run total to
  runs.cost_usd, tasks.cost_usd, and agents.spent_monthly_usd.
- Drop `costUsd` from `finishRun` input, the worker `run-finished`
  ChildMessage shape, the automation `RunFinishedEvent` type, and every
  caller + test fixture. Cost is owned by the per-message path.
- worker-runner stops extracting `usage.cost.total` — token counts stay
  for session accumulation, cost comes from pi-ai via CostWriter.
- run-finished-helpers now gates budget incident checks on
  `event.status === "succeeded"` instead of a cost field that no longer
  exists on the event.
- automation-run-completion fires `budget_threshold` on every succeeded
  run (post-run spent_monthly_usd already reflects the SUM path).
- Update budget tests to write a `recordCostEntry` before `finishRun`
  to simulate the per-message cost path; drop `costUsd` from cross-
  subsystem, edge-case-budget, database, and related fixtures.

One writer, one reader, one source of truth. No drift.
2026-04-17 15:04:22 -04:00

260 lines
7.0 KiB
TypeScript

import { mkdtemp, rm } from "node:fs/promises";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { afterEach, beforeAll, describe, expect, it, vi } from "vitest";
beforeAll(() => {
process.env.AGENTCOMPANY_ALLOW_INSECURE_VAULT = "true";
});
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";
const tempDirs: string[] = [];
async function createTempDir(prefix: string) {
const dir = await mkdtemp(join(tmpdir(), prefix));
tempDirs.push(dir);
return dir;
}
afterEach(async () => {
await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
});
describe("AppDatabase", () => {
it("persists core entities, secrets, run logs, and backup restore state", async () => {
const profileDir = await createTempDir("agentcompany-profile-");
const backupDir = await createTempDir("agentcompany-backup-");
const db = new AppDatabase(profileDir);
db.init();
const companyId = db.saveCompany({
name: "Orbit Foundry",
description: "Test company",
status: "active",
});
const workspaceId = db.saveWorkspace({
companyId,
name: "Primary",
localPath: profileDir,
repoUrl: "https://example.com/repo.git",
repoRef: "main",
isPrimary: true,
});
const agentId = db.saveAgent({
companyId,
name: "Helios",
role: "CEO",
title: "Lead agent",
status: "active",
reportsTo: null,
connectorId: "openai-codex",
workspaceId,
model: null,
capabilities: "Operate the company",
budgetMonthlyUsd: 200,
});
const goalId = db.saveGoal({
companyId,
title: "Ship the desktop app",
description: "Reach packaged desktop parity.",
ownerAgentId: agentId,
status: "active",
});
const projectId = db.saveProject({
companyId,
goalId,
name: "Desktop rebuild",
description: "Rebuild as Electron desktop app.",
leadAgentId: agentId,
status: "in_progress",
});
const taskId = db.saveTask({
companyId,
projectId,
goalId,
title: "Validate connector state",
description: "Run a health check and capture output.",
assigneeAgentId: agentId,
workspaceId,
priority: "high",
status: "todo",
});
db.saveSecret({
companyId,
name: "vault-openai",
description: "Test key",
value: "sk-test-123",
});
db.saveConnector({
id: "openai-codex",
command: "codex",
model: null,
envBindingText: "OPENAI_API_KEY=vault-openai",
notes: "Connector note",
});
expect(db.resolveSecretEnv("OPENAI_API_KEY=vault-openai", companyId)).toEqual({
OPENAI_API_KEY: "sk-test-123",
});
const run = db.createRun({
companyId,
taskId,
agentId,
workspaceId,
connectorId: "openai-codex",
});
await db.appendRunLog(run.id, "stdout", "hello from run\n");
expect(db.readRunLog(run.id)).toContain("hello from run");
db.markRunStatus(run.id, "running", "Connector launched");
db.finishRun({
runId: run.id,
status: "succeeded",
summary: "Completed successfully",
errorMessage: null,
exitCode: 0,
signal: null,
model: "gpt-5-codex",
sessionDisplayId: "session-1",
unattributedCost: false,
});
const snapshotBeforeBackup = db.listSnapshot();
expect(snapshotBeforeBackup.companies).toHaveLength(1);
expect(snapshotBeforeBackup.runs[0]?.status).toBe("succeeded");
expect(snapshotBeforeBackup.tasks[0]?.activeRunId).toBeNull();
await db.exportBackup(backupDir);
db.saveCompany({
name: "Temporary company",
description: "Should disappear after restore",
status: "active",
});
expect(db.listSnapshot().companies).toHaveLength(2);
await db.restoreFromBackup(backupDir);
const restoredSnapshot = db.listSnapshot();
expect(restoredSnapshot.companies).toHaveLength(1);
expect(restoredSnapshot.companies[0]?.name).toBe("Orbit Foundry");
expect(restoredSnapshot.runs[0]?.summary).toBe("Completed successfully");
db.close();
});
it("returns a renderer snapshot scoped to the current company while preserving global selectors", async () => {
const profileDir = await createTempDir("agentcompany-renderer-snapshot-");
const db = new AppDatabase(profileDir);
db.init();
const alphaCompanyId = db.saveCompany({
name: "Alpha",
description: "Alpha company",
status: "active",
});
const betaCompanyId = db.saveCompany({
name: "Beta",
description: "Beta company",
status: "active",
});
const alphaWorkspaceId = db.saveWorkspace({
companyId: alphaCompanyId,
name: "Alpha Workspace",
localPath: profileDir,
repoUrl: "",
repoRef: "",
isPrimary: true,
});
const betaWorkspaceId = db.saveWorkspace({
companyId: betaCompanyId,
name: "Beta Workspace",
localPath: profileDir,
repoUrl: "",
repoRef: "",
isPrimary: true,
});
const alphaAgentId = db.saveAgent({
companyId: alphaCompanyId,
name: "Alpha CEO",
role: "CEO",
title: "",
status: "active",
reportsTo: null,
connectorId: "openai-codex",
workspaceId: alphaWorkspaceId,
model: null,
capabilities: "Lead alpha",
budgetMonthlyUsd: 100,
});
db.saveAgent({
companyId: betaCompanyId,
name: "Beta CEO",
role: "CEO",
title: "",
status: "active",
reportsTo: null,
connectorId: "openai-codex",
workspaceId: betaWorkspaceId,
model: null,
capabilities: "Lead beta",
budgetMonthlyUsd: 100,
});
db.saveTask({
companyId: alphaCompanyId,
title: "Alpha task",
description: "Alpha work",
assigneeAgentId: alphaAgentId,
workspaceId: alphaWorkspaceId,
priority: "medium",
status: "todo",
});
db.saveTask({
companyId: betaCompanyId,
title: "Beta task",
description: "Beta work",
assigneeAgentId: null,
workspaceId: betaWorkspaceId,
priority: "medium",
status: "todo",
});
db.setCurrentCompany(alphaCompanyId);
const rendererSnapshot = db.listRendererSnapshot();
expect(rendererSnapshot.currentCompanyId).toBe(alphaCompanyId);
expect(rendererSnapshot.companies).toHaveLength(2);
expect(rendererSnapshot.connectors.length).toBeGreaterThan(0);
expect(rendererSnapshot.workspaces.map((workspace) => workspace.companyId)).toEqual([alphaCompanyId]);
expect(rendererSnapshot.agents.every((agent) => agent.companyId === alphaCompanyId)).toBe(true);
expect(rendererSnapshot.tasks.every((task) => task.companyId === alphaCompanyId)).toBe(true);
db.close();
});
});