diff --git a/Cargo.lock b/Cargo.lock index e4e138ffa..f88b44544 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1536,7 +1536,7 @@ dependencies = [ [[package]] name = "fabro-agent" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "async-trait", @@ -1575,7 +1575,7 @@ dependencies = [ [[package]] name = "fabro-api" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "chrono", "fabro-config", @@ -1596,7 +1596,7 @@ dependencies = [ [[package]] name = "fabro-auth" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "async-trait", @@ -1620,7 +1620,7 @@ dependencies = [ [[package]] name = "fabro-checkpoint" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "chrono", "fabro-config", @@ -1636,7 +1636,7 @@ dependencies = [ [[package]] name = "fabro-cli" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "assert_cmd", @@ -1732,7 +1732,7 @@ dependencies = [ [[package]] name = "fabro-client" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "bytes", @@ -1761,7 +1761,7 @@ dependencies = [ [[package]] name = "fabro-config" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "chrono", @@ -1788,7 +1788,7 @@ dependencies = [ [[package]] name = "fabro-core" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "async-trait", "fabro-types", @@ -1803,7 +1803,7 @@ dependencies = [ [[package]] name = "fabro-dev" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "assert_cmd", @@ -1823,7 +1823,7 @@ dependencies = [ [[package]] name = "fabro-devcontainer" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "fabro-http", "fabro-static", @@ -1840,7 +1840,7 @@ dependencies = [ [[package]] name = "fabro-dump" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "bytes", @@ -1854,7 +1854,7 @@ dependencies = [ [[package]] name = "fabro-github" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "base64", @@ -1876,7 +1876,7 @@ dependencies = [ [[package]] name = "fabro-graphviz" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "fabro-types", @@ -1890,7 +1890,7 @@ dependencies = [ [[package]] name = "fabro-hooks" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "async-trait", "fabro-agent", @@ -1914,7 +1914,7 @@ dependencies = [ [[package]] name = "fabro-http" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "fabro-static", "http", @@ -1924,7 +1924,7 @@ dependencies = [ [[package]] name = "fabro-install" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "base64", @@ -1939,7 +1939,7 @@ dependencies = [ [[package]] name = "fabro-interview" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "async-trait", "dialoguer", @@ -1954,7 +1954,7 @@ dependencies = [ [[package]] name = "fabro-llm" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "async-trait", @@ -1986,7 +1986,7 @@ dependencies = [ [[package]] name = "fabro-macros" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "clap", "fabro-options-metadata", @@ -1997,7 +1997,7 @@ dependencies = [ [[package]] name = "fabro-mcp" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "fabro-config", @@ -2013,7 +2013,7 @@ dependencies = [ [[package]] name = "fabro-model" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "fabro-static", "insta", @@ -2024,7 +2024,7 @@ dependencies = [ [[package]] name = "fabro-oauth" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "axum", @@ -2046,7 +2046,7 @@ dependencies = [ [[package]] name = "fabro-options-metadata" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "serde", "serde_json", @@ -2054,7 +2054,7 @@ dependencies = [ [[package]] name = "fabro-proc" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "cc", "libc", @@ -2063,7 +2063,7 @@ dependencies = [ [[package]] name = "fabro-redact" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "aho-corasick", "ref-cast", @@ -2079,7 +2079,7 @@ dependencies = [ [[package]] name = "fabro-retro" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "chrono", @@ -2098,7 +2098,7 @@ dependencies = [ [[package]] name = "fabro-sandbox" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "async-trait", @@ -2138,7 +2138,7 @@ dependencies = [ [[package]] name = "fabro-server" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "async-trait", @@ -2219,7 +2219,7 @@ dependencies = [ [[package]] name = "fabro-slack" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "fabro-http", "fabro-interview", @@ -2240,18 +2240,18 @@ dependencies = [ [[package]] name = "fabro-spa" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "rust-embed", ] [[package]] name = "fabro-static" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" [[package]] name = "fabro-store" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "async-trait", "bytes", @@ -2278,7 +2278,7 @@ dependencies = [ [[package]] name = "fabro-telemetry" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "base64", @@ -2304,7 +2304,7 @@ dependencies = [ [[package]] name = "fabro-template" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "fabro-util", @@ -2316,7 +2316,7 @@ dependencies = [ [[package]] name = "fabro-test" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "assert_cmd", "axum", @@ -2339,7 +2339,7 @@ dependencies = [ [[package]] name = "fabro-tracker" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "async-trait", @@ -2353,7 +2353,7 @@ dependencies = [ [[package]] name = "fabro-types" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "chrono", "clap", @@ -2374,7 +2374,7 @@ dependencies = [ [[package]] name = "fabro-util" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "console 0.15.11", @@ -2394,7 +2394,7 @@ dependencies = [ [[package]] name = "fabro-validate" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "fabro-graphviz", "fabro-model", @@ -2404,7 +2404,7 @@ dependencies = [ [[package]] name = "fabro-vault" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "chrono", "fabro-types", @@ -2416,7 +2416,7 @@ dependencies = [ [[package]] name = "fabro-workflow" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "assert_cmd", @@ -7120,7 +7120,7 @@ dependencies = [ [[package]] name = "twin-github" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "axum", "base64", @@ -7139,7 +7139,7 @@ dependencies = [ [[package]] name = "twin-openai" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" dependencies = [ "anyhow", "async-stream", diff --git a/Cargo.toml b/Cargo.toml index f7897daf1..52d5940be 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,7 +5,7 @@ resolver = "2" [workspace.package] edition = "2021" -version = "0.223.0-nightly.0" +version = "0.224.0-nightly.0" license = "MIT" [workspace.dependencies] diff --git a/apps/fabro-web/app/components/stage-sidebar.tsx b/apps/fabro-web/app/components/stage-sidebar.tsx index 36d1d5d0e..e2cfe88f8 100644 --- a/apps/fabro-web/app/components/stage-sidebar.tsx +++ b/apps/fabro-web/app/components/stage-sidebar.tsx @@ -11,7 +11,7 @@ import { } from "@heroicons/react/24/solid"; import { Bars3BottomLeftIcon, DocumentTextIcon, MapIcon } from "@heroicons/react/24/outline"; import { formatDurationSecs } from "../lib/format"; -import { ACTIVE_STAGE_STATES } from "../lib/stage-sidebar"; +import { ACTIVE_STAGE_STATES, formatStageLabel } from "../lib/stage-sidebar"; import { useTickingNow } from "../lib/time"; export interface Stage { @@ -19,7 +19,8 @@ export interface Stage { name: string; status: StageState; duration: string; - dotId?: string; + nodeId: string; + visit: number; } export const statusConfig: Record; color: string }> = { @@ -97,7 +98,7 @@ export function StageSidebar({ stages, runId, selectedStageId, activeLink }: Sta }`} > - {stage.name} + {formatStageLabel(stage)} {stageDuration(stage)} @@ -153,4 +154,4 @@ export function StageSidebar({ stages, runId, selectedStageId, activeLink }: Sta ); -} \ No newline at end of file +} diff --git a/apps/fabro-web/app/lib/run-events.test.tsx b/apps/fabro-web/app/lib/run-events.test.tsx index c16eb20f9..a86e877e5 100644 --- a/apps/fabro-web/app/lib/run-events.test.tsx +++ b/apps/fabro-web/app/lib/run-events.test.tsx @@ -51,14 +51,14 @@ describe("queryKeysForRunEvent", () => { }); test("stage.retrying invalidates stages, billing, events, and stage turns", () => { - expect(queryKeysForRunEvent("run-1", "stage.retrying", "stage-7")).toEqual([ + expect(queryKeysForRunEvent("run-1", "stage.retrying", "verify@2")).toEqual([ queryKeys.runs.stages("run-1"), queryKeys.runs.billing("run-1"), queryKeys.runs.events("run-1", 1000), queryKeys.runs.graph("run-1", "LR"), queryKeys.runs.graph("run-1", "TB"), queryKeys.runs.detail("run-1"), - queryKeys.runs.stageTurns("run-1", "stage-7"), + queryKeys.runs.stageTurns("run-1", "verify@2"), ]); }); }); @@ -184,6 +184,63 @@ describe("subscribeToRunEvents", () => { coordinator.close(); }); + test("envelope with suffixed stage_id invalidates stageTurns(runId, stageId)", async () => { + const source = new FakeEventSource(); + const keys: string[] = []; + const coordinator = createCoordinator(() => source); + const cleanup = subscribeToRunEvents( + "run-stage", + (key) => { + keys.push(key); + return Promise.resolve(); + }, + () => source, + { debounceMs: 0, coordinator }, + ); + + await waitFor(() => source.onmessage !== null); + source.emit({ + event: "stage.retrying", + run_id: "run-stage", + stage_id: "verify@2", + node_id: "verify", + }); + + expect(keys).toContain(queryKeys.runs.stageTurns("run-stage", "verify@2")); + expect(keys).toContain(queryKeys.runs.stages("run-stage")); + expect(keys).toContain(queryKeys.runs.events("run-stage", 1000)); + expect(keys).toContain(queryKeys.runs.graph("run-stage", "LR")); + expect(keys).toContain(queryKeys.runs.detail("run-stage")); + expect(keys).not.toContain(queryKeys.runs.stageTurns("run-stage", "verify")); + + cleanup(); + coordinator.close(); + }); + + test("falls back to node_id when an event has no stage_id", async () => { + const source = new FakeEventSource(); + const keys: string[] = []; + const coordinator = createCoordinator(() => source); + const cleanup = subscribeToRunEvents( + "run-stage-node", + (key) => { + keys.push(key); + return Promise.resolve(); + }, + () => source, + { debounceMs: 0, coordinator }, + ); + + await waitFor(() => source.onmessage !== null); + source.emit({ event: "stage.started", run_id: "run-stage-node", node_id: "verify" }); + + expect(keys).toContain(queryKeys.runs.stageTurns("run-stage-node", "verify")); + expect(keys).toContain(queryKeys.runs.stages("run-stage-node")); + + cleanup(); + coordinator.close(); + }); + test("fallback malformed events are ignored and StrictMode-style cleanup does not underflow", () => { const firstSource = new FakeEventSource(); const secondSource = new FakeEventSource(); @@ -251,4 +308,4 @@ async function waitFor(condition: () => boolean, timeoutMs = 200) { await new Promise((resolve) => setTimeout(resolve, 2)); } throw new Error("condition did not become true before timeout"); -} \ No newline at end of file +} diff --git a/apps/fabro-web/app/lib/run-events.ts b/apps/fabro-web/app/lib/run-events.ts index c06c251a0..944303d2a 100644 --- a/apps/fabro-web/app/lib/run-events.ts +++ b/apps/fabro-web/app/lib/run-events.ts @@ -19,6 +19,7 @@ interface RunEventPayload extends EventPayload { event?: string; run_id?: string; node_id?: string; + stage_id?: string; properties?: Record; } @@ -172,6 +173,7 @@ function resyncKeysForRun(runId: string) { } function stageIdFromPayload(payload: RunEventPayload): string | undefined { + if (typeof payload.stage_id === "string") return payload.stage_id; if (typeof payload.node_id === "string") return payload.node_id; const nodeId = payload.properties?.node_id; return typeof nodeId === "string" ? nodeId : undefined; diff --git a/apps/fabro-web/app/lib/stage-sidebar.test.ts b/apps/fabro-web/app/lib/stage-sidebar.test.ts new file mode 100644 index 000000000..9e8891645 --- /dev/null +++ b/apps/fabro-web/app/lib/stage-sidebar.test.ts @@ -0,0 +1,172 @@ +import { describe, expect, test } from "bun:test"; +import type { PaginatedRunStageList, StageState } from "@qltysh/fabro-api-client"; + +import type { Stage } from "../components/stage-sidebar"; +import { aggregateGraphNodeStatus, formatStageLabel, mapRunStagesToSidebarStages } from "./stage-sidebar"; + +function makeStage(nodeId: string, visit: number, status: StageState): Stage { + return { + id: `${nodeId}@${visit}`, + name: nodeId, + nodeId, + visit, + status, + duration: "--", + }; +} + +describe("mapRunStagesToSidebarStages", () => { + test("maps two visits of the same node to distinct sidebar entries", () => { + const stages: PaginatedRunStageList = { + data: [ + { + id: "apply-changes@1", + name: "Apply Changes", + status: "succeeded", + duration_secs: 12.5, + node_id: "apply", + visit: 1, + }, + { + id: "apply-changes@2", + name: "Apply Changes", + status: "running", + node_id: "apply", + visit: 2, + }, + ], + meta: { has_more: false }, + }; + + const result = mapRunStagesToSidebarStages(stages); + expect(result).toHaveLength(2); + + expect(result[0].id).toBe("apply-changes@1"); + expect(result[0].nodeId).toBe("apply"); + expect(result[0].visit).toBe(1); + expect(formatStageLabel(result[0])).toBe("Apply Changes"); + + expect(result[1].id).toBe("apply-changes@2"); + expect(result[1].nodeId).toBe("apply"); + expect(result[1].visit).toBe(2); + expect(formatStageLabel(result[1])).toBe("Apply Changes (2)"); + }); + + test("filters by node_id (suffixed start@1 / exit@1 are still hidden)", () => { + const stages: PaginatedRunStageList = { + data: [ + { + id: "start@1", + name: "start", + status: "succeeded", + node_id: "start", + visit: 1, + }, + { + id: "verify@1", + name: "verify", + status: "succeeded", + node_id: "verify", + visit: 1, + }, + { + id: "exit@1", + name: "exit", + status: "succeeded", + node_id: "exit", + visit: 1, + }, + ], + meta: { has_more: false }, + }; + + const result = mapRunStagesToSidebarStages(stages); + expect(result.map((s) => s.id)).toEqual(["verify@1"]); + }); + + test("missing duration renders as '--'", () => { + const stages: PaginatedRunStageList = { + data: [ + { + id: "verify@1", + name: "verify", + status: "running", + node_id: "verify", + visit: 1, + }, + ], + meta: { has_more: false }, + }; + + expect(mapRunStagesToSidebarStages(stages)[0].duration).toBe("--"); + }); +}); + +describe("aggregateGraphNodeStatus", () => { + test("(failed, running) renders as running and clicks open the latest visit", () => { + const result = aggregateGraphNodeStatus([ + makeStage("verify", 1, "failed"), + makeStage("verify", 2, "running"), + ]); + expect(result.get("verify")).toEqual({ + displayStatus: "running", + latestStageId: "verify@2", + }); + }); + + test("(failed, succeeded) renders as succeeded — failure-then-fix shows healed", () => { + const result = aggregateGraphNodeStatus([ + makeStage("verify", 1, "failed"), + makeStage("verify", 2, "succeeded"), + ]); + expect(result.get("verify")).toEqual({ + displayStatus: "succeeded", + latestStageId: "verify@2", + }); + }); + + test("(succeeded, failed) renders as failed and clicks open the latest visit", () => { + const result = aggregateGraphNodeStatus([ + makeStage("verify", 1, "succeeded"), + makeStage("verify", 2, "failed"), + ]); + expect(result.get("verify")).toEqual({ + displayStatus: "failed", + latestStageId: "verify@2", + }); + }); + + test("(running, retrying) — latest active wins", () => { + const result = aggregateGraphNodeStatus([ + makeStage("verify", 1, "running"), + makeStage("verify", 2, "retrying"), + ]); + expect(result.get("verify")).toEqual({ + displayStatus: "retrying", + latestStageId: "verify@2", + }); + }); + + test("orders by visit even when input is shuffled", () => { + const result = aggregateGraphNodeStatus([ + makeStage("verify", 2, "running"), + makeStage("verify", 1, "failed"), + ]); + expect(result.get("verify")?.latestStageId).toBe("verify@2"); + }); + + test("single visit per node is unaffected", () => { + const result = aggregateGraphNodeStatus([ + makeStage("plan", 1, "succeeded"), + makeStage("apply", 1, "running"), + ]); + expect(result.get("plan")).toEqual({ + displayStatus: "succeeded", + latestStageId: "plan@1", + }); + expect(result.get("apply")).toEqual({ + displayStatus: "running", + latestStageId: "apply@1", + }); + }); +}); diff --git a/apps/fabro-web/app/lib/stage-sidebar.ts b/apps/fabro-web/app/lib/stage-sidebar.ts index 8ddccb267..84a5de2de 100644 --- a/apps/fabro-web/app/lib/stage-sidebar.ts +++ b/apps/fabro-web/app/lib/stage-sidebar.ts @@ -19,18 +19,65 @@ export const SUCCEEDED_STAGE_STATES: ReadonlySet = new Set([ StageState.PARTIALLY_SUCCEEDED, ]); +/** + * Display label for a stage. Suffixes `(N)` for visits > 1 so a looped node + * (e.g. `verify`) renders as `verify`, `verify (2)`, `verify (3)` in the + * sidebar and stage header. + */ +export function formatStageLabel(stage: { name: string; visit: number }): string { + return stage.visit > 1 ? `${stage.name} (${stage.visit})` : stage.name; +} + export function mapRunStagesToSidebarStages( stagesResult: PaginatedRunStageList | null | undefined, ): Stage[] { return (stagesResult?.data ?? []) - .filter((stage) => isVisibleStage(stage.id)) + .filter((stage) => isVisibleStage(stage.node_id)) .map((stage) => ({ id: stage.id, name: stage.name, - dotId: stage.dot_id ?? stage.id, + nodeId: stage.node_id, + visit: stage.visit, status: stage.status, duration: stage.duration_secs != null ? formatDurationSecs(stage.duration_secs) : "--", })); } + +/** + * Aggregate per-node display state for the workflow graph. + * + * Status policy: if any visit is active (running/retrying), the node renders + * that active state (latest active visit wins). Otherwise the node renders + * the latest visit's terminal state. The click target is always the latest + * visit's stageId. + */ +export function aggregateGraphNodeStatus(stages: readonly Stage[]): Map< + string, + { displayStatus: StageState; latestStageId: string } +> { + // Single pass per nodeId: track the visit with the highest `visit` overall + // (drives click target + terminal status) and the highest-visit *active* + // stage (drives display when any visit is in flight). + const latest = new Map(); + const latestActive = new Map(); + for (const stage of stages) { + const prevLatest = latest.get(stage.nodeId); + if (!prevLatest || stage.visit > prevLatest.visit) { + latest.set(stage.nodeId, stage); + } + if (ACTIVE_STAGE_STATES.has(stage.status)) { + const prevActive = latestActive.get(stage.nodeId); + if (!prevActive || stage.visit > prevActive.visit) { + latestActive.set(stage.nodeId, stage); + } + } + } + const result = new Map(); + for (const [nodeId, latestStage] of latest) { + const display = latestActive.get(nodeId) ?? latestStage; + result.set(nodeId, { displayStatus: display.status, latestStageId: latestStage.id }); + } + return result; +} diff --git a/apps/fabro-web/app/routes/run-overview.tsx b/apps/fabro-web/app/routes/run-overview.tsx index bf0c0dd93..48066eae0 100644 --- a/apps/fabro-web/app/routes/run-overview.tsx +++ b/apps/fabro-web/app/routes/run-overview.tsx @@ -3,7 +3,6 @@ import { useNavigate, useParams } from "react-router"; import { graphTheme } from "../lib/graph-theme"; import { useRun, useRunGraph, useRunStages } from "../lib/queries"; import { StageSidebar } from "../components/stage-sidebar"; -import type { Stage } from "../components/stage-sidebar"; import { GRAPH_DEFAULT_ZOOM_INDEX, GRAPH_ZOOM_STEPS, @@ -13,6 +12,7 @@ import { EmptyState } from "../components/state"; import { ACTIVE_STAGE_STATES, SUCCEEDED_STAGE_STATES, + aggregateGraphNodeStatus, mapRunStagesToSidebarStages, } from "../lib/stage-sidebar"; @@ -54,27 +54,27 @@ export default function RunOverview() { const inner = innerRef.current; if (!inner || !graphSvg) return; - let cancelled = false; - (async () => { - if (cancelled) return; inner.innerHTML = graphSvg; const svg = inner.querySelector("svg"); if (!svg) return; svgRef.current = svg; const gt = graphTheme; - const runningDotIds = new Set( - stages.filter((s: Stage) => ACTIVE_STAGE_STATES.has(s.status)).map((s: Stage) => s.dotId ?? s.id), - ); - const failedDotIds = new Set( - stages.filter((s: Stage) => s.status === "failed").map((s: Stage) => s.dotId ?? s.id), - ); - const completedDotIds = new Set( - stages.filter((s: Stage) => SUCCEEDED_STAGE_STATES.has(s.status)).map((s: Stage) => s.dotId ?? s.id), - ); - const dotIdToStageId = new Map( - stages.map((s: Stage) => [s.dotId ?? s.id, s.id]), - ); + const aggregated = aggregateGraphNodeStatus(stages); + const runningDotIds = new Set(); + const failedDotIds = new Set(); + const completedDotIds = new Set(); + const dotIdToStageId = new Map(); + for (const [nodeId, { displayStatus, latestStageId }] of aggregated) { + dotIdToStageId.set(nodeId, latestStageId); + if (ACTIVE_STAGE_STATES.has(displayStatus)) { + runningDotIds.add(nodeId); + } else if (displayStatus === "failed") { + failedDotIds.add(nodeId); + } else if (SUCCEEDED_STAGE_STATES.has(displayStatus)) { + completedDotIds.add(nodeId); + } + } const ns = "http://www.w3.org/2000/svg"; for (const group of svg.querySelectorAll(".node")) { @@ -148,8 +148,6 @@ export default function RunOverview() { } } } - })(); - return () => { cancelled = true; }; }, [stages, graphSvg, id, navigate, terminalOutcome]); const onPointerDown = useCallback((e: React.PointerEvent) => { diff --git a/apps/fabro-web/app/routes/run-stages.test.ts b/apps/fabro-web/app/routes/run-stages.test.ts index d10b5b2d2..704c5cac8 100644 --- a/apps/fabro-web/app/routes/run-stages.test.ts +++ b/apps/fabro-web/app/routes/run-stages.test.ts @@ -1,6 +1,7 @@ import { describe, expect, test } from "bun:test"; +import type { EventEnvelope } from "@qltysh/fabro-api-client"; -import { isSafeMarkdownHref } from "./run-stages"; +import { isSafeMarkdownHref, turnsFromEvents } from "./run-stages"; describe("isSafeMarkdownHref", () => { test("rejects protocol-relative URLs", () => { @@ -15,3 +16,96 @@ describe("isSafeMarkdownHref", () => { expect(isSafeMarkdownHref("mailto:test@example.com")).toBe(true); }); }); + +function makeEnvelope(overrides: Partial): EventEnvelope { + return { + seq: 1, + id: "evt", + ts: "2026-01-01T00:00:00Z", + run_id: "run-1", + event: "stage.prompt", + ...overrides, + } as EventEnvelope; +} + +describe("turnsFromEvents", () => { + test("filters events by stage_id (verify@1 vs verify@2 do not cross-contaminate)", () => { + const events: EventEnvelope[] = [ + makeEnvelope({ + seq: 1, + event: "stage.prompt", + stage_id: "verify@1", + node_id: "verify", + properties: { text: "first visit prompt" }, + }), + makeEnvelope({ + seq: 2, + event: "stage.prompt", + stage_id: "verify@2", + node_id: "verify", + properties: { text: "second visit prompt" }, + }), + makeEnvelope({ + seq: 3, + event: "agent.message", + stage_id: "verify@1", + node_id: "verify", + properties: { text: "first visit reply" }, + }), + makeEnvelope({ + seq: 4, + event: "agent.message", + stage_id: "verify@2", + node_id: "verify", + properties: { text: "second visit reply" }, + }), + ]; + + const firstVisit = turnsFromEvents(events, "verify@1"); + expect(firstVisit).toEqual([ + { kind: "system", content: "first visit prompt" }, + { kind: "assistant", content: "first visit reply" }, + ]); + + const secondVisit = turnsFromEvents(events, "verify@2"); + expect(secondVisit).toEqual([ + { kind: "system", content: "second visit prompt" }, + { kind: "assistant", content: "second visit reply" }, + ]); + }); + + test("command turn carries the requested stage_id, no @1 fallback", () => { + const events: EventEnvelope[] = [ + makeEnvelope({ + seq: 1, + event: "command.started", + stage_id: "verify@2", + node_id: "verify", + properties: { script: "echo hi", language: "shell" }, + }), + makeEnvelope({ + seq: 2, + event: "command.completed", + stage_id: "verify@2", + node_id: "verify", + properties: { + stdout: "hi", + stderr: "", + exit_code: 0, + duration_ms: 5, + termination: "exited", + }, + }), + ]; + + const turns = turnsFromEvents(events, "verify@2"); + expect(turns).toHaveLength(1); + const turn = turns[0]; + expect(turn.kind).toBe("command"); + if (turn.kind === "command") { + expect(turn.stageId).toBe("verify@2"); + expect(turn.script).toBe("echo hi"); + expect(turn.running).toBe(false); + } + }); +}); diff --git a/apps/fabro-web/app/routes/run-stages.tsx b/apps/fabro-web/app/routes/run-stages.tsx index ab70b98cd..eac535144 100644 --- a/apps/fabro-web/app/routes/run-stages.tsx +++ b/apps/fabro-web/app/routes/run-stages.tsx @@ -42,7 +42,7 @@ import { CopyButton } from "../components/ui"; import { formatDurationSecs } from "../lib/format"; import { useTickingNow } from "../lib/time"; import { fetchRunCommandLog, useRunEventsList, useRunStageTurns, useRunStages } from "../lib/queries"; -import { mapRunStagesToSidebarStages } from "../lib/stage-sidebar"; +import { ACTIVE_STAGE_STATES, formatStageLabel, mapRunStagesToSidebarStages } from "../lib/stage-sidebar"; import { getNumber, getString, type UnknownRecord } from "../lib/unknown"; import { CommandOutputStream, @@ -69,8 +69,8 @@ function readTermination(props: UnknownRecord): CommandTermination { return CommandTermination.EXITED; } -function turnsFromEvents(events: EventEnvelope[], stageId: string): TurnType[] { - const stageEvents = events.filter((e) => e.node_id === stageId); +export function turnsFromEvents(events: EventEnvelope[], stageId: string): TurnType[] { + const stageEvents = events.filter((e) => e.stage_id === stageId); const turns: TurnType[] = []; // Collect tool pairs: started → completed const pendingTools = new Map(); @@ -115,7 +115,7 @@ function turnsFromEvents(events: EventEnvelope[], stageId: string): TurnType[] { } case "command.started": { pendingCommand = { - stageId: e.stage_id ?? `${stageId}@1`, + stageId, script: getString(props, "script") ?? "", language: getString(props, "language") ?? "shell", }; @@ -124,7 +124,7 @@ function turnsFromEvents(events: EventEnvelope[], stageId: string): TurnType[] { case "command.completed": { turns.push({ kind: "command", - stageId: pendingCommand?.stageId ?? e.stage_id ?? `${stageId}@1`, + stageId: pendingCommand?.stageId ?? stageId, script: pendingCommand?.script ?? "", language: pendingCommand?.language ?? "shell", stdout: getString(props, "stdout") ?? "", @@ -610,7 +610,7 @@ export default function RunStages() { () => mapTurns(turnsQuery.data, eventsQuery.data, selectedStage?.id), [eventsQuery.data, selectedStage?.id, turnsQuery.data], ); - const isRunning = selectedStage?.status === "running"; + const isActive = selectedStage ? ACTIVE_STAGE_STATES.has(selectedStage.status) : false; if (!id || !stages.length) { return ( @@ -632,11 +632,13 @@ export default function RunStages() {
- -

{selectedStage.name}

+ +

+ {formatStageLabel(selectedStage)} +

diff --git a/docs/public/api-reference/fabro-api.yaml b/docs/public/api-reference/fabro-api.yaml index 18ccca73c..17ab1ae22 100644 --- a/docs/public/api-reference/fabro-api.yaml +++ b/docs/public/api-reference/fabro-api.yaml @@ -6174,7 +6174,7 @@ components: example: 3501.0 BillingStageRef: - description: Reference to a billing stage. + description: Reference to a workflow node in a billing stage row. type: object required: - id @@ -6334,11 +6334,13 @@ components: - id - name - status + - node_id + - visit properties: id: type: string - description: Unique stage identifier within the run. - example: propose-changes + description: StageId in "node_id@visit" form, e.g. verify@2. + example: verify@2 name: type: string description: Human-readable stage name. @@ -6349,10 +6351,16 @@ components: type: number description: Time spent in this stage, in seconds. example: 154.0 - dot_id: + node_id: type: string - description: Node identifier in the Graphviz graph source. - example: propose + description: Node id in the workflow graph; multiple stages with different visits share the same node_id. + example: verify + visit: + type: integer + format: uint32 + minimum: 1 + description: 1-based visit count; bumped each time the workflow re-enters this node. + example: 2 started_at: type: ["string", "null"] format: date-time @@ -6627,7 +6635,7 @@ components: # ── Billing Schemas ────────────────────────────────────────────────── RunBillingStage: - description: Token counts and billed totals for a single stage within a run. + description: Token counts and billed totals for one workflow node within a run. Rows are grouped by node; billing and runtime sum every visit of that node. type: object required: - stage @@ -6638,7 +6646,7 @@ components: stage: $ref: "#/components/schemas/BillingStageRef" model: - description: Model used for this stage; null for non-LLM stages. + description: Latest usage-bearing visit model for this node; null when no visit used an LLM model. oneOf: - $ref: "#/components/schemas/ModelReference" - type: "null" @@ -6646,7 +6654,7 @@ components: $ref: "#/components/schemas/BilledTokenCounts" runtime_secs: type: number - description: Wall-clock runtime in seconds. + description: Wall-clock runtime in seconds, summed across every visit of this node. example: 154.0 started_at: type: ["string", "null"] @@ -6717,7 +6725,7 @@ components: $ref: "#/components/schemas/ModelReference" stages: type: integer - description: Number of stages that used this model. + description: Number of usage-bearing stage visits that used this model. example: 2 billing: $ref: "#/components/schemas/BilledTokenCounts" @@ -6732,7 +6740,7 @@ components: properties: stages: type: array - description: Per-stage billing breakdown. + description: Per-node billing breakdown. Each row sums billing and runtime across all visits of that node. items: $ref: "#/components/schemas/RunBillingStage" totals: @@ -8352,4 +8360,4 @@ components: login: type: string description: User's login identifier (e.g. GitHub username). - example: octocat \ No newline at end of file + example: octocat diff --git a/docs/superpowers/plans/2026-05-04-cross-tab-sse-coordination.md b/docs/superpowers/plans/2026-05-04-cross-tab-sse-coordination.md new file mode 100644 index 000000000..115ff9326 --- /dev/null +++ b/docs/superpowers/plans/2026-05-04-cross-tab-sse-coordination.md @@ -0,0 +1,106 @@ +# Cross-Tab SSE Coordination Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Coordinate Fabro web SSE subscriptions across tabs so one browser profile/origin opens at most one UI-owned `/api/v1/attach` EventSource in steady state when `BroadcastChannel` is available; brief overlap during election/takeover is tolerated and deduped. + +**Architecture:** Add a browser-side SSE coordinator that elects one tab as leader, has that leader own the global EventSource, and broadcasts parsed run events to follower tabs over `BroadcastChannel`. Existing board and run-detail invalidation logic becomes a consumer of that global event feed, with the current per-tab SSE behavior preserved as a compatibility fallback. + +**Tech Stack:** React, SWR, browser `BroadcastChannel`, browser `EventSource`, Bun tests, existing Fabro web API query keys. + +--- + +## Summary + +Build a browser-side SSE coordinator so Fabro web opens at most one `/api/v1/attach` EventSource per origin/browser profile in steady state when `BroadcastChannel` is available. Temporary duplicate leaders may exist during election/takeover, but event dedupe and generation checks make the overlap harmless and short-lived. The global stream becomes a shared cache-invalidation feed for both the runs board and run detail pages. No server API, OpenAPI, or Rust streaming contract changes are part of v1. + +This supersedes the earlier web SSE limitation documented in `docs/plans/2026-04-19-002-feat-web-ui-lifecycle-actions-plan.md`: the old shared hook was code reuse only; this plan adds actual socket deduplication. + +## Implementation Changes + +- Add `apps/fabro-web/app/lib/cross-tab-sse.ts`. + - Export `subscribeToCrossTabSse(...)` with the same invalidation style as `subscribeToSharedEventSource`, a `resyncKeys` callback for gap recovery, and a `fallbackSubscribe` callback used when cross-tab coordination is unavailable. + - Use `BroadcastChannel` name `fabro:sse:v1`. + - Generate `tabId` with `crypto.randomUUID()` and a safe random fallback. + - Open one leader-owned `EventSource` to `queryKeys.system.attach()` (`/api/v1/attach`) in steady state. + - Leader dispatches each parsed `EventEnvelope` locally and broadcasts it to follower tabs. + - Followers do not open EventSource while a valid visible leader heartbeat exists; lower lexical `tabId` does not preempt a healthy visible leader. + +- Implement leader election in the cross-tab module. + - Constants: `HEARTBEAT_MS = 1000`, `LEADER_STALE_MS = 4000`, `ELECTION_JITTER_MS = 150`. + - Messages: `hello`, `heartbeat`, `candidate`, `leader-changed`, `release`, `resync`, `event`. + - Define a typed message union. Every message includes `type`, `version: 1`, `tabId`, and `sentAt`. + - `heartbeat`: `{ type, version, tabId, sentAt, leaderId, generation, visibility }`. + - `candidate`: `{ type, version, tabId, sentAt, candidateId: tabId, candidateGeneration, visibility, observedLeaderId, observedGeneration, reason }`, where `reason` is `"hidden-leader" | "stale-leader" | "release" | "no-leader"`. + - `leader-changed`: `{ type, version, tabId, sentAt, leaderId, generation, visibility }`. + - `release`: `{ type, version, tabId, sentAt, leaderId, generation }`. + - `resync`: `{ type, version, tabId, sentAt, leaderId, generation, reason }`. + - `event`: `{ type, version, tabId, sentAt, leaderId, generation, payload }`. + - Use the candidacy phase for all leadership changes: hidden-leader takeover, stale-leader recovery, leader release, and no-leader startup. + - A candidate sets `candidateGeneration = observedGeneration + 1`, broadcasts `candidate`, waits jitter, and opens EventSource only if no higher-priority candidate for the same `candidateGeneration` appears. + - Candidate priority is election-scoped: visible candidates outrank hidden candidates; for equal visibility, lower lexical `candidateId` wins. This priority resolves elections and same-generation split brain only; it is not a reason to preempt a fresh visible leader. + - When a visible follower observes a fresh hidden leader heartbeat, it enters candidacy with `reason: "hidden-leader"`. + - When tabs detect a stale leader, leader release, or no known leader, they enter the same candidacy flow with the matching `reason`. + - If two visible candidates race for the same observed leader/generation, the lower lexical `candidateId` wins. + - Current leaders release when they observe a candidate whose `observedLeaderId` matches their `leaderId` and whose `observedGeneration` is current or newer. + - If same-generation split brain still occurs, lower-priority leaders release when they observe a same-generation higher-priority leader heartbeat or `leader-changed`. + - Each new leader uses `candidateGeneration`, broadcasts `leader-changed`, and followers ignore heartbeats/events from non-current leaders or stale generations. + - Hidden leader keeps the stream only when no visible candidate takes over. + - On `pagehide`/last local unsubscribe, a leader closes EventSource and broadcasts `release`. + - Brief split brain is tolerated; dedupe events by `payload.id`, falling back to `${run_id}:${seq}:${event}`. + - Keep dedupe bounded with a recent-event cache: max 1000 IDs and 5 minute TTL. Evict oldest entries when the max is exceeded and prune expired entries during event handling. Duplicate invalidations after eviction are acceptable; unbounded growth is not. + +- Migrate consumers. + - `apps/fabro-web/app/lib/board-events.ts`: subscribe through the cross-tab global stream; keep existing board event allowlist. + - `apps/fabro-web/app/lib/run-events.ts`: subscribe through the same global stream, filter by `payload.run_id === runId`, and reuse `queryKeysForRunEvent`. + - In coordinated mode, run detail pages stay subscribed while mounted, including terminal runs, so post-terminal archive/unarchive changes can reconcile live. + - Do not close the global stream on `run.completed` / `run.failed`; terminal events only invalidate run-scoped keys. + - Keep `subscribeToSharedEventSource` in `apps/fabro-web/app/lib/sse.ts` for fallback and existing local sharing behavior. + +- Gap and fallback behavior. + - If `BroadcastChannel` is unavailable or throws, call each subscriber's `fallbackSubscribe`. + - Board fallback uses the existing global `/api/v1/attach` path. + - Run detail fallback preserves the existing run-scoped `/api/v1/runs/:id/attach` path, so the old terminal-tab stale limitation remains only in fallback mode. + - Do not add replay to `/api/v1/attach`. + - On leader takeover, stale leader timeout, leader release, and new leader generation, broadcast `resync` or `leader-changed` so every tab with active local subscriptions runs its own `resyncKeys`. + - On `visibilitychange` back to visible without leadership change, run only that tab's local `resyncKeys`; do not broadcast cross-tab resync. + - Board `resyncKeys`: `queryKeys.boards.runs()`. + - Run `resyncKeys`: detail, files, billing, stages, events, LR graph, TB graph, and questions for that run. + +## Tests + +- Add `apps/fabro-web/app/lib/cross-tab-sse.test.ts` with fake `BroadcastChannel`, fake `EventSource`, and fake timers. + - One leader opens `/api/v1/attach`; followers open no EventSource. + - Leader broadcasts an event and all local subscribers receive invalidations. + - Run subscribers ignore events for other `run_id` values. + - Board and run subscriptions can coexist on the same global stream. + - Temporary duplicate leaders are allowed only during election/takeover and converge back to one leader. + - With fresh hidden-leader heartbeats, a visible tab broadcasts candidacy, hidden leader closes, visible tab opens `/api/v1/attach`, and followers resync. + - Two visible candidates racing for the same hidden leader resolve to the lower lexical `candidateId`. + - Two tabs detect the same stale leader simultaneously; only the winning candidate opens `/api/v1/attach` after jitter. + - Same-generation split brain converges to one leader by visibility, then lexical `tabId`. + - A fresh visible leader exists; a new visible follower with a lower lexical `tabId` joins and does not take leadership. + - A hidden leader does not release for a candidate that references an old `observedLeaderId` or stale `observedGeneration`. + - Stale heartbeat triggers takeover and every active tab calls its own `resyncKeys`. + - A follower tab calls its own board/run `resyncKeys` after another tab becomes leader. + - Duplicate event IDs are ignored. + - The recent-event dedupe cache evicts by TTL and max-size bound; duplicate invalidation may recur only after eviction. + - Stale heartbeat/event messages from an old leader/generation are ignored after takeover. + - Last unsubscribe closes leader EventSource and sends `release`. + - Missing/broken BroadcastChannel uses the per-tab fallback. + +- Update existing tests. + - `apps/fabro-web/app/lib/board-events.test.tsx`: assert coordinated mode uses `/api/v1/attach` once and fallback preserves current behavior. + - `apps/fabro-web/app/lib/run-events.test.tsx`: assert coordinated mode filters by `run_id`; fallback preserves current `/runs/:id/attach` behavior; terminal run events do not close the global coordinator; terminal detail tabs still receive archive/unarchive invalidations while mounted in coordinated mode. + +- Verification commands. + - `cd apps/fabro-web && bun test app/lib/cross-tab-sse.test.ts app/lib/board-events.test.tsx app/lib/run-events.test.tsx` + - `cd apps/fabro-web && bun run typecheck` + - Optional manual check: open 8 run-detail tabs for active runs; after election settles and with BroadcastChannel available, Chrome should show one active `/api/v1/attach` EventStream across the participating tabs and no UI-created `/api/v1/runs/:id/attach` streams. + +## Assumptions + +- Treat SSE as live cache invalidation, not an exact gapless event log. +- Preserve current run-specific attach endpoint for CLI, API clients, and fallback only. +- Do not touch unrelated dirty worktree files; intentionally replace or adapt the earlier interrupted SSE test edits as part of the new test suite. +- No docs or public API updates are required for v1 because this is an internal web transport change. diff --git a/lib/crates/fabro-dump/src/lib.rs b/lib/crates/fabro-dump/src/lib.rs index 69a670565..5b2da9403 100644 --- a/lib/crates/fabro-dump/src/lib.rs +++ b/lib/crates/fabro-dump/src/lib.rs @@ -66,18 +66,13 @@ impl RunDump { entries.push(RunDumpEntry::text("graph.fabro", graph_source.clone())); } - let mut stages: Vec<_> = state.iter_stages().collect(); + let stages: Vec<_> = state.iter_stages().collect(); if stages.len() > MAX_STAGES_IN_DUMP { bail!( "run dump supports at most {MAX_STAGES_IN_DUMP} stages with the current path prefix width (got {})", stages.len() ); } - stages.sort_by(|(left_id, left), (right_id, right)| { - left.first_event_seq - .cmp(&right.first_event_seq) - .then_with(|| left_id.cmp(right_id)) - }); let mut stage_ranks = HashMap::new(); for (index, (stage_id, _)) in stages.iter().enumerate() { diff --git a/lib/crates/fabro-server/src/demo/mod.rs b/lib/crates/fabro-server/src/demo/mod.rs index a51cb4050..29b975160 100644 --- a/lib/crates/fabro-server/src/demo/mod.rs +++ b/lib/crates/fabro-server/src/demo/mod.rs @@ -784,9 +784,10 @@ mod runs { RunNamespace, RunPrepareSettings, RunSandboxSettings, }; use fabro_types::settings::{InterpString, ProjectNamespace, WorkflowNamespace}; - use fabro_types::{RunId, WorkflowSettings}; + use fabro_types::{RunId, StageId, WorkflowSettings}; use super::ts; + use crate::server::run_stage_from_stage_id; fn labels(entries: &[(&str, &str)]) -> HashMap { entries @@ -1181,38 +1182,41 @@ mod runs { pub(super) fn stages() -> Vec { vec![ - RunStage { - id: "detect-drift".into(), - name: "Detect Drift".into(), - status: StageState::Succeeded, - duration_secs: Some(72.0), - dot_id: Some("detect".into()), - started_at: None, - }, - RunStage { - id: "propose-changes".into(), - name: "Propose Changes".into(), - status: StageState::Succeeded, - duration_secs: Some(154.0), - dot_id: Some("propose".into()), - started_at: None, - }, - RunStage { - id: "review-changes".into(), - name: "Review Changes".into(), - status: StageState::Succeeded, - duration_secs: Some(45.0), - dot_id: Some("review".into()), - started_at: None, - }, - RunStage { - id: "apply-changes".into(), - name: "Apply Changes".into(), - status: StageState::Running, - duration_secs: Some(118.0), - dot_id: Some("apply".into()), - started_at: None, - }, + run_stage_from_stage_id( + &StageId::new("detect-drift", 1), + "Detect Drift", + StageState::Succeeded, + Some(72.0), + None, + ), + run_stage_from_stage_id( + &StageId::new("propose-changes", 1), + "Propose Changes", + StageState::Succeeded, + Some(154.0), + None, + ), + run_stage_from_stage_id( + &StageId::new("review-changes", 1), + "Review Changes", + StageState::Succeeded, + Some(45.0), + None, + ), + run_stage_from_stage_id( + &StageId::new("apply-changes", 1), + "Apply Changes", + StageState::Succeeded, + Some(118.0), + None, + ), + run_stage_from_stage_id( + &StageId::new("apply-changes", 2), + "Apply Changes", + StageState::Running, + None, + None, + ), ] } diff --git a/lib/crates/fabro-server/src/server.rs b/lib/crates/fabro-server/src/server.rs index eb35ef8b0..f86f23e9c 100644 --- a/lib/crates/fabro-server/src/server.rs +++ b/lib/crates/fabro-server/src/server.rs @@ -54,7 +54,7 @@ use fabro_llm::types::{ ContentPart, FinishReason, Message as LlmMessage, Request as LlmRequest, Role, ToolChoice, ToolDefinition, }; -use fabro_model::{BilledModelUsage, BilledTokenCounts, Catalog, ModelTestMode, Provider}; +use fabro_model::{BilledTokenCounts, Catalog, ModelTestMode, Provider}; use fabro_redact::redact_jsonl_line; use fabro_sandbox::daytona::{self, DaytonaSandbox}; use fabro_sandbox::reconnect::reconnect; @@ -536,17 +536,50 @@ pub(crate) struct ResolvedAppStateSettings { pub(crate) manifest_run_settings: std::result::Result, } -fn accumulate_model_billing(entry: &mut ModelBillingTotals, usage: &BilledModelUsage) { - let tokens = usage.tokens(); - entry.stages += 1; - entry.billing.input_tokens += tokens.input_tokens; - entry.billing.output_tokens += tokens.output_tokens; - entry.billing.reasoning_tokens += tokens.reasoning_tokens; - entry.billing.cache_read_tokens += tokens.cache_read_tokens; - entry.billing.cache_write_tokens += tokens.cache_write_tokens; - entry.billing.total_tokens += tokens.total_tokens(); - if let Some(value) = usage.total_usd_micros { - *entry.billing.total_usd_micros.get_or_insert(0) += value; +fn accumulate_billed_token_counts(target: &mut BilledTokenCounts, source: &BilledTokenCounts) { + target.input_tokens += source.input_tokens; + target.output_tokens += source.output_tokens; + target.reasoning_tokens += source.reasoning_tokens; + target.cache_read_tokens += source.cache_read_tokens; + target.cache_write_tokens += source.cache_write_tokens; + target.total_tokens += source.total_tokens; + if let Some(value) = source.total_usd_micros { + *target.total_usd_micros.get_or_insert(0) += value; + } +} + +fn accumulate_billing_rollup( + accumulator: &mut BillingAccumulator, + rollup: &fabro_workflow::ProjectionBillingRollup, +) { + accumulator.total_runs += 1; + accumulator.total_runtime_secs += rollup.runtime_ms as f64 / 1000.0; + for model in &rollup.by_model { + let entry = accumulator + .by_model + .entry(model.model_id.clone()) + .or_default(); + entry.stages += model.stages; + accumulate_billed_token_counts(&mut entry.billing, &model.billing); + } +} + +pub(crate) fn run_stage_from_stage_id( + stage_id: &StageId, + name: impl Into, + status: StageState, + duration_secs: Option, + started_at: Option>, +) -> RunStage { + RunStage { + id: stage_id.to_string(), + name: name.into(), + status, + duration_secs, + node_id: stage_id.node_id().to_string(), + visit: std::num::NonZeroU32::new(stage_id.visit()) + .expect("StageId stores a non-zero visit"), + started_at, } } @@ -2776,9 +2809,9 @@ async fn execute_run_in_process(state: Arc, run_id: RunId) { } } - // Save final checkpoint - let checkpoint = match run_store.state().await { - Ok(state) => state.checkpoint, + // Save final projection + let final_projection = match run_store.state().await { + Ok(state) => Some(state), Err(err) => { tracing::warn!(run_id = %run_id, error = %err, "Failed to load run state from store"); None @@ -2786,32 +2819,17 @@ async fn execute_run_in_process(state: Arc, run_id: RunId) { }; // Accumulate aggregate usage after execution completes. - if let Some(ref cp) = checkpoint { - let stage_durations = match run_store.list_events().await { - Ok(events) => fabro_workflow::extract_stage_durations_from_events(&events), - Err(err) => { - tracing::warn!(run_id = %run_id, error = %err, "Failed to load run events from store"); - HashMap::default() - } - }; - let mut agg = state - .aggregate_billing - .lock() - .expect("aggregate_billing lock poisoned"); - agg.total_runs += 1; - let mut run_runtime: f64 = 0.0; - for (node_id, outcome) in &cp.node_outcomes { - if let Some(usage) = &outcome.usage { - let entry = agg - .by_model - .entry(usage.model_id().to_string()) - .or_default(); - accumulate_model_billing(entry, usage); - } - let duration_ms = stage_durations.get(node_id).copied().unwrap_or(0); - run_runtime += duration_ms as f64 / 1000.0; + if let Some(ref projection) = final_projection { + if projection.checkpoint.is_some() { + let mut agg = state + .aggregate_billing + .lock() + .expect("aggregate_billing lock poisoned"); + accumulate_billing_rollup( + &mut agg, + &fabro_workflow::billing_rollup_from_projection(projection), + ); } - agg.total_runtime_secs += run_runtime; } let mut runs = state.runs.lock().expect("runs lock poisoned"); @@ -2860,7 +2878,9 @@ async fn execute_run_in_process(state: Arc, run_id: RunId) { }; } } - managed_run.checkpoint = checkpoint; + managed_run.checkpoint = final_projection + .as_ref() + .and_then(|projection| projection.checkpoint.clone()); managed_run.run_dir = Some(run_dir); clear_live_run_state(managed_run); } @@ -3103,32 +3123,15 @@ async fn execute_run_subprocess(state: Arc, run_id: RunId) { } }; - if let Some(ref checkpoint) = final_state.checkpoint { - let stage_durations = match run_store.list_events().await { - Ok(events) => fabro_workflow::extract_stage_durations_from_events(&events), - Err(err) => { - tracing::warn!(run_id = %run_id, error = %err, "Failed to load run events from store"); - HashMap::default() - } - }; + if final_state.checkpoint.is_some() { let mut agg = state .aggregate_billing .lock() .expect("aggregate_billing lock poisoned"); - agg.total_runs += 1; - let mut run_runtime: f64 = 0.0; - for (node_id, outcome) in &checkpoint.node_outcomes { - if let Some(usage) = &outcome.usage { - let entry = agg - .by_model - .entry(usage.model_id().to_string()) - .or_default(); - accumulate_model_billing(entry, usage); - } - let duration_ms = stage_durations.get(node_id).copied().unwrap_or(0); - run_runtime += duration_ms as f64 / 1000.0; - } - agg.total_runtime_secs += run_runtime; + accumulate_billing_rollup( + &mut agg, + &fabro_workflow::billing_rollup_from_projection(&final_state), + ); } let mut runs = state.runs.lock().expect("runs lock poisoned"); diff --git a/lib/crates/fabro-server/src/server/handler/billing.rs b/lib/crates/fabro-server/src/server/handler/billing.rs index 7506cf457..5c000869c 100644 --- a/lib/crates/fabro-server/src/server/handler/billing.rs +++ b/lib/crates/fabro-server/src/server/handler/billing.rs @@ -1,13 +1,14 @@ +use std::collections::HashMap; use std::sync::Arc; -use chrono::Utc; -use fabro_types::{StageId, StageProjection}; +use chrono::{DateTime, Utc}; +use fabro_types::{RunProjection, StageProjection, StageState}; use super::super::{ - ApiError, AppState, BilledTokenCounts, BillingByModel, BillingStageRef, HashMap, IntoResponse, - Json, ListResponse, ModelBillingTotals, ModelReference, PaginationParams, Path, Query, - RequiredUser, Response, Router, RunBilling, RunBillingStage, RunBillingTotals, RunId, RunStage, - State, StatusCode, accumulate_model_billing, get, parse_run_id_path, + ApiError, AppState, BillingByModel, BillingStageRef, IntoResponse, Json, ListResponse, + ModelReference, PaginationParams, Path, Query, RequiredUser, Response, Router, RunBilling, + RunBillingStage, RunBillingTotals, RunId, State, StatusCode, get, parse_run_id_path, + run_stage_from_stage_id, }; pub(super) fn routes() -> Router> { @@ -16,57 +17,6 @@ pub(super) fn routes() -> Router> { .route("/runs/{id}/billing", get(get_run_billing)) } -/// One row per `node_id`, latest visit wins. -/// -/// Mirrors the aggregation rule used in `fabro_workflow::pipeline::finalize`: -/// the displayed row uses the latest visit's data, but the row's sort key is -/// the minimum `first_event_seq` across all visits of that node — i.e. the -/// node's first appearance in the event log. This produces the same A, B -/// order for an A → B → A loop that finalize produces. -struct DedupedStage<'a> { - node_id: String, - stage: &'a StageProjection, - sort_key_first_event: u32, -} - -fn dedupe_by_node_id<'a>( - stages: impl IntoIterator, -) -> Vec> { - let mut by_node: HashMap<&'a str, (u32, u32, &'a StageProjection)> = HashMap::new(); - for (stage_id, stage) in stages { - let node_id = stage_id.node_id(); - let visit = stage_id.visit(); - let first_event = stage.first_event_seq.get(); - by_node - .entry(node_id) - .and_modify(|entry| { - if first_event < entry.0 { - entry.0 = first_event; - } - if visit >= entry.1 { - entry.1 = visit; - entry.2 = stage; - } - }) - .or_insert((first_event, visit, stage)); - } - - let mut deduped: Vec> = by_node - .into_iter() - .map(|(node_id, (first_event, _visit, stage))| DedupedStage { - node_id: node_id.to_string(), - stage, - sort_key_first_event: first_event, - }) - .collect(); - deduped.sort_by(|a, b| { - a.sort_key_first_event - .cmp(&b.sort_key_first_event) - .then_with(|| a.node_id.cmp(&b.node_id)) - }); - deduped -} - async fn list_run_stages( _auth: RequiredUser, State(state): State>, @@ -90,20 +40,18 @@ async fn list_run_stages( }; let now = Utc::now(); - let stages: Vec = dedupe_by_node_id(projection.iter_stages()) - .into_iter() - .map(|entry| { - let DedupedStage { node_id, stage, .. } = entry; - RunStage { - id: node_id.clone(), - name: node_id.clone(), - status: stage.effective_state(), - duration_secs: stage.runtime_secs(now), - dot_id: Some(node_id), - started_at: stage.started_at, - } + let stages = projection + .iter_stages() + .map(|(stage_id, stage)| { + run_stage_from_stage_id( + stage_id, + stage_id.node_id().to_string(), + stage.effective_state(), + stage.runtime_secs(now), + stage.started_at, + ) }) - .collect(); + .collect::>(); (StatusCode::OK, Json(ListResponse::new(stages))).into_response() } @@ -128,69 +76,45 @@ async fn get_run_billing( } }; - let now = Utc::now(); - - let mut by_model_totals = HashMap::::new(); - let mut runtime_secs = 0.0_f64; - let mut stages = Vec::new(); - - for entry in dedupe_by_node_id(projection.iter_stages()) { - let DedupedStage { node_id, stage, .. } = entry; - - let row_runtime = stage.runtime_secs(now).unwrap_or(0.0); - runtime_secs += row_runtime; - - let (billing, model) = if let Some(usage) = stage.usage.as_ref() { - let billing = BilledTokenCounts::from_billed_usage(std::slice::from_ref(usage)); - let model_id = usage.model_id(); - let model_totals = match by_model_totals.get_mut(model_id) { - Some(totals) => totals, - None => by_model_totals.entry(model_id.to_string()).or_default(), - }; - accumulate_model_billing(model_totals, usage); - ( - billing, - Some(ModelReference { - id: model_id.to_string(), - }), - ) - } else { - (BilledTokenCounts::default(), None) - }; - - stages.push(RunBillingStage { - billing, - model, - runtime_secs: row_runtime, - stage: BillingStageRef { - id: node_id.clone(), - name: node_id, + let rollup = fabro_workflow::billing_rollup_from_projection(&projection); + let by_model = rollup + .by_model + .iter() + .map(|model| BillingByModel { + billing: model.billing.clone(), + model: ModelReference { + id: model.model_id.clone(), }, - started_at: stage.started_at, - state: Some(stage.effective_state()), - }); - } + stages: model.stages, + }) + .collect::>(); - // Grand totals are the sum of the per-model totals we already accumulated. - let mut totals = BilledTokenCounts::default(); - for model_totals in by_model_totals.values() { - totals.input_tokens += model_totals.billing.input_tokens; - totals.output_tokens += model_totals.billing.output_tokens; - totals.reasoning_tokens += model_totals.billing.reasoning_tokens; - totals.cache_read_tokens += model_totals.billing.cache_read_tokens; - totals.cache_write_tokens += model_totals.billing.cache_write_tokens; - totals.total_tokens += model_totals.billing.total_tokens; - if let Some(value) = model_totals.billing.total_usd_micros { - *totals.total_usd_micros.get_or_insert(0) += value; - } - } - - let by_model = by_model_totals + let rollup_by_node = rollup + .stages + .iter() + .map(|stage| (stage.node_id.as_str(), stage)) + .collect::>(); + let live_rows = live_billing_rows(&projection, Utc::now()); + let runtime_secs = live_rows.iter().map(|row| row.runtime_secs).sum::(); + let stages = live_rows .into_iter() - .map(|(model, totals)| BillingByModel { - billing: totals.billing, - model: ModelReference { id: model }, - stages: totals.stages, + .map(|row| { + let rollup_stage = rollup_by_node.get(row.node_id.as_str()); + RunBillingStage { + billing: rollup_stage + .map(|stage| stage.billing.clone()) + .unwrap_or_default(), + model: rollup_stage + .and_then(|stage| stage.model_id.as_ref()) + .map(|id| ModelReference { id: id.clone() }), + runtime_secs: row.runtime_secs, + stage: BillingStageRef { + id: row.node_id.clone(), + name: row.node_id, + }, + started_at: row.started_at, + state: row.state, + } }) .collect::>(); @@ -198,16 +122,80 @@ async fn get_run_billing( by_model, stages, totals: RunBillingTotals { - cache_read_tokens: totals.cache_read_tokens, - cache_write_tokens: totals.cache_write_tokens, - input_tokens: totals.input_tokens, - output_tokens: totals.output_tokens, - reasoning_tokens: totals.reasoning_tokens, + cache_read_tokens: rollup.totals.cache_read_tokens, + cache_write_tokens: rollup.totals.cache_write_tokens, + input_tokens: rollup.totals.input_tokens, + output_tokens: rollup.totals.output_tokens, + reasoning_tokens: rollup.totals.reasoning_tokens, runtime_secs, - total_tokens: totals.total_tokens, - total_usd_micros: totals.total_usd_micros, + total_tokens: rollup.totals.total_tokens, + total_usd_micros: rollup.totals.total_usd_micros, }, }; (StatusCode::OK, Json(response)).into_response() } + +struct LiveBillingRow { + node_id: String, + runtime_secs: f64, + started_at: Option>, + state: Option, + latest_visit: u32, +} + +fn live_billing_rows(projection: &RunProjection, now: DateTime) -> Vec { + let mut row_indices = HashMap::::new(); + let mut rows = Vec::::new(); + + for (stage_id, stage) in projection.iter_stages() { + let node_id = stage_id.node_id(); + if is_exit_stage(projection, node_id) || !stage_has_billing_row(stage) { + continue; + } + + let index = *row_indices.entry(node_id.to_string()).or_insert_with(|| { + let index = rows.len(); + rows.push(LiveBillingRow { + node_id: node_id.to_string(), + runtime_secs: 0.0, + started_at: None, + state: None, + latest_visit: 0, + }); + index + }); + let row = &mut rows[index]; + row.runtime_secs += billing_runtime_secs(stage, now).unwrap_or(0.0); + + if stage_id.visit() >= row.latest_visit { + row.latest_visit = stage_id.visit(); + row.started_at = stage.started_at; + row.state = Some(stage.effective_state()); + } + } + + rows +} + +fn billing_runtime_secs(stage: &StageProjection, now: DateTime) -> Option { + stage + .duration_ms + .map(|ms| ms as f64 / 1000.0) + .or_else(|| stage.runtime_secs(now)) +} + +fn stage_has_billing_row(stage: &StageProjection) -> bool { + stage.completion.is_some() + || stage.duration_ms.is_some() + || stage.usage.is_some() + || stage.started_at.is_some() + || stage.state.is_some() +} + +fn is_exit_stage(projection: &RunProjection, node_id: &str) -> bool { + projection + .spec() + .and_then(|spec| spec.graph().nodes.get(node_id)) + .is_some_and(|node| node.handler_type() == Some("exit")) +} diff --git a/lib/crates/fabro-server/src/server/tests.rs b/lib/crates/fabro-server/src/server/tests.rs index 47fd611cb..aadc30c8e 100644 --- a/lib/crates/fabro-server/src/server/tests.rs +++ b/lib/crates/fabro-server/src/server/tests.rs @@ -18,8 +18,8 @@ use fabro_model::Provider; use fabro_types::settings::ServerAuthMethod; use fabro_types::{ AttrValue, AuthMethod, CommandTermination, FailureCategory, FailureDetail, Graph, - InterviewQuestionRecord, Outcome, QuestionType, RunBlobId, RunId, RunSpec, StageOutcome, - SystemActorKind, fixtures, + InterviewQuestionRecord, Outcome, QuestionType, RunBlobId, RunId, RunSpec, SystemActorKind, + fixtures, }; use fabro_util::check_report::CheckStatus; use httpmock::Method::{GET, POST}; @@ -2112,6 +2112,30 @@ async fn create_durable_run_with_events( } } +/// Append a stage lifecycle event with an explicit `StageScope`, so the +/// stored envelope carries the full `stage_id` (`node_id@visit`). The bare +/// [`workflow_event::append_event`] helper only writes `node_id` because +/// stage lifecycle variants don't carry visit in their payload — production +/// always emits via `Emitter::emit_scoped`. +async fn append_scoped_stage_event( + state: &Arc, + run_id: RunId, + node_id: &str, + visit: u32, + event: &workflow_event::Event, +) { + let scope = fabro_workflow::event::StageScope { + node_id: node_id.to_string(), + visit, + parallel_group_id: None, + parallel_branch_id: None, + }; + let stored = fabro_workflow::event::to_run_event_at(&run_id, event, Utc::now(), Some(&scope)); + let payload = fabro_workflow::event::build_redacted_event_payload(&stored, &run_id).unwrap(); + let run_store = state.store.open_run(&run_id).await.unwrap(); + run_store.append_event(&payload).await.unwrap(); +} + fn stage_status<'a>(body: &'a serde_json::Value, id: &str) -> &'a str { body["data"] .as_array() @@ -2134,7 +2158,14 @@ async fn list_run_stages_projects_retrying_until_completion() { }, workflow_event::Event::RunStarting, workflow_event::Event::RunRunning, - workflow_event::Event::StageStarted { + ]) + .await; + append_scoped_stage_event( + &state, + run_id, + "setup", + 1, + &workflow_event::Event::StageStarted { node_id: "setup".to_string(), name: "Setup".to_string(), index: 0, @@ -2142,7 +2173,14 @@ async fn list_run_stages_projects_retrying_until_completion() { attempt: 1, max_attempts: 1, }, - workflow_event::Event::StageCompleted { + ) + .await; + append_scoped_stage_event( + &state, + run_id, + "setup", + 1, + &workflow_event::Event::StageCompleted { node_id: "setup".to_string(), name: "Setup".to_string(), index: 0, @@ -2164,7 +2202,14 @@ async fn list_run_stages_projects_retrying_until_completion() { attempt: 1, max_attempts: 1, }, - workflow_event::Event::StageStarted { + ) + .await; + append_scoped_stage_event( + &state, + run_id, + "work", + 1, + &workflow_event::Event::StageStarted { node_id: "work".to_string(), name: "Work".to_string(), index: 1, @@ -2172,16 +2217,31 @@ async fn list_run_stages_projects_retrying_until_completion() { attempt: 1, max_attempts: 3, }, - workflow_event::Event::StageFailed { + ) + .await; + append_scoped_stage_event( + &state, + run_id, + "work", + 1, + &workflow_event::Event::StageFailed { node_id: "work".to_string(), name: "Work".to_string(), index: 1, failure: FailureDetail::new("try again", FailureCategory::TransientInfra), will_retry: true, duration_ms: 10, + billing: None, actor: None, }, - workflow_event::Event::StageRetrying { + ) + .await; + append_scoped_stage_event( + &state, + run_id, + "work", + 1, + &workflow_event::Event::StageRetrying { node_id: "work".to_string(), name: "Work".to_string(), index: 1, @@ -2189,41 +2249,9 @@ async fn list_run_stages_projects_retrying_until_completion() { max_attempts: 3, delay_ms: 100, }, - ]) + ) .await; - let mut node_outcomes = HashMap::new(); - node_outcomes.insert("setup".to_string(), Outcome::success()); - let mut checkpoint = Checkpoint { - timestamp: Utc::now(), - current_node: "setup".to_string(), - completed_nodes: vec!["setup".to_string()], - node_retries: HashMap::new(), - context_values: HashMap::new(), - node_outcomes, - next_node_id: Some("work".to_string()), - git_commit_sha: None, - loop_failure_signatures: HashMap::new(), - restart_failure_signatures: HashMap::new(), - node_visits: HashMap::new(), - }; - - let run_dir = std::env::temp_dir().join(format!("fabro-server-test-{run_id}")); - std::fs::create_dir_all(&run_dir).unwrap(); - let mut managed = managed_run( - MINIMAL_DOT.to_string(), - RunStatus::Running, - Utc::now(), - run_dir, - RunExecutionMode::Start, - ); - managed.checkpoint = Some(checkpoint.clone()); - state - .runs - .lock() - .expect("runs lock poisoned") - .insert(run_id, managed); - let response = app .clone() .oneshot( @@ -2236,29 +2264,14 @@ async fn list_run_stages_projects_retrying_until_completion() { .await .unwrap(); let body = response_json!(response, StatusCode::OK).await; - assert_eq!(stage_status(&body, "setup"), "succeeded"); - assert_eq!(stage_status(&body, "work"), "retrying"); + assert_eq!(stage_status(&body, "setup@1"), "succeeded"); + assert_eq!(stage_status(&body, "work@1"), "retrying"); - let mut work_outcome = Outcome::success(); - work_outcome.status = StageOutcome::PartiallySucceeded; - checkpoint.completed_nodes.push("work".to_string()); - checkpoint - .node_outcomes - .insert("work".to_string(), work_outcome); - checkpoint.current_node = "work".to_string(); - checkpoint.next_node_id = Some("exit".to_string()); - state - .runs - .lock() - .expect("runs lock poisoned") - .get_mut(&run_id) - .unwrap() - .checkpoint = Some(checkpoint); - - let run_store = state.store.open_run(&run_id).await.unwrap(); - workflow_event::append_event( - &run_store, - &run_id, + append_scoped_stage_event( + &state, + run_id, + "work", + 1, &workflow_event::Event::StageCompleted { node_id: "work".to_string(), name: "Work".to_string(), @@ -2282,8 +2295,7 @@ async fn list_run_stages_projects_retrying_until_completion() { max_attempts: 3, }, ) - .await - .unwrap(); + .await; let response = app .oneshot( @@ -2296,7 +2308,574 @@ async fn list_run_stages_projects_retrying_until_completion() { .await .unwrap(); let body = response_json!(response, StatusCode::OK).await; - assert_eq!(stage_status(&body, "work"), "partially_succeeded"); + assert_eq!(stage_status(&body, "work@1"), "partially_succeeded"); +} + +fn stage_entry<'a>(body: &'a serde_json::Value, id: &str) -> &'a serde_json::Value { + body["data"] + .as_array() + .unwrap() + .iter() + .find(|stage| stage["id"] == id) + .unwrap_or_else(|| panic!("stage {id} not found in {body:#?}")) +} + +fn test_billed_usage( + model_id: &str, + input_tokens: i64, + output_tokens: i64, +) -> fabro_model::BilledModelUsage { + serde_json::from_value(json!({ + "input": { + "usage": { + "model": { + "provider": "openai", + "model_id": model_id + }, + "tokens": { + "input_tokens": input_tokens, + "output_tokens": output_tokens + } + }, + "facts": { + "provider": "open_ai" + } + }, + "total_usd_micros": input_tokens + output_tokens + })) + .unwrap() +} + +#[tokio::test] +async fn list_run_stages_distinguishes_visits() { + let state = test_app_state_with_isolated_storage(); + let app = crate::test_support::build_test_router(Arc::clone(&state)); + let run_id = RunId::new(); + + create_durable_run_with_events(&state, run_id, &[ + workflow_event::Event::RunSubmitted { + definition_blob: None, + }, + workflow_event::Event::RunStarting, + workflow_event::Event::RunRunning, + ]) + .await; + + // First visit of `verify` — failed. + append_scoped_stage_event( + &state, + run_id, + "verify", + 1, + &workflow_event::Event::StageStarted { + node_id: "verify".to_string(), + name: "Verify".to_string(), + index: 1, + handler_type: "command".to_string(), + attempt: 1, + max_attempts: 1, + }, + ) + .await; + append_scoped_stage_event( + &state, + run_id, + "verify", + 1, + &workflow_event::Event::StageCompleted { + node_id: "verify".to_string(), + name: "Verify".to_string(), + index: 1, + duration_ms: 1500, + status: "failed".to_string(), + preferred_label: None, + suggested_next_ids: Vec::new(), + billing: None, + failure: None, + notes: None, + files_touched: Vec::new(), + context_updates: None, + jump_to_node: None, + context_values: None, + node_visits: None, + loop_failure_signatures: None, + restart_failure_signatures: None, + response: None, + attempt: 1, + max_attempts: 1, + }, + ) + .await; + + // Second visit of `verify` — running. + append_scoped_stage_event( + &state, + run_id, + "verify", + 2, + &workflow_event::Event::StageStarted { + node_id: "verify".to_string(), + name: "Verify".to_string(), + index: 1, + handler_type: "command".to_string(), + attempt: 1, + max_attempts: 1, + }, + ) + .await; + + let response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri(api(&format!("/runs/{run_id}/stages"))) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json!(response, StatusCode::OK).await; + + let data = body["data"].as_array().unwrap(); + let verify_entries: Vec<_> = data.iter().filter(|s| s["node_id"] == "verify").collect(); + assert_eq!(verify_entries.len(), 2, "expected two verify visits"); + + let first = stage_entry(&body, "verify@1"); + assert_eq!(first["node_id"], "verify"); + assert_eq!(first["visit"], 1); + assert_eq!(first["status"], "failed"); + assert_eq!(first["duration_secs"], 1.5); + + let second = stage_entry(&body, "verify@2"); + assert_eq!(second["node_id"], "verify"); + assert_eq!(second["visit"], 2); + assert_eq!(second["status"], "running"); + + // Old `dot_id` field must be gone. + assert!(first.get("dot_id").is_none(), "dot_id should be removed"); +} + +/// `checkpoint.completed_nodes` records every visit, so a looped node appears +/// once per re-entry. Billing must dedup so a retried node renders as one row +/// and `runtime_secs` is summed across all visits exactly once. +#[tokio::test] +async fn run_billing_dedups_retried_nodes_and_sums_their_durations() { + let state = test_app_state_with_isolated_storage(); + let app = crate::test_support::build_test_router(Arc::clone(&state)); + let run_id = RunId::new(); + + create_durable_run_with_events(&state, run_id, &[ + workflow_event::Event::RunSubmitted { + definition_blob: None, + }, + workflow_event::Event::RunStarting, + workflow_event::Event::RunRunning, + ]) + .await; + + // Visit 1 of `verify` — completed in 1.5s. + append_scoped_stage_event( + &state, + run_id, + "verify", + 1, + &workflow_event::Event::StageCompleted { + node_id: "verify".to_string(), + name: "Verify".to_string(), + index: 1, + duration_ms: 1500, + status: "failed".to_string(), + preferred_label: None, + suggested_next_ids: Vec::new(), + billing: None, + failure: None, + notes: None, + files_touched: Vec::new(), + context_updates: None, + jump_to_node: None, + context_values: None, + node_visits: None, + loop_failure_signatures: None, + restart_failure_signatures: None, + response: None, + attempt: 1, + max_attempts: 1, + }, + ) + .await; + + // Visit 2 of `verify` — completed in 0.8s. + append_scoped_stage_event( + &state, + run_id, + "verify", + 2, + &workflow_event::Event::StageCompleted { + node_id: "verify".to_string(), + name: "Verify".to_string(), + index: 1, + duration_ms: 800, + status: "succeeded".to_string(), + preferred_label: None, + suggested_next_ids: Vec::new(), + billing: None, + failure: None, + notes: None, + files_touched: Vec::new(), + context_updates: None, + jump_to_node: None, + context_values: None, + node_visits: None, + loop_failure_signatures: None, + restart_failure_signatures: None, + response: None, + attempt: 1, + max_attempts: 1, + }, + ) + .await; + + // Checkpoint records `verify` twice (once per visit) — this is what makes + // the dedup necessary. + let run_store = state.store.open_run(&run_id).await.unwrap(); + workflow_event::append_event( + &run_store, + &run_id, + &workflow_event::Event::CheckpointCompleted { + node_id: "verify".to_string(), + status: "running".to_string(), + current_node: "verify".to_string(), + completed_nodes: vec!["verify".to_string(), "verify".to_string()], + node_retries: std::collections::BTreeMap::new(), + context_values: std::collections::BTreeMap::new(), + node_outcomes: std::collections::BTreeMap::from([( + "verify".to_string(), + Outcome::default(), + )]), + next_node_id: Some("done".to_string()), + git_commit_sha: None, + loop_failure_signatures: std::collections::BTreeMap::new(), + restart_failure_signatures: std::collections::BTreeMap::new(), + node_visits: std::collections::BTreeMap::from([("verify".to_string(), 2usize)]), + diff: None, + }, + ) + .await + .unwrap(); + + let response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri(api(&format!("/runs/{run_id}/billing"))) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json!(response, StatusCode::OK).await; + + let stages = body["stages"].as_array().unwrap(); + assert_eq!( + stages.len(), + 1, + "expected one row for the retried verify node" + ); + assert_eq!(stages[0]["stage"]["id"], "verify"); + // Duration on the row is the sum across visits (1.5s + 0.8s = 2.3s). + assert!( + (stages[0]["runtime_secs"].as_f64().unwrap() - 2.3).abs() < f64::EPSILON, + "row runtime_secs should sum visits, got {}", + stages[0]["runtime_secs"] + ); + + // Totals must not double-count: a single 2.3s, not 4.6s. + assert!( + (body["totals"]["runtime_secs"].as_f64().unwrap() - 2.3).abs() < f64::EPSILON, + "totals.runtime_secs should sum visits exactly once, got {}", + body["totals"]["runtime_secs"] + ); +} + +#[tokio::test] +async fn run_billing_sums_usage_across_retry_visits_and_uses_latest_model() { + let state = test_app_state_with_isolated_storage(); + let app = crate::test_support::build_test_router(Arc::clone(&state)); + let run_id = RunId::new(); + let failed_usage = test_billed_usage("gpt-old", 100, 10); + let success_usage = test_billed_usage("gpt-new", 200, 20); + + create_durable_run_with_events(&state, run_id, &[ + workflow_event::Event::RunSubmitted { + definition_blob: None, + }, + workflow_event::Event::RunStarting, + workflow_event::Event::RunRunning, + ]) + .await; + + append_scoped_stage_event( + &state, + run_id, + "verify", + 1, + &workflow_event::Event::StageFailed { + node_id: "verify".to_string(), + name: "Verify".to_string(), + index: 1, + failure: FailureDetail::new("try again", FailureCategory::TransientInfra), + will_retry: true, + duration_ms: 1200, + billing: Some(failed_usage), + actor: None, + }, + ) + .await; + append_scoped_stage_event( + &state, + run_id, + "verify", + 2, + &workflow_event::Event::StageCompleted { + node_id: "verify".to_string(), + name: "Verify".to_string(), + index: 1, + duration_ms: 800, + status: "succeeded".to_string(), + preferred_label: None, + suggested_next_ids: Vec::new(), + billing: Some(success_usage.clone()), + failure: None, + notes: None, + files_touched: Vec::new(), + context_updates: None, + jump_to_node: None, + context_values: None, + node_visits: None, + loop_failure_signatures: None, + restart_failure_signatures: None, + response: None, + attempt: 2, + max_attempts: 2, + }, + ) + .await; + + let mut latest_outcome: Outcome> = Outcome::success(); + latest_outcome.usage = Some(success_usage); + latest_outcome.duration_ms = Some(800); + let run_store = state.store.open_run(&run_id).await.unwrap(); + workflow_event::append_event( + &run_store, + &run_id, + &workflow_event::Event::CheckpointCompleted { + node_id: "verify".to_string(), + status: "running".to_string(), + current_node: "verify".to_string(), + completed_nodes: vec!["verify".to_string(), "verify".to_string()], + node_retries: std::collections::BTreeMap::from([("verify".to_string(), 2)]), + context_values: std::collections::BTreeMap::new(), + node_outcomes: std::collections::BTreeMap::from([( + "verify".to_string(), + latest_outcome, + )]), + next_node_id: None, + git_commit_sha: None, + loop_failure_signatures: std::collections::BTreeMap::new(), + restart_failure_signatures: std::collections::BTreeMap::new(), + node_visits: std::collections::BTreeMap::from([("verify".to_string(), 2usize)]), + diff: None, + }, + ) + .await + .unwrap(); + + let response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri(api(&format!("/runs/{run_id}/billing"))) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json!(response, StatusCode::OK).await; + + let stages = body["stages"].as_array().unwrap(); + assert_eq!(stages.len(), 1); + assert_eq!(stages[0]["stage"]["id"], "verify"); + assert_eq!(stages[0]["model"]["id"], "gpt-new"); + assert_eq!(stages[0]["billing"]["input_tokens"], 300); + assert_eq!(stages[0]["billing"]["output_tokens"], 30); + assert_eq!(stages[0]["billing"]["total_usd_micros"], 330); + assert!((stages[0]["runtime_secs"].as_f64().unwrap() - 2.0).abs() < f64::EPSILON); + + assert_eq!(body["totals"]["input_tokens"], 300); + assert_eq!(body["totals"]["output_tokens"], 30); + assert_eq!(body["totals"]["total_usd_micros"], 330); + assert!((body["totals"]["runtime_secs"].as_f64().unwrap() - 2.0).abs() < f64::EPSILON); + + let by_model = body["by_model"].as_array().unwrap(); + assert_eq!(by_model.len(), 2); + let old_model = by_model + .iter() + .find(|entry| entry["model"]["id"] == "gpt-old") + .unwrap(); + let new_model = by_model + .iter() + .find(|entry| entry["model"]["id"] == "gpt-new") + .unwrap(); + assert_eq!(old_model["stages"], 1); + assert_eq!(old_model["billing"]["input_tokens"], 100); + assert_eq!(new_model["stages"], 1); + assert_eq!(new_model["billing"]["input_tokens"], 200); +} + +#[tokio::test] +async fn list_run_stages_shows_retrying_after_failed_event() { + let state = test_app_state_with_isolated_storage(); + let app = crate::test_support::build_test_router(Arc::clone(&state)); + let run_id = RunId::new(); + + create_durable_run_with_events(&state, run_id, &[ + workflow_event::Event::RunSubmitted { + definition_blob: None, + }, + workflow_event::Event::RunStarting, + workflow_event::Event::RunRunning, + ]) + .await; + + append_scoped_stage_event( + &state, + run_id, + "work", + 1, + &workflow_event::Event::StageStarted { + node_id: "work".to_string(), + name: "Work".to_string(), + index: 0, + handler_type: "command".to_string(), + attempt: 1, + max_attempts: 3, + }, + ) + .await; + append_scoped_stage_event( + &state, + run_id, + "work", + 1, + &workflow_event::Event::StageFailed { + node_id: "work".to_string(), + name: "Work".to_string(), + index: 0, + failure: FailureDetail::new("flake", FailureCategory::TransientInfra), + will_retry: true, + duration_ms: 5, + billing: None, + actor: None, + }, + ) + .await; + append_scoped_stage_event( + &state, + run_id, + "work", + 1, + &workflow_event::Event::StageRetrying { + node_id: "work".to_string(), + name: "Work".to_string(), + index: 0, + attempt: 2, + max_attempts: 3, + delay_ms: 50, + }, + ) + .await; + + let response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri(api(&format!("/runs/{run_id}/stages"))) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json!(response, StatusCode::OK).await; + assert_eq!(stage_status(&body, "work@1"), "retrying"); +} + +#[tokio::test] +async fn list_run_stages_shows_retrying_when_failed_will_retry() { + let state = test_app_state_with_isolated_storage(); + let app = crate::test_support::build_test_router(Arc::clone(&state)); + let run_id = RunId::new(); + + create_durable_run_with_events(&state, run_id, &[ + workflow_event::Event::RunSubmitted { + definition_blob: None, + }, + workflow_event::Event::RunStarting, + workflow_event::Event::RunRunning, + ]) + .await; + + append_scoped_stage_event( + &state, + run_id, + "work", + 1, + &workflow_event::Event::StageStarted { + node_id: "work".to_string(), + name: "Work".to_string(), + index: 0, + handler_type: "command".to_string(), + attempt: 1, + max_attempts: 3, + }, + ) + .await; + // Only StageFailed, no StageRetrying yet — should still render retrying + // because props.will_retry is true. + append_scoped_stage_event( + &state, + run_id, + "work", + 1, + &workflow_event::Event::StageFailed { + node_id: "work".to_string(), + name: "Work".to_string(), + index: 0, + failure: FailureDetail::new("flake", FailureCategory::TransientInfra), + will_retry: true, + duration_ms: 5, + billing: None, + actor: None, + }, + ) + .await; + + let response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri(api(&format!("/runs/{run_id}/stages"))) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json!(response, StatusCode::OK).await; + assert_eq!(stage_status(&body, "work@1"), "retrying"); } #[tokio::test] @@ -2326,6 +2905,7 @@ async fn run_billing_retried_node_then_succeeded_emits_one_row_with_final_attemp failure: FailureDetail::new("transient", FailureCategory::TransientInfra), will_retry: true, duration_ms: 10, + billing: None, actor: None, }, workflow_event::Event::StageRetrying { @@ -2438,7 +3018,7 @@ fn revisit_test_completed_with_visit( } #[tokio::test] -async fn run_billing_revisited_node_collapses_to_two_rows_with_latest_visit_data() { +async fn run_billing_revisited_node_collapses_to_two_rows_with_summed_visit_duration() { let state = test_app_state_with_isolated_storage(); let app = crate::test_support::build_test_router(Arc::clone(&state)); let run_id = RunId::new(); @@ -2480,8 +3060,8 @@ async fn run_billing_revisited_node_collapses_to_two_rows_with_latest_visit_data assert_eq!(stages[1]["stage"]["id"], "b"); let a_runtime = stages[0]["runtime_secs"].as_f64().unwrap(); assert!( - (a_runtime - 0.099).abs() < f64::EPSILON, - "A should carry latest visit's duration (99ms), got {a_runtime}" + (a_runtime - 0.1).abs() < f64::EPSILON, + "A should sum both visit durations (1ms + 99ms), got {a_runtime}" ); let b_runtime = stages[1]["runtime_secs"].as_f64().unwrap(); assert!( @@ -3886,7 +4466,13 @@ async fn create_run_pull_request_creates_and_persists_record() { .header("authorization", "Bearer openai-key"); then.status(200) .header("content-type", "application/json") - .json_body(openai_responses_payload("Narrative from mock.")); + .json_body(openai_responses_payload( + &serde_json::to_string(&json!({ + "title": "Mock title", + "body": "Narrative from mock.", + })) + .unwrap(), + )); }) .await; let openai_base_url = llm.url("/v1"); @@ -6050,6 +6636,62 @@ async fn get_aggregate_billing_returns_zeros_initially() { assert!(body["by_model"].as_array().unwrap().is_empty()); } +#[test] +fn aggregate_billing_counts_projection_rollup_usage_visits() { + let mut accumulator = BillingAccumulator::default(); + let rollup = fabro_workflow::ProjectionBillingRollup { + stages: Vec::new(), + totals: BilledTokenCounts { + input_tokens: 300, + output_tokens: 30, + total_tokens: 330, + reasoning_tokens: 0, + cache_read_tokens: 0, + cache_write_tokens: 0, + total_usd_micros: Some(330), + }, + by_model: vec![ + fabro_workflow::ProjectionBillingByModel { + model_id: "gpt-old".to_string(), + stages: 1, + billing: BilledTokenCounts { + input_tokens: 100, + output_tokens: 10, + total_tokens: 110, + reasoning_tokens: 0, + cache_read_tokens: 0, + cache_write_tokens: 0, + total_usd_micros: Some(110), + }, + }, + fabro_workflow::ProjectionBillingByModel { + model_id: "gpt-new".to_string(), + stages: 1, + billing: BilledTokenCounts { + input_tokens: 200, + output_tokens: 20, + total_tokens: 220, + reasoning_tokens: 0, + cache_read_tokens: 0, + cache_write_tokens: 0, + total_usd_micros: Some(220), + }, + }, + ], + runtime_ms: 2000, + billed_visit_count: 2, + }; + + accumulate_billing_rollup(&mut accumulator, &rollup); + + assert_eq!(accumulator.total_runs, 1); + assert_eq!(accumulator.total_runtime_secs, 2.0); + assert_eq!(accumulator.by_model["gpt-old"].stages, 1); + assert_eq!(accumulator.by_model["gpt-old"].billing.input_tokens, 100); + assert_eq!(accumulator.by_model["gpt-new"].stages, 1); + assert_eq!(accumulator.by_model["gpt-new"].billing.input_tokens, 200); +} + #[tokio::test] async fn post_runs_returns_submitted_status() { let state = test_app_state(); diff --git a/lib/crates/fabro-store/src/artifact_store.rs b/lib/crates/fabro-store/src/artifact_store.rs index 4cd4f85e1..5be00d792 100644 --- a/lib/crates/fabro-store/src/artifact_store.rs +++ b/lib/crates/fabro-store/src/artifact_store.rs @@ -297,8 +297,13 @@ fn decode_artifact_location( )) })?; let (retry, filename) = decode_retry_and_filename(location, &mut parts)?; + let stage_id = StageId::try_new(node_id, visit).map_err(|err| { + Error::Other(format!( + "artifact location {location} has an invalid stage id: {err}" + )) + })?; Ok(NodeArtifact { - node: StageId::new(node_id, visit), + node: stage_id, retry, filename, size, diff --git a/lib/crates/fabro-store/src/run_state.rs b/lib/crates/fabro-store/src/run_state.rs index c6962746e..ae23ff9fd 100644 --- a/lib/crates/fabro-store/src/run_state.rs +++ b/lib/crates/fabro-store/src/run_state.rs @@ -9,8 +9,8 @@ use fabro_types::run_event::{ use fabro_types::{ BilledModelUsage, Checkpoint, Conclusion, EventBody, FailureSignature, InterviewQuestionRecord, Outcome, PendingInterviewRecord, PullRequestRecord, RunControlAction, RunEvent, RunId, - RunProjection, RunSpec, RunStatus, RunSummary, SandboxRecord, StageCompletion, StageOutcome, - StageProjection, StageState, StartRecord, TerminalStatus, first_event_seq, + RunProjection, RunSpec, RunStatus, RunSummary, SandboxRecord, StageCompletion, StageId, + StageOutcome, StageProjection, StageState, StartRecord, TerminalStatus, first_event_seq, }; use fabro_util::error::render_with_causes; use serde_json::Value; @@ -298,67 +298,68 @@ impl RunProjectionReducer for RunProjection { stage.begin_attempt(ts); } EventBody::StageRetrying(_) => { - let Some(stage) = stage_at_current_visit(self, stored, event.seq) else { + let Some(stage) = stage_at_stored_or_current_visit(self, stored, event.seq) else { return Ok(()); }; stage.state = Some(StageState::Retrying); } EventBody::StagePrompt(props) => { - let Some(stage) = stage_at_visit(self, stored, props.visit, event.seq) else { + let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) + else { return Ok(()); }; stage.prompt = Some(props.text.clone()); stage.provider_used = provider_used_from_prompt(props); } EventBody::PromptCompleted(props) => { - let Some(stage) = stage_at_current_visit(self, stored, event.seq) else { + let Some(stage) = stage_at_stored_or_current_visit(self, stored, event.seq) else { return Ok(()); }; stage.response = Some(props.response.clone()); } EventBody::StageCompleted(props) => { - let Some(node_id) = stored.node_id.as_deref() else { - return Ok(()); - }; - let visit = stage_visit(node_id, props.node_visits.as_ref(), self).unwrap_or(1); let response = props.response.clone(); let outcome = stage_outcome_from_props(props); let completion = stage_completion_from_outcome(&outcome, ts); - let usage = props.billing.clone(); - let duration_ms = props.duration_ms; - let terminal_state = StageState::from(outcome.status); - let stage = self.stage_entry(node_id, visit, first_event_seq(event.seq)); + let Some(stage) = + stage_at_completed_visit(self, stored, props.node_visits.as_ref(), event.seq) + else { + return Ok(()); + }; stage.response = response; stage.completion = Some(completion); - stage.duration_ms = Some(duration_ms); - stage.usage = usage; - stage.state = Some(terminal_state); + stage.duration_ms = Some(props.duration_ms); + stage.usage.clone_from(&props.billing); + stage.state = Some(StageState::from(outcome.status)); } EventBody::StageFailed(props) => { let failure_reason = props.failure.as_ref().map(|detail| detail.message.clone()); - let duration_ms = props.duration_ms; - let Some(stage) = stage_at_current_visit(self, stored, event.seq) else { + let Some(stage) = stage_at_stored_or_current_visit(self, stored, event.seq) else { return Ok(()); }; + let outcome = StageOutcome::Failed { + retry_requested: props.will_retry, + }; stage.completion = Some(StageCompletion { - outcome: StageOutcome::Failed { - retry_requested: false, - }, + outcome, notes: None, failure_reason, timestamp: ts, }); - stage.duration_ms = Some(duration_ms); - stage.state = Some(StageState::Failed); + stage.duration_ms = Some(props.duration_ms); + stage.usage.clone_from(&props.billing); + stage.state = Some(StageState::from(outcome)); } EventBody::AgentSessionStarted(props) => { - let Some(stage) = stage_at_visit(self, stored, props.visit, event.seq) else { + let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) + else { return Ok(()); }; stage.provider_used = Some(provider_used_from_agent_session_started(props)); } EventBody::AgentCliStarted(props) => { - let Some(stage) = stage_at_visit(self, stored, props.visit, event.seq) else { + let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) + else { return Ok(()); }; stage.provider_used = Some(provider_used_from_agent_cli_started(props)); @@ -367,7 +368,7 @@ impl RunProjectionReducer for RunProjection { let script_invocation = serde_json::to_value(props).map_err(|err| { Error::InvalidEvent(format!("invalid command.started payload: {err}")) })?; - let Some(stage) = stage_at_current_visit(self, stored, event.seq) else { + let Some(stage) = stage_at_stored_or_current_visit(self, stored, event.seq) else { return Ok(()); }; stage.script_invocation = Some(script_invocation); @@ -376,7 +377,7 @@ impl RunProjectionReducer for RunProjection { let script_timing = serde_json::to_value(props).map_err(|err| { Error::InvalidEvent(format!("invalid command.completed payload: {err}")) })?; - let Some(stage) = stage_at_current_visit(self, stored, event.seq) else { + let Some(stage) = stage_at_stored_or_current_visit(self, stored, event.seq) else { return Ok(()); }; stage.stdout = Some(props.stdout.clone()); @@ -392,7 +393,7 @@ impl RunProjectionReducer for RunProjection { let parallel_results = serde_json::to_value(&props.results).map_err(|err| { Error::InvalidEvent(format!("invalid parallel.completed payload: {err}")) })?; - let Some(stage) = stage_at_current_visit(self, stored, event.seq) else { + let Some(stage) = stage_at_stored_or_current_visit(self, stored, event.seq) else { return Ok(()); }; stage.parallel_results = Some(parallel_results); @@ -410,6 +411,9 @@ fn stage_at_visit<'a>( visit: u32, seq: u32, ) -> Option<&'a mut StageProjection> { + if visit == 0 { + return None; + } let node_id = stored.node_id.as_deref()?; Some(state.stage_entry(node_id, visit, first_event_seq(seq))) } @@ -424,6 +428,51 @@ fn stage_at_current_visit<'a>( Some(state.stage_entry(node_id, visit, first_event_seq(seq))) } +fn stage_at_stored_stage_id<'a>( + state: &'a mut RunProjection, + stage_id: &StageId, + seq: u32, +) -> &'a mut StageProjection { + state.stage_entry(stage_id.node_id(), stage_id.visit(), first_event_seq(seq)) +} + +fn stage_at_stored_or_visit<'a>( + state: &'a mut RunProjection, + stored: &RunEvent, + visit: u32, + seq: u32, +) -> Option<&'a mut StageProjection> { + if let Some(stage_id) = stored.stage_id.as_ref() { + return Some(stage_at_stored_stage_id(state, stage_id, seq)); + } + stage_at_visit(state, stored, visit, seq) +} + +fn stage_at_stored_or_current_visit<'a>( + state: &'a mut RunProjection, + stored: &RunEvent, + seq: u32, +) -> Option<&'a mut StageProjection> { + if let Some(stage_id) = stored.stage_id.as_ref() { + return Some(stage_at_stored_stage_id(state, stage_id, seq)); + } + stage_at_current_visit(state, stored, seq) +} + +fn stage_at_completed_visit<'a>( + state: &'a mut RunProjection, + stored: &RunEvent, + node_visits: Option<&BTreeMap>, + seq: u32, +) -> Option<&'a mut StageProjection> { + if let Some(stage_id) = stored.stage_id.as_ref() { + return Some(stage_at_stored_stage_id(state, stage_id, seq)); + } + let node_id = stored.node_id.as_deref()?; + let visit = stage_visit(node_id, node_visits, state).unwrap_or(1); + Some(state.stage_entry(node_id, visit, first_event_seq(seq))) +} + pub(crate) fn build_summary(state: &RunProjection, run_id: &RunId) -> RunSummary { let workflow_name = state.spec.as_ref().map(|spec| { if spec.graph.name.is_empty() { @@ -545,6 +594,7 @@ fn stage_visit( node_visits .and_then(|visits| visits.get(node_id).copied()) .and_then(|visit| u32::try_from(visit).ok()) + .filter(|visit| *visit > 0) .or_else(|| state.current_visit_for(node_id)) } @@ -669,6 +719,32 @@ mod tests { event } + fn test_usage(model_id: &str, input_tokens: i64, output_tokens: i64) -> BilledModelUsage { + serde_json::from_value(json!({ + "input": { + "usage": { + "model": { + "provider": "openai", + "model_id": model_id + }, + "tokens": { + "input_tokens": input_tokens, + "output_tokens": output_tokens + } + }, + "facts": { + "provider": "open_ai" + } + }, + "total_usd_micros": input_tokens + output_tokens + })) + .unwrap() + } + + fn usage_json(usage: &BilledModelUsage) -> serde_json::Value { + serde_json::to_value(usage).unwrap() + } + fn test_raw_event( seq: u32, event: &str, @@ -879,6 +955,210 @@ mod tests { assert_eq!(stage.prompt.as_deref(), Some("prompt")); } + #[test] + fn stage_completed_event_captures_duration_and_usage_per_visit() { + let mut state = RunProjection::default(); + let usage = test_usage("gpt-5.2", 123, 45); + + state + .apply_event(&test_event( + 3, + EventBody::StageCompleted(StageCompletedProps { + index: 0, + duration_ms: 789, + status: StageOutcome::Succeeded, + preferred_label: None, + suggested_next_ids: Vec::new(), + billing: Some(usage.clone()), + failure: None, + notes: None, + files_touched: Vec::new(), + context_updates: None, + jump_to_node: None, + context_values: None, + node_visits: None, + loop_failure_signatures: None, + restart_failure_signatures: None, + response: Some("done".to_string()), + attempt: 1, + max_attempts: 1, + }), + Some("build"), + )) + .unwrap(); + + let stage = state.stage(&StageId::new("build", 1)).unwrap(); + assert_eq!(stage.duration_ms, Some(789)); + assert_eq!(stage.usage.as_ref(), Some(&usage)); + } + + #[test] + fn stage_failed_event_captures_duration_and_usage_per_visit() { + let mut state = RunProjection::default(); + let stage_id = StageId::new("build", 1); + let usage = test_usage("gpt-5.2", 321, 54); + + state + .apply_event(&test_stage_event( + 2, + EventBody::StageStarted(StageStartedProps { + index: 0, + handler_type: "agent".to_string(), + attempt: 1, + max_attempts: 1, + }), + stage_id.clone(), + )) + .unwrap(); + state + .apply_event(&test_raw_event( + 3, + "stage.failed", + &json!({ + "index": 0, + "failure": { + "message": "provider failed", + "failure_class": "transient_infra" + }, + "will_retry": false, + "duration_ms": 654, + "billing": usage_json(&usage) + }), + Some("build"), + )) + .unwrap(); + + let stage = state.stage(&stage_id).unwrap(); + assert_eq!(stage.duration_ms, Some(654)); + assert_eq!(stage.usage.as_ref(), Some(&usage)); + } + + #[test] + fn two_visits_of_one_node_retain_distinct_usage() { + let mut state = RunProjection::default(); + let first_usage = test_usage("gpt-5.2", 100, 10); + let second_usage = test_usage("gpt-5.2", 200, 20); + + for (seq, visit, duration_ms, usage) in [ + (3, 1usize, 111, first_usage.clone()), + (4, 2usize, 222, second_usage.clone()), + ] { + state + .apply_event(&test_event( + seq, + EventBody::StageCompleted(StageCompletedProps { + index: 0, + duration_ms, + status: StageOutcome::Succeeded, + preferred_label: None, + suggested_next_ids: Vec::new(), + billing: Some(usage), + failure: None, + notes: None, + files_touched: Vec::new(), + context_updates: None, + jump_to_node: None, + context_values: None, + node_visits: Some(BTreeMap::from([("build".to_string(), visit)])), + loop_failure_signatures: None, + restart_failure_signatures: None, + response: None, + attempt: 1, + max_attempts: 1, + }), + Some("build"), + )) + .unwrap(); + } + + let first_stage = state.stage(&StageId::new("build", 1)).unwrap(); + let second_stage = state.stage(&StageId::new("build", 2)).unwrap(); + assert_eq!(first_stage.duration_ms, Some(111)); + assert_eq!(first_stage.usage.as_ref(), Some(&first_usage)); + assert_eq!(second_stage.duration_ms, Some(222)); + assert_eq!(second_stage.usage.as_ref(), Some(&second_usage)); + } + + #[test] + fn stage_completed_prefers_stored_stage_id_over_legacy_node_visits() { + let mut state = RunProjection::default(); + let usage = test_usage("gpt-5.2", 300, 30); + let scoped_stage_id = StageId::new("build", 2); + + state + .apply_event(&test_stage_event( + 3, + EventBody::StageCompleted(StageCompletedProps { + index: 0, + duration_ms: 333, + status: StageOutcome::Succeeded, + preferred_label: None, + suggested_next_ids: Vec::new(), + billing: Some(usage.clone()), + failure: None, + notes: None, + files_touched: Vec::new(), + context_updates: None, + jump_to_node: None, + context_values: None, + node_visits: Some(BTreeMap::from([("build".to_string(), 1usize)])), + loop_failure_signatures: None, + restart_failure_signatures: None, + response: Some("done".to_string()), + attempt: 1, + max_attempts: 1, + }), + scoped_stage_id.clone(), + )) + .unwrap(); + + assert!( + state.stage(&StageId::new("build", 1)).is_none(), + "legacy node_visits must not override stored stage_id" + ); + let stage = state.stage(&scoped_stage_id).unwrap(); + assert_eq!(stage.duration_ms, Some(333)); + assert_eq!(stage.usage.as_ref(), Some(&usage)); + assert_eq!(stage.response.as_deref(), Some("done")); + } + + #[test] + fn stage_failed_prefers_stored_stage_id_and_preserves_retry_request() { + let mut state = RunProjection::default(); + let usage = test_usage("gpt-5.2", 400, 40); + let scoped_stage_id = StageId::new("build", 2); + + state + .apply_event(&test_stage_event( + 3, + EventBody::StageFailed(StageFailedProps { + index: 0, + failure: Some(fabro_types::FailureDetail::new( + "try again", + fabro_types::FailureCategory::TransientInfra, + )), + will_retry: true, + duration_ms: 444, + billing: Some(usage.clone()), + }), + scoped_stage_id.clone(), + )) + .unwrap(); + + assert!( + state.stage(&StageId::new("build", 1)).is_none(), + "current-visit fallback must not override stored stage_id" + ); + let stage = state.stage(&scoped_stage_id).unwrap(); + assert_eq!(stage.duration_ms, Some(444)); + assert_eq!(stage.usage.as_ref(), Some(&usage)); + let completion = stage.completion.as_ref().unwrap(); + assert_eq!(completion.outcome, StageOutcome::Failed { + retry_requested: true, + }); + assert_eq!(completion.failure_reason.as_deref(), Some("try again")); + } + #[test] fn checkpoint_completed_creates_projection_entry_for_skipped_stage() { let mut state = RunProjection::default(); @@ -1516,12 +1796,13 @@ mod tests { } } - fn failed_props(duration_ms: u64) -> StageFailedProps { + fn failed_props(duration_ms: u64, will_retry: bool) -> StageFailedProps { StageFailedProps { index: 0, failure: Some(FailureDetail::new("boom", FailureCategory::TransientInfra)), - will_retry: true, + will_retry, duration_ms, + billing: None, } } @@ -1644,7 +1925,7 @@ mod tests { state .apply_event(&test_event( 2, - EventBody::StageFailed(failed_props(10)), + EventBody::StageFailed(failed_props(10, false)), Some("build"), )) .unwrap(); @@ -1669,7 +1950,7 @@ mod tests { state .apply_event(&test_event( 2, - EventBody::StageFailed(failed_props(10)), + EventBody::StageFailed(failed_props(10, true)), Some("build"), )) .unwrap(); @@ -1700,7 +1981,7 @@ mod tests { state .apply_event(&test_event( 2, - EventBody::StageFailed(failed_props(10)), + EventBody::StageFailed(failed_props(10, true)), Some("build"), )) .unwrap(); diff --git a/lib/crates/fabro-store/tests/serializable_projection.rs b/lib/crates/fabro-store/tests/serializable_projection.rs index 6d51065ba..8809acda2 100644 --- a/lib/crates/fabro-store/tests/serializable_projection.rs +++ b/lib/crates/fabro-store/tests/serializable_projection.rs @@ -5,8 +5,8 @@ use fabro_store::{RunProjection, SerializableProjection, StageId}; use fabro_types::graph::Graph; use fabro_types::run::RunSpec; use fabro_types::{ - Checkpoint, RunStatus, SandboxRecord, StageCompletion, StageOutcome, StartRecord, - TerminalStatus, WorkflowSettings, first_event_seq, fixtures, + BilledModelUsage, Checkpoint, RunStatus, SandboxRecord, StageCompletion, StageOutcome, + StartRecord, TerminalStatus, WorkflowSettings, first_event_seq, fixtures, }; use serde_json::json; @@ -52,6 +52,28 @@ fn sample_checkpoint() -> Checkpoint { } } +fn sample_usage() -> BilledModelUsage { + serde_json::from_value(json!({ + "input": { + "usage": { + "model": { + "provider": "openai", + "model_id": "gpt-5.2" + }, + "tokens": { + "input_tokens": 123, + "output_tokens": 45 + } + }, + "facts": { + "provider": "open_ai" + } + }, + "total_usd_micros": 168 + })) + .expect("sample usage should deserialize") +} + #[test] fn serializable_projection_round_trips_and_trims_bulky_node_fields() { let stage_id = StageId::new("build", 2); @@ -94,11 +116,17 @@ fn serializable_projection_round_trips_and_trims_bulky_node_fields() { stage.script_invocation = Some(json!({ "command": "cargo test" })); stage.script_timing = Some(json!({ "duration_ms": 10 })); stage.parallel_results = Some(json!([{ "stage": "fanout@1" }])); + stage.duration_ms = Some(1234); + stage.usage = Some(sample_usage()); stage.stdout = Some("stdout".to_string()); stage.stderr = Some("stderr".to_string()); let serialized = serde_json::to_value(SerializableProjection(&projection)) .expect("projection should serialize"); + assert!( + serialized["stages"]["build@2"].get("usage").is_none(), + "stage usage is server-internal and should not be serialized" + ); let round_tripped: RunProjection = serde_json::from_value(serialized).expect("serialized projection should deserialize"); let node = round_tripped.stage(&stage_id).expect("node should remain"); @@ -138,6 +166,8 @@ fn serializable_projection_round_trips_and_trims_bulky_node_fields() { node.parallel_results, Some(json!([{ "stage": "fanout@1" }])) ); + assert_eq!(node.duration_ms, Some(1234)); + assert_eq!(node.usage, None); } #[test] diff --git a/lib/crates/fabro-types/src/lib.rs b/lib/crates/fabro-types/src/lib.rs index c04702c1e..22675855c 100644 --- a/lib/crates/fabro-types/src/lib.rs +++ b/lib/crates/fabro-types/src/lib.rs @@ -80,7 +80,7 @@ pub use run_summary::RunSummary; pub use sandbox_record::SandboxRecord; pub use secret::{SecretMetadata, SecretType}; pub use stage_completion::StageCompletion; -pub use stage_id::{ParallelBranchId, StageId}; +pub use stage_id::{InvalidStageVisit, ParallelBranchId, StageId}; pub use start::StartRecord; pub use status::{ BlockedReason, FailureReason, InvalidTransition, ParseFailureReasonError, diff --git a/lib/crates/fabro-types/src/outcome.rs b/lib/crates/fabro-types/src/outcome.rs index af34cd10b..8ef0c67dc 100644 --- a/lib/crates/fabro-types/src/outcome.rs +++ b/lib/crates/fabro-types/src/outcome.rs @@ -138,12 +138,32 @@ impl From for StageState { match outcome { StageOutcome::Succeeded => Self::Succeeded, StageOutcome::PartiallySucceeded => Self::PartiallySucceeded, - StageOutcome::Failed { .. } => Self::Failed, + StageOutcome::Failed { + retry_requested: true, + } => Self::Retrying, + StageOutcome::Failed { + retry_requested: false, + } => Self::Failed, StageOutcome::Skipped => Self::Skipped, } } } +#[cfg(test)] +mod stage_state_tests { + use super::{StageOutcome, StageState}; + + #[test] + fn retry_requested_failure_projects_as_retrying() { + assert_eq!( + StageState::from(StageOutcome::Failed { + retry_requested: true, + }), + StageState::Retrying + ); + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum FailureCategory { @@ -340,9 +360,10 @@ mod tests { StageState::from(StageOutcome::Failed { retry_requested: true, }), - StageState::Failed + StageState::Retrying ); assert!(StageState::Cancelled.is_terminal()); + assert!(!StageState::Retrying.is_terminal()); assert!(!StageState::Running.is_terminal()); } } diff --git a/lib/crates/fabro-types/src/run_event/stage.rs b/lib/crates/fabro-types/src/run_event/stage.rs index 9f1781d82..1b6609884 100644 --- a/lib/crates/fabro-types/src/run_event/stage.rs +++ b/lib/crates/fabro-types/src/run_event/stage.rs @@ -57,6 +57,8 @@ pub struct StageFailedProps { pub will_retry: bool, #[serde(default)] pub duration_ms: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub billing: Option, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] diff --git a/lib/crates/fabro-types/src/run_projection.rs b/lib/crates/fabro-types/src/run_projection.rs index b8aca23b3..90da9b59b 100644 --- a/lib/crates/fabro-types/src/run_projection.rs +++ b/lib/crates/fabro-types/src/run_projection.rs @@ -89,6 +89,8 @@ impl StageProjection { prompt: None, response: None, completion: None, + duration_ms: None, + usage: None, provider_used: None, diff: None, script_invocation: None, @@ -102,8 +104,6 @@ impl StageProjection { live_streaming: None, termination: None, started_at: None, - duration_ms: None, - usage: None, state: None, } } @@ -159,12 +159,32 @@ impl RunProjection { self.stages.get(stage) } + /// Iterate stages in `first_event_seq` order (the chronological order in + /// which each stage's first lifecycle event was recorded). Internal + /// storage is a `HashMap`, so iteration would otherwise be + /// non-deterministic; every caller wants chronological order, so we sort + /// here once instead of asking each caller to remember. pub fn iter_stages(&self) -> impl Iterator { - self.stages.iter() + let mut entries: Vec<(&StageId, &StageProjection)> = self.stages.iter().collect(); + entries.sort_by(|(left_id, left_stage), (right_id, right_stage)| { + left_stage + .first_event_seq + .cmp(&right_stage.first_event_seq) + .then_with(|| left_id.cmp(right_id)) + }); + entries.into_iter() } + /// Mutable counterpart of [`iter_stages`]. Same chronological ordering. pub fn iter_stages_mut(&mut self) -> impl Iterator { - self.stages.iter_mut() + let mut entries: Vec<(&StageId, &mut StageProjection)> = self.stages.iter_mut().collect(); + entries.sort_by(|(left_id, left_stage), (right_id, right_stage)| { + left_stage + .first_event_seq + .cmp(&right_stage.first_event_seq) + .then_with(|| left_id.cmp(right_id)) + }); + entries.into_iter() } pub fn is_empty(&self) -> bool { @@ -246,3 +266,90 @@ impl RunProjection { } } } + +#[cfg(test)] +mod iter_stages_tests { + use std::num::NonZeroU32; + + use super::RunProjection; + + fn seq(n: u32) -> NonZeroU32 { + NonZeroU32::new(n).unwrap() + } + + #[test] + fn iter_stages_yields_chronological_order_across_nodes() { + let mut p = RunProjection::default(); + // Insert in non-monotonic seq order to exercise the sort. + p.stage_entry("c", 1, seq(30)); + p.stage_entry("a", 1, seq(10)); + p.stage_entry("b", 1, seq(20)); + + let order: Vec<&str> = p + .iter_stages() + .map(|(stage_id, _)| stage_id.node_id()) + .collect(); + assert_eq!(order, vec!["a", "b", "c"]); + } + + #[test] + fn iter_stages_orders_visits_within_a_node() { + let mut p = RunProjection::default(); + // Visit 2 inserted first; visit 1's earlier first_event_seq must still + // win the chronological ordering. + p.stage_entry("verify", 2, seq(50)); + p.stage_entry("verify", 1, seq(20)); + + let visits: Vec = p + .iter_stages() + .map(|(stage_id, _)| stage_id.visit()) + .collect(); + assert_eq!(visits, vec![1, 2]); + } + + #[test] + fn iter_stages_mut_yields_chronological_order() { + let mut p = RunProjection::default(); + p.stage_entry("c", 1, seq(30)); + p.stage_entry("a", 1, seq(10)); + p.stage_entry("b", 1, seq(20)); + + let order: Vec = p + .iter_stages_mut() + .map(|(stage_id, _)| stage_id.node_id().to_string()) + .collect(); + assert_eq!(order, vec!["a", "b", "c"]); + } + + #[test] + fn iter_stages_tie_breaks_same_first_event_seq_by_stage_id() { + for _ in 0..128 { + let mut p = RunProjection::default(); + p.stage_entry("verify", 2, seq(10)); + p.stage_entry("build", 1, seq(10)); + p.stage_entry("verify", 1, seq(10)); + + let order: Vec = p + .iter_stages() + .map(|(stage_id, _)| stage_id.to_string()) + .collect(); + assert_eq!(order, vec!["build@1", "verify@1", "verify@2"]); + } + } + + #[test] + fn iter_stages_mut_tie_breaks_same_first_event_seq_by_stage_id() { + for _ in 0..128 { + let mut p = RunProjection::default(); + p.stage_entry("verify", 2, seq(10)); + p.stage_entry("build", 1, seq(10)); + p.stage_entry("verify", 1, seq(10)); + + let order: Vec = p + .iter_stages_mut() + .map(|(stage_id, _)| stage_id.to_string()) + .collect(); + assert_eq!(order, vec!["build@1", "verify@1", "verify@2"]); + } + } +} diff --git a/lib/crates/fabro-types/src/stage_id.rs b/lib/crates/fabro-types/src/stage_id.rs index baae58864..91711cb79 100644 --- a/lib/crates/fabro-types/src/stage_id.rs +++ b/lib/crates/fabro-types/src/stage_id.rs @@ -1,4 +1,5 @@ use std::fmt; +use std::num::NonZeroU32; use std::str::FromStr; use serde::de::Error as _; @@ -7,16 +8,21 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer}; #[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] pub struct StageId { node_id: String, - visit: u32, + visit: NonZeroU32, } impl StageId { #[must_use] pub fn new(node_id: impl Into, visit: u32) -> Self { - Self { + Self::try_new(node_id, visit).expect("stage id visit must be greater than zero") + } + + pub fn try_new(node_id: impl Into, visit: u32) -> Result { + let visit = NonZeroU32::new(visit).ok_or(InvalidStageVisit)?; + Ok(Self { node_id: node_id.into(), visit, - } + }) } #[must_use] @@ -26,7 +32,7 @@ impl StageId { #[must_use] pub fn visit(&self) -> u32 { - self.visit + self.visit.get() } } @@ -47,6 +53,17 @@ impl fmt::Display for ParseStageIdError { impl std::error::Error for ParseStageIdError {} +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct InvalidStageVisit; + +impl fmt::Display for InvalidStageVisit { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("stage id visit must be greater than zero") + } +} + +impl std::error::Error for InvalidStageVisit {} + impl FromStr for StageId { type Err = ParseStageIdError; @@ -67,7 +84,7 @@ impl FromStr for StageId { let visit = visit .parse() .map_err(|err| ParseStageIdError(format!("invalid stage id visit: {err}")))?; - Ok(Self::new(node_id, visit)) + Self::try_new(node_id, visit).map_err(|err| ParseStageIdError(err.to_string())) } } @@ -224,6 +241,18 @@ mod tests { assert!(err.to_string().starts_with("invalid stage id visit:")); } + #[test] + fn parse_rejects_zero_visit() { + let err = "code@0".parse::().unwrap_err(); + assert_eq!(err.to_string(), "stage id visit must be greater than zero"); + } + + #[test] + fn try_new_rejects_zero_visit() { + let err = StageId::try_new("code", 0).unwrap_err(); + assert_eq!(err.to_string(), "stage id visit must be greater than zero"); + } + #[test] fn parse_rejects_empty_node_id() { let err = "@3".parse::().unwrap_err(); diff --git a/lib/crates/fabro-workflow/src/billing_rollup.rs b/lib/crates/fabro-workflow/src/billing_rollup.rs new file mode 100644 index 000000000..03d5207d3 --- /dev/null +++ b/lib/crates/fabro-workflow/src/billing_rollup.rs @@ -0,0 +1,291 @@ +use std::collections::{BTreeMap, HashMap}; + +use fabro_types::{BilledModelUsage, BilledTokenCounts, RunProjection}; + +#[derive(Debug, Clone, PartialEq)] +pub struct ProjectionBillingStage { + pub node_id: String, + pub billing: BilledTokenCounts, + pub duration_ms: u64, + pub model_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProjectionBillingByModel { + pub model_id: String, + pub stages: i64, + pub billing: BilledTokenCounts, +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct ProjectionBillingRollup { + pub stages: Vec, + pub totals: BilledTokenCounts, + pub by_model: Vec, + pub runtime_ms: u64, + pub billed_visit_count: usize, +} + +impl ProjectionBillingRollup { + #[must_use] + pub fn billing_if_present(&self) -> Option { + (self.billed_visit_count > 0).then(|| self.totals.clone()) + } +} + +#[must_use] +pub fn billing_rollup_from_projection(projection: &RunProjection) -> ProjectionBillingRollup { + let mut stage_indices = HashMap::::new(); + let mut stages = Vec::::new(); + let mut by_model = BTreeMap::::new(); + let mut totals = BilledTokenCounts::default(); + let mut runtime_ms = 0_u64; + let mut billed_visit_count = 0_usize; + + for (stage_id, stage) in projection.iter_stages() { + if is_exit_stage(projection, stage_id.node_id()) { + continue; + } + if stage.completion.is_none() && stage.duration_ms.is_none() && stage.usage.is_none() { + continue; + } + + let node_id = stage_id.node_id(); + let index = *stage_indices.entry(node_id.to_string()).or_insert_with(|| { + let index = stages.len(); + stages.push(ProjectionBillingStage { + node_id: node_id.to_string(), + billing: BilledTokenCounts::default(), + duration_ms: 0, + model_id: None, + }); + index + }); + let row = &mut stages[index]; + + if let Some(duration_ms) = stage.duration_ms { + row.duration_ms = row.duration_ms.saturating_add(duration_ms); + runtime_ms = runtime_ms.saturating_add(duration_ms); + } + + if let Some(usage) = stage.usage.as_ref() { + billed_visit_count += 1; + row.model_id = Some(usage.model_id().to_string()); + accumulate_usage(&mut row.billing, usage); + accumulate_usage(&mut totals, usage); + + let model_id = usage.model_id().to_string(); + let model_entry = + by_model + .entry(model_id.clone()) + .or_insert_with(|| ProjectionBillingByModel { + model_id, + stages: 0, + billing: BilledTokenCounts::default(), + }); + model_entry.stages += 1; + accumulate_usage(&mut model_entry.billing, usage); + } + } + + ProjectionBillingRollup { + stages, + totals, + by_model: by_model.into_values().collect(), + runtime_ms, + billed_visit_count, + } +} + +fn is_exit_stage(projection: &RunProjection, node_id: &str) -> bool { + projection + .spec() + .and_then(|spec| spec.graph().nodes.get(node_id)) + .is_some_and(|node| node.handler_type() == Some("exit")) +} + +fn accumulate_usage(counts: &mut BilledTokenCounts, usage: &BilledModelUsage) { + let tokens = usage.tokens(); + counts.input_tokens += tokens.input_tokens; + counts.output_tokens += tokens.output_tokens; + counts.reasoning_tokens += tokens.reasoning_tokens; + counts.cache_read_tokens += tokens.cache_read_tokens; + counts.cache_write_tokens += tokens.cache_write_tokens; + counts.total_tokens += tokens.total_tokens(); + if let Some(value) = usage.total_usd_micros { + *counts.total_usd_micros.get_or_insert(0) += value; + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use fabro_types::{ + AttrValue, BilledModelUsage, Graph, Node, RunProjection, RunSpec, StageCompletion, + StageOutcome, WorkflowSettings, first_event_seq, fixtures, + }; + use serde_json::json; + + use super::billing_rollup_from_projection; + + fn test_usage(model_id: &str, input_tokens: i64, output_tokens: i64) -> BilledModelUsage { + serde_json::from_value(json!({ + "input": { + "usage": { + "model": { + "provider": "openai", + "model_id": model_id + }, + "tokens": { + "input_tokens": input_tokens, + "output_tokens": output_tokens + } + }, + "facts": { + "provider": "open_ai" + } + }, + "total_usd_micros": input_tokens + output_tokens + })) + .unwrap() + } + + #[test] + fn rollup_groups_stage_rows_by_node_and_sums_retry_visit_usage() { + let mut projection = RunProjection::default(); + let failed_usage = test_usage("gpt-old", 100, 10); + let success_usage = test_usage("gpt-new", 200, 20); + let first = projection.stage_entry("verify", 1, first_event_seq(1)); + first.duration_ms = Some(1200); + first.usage = Some(failed_usage); + first.completion = Some(StageCompletion { + outcome: StageOutcome::Failed { + retry_requested: true, + }, + notes: None, + failure_reason: Some("try again".to_string()), + timestamp: chrono::Utc::now(), + }); + let second = projection.stage_entry("verify", 2, first_event_seq(2)); + second.duration_ms = Some(800); + second.usage = Some(success_usage); + second.completion = Some(StageCompletion { + outcome: StageOutcome::Succeeded, + notes: None, + failure_reason: None, + timestamp: chrono::Utc::now(), + }); + + let rollup = billing_rollup_from_projection(&projection); + + assert_eq!(rollup.stages.len(), 1); + assert_eq!(rollup.stages[0].node_id, "verify"); + assert_eq!(rollup.stages[0].model_id.as_deref(), Some("gpt-new")); + assert_eq!(rollup.stages[0].duration_ms, 2000); + assert_eq!(rollup.stages[0].billing.input_tokens, 300); + assert_eq!(rollup.stages[0].billing.output_tokens, 30); + assert_eq!(rollup.stages[0].billing.total_usd_micros, Some(330)); + + assert_eq!(rollup.runtime_ms, 2000); + assert_eq!(rollup.totals.input_tokens, 300); + assert_eq!(rollup.totals.output_tokens, 30); + assert_eq!(rollup.totals.total_usd_micros, Some(330)); + assert_eq!(rollup.billed_visit_count, 2); + + assert_eq!(rollup.by_model.len(), 2); + assert_eq!(rollup.by_model[0].model_id, "gpt-new"); + assert_eq!(rollup.by_model[0].stages, 1); + assert_eq!(rollup.by_model[0].billing.input_tokens, 200); + assert_eq!(rollup.by_model[1].model_id, "gpt-old"); + assert_eq!(rollup.by_model[1].stages, 1); + assert_eq!(rollup.by_model[1].billing.input_tokens, 100); + } + + #[test] + fn rollup_includes_completed_non_llm_stage_rows_with_zero_billing() { + let mut projection = RunProjection::default(); + let stage = projection.stage_entry("start", 1, first_event_seq(1)); + stage.duration_ms = Some(25); + stage.completion = Some(StageCompletion { + outcome: StageOutcome::Succeeded, + notes: None, + failure_reason: None, + timestamp: chrono::Utc::now(), + }); + + let rollup = billing_rollup_from_projection(&projection); + + assert_eq!(rollup.stages.len(), 1); + assert_eq!(rollup.stages[0].node_id, "start"); + assert_eq!(rollup.stages[0].duration_ms, 25); + assert!(rollup.stages[0].model_id.is_none()); + assert_eq!(rollup.stages[0].billing.input_tokens, 0); + assert_eq!(rollup.runtime_ms, 25); + assert!(rollup.by_model.is_empty()); + assert!(rollup.billing_if_present().is_none()); + } + + #[test] + fn rollup_excludes_terminal_exit_stage_rows() { + let mut projection = RunProjection::default(); + projection.spec = Some(run_spec_with_exit_node()); + let start = projection.stage_entry("start", 1, first_event_seq(1)); + start.duration_ms = Some(25); + start.completion = Some(StageCompletion { + outcome: StageOutcome::Succeeded, + notes: None, + failure_reason: None, + timestamp: chrono::Utc::now(), + }); + let exit = projection.stage_entry("exit", 1, first_event_seq(2)); + exit.duration_ms = Some(7); + exit.completion = Some(StageCompletion { + outcome: StageOutcome::Succeeded, + notes: None, + failure_reason: None, + timestamp: chrono::Utc::now(), + }); + + let rollup = billing_rollup_from_projection(&projection); + + assert_eq!(rollup.stages.len(), 1); + assert_eq!(rollup.stages[0].node_id, "start"); + assert_eq!(rollup.runtime_ms, 25); + } + + fn run_spec_with_exit_node() -> RunSpec { + let mut graph = Graph::new("test"); + graph.nodes.insert("start".to_string(), { + let mut node = Node::new("start"); + node.attrs.insert( + "shape".to_string(), + AttrValue::String("Mdiamond".to_string()), + ); + node + }); + graph.nodes.insert("exit".to_string(), { + let mut node = Node::new("exit"); + node.attrs.insert( + "shape".to_string(), + AttrValue::String("Msquare".to_string()), + ); + node + }); + + RunSpec { + run_id: fixtures::RUN_1, + settings: WorkflowSettings::default(), + graph, + workflow_slug: None, + source_directory: None, + labels: HashMap::new(), + provenance: None, + manifest_blob: None, + definition_blob: None, + git: None, + fork_source_ref: None, + in_place: false, + } + } +} diff --git a/lib/crates/fabro-workflow/src/error.rs b/lib/crates/fabro-workflow/src/error.rs index 1d08b2777..6a3ff0d06 100644 --- a/lib/crates/fabro-workflow/src/error.rs +++ b/lib/crates/fabro-workflow/src/error.rs @@ -1838,6 +1838,7 @@ mod tests { failure: failure.clone(), will_retry: false, duration_ms: 0, + billing: None, actor: None, }; diff --git a/lib/crates/fabro-workflow/src/event/convert.rs b/lib/crates/fabro-workflow/src/event/convert.rs index 815aaec11..4aa9df5b1 100644 --- a/lib/crates/fabro-workflow/src/event/convert.rs +++ b/lib/crates/fabro-workflow/src/event/convert.rs @@ -290,12 +290,14 @@ fn event_body_from_event(event: &Event) -> EventBody { failure, will_retry, duration_ms, + billing, .. } => EventBody::StageFailed(fabro_types::StageFailedProps { index: *index, failure: Some(failure.clone()), will_retry: *will_retry, duration_ms: *duration_ms, + billing: billing.clone(), }), Event::StageRetrying { index, @@ -1178,7 +1180,7 @@ mod tests { use crate::error::Error; use crate::event::test_support::user_principal; use crate::event::{Event, StageScope}; - use crate::outcome::FailureDetail; + use crate::outcome::{BilledModelUsage, FailureDetail}; #[derive(Debug)] struct EventTestCause; @@ -1200,6 +1202,28 @@ mod tests { } } + fn test_usage(model_id: &str, input_tokens: i64, output_tokens: i64) -> BilledModelUsage { + serde_json::from_value(serde_json::json!({ + "input": { + "usage": { + "model": { + "provider": "openai", + "model_id": model_id + }, + "tokens": { + "input_tokens": input_tokens, + "output_tokens": output_tokens + } + }, + "facts": { + "provider": "open_ai" + } + }, + "total_usd_micros": input_tokens + output_tokens + })) + .unwrap() + } + #[test] fn run_event_stage_completed_places_node_fields_in_header() { let stored = to_run_event_at( @@ -1279,6 +1303,7 @@ mod tests { #[test] fn run_event_stage_failure_keeps_failure_detail() { + let usage = test_usage("gpt-5.2", 321, 54); let stored = to_run_event(&fixtures::RUN_3, &Event::StageFailed { node_id: "code".to_string(), name: "Code".to_string(), @@ -1289,6 +1314,7 @@ mod tests { ), will_retry: true, duration_ms: 5000, + billing: Some(usage.clone()), actor: None, }); @@ -1297,6 +1323,7 @@ mod tests { assert_eq!(properties["failure"]["message"], "lint failed"); assert_eq!(properties["failure"]["failure_class"], "deterministic"); assert_eq!(properties["will_retry"], true); + assert_eq!(properties["billing"], serde_json::to_value(&usage).unwrap()); } #[test] diff --git a/lib/crates/fabro-workflow/src/event/events.rs b/lib/crates/fabro-workflow/src/event/events.rs index a3b2e2b5b..c753e20c6 100644 --- a/lib/crates/fabro-workflow/src/event/events.rs +++ b/lib/crates/fabro-workflow/src/event/events.rs @@ -206,6 +206,7 @@ pub enum Event { failure: FailureDetail, will_retry: bool, duration_ms: u64, + billing: Option, #[serde(default, skip_serializing_if = "Option::is_none")] actor: Option, }, diff --git a/lib/crates/fabro-workflow/src/lib.rs b/lib/crates/fabro-workflow/src/lib.rs index 9af074d97..5478e611d 100644 --- a/lib/crates/fabro-workflow/src/lib.rs +++ b/lib/crates/fabro-workflow/src/lib.rs @@ -20,6 +20,7 @@ use std::sync::Arc; use fabro_retro::retro::CompletedStage; use fabro_store::EventEnvelope; +use fabro_types::{EventBody, StageId}; /// Callback invoked when a workflow node starts executing. pub type OnNodeCallback = Option>; @@ -86,34 +87,190 @@ pub fn build_completed_stages(cp: &records::Checkpoint, run_failed: bool) -> Vec stages } -pub fn extract_stage_durations_from_events(events: &[EventEnvelope]) -> HashMap { +/// Extract the `duration_ms` from a `stage.completed` / `stage.failed` +/// event body, or `None` for any other variant. +fn stage_completion_duration_ms(body: &EventBody) -> Option { + match body { + EventBody::StageCompleted(props) => Some(props.duration_ms), + EventBody::StageFailed(props) => Some(props.duration_ms), + _ => None, + } +} + +/// Extract per-stage (node_id, visit) durations from `stage.completed` / +/// `stage.failed` events. Keys on the full [`StageId`] so multi-visit stages +/// (e.g. a looped `verify` node) keep distinct durations. +/// +/// This is the canonical primitive; [`total_stage_duration_by_node`] and +/// [`latest_stage_duration_by_node`] are explicit rollups built on top of it. +pub fn extract_stage_durations_by_stage_id(events: &[EventEnvelope]) -> HashMap { let mut durations = HashMap::new(); for envelope in events { - let event = &envelope.event; - let event_name = event.event_name(); - if event_name != "stage.completed" && event_name != "stage.failed" { - continue; - } - let Some(node_id) = event.node_id.as_deref() else { + let Some(duration_ms) = stage_completion_duration_ms(&envelope.event.body) else { continue; }; - let Some(duration_ms) = event - .properties() - .ok() - .and_then(|properties| properties.get("duration_ms").cloned()) - .and_then(|duration| duration.as_u64()) - else { + let Some(stage_id) = envelope.event.stage_id.as_ref() else { continue; }; - durations.insert(node_id.to_string(), duration_ms); + durations.insert(stage_id.clone(), duration_ms); } durations } +/// Total duration spent in each node, summed across every visit. Use for +/// billing/usage where a retried node should count its full time. +pub fn total_stage_duration_by_node(events: &[EventEnvelope]) -> HashMap { + let mut totals: HashMap = HashMap::new(); + for (stage_id, duration_ms) in extract_stage_durations_by_stage_id(events) { + *totals.entry(stage_id.node_id().to_string()).or_default() += duration_ms; + } + totals +} + +/// Duration of each node's most recent visit (the highest visit number). Use +/// for run summaries and retros where the table shows one row per node and +/// "the last attempt" is the right representative. +pub fn latest_stage_duration_by_node(events: &[EventEnvelope]) -> HashMap { + let mut entries: Vec<(StageId, u64)> = extract_stage_durations_by_stage_id(events) + .into_iter() + .collect(); + entries.sort_by_key(|(stage_id, _)| stage_id.visit()); + let mut latest = HashMap::new(); + for (stage_id, duration_ms) in entries { + latest.insert(stage_id.node_id().to_string(), duration_ms); + } + latest +} + +#[cfg(test)] +mod duration_tests { + use chrono::{TimeZone, Utc}; + use fabro_store::EventEnvelope; + use fabro_types::run_event::{StageCompletedProps, StageFailedProps}; + use fabro_types::{EventBody, RunEvent, StageId, StageOutcome, fixtures}; + + use super::{ + extract_stage_durations_by_stage_id, latest_stage_duration_by_node, + total_stage_duration_by_node, + }; + + fn completed_event(seq: u32, node: &str, visit: u32, duration_ms: u64) -> EventEnvelope { + let event = RunEvent { + id: format!("evt_{seq}"), + ts: Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap(), + run_id: fixtures::RUN_1, + node_id: Some(node.to_string()), + node_label: None, + stage_id: Some(StageId::new(node, visit)), + parallel_group_id: None, + parallel_branch_id: None, + session_id: None, + parent_session_id: None, + tool_call_id: None, + actor: None, + body: EventBody::StageCompleted(StageCompletedProps { + index: 0, + duration_ms, + status: StageOutcome::Succeeded, + preferred_label: None, + suggested_next_ids: vec![], + billing: None, + failure: None, + notes: None, + files_touched: vec![], + context_updates: None, + jump_to_node: None, + context_values: None, + node_visits: None, + loop_failure_signatures: None, + restart_failure_signatures: None, + response: None, + attempt: 1, + max_attempts: 1, + }), + }; + EventEnvelope { seq, event } + } + + fn failed_event(seq: u32, node: &str, visit: u32, duration_ms: u64) -> EventEnvelope { + let event = RunEvent { + id: format!("evt_{seq}"), + ts: Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap(), + run_id: fixtures::RUN_1, + node_id: Some(node.to_string()), + node_label: None, + stage_id: Some(StageId::new(node, visit)), + parallel_group_id: None, + parallel_branch_id: None, + session_id: None, + parent_session_id: None, + tool_call_id: None, + actor: None, + body: EventBody::StageFailed(StageFailedProps { + index: 0, + failure: None, + will_retry: true, + duration_ms, + billing: None, + }), + }; + EventEnvelope { seq, event } + } + + #[test] + fn extract_keys_durations_by_full_stage_id() { + let events = vec![ + completed_event(1, "verify", 1, 100), + completed_event(2, "verify", 2, 200), + ]; + let durations = extract_stage_durations_by_stage_id(&events); + assert_eq!( + durations.get(&StageId::new("verify", 1)).copied(), + Some(100) + ); + assert_eq!( + durations.get(&StageId::new("verify", 2)).copied(), + Some(200) + ); + } + + #[test] + fn total_sums_across_visits_per_node() { + let events = vec![ + completed_event(1, "verify", 1, 100), + completed_event(2, "verify", 2, 200), + completed_event(3, "build", 1, 50), + ]; + let totals = total_stage_duration_by_node(&events); + assert_eq!(totals.get("verify").copied(), Some(300)); + assert_eq!(totals.get("build").copied(), Some(50)); + } + + #[test] + fn latest_picks_highest_visit_regardless_of_input_order() { + // Visit 2 appears in the events vector before visit 1; the result + // must still reflect visit 2's duration (the latest visit). + let events = vec![ + completed_event(1, "verify", 2, 999), + completed_event(2, "verify", 1, 100), + ]; + let latest = latest_stage_duration_by_node(&events); + assert_eq!(latest.get("verify").copied(), Some(999)); + } + + #[test] + fn stage_failed_durations_are_included() { + let events = vec![failed_event(1, "verify", 1, 75)]; + let durations = extract_stage_durations_by_stage_id(&events); + assert_eq!(durations.get(&StageId::new("verify", 1)).copied(), Some(75)); + } +} + #[doc(hidden)] pub mod artifact; pub mod artifact_snapshot; pub mod artifact_upload; +pub mod billing_rollup; pub mod command_log; pub(crate) mod condition; pub mod context; @@ -142,6 +299,10 @@ pub mod run_control; pub(crate) mod run_dir; pub mod run_lookup; +pub use billing_rollup::{ + ProjectionBillingByModel, ProjectionBillingRollup, ProjectionBillingStage, + billing_rollup_from_projection, +}; pub use error::{Error, FailureCategory, FailureSignature, FailureSignatureExt, Result}; pub use manifest_path::ManifestPath; pub mod run_materialization; diff --git a/lib/crates/fabro-workflow/src/lifecycle/event.rs b/lib/crates/fabro-workflow/src/lifecycle/event.rs index 482f0311a..82dbd0501 100644 --- a/lib/crates/fabro-workflow/src/lifecycle/event.rs +++ b/lib/crates/fabro-workflow/src/lifecycle/event.rs @@ -224,6 +224,7 @@ impl RunLifecycle for EventLifecycle { failure, will_retry: true, duration_ms, + billing: outcome.usage.clone(), actor, }, &scope, @@ -275,6 +276,7 @@ impl RunLifecycle for EventLifecycle { failure, will_retry: false, duration_ms, + billing: outcome.usage.clone(), actor, }, &scope, diff --git a/lib/crates/fabro-workflow/src/pipeline/finalize.rs b/lib/crates/fabro-workflow/src/pipeline/finalize.rs index 921697369..927d2d833 100644 --- a/lib/crates/fabro-workflow/src/pipeline/finalize.rs +++ b/lib/crates/fabro-workflow/src/pipeline/finalize.rs @@ -19,6 +19,7 @@ use crate::run_status::{FailureReason, RunStatus, SuccessReason}; use crate::runtime_store::RunStoreHandle; use crate::sandbox_git::git_diff_with_timeout; use crate::services::RunServices; +use crate::{ProjectionBillingRollup, billing_rollup_from_projection}; pub fn classify_engine_result( engine_result: &Result, @@ -68,22 +69,22 @@ pub(crate) async fn build_conclusion_from_store( run_duration_ms: u64, final_git_commit_sha: Option, ) -> Conclusion { - let (state_result, events_result) = tokio::join!(run_store.state(), run_store.list_events()); - let projection = state_result.ok(); + let projection = run_store.state().await.ok(); let projection_order = projection .as_ref() .map(stage_projection_order) .unwrap_or_default(); + let projection_billing = projection + .as_ref() + .map(billing_rollup_from_projection) + .unwrap_or_default(); let checkpoint = projection .as_ref() .and_then(|state| state.checkpoint.as_ref()); - let stage_durations = events_result - .map(|events| crate::extract_stage_durations_from_events(&events)) - .unwrap_or_default(); build_conclusion_from_parts( checkpoint, - &stage_durations, + &projection_billing, &projection_order, status, failure_reason, @@ -94,7 +95,7 @@ pub(crate) async fn build_conclusion_from_store( fn build_conclusion_from_parts( checkpoint: Option<&Checkpoint>, - stage_durations: &HashMap, + projection_billing: &ProjectionBillingRollup, projection_order: &HashMap, status: StageOutcome, failure_reason: Option, @@ -105,6 +106,11 @@ fn build_conclusion_from_parts( // while the other checkpoint maps are keyed by node_id. Dedupe to one row // per node so the stages table matches the deduped billing total. let (stages, total_retries) = if let Some(cp) = checkpoint { + let billing_by_node = projection_billing + .stages + .iter() + .map(|stage| (stage.node_id.as_str(), stage)) + .collect::>(); let mut stage_rows = Vec::new(); let mut seen = std::collections::HashSet::new(); let mut retries_sum: u32 = 0; @@ -130,7 +136,6 @@ fn build_conclusion_from_parts( } for (original_checkpoint_order, node_id) in stage_order { - let outcome = cp.node_outcomes.get(node_id); let retries = cp .node_retries .get(node_id) @@ -138,14 +143,13 @@ fn build_conclusion_from_parts( .unwrap_or(1) .saturating_sub(1); retries_sum += retries; + let billing = billing_by_node.get(node_id); let summary = StageSummary { stage_id: node_id.to_string(), stage_label: node_id.to_string(), - duration_ms: stage_durations.get(node_id).copied().unwrap_or(0), - billing_usd_micros: outcome - .and_then(|o| o.usage.as_ref()) - .and_then(|usage| usage.total_usd_micros), + duration_ms: billing.map_or(0, |stage| stage.duration_ms), + billing_usd_micros: billing.and_then(|stage| stage.billing.total_usd_micros), retries, }; stage_rows.push(( @@ -176,7 +180,7 @@ fn build_conclusion_from_parts( failure_reason, final_git_commit_sha, stages, - billing: checkpoint.and_then(billing_from_checkpoint), + billing: projection_billing.billing_if_present(), total_retries, } } @@ -391,15 +395,8 @@ async fn compute_final_patch( } } -/// Iterates `node_outcomes.values()` rather than `completed_nodes` to avoid -/// over-counting the last visit's usage on looping workflows. -pub(crate) fn billing_from_checkpoint(cp: &Checkpoint) -> Option { - let usage: Vec<_> = cp - .node_outcomes - .values() - .filter_map(|o| o.usage.clone()) - .collect(); - (!usage.is_empty()).then(|| BilledTokenCounts::from_billed_usage(&usage)) +pub(crate) fn billing_from_projection(projection: &RunProjection) -> Option { + billing_rollup_from_projection(projection).billing_if_present() } pub(crate) fn build_terminal_event( @@ -503,7 +500,6 @@ pub async fn finalize(retroed: Retroed, options: &FinalizeOptions) -> Result Result BilledModelUsage { + serde_json::from_value(serde_json::json!({ + "input": { + "usage": { + "model": { + "provider": "openai", + "model_id": model_id + }, + "tokens": { + "input_tokens": input_tokens, + "output_tokens": output_tokens + } + }, + "facts": { + "provider": "open_ai" + } + }, + "total_usd_micros": input_tokens + output_tokens + })) + .unwrap() + } + #[test] fn conclusion_stage_order_follows_projection_first_event_order() { let mut projection = RunProjection::default(); @@ -753,7 +776,7 @@ mod tests { let conclusion = build_conclusion_from_parts( Some(&checkpoint), - &HashMap::new(), + &ProjectionBillingRollup::default(), &projection_order, StageOutcome::Succeeded, None, @@ -788,7 +811,7 @@ mod tests { let conclusion = build_conclusion_from_parts( Some(&checkpoint), - &HashMap::new(), + &ProjectionBillingRollup::default(), &projection_order, StageOutcome::Succeeded, None, @@ -804,6 +827,66 @@ mod tests { assert_eq!(stage_ids, vec!["skipped", "finished"]); } + #[test] + fn conclusion_billing_sums_retry_visit_usage_from_projection() { + let mut projection = RunProjection::default(); + let failed_usage = test_usage("gpt-old", 100, 10); + let success_usage = test_usage("gpt-new", 200, 20); + let failed = projection.stage_entry("verify", 1, first_event_seq(1)); + failed.duration_ms = Some(1200); + failed.usage = Some(failed_usage); + failed.completion = Some(StageCompletion { + outcome: StageOutcome::Failed { + retry_requested: true, + }, + notes: None, + failure_reason: Some("try again".to_string()), + timestamp: chrono::Utc::now(), + }); + let succeeded = projection.stage_entry("verify", 2, first_event_seq(2)); + succeeded.duration_ms = Some(800); + succeeded.usage = Some(success_usage.clone()); + succeeded.completion = Some(StageCompletion { + outcome: StageOutcome::Succeeded, + notes: None, + failure_reason: None, + timestamp: chrono::Utc::now(), + }); + + let projection_order = stage_projection_order(&projection); + let projection_billing = billing_rollup_from_projection(&projection); + let mut latest_outcome = Outcome::success(); + latest_outcome.usage = Some(success_usage); + latest_outcome.duration_ms = Some(800); + let mut checkpoint = checkpoint_with( + vec!["verify", "verify"], + HashMap::from([("verify".to_string(), latest_outcome)]), + ); + checkpoint.node_retries.insert("verify".to_string(), 2); + + let conclusion = build_conclusion_from_parts( + Some(&checkpoint), + &projection_billing, + &projection_order, + StageOutcome::Succeeded, + None, + 10, + None, + ); + + assert_eq!(conclusion.billing.as_ref().unwrap().input_tokens, 300); + assert_eq!(conclusion.billing.as_ref().unwrap().output_tokens, 30); + assert_eq!( + conclusion.billing.as_ref().unwrap().total_usd_micros, + Some(330) + ); + assert_eq!(conclusion.stages.len(), 1); + assert_eq!(conclusion.stages[0].stage_id, "verify"); + assert_eq!(conclusion.stages[0].duration_ms, 2000); + assert_eq!(conclusion.stages[0].billing_usd_micros, Some(330)); + assert_eq!(conclusion.stages[0].retries, 1); + } + fn test_services( run_store: RunStoreHandle, emitter: Arc, diff --git a/lib/crates/fabro-workflow/src/pipeline/mod.rs b/lib/crates/fabro-workflow/src/pipeline/mod.rs index 71aefbef9..424dd69bd 100644 --- a/lib/crates/fabro-workflow/src/pipeline/mod.rs +++ b/lib/crates/fabro-workflow/src/pipeline/mod.rs @@ -12,14 +12,15 @@ mod validate; pub use execute::execute; pub use fabro_types::PullRequestRecord; pub(crate) use finalize::{ - billing_from_checkpoint, build_conclusion_from_store, build_terminal_event, + billing_from_projection, build_conclusion_from_store, build_terminal_event, }; pub use finalize::{classify_engine_result, finalize, write_finalize_commit}; pub use initialize::initialize; pub use parse::parse; pub(crate) use persist::persist; pub use pull_request::{ - AutoMergeOptions, OpenPullRequestRequest, build_pr_body, maybe_open_pull_request, pull_request, + AutoMergeOptions, OpenPullRequestRequest, PrContent, build_pr_content, maybe_open_pull_request, + pull_request, }; pub use retro::{retro, run_retro}; pub use transform::transform; diff --git a/lib/crates/fabro-workflow/src/pipeline/prompts/pr_body.md b/lib/crates/fabro-workflow/src/pipeline/prompts/pr_body.md new file mode 100644 index 000000000..9175e62db --- /dev/null +++ b/lib/crates/fabro-workflow/src/pipeline/prompts/pr_body.md @@ -0,0 +1,54 @@ +You are writing a pull request title and description for a code change produced by an AI workflow. + +OUTPUT FORMAT +Return a JSON object with exactly two fields: +- "title": a one-line title, max 72 characters, no trailing period. +- "body": the markdown body as described below. + +DO NOT INCLUDE in the body +- A `#` or `##` title heading at the top -- the title goes in the `title` field. +- A "Retro" section, "Fabro Details" section, cost/duration table, or "Generated with" footer -- those are appended programmatically after your output. +- The full plan text -- the full plan is appended programmatically as a
block. +- Bare `#1`, `#2` list prefixes -- GitHub auto-links those as issue references. Use plain `1.`, `2.` instead. +- A test plan unless the testing approach is non-obvious. + +SIZE THE BODY TO THE CHANGE +First classify along two axes from the diff: +- Size: how many files changed, how large the diff is. +- Complexity: trivial (rename / typo / dep bump / config) vs. design decisions / new patterns / cross-cutting concerns. + +Then write at the matching depth: + +| Profile | Body shape | +|---|---| +| Small + simple (typo, config, dep bump) | 1-2 sentences, no headers, total under ~300 characters | +| Small + non-trivial (targeted bugfix, behavioral change) | Short "Problem / Fix" narrative, 3-5 sentences. No headers unless two distinct concerns. | +| Medium feature or refactor | Summary paragraph, then a section explaining what changed and why. Call out design decisions. | +| Large or architecturally significant | Full narrative: problem context, approach chosen (and why), key decisions, migration/rollback notes if relevant. | +| Performance improvement | Include before/after measurements if available. A markdown table works well here. | + +Brevity matters for small changes. A 3-line bugfix with a 20-line description signals miscalibration. When in doubt, shorter is better -- reviewers can read the diff. + +WRITING PRINCIPLES +- Lead with value: the first sentence tells the reviewer *why this PR exists*, not *what files changed*. +- Describe the net result, not the journey: skip intermediate failures, debugging steps, and refactors done during development. +- Trust the final diff: if the goal or plan disagree with the diff, the diff is authoritative. +- Explain the non-obvious: spend description space on what the diff doesn't show -- why this approach, what was rejected, what to look at first. +- Use structure when it earns its keep: no empty sections, no template headers without content. +- If the body uses any `##` heading, the opening summary must also be under a heading (e.g. `## Summary`); otherwise a bare paragraph is fine. + +PLAN SUMMARY +The full plan is attached separately as a
block, so do not restate it. Include a brief `### Plan Summary` with bullet points only when the change is medium or larger in the sizing matrix above. Skip it for small changes. + +VISUAL AIDS +Include a visual aid only when a reviewer would struggle to reconstruct the mental model from prose alone -- based on what changes structurally, not on PR size. Skip for trivial / mechanical changes, or when prose already communicates clearly. + +| PR changes... | Visual aid | +|---|---| +| 3+ interacting components or services | Mermaid component / interaction diagram | +| Multi-step workflow or pipeline with non-obvious sequencing | Mermaid flow diagram | +| 3+ behavioral modes or variants | Markdown comparison table | +| Before/after data or trade-offs | Markdown table | +| Data model changes with 3+ related entities | Mermaid ERD | + +Mermaid: prefer `TB` direction, <=10 nodes typical. Place inline at the point of relevance, not in a separate "Diagrams" section. diff --git a/lib/crates/fabro-workflow/src/pipeline/pull_request.rs b/lib/crates/fabro-workflow/src/pipeline/pull_request.rs index 5b9355ccf..35f27aa92 100644 --- a/lib/crates/fabro-workflow/src/pipeline/pull_request.rs +++ b/lib/crates/fabro-workflow/src/pipeline/pull_request.rs @@ -1,16 +1,17 @@ -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use fabro_auth::CredentialSource; use fabro_github::{self as github_app, ssh_url_to_https}; use fabro_graphviz::parser; use fabro_llm::client::Client; -use fabro_llm::generate::{GenerateParams, generate}; +use fabro_llm::generate::{GenerateParams, generate_object}; +use fabro_model::Catalog; use fabro_retro::retro::Retro; use fabro_store::RunProjection; use fabro_types::PullRequestRecord; use fabro_types::settings::run::MergeStrategy; use fabro_util::text::strip_goal_decoration; -use tracing::{debug, info}; +use tracing::{debug, info, warn}; use super::types::{Concluded, Finalized, PullRequestOptions}; use crate::event::{Event, RunNoticeLevel}; @@ -18,16 +19,115 @@ use crate::outcome::{StageOutcome, format_cost as outcome_format_cost}; use crate::records::{Conclusion, RunSpec}; use crate::runtime_store::RunStoreHandle; +/// Maximum length of a PR title (Unicode scalar values). +const PR_TITLE_MAX_CHARS: usize = 72; + +/// Structured output schema for the LLM-generated PR title and body. +static PR_CONTENT_SCHEMA: LazyLock = LazyLock::new(|| { + serde_json::json!({ + "type": "object", + "properties": { + "title": { "type": "string" }, + "body": { "type": "string" } + }, + "required": ["title", "body"], + "additionalProperties": false + }) +}); + +/// Complete pull request content generated for a workflow run. +#[derive(Debug, serde::Deserialize)] +pub struct PrContent { + pub title: String, + pub body: String, +} + +/// System prompt that instructs the LLM how to write a Fabro PR title and +/// body. The trailing programmatic sections (Plan `
`, Retro, +/// Fabro Details, footer) are appended after the LLM body — the prompt +/// explicitly forbids the LLM from duplicating them. +const PR_BODY_SYSTEM_PROMPT: &str = include_str!("prompts/pr_body.md"); + +const DEFAULT_PR_TITLE: &str = "Update workflow output"; +const EMPTY_BODY_NOTICE: &str = "> _The LLM did not produce a description for this change. The diff and the appended details are the source of truth for review._"; + +/// Truncation budget for the LLM prompt's plan / diff sections. +#[derive(Debug, PartialEq, Eq)] +struct TruncationCaps { + plan: usize, + diff: usize, +} + +const DIFF_HARD_CAP: usize = 500_000; +const PLAN_HARD_CAP: usize = 100_000; +const DIFF_FRACTION_NUM: usize = 4; +const PLAN_FRACTION_NUM: usize = 1; +const FRACTION_DEN: usize = 10; +const UNKNOWN_MODEL_CTX: usize = 200_000; + +/// Resolve truncation caps based on the model's context window. Unknown +/// models use the baseline 200k context-window assumption. +fn truncation_caps(model: &str) -> TruncationCaps { + let ctx = Catalog::builtin() + .get(model) + .and_then(|m| usize::try_from(m.context_window()).ok()) + .unwrap_or(UNKNOWN_MODEL_CTX); + + truncation_caps_for_context_window(ctx) +} + +fn truncation_caps_for_context_window(ctx: usize) -> TruncationCaps { + TruncationCaps { + diff: ctx + .saturating_mul(DIFF_FRACTION_NUM) + .checked_div(FRACTION_DEN) + .unwrap_or(DIFF_HARD_CAP) + .min(DIFF_HARD_CAP), + plan: ctx + .saturating_mul(PLAN_FRACTION_NUM) + .checked_div(FRACTION_DEN) + .unwrap_or(PLAN_HARD_CAP) + .min(PLAN_HARD_CAP), + } +} + +/// Truncate `s` to at most `max` Unicode scalar values without splitting a +/// UTF-8 sequence. +fn truncate_chars(s: &str, max: usize) -> &str { + s.char_indices() + .nth(max) + .map_or(s, |(boundary, _)| &s[..boundary]) +} + +/// Truncate `s` to at most `max` Unicode scalar values, replacing the +/// trailing char with `…` when truncation occurs. +fn truncate_with_ellipsis(s: &str, max: usize) -> String { + if s.chars().count() > max { + let truncated: String = s.chars().take(max - 1).collect(); + format!("{truncated}\u{2026}") + } else { + s.to_string() + } +} + +/// Cap a PR title at [`PR_TITLE_MAX_CHARS`]. +fn enforce_title_cap(title: &str) -> String { + truncate_with_ellipsis(title, PR_TITLE_MAX_CHARS) +} + /// Derive a PR title from the workflow goal. /// -/// Uses the first line, truncated to 120 characters for readability. +/// Uses the first line, truncated to the same cap as LLM-generated titles. fn pr_title_from_goal(goal: &str) -> String { - let stripped = strip_goal_decoration(goal); - if stripped.chars().count() > 120 { - let truncated: String = stripped.chars().take(119).collect(); - format!("{truncated}…") + truncate_with_ellipsis(strip_goal_decoration(goal), PR_TITLE_MAX_CHARS) +} + +fn fallback_pr_title(goal: &str) -> String { + let title = pr_title_from_goal(goal); + if title.trim().is_empty() { + DEFAULT_PR_TITLE.to_string() } else { - stripped.to_string() + title } } @@ -286,70 +386,43 @@ async fn load_pull_request_diff(run_store: &RunStoreHandle) -> String { .unwrap_or_default() } -/// Build a complete PR body by combining LLM-generated narrative with -/// programmatic sections (plan, retro, fabro details). -pub async fn build_pr_body( +/// Build complete PR content by combining LLM-generated narrative with +/// deterministic fallbacks and programmatic sections. +pub async fn build_pr_content( diff: &str, goal: &str, model: &str, run_store: &RunStoreHandle, llm_source: &dyn CredentialSource, conclusion: Option<&Conclusion>, -) -> Result { + run_state: Option<&RunProjection>, +) -> Result { let client = Client::from_source(llm_source) .await .map_err(|e| format!("Failed to create LLM client: {e}"))?; - build_pr_body_with_client(diff, goal, model, run_store, conclusion, Arc::new(client)).await -} - -async fn build_pr_body_with_client( - diff: &str, - goal: &str, - model: &str, - run_store: &RunStoreHandle, - conclusion: Option<&Conclusion>, - client: Arc, -) -> Result { - build_pr_body_with_client_and_state(diff, goal, model, run_store, conclusion, client, None) - .await -} - -async fn build_pr_body_with_source_and_state( - diff: &str, - goal: &str, - model: &str, - run_store: &RunStoreHandle, - llm_source: &dyn CredentialSource, - conclusion: Option<&Conclusion>, - run_state: Option<&fabro_store::RunProjection>, -) -> Result { - let client = Client::from_source(llm_source) - .await - .map_err(|e| format!("Failed to create LLM client: {e}"))?; - - build_pr_body_with_client_and_state( + build_pr_content_with_client( diff, goal, model, run_store, conclusion, - Arc::new(client), run_state, + Arc::new(client), ) .await } -async fn build_pr_body_with_client_and_state( +async fn build_pr_content_with_client( diff: &str, goal: &str, model: &str, run_store: &RunStoreHandle, conclusion: Option<&Conclusion>, + run_state: Option<&RunProjection>, client: Arc, - run_state: Option<&fabro_store::RunProjection>, -) -> Result { - info!("Building PR body"); +) -> Result { + info!("Building PR content"); let loaded_run_state = if run_state.is_none() { run_store @@ -369,29 +442,11 @@ async fn build_pr_body_with_client_and_state( let run_spec = run_state.and_then(|state| state.spec.clone()); let dot_source = run_state.and_then(|state| state.graph_source.clone()); - // Build LLM prompt - let system = if plan_text.is_some() { - "Write a PR description with: (1) 2-3 concise paragraphs explaining the change, then (2) a '### Plan Summary' section with bullet points summarizing the plan. Do not include a title. Do not include the full plan.".to_string() - } else { - "Write a concise PR description in 2-3 paragraphs explaining the change. Do not include a title.".to_string() - }; - - // Truncate diff to fit context windows (~50k chars) - let max_diff_len = 50_000; - let truncated_diff = if diff.len() > max_diff_len { - &diff[..diff.floor_char_boundary(max_diff_len)] - } else { - diff - }; + let caps = truncation_caps(model); + let truncated_diff = truncate_chars(diff, caps.diff); let prompt = if let Some(ref plan) = plan_text { - // Truncate plan for LLM context (~20k chars) - let max_plan_len = 20_000; - let truncated_plan = if plan.len() > max_plan_len { - &plan[..plan.floor_char_boundary(max_plan_len)] - } else { - plan.as_str() - }; + let truncated_plan = truncate_chars(plan, caps.plan); format!( "Goal: {goal}\n\nPlan:\n```\n{truncated_plan}\n```\n\nDiff:\n```\n{truncated_diff}\n```" ) @@ -400,14 +455,32 @@ async fn build_pr_body_with_client_and_state( }; let params = GenerateParams::new(model, client) - .system(system) + .system(PR_BODY_SYSTEM_PROMPT) .prompt(prompt); - let result = generate(params) + let result = generate_object(params, PR_CONTENT_SCHEMA.clone()) .await .map_err(|e| format!("LLM generation failed: {e}"))?; - let llm_output = result.response.text(); + let output = result + .output + .ok_or_else(|| "LLM generation returned no structured output".to_string())?; + let generated: PrContent = serde_json::from_value(output) + .map_err(|e| format!("Failed to deserialize PR content: {e}"))?; + + let title = if generated.title.trim().is_empty() { + fallback_pr_title(goal) + } else { + generated.title.trim().to_string() + }; + let title = enforce_title_cap(&title); + + let llm_body = if generated.body.trim().is_empty() { + warn!(model = %model, "LLM generated empty PR body; using skeleton PR body"); + EMPTY_BODY_NOTICE.to_string() + } else { + generated.body + }; let retro_section = retro.as_ref().map(format_retro_section).unwrap_or_default(); let arc_details_section = conclusion @@ -416,15 +489,15 @@ async fn build_pr_body_with_client_and_state( .unwrap_or_default(); let body = assemble_pr_body( - &llm_output, + &llm_body, plan_text.as_deref(), &retro_section, &arc_details_section, ); - info!("PR body generated"); + info!("PR content generated"); - Ok(body) + Ok(PrContent { title, body }) } /// Auto-merge configuration for a pull request. @@ -446,7 +519,7 @@ pub struct OpenPullRequestRequest<'a> { pub run_store: &'a RunStoreHandle, pub llm_source: &'a dyn CredentialSource, pub conclusion: Option<&'a Conclusion>, - pub run_state: Option<&'a fabro_store::RunProjection>, + pub run_state: Option<&'a RunProjection>, } /// Optionally open a pull request after a successful workflow run. @@ -465,7 +538,7 @@ pub async fn maybe_open_pull_request( let (owner, repo) = github_app::parse_github_owner_repo(&https_url).map_err(|err| format!("{err:#}"))?; - let body = build_pr_body_with_source_and_state( + let content = build_pr_content( req.diff, req.goal, req.model, @@ -476,9 +549,8 @@ pub async fn maybe_open_pull_request( ) .await .map_err(|err| format!("{err:#}"))?; - let body = truncate_pr_body(&body); - - let title = pr_title_from_goal(req.goal); + let body = truncate_pr_body(&content.body); + let title = content.title; let created = github_app::create_pull_request( &req.github, @@ -782,6 +854,16 @@ mod tests { }) } + /// JSON string the MockProvider/openai mock returns to simulate the + /// structured-output response for `(title, body)`. + fn pr_content_json(title: &str, body: &str) -> String { + serde_json::to_string(&serde_json::json!({ + "title": title, + "body": body, + })) + .unwrap() + } + fn make_test_conclusion() -> Conclusion { Conclusion { timestamp: Utc::now(), @@ -1123,20 +1205,25 @@ mod tests { } #[tokio::test] - async fn build_pr_body_uses_in_memory_conclusion() { + async fn build_pr_content_uses_in_memory_conclusion() { let store = test_store(); let run_store = store.create_run(&fixtures::RUN_1).await.unwrap(); - let body = build_pr_body_with_client( + let PrContent { title, body } = build_pr_content_with_client( "diff --git a/src/lib.rs b/src/lib.rs\n+fn new_feature() {}\n", "Implement feature", "mock-model", &run_store.clone().into(), Some(&make_test_conclusion()), - explicit_client("mock", "Narrative from mock."), + None, + explicit_client( + "mock", + &pr_content_json("Mock title", "Narrative from mock."), + ), ) .await .unwrap(); + assert_eq!(title, "Mock title"); assert!(body.contains("Narrative from mock.")); assert!(body.contains("### Fabro Details")); assert!(body.contains("Ran 3 stages in 2m 30s for $0.42")); @@ -1144,7 +1231,7 @@ mod tests { } #[tokio::test] - async fn build_pr_body_uses_store_records_without_legacy_files() { + async fn build_pr_content_uses_store_records_without_legacy_files() { let store = test_store(); let run_store = store.create_run(&fixtures::RUN_1).await.unwrap(); @@ -1196,16 +1283,21 @@ mod tests { .await .unwrap(); - let body = build_pr_body_with_client( + let body = build_pr_content_with_client( "diff --git a/src/lib.rs b/src/lib.rs\n+fn new_feature() {}\n", "Implement feature", "mock-model", &run_store.clone().into(), Some(&make_test_conclusion()), - explicit_client("mock", "Narrative from mock."), + None, + explicit_client( + "mock", + &pr_content_json("Mock title", "Narrative from mock."), + ), ) .await - .unwrap(); + .unwrap() + .body; assert!(body.contains("Narrative from mock.")); assert!(body.contains("### Retro")); @@ -1214,7 +1306,7 @@ mod tests { } #[tokio::test] - async fn build_pr_body_uses_plan_text_from_store_without_response_md() { + async fn build_pr_content_uses_plan_text_from_store_without_response_md() { let store = test_store(); let run_store = store.create_run(&fixtures::RUN_1).await.unwrap(); @@ -1283,42 +1375,52 @@ mod tests { .await .unwrap(); - let body = build_pr_body_with_client( + let body = build_pr_content_with_client( "diff --git a/src/lib.rs b/src/lib.rs\n+fn new_feature() {}\n", "Implement feature", "mock-model", &run_store.clone().into(), Some(&make_test_conclusion()), - explicit_client("mock", "Narrative from mock."), + None, + explicit_client( + "mock", + &pr_content_json("Mock title", "Narrative from mock."), + ), ) .await - .unwrap(); + .unwrap() + .body; assert!(body.contains("Full plan")); assert!(body.contains("Plan from store")); } #[tokio::test] - async fn build_pr_body_uses_explicit_llm_client() { + async fn build_pr_content_uses_explicit_llm_client() { let store = test_store(); let run_store = store.create_run(&fixtures::RUN_1).await.unwrap(); - let body = build_pr_body_with_client( + let body = build_pr_content_with_client( "diff --git a/src/lib.rs b/src/lib.rs\n+fn new_feature() {}\n", "Implement feature", "gpt-5.4", &run_store.clone().into(), Some(&make_test_conclusion()), - explicit_client("openai", "Narrative from explicit client."), + None, + explicit_client( + "openai", + &pr_content_json("Explicit title", "Narrative from explicit client."), + ), ) .await - .unwrap(); + .unwrap() + .body; assert!(body.contains("Narrative from explicit client.")); assert!(!body.contains("Narrative from mock.")); } #[tokio::test] - async fn build_pr_body_uses_vault_only_openai_codex_source() { + async fn build_pr_content_uses_vault_only_openai_codex_source() { let server = MockServer::start_async().await; let response_mock = server .mock_async(|when, then| { @@ -1327,7 +1429,10 @@ mod tests { .header("authorization", "Bearer vault-openai-key"); then.status(200) .header("content-type", "application/json") - .json_body(openai_responses_payload("Narrative from vault source.")); + .json_body(openai_responses_payload(&pr_content_json( + "Vault title", + "Narrative from vault source.", + ))); }) .await; @@ -1355,17 +1460,19 @@ mod tests { let run_store = store.create_run(&fixtures::RUN_1).await.unwrap(); let run_store_handle: RunStoreHandle = run_store.into(); - let body = build_pr_body( + let PrContent { title, body } = build_pr_content( "diff --git a/src/lib.rs b/src/lib.rs\n+fn new_feature() {}\n", "Implement feature", "gpt-5.4", &run_store_handle, llm_source.as_ref(), Some(&make_test_conclusion()), + None, ) .await .unwrap(); + assert_eq!(title, "Vault title"); assert!(body.contains("Narrative from vault source.")); response_mock.assert_async().await; } @@ -1462,7 +1569,7 @@ mod tests { fn pr_title_truncates_long_line() { let long = "x".repeat(300); let title = pr_title_from_goal(&long); - assert_eq!(title.chars().count(), 120); + assert_eq!(title.chars().count(), 72); assert!(title.ends_with('…')); } @@ -1485,6 +1592,42 @@ mod tests { assert_eq!(pr_title_from_goal("Fix bug"), "Fix bug"); } + #[test] + fn truncation_caps_scale_with_context_window_and_clamp() { + assert_eq!( + truncation_caps_for_context_window(100_000), + TruncationCaps { + diff: 40_000, + plan: 10_000, + } + ); + assert_eq!( + truncation_caps_for_context_window(200_000), + TruncationCaps { + diff: 80_000, + plan: 20_000, + } + ); + assert_eq!( + truncation_caps_for_context_window(1_000_000), + TruncationCaps { + diff: 400_000, + plan: 100_000, + } + ); + assert_eq!( + truncation_caps_for_context_window(10_000_000), + TruncationCaps { + diff: 500_000, + plan: 100_000, + } + ); + assert_eq!(truncation_caps("unknown-model"), TruncationCaps { + diff: 80_000, + plan: 20_000, + }); + } + #[tokio::test] async fn empty_diff_returns_none() { let store = test_store(); @@ -1575,4 +1718,373 @@ mod tests { assert!(diff.contains("from_store")); } + + // ── Structured-output PR content tests ────────────────────────────── + + /// MockProvider returns an over-long title; builder must cap it at 72 + /// chars and end with `…`. Exercises [`enforce_title_cap`] inside + /// [`build_pr_content_with_client`]. + #[tokio::test] + async fn build_pr_content_truncates_long_title() { + let store = test_store(); + let run_store = store.create_run(&fixtures::RUN_1).await.unwrap(); + let long_title = "x".repeat(200); + let payload = pr_content_json(&long_title, "Body content."); + let title = build_pr_content_with_client( + "diff --git a/src/lib.rs b/src/lib.rs\n+fn x() {}\n", + "Implement feature", + "mock-model", + &run_store.clone().into(), + Some(&make_test_conclusion()), + None, + explicit_client("mock", &payload), + ) + .await + .unwrap() + .title; + + assert_eq!(title.chars().count(), 72); + assert!(title.ends_with('\u{2026}')); + } + + #[tokio::test] + async fn build_pr_content_uses_default_title_when_generated_and_goal_titles_empty() { + let store = test_store(); + let run_store = store.create_run(&fixtures::RUN_1).await.unwrap(); + let payload = pr_content_json("", "Body content."); + let title = build_pr_content_with_client( + "diff --git a/src/lib.rs b/src/lib.rs\n+fn x() {}\n", + "## Plan:", + "mock-model", + &run_store.clone().into(), + Some(&make_test_conclusion()), + None, + explicit_client("mock", &payload), + ) + .await + .unwrap() + .title; + + assert_eq!(title, DEFAULT_PR_TITLE); + } + + /// Empty or whitespace-only bodies use the skeleton fallback instead of + /// aborting PR creation. + #[tokio::test] + async fn build_pr_content_uses_skeleton_when_body_empty() { + let store = test_store(); + let run_store = store.create_run(&fixtures::RUN_1).await.unwrap(); + + let run_spec = RunSpec { + run_id: fixtures::RUN_1, + settings: fabro_types::WorkflowSettings::default(), + graph: Graph::new("test"), + workflow_slug: Some("test".to_string()), + source_directory: Some("/tmp/project".to_string()), + git: None, + labels: HashMap::new(), + provenance: None, + manifest_blob: None, + definition_blob: None, + fork_source_ref: None, + in_place: false, + }; + append_event(&run_store, &fixtures::RUN_1, &Event::RunCreated { + run_id: fixtures::RUN_1, + settings: serde_json::to_value(&run_spec.settings).unwrap(), + graph: serde_json::to_value(&run_spec.graph).unwrap(), + workflow_source: Some("digraph test { plan -> code }".to_string()), + workflow_config: None, + labels: run_spec.labels.clone().into_iter().collect(), + run_dir: "/tmp/project".to_string(), + source_directory: run_spec.source_directory.clone(), + workflow_slug: run_spec.workflow_slug.clone(), + db_prefix: None, + provenance: None, + manifest_blob: None, + git: None, + fork_source_ref: None, + in_place: false, + web_url: None, + }) + .await + .unwrap(); + append_event(&run_store, &fixtures::RUN_1, &Event::StageCompleted { + node_id: "plan".to_string(), + name: "plan".to_string(), + index: 0, + duration_ms: 1, + status: "succeeded".to_string(), + preferred_label: None, + suggested_next_ids: vec![], + billing: None, + failure: None, + notes: None, + files_touched: vec![], + context_updates: None, + jump_to_node: None, + context_values: None, + node_visits: None, + loop_failure_signatures: None, + restart_failure_signatures: None, + response: Some("Plan from store".to_string()), + attempt: 1, + max_attempts: 1, + }) + .await + .unwrap(); + append_event(&run_store, &fixtures::RUN_1, &Event::RetroCompleted { + duration_ms: 1, + response: Some(String::new()), + retro: Some(serde_json::to_value(make_test_retro()).unwrap()), + }) + .await + .unwrap(); + + let payload = pr_content_json("Mock", " \n"); + let body = build_pr_content_with_client( + "diff --git a/src/lib.rs b/src/lib.rs\n+fn x() {}\n", + "Implement feature", + "mock-model", + &run_store.clone().into(), + Some(&make_test_conclusion()), + None, + explicit_client("mock", &payload), + ) + .await + .unwrap() + .body; + + assert!(body.contains("The LLM did not produce a description")); + assert!(body.contains("Full plan")); + assert!(body.contains("Plan from store")); + assert!(body.contains("### Retro")); + assert!(body.contains("### Fabro Details")); + assert!(body.contains("Generated with [Fabro](https://fabro.sh)")); + } + + // ── maybe_open_pull_request fallback tests ────────────────────────── + + /// Set of mock servers and credentials for the `maybe_open_pull_request` + /// fallback path. The builder's `Client::from_source` rebuilds the LLM + /// client from the credential source, so the in-process MockProvider + /// cannot intercept — we mock the OpenAI HTTP endpoint instead. + struct FallbackHarness { + _vault_dir: tempfile::TempDir, + // Held to keep the mock listener alive for the duration of the test; + // the test interacts with it via `Client::from_source` (which goes + // out via HTTP to the mock URL stored in `llm_source`). + openai_server: MockServer, + github_server: MockServer, + openai_mock_id: usize, + github_mock_id: usize, + llm_source: Arc, + creds: fabro_github::GitHubCredentials, + run_store: RunStoreHandle, + } + + impl FallbackHarness { + async fn assert_mocks_called_once(&self) { + httpmock::Mock::new(self.openai_mock_id, &self.openai_server) + .assert_async() + .await; + httpmock::Mock::new(self.github_mock_id, &self.github_server) + .assert_async() + .await; + } + } + + /// Stand up an OpenAI mock that returns the given structured-output + /// payload, a GitHub mock that accepts a PR creation, a vault-backed + /// credential source, and a run store seeded with a non-empty + /// `final_patch`. + async fn setup_fallback_test_harness(openai_payload_text: &str) -> FallbackHarness { + let openai_server = MockServer::start_async().await; + let openai_mock = openai_server + .mock_async(|when, then| { + when.method(POST) + .path("/v1/responses") + .header("authorization", "Bearer vault-openai-key"); + then.status(200) + .header("content-type", "application/json") + .json_body(openai_responses_payload(openai_payload_text)); + }) + .await; + + let github_server = MockServer::start_async().await; + let github_mock = github_server + .mock_async(|when, then| { + when.method(POST) + .path("/repos/owner/repo/pulls") + .header("authorization", "Bearer test-token"); + then.status(201) + .header("content-type", "application/json") + .json_body(serde_json::json!({ + "number": 1, + "html_url": "https://example.test/owner/repo/pull/1", + "node_id": "PR_kwTest1", + })); + }) + .await; + + let vault_dir = tempfile::tempdir().unwrap(); + let mut vault = Vault::load(vault_dir.path().join("secrets.json")).unwrap(); + vault + .set( + "openai_codex", + &serde_json::to_string(&openai_api_key_credential("vault-openai-key")).unwrap(), + SecretType::Credential, + None, + ) + .unwrap(); + let base_url = openai_server.url("/v1"); + let llm_source: Arc = + Arc::new(VaultCredentialSource::with_env_lookup( + Arc::new(AsyncRwLock::new(vault)), + move |name| match name { + "OPENAI_BASE_URL" => Some(base_url.clone()), + _ => None, + }, + )); + + let creds = fabro_github::GitHubCredentials::Token("test-token".to_string()); + + let store = test_store(); + let run_store = store.create_run(&fixtures::RUN_1).await.unwrap(); + // Seed a non-empty `final_patch` so `load_pull_request_diff` returns + // diff content and the early-return for empty diffs does not fire. + let run_spec = RunSpec { + run_id: fixtures::RUN_1, + settings: fabro_types::WorkflowSettings::default(), + graph: Graph::new("test"), + workflow_slug: None, + source_directory: None, + git: None, + labels: HashMap::new(), + provenance: None, + manifest_blob: None, + definition_blob: None, + fork_source_ref: None, + in_place: false, + }; + append_event(&run_store, &fixtures::RUN_1, &Event::RunCreated { + run_id: fixtures::RUN_1, + settings: serde_json::to_value(&run_spec.settings).unwrap(), + graph: serde_json::to_value(&run_spec.graph).unwrap(), + workflow_source: None, + workflow_config: None, + labels: run_spec.labels.clone().into_iter().collect(), + run_dir: "/tmp/x".to_string(), + source_directory: None, + workflow_slug: None, + db_prefix: None, + provenance: None, + manifest_blob: None, + git: None, + fork_source_ref: None, + in_place: false, + web_url: None, + }) + .await + .unwrap(); + append_event(&run_store, &fixtures::RUN_1, &Event::WorkflowRunCompleted { + duration_ms: 1, + artifact_count: 0, + status: "succeeded".to_string(), + reason: SuccessReason::Completed, + total_usd_micros: None, + final_git_commit_sha: None, + final_patch: Some( + "diff --git a/src/lib.rs b/src/lib.rs\n+fn from_store() {}\n".to_string(), + ), + billing: None, + }) + .await + .unwrap(); + + let openai_mock_id = openai_mock.id; + let github_mock_id = github_mock.id; + + FallbackHarness { + _vault_dir: vault_dir, + openai_server, + github_server, + openai_mock_id, + github_mock_id, + llm_source, + creds, + run_store: run_store.into(), + } + } + + /// LLM returns a usable body but an empty title; the content builder + /// falls back to `pr_title_from_goal` (first line, decoration stripped) + /// and PR creation succeeds with that title. + #[tokio::test] + async fn maybe_open_pull_request_falls_back_to_goal_title_when_llm_returns_empty_title() { + let payload = pr_content_json("", "Narrative."); + let harness = setup_fallback_test_harness(&payload).await; + + let github_base_url = harness.github_server.url(""); + let github = github_app::GitHubContext::new(&harness.creds, &github_base_url); + + let result = maybe_open_pull_request(OpenPullRequestRequest { + github, + origin_url: "https://github.com/owner/repo.git", + base_branch: "main", + head_branch: "fabro/run/123", + goal: "Fix telemetry leak\n\ndetails...", + diff: "diff --git a/src/lib.rs b/src/lib.rs\n+fn x() {}\n", + model: "gpt-5.4", + draft: false, + auto_merge: None, + run_store: &harness.run_store, + llm_source: harness.llm_source.as_ref(), + conclusion: None, + run_state: None, + }) + .await + .expect("PR creation should succeed"); + + let record = result.expect("PR record should be Some"); + assert_eq!(record.title, "Fix telemetry leak"); + harness.assert_mocks_called_once().await; + } + + /// LLM returns an empty title; the content builder fallback still caps + /// the deterministic goal title at 72 chars ending with `…`. + #[tokio::test] + async fn maybe_open_pull_request_caps_fallback_title_at_72_chars() { + let payload = pr_content_json("", "Narrative."); + let harness = setup_fallback_test_harness(&payload).await; + + let github_base_url = harness.github_server.url(""); + let github = github_app::GitHubContext::new(&harness.creds, &github_base_url); + + // Single ~200-char line, no `Plan:` / heading prefix, no newlines. + let goal = "x".repeat(200); + + let result = maybe_open_pull_request(OpenPullRequestRequest { + github, + origin_url: "https://github.com/owner/repo.git", + base_branch: "main", + head_branch: "fabro/run/123", + goal: &goal, + diff: "diff --git a/src/lib.rs b/src/lib.rs\n+fn x() {}\n", + model: "gpt-5.4", + draft: false, + auto_merge: None, + run_store: &harness.run_store, + llm_source: harness.llm_source.as_ref(), + conclusion: None, + run_state: None, + }) + .await + .expect("PR creation should succeed"); + + let record = result.expect("PR record should be Some"); + assert_eq!(record.title.chars().count(), 72); + assert!(record.title.ends_with('\u{2026}')); + harness.assert_mocks_called_once().await; + } } diff --git a/lib/crates/fabro-workflow/src/pipeline/retro.rs b/lib/crates/fabro-workflow/src/pipeline/retro.rs index 649435c2e..d4c492373 100644 --- a/lib/crates/fabro-workflow/src/pipeline/retro.rs +++ b/lib/crates/fabro-workflow/src/pipeline/retro.rs @@ -53,7 +53,7 @@ pub async fn run_retro(options: &RetroOptions, dry_run: bool) -> Option { return None; } }; - let stage_durations = crate::extract_stage_durations_from_events(&events); + let stage_durations = crate::latest_stage_duration_by_node(&events); let mut retro = derive_retro( options.run_id, &options.workflow_name, diff --git a/lib/crates/fabro-workflow/src/pull_request.rs b/lib/crates/fabro-workflow/src/pull_request.rs index 3989bf63c..3b65dfcbc 100644 --- a/lib/crates/fabro-workflow/src/pull_request.rs +++ b/lib/crates/fabro-workflow/src/pull_request.rs @@ -1,4 +1,4 @@ pub use crate::pipeline::{ - AutoMergeOptions, OpenPullRequestRequest, PullRequestRecord, build_pr_body, + AutoMergeOptions, OpenPullRequestRequest, PrContent, PullRequestRecord, build_pr_content, maybe_open_pull_request, }; diff --git a/lib/crates/fabro-workflow/src/test_support.rs b/lib/crates/fabro-workflow/src/test_support.rs index 12aa40f10..0bf6377a6 100644 --- a/lib/crates/fabro-workflow/src/test_support.rs +++ b/lib/crates/fabro-workflow/src/test_support.rs @@ -17,7 +17,7 @@ use crate::handler::HandlerRegistry; use crate::outcome::Outcome; use crate::pipeline; use crate::pipeline::types::{Executed, Initialized}; -use crate::pipeline::{billing_from_checkpoint, build_terminal_event}; +use crate::pipeline::{billing_from_projection, build_terminal_event}; use crate::records::Checkpoint; use crate::run_metadata::RunMetadataRuntime; use crate::run_options::RunOptions; @@ -36,10 +36,7 @@ async fn execute_and_emit_terminal(initialized: InitializedState) -> Executed { let executed = Box::pin(pipeline::execute(initialized.initialized)).await; initialized.store_logger.flush().await; let state = executed.engine.run.run_store.state().await.ok(); - let billing = state - .as_ref() - .and_then(|s| s.checkpoint.as_ref()) - .and_then(billing_from_checkpoint); + let billing = state.as_ref().and_then(billing_from_projection); let event = build_terminal_event( &executed.outcome, executed.duration_ms, diff --git a/lib/crates/fabro-workflow/tests/it/integration.rs b/lib/crates/fabro-workflow/tests/it/integration.rs index 002d473da..f6ece331b 100644 --- a/lib/crates/fabro-workflow/tests/it/integration.rs +++ b/lib/crates/fabro-workflow/tests/it/integration.rs @@ -6809,7 +6809,13 @@ async fn workflow_run_with_vault_only_openai_codex_builds_pr_body() { .header("authorization", "Bearer vault-openai-key"); then.status(200) .header("content-type", "application/json") - .json_body(openai_responses_payload("Narrative from vault source.")); + .json_body(openai_responses_payload( + &serde_json::to_string(&serde_json::json!({ + "title": "Vault title", + "body": "Narrative from vault source.", + })) + .unwrap(), + )); }) .await; @@ -6890,7 +6896,7 @@ async fn workflow_run_with_vault_only_openai_codex_builds_pr_body() { let run_store = store.open_run_reader(&run_options.run_id).await.unwrap(); let run_store_handle: fabro_workflow::runtime_store::RunStoreHandle = run_store.into(); - let body = fabro_workflow::pull_request::build_pr_body( + let content = fabro_workflow::pull_request::build_pr_content( "diff --git a/src/lib.rs b/src/lib.rs\n+fn new_feature() {}\n", "Implement feature", "gpt-5.4", @@ -6906,11 +6912,13 @@ async fn workflow_run_with_vault_only_openai_codex_builds_pr_body() { billing: None, total_retries: 0, }), + None, ) .await .expect("PR body should build from vault-only credentials"); - assert!(body.contains("Narrative from vault source.")); + assert_eq!(content.title, "Vault title"); + assert!(content.body.contains("Narrative from vault source.")); response_mock.assert_async().await; } diff --git a/lib/packages/fabro-api-client/src/models/billing-by-model.ts b/lib/packages/fabro-api-client/src/models/billing-by-model.ts index d8d802575..74f074401 100644 --- a/lib/packages/fabro-api-client/src/models/billing-by-model.ts +++ b/lib/packages/fabro-api-client/src/models/billing-by-model.ts @@ -26,7 +26,7 @@ import type { ModelReference } from './model-reference'; export interface BillingByModel { 'model': ModelReference; /** - * Number of stages that used this model. + * Number of usage-bearing stage visits that used this model. */ 'stages': number; 'billing': BilledTokenCounts; diff --git a/lib/packages/fabro-api-client/src/models/billing-stage-ref.ts b/lib/packages/fabro-api-client/src/models/billing-stage-ref.ts index bcc94a656..7b87ef631 100644 --- a/lib/packages/fabro-api-client/src/models/billing-stage-ref.ts +++ b/lib/packages/fabro-api-client/src/models/billing-stage-ref.ts @@ -15,7 +15,7 @@ /** - * Reference to a billing stage. + * Reference to a workflow node in a billing stage row. */ export interface BillingStageRef { /** diff --git a/lib/packages/fabro-api-client/src/models/run-billing-stage.ts b/lib/packages/fabro-api-client/src/models/run-billing-stage.ts index aad11d6e6..acdf57c97 100644 --- a/lib/packages/fabro-api-client/src/models/run-billing-stage.ts +++ b/lib/packages/fabro-api-client/src/models/run-billing-stage.ts @@ -27,22 +27,19 @@ import type { ModelReference } from './model-reference'; import type { StageState } from './stage-state'; /** - * Token counts and billed totals for a single stage within a run. + * Token counts and billed totals for one workflow node within a run. Rows are grouped by node; billing and runtime sum every visit of that node. */ export interface RunBillingStage { 'stage': BillingStageRef; 'model': ModelReference | null; 'billing': BilledTokenCounts; /** - * Wall-clock runtime in seconds. + * Wall-clock runtime in seconds, summed across every visit of this node. */ 'runtime_secs': number; /** * Wall-clock time the latest attempt of this stage started, if known. */ 'started_at'?: string | null; - /** - * Lifecycle state of the stage. Use to detect in-flight rows for client-side runtime ticking. - */ 'state'?: StageState | null; } diff --git a/lib/packages/fabro-api-client/src/models/run-billing.ts b/lib/packages/fabro-api-client/src/models/run-billing.ts index 88544b1cc..09cfed9e3 100644 --- a/lib/packages/fabro-api-client/src/models/run-billing.ts +++ b/lib/packages/fabro-api-client/src/models/run-billing.ts @@ -28,7 +28,7 @@ import type { RunBillingTotals } from './run-billing-totals'; */ export interface RunBilling { /** - * Per-stage billing breakdown. + * Per-node billing breakdown. Each row sums billing and runtime across all visits of that node. */ 'stages': Array; 'totals': RunBillingTotals; diff --git a/lib/packages/fabro-api-client/src/models/run-stage.ts b/lib/packages/fabro-api-client/src/models/run-stage.ts index 2e5e82f0c..060a6166e 100644 --- a/lib/packages/fabro-api-client/src/models/run-stage.ts +++ b/lib/packages/fabro-api-client/src/models/run-stage.ts @@ -22,7 +22,7 @@ import type { StageState } from './stage-state'; */ export interface RunStage { /** - * Unique stage identifier within the run. + * StageId in \"node_id@visit\" form, e.g. verify@2. */ 'id': string; /** @@ -35,13 +35,15 @@ export interface RunStage { */ 'duration_secs'?: number; /** - * Node identifier in the Graphviz graph source. + * Node id in the workflow graph; multiple stages with different visits share the same node_id. */ - 'dot_id'?: string; + 'node_id': string; + /** + * 1-based visit count; bumped each time the workflow re-enters this node. + */ + 'visit': number; /** * Wall-clock time the latest attempt of this stage started, if known. */ 'started_at'?: string | null; } - - diff --git a/lib/packages/fabro-api-client/src/models/stage-projection.ts b/lib/packages/fabro-api-client/src/models/stage-projection.ts index 7f18d6bc4..139fdd437 100644 --- a/lib/packages/fabro-api-client/src/models/stage-projection.ts +++ b/lib/packages/fabro-api-client/src/models/stage-projection.ts @@ -63,9 +63,5 @@ export interface StageProjection { * Wall-clock duration of the stage\'s latest terminal attempt, if known. */ 'duration_ms'?: number | null; - /** - * Lifecycle state of the stage projection. - */ 'state'?: StageState | null; } -