diff --git a/docs/architecture.md b/docs/architecture.md index 6c66084..e03aaa9 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -60,7 +60,8 @@ Operator (SPA) approves/denies via | `src/jobs/types.rs` | `TriggerReason`, `ClaudeOutput`. `TriggerReason::event_id(project_id)` builds the `seen_events` dedupe key, prefixed with `project_id` (#159 H1) then discriminated per variant: `Issue` hashes title+description, `ReviewMR` keys on `head_sha` (#152 W6), `FixReview` keys on `review_id` (GitHub's review `html_url`, GitLab's MR `updated_at` — GitLab reviews carry no body, so every approve/unapprove used to hash to the same key, #159 M1), and `MRComment`/`IssueComment` key on the comment's own `url` rather than hashing its text, so an intentional identical repeat isn't deduped as a redelivery (#159 M2); each falls back to its pre-fix hash when the provider omitted the discriminator | | `src/jobs/store.rs` | `TaskStore` core — the run loop (`confirm_task` + branch-conflict guard), `set_pid`/`set_session_id`, hub handle, and the **queue-admission** plumbing (`admit_lock`, `request_admit`, `run_admission_loop`). `confirm_task` claims the task with an atomic **compare-and-swap** (`claim_pending`: `agent_state` cold→pending, one winner) so two racing confirms can't both spawn an agent into one worktree (#142 F4); the spawned run future then **registers its own admission entry as its first action** (over a oneshot-delivered abort handle) and unregisters as its last, so a fast-failing run can't have the parent register a permanently-stale entry that inflates the admission gauge forever (#142 F5). `mark_seen`/`unmark_seen` claim and (on a failed dispatch) release the dedupe key so a redelivery isn't dropped (#142 F9). Lifecycle/CRUD/query helpers live in the sibling modules below (store.rs was split at the 400-line cap) | | `src/jobs/edits.rs` | `TaskStore::update_task` (the operator PATCH, incl. queue/priority edits) + the `TaskEdits` payload — split out of `store.rs` at the cap | -| `src/jobs/results.rs` | `TaskStore` per-run metrics persistence into `task_sessions` (`replace_result`/`save_result`/`save_error_result`) — split out of `store.rs` at the cap | +| `src/jobs/results.rs` | `TaskStore` per-run metrics persistence into `task_sessions` (`replace_result`/`save_result`/`save_error_result`) — also sets the two cache-token columns (`cache_creation_input_tokens`/`cache_read_input_tokens`), reading `ClaudeOutput.usage` — split out of `store.rs` at the cap | +| `src/jobs/usage_backfill.rs` | `backfill_session_usage` — startup sweep (#163) that recovers historical `task_sessions` token counts (pre-fix rows stuck at `NULL`/0) from the raw `result` events already persisted in `events.payload`; modeled on `recover_orphans`/`ProjectDb::sweep_orphans` | | `src/jobs/lifecycle.rs` | `derive_agent_state` (read-time overlay: `is_running` → `running`, `is_warm` → `warm`, else the durable `cold\|pending\|failed`), state-name constants, the `migrate_status` backfill mapping (testable), and the durable-state transition helpers on `TaskStore` — `recover_orphans` (also best-effort SIGKILLs each orphan's leftover agent **process group** by its recorded pid before nulling it, guarding pid reuse via a `/proc//cmdline` sniff, so a previous hard-SIGKILL/OOM doesn't leave live agents in worktrees a recovered task may re-enter — #142 F8), `claim_pending` (the confirm CAS, #142 F4), `set_states`, `finish_task(agent_state, task_state, note)` (clears the recorded pid at a terminal finish so a later pid-reuse can't mistarget the reaper/group-kill — #142 F17), `mark_branch_released` (issue/PR-close → terminal `released`), `note_task_result`, `publish_state` | | `src/jobs/control.rs` | `TaskStore` operator controls — `kill_task` (Pause → durable `cold` / `task_state` `working_on` + "paused by operator" note), `delete_task` (also **releases the task's `seen_events` claim** via `unmark_seen`, recomputed from its stored trigger, before reclaiming the branch worktree via `reclaim_worktree_for_task` — #159 M6, #103), `continue_task` (Resume), `push_message`. Pause/delete additionally **SIGKILL the whole agent process group** by the recorded pid (`kill_task_process_group`), not just the `kill_on_drop` leader, so backgrounded grandchildren die too — and delete's reclaim **skips (keeps) the worktree when a live sibling shares the branch** (`branch_is_live`), since a retry may be mid-run in it (#142 F7/F11). `continue_task` resets a resumed task to `cold` (not `pending`) via a **conditional UPDATE that refuses when `agent_state` is already `pending`** — an unconditional reset let a second concurrent `continue_task`/`push_message` stomp a spawn already claimed by `confirm_task`'s CAS (a still-unregistered #83 window) and double-spawn the agent (#159 H2) — so a bailed `confirm_task` leaves it `cold`, never a stuck `pending` that deadlocks its branch (#142 F10). `push_message`'s cold path now also accepts a never-started (`pending`, no `session_id`) task, since `confirm_task` doesn't need one (#159 M5). Both pause/delete and session-end auto-deny any still-`pending` `auth_requests` for the task (`deny_pending_auth`), so a parked approval that can no longer be answered clears from the UI (#142 F15) | | `src/jobs/pending_message.rs` | `TaskStore::take_pending_message` (atomic read-and-clear-if-unchanged — replaces the old read-then-unconditional-NULL that could drop a comment appended mid-spawn, #159 H2) and `append_pending_message` (persist-then-merge, shared by `push_message`'s cold path and the turn loop's stray-message drain, #159 M3); split out of `control.rs` at the 400-line cap | @@ -144,7 +145,7 @@ Tables (current set, see migration files for canonical schemas): - `queues` — the **task queue** catalog (#39): `id`, `name`, `priority` (smallint, queue-vs-queue ordering, higher = sooner), `created_at`. An operator-orderable backlog bucket; the migration seeds one `Default` queue (stable UUID) so a fresh install can enqueue with no setup. The scheduling logic lives in `jobs/queue.rs` - `tasks` — one row per work item, with **two orthogonal state axes**. `task_state` (persisted, operator-owned lifecycle: `pending|working_on|completed|paused|failed`) is auto-advanced by the runner but freely PATCH-overridable. `paused` marks a warm session whose last turn ended before finishing (a non-`success` `result` subtype like `error_max_turns`) — resumable, and distinct from `completed` so turn-end isn't mistaken for "done" (#106). One further **terminal** value, `released`, is **system-set only** (not PATCH-overridable, not in the operator `TASK_STATES` set): the runner's release path (`mark_branch_released`) moves a task there when its issue/PR closes and the worktree is reclaimed, keeping the row + `task_sessions`/events for history rather than deleting them (#103). `agent_state` is **derived at read time** by `derive_agent_state` (`cold|warm|pending|running|failed`) — its persisted backing column narrows to `cold|pending|failed`, and `warm`/`running` are overlaid from the hub's in-memory liveness flags, never stored. There is no `killed` state — the disposition reason (paused / budget / orphan) is recorded as `task_sessions` text instead. `pending_message` carries a queued follow-up for the resume path. **Normalized:** a task carries only `project_id` (NOT NULL, `ON DELETE CASCADE`) plus `branch` and the working `session_id`; remote URL, default branch, provider and owning service are resolved at run time through `project_id → projects → service` (not duplicated on the row), and the API re-derives them onto `TaskView` so the SPA still sees them. `model_id` (nullable FK → `models`) is the task's selected model (see `service_models` above). **Queue link (#39):** `queue_id` (nullable FK → `queues`, `ON DELETE SET NULL` — deleting a queue un-queues its tasks) + `priority` (smallint, default 0, in-queue sort key); indexed `(queue_id, priority)`. The scheduler only pulls tasks with `queue_id` set; both are operator-editable via PATCH **while pending**. "Filter out completed" is the `task_state = pending` predicate the picker uses - `events` — durable hub-frame stream; one row per frame of any kind (agent event, `auth_request`, `status`) distinguished by the `kind` column, PK `(task_id, seq)`, each frame consuming a unique seq, append-only, batch-inserted (100 at a time) by the live hub, cascade-deleted with the task. **Turns are derived, not stored:** the SPA folds this flat stream into per-turn input/output (`useTurns.ts`) from the `type:"user"` (input, via `--replay-user-messages`) and `type:"result"` (output) markers already present — no `turn_index` column, no turns table, no second writer in the hot path -- `task_sessions` — per-run cost / turns / tokens / result text + the run's claude `session_id` and `created_at`. **1:N with tasks**: turns within one agent run (same `session_id`) accumulate into a single row; a new run (fresh `session_id` — a resume after kill, or a retry) starts a new row, so run history is preserved. `get_task` / the detail view surface the most recent row +- `task_sessions` — per-run cost / turns / tokens / result text + the run's claude `session_id` and `created_at`. **1:N with tasks**: turns within one agent run (same `session_id`) accumulate into a single row; a new run (fresh `session_id` — a resume after kill, or a retry) starts a new row, so run history is preserved. `get_task` / the detail view surface every row (newest first), with the newest also exposed as the single-row `result` summary for the existing UI. `input_tokens`/`output_tokens` plus the two nullable `cache_creation_input_tokens`/`cache_read_input_tokens` columns (#163) come from the `claude` CLI's `result` event `usage.*` object — nested, not top-level (a pre-#163 bug read them from the wrong shape and always got 0). **`NULL` on the two cache columns is the "not yet backfilled" marker**: every write path since the fix sets `Some(...)` (error paths `Some(0)`), and a startup sweep (`src/jobs/usage_backfill.rs::backfill_session_usage`, spawned from `main.rs`) recovers historical rows from the raw `result` events already persisted in `events.payload`, taking the last event per `session_id` (highest `seq` — the `result` event's counters are session-cumulative, not a per-turn delta) and zeroing rows with no matching event. The sweep is idempotent — every row it touches leaves the `NULL` selection, so a re-run (or a redeploy) touches nothing once caught up - `projects` — discovered repos, per-project `allowed_operations` glob list and `env_file` (a `.env`-style minijinja template injected as env vars at agent spawn). Per-trigger-type model overrides live in the sibling `project_models` table (see below), surfaced as `ProjectDetailResponse.models` - `project_branches` — branches the agent has touched, with `issue_iid` / `pr_iid` linkage and status - `auth_requests` — operator-approval items raised by the in-process permission handler @@ -163,9 +164,9 @@ Bearer-auth gates `/api/*` (and the SPA, when `API_BEARER_TOKEN` is set). `/webh | `GET` | `/ws` | **Single app-wide** WebSocket live stream (multiplexes all tasks). In-band auth (first frame = token). Outbound `Envelope` frames (`task_id`/`event`\|`auth_request`\|`status`); inbound `{kind: chat\|redefine\|stop, task_id}` routed to the agent stdin | | `GET` | `/api/tasks` | optional `?task_state=` (SQL `WHERE` on the persisted column) and `?agent_state=` (derived per task, filtered in-memory — acceptable at single-operator scale). Each task emits both `task_state` and the derived `agent_state`; no `live` boolean | | `POST` | `/api/tasks` | operator-driven dispatch: `{ project_id, trigger: TriggerReason }` → pending task (use when the webhook missed/was filtered) | -| `GET` | `/api/tasks/stats` | time spent per `?group_by=project\|service\|branch\|trigger_type` within `?from=`/`?to=` (default last 30d). Optional `?service_id=`/`?project_id=`/`?branch=` narrow the set before aggregation (`service_id` resolved via the project, the other two are `tasks` columns). Running tasks counted as `now - started_at`. | +| `GET` | `/api/tasks/stats` | time spent **and token/cost usage** per `?group_by=project\|service\|branch\|trigger_type` within `?from=`/`?to=` (default last 30d). Optional `?service_id=`/`?project_id=`/`?branch=` narrow the set before aggregation (`service_id` resolved via the project, the other two are `tasks` columns). Running tasks counted as `now - started_at`. Each row and the response also carry `cost_usd`/`input_tokens`/`output_tokens`/`cache_creation_input_tokens`/`cache_read_input_tokens` (`total_*` at the top level), summed across every `task_sessions` row of the tasks in that bucket — attributed to the owning **task's** `created_at` window, not the session's own timestamp, matching time-spent | | `POST` | `/api/tasks/bulk` | apply one action across many tasks: `{ action: run\|pause\|resume\|delete, ids: [..] }`. `run` confirms a pending task or resumes a paused one (picked per task). Reuses the per-task store methods and returns `{ succeeded: [id], failed: [{ id, error }] }` — a bad row doesn't sink the batch. Backs the SPA's multi-select toolbar | -| `GET`/`PATCH`/`DELETE` | `/api/tasks/{id}` | detail + latest session; PATCH sets `task_state` to any of `pending\|working_on\|completed\|paused\|failed` on **any** task (operator override of the auto-advance), and edits the run *inputs* — `branch`, the trigger's `title`/`description` (patched into `trigger_data`, which drives the prompt), plus the queue fields `queue_id` (double-option: omit = unchanged, `null` = dequeue, id = enqueue) and `priority` — only while `task_state == pending` (before the task is tied to a run); the branch may not equal the project's default branch, and `title`/`description` are accepted only for triggers that carry those fields. Enqueuing/reprioritizing fires `try_admit_next`. DELETE force-kills if running, **reclaims the branch worktree** (best-effort, mirroring the close/merge release — #103), then drops the row (and backfills a freed slot from the queue) | +| `GET`/`PATCH`/`DELETE` | `/api/tasks/{id}` | detail: `result` (the newest `task_sessions` row, kept for the existing UI) plus `sessions` (every row for the task, newest first — the full per-run cost/token breakdown, #163); PATCH sets `task_state` to any of `pending\|working_on\|completed\|paused\|failed` on **any** task (operator override of the auto-advance), and edits the run *inputs* — `branch`, the trigger's `title`/`description` (patched into `trigger_data`, which drives the prompt), plus the queue fields `queue_id` (double-option: omit = unchanged, `null` = dequeue, id = enqueue) and `priority` — only while `task_state == pending` (before the task is tied to a run); the branch may not equal the project's default branch, and `title`/`description` are accepted only for triggers that carry those fields. Enqueuing/reprioritizing fires `try_admit_next`. DELETE force-kills if running, **reclaims the branch worktree** (best-effort, mirroring the close/merge release — #103), then drops the row (and backfills a freed slot from the queue) | | `POST` | `/api/tasks/{id}/confirm` | pending → queued (durable `agent_state` pending) | | `POST` | `/api/tasks/{id}/retry` | clone the task as a new row | | `POST` | `/api/tasks/{id}/kill` | SIGKILL; preserves session_id for Resume | diff --git a/frontend/src/components/TaskResultPanel.test.ts b/frontend/src/components/TaskResultPanel.test.ts new file mode 100644 index 0000000..46dc1a0 --- /dev/null +++ b/frontend/src/components/TaskResultPanel.test.ts @@ -0,0 +1,66 @@ +import { describe, it, expect } from "vitest"; +import { mount } from "@vue/test-utils"; +import TaskResultPanel from "./TaskResultPanel.vue"; +import type { TaskResult } from "../types/api"; + +function makeResult(over: Partial = {}): TaskResult { + return { + id: "r1", + task_id: "t1", + cost_usd: 0.1234, + input_tokens: 100, + output_tokens: 200, + cache_creation_input_tokens: 50, + cache_read_input_tokens: 75, + num_turns: 3, + is_error: false, + result_text: "done", + session_id: "abcd1234efgh5678", + created_at: "2026-01-01T00:00:00Z", + ...over, + }; +} + +describe("TaskResultPanel", () => { + it("renders nothing when result is null", () => { + const w = mount(TaskResultPanel, { props: { result: null, sessions: [] } }); + expect(w.text()).toBe(""); + }); + + it("renders the summary for a single session", () => { + const result = makeResult(); + const w = mount(TaskResultPanel, { props: { result, sessions: [result] } }); + expect(w.text()).toContain("0.1234"); + expect(w.text()).toContain("3"); + expect(w.text()).toContain("100"); + expect(w.text()).toContain("200"); + }); + + it("renders — for null cache token fields instead of 0", () => { + const result = makeResult({ cache_creation_input_tokens: null, cache_read_input_tokens: null }); + const w = mount(TaskResultPanel, { props: { result, sessions: [result] } }); + expect(w.text()).toContain("—"); + expect(w.text()).not.toMatch(/Cache read.*\b0\b/); + }); + + it("renders the per-session table when there is more than one session", () => { + const s1 = makeResult({ id: "r1", session_id: "session-one-aaa", cost_usd: 0.1 }); + const s2 = makeResult({ id: "r2", session_id: "session-two-bbb", cost_usd: 0.2 }); + const w = mount(TaskResultPanel, { props: { result: s1, sessions: [s1, s2] } }); + expect(w.find("table").exists()).toBe(true); + expect(w.text()).toContain("0.1000"); + expect(w.text()).toContain("0.2000"); + }); + + it("does not render the per-session table for a single session", () => { + const result = makeResult(); + const w = mount(TaskResultPanel, { props: { result, sessions: [result] } }); + expect(w.find("table").exists()).toBe(false); + }); + + it("does not render the per-session table for zero sessions", () => { + const result = makeResult(); + const w = mount(TaskResultPanel, { props: { result, sessions: [] } }); + expect(w.find("table").exists()).toBe(false); + }); +}); diff --git a/frontend/src/components/TaskResultPanel.vue b/frontend/src/components/TaskResultPanel.vue new file mode 100644 index 0000000..f3755eb --- /dev/null +++ b/frontend/src/components/TaskResultPanel.vue @@ -0,0 +1,86 @@ + + + diff --git a/frontend/src/composables/useTaskDetail.test.ts b/frontend/src/composables/useTaskDetail.test.ts index 1c7d830..0989aa6 100644 --- a/frontend/src/composables/useTaskDetail.test.ts +++ b/frontend/src/composables/useTaskDetail.test.ts @@ -46,6 +46,7 @@ function taskDetail(over: Partial = {}): TaskDetail { queue_id: null, priority: 0, result: null, + sessions: [], work_dir: null, ...over, }; diff --git a/frontend/src/stores/tasks.test.ts b/frontend/src/stores/tasks.test.ts index 96ae0fa..89cb555 100644 --- a/frontend/src/stores/tasks.test.ts +++ b/frontend/src/stores/tasks.test.ts @@ -35,6 +35,7 @@ function detail(over: Partial = {}): TaskDetail { queue_id: null, priority: 0, result: null, + sessions: [], work_dir: null, ...over, }; diff --git a/frontend/src/types/api.ts b/frontend/src/types/api.ts index d7af5fc..814d618 100644 --- a/frontend/src/types/api.ts +++ b/frontend/src/types/api.ts @@ -80,6 +80,10 @@ export interface TaskResult { cost_usd: number; input_tokens: number; output_tokens: number; + // Nullable: historical rows may predate the cache-token backfill — null means + // "unknown", not zero. + cache_creation_input_tokens: number | null; + cache_read_input_tokens: number | null; num_turns: number; is_error: boolean; result_text: string; @@ -89,6 +93,9 @@ export interface TaskResult { export interface TaskDetail extends Task { result: TaskResult | null; + // Every session/turn result for this task, newest-first. `result` above is + // just `sessions[0]`, kept for backward compat. + sessions: TaskResult[]; work_dir: string | null; } @@ -329,6 +336,11 @@ export interface StatsRow { label: string; task_count: number; total_secs: number; + cost_usd: number; + input_tokens: number; + output_tokens: number; + cache_creation_input_tokens: number; + cache_read_input_tokens: number; } export interface StatsResponse { @@ -338,6 +350,11 @@ export interface StatsResponse { total_tasks: number; total_secs: number; rows: StatsRow[]; + total_cost_usd: number; + total_input_tokens: number; + total_output_tokens: number; + total_cache_creation_input_tokens: number; + total_cache_read_input_tokens: number; } export interface StatsQuery { diff --git a/frontend/src/util/formatTokens.test.ts b/frontend/src/util/formatTokens.test.ts new file mode 100644 index 0000000..59591fd --- /dev/null +++ b/frontend/src/util/formatTokens.test.ts @@ -0,0 +1,31 @@ +import { describe, it, expect } from "vitest"; +import { formatTokens } from "./formatTokens"; + +describe("formatTokens", () => { + it("passes sub-1000 counts through unchanged", () => { + expect(formatTokens(999)).toBe("999"); + expect(formatTokens(1)).toBe("1"); + }); + + it("renders 0 as 0, not the sentinel", () => { + expect(formatTokens(0)).toBe("0"); + }); + + it("scales to K", () => { + expect(formatTokens(1234)).toBe("1.2K"); + }); + + it("scales to M", () => { + expect(formatTokens(1234567)).toBe("1.2M"); + }); + + it("renders null as the sentinel", () => { + expect(formatTokens(null)).toBe("—"); + }); + + it("renders negative or non-finite as the sentinel", () => { + expect(formatTokens(-1)).toBe("—"); + expect(formatTokens(NaN)).toBe("—"); + expect(formatTokens(Infinity)).toBe("—"); + }); +}); diff --git a/frontend/src/util/formatTokens.ts b/frontend/src/util/formatTokens.ts new file mode 100644 index 0000000..4802015 --- /dev/null +++ b/frontend/src/util/formatTokens.ts @@ -0,0 +1,12 @@ +// Compact token-count formatter — "999", "1.2K", "1.2M", "—". +// Mirrors formatSecs's sentinel convention (util/duration.ts): "—" is reserved +// for unset/invalid input (null, negative, non-finite); a genuine 0 tokens +// still renders as "0". null covers the cache-token fields, which read as +// "unknown" (not zero) on rows predating the backfill. + +export function formatTokens(n: number | null): string { + if (n === null || !Number.isFinite(n) || n < 0) return "—"; + if (n < 1000) return `${n}`; + if (n < 1_000_000) return `${(n / 1000).toFixed(1)}K`; + return `${(n / 1_000_000).toFixed(1)}M`; +} diff --git a/frontend/src/views/Stats.vue b/frontend/src/views/Stats.vue index a36c798..fb9cd04 100644 --- a/frontend/src/views/Stats.vue +++ b/frontend/src/views/Stats.vue @@ -2,6 +2,7 @@ import { computed, onMounted, ref, watch } from "vue"; import { tasksApi } from "../api/tasks"; import { formatSecs } from "../util/duration"; +import { formatTokens } from "../util/formatTokens"; import { useTaskFilters } from "../util/taskFilters"; import { useServicesStore } from "../stores/services"; import type { StatsGroupBy, StatsResponse, Task } from "../types/api"; @@ -197,7 +198,7 @@ function quickRange(days: number) {
Loading…
-
+

Total tasks

{{ data.total_tasks }}

@@ -210,6 +211,16 @@ function quickRange(days: number) {

Groups

{{ data.rows.length }}

+
+

Total cost

+

${{ data.total_cost_usd.toFixed(2) }}

+
+
+

Total tokens

+

+ {{ formatTokens(data.total_input_tokens) }} in / {{ formatTokens(data.total_output_tokens) }} out +

+

{{ totalLine }}

@@ -221,6 +232,8 @@ function quickRange(days: number) { {{ groupOptions.find((g) => g.value === groupBy)?.label }} Tasks Spent + Cost + Tokens (in/out) Share @@ -229,6 +242,13 @@ function quickRange(days: number) { {{ row.label }} {{ row.task_count }} {{ formatSecs(row.total_secs) }} + ${{ row.cost_usd.toFixed(2) }} + + {{ formatTokens(row.input_tokens) }} / {{ formatTokens(row.output_tokens) }} + {{ data.total_secs > 0 @@ -238,7 +258,7 @@ function quickRange(days: number) { - + No tasks in this window. diff --git a/frontend/src/views/TaskDetail.vue b/frontend/src/views/TaskDetail.vue index 0e91804..893530e 100644 --- a/frontend/src/views/TaskDetail.vue +++ b/frontend/src/views/TaskDetail.vue @@ -4,10 +4,10 @@ import StatusPill from "../components/StatusPill.vue"; import ProviderBadge from "../components/ProviderBadge.vue"; import TurnTimeline from "../components/TurnTimeline.vue"; import TriggerView from "../components/TriggerView.vue"; -import MarkdownView from "../components/MarkdownView.vue"; import DiffView from "../components/DiffView.vue"; import Accordion from "../components/Accordion.vue"; import TaskEditPanel from "../components/TaskEditPanel.vue"; +import TaskResultPanel from "../components/TaskResultPanel.vue"; import TaskOutlinePanel from "../components/TaskOutlinePanel.vue"; import AuthApprovalForm from "../components/AuthApprovalForm.vue"; import AnimatedNumber from "../components/AnimatedNumber.vue"; @@ -317,33 +317,11 @@ watch(pendingApprovals, (p) => { -
-

Result

-
-
-
Cost
-
${{ store.detail.result.cost_usd.toFixed(4) }}
-
-
-
Turns
-
{{ store.detail.result.num_turns }}
-
-
-
Tokens in / out
-
{{ store.detail.result.input_tokens }} / {{ store.detail.result.output_tokens }}
-
-
-
- 
{{ store.detail.result.result_text }}
- -
-
+ Result<(), DbErr> { + // SQLite allows only one ADD COLUMN per ALTER TABLE statement, so these + // are two separate calls (matches m20260601_000011's precedent). + manager + .alter_table( + Table::alter() + .table(TaskSessions::Table) + .add_column( + ColumnDef::new(TaskSessionsExt::CacheCreationInputTokens) + .big_integer() + .null(), + ) + .to_owned(), + ) + .await?; + manager + .alter_table( + Table::alter() + .table(TaskSessions::Table) + .add_column( + ColumnDef::new(TaskSessionsExt::CacheReadInputTokens) + .big_integer() + .null(), + ) + .to_owned(), + ) + .await + } + + async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager + .alter_table( + Table::alter() + .table(TaskSessions::Table) + .drop_column(TaskSessionsExt::CacheCreationInputTokens) + .to_owned(), + ) + .await?; + manager + .alter_table( + Table::alter() + .table(TaskSessions::Table) + .drop_column(TaskSessionsExt::CacheReadInputTokens) + .to_owned(), + ) + .await + } +} + +/// The table was created as `task_results` (see +/// `m20260415_000002_create_task_results`) and renamed to `task_sessions` by +/// raw SQL in `m20260609_000022_normalize_tasks_and_runs` — that migration's +/// `TaskResults` iden now names a stale table, so this migration declares its +/// own iden against the current name instead of reusing it. +#[derive(DeriveIden)] +enum TaskSessions { + Table, +} + +#[derive(DeriveIden)] +enum TaskSessionsExt { + CacheCreationInputTokens, + CacheReadInputTokens, +} diff --git a/src/agent/claude.rs b/src/agent/claude.rs index d812ee7..1191d8d 100644 --- a/src/agent/claude.rs +++ b/src/agent/claude.rs @@ -338,12 +338,14 @@ mod tests { #[test] fn parse_result_picks_the_result_event() { + // `modelUsage` is an unknown field to `ClaudeOutput` — proves unknown + // fields are tolerated/dropped, not just ignored by luck. let stdout = concat!( r#"{"type":"system","subtype":"init","session_id":"abc"}"#, "\n", r#"{"type":"assistant","message":{"usage":{"output_tokens":5}}}"#, "\n", - r#"{"type":"result","result":"done","session_id":"abc","total_cost_usd":0.12,"is_error":false,"num_turns":3,"output_tokens":42}"#, + r#"{"type":"result","result":"done","session_id":"abc","total_cost_usd":0.12,"is_error":false,"num_turns":3,"usage":{"input_tokens":100,"output_tokens":42,"cache_creation_input_tokens":5,"cache_read_input_tokens":3},"modelUsage":{"claude-opus-4-8":{"inputTokens":100,"outputTokens":42}}}"#, "\n", ); let out = ClaudeCode.parse_result(stdout).unwrap(); @@ -351,7 +353,22 @@ mod tests { assert_eq!(out.session_id, "abc"); assert!(!out.is_error); assert_eq!(out.num_turns, 3); - assert_eq!(out.output_tokens, 42); + assert_eq!(out.usage.input_tokens, 100); + assert_eq!(out.usage.output_tokens, 42); + assert_eq!(out.usage.cache_creation_input_tokens, 5); + assert_eq!(out.usage.cache_read_input_tokens, 3); + } + + /// A `result` event that omits `usage` entirely (older/degraded CLI output) + /// parses to `Usage::default()` (all zeros) rather than failing. + #[test] + fn parse_result_defaults_usage_when_absent() { + let stdout = r#"{"type":"result","result":"done","session_id":"abc","total_cost_usd":0.0,"is_error":false,"num_turns":1}"#; + let out = ClaudeCode.parse_result(stdout).unwrap(); + assert_eq!(out.usage.input_tokens, 0); + assert_eq!(out.usage.output_tokens, 0); + assert_eq!(out.usage.cache_creation_input_tokens, 0); + assert_eq!(out.usage.cache_read_input_tokens, 0); } #[test] diff --git a/src/api/handlers/mod.rs b/src/api/handlers/mod.rs index 7adc3ef..1378533 100644 --- a/src/api/handlers/mod.rs +++ b/src/api/handlers/mod.rs @@ -39,7 +39,11 @@ pub struct TaskView { pub struct TaskDetail { #[serde(flatten)] pub task: TaskView, + /// The newest session's summary — kept for the existing UI contract. pub result: Option, + /// Every `task_sessions` row for this task (1:N — one per agent run/resume), + /// newest first; `result` is `sessions.first()`. + pub sessions: Vec, /// Absolute path to this task's git worktree on the agent host. /// `None` if the task's service can no longer be resolved. pub work_dir: Option, diff --git a/src/api/handlers/query.rs b/src/api/handlers/query.rs index 6cedecb..cd7b819 100644 --- a/src/api/handlers/query.rs +++ b/src/api/handlers/query.rs @@ -45,7 +45,7 @@ pub async fn get_task( State(state): State, Path(id): Path, ) -> Result, StatusCode> { - let (task, result) = state + let (task, sessions) = state .task_store .get_task(id) .await @@ -59,7 +59,8 @@ pub async fn get_task( let task = task_view(task, &state).await; Ok(Json(TaskDetail { task, - result, + result: sessions.first().cloned(), + sessions, work_dir, })) } diff --git a/src/api/stats.rs b/src/api/stats.rs index 7a8d5b6..2f7006e 100644 --- a/src/api/stats.rs +++ b/src/api/stats.rs @@ -17,7 +17,11 @@ use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::AppState; -use crate::entity::tasks; +use crate::entity::{task_sessions, tasks}; + +/// SQLite's default bind-param cap is 999; chunk `task_id IN (...)` queries +/// well under that so a large task set doesn't blow the limit. +const ID_CHUNK_SIZE: usize = 500; #[derive(Deserialize)] pub struct StatsQuery { @@ -40,6 +44,11 @@ pub struct StatsRow { pub label: String, pub task_count: i64, pub total_secs: i64, + pub cost_usd: f64, + pub input_tokens: i64, + pub output_tokens: i64, + pub cache_creation_input_tokens: i64, + pub cache_read_input_tokens: i64, } #[derive(Serialize)] @@ -49,9 +58,43 @@ pub struct StatsResponse { pub group_by: String, pub total_tasks: i64, pub total_secs: i64, + pub total_cost_usd: f64, + pub total_input_tokens: i64, + pub total_output_tokens: i64, + pub total_cache_creation_input_tokens: i64, + pub total_cache_read_input_tokens: i64, pub rows: Vec, } +/// A task's usage summed across every `task_sessions` row it has (multiple rows +/// per task happen across resumes/retries — each row is already that run's +/// cumulative total, so summing rows for one task is correct). Attributed to +/// the owning **task's** `created_at` window, not the session's own timestamp — +/// same as time-spent above. +#[derive(Default, Clone, Copy)] +struct UsageAgg { + cost_usd: f64, + input_tokens: i64, + output_tokens: i64, + cache_creation_input_tokens: i64, + cache_read_input_tokens: i64, +} + +/// Fold session rows into a per-task usage total. Pure/no-DB so it's unit +/// testable independently of the aggregation query above. +fn fold_sessions(sessions: &[task_sessions::Model]) -> HashMap { + let mut out: HashMap = HashMap::new(); + for s in sessions { + let entry = out.entry(s.task_id).or_default(); + entry.cost_usd += s.cost_usd; + entry.input_tokens += s.input_tokens; + entry.output_tokens += s.output_tokens; + entry.cache_creation_input_tokens += s.cache_creation_input_tokens.unwrap_or(0); + entry.cache_read_input_tokens += s.cache_read_input_tokens.unwrap_or(0); + } + out +} + pub async fn task_stats( State(state): State, Query(q): Query, @@ -80,6 +123,22 @@ pub async fn task_stats( .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + // Usage/cost is attributed to the owning task's `created_at` window (this + // query's `rows`), not the session's own timestamp — same as time-spent. + // Chunked so a large task set doesn't blow SQLite's default 999 bind-param + // cap on the `IN (...)`. + let task_ids: Vec = rows.iter().map(|t| t.id).collect(); + let mut sessions: Vec = Vec::new(); + for chunk in task_ids.chunks(ID_CHUNK_SIZE) { + let mut found = task_sessions::Entity::find() + .filter(task_sessions::Column::TaskId.is_in(chunk.to_vec())) + .all(state.task_store.db()) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + sessions.append(&mut found); + } + let usage_by_task = fold_sessions(&sessions); + // Service slugs for service-grouped labels. let services = state .service_store @@ -107,6 +166,11 @@ pub async fn task_stats( // Accumulated in-loop so the service filter (applied here, not in SQL) keeps // the headline count consistent with the rows. let mut total_tasks: i64 = 0; + let mut total_cost_usd: f64 = 0.0; + let mut total_input_tokens: i64 = 0; + let mut total_output_tokens: i64 = 0; + let mut total_cache_creation_input_tokens: i64 = 0; + let mut total_cache_read_input_tokens: i64 = 0; for t in rows { if let Some(want) = q.service_id { @@ -118,6 +182,12 @@ pub async fn task_stats( total_tasks += 1; let secs = duration_secs(&t, now); total_secs += secs; + let usage = usage_by_task.get(&t.id).copied().unwrap_or_default(); + total_cost_usd += usage.cost_usd; + total_input_tokens += usage.input_tokens; + total_output_tokens += usage.output_tokens; + total_cache_creation_input_tokens += usage.cache_creation_input_tokens; + total_cache_read_input_tokens += usage.cache_read_input_tokens; let project_name = project_meta .get(&t.project_id) .map(|(name, _)| name.clone()) @@ -150,9 +220,19 @@ pub async fn task_stats( label, task_count: 0, total_secs: 0, + cost_usd: 0.0, + input_tokens: 0, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, }); entry.task_count += 1; entry.total_secs += secs; + entry.cost_usd += usage.cost_usd; + entry.input_tokens += usage.input_tokens; + entry.output_tokens += usage.output_tokens; + entry.cache_creation_input_tokens += usage.cache_creation_input_tokens; + entry.cache_read_input_tokens += usage.cache_read_input_tokens; } let mut rows: Vec = buckets.into_values().collect(); @@ -164,6 +244,11 @@ pub async fn task_stats( group_by: group_by.into(), total_tasks, total_secs, + total_cost_usd, + total_input_tokens, + total_output_tokens, + total_cache_creation_input_tokens, + total_cache_read_input_tokens, rows, })) } @@ -184,3 +269,76 @@ fn duration_secs(t: &tasks::Model, now: DateTime) -> i64 { }; (end - started).num_seconds().max(0) } + +#[cfg(test)] +mod tests { + use super::*; + + fn session( + task_id: Uuid, + cost: f64, + input: i64, + output: i64, + cache_create: Option, + cache_read: Option, + ) -> task_sessions::Model { + task_sessions::Model { + id: Uuid::new_v4(), + task_id, + cost_usd: cost, + input_tokens: input, + output_tokens: output, + cache_creation_input_tokens: cache_create, + cache_read_input_tokens: cache_read, + num_turns: 1, + is_error: false, + result_text: String::new(), + session_id: Uuid::new_v4().to_string(), + created_at: Utc::now().into(), + } + } + + /// Multiple session rows for one task_id (a resume/retry history) sum + /// correctly across every field. + #[test] + fn fold_sessions_sums_multiple_rows_per_task() { + let task_id = Uuid::new_v4(); + let sessions = vec![ + session(task_id, 1.0, 10, 20, Some(1), Some(2)), + session(task_id, 2.5, 30, 40, Some(3), Some(4)), + ]; + let agg = fold_sessions(&sessions); + let entry = agg.get(&task_id).expect("task must be present"); + assert_eq!(entry.cost_usd, 3.5); + assert_eq!(entry.input_tokens, 40); + assert_eq!(entry.output_tokens, 60); + assert_eq!(entry.cache_creation_input_tokens, 4); + assert_eq!(entry.cache_read_input_tokens, 6); + } + + /// `None` cache fields (a pre-backfill row) contribute 0, not a panic/NaN. + #[test] + fn fold_sessions_none_cache_fields_contribute_zero() { + let task_id = Uuid::new_v4(); + let sessions = vec![session(task_id, 1.0, 10, 20, None, None)]; + let agg = fold_sessions(&sessions); + let entry = agg.get(&task_id).unwrap(); + assert_eq!(entry.cache_creation_input_tokens, 0); + assert_eq!(entry.cache_read_input_tokens, 0); + } + + /// Distinct task_ids get distinct entries, not merged into one. + #[test] + fn fold_sessions_keeps_tasks_separate() { + let a = Uuid::new_v4(); + let b = Uuid::new_v4(); + let sessions = vec![ + session(a, 1.0, 1, 1, Some(0), Some(0)), + session(b, 2.0, 2, 2, Some(0), Some(0)), + ]; + let agg = fold_sessions(&sessions); + assert_eq!(agg.len(), 2); + assert_eq!(agg.get(&a).unwrap().cost_usd, 1.0); + assert_eq!(agg.get(&b).unwrap().cost_usd, 2.0); + } +} diff --git a/src/entity/task_sessions.rs b/src/entity/task_sessions.rs index a70a62b..26c5e9a 100644 --- a/src/entity/task_sessions.rs +++ b/src/entity/task_sessions.rs @@ -10,6 +10,11 @@ pub struct Model { pub cost_usd: f64, pub input_tokens: i64, pub output_tokens: i64, + /// NULL = not yet backfilled (pre-fix historical row); every write path + /// since the #163 fix sets `Some(...)` (error paths `Some(0)`). See + /// `src/jobs/usage_backfill.rs`. + pub cache_creation_input_tokens: Option, + pub cache_read_input_tokens: Option, pub num_turns: i32, pub is_error: bool, #[sea_orm(column_type = "Text")] diff --git a/src/jobs/lifecycle.rs b/src/jobs/lifecycle.rs index c73239b..da91255 100644 --- a/src/jobs/lifecycle.rs +++ b/src/jobs/lifecycle.rs @@ -322,6 +322,8 @@ impl TaskStore { cost_usd: Set(0.0), input_tokens: Set(0), output_tokens: Set(0), + cache_creation_input_tokens: Set(Some(0)), + cache_read_input_tokens: Set(Some(0)), num_turns: Set(0), is_error: Set(false), result_text: Set(text.to_string()), @@ -410,7 +412,7 @@ mod tests { use crate::jobs::hub::LiveSessions; use crate::jobs::registry::RunningTasks; use crate::jobs::store::{TaskStore, TaskStoreDeps}; - use crate::jobs::types::{ClaudeOutput, TriggerReason}; + use crate::jobs::types::{ClaudeOutput, TriggerReason, Usage}; use crate::project::{ BranchStatus, NewBranchEntry, NewProjectConfig, ProjectStore, ProviderKind, }; @@ -590,8 +592,14 @@ mod tests { .finish_task(id, "failed", "failed", Some("killed: token budget")) .await .unwrap(); - let (_, result) = store.get_task(id).await.unwrap().unwrap(); - assert!(result.unwrap().result_text.contains("token budget")); + let (_, sessions) = store.get_task(id).await.unwrap().unwrap(); + assert!( + sessions + .first() + .unwrap() + .result_text + .contains("token budget") + ); drop(store); drop(db); @@ -636,23 +644,23 @@ mod tests { let n = store.recover_orphans().await.unwrap(); assert_eq!(n, 2, "in-flight + confirmed-unstarted are orphans"); - let (t, result) = store.get_task(inflight).await.unwrap().unwrap(); + let (t, sessions) = store.get_task(inflight).await.unwrap().unwrap(); assert_eq!( (t.agent_state.as_str(), t.task_state.as_str()), ("failed", "failed") ); assert!(t.finished_at.is_some()); - assert!(result.unwrap().result_text.contains("orphan")); + assert!(sessions.first().unwrap().result_text.contains("orphan")); // Confirmed-but-unstarted → rewound to cold/pending (Run shows again), no // finished_at, no orphan note. - let (c, c_result) = store.get_task(confirmed).await.unwrap().unwrap(); + let (c, c_sessions) = store.get_task(confirmed).await.unwrap().unwrap(); assert_eq!( (c.agent_state.as_str(), c.task_state.as_str()), ("cold", "pending") ); assert!(c.finished_at.is_none()); - assert!(c_result.is_none()); + assert!(c_sessions.is_empty()); let (p, _) = store.get_task(pending).await.unwrap().unwrap(); assert_eq!( @@ -791,8 +799,11 @@ mod tests { is_error: false, num_turns: 3, subtype: "success".into(), - input_tokens: 100, - output_tokens: 50, + usage: Usage { + input_tokens: 100, + output_tokens: 50, + ..Default::default() + }, }, ) .await diff --git a/src/jobs/mod.rs b/src/jobs/mod.rs index cfaada8..1f511c8 100644 --- a/src/jobs/mod.rs +++ b/src/jobs/mod.rs @@ -24,3 +24,4 @@ pub mod turn; pub mod turn_kill; pub mod turn_loop; pub mod types; +pub mod usage_backfill; diff --git a/src/jobs/queries.rs b/src/jobs/queries.rs index 2c2c9c3..1eb282e 100644 --- a/src/jobs/queries.rs +++ b/src/jobs/queries.rs @@ -232,10 +232,13 @@ impl TaskStore { .context("loading task events") } + /// A task alongside every `task_sessions` row it has accumulated (1:N — + /// one row per agent run/resume), newest first. Callers that only want the + /// most recent run's summary use `.first()`. pub async fn get_task( &self, task_id: Uuid, - ) -> Result)>> { + ) -> Result)>> { let task = tasks::Entity::find_by_id(task_id) .one(self.db()) .await @@ -245,14 +248,13 @@ impl TaskStore { return Ok(None); }; - // Sessions are 1:N; surface the most recent run for the detail view. - let result = task_sessions::Entity::find() + let sessions = task_sessions::Entity::find() .filter(task_sessions::Column::TaskId.eq(task_id)) .order_by_desc(task_sessions::Column::CreatedAt) - .one(self.db()) + .all(self.db()) .await .context("db error")?; - Ok(Some((task, result))) + Ok(Some((task, sessions))) } } diff --git a/src/jobs/results.rs b/src/jobs/results.rs index 49eeaec..498de11 100644 --- a/src/jobs/results.rs +++ b/src/jobs/results.rs @@ -25,8 +25,11 @@ impl TaskStore { { let mut active: task_sessions::ActiveModel = existing.into(); active.cost_usd = Set(output.total_cost_usd); - active.input_tokens = Set(output.input_tokens as i64); - active.output_tokens = Set(output.output_tokens as i64); + active.input_tokens = Set(output.usage.input_tokens as i64); + active.output_tokens = Set(output.usage.output_tokens as i64); + active.cache_creation_input_tokens = + Set(Some(output.usage.cache_creation_input_tokens as i64)); + active.cache_read_input_tokens = Set(Some(output.usage.cache_read_input_tokens as i64)); active.num_turns = Set(output.num_turns as i32); active.is_error = Set(output.is_error); active.result_text = Set(output.result.clone()); @@ -44,8 +47,10 @@ impl TaskStore { id: Set(Uuid::new_v4()), task_id: Set(task_id), cost_usd: Set(output.total_cost_usd), - input_tokens: Set(output.input_tokens as i64), - output_tokens: Set(output.output_tokens as i64), + input_tokens: Set(output.usage.input_tokens as i64), + output_tokens: Set(output.usage.output_tokens as i64), + cache_creation_input_tokens: Set(Some(output.usage.cache_creation_input_tokens as i64)), + cache_read_input_tokens: Set(Some(output.usage.cache_read_input_tokens as i64)), num_turns: Set(output.num_turns as i32), is_error: Set(output.is_error), result_text: Set(output.result.clone()), @@ -68,6 +73,8 @@ impl TaskStore { cost_usd: Set(0.0), input_tokens: Set(0), output_tokens: Set(0), + cache_creation_input_tokens: Set(Some(0)), + cache_read_input_tokens: Set(Some(0)), num_turns: Set(0), is_error: Set(true), result_text: Set(error.to_string()), @@ -83,3 +90,262 @@ impl TaskStore { Ok(()) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use sea_orm::DatabaseConnection; + + use super::*; + use crate::auth::store::AuthStore; + use crate::auth::waiter::AuthWaiter; + use crate::config::Config; + use crate::jobs::db_fixture::{drop_db, fresh_db}; + use crate::jobs::hub::LiveSessions; + use crate::jobs::registry::RunningTasks; + use crate::jobs::store::TaskStoreDeps; + use crate::jobs::types::{TriggerReason, Usage}; + use crate::project::{NewProjectConfig, ProjectStore, ProviderKind}; + use crate::provider::ProviderRegistry; + use crate::service::{NewService, ServiceStore}; + use crate::workspace::Workspace; + + fn store_with(db: &DatabaseConnection) -> Arc { + let config = Config { + api_bearer_token: None, + database_url: String::new(), + repo_base_path: "/tmp/agent-test".into(), + max_concurrent_jobs: 3, + listen_addr: String::new(), + public_base_url: None, + task_token_budget: 1_000_000, + operator_approval_timeout_secs: 0, + job_timeout_secs: 0, + warm_idle_timeout_secs: 0, + nvm_dir: None, + project_db_admin_url: None, + project_db_host_for_agent: None, + }; + Arc::new(TaskStore::new(TaskStoreDeps { + db: db.clone(), + config, + providers: ProviderRegistry::new(ServiceStore::new(db.clone())), + project_store: Arc::new(ProjectStore::new(db.clone())), + workspace: Arc::new(Workspace::new("/tmp/agent-test")), + running: RunningTasks::new(), + hub: LiveSessions::new(db.clone()), + auth_store: Arc::new(AuthStore::new(db.clone())), + auth_waiter: AuthWaiter::new(), + })) + } + + async fn seed_service(db: &DatabaseConnection) -> Uuid { + ServiceStore::new(db.clone()) + .create(NewService { + kind: ProviderKind::Github, + slug: format!("svc-{}", Uuid::new_v4().simple()), + display_name: "test".into(), + base_url: "https://example.test".into(), + token: "t".into(), + webhook_secret: "s".into(), + bot_username: "bot".into(), + autofire: false, + auth_kind: Default::default(), + app_credentials: None, + trigger_mode: Default::default(), + trigger_label: String::new(), + models: None, + triggers: None, + }) + .await + .expect("seed service") + .id + } + + async fn seed_project(db: &DatabaseConnection, service_id: Uuid) -> Uuid { + ProjectStore::new(db.clone()) + .upsert_project(NewProjectConfig { + provider: ProviderKind::Github, + service_id, + project_slug: format!("p-{}", Uuid::new_v4().simple()), + full_name: "acme/widgets".into(), + remote_url: "git@x:acme/widgets.git".into(), + default_branch: "main".into(), + my_username: "bot".into(), + }) + .await + .expect("seed project") + .0 + .id + } + + async fn seed_task(db: &DatabaseConnection, store: &Arc) -> Uuid { + let svc = seed_service(db).await; + let project_id = seed_project(db, svc).await; + store + .create_task( + TriggerReason::Issue { + iid: 1, + title: "t".into(), + description: String::new(), + url: "http://x/1".into(), + }, + project_id, + ) + .await + .expect("create task") + } + + fn output(session_id: &str, cost: f64, turns: u32) -> ClaudeOutput { + ClaudeOutput { + result: "done".into(), + session_id: session_id.into(), + total_cost_usd: cost, + is_error: false, + num_turns: turns, + subtype: "success".into(), + usage: Usage { + input_tokens: 100, + output_tokens: 50, + cache_creation_input_tokens: 10, + cache_read_input_tokens: 5, + }, + } + } + + async fn session_row( + db: &DatabaseConnection, + task_id: Uuid, + session_id: &str, + ) -> task_sessions::Model { + task_sessions::Entity::find() + .filter(task_sessions::Column::TaskId.eq(task_id)) + .filter(task_sessions::Column::SessionId.eq(session_id)) + .one(db) + .await + .unwrap() + .expect("session row must exist") + } + + /// A fresh session_id sets all four token fields from `ClaudeOutput.usage`. + #[tokio::test] + async fn replace_result_fresh_session_sets_all_token_fields() { + let Some((db, name, admin)) = fresh_db().await else { + eprintln!("DATABASE_URL not set; skipping"); + return; + }; + let store = store_with(&db); + let task_id = seed_task(&db, &store).await; + + store + .replace_result(task_id, &output("sess-1", 1.5, 2)) + .await + .unwrap(); + + let row = session_row(&db, task_id, "sess-1").await; + assert_eq!(row.input_tokens, 100); + assert_eq!(row.output_tokens, 50); + assert_eq!(row.cache_creation_input_tokens, Some(10)); + assert_eq!(row.cache_read_input_tokens, Some(5)); + + drop(store); + drop(db); + drop_db(&admin, &name).await; + } + + /// Calling `replace_result` twice with the same session_id updates the one + /// row in place rather than creating a second — matches the overwrite (not + /// accumulate) semantics. + #[tokio::test] + async fn replace_result_same_session_updates_in_place() { + let Some((db, name, admin)) = fresh_db().await else { + return; + }; + let store = store_with(&db); + let task_id = seed_task(&db, &store).await; + + store + .replace_result(task_id, &output("sess-1", 1.0, 1)) + .await + .unwrap(); + store + .replace_result(task_id, &output("sess-1", 3.0, 2)) + .await + .unwrap(); + + let count = task_sessions::Entity::find() + .filter(task_sessions::Column::TaskId.eq(task_id)) + .filter(task_sessions::Column::SessionId.eq("sess-1")) + .all(&db) + .await + .unwrap() + .len(); + assert_eq!(count, 1, "second call updates, doesn't insert"); + + let row = session_row(&db, task_id, "sess-1").await; + assert_eq!(row.cost_usd, 3.0); + assert_eq!(row.num_turns, 2); + + drop(store); + drop(db); + drop_db(&admin, &name).await; + } + + /// A different session_id for the same task creates a second row. + #[tokio::test] + async fn replace_result_new_session_creates_second_row() { + let Some((db, name, admin)) = fresh_db().await else { + return; + }; + let store = store_with(&db); + let task_id = seed_task(&db, &store).await; + + store + .replace_result(task_id, &output("sess-1", 1.0, 1)) + .await + .unwrap(); + store + .replace_result(task_id, &output("sess-2", 2.0, 1)) + .await + .unwrap(); + + let count = task_sessions::Entity::find() + .filter(task_sessions::Column::TaskId.eq(task_id)) + .all(&db) + .await + .unwrap() + .len(); + assert_eq!(count, 2, "a new session_id starts a new row"); + + drop(store); + drop(db); + drop_db(&admin, &name).await; + } + + /// `save_error_result` sets both new cache columns to `Some(0)`, not NULL — + /// the write-path invariant the backfill sweep's NULL-selection relies on. + #[tokio::test] + async fn save_error_result_sets_cache_columns_to_zero() { + let Some((db, name, admin)) = fresh_db().await else { + return; + }; + let store = store_with(&db); + let task_id = seed_task(&db, &store).await; + + store.save_error_result(task_id, "boom").await.unwrap(); + + let row = task_sessions::Entity::find() + .filter(task_sessions::Column::TaskId.eq(task_id)) + .one(&db) + .await + .unwrap() + .expect("error row must exist"); + assert_eq!(row.cache_creation_input_tokens, Some(0)); + assert_eq!(row.cache_read_input_tokens, Some(0)); + + drop(store); + drop(db); + drop_db(&admin, &name).await; + } +} diff --git a/src/jobs/turn.rs b/src/jobs/turn.rs index 895ac6a..9825a3e 100644 --- a/src/jobs/turn.rs +++ b/src/jobs/turn.rs @@ -255,8 +255,11 @@ mod tests { is_error, num_turns: 3, subtype: subtype.to_string(), - input_tokens: 0, - output_tokens: 0, + usage: crate::jobs::types::Usage { + input_tokens: 0, + output_tokens: 0, + ..Default::default() + }, } } diff --git a/src/jobs/types.rs b/src/jobs/types.rs index 159ba0e..5c0be47 100644 --- a/src/jobs/types.rs +++ b/src/jobs/types.rs @@ -184,10 +184,26 @@ pub struct ClaudeOutput { /// backend omits it — then `is_error` is the fallback signal (#106). #[serde(default)] pub subtype: String, + /// The real `claude` CLI nests token counts under `usage.*`, not top-level + /// (issue #163 — before this fix, top-level `input_tokens`/`output_tokens` + /// fields here always deserialized to 0 in production). + #[serde(default)] + pub usage: Usage, +} + +/// Token accounting from the `result` event's `usage` object. Cumulative for +/// the whole session as of that event, not a per-turn delta (mirrors +/// `total_cost_usd`/`num_turns` on `ClaudeOutput`). +#[derive(Debug, Default, Deserialize)] +pub struct Usage { #[serde(default)] pub input_tokens: u64, #[serde(default)] pub output_tokens: u64, + #[serde(default)] + pub cache_creation_input_tokens: u64, + #[serde(default)] + pub cache_read_input_tokens: u64, } impl ClaudeOutput { @@ -209,7 +225,7 @@ impl ClaudeOutput { #[cfg(test)] mod tests { - use super::{ClaudeOutput, TriggerReason}; + use super::{ClaudeOutput, TriggerReason, Usage}; use uuid::Uuid; fn review_mr(iid: u64, head_sha: &str) -> TriggerReason { @@ -294,8 +310,7 @@ mod tests { is_error, num_turns: 1, subtype: subtype.into(), - input_tokens: 0, - output_tokens: 0, + usage: Usage::default(), } } diff --git a/src/jobs/usage_backfill.rs b/src/jobs/usage_backfill.rs new file mode 100644 index 0000000..97fd763 --- /dev/null +++ b/src/jobs/usage_backfill.rs @@ -0,0 +1,365 @@ +//! Startup backfill sweep for `task_sessions` rows whose cache-token columns +//! predate the #163 fix (the `claude` CLI nests token counts under +//! `usage.*`, but `ClaudeOutput` used to declare them top-level, so they always +//! deserialized to 0 in production). The raw `result` stream-json event was +//! persisted verbatim in `events.payload` on every historical turn, so this +//! recovers the real counts instead of leaving old rows stuck at 0 forever. +//! +//! Modeled on the other best-effort startup sweeps (`TaskStore::recover_orphans`, +//! `ProjectDb::sweep_orphans`): non-fatal, logged, `tokio::spawn`ed from +//! `main.rs` so it never blocks startup. + +use std::collections::HashMap; + +use anyhow::{Context, Result}; +use sea_orm::*; +use uuid::Uuid; + +use crate::entity::{events, task_sessions}; +use crate::jobs::types::ClaudeOutput; + +/// `events` rows are paged per task at this size — large enough to make a dent +/// per round-trip, small enough to bound memory for a long-lived task. +const PAGE_SIZE: u64 = 500; + +/// Backfill every `task_sessions` row still missing its cache-token columns +/// (`cache_read_input_tokens IS NULL` is the "not yet backfilled" marker) from +/// the raw `result` events already persisted for that row's task. Idempotent: +/// every row the sweep visits is rewritten with `Some(...)` — a match from a +/// `result` event, else `Some(0)` — so it leaves the NULL set and a re-run +/// touches zero rows. Returns the number of rows updated. +pub async fn backfill_session_usage(db: &DatabaseConnection) -> Result { + let pending = task_sessions::Entity::find() + .filter(task_sessions::Column::CacheReadInputTokens.is_null()) + .all(db) + .await + .context("loading task_sessions rows pending usage backfill")?; + + if pending.is_empty() { + return Ok(0); + } + + let mut by_task: HashMap> = HashMap::new(); + for row in pending { + by_task.entry(row.task_id).or_default().push(row); + } + + let mut updated = 0u64; + for (task_id, rows) in by_task { + let results = scan_result_events(db, task_id).await?; + for row in rows { + let output = results + .get(&row.session_id) + .and_then(|payload| serde_json::from_value::(payload.clone()).ok()); + + let mut active: task_sessions::ActiveModel = row.into(); + match output { + Some(out) => { + active.input_tokens = Set(out.usage.input_tokens as i64); + active.output_tokens = Set(out.usage.output_tokens as i64); + active.cache_creation_input_tokens = + Set(Some(out.usage.cache_creation_input_tokens as i64)); + active.cache_read_input_tokens = + Set(Some(out.usage.cache_read_input_tokens as i64)); + active.cost_usd = Set(out.total_cost_usd); + active.num_turns = Set(out.num_turns as i32); + } + None => { + // No matching result event covers this session — nothing to + // recover; zero (not left NULL) so this row leaves the sweep's + // selection for good. + active.cache_creation_input_tokens = Set(Some(0)); + active.cache_read_input_tokens = Set(Some(0)); + } + } + active + .update(db) + .await + .context("backfilling task_sessions row")?; + updated += 1; + } + } + + Ok(updated) +} + +/// Page through `task_id`'s persisted `events` (agent-stream frames only — +/// `kind = "event"`) in `seq` order, folding every `result`-type payload into a +/// `session_id → payload` map. Later events overwrite earlier ones for the same +/// `session_id` (last-wins by `seq`), matching `replace_result`'s overwrite (not +/// accumulate) semantics — the `result` event's `usage`/`total_cost_usd`/ +/// `num_turns` are already session-cumulative as of that turn. +async fn scan_result_events( + db: &DatabaseConnection, + task_id: Uuid, +) -> Result> { + let mut out = HashMap::new(); + let mut last_seq = -1i64; + loop { + let page = events::Entity::find() + .filter(events::Column::TaskId.eq(task_id)) + .filter(events::Column::Kind.eq("event")) + .filter(events::Column::Seq.gt(last_seq)) + .order_by_asc(events::Column::Seq) + .limit(PAGE_SIZE) + .all(db) + .await + .context("paging task events for usage backfill")?; + + let page_len = page.len() as u64; + for row in &page { + last_seq = row.seq; + if row.payload.get("type").and_then(|t| t.as_str()) == Some("result") + && let Some(sid) = row.payload.get("session_id").and_then(|s| s.as_str()) + { + out.insert(sid.to_string(), row.payload.clone()); + } + } + + if page_len < PAGE_SIZE { + break; + } + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::auth::store::AuthStore; + use crate::auth::waiter::AuthWaiter; + use crate::config::Config; + use crate::jobs::db_fixture::{drop_db, fresh_db}; + use crate::jobs::hub::LiveSessions; + use crate::jobs::registry::RunningTasks; + use crate::jobs::store::{TaskStore, TaskStoreDeps}; + use crate::jobs::types::TriggerReason; + use crate::project::{NewProjectConfig, ProjectStore, ProviderKind}; + use crate::provider::ProviderRegistry; + use crate::service::{NewService, ServiceStore}; + use crate::workspace::Workspace; + + fn store_with(db: &DatabaseConnection) -> Arc { + let config = Config { + api_bearer_token: None, + database_url: String::new(), + repo_base_path: "/tmp/agent-test".into(), + max_concurrent_jobs: 3, + listen_addr: String::new(), + public_base_url: None, + task_token_budget: 1_000_000, + operator_approval_timeout_secs: 0, + job_timeout_secs: 0, + warm_idle_timeout_secs: 0, + nvm_dir: None, + project_db_admin_url: None, + project_db_host_for_agent: None, + }; + Arc::new(TaskStore::new(TaskStoreDeps { + db: db.clone(), + config, + providers: ProviderRegistry::new(ServiceStore::new(db.clone())), + project_store: Arc::new(ProjectStore::new(db.clone())), + workspace: Arc::new(Workspace::new("/tmp/agent-test")), + running: RunningTasks::new(), + hub: LiveSessions::new(db.clone()), + auth_store: Arc::new(AuthStore::new(db.clone())), + auth_waiter: AuthWaiter::new(), + })) + } + + async fn seed_task(db: &DatabaseConnection, store: &Arc) -> Uuid { + let service_id = ServiceStore::new(db.clone()) + .create(NewService { + kind: ProviderKind::Github, + slug: format!("svc-{}", Uuid::new_v4().simple()), + display_name: "test".into(), + base_url: "https://example.test".into(), + token: "t".into(), + webhook_secret: "s".into(), + bot_username: "bot".into(), + autofire: false, + auth_kind: Default::default(), + app_credentials: None, + trigger_mode: Default::default(), + trigger_label: String::new(), + models: None, + triggers: None, + }) + .await + .expect("seed service") + .id; + let project_id = ProjectStore::new(db.clone()) + .upsert_project(NewProjectConfig { + provider: ProviderKind::Github, + service_id, + project_slug: format!("p-{}", Uuid::new_v4().simple()), + full_name: "acme/widgets".into(), + remote_url: "git@x:acme/widgets.git".into(), + default_branch: "main".into(), + my_username: "bot".into(), + }) + .await + .expect("seed project") + .0 + .id; + store + .create_task( + TriggerReason::Issue { + iid: 1, + title: "t".into(), + description: String::new(), + url: "http://x/1".into(), + }, + project_id, + ) + .await + .expect("create task") + } + + /// Insert a NULL-cache-columns `task_sessions` row directly (simulating a + /// pre-#163 row) for `session_id`. + async fn seed_pending_session( + db: &DatabaseConnection, + task_id: Uuid, + session_id: &str, + ) -> Uuid { + let id = Uuid::new_v4(); + task_sessions::ActiveModel { + id: Set(id), + task_id: Set(task_id), + cost_usd: Set(0.0), + input_tokens: Set(0), + output_tokens: Set(0), + cache_creation_input_tokens: Set(None), + cache_read_input_tokens: Set(None), + num_turns: Set(0), + is_error: Set(false), + result_text: Set(String::new()), + session_id: Set(session_id.to_string()), + created_at: Set(chrono::Utc::now().into()), + } + .insert(db) + .await + .expect("seed pending session row"); + id + } + + async fn seed_result_event( + db: &DatabaseConnection, + task_id: Uuid, + seq: i64, + session_id: &str, + input_tokens: u64, + output_tokens: u64, + ) { + let payload = serde_json::json!({ + "type": "result", + "result": "done", + "session_id": session_id, + "total_cost_usd": 1.5, + "is_error": false, + "num_turns": 2, + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cache_creation_input_tokens": 7, + "cache_read_input_tokens": 4, + }, + }); + events::ActiveModel { + task_id: Set(task_id), + seq: Set(seq), + kind: Set("event".to_string()), + payload: Set(payload), + } + .insert(db) + .await + .expect("seed result event"); + } + + /// Two `result` events for one session — the sweep must pick the last one + /// (highest seq), not the first. + #[tokio::test] + async fn backfill_picks_the_last_result_event_by_seq() { + let Some((db, name, admin)) = fresh_db().await else { + eprintln!("DATABASE_URL not set; skipping"); + return; + }; + let store = store_with(&db); + let task_id = seed_task(&db, &store).await; + let row_id = seed_pending_session(&db, task_id, "sess-1").await; + seed_result_event(&db, task_id, 0, "sess-1", 10, 5).await; + seed_result_event(&db, task_id, 1, "sess-1", 100, 50).await; + + let n = backfill_session_usage(&db).await.unwrap(); + assert_eq!(n, 1); + + let row = task_sessions::Entity::find_by_id(row_id) + .one(&db) + .await + .unwrap() + .unwrap(); + assert_eq!(row.input_tokens, 100, "last event by seq wins"); + assert_eq!(row.output_tokens, 50); + assert_eq!(row.cache_creation_input_tokens, Some(7)); + assert_eq!(row.cache_read_input_tokens, Some(4)); + + drop(store); + drop(db); + drop_db(&admin, &name).await; + } + + /// A session with no matching result event gets `Some(0)` cache columns, + /// not left NULL. + #[tokio::test] + async fn backfill_zeroes_sessions_with_no_matching_event() { + let Some((db, name, admin)) = fresh_db().await else { + return; + }; + let store = store_with(&db); + let task_id = seed_task(&db, &store).await; + let row_id = seed_pending_session(&db, task_id, "sess-orphan").await; + // Some other session's event exists, but none for "sess-orphan". + seed_result_event(&db, task_id, 0, "sess-other", 10, 5).await; + + let n = backfill_session_usage(&db).await.unwrap(); + assert_eq!(n, 1); + + let row = task_sessions::Entity::find_by_id(row_id) + .one(&db) + .await + .unwrap() + .unwrap(); + assert_eq!(row.cache_creation_input_tokens, Some(0)); + assert_eq!(row.cache_read_input_tokens, Some(0)); + + drop(store); + drop(db); + drop_db(&admin, &name).await; + } + + /// Running the sweep a second time touches 0 rows — every row it visited + /// left the NULL selection on the first pass. + #[tokio::test] + async fn backfill_is_idempotent() { + let Some((db, name, admin)) = fresh_db().await else { + return; + }; + let store = store_with(&db); + let task_id = seed_task(&db, &store).await; + seed_pending_session(&db, task_id, "sess-1").await; + seed_result_event(&db, task_id, 0, "sess-1", 10, 5).await; + + let first = backfill_session_usage(&db).await.unwrap(); + assert_eq!(first, 1); + let second = backfill_session_usage(&db).await.unwrap(); + assert_eq!(second, 0, "a re-run must touch nothing once caught up"); + + drop(store); + drop(db); + drop_db(&admin, &name).await; + } +} diff --git a/src/main.rs b/src/main.rs index 560d3be..6473bce 100644 --- a/src/main.rs +++ b/src/main.rs @@ -86,6 +86,23 @@ async fn main() -> anyhow::Result<()> { } } + // Historical `task_sessions` rows predate the #163 fix that reads token + // counts from the CLI's actual nested `usage.*` shape — backfill them from + // the raw `result` events already persisted in `events`. Backgrounded so a + // large backlog never delays serving. + { + let db = task_store.db().clone(); + tokio::spawn(async move { + match agent::jobs::usage_backfill::backfill_session_usage(&db).await { + Ok(0) => {} + Ok(n) => tracing::info!(backfilled = n, "backfilled task_sessions token usage"), + Err(e) => { + tracing::warn!(error = %e, "failed to backfill task_sessions token usage") + } + } + }); + } + // The admission loop pulls queued tasks into free slots whenever a slot // frees (signalled via request_admit). It owns the recursive confirm path // that the run future's closure can't call directly.