fix(runtime): classify run usage status

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,

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'
+        `);
+      },
+    },
+  ];
+}

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'

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

src/main/database/snapshot.ts

@@ -91,7 +91,7 @@ export function buildSnapshotMethods(
         ...companyArgs,
       ),
       runs: readAll<RunRecord>(
-        `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<CostRecord>(

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);

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";
+}

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,

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(),

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,

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;

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,
     });

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"

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;