fix(runtime): classify run usage status
This commit is contained in:
1 parent
d971ca7929
commit
eefb515aa2
17 files changed
+308
-26
No files matched your search
@@ -16,6 +16,7 @@ import type { CredentialResolver } from "@main/credentials";
|
|||||||
import { getConnectorDefinition } from "@main/connectors";
|
import { getConnectorDefinition } from "@main/connectors";
|
||||||
import { projectInboxMessageToPi } from "./inbox-projection";
|
import { projectInboxMessageToPi } from "./inbox-projection";
|
||||||
import { projectAgentEventToLogChunk, extractTextBody } from "./agent-event-projection";
|
import { projectAgentEventToLogChunk, extractTextBody } from "./agent-event-projection";
|
||||||
|
import { hasAttributableAssistantUsage } from "../run-usage-status";
|
||||||
import { captureBaselineIfRequired, reconcileWorkspaceArtifacts } from "./artifact-projector";
|
import { captureBaselineIfRequired, reconcileWorkspaceArtifacts } from "./artifact-projector";
|
||||||
import {
|
import {
|
||||||
handleBeforeToolCall,
|
handleBeforeToolCall,
|
||||||
@@ -357,9 +358,16 @@ export class RuntimeBridge {
|
|||||||
if (envelope.event.type === "message_end") {
|
if (envelope.event.type === "message_end") {
|
||||||
const ctx = this.activeRuns.get(envelope.runId);
|
const ctx = this.activeRuns.get(envelope.runId);
|
||||||
const messageRaw = envelope.event.message as
|
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;
|
| undefined;
|
||||||
if (ctx && messageRaw?.role === "assistant" && messageRaw.usage) {
|
if (ctx && hasAttributableAssistantUsage(messageRaw)) {
|
||||||
try {
|
try {
|
||||||
this.costWriter.observeAssistantMessage({
|
this.costWriter.observeAssistantMessage({
|
||||||
runId: ctx.runId,
|
runId: ctx.runId,
|
||||||
|
|||||||
@@ -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'
|
||||||
|
`);
|
||||||
|
},
|
||||||
|
},
|
||||||
|
];
|
||||||
|
}
|
||||||
@@ -2,12 +2,13 @@ import { randomUUID } from "node:crypto";
|
|||||||
import { writeFileSync } from "node:fs";
|
import { writeFileSync } from "node:fs";
|
||||||
import { join } from "node:path";
|
import { join } from "node:path";
|
||||||
import type { DatabaseSync, SQLInputValue } from "node:sqlite";
|
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 { type ReadAll, type Transaction, nowIso, toBoolean } from "./helpers";
|
||||||
import { buildRunLogMethods } from "./run-logs";
|
import { buildRunLogMethods } from "./run-logs";
|
||||||
import { buildRunEventMethods } from "./run-events";
|
import { buildRunEventMethods } from "./run-events";
|
||||||
import { buildCostEntryMethods } from "./cost-entries";
|
import { buildCostEntryMethods } from "./cost-entries";
|
||||||
import { checkoutTaskInTransaction } from "./task-checkout";
|
import { checkoutTaskInTransaction } from "./task-checkout";
|
||||||
|
import { deriveLegacyRunUsageStatus } from "../run-usage-status";
|
||||||
|
|
||||||
class TaskCheckoutConflict extends Error {
|
class TaskCheckoutConflict extends Error {
|
||||||
constructor(readonly taskId: string) {
|
constructor(readonly taskId: string) {
|
||||||
@@ -58,7 +59,7 @@ export function buildRunMethods(
|
|||||||
getRun(id: string): RunRecord {
|
getRun(id: string): RunRecord {
|
||||||
const row = getDb()
|
const row = getDb()
|
||||||
.prepare(
|
.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;
|
.get(id) as RunRecord | undefined;
|
||||||
if (!row) throw new Error(`Run not found: ${id}`);
|
if (!row) throw new Error(`Run not found: ${id}`);
|
||||||
@@ -150,6 +151,7 @@ export function buildRunMethods(
|
|||||||
signal: string | null;
|
signal: string | null;
|
||||||
model: string | null;
|
model: string | null;
|
||||||
unattributedCost: boolean;
|
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
|
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 }
|
| { taskId: string | null; agentId: string; companyId: string; connectorId: ConnectorId; startedAt: string | null }
|
||||||
@@ -165,6 +167,10 @@ export function buildRunMethods(
|
|||||||
const runNeverStarted =
|
const runNeverStarted =
|
||||||
run.startedAt === null && (input.status === "failed" || input.status === "cancelled");
|
run.startedAt === null && (input.status === "failed" || input.status === "cancelled");
|
||||||
const ts = nowIso();
|
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`.
|
// Cost is owned by CostWriter — per-message rows in `cost_entries`.
|
||||||
// finishRun reads the aggregate instead of writing another row, so
|
// finishRun reads the aggregate instead of writing another row, so
|
||||||
@@ -189,12 +195,12 @@ export function buildRunMethods(
|
|||||||
transaction(() => {
|
transaction(() => {
|
||||||
getDb()
|
getDb()
|
||||||
.prepare(
|
.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(
|
.run(
|
||||||
input.status, input.summary ?? "", input.errorMessage, input.exitCode,
|
input.status, input.summary ?? "", input.errorMessage, input.exitCode,
|
||||||
input.signal, input.model, runCostUsd,
|
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) {
|
if (run.taskId !== null) {
|
||||||
getDb().prepare("update tasks set active_run_id = null, updated_at = ? where id = ?").run(ts, run.taskId);
|
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();
|
const ts = nowIso();
|
||||||
getDb()
|
getDb()
|
||||||
.prepare(
|
.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) {
|
if (taskId !== null) {
|
||||||
getDb().prepare("update tasks set active_run_id = null, updated_at = ? where id = ?").run(ts, taskId);
|
getDb().prepare("update tasks set active_run_id = null, updated_at = ? where id = ?").run(ts, taskId);
|
||||||
getDb()
|
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,
|
`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,
|
connector_id as connectorId, status, summary, error_message as errorMessage, exit_code as exitCode,
|
||||||
signal, model, cost_usd as costUsd,
|
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
|
log_path as logPath, created_at as createdAt, updated_at as updatedAt
|
||||||
from runs
|
from runs
|
||||||
where status = 'queued'
|
where status = 'queued'
|
||||||
|
|||||||
@@ -19,6 +19,7 @@ import { createMigrationsV50V50 } from "./migrations/v50-v50";
|
|||||||
import { createMigrationsV51V51 } from "./migrations/v51-v51";
|
import { createMigrationsV51V51 } from "./migrations/v51-v51";
|
||||||
import { createMigrationsV52V52 } from "./migrations/v52-v52";
|
import { createMigrationsV52V52 } from "./migrations/v52-v52";
|
||||||
import { createMigrationsV53V53 } from "./migrations/v53-v53";
|
import { createMigrationsV53V53 } from "./migrations/v53-v53";
|
||||||
|
import { createMigrationsV54V54 } from "./migrations/v54-v54";
|
||||||
import {
|
import {
|
||||||
columnExists,
|
columnExists,
|
||||||
nowIso,
|
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", "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", "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("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");
|
addColumnIfMissing("cost_entries", "project_id", "alter table cost_entries add column project_id text");
|
||||||
db.exec(`
|
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);
|
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 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 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 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 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 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));
|
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),
|
...createMigrationsV51V51(deps),
|
||||||
...createMigrationsV52V52(deps),
|
...createMigrationsV52V52(deps),
|
||||||
...createMigrationsV53V53(deps),
|
...createMigrationsV53V53(deps),
|
||||||
|
...createMigrationsV54V54(deps),
|
||||||
];
|
];
|
||||||
|
|
||||||
// Pre-migration snapshot. If any migration is pending we take a full
|
// Pre-migration snapshot. If any migration is pending we take a full
|
||||||
|
|||||||
@@ -91,7 +91,7 @@ export function buildSnapshotMethods(
|
|||||||
...companyArgs,
|
...companyArgs,
|
||||||
),
|
),
|
||||||
runs: readAll<RunRecord>(
|
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,
|
...companyArgs,
|
||||||
).map((record) => ({ ...record, unattributedCost: Boolean(record.unattributedCost) })) as RunRecord[],
|
).map((record) => ({ ...record, unattributedCost: Boolean(record.unattributedCost) })) as RunRecord[],
|
||||||
costs: readAll<CostRecord>(
|
costs: readAll<CostRecord>(
|
||||||
|
|||||||
@@ -156,9 +156,9 @@ export function buildTaskCrudMethods(
|
|||||||
const ts = nowIso();
|
const ts = nowIso();
|
||||||
getDb()
|
getDb()
|
||||||
.prepare(
|
.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()
|
getDb()
|
||||||
.prepare("update tasks set active_run_id = null, updated_at = ? where id = ?")
|
.prepare("update tasks set active_run_id = null, updated_at = ? where id = ?")
|
||||||
.run(ts, run.taskId);
|
.run(ts, run.taskId);
|
||||||
|
|||||||
@@ -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";
|
||||||
|
}
|
||||||
@@ -107,6 +107,7 @@ function cancelRun(runId: string) {
|
|||||||
exitCode: null,
|
exitCode: null,
|
||||||
signal: null,
|
signal: null,
|
||||||
model: null,
|
model: null,
|
||||||
|
usageStatus: "aborted_before_usage",
|
||||||
unattributedCost: true,
|
unattributedCost: true,
|
||||||
inputTokens: null,
|
inputTokens: null,
|
||||||
outputTokens: null,
|
outputTokens: null,
|
||||||
|
|||||||
@@ -1,12 +1,13 @@
|
|||||||
import { z } from "zod";
|
import { z } from "zod";
|
||||||
import type { AppDatabase } from "../../database";
|
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 { NotificationOptions } from "../../notification-service";
|
||||||
import type { createBudgetEnforcementService } from "../../budget-enforcement-service";
|
import type { createBudgetEnforcementService } from "../../budget-enforcement-service";
|
||||||
import type { createSessionManagementService } from "../../session-management-service";
|
import type { createSessionManagementService } from "../../session-management-service";
|
||||||
import type { RuntimeBridge, RuntimeMessageInjector } from "../../agent-runtime";
|
import type { RuntimeBridge, RuntimeMessageInjector } from "../../agent-runtime";
|
||||||
|
|
||||||
const runStatusSchema = z.enum(["queued", "running", "succeeded", "failed", "cancelled", "timed_out", "interrupted"]);
|
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.
|
* Single source of truth for worker → main RPC method names.
|
||||||
@@ -51,6 +52,7 @@ export const workerMessageSchema = z.union([
|
|||||||
exitCode: z.number().int().nullable(),
|
exitCode: z.number().int().nullable(),
|
||||||
signal: z.string().nullable(),
|
signal: z.string().nullable(),
|
||||||
model: z.string().nullable(),
|
model: z.string().nullable(),
|
||||||
|
usageStatus: runUsageStatusSchema,
|
||||||
unattributedCost: z.boolean(),
|
unattributedCost: z.boolean(),
|
||||||
inputTokens: z.number().nullable(),
|
inputTokens: z.number().nullable(),
|
||||||
outputTokens: z.number().nullable(),
|
outputTokens: z.number().nullable(),
|
||||||
|
|||||||
@@ -47,6 +47,7 @@ export function finalizeRun(
|
|||||||
exitCode: event.exitCode,
|
exitCode: event.exitCode,
|
||||||
signal: event.signal,
|
signal: event.signal,
|
||||||
model: event.model,
|
model: event.model,
|
||||||
|
usageStatus: event.usageStatus,
|
||||||
unattributedCost: event.unattributedCost,
|
unattributedCost: event.unattributedCost,
|
||||||
inputTokens: event.inputTokens,
|
inputTokens: event.inputTokens,
|
||||||
outputTokens: event.outputTokens,
|
outputTokens: event.outputTokens,
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
import type { AgentMessage } from "@mariozechner/pi-agent-core";
|
import type { AgentMessage } from "@mariozechner/pi-agent-core";
|
||||||
import type { RunStatus } from "@shared/types";
|
import type { RunStatus, RunUsageStatus } from "@shared/types";
|
||||||
|
|
||||||
export interface RunJob {
|
export interface RunJob {
|
||||||
runId: string;
|
runId: string;
|
||||||
@@ -64,6 +64,7 @@ export type ChildMessage =
|
|||||||
exitCode: number | null;
|
exitCode: number | null;
|
||||||
signal: string | null;
|
signal: string | null;
|
||||||
model: string | null;
|
model: string | null;
|
||||||
|
usageStatus: RunUsageStatus;
|
||||||
unattributedCost: boolean;
|
unattributedCost: boolean;
|
||||||
inputTokens: number | null;
|
inputTokens: number | null;
|
||||||
outputTokens: number | null;
|
outputTokens: number | null;
|
||||||
|
|||||||
@@ -8,6 +8,12 @@ import { createPiAgent, resolveThinkingLevel } from "./pi-agent-factory";
|
|||||||
import type { WorkerRpcClient } from "../entry/rpc";
|
import type { WorkerRpcClient } from "../entry/rpc";
|
||||||
import { createHireAgentTool, HIRE_AGENT_TOOL_NAME } from "./internal-tools/hire-agent-tool";
|
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 { createFileApprovalTool, FILE_APPROVAL_TOOL_NAME } from "./internal-tools/file-approval-tool";
|
||||||
|
import {
|
||||||
|
classifyRunUsage,
|
||||||
|
hasAttributableAssistantUsage,
|
||||||
|
isTerminalAssistantStopReason,
|
||||||
|
type TerminalAssistantStopReason,
|
||||||
|
} from "../../run-usage-status";
|
||||||
|
|
||||||
interface PrepareRunResult {
|
interface PrepareRunResult {
|
||||||
provider: string;
|
provider: string;
|
||||||
@@ -100,6 +106,8 @@ export class WorkerRunner {
|
|||||||
inputTokens: null,
|
inputTokens: null,
|
||||||
outputTokens: null,
|
outputTokens: null,
|
||||||
};
|
};
|
||||||
|
let sawAttributableUsage = false;
|
||||||
|
let terminalStopReason: TerminalAssistantStopReason | null = null;
|
||||||
let errorMessage: string | null = null;
|
let errorMessage: string | null = null;
|
||||||
|
|
||||||
// Inject the AgentCompany job env (AGENT_COMPANY_API_URL / API_KEY /
|
// Inject the AgentCompany job env (AGENT_COMPANY_API_URL / API_KEY /
|
||||||
@@ -217,12 +225,12 @@ export class WorkerRunner {
|
|||||||
errorMessage?: string;
|
errorMessage?: string;
|
||||||
usage?: { input?: number; output?: number; cost?: { total?: number } };
|
usage?: { input?: number; output?: number; cost?: { total?: number } };
|
||||||
};
|
};
|
||||||
if (
|
if (maybeAssistant.role === "assistant" && isTerminalAssistantStopReason(maybeAssistant.stopReason)) {
|
||||||
maybeAssistant.role === "assistant"
|
terminalStopReason = maybeAssistant.stopReason;
|
||||||
&& (maybeAssistant.stopReason === "error" || maybeAssistant.stopReason === "aborted")
|
|
||||||
) {
|
|
||||||
errorMessage = maybeAssistant.errorMessage ?? `Agent stopped with ${maybeAssistant.stopReason}.`;
|
errorMessage = maybeAssistant.errorMessage ?? `Agent stopped with ${maybeAssistant.stopReason}.`;
|
||||||
} else if (maybeAssistant.role === "assistant" && maybeAssistant.usage) {
|
}
|
||||||
|
if (hasAttributableAssistantUsage(maybeAssistant)) {
|
||||||
|
sawAttributableUsage = true;
|
||||||
finalUsage = {
|
finalUsage = {
|
||||||
inputTokens: typeof maybeAssistant.usage.input === "number" ? maybeAssistant.usage.input : finalUsage.inputTokens,
|
inputTokens: typeof maybeAssistant.usage.input === "number" ? maybeAssistant.usage.input : finalUsage.inputTokens,
|
||||||
outputTokens: typeof maybeAssistant.usage.output === "number" ? maybeAssistant.usage.output : finalUsage.outputTokens,
|
outputTokens: typeof maybeAssistant.usage.output === "number" ? maybeAssistant.usage.output : finalUsage.outputTokens,
|
||||||
@@ -265,6 +273,12 @@ export class WorkerRunner {
|
|||||||
: errorMessage
|
: errorMessage
|
||||||
? "failed"
|
? "failed"
|
||||||
: "succeeded";
|
: "succeeded";
|
||||||
|
const usageStatus = classifyRunUsage({
|
||||||
|
runStatus: status,
|
||||||
|
errorMessage,
|
||||||
|
sawAttributableUsage,
|
||||||
|
terminalStopReason,
|
||||||
|
});
|
||||||
|
|
||||||
post({
|
post({
|
||||||
type: "run-finished",
|
type: "run-finished",
|
||||||
@@ -275,12 +289,8 @@ export class WorkerRunner {
|
|||||||
exitCode: null,
|
exitCode: null,
|
||||||
signal: null,
|
signal: null,
|
||||||
model: resolvedModelId,
|
model: resolvedModelId,
|
||||||
// Worker no longer owns cost attribution — CostWriter (main-side)
|
usageStatus,
|
||||||
// writes per-message rows, finishRun sums them. This flag stays true
|
unattributedCost: usageStatus !== "ok",
|
||||||
// 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,
|
|
||||||
inputTokens: finalUsage.inputTokens,
|
inputTokens: finalUsage.inputTokens,
|
||||||
outputTokens: finalUsage.outputTokens,
|
outputTokens: finalUsage.outputTokens,
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -145,6 +145,15 @@ export type RunStatus =
|
|||||||
| "timed_out"
|
| "timed_out"
|
||||||
| "interrupted";
|
| "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 =
|
export type DepartmentId =
|
||||||
| "executive"
|
| "executive"
|
||||||
| "engineering"
|
| "engineering"
|
||||||
|
|||||||
@@ -14,6 +14,7 @@ import type {
|
|||||||
ApprovalState,
|
ApprovalState,
|
||||||
ApprovalType,
|
ApprovalType,
|
||||||
RunStatus,
|
RunStatus,
|
||||||
|
RunUsageStatus,
|
||||||
SkillSourceType,
|
SkillSourceType,
|
||||||
TrustLevel,
|
TrustLevel,
|
||||||
} from "./enums";
|
} from "./enums";
|
||||||
@@ -201,6 +202,7 @@ export interface RunRecord {
|
|||||||
model: string | null;
|
model: string | null;
|
||||||
costUsd: number | null;
|
costUsd: number | null;
|
||||||
unattributedCost: boolean;
|
unattributedCost: boolean;
|
||||||
|
usageStatus: RunUsageStatus;
|
||||||
startedAt: string | null;
|
startedAt: string | null;
|
||||||
finishedAt: string | null;
|
finishedAt: string | null;
|
||||||
logPath: string | null;
|
logPath: string | null;
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ function makeRun(overrides: Partial<RunRecord> = {}): RunRecord {
|
|||||||
model: null,
|
model: null,
|
||||||
costUsd: null,
|
costUsd: null,
|
||||||
unattributedCost: false,
|
unattributedCost: false,
|
||||||
|
usageStatus: "ok",
|
||||||
startedAt: null,
|
startedAt: null,
|
||||||
finishedAt: null,
|
finishedAt: null,
|
||||||
logPath: null,
|
logPath: null,
|
||||||
|
|||||||
@@ -304,6 +304,60 @@ describe("AppDatabase", () => {
|
|||||||
expect(db.entityCrud.getAgent(agentId).spentMonthlyUsd).toBe(12);
|
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", () => {
|
it("activates pending hires when the corresponding approval is granted", () => {
|
||||||
const companyId = db.company.saveCompany({
|
const companyId = db.company.saveCompany({
|
||||||
name: "Hiring Company",
|
name: "Hiring Company",
|
||||||
|
|||||||
@@ -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");
|
||||||
|
});
|
||||||
|
});
|
||||||
Reference in new issue
Block a user