diff --git a/src/main/agent-runtime/runtime-bridge.ts b/src/main/agent-runtime/runtime-bridge.ts index 164442a2..a18138e7 100644 --- a/src/main/agent-runtime/runtime-bridge.ts +++ b/src/main/agent-runtime/runtime-bridge.ts @@ -16,6 +16,7 @@ import type { CredentialResolver } from "@main/credentials"; import { getConnectorDefinition } from "@main/connectors"; import { projectInboxMessageToPi } from "./inbox-projection"; import { projectAgentEventToLogChunk, extractTextBody } from "./agent-event-projection"; +import { hasAttributableAssistantUsage } from "../run-usage-status"; import { captureBaselineIfRequired, reconcileWorkspaceArtifacts } from "./artifact-projector"; import { handleBeforeToolCall, @@ -357,9 +358,16 @@ export class RuntimeBridge { if (envelope.event.type === "message_end") { const ctx = this.activeRuns.get(envelope.runId); const messageRaw = envelope.event.message as - | { role?: string; usage?: unknown; content?: unknown; toolCallId?: string; timestamp?: number } + | { + role?: string; + stopReason?: string; + usage?: { input?: unknown; output?: unknown }; + content?: unknown; + toolCallId?: string; + timestamp?: number; + } | undefined; - if (ctx && messageRaw?.role === "assistant" && messageRaw.usage) { + if (ctx && hasAttributableAssistantUsage(messageRaw)) { try { this.costWriter.observeAssistantMessage({ runId: ctx.runId, diff --git a/src/main/database/migrations/v54-v54.ts b/src/main/database/migrations/v54-v54.ts new file mode 100644 index 00000000..ec58877e --- /dev/null +++ b/src/main/database/migrations/v54-v54.ts @@ -0,0 +1,34 @@ +import log from "electron-log/main.js"; +import type { Migration, MigrationDeps } from "./types"; +import { columnExists, tableExists } from "./types"; + +export function createMigrationsV54V54(deps: MigrationDeps): Migration[] { + const { db } = deps; + return [ + { + version: 54, + name: "runs_usage_status", + fn: () => { + if (!tableExists(db, "runs")) { + log.info("[migration v54] runs table missing; nothing to do"); + return; + } + + if (!columnExists(db, "runs", "usage_status")) { + db.exec("ALTER TABLE runs ADD COLUMN usage_status TEXT NOT NULL DEFAULT 'pre_first_message'"); + log.info("[migration v54] runs.usage_status column added"); + } + + db.exec(` + UPDATE runs + SET usage_status = CASE + WHEN unattributed_cost = 0 THEN 'ok' + WHEN status IN ('cancelled', 'interrupted') THEN 'aborted_before_usage' + ELSE 'pre_first_message' + END + WHERE usage_status = 'pre_first_message' + `); + }, + }, + ]; +} diff --git a/src/main/database/runs.ts b/src/main/database/runs.ts index 03aa422e..c1f1077d 100644 --- a/src/main/database/runs.ts +++ b/src/main/database/runs.ts @@ -2,12 +2,13 @@ import { randomUUID } from "node:crypto"; import { writeFileSync } from "node:fs"; import { join } from "node:path"; import type { DatabaseSync, SQLInputValue } from "node:sqlite"; -import type { ConnectorId, RunRecord } from "@shared/types"; +import type { ConnectorId, RunRecord, RunUsageStatus } from "@shared/types"; import { type ReadAll, type Transaction, nowIso, toBoolean } from "./helpers"; import { buildRunLogMethods } from "./run-logs"; import { buildRunEventMethods } from "./run-events"; import { buildCostEntryMethods } from "./cost-entries"; import { checkoutTaskInTransaction } from "./task-checkout"; +import { deriveLegacyRunUsageStatus } from "../run-usage-status"; class TaskCheckoutConflict extends Error { constructor(readonly taskId: string) { @@ -58,7 +59,7 @@ export function buildRunMethods( getRun(id: string): RunRecord { const row = getDb() .prepare( - "select id, company_id as companyId, task_id as taskId, agent_id as agentId, workspace_id as workspaceId, connector_id as connectorId, status, summary, error_message as errorMessage, exit_code as exitCode, signal, model, cost_usd as costUsd, unattributed_cost as unattributedCost, started_at as startedAt, finished_at as finishedAt, log_path as logPath, created_at as createdAt, updated_at as updatedAt from runs where id = ?", + "select id, company_id as companyId, task_id as taskId, agent_id as agentId, workspace_id as workspaceId, connector_id as connectorId, status, summary, error_message as errorMessage, exit_code as exitCode, signal, model, cost_usd as costUsd, unattributed_cost as unattributedCost, usage_status as usageStatus, started_at as startedAt, finished_at as finishedAt, log_path as logPath, created_at as createdAt, updated_at as updatedAt from runs where id = ?", ) .get(id) as RunRecord | undefined; if (!row) throw new Error(`Run not found: ${id}`); @@ -150,6 +151,7 @@ export function buildRunMethods( signal: string | null; model: string | null; unattributedCost: boolean; + usageStatus?: RunUsageStatus; }) { const run = getDb().prepare("select task_id as taskId, agent_id as agentId, company_id as companyId, connector_id as connectorId, started_at as startedAt from runs where id = ?").get(input.runId) as | { taskId: string | null; agentId: string; companyId: string; connectorId: ConnectorId; startedAt: string | null } @@ -165,6 +167,10 @@ export function buildRunMethods( const runNeverStarted = run.startedAt === null && (input.status === "failed" || input.status === "cancelled"); const ts = nowIso(); + const usageStatus = input.usageStatus ?? deriveLegacyRunUsageStatus({ + runStatus: input.status, + unattributedCost: input.unattributedCost, + }); // Cost is owned by CostWriter — per-message rows in `cost_entries`. // finishRun reads the aggregate instead of writing another row, so @@ -189,12 +195,12 @@ export function buildRunMethods( transaction(() => { getDb() .prepare( - "update runs set status = ?, summary = ?, error_message = ?, exit_code = ?, signal = ?, model = ?, cost_usd = ?, unattributed_cost = ?, finished_at = ?, updated_at = ? where id = ?", + "update runs set status = ?, summary = ?, error_message = ?, exit_code = ?, signal = ?, model = ?, cost_usd = ?, unattributed_cost = ?, usage_status = ?, finished_at = ?, updated_at = ? where id = ?", ) .run( input.status, input.summary ?? "", input.errorMessage, input.exitCode, input.signal, input.model, runCostUsd, - input.unattributedCost ? 1 : 0, ts, ts, input.runId, + input.unattributedCost ? 1 : 0, usageStatus, ts, ts, input.runId, ); if (run.taskId !== null) { getDb().prepare("update tasks set active_run_id = null, updated_at = ? where id = ?").run(ts, run.taskId); @@ -262,9 +268,9 @@ export function buildRunMethods( const ts = nowIso(); getDb() .prepare( - "update runs set status = 'cancelled', summary = ?, error_message = null, exit_code = null, signal = null, model = null, cost_usd = null, unattributed_cost = 1, finished_at = ?, updated_at = ? where id = ?", + "update runs set status = 'cancelled', summary = ?, error_message = null, exit_code = null, signal = null, model = null, cost_usd = null, unattributed_cost = 1, usage_status = ?, finished_at = ?, updated_at = ? where id = ?", ) - .run(summary, ts, ts, runId); + .run(summary, "aborted_before_usage", ts, ts, runId); if (taskId !== null) { getDb().prepare("update tasks set active_run_id = null, updated_at = ? where id = ?").run(ts, taskId); getDb() @@ -309,7 +315,7 @@ export function buildRunMethods( `select id, company_id as companyId, task_id as taskId, agent_id as agentId, workspace_id as workspaceId, connector_id as connectorId, status, summary, error_message as errorMessage, exit_code as exitCode, signal, model, cost_usd as costUsd, - unattributed_cost as unattributedCost, started_at as startedAt, finished_at as finishedAt, + unattributed_cost as unattributedCost, usage_status as usageStatus, started_at as startedAt, finished_at as finishedAt, log_path as logPath, created_at as createdAt, updated_at as updatedAt from runs where status = 'queued' diff --git a/src/main/database/schema.ts b/src/main/database/schema.ts index a2991d0e..9e318d15 100644 --- a/src/main/database/schema.ts +++ b/src/main/database/schema.ts @@ -19,6 +19,7 @@ import { createMigrationsV50V50 } from "./migrations/v50-v50"; import { createMigrationsV51V51 } from "./migrations/v51-v51"; import { createMigrationsV52V52 } from "./migrations/v52-v52"; import { createMigrationsV53V53 } from "./migrations/v53-v53"; +import { createMigrationsV54V54 } from "./migrations/v54-v54"; import { columnExists, nowIso, @@ -49,6 +50,7 @@ export function ensureOptionalColumns(db: DatabaseSync) { addColumnIfMissing("tasks", "checkout_run_id", "alter table tasks add column checkout_run_id text"); addColumnIfMissing("tasks", "execution_locked_at", "alter table tasks add column execution_locked_at text"); addColumnIfMissing("tasks", "active_session_id", "alter table tasks add column active_session_id text"); + addColumnIfMissing("runs", "usage_status", "alter table runs add column usage_status text not null default 'pre_first_message'"); addColumnIfMissing("cost_entries", "project_id", "alter table cost_entries add column project_id text"); db.exec(` create table if not exists task_comments (id text primary key, company_id text not null, task_id text not null, author_agent_id text, author_name text not null, body text not null, created_at text not null); @@ -121,7 +123,7 @@ export function initializeDatabase( create table if not exists projects (id text primary key, company_id text not null, goal_id text, name text not null, description text not null, lead_agent_id text, status text not null, target_date text, created_at text not null, updated_at text not null); create table if not exists tasks (id text primary key, company_id text not null, project_id text, goal_id text, parent_id text, title text not null, description text not null, assignee_agent_id text, workspace_id text, status text not null, priority text not null, active_run_id text, session_json text, session_display_id text, cost_usd real not null default 0, created_at text not null, updated_at text not null); create table if not exists approvals (id text primary key, company_id text not null, related_task_id text, requested_by_agent_id text, type text not null, payload_summary text not null, impact_summary text not null, payload_json text, decision_note text not null default '', state text not null, created_at text not null, updated_at text not null); - create table if not exists runs (id text primary key, company_id text not null, task_id text not null, agent_id text not null, workspace_id text, connector_id text not null, status text not null, summary text not null, error_message text, exit_code integer, signal text, model text, session_display_id text, cost_usd real, unattributed_cost integer not null default 0, started_at text, finished_at text, log_path text, created_at text not null, updated_at text not null); + create table if not exists runs (id text primary key, company_id text not null, task_id text not null, agent_id text not null, workspace_id text, connector_id text not null, status text not null, summary text not null, error_message text, exit_code integer, signal text, model text, session_display_id text, cost_usd real, unattributed_cost integer not null default 0, usage_status text not null default 'pre_first_message', started_at text, finished_at text, log_path text, created_at text not null, updated_at text not null); create table if not exists cost_entries (id text primary key, company_id text not null, run_id text not null, task_id text, agent_id text, connector_id text not null, amount_usd real, unattributed integer not null default 0, created_at text not null); create table if not exists activity (id text primary key, company_id text not null, actor text not null, action text not null, entity_type text not null, entity_id text not null, detail text not null, created_at text not null); create table if not exists secrets (id text primary key, company_id text not null, name text not null, description text not null, backend text not null, ciphertext blob not null, updated_at text not null, created_at text not null, unique(company_id, name)); @@ -218,6 +220,7 @@ export function runMigrations( ...createMigrationsV51V51(deps), ...createMigrationsV52V52(deps), ...createMigrationsV53V53(deps), + ...createMigrationsV54V54(deps), ]; // Pre-migration snapshot. If any migration is pending we take a full diff --git a/src/main/database/snapshot.ts b/src/main/database/snapshot.ts index 71412fa5..1b40a4fe 100644 --- a/src/main/database/snapshot.ts +++ b/src/main/database/snapshot.ts @@ -91,7 +91,7 @@ export function buildSnapshotMethods( ...companyArgs, ), runs: readAll( - `select id, company_id as companyId, task_id as taskId, agent_id as agentId, workspace_id as workspaceId, connector_id as connectorId, status, summary, error_message as errorMessage, exit_code as exitCode, signal, model, cost_usd as costUsd, unattributed_cost as unattributedCost, started_at as startedAt, finished_at as finishedAt, log_path as logPath, created_at as createdAt, updated_at as updatedAt from runs${companyWhere} order by created_at desc limit ${companyId ? 2000 : 2000}`, + `select id, company_id as companyId, task_id as taskId, agent_id as agentId, workspace_id as workspaceId, connector_id as connectorId, status, summary, error_message as errorMessage, exit_code as exitCode, signal, model, cost_usd as costUsd, unattributed_cost as unattributedCost, usage_status as usageStatus, started_at as startedAt, finished_at as finishedAt, log_path as logPath, created_at as createdAt, updated_at as updatedAt from runs${companyWhere} order by created_at desc limit ${companyId ? 2000 : 2000}`, ...companyArgs, ).map((record) => ({ ...record, unattributedCost: Boolean(record.unattributedCost) })) as RunRecord[], costs: readAll( diff --git a/src/main/database/tasks-crud.ts b/src/main/database/tasks-crud.ts index 1372b245..c8861b7f 100644 --- a/src/main/database/tasks-crud.ts +++ b/src/main/database/tasks-crud.ts @@ -156,9 +156,9 @@ export function buildTaskCrudMethods( const ts = nowIso(); getDb() .prepare( - "update runs set status = 'cancelled', summary = ?, error_message = null, exit_code = null, signal = null, model = null, session_display_id = null, cost_usd = null, unattributed_cost = 1, finished_at = ?, updated_at = ? where id = ?", + "update runs set status = 'cancelled', summary = ?, error_message = null, exit_code = null, signal = null, model = null, session_display_id = null, cost_usd = null, unattributed_cost = 1, usage_status = ?, finished_at = ?, updated_at = ? where id = ?", ) - .run("Task deleted.", ts, ts, run.id); + .run("Task deleted.", "aborted_before_usage", ts, ts, run.id); getDb() .prepare("update tasks set active_run_id = null, updated_at = ? where id = ?") .run(ts, run.taskId); diff --git a/src/main/run-usage-status.ts b/src/main/run-usage-status.ts new file mode 100644 index 00000000..096707cd --- /dev/null +++ b/src/main/run-usage-status.ts @@ -0,0 +1,66 @@ +import type { RunStatus, RunUsageStatus } from "@shared/types"; + +export type TerminalAssistantStopReason = "error" | "aborted"; + +export interface AssistantUsageProbe { + role?: string; + stopReason?: string; + usage?: { + input?: unknown; + output?: unknown; + }; +} + +export interface AttributableAssistantUsageProbe extends AssistantUsageProbe { + role: "assistant"; + usage: { + input: number; + output: number; + }; +} + +export interface RunUsageClassificationInput { + runStatus: RunStatus; + errorMessage: string | null; + sawAttributableUsage: boolean; + terminalStopReason: TerminalAssistantStopReason | null; +} + +export function isTerminalAssistantStopReason(reason: string | undefined): reason is TerminalAssistantStopReason { + return reason === "error" || reason === "aborted"; +} + +export function hasAttributableAssistantUsage( + message: AssistantUsageProbe | undefined, +): message is AttributableAssistantUsageProbe { + if (message?.role !== "assistant") return false; + if (typeof message.usage?.input !== "number" || typeof message.usage.output !== "number") return false; + if ( + isTerminalAssistantStopReason(message.stopReason) + && message.usage.input === 0 + && message.usage.output === 0 + ) { + return false; + } + return true; +} + +export function classifyRunUsage(input: RunUsageClassificationInput): RunUsageStatus { + if (input.terminalStopReason === "error") return "error_stop"; + if (input.terminalStopReason === "aborted") { + return input.sawAttributableUsage ? "ok" : "aborted_before_usage"; + } + if (input.sawAttributableUsage) return "ok"; + if (input.runStatus === "cancelled" || input.runStatus === "interrupted") return "aborted_before_usage"; + if (input.errorMessage || input.runStatus === "failed" || input.runStatus === "timed_out") return "pre_first_message"; + return "pre_first_message"; +} + +export function deriveLegacyRunUsageStatus(input: { + runStatus: RunStatus; + unattributedCost: boolean; +}): RunUsageStatus { + if (!input.unattributedCost) return "ok"; + if (input.runStatus === "cancelled" || input.runStatus === "interrupted") return "aborted_before_usage"; + return "pre_first_message"; +} diff --git a/src/main/worker/entry/index.ts b/src/main/worker/entry/index.ts index 99432b39..320c6ad3 100644 --- a/src/main/worker/entry/index.ts +++ b/src/main/worker/entry/index.ts @@ -107,6 +107,7 @@ function cancelRun(runId: string) { exitCode: null, signal: null, model: null, + usageStatus: "aborted_before_usage", unattributedCost: true, inputTokens: null, outputTokens: null, diff --git a/src/main/worker/entry/types.ts b/src/main/worker/entry/types.ts index 70237e63..9945ea20 100644 --- a/src/main/worker/entry/types.ts +++ b/src/main/worker/entry/types.ts @@ -1,12 +1,13 @@ import { z } from "zod"; import type { AppDatabase } from "../../database"; -import type { DesktopEvent, RunRecord } from "@shared/types"; +import { RUN_USAGE_STATUSES, type DesktopEvent, type RunRecord } from "@shared/types"; import type { NotificationOptions } from "../../notification-service"; import type { createBudgetEnforcementService } from "../../budget-enforcement-service"; import type { createSessionManagementService } from "../../session-management-service"; import type { RuntimeBridge, RuntimeMessageInjector } from "../../agent-runtime"; const runStatusSchema = z.enum(["queued", "running", "succeeded", "failed", "cancelled", "timed_out", "interrupted"]); +const runUsageStatusSchema = z.enum(RUN_USAGE_STATUSES); /** * Single source of truth for worker → main RPC method names. @@ -51,6 +52,7 @@ export const workerMessageSchema = z.union([ exitCode: z.number().int().nullable(), signal: z.string().nullable(), model: z.string().nullable(), + usageStatus: runUsageStatusSchema, unattributedCost: z.boolean(), inputTokens: z.number().nullable(), outputTokens: z.number().nullable(), diff --git a/src/main/worker/handlers/run-finished-helpers.ts b/src/main/worker/handlers/run-finished-helpers.ts index 68dfde39..5fab1382 100644 --- a/src/main/worker/handlers/run-finished-helpers.ts +++ b/src/main/worker/handlers/run-finished-helpers.ts @@ -47,6 +47,7 @@ export function finalizeRun( exitCode: event.exitCode, signal: event.signal, model: event.model, + usageStatus: event.usageStatus, unattributedCost: event.unattributedCost, inputTokens: event.inputTokens, outputTokens: event.outputTokens, diff --git a/src/main/worker/runtime/helpers.ts b/src/main/worker/runtime/helpers.ts index cfd641bf..13baf684 100644 --- a/src/main/worker/runtime/helpers.ts +++ b/src/main/worker/runtime/helpers.ts @@ -1,5 +1,5 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; -import type { RunStatus } from "@shared/types"; +import type { RunStatus, RunUsageStatus } from "@shared/types"; export interface RunJob { runId: string; @@ -64,6 +64,7 @@ export type ChildMessage = exitCode: number | null; signal: string | null; model: string | null; + usageStatus: RunUsageStatus; unattributedCost: boolean; inputTokens: number | null; outputTokens: number | null; diff --git a/src/main/worker/runtime/index.ts b/src/main/worker/runtime/index.ts index 43283bea..388b0e64 100644 --- a/src/main/worker/runtime/index.ts +++ b/src/main/worker/runtime/index.ts @@ -8,6 +8,12 @@ import { createPiAgent, resolveThinkingLevel } from "./pi-agent-factory"; import type { WorkerRpcClient } from "../entry/rpc"; import { createHireAgentTool, HIRE_AGENT_TOOL_NAME } from "./internal-tools/hire-agent-tool"; import { createFileApprovalTool, FILE_APPROVAL_TOOL_NAME } from "./internal-tools/file-approval-tool"; +import { + classifyRunUsage, + hasAttributableAssistantUsage, + isTerminalAssistantStopReason, + type TerminalAssistantStopReason, +} from "../../run-usage-status"; interface PrepareRunResult { provider: string; @@ -100,6 +106,8 @@ export class WorkerRunner { inputTokens: null, outputTokens: null, }; + let sawAttributableUsage = false; + let terminalStopReason: TerminalAssistantStopReason | null = null; let errorMessage: string | null = null; // Inject the AgentCompany job env (AGENT_COMPANY_API_URL / API_KEY / @@ -217,12 +225,12 @@ export class WorkerRunner { errorMessage?: string; usage?: { input?: number; output?: number; cost?: { total?: number } }; }; - if ( - maybeAssistant.role === "assistant" - && (maybeAssistant.stopReason === "error" || maybeAssistant.stopReason === "aborted") - ) { + if (maybeAssistant.role === "assistant" && isTerminalAssistantStopReason(maybeAssistant.stopReason)) { + terminalStopReason = maybeAssistant.stopReason; errorMessage = maybeAssistant.errorMessage ?? `Agent stopped with ${maybeAssistant.stopReason}.`; - } else if (maybeAssistant.role === "assistant" && maybeAssistant.usage) { + } + if (hasAttributableAssistantUsage(maybeAssistant)) { + sawAttributableUsage = true; finalUsage = { inputTokens: typeof maybeAssistant.usage.input === "number" ? maybeAssistant.usage.input : finalUsage.inputTokens, outputTokens: typeof maybeAssistant.usage.output === "number" ? maybeAssistant.usage.output : finalUsage.outputTokens, @@ -265,6 +273,12 @@ export class WorkerRunner { : errorMessage ? "failed" : "succeeded"; + const usageStatus = classifyRunUsage({ + runStatus: status, + errorMessage, + sawAttributableUsage, + terminalStopReason, + }); post({ type: "run-finished", @@ -275,12 +289,8 @@ export class WorkerRunner { exitCode: null, signal: null, model: resolvedModelId, - // Worker no longer owns cost attribution — CostWriter (main-side) - // writes per-message rows, finishRun sums them. This flag stays true - // only for run paths that never had a chance to emit usage (errors - // before the first message_end); per-message drift detection belongs - // elsewhere. - unattributedCost: finalUsage.inputTokens === null, + usageStatus, + unattributedCost: usageStatus !== "ok", inputTokens: finalUsage.inputTokens, outputTokens: finalUsage.outputTokens, }); diff --git a/src/shared/types/enums.ts b/src/shared/types/enums.ts index 4ab04914..63077e7b 100644 --- a/src/shared/types/enums.ts +++ b/src/shared/types/enums.ts @@ -145,6 +145,15 @@ export type RunStatus = | "timed_out" | "interrupted"; +export const RUN_USAGE_STATUSES = [ + "ok", + "pre_first_message", + "error_stop", + "aborted_before_usage", +] as const; + +export type RunUsageStatus = (typeof RUN_USAGE_STATUSES)[number]; + export type DepartmentId = | "executive" | "engineering" diff --git a/src/shared/types/records-core.ts b/src/shared/types/records-core.ts index f334aa2c..9b6e86f9 100644 --- a/src/shared/types/records-core.ts +++ b/src/shared/types/records-core.ts @@ -14,6 +14,7 @@ import type { ApprovalState, ApprovalType, RunStatus, + RunUsageStatus, SkillSourceType, TrustLevel, } from "./enums"; @@ -201,6 +202,7 @@ export interface RunRecord { model: string | null; costUsd: number | null; unattributedCost: boolean; + usageStatus: RunUsageStatus; startedAt: string | null; finishedAt: string | null; logPath: string | null; diff --git a/tests/unit/dashboard-charts.test.ts b/tests/unit/dashboard-charts.test.ts index 43875db5..28fbb3c7 100644 --- a/tests/unit/dashboard-charts.test.ts +++ b/tests/unit/dashboard-charts.test.ts @@ -29,6 +29,7 @@ function makeRun(overrides: Partial = {}): RunRecord { model: null, costUsd: null, unattributedCost: false, + usageStatus: "ok", startedAt: null, finishedAt: null, logPath: null, diff --git a/tests/unit/database.test.ts b/tests/unit/database.test.ts index f1738d91..2cc3dcf7 100644 --- a/tests/unit/database.test.ts +++ b/tests/unit/database.test.ts @@ -304,6 +304,60 @@ describe("AppDatabase", () => { expect(db.entityCrud.getAgent(agentId).spentMonthlyUsd).toBe(12); }); + it("persists explicit run usage status", () => { + const companyId = db.company.saveCompany({ + name: "Usage Status Company", + description: "Tracks run usage status", + status: "active", + }); + const workspaceId = db.workspacesMod.saveWorkspace({ + companyId, + name: "Usage Workspace", + localPath: "/tmp/usage-status", + repoUrl: "", + repoRef: "", + isPrimary: true, + }); + const agentId = db.agents.saveAgent({ + companyId, + name: "Usage Agent", + role: "Engineer", + title: "Engineer", + status: "active", + reportsTo: null, + connectorId: "openai-codex", + workspaceId, + model: "gpt-5", + capabilities: "Run tasks", + budgetMonthlyUsd: 100, + }); + const taskId = db.tasks.saveTask({ + companyId, + title: "Usage task", + description: "Exercise explicit usage status.", + assigneeAgentId: agentId, + workspaceId, + priority: "medium", + status: "in_progress", + }); + const run = db.runs.createRun({ companyId, taskId, agentId, workspaceId, connectorId: "openai-codex" }); + + db.runs.finishRun({ + runId: run.id, + status: "failed", + summary: null, + errorMessage: "stream setup failed", + exitCode: null, + signal: null, + model: "gpt-5", + unattributedCost: true, + usageStatus: "error_stop", + }); + + expect(db.runs.getRun(run.id).usageStatus).toBe("error_stop"); + expect(db.snap.listSnapshot().runs.find((record) => record.id === run.id)?.usageStatus).toBe("error_stop"); + }); + it("activates pending hires when the corresponding approval is granted", () => { const companyId = db.company.saveCompany({ name: "Hiring Company", diff --git a/tests/unit/run-usage-status.test.ts b/tests/unit/run-usage-status.test.ts new file mode 100644 index 00000000..0210a534 --- /dev/null +++ b/tests/unit/run-usage-status.test.ts @@ -0,0 +1,84 @@ +import { describe, expect, it } from "vitest"; +import { + classifyRunUsage, + deriveLegacyRunUsageStatus, + hasAttributableAssistantUsage, +} from "../../src/main/run-usage-status"; + +describe("run usage status", () => { + it("does not treat pi error EMPTY_USAGE as attributable cost", () => { + expect( + hasAttributableAssistantUsage({ + role: "assistant", + stopReason: "error", + usage: { input: 0, output: 0 }, + }), + ).toBe(false); + }); + + it("accepts normal assistant usage as attributable", () => { + expect( + hasAttributableAssistantUsage({ + role: "assistant", + stopReason: "stop", + usage: { input: 12, output: 7 }, + }), + ).toBe(true); + }); + + it("keeps non-zero terminal usage attributable", () => { + expect( + hasAttributableAssistantUsage({ + role: "assistant", + stopReason: "error", + usage: { input: 12, output: 1 }, + }), + ).toBe(true); + }); + + it("classifies terminal usage states explicitly", () => { + expect( + classifyRunUsage({ + runStatus: "failed", + errorMessage: "stream setup failed", + sawAttributableUsage: false, + terminalStopReason: "error", + }), + ).toBe("error_stop"); + + expect( + classifyRunUsage({ + runStatus: "cancelled", + errorMessage: null, + sawAttributableUsage: false, + terminalStopReason: "aborted", + }), + ).toBe("aborted_before_usage"); + + expect( + classifyRunUsage({ + runStatus: "failed", + errorMessage: "prepareRun failed", + sawAttributableUsage: false, + terminalStopReason: null, + }), + ).toBe("pre_first_message"); + + expect( + classifyRunUsage({ + runStatus: "succeeded", + errorMessage: null, + sawAttributableUsage: true, + terminalStopReason: null, + }), + ).toBe("ok"); + }); + + it("maps legacy finishRun inputs into explicit usage status", () => { + expect(deriveLegacyRunUsageStatus({ runStatus: "succeeded", unattributedCost: false })).toBe("ok"); + expect(deriveLegacyRunUsageStatus({ runStatus: "interrupted", unattributedCost: true })).toBe( + "aborted_before_usage", + ); + expect(deriveLegacyRunUsageStatus({ runStatus: "failed", unattributedCost: true })).toBe("pre_first_message"); + }); +});