@wichayutdew/pi-workflows 3.0.0 → 3.2.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@wichayutdew/pi-workflows",
3
- "version": "3.0.0",
3
+ "version": "3.2.0",
4
4
  "description": "A declarative, pauseable workflow harness for Pi",
5
5
  "license": "MIT",
6
6
  "repository": {
@@ -33,6 +33,7 @@ const completedStep = (
33
33
  ...(run.currentStepOmittedAttempts
34
34
  ? { omittedAttempts: run.currentStepOmittedAttempts }
35
35
  : {}),
36
+ ...(run.currentStepUsage ? { usage: run.currentStepUsage } : {}),
36
37
  completedAt: now,
37
38
  });
38
39
 
@@ -126,6 +127,7 @@ export const advanceRun = (
126
127
  history: [...run.history, completed],
127
128
  currentStepAttempts: undefined,
128
129
  currentStepOmittedAttempts: undefined,
130
+ currentStepUsage: undefined,
129
131
  ...(cwd ? { cwd } : {}),
130
132
  ...(effects.workspaceCwd ? { restartWorkspaceCwd: undefined } : {}),
131
133
  stepHandoff: summary,
@@ -175,6 +177,7 @@ export const advanceRun = (
175
177
  history: [...run.history, completed],
176
178
  currentStepAttempts: undefined,
177
179
  currentStepOmittedAttempts: undefined,
180
+ currentStepUsage: undefined,
178
181
  ...(cwd ? { cwd } : {}),
179
182
  ...(effects.workspaceCwd ? { restartWorkspaceCwd: undefined } : {}),
180
183
  stepHandoff: preservesGateRevisionContext ? run.stepHandoff : summary,
@@ -100,6 +100,7 @@ export const reconcileRun = (
100
100
  history: retainedHistory,
101
101
  currentStepAttempts: changedEntry.attempts,
102
102
  currentStepOmittedAttempts: changedEntry.omittedAttempts,
103
+ currentStepUsage: changedEntry.usage,
103
104
  visits: rebuildVisits(retainedHistory, restartedStep),
104
105
  cwd: retainedWorkspaceCwd(reconciledRun, retainedHistory),
105
106
  reviewedArtifact: reviewedApproval?.artifact ?? '',
@@ -22,6 +22,12 @@ import {
22
22
  type WorkflowRunStatus,
23
23
  } from './state-types.ts';
24
24
  import { workflowTraceChars } from './step-trace.ts';
25
+ import {
26
+ emptyUsageAggregate,
27
+ isUsageAggregate,
28
+ mergeUsage,
29
+ type UsageAggregate,
30
+ } from './usage.ts';
25
31
 
26
32
  type UnknownRecord = Readonly<Record<string, unknown>>;
27
33
 
@@ -140,6 +146,7 @@ const isStepExecutionAttempt = (
140
146
  (value.taskTruncated === true) !==
141
147
  (typeof value.omittedTaskChars === 'number') ||
142
148
  typeof value.startedAt !== 'number' ||
149
+ (value.usage !== undefined && !isUsageAggregate(value.usage)) ||
143
150
  (value.result !== undefined && !isStepAttemptResult(value.result)) ||
144
151
  (value.gateDecision !== undefined &&
145
152
  !isStepGateDecision(value.gateDecision))
@@ -202,6 +209,20 @@ const isStepExecutionAttempts = (
202
209
  );
203
210
  });
204
211
 
212
+ const usageMatchesAttempts = (
213
+ attempts: ReadonlyArray<StepExecutionAttempt> | undefined,
214
+ aggregate: UsageAggregate | undefined,
215
+ omittedAttempts: unknown,
216
+ ): boolean => {
217
+ if (!aggregate || omittedAttempts !== undefined) return true;
218
+ const fromAttempts = (attempts ?? []).reduce(
219
+ (total, attempt) =>
220
+ attempt.usage ? mergeUsage(total, attempt.usage.models) : total,
221
+ emptyUsageAggregate(),
222
+ );
223
+ return JSON.stringify(fromAttempts) === JSON.stringify(aggregate);
224
+ };
225
+
205
226
  const isStepHistoryEntry = (value: unknown): value is StepHistoryEntry =>
206
227
  isRecord(value) &&
207
228
  typeof value.stepId === 'string' &&
@@ -215,8 +236,11 @@ const isStepHistoryEntry = (value: unknown): value is StepHistoryEntry =>
215
236
  value.artifact === value.approval.artifact)) &&
216
237
  (value.attempts === undefined || isStepExecutionAttempts(value.attempts)) &&
217
238
  (value.omittedAttempts === undefined ||
218
- (Number.isSafeInteger(value.omittedAttempts) &&
219
- (value.omittedAttempts as number) > 0)) &&
239
+ (typeof value.omittedAttempts === 'number' &&
240
+ Number.isSafeInteger(value.omittedAttempts) &&
241
+ value.omittedAttempts > 0)) &&
242
+ (value.usage === undefined || isUsageAggregate(value.usage)) &&
243
+ usageMatchesAttempts(value.attempts, value.usage, value.omittedAttempts) &&
220
244
  typeof value.completedAt === 'number';
221
245
 
222
246
  const isGateResolution = (value: unknown): value is GateResolution =>
@@ -336,6 +360,8 @@ export const isWorkflowRun = (value: unknown): value is WorkflowRun => {
336
360
  (value.currentStepOmittedAttempts === undefined ||
337
361
  (Number.isSafeInteger(value.currentStepOmittedAttempts) &&
338
362
  (value.currentStepOmittedAttempts as number) > 0)) &&
363
+ (value.currentStepUsage === undefined ||
364
+ isUsageAggregate(value.currentStepUsage)) &&
339
365
  isVisitCounts(value.visits) &&
340
366
  typeof value.startedAt === 'number' &&
341
367
  typeof value.updatedAt === 'number' &&
@@ -362,6 +388,16 @@ export const isWorkflowRun = (value: unknown): value is WorkflowRun => {
362
388
  value.pausedFrom === 'awaiting-gate');
363
389
  if (!hasValidOptionalFields) return false;
364
390
 
391
+ if (
392
+ !usageMatchesAttempts(
393
+ value.currentStepAttempts as
394
+ ReadonlyArray<StepExecutionAttempt> | undefined,
395
+ value.currentStepUsage as UsageAggregate | undefined,
396
+ value.currentStepOmittedAttempts,
397
+ )
398
+ )
399
+ return false;
400
+
365
401
  const pendingGate = value.pendingGate;
366
402
  if (pendingGate !== undefined && !isPendingGate(pendingGate)) return false;
367
403
 
@@ -1,3 +1,5 @@
1
+ import type { UsageAggregate } from './usage.ts';
2
+
1
3
  export const RUN_STATE_VERSION = 1 as const;
2
4
  export const MAX_GATE_FEEDBACK_CHARS = 50_000;
3
5
  export const MAX_RESUME_INPUT_CHARS = 16_000;
@@ -67,6 +69,8 @@ export type StepExecutionAttempt =
67
69
  readonly logTruncated?: true | undefined;
68
70
  readonly omittedLogEvents?: number | undefined;
69
71
  readonly startedAt: number;
72
+ /** Finalized Pi usage for this exact attempt, if available. */
73
+ readonly usage?: UsageAggregate | undefined;
70
74
  readonly result?: StepAttemptResult | undefined;
71
75
  readonly gateDecision?: StepGateDecision | undefined;
72
76
  }
@@ -81,6 +85,8 @@ export type StepExecutionAttempt =
81
85
  readonly taskTruncated?: true | undefined;
82
86
  readonly omittedTaskChars?: number | undefined;
83
87
  readonly startedAt: number;
88
+ /** Finalized Pi usage for this exact attempt, if available. */
89
+ readonly usage?: UsageAggregate | undefined;
84
90
  readonly transcript?: SubagentTranscriptReference | undefined;
85
91
  readonly result?: StepAttemptResult | undefined;
86
92
  readonly gateDecision?: StepGateDecision | undefined;
@@ -101,6 +107,8 @@ export type StepHistoryEntry = {
101
107
  readonly attempts?: ReadonlyArray<StepExecutionAttempt> | undefined;
102
108
  /** Older attempts compacted to keep the checkpoint bounded. */
103
109
  readonly omittedAttempts?: number | undefined;
110
+ /** Aggregate usage for this step, including compacted attempts. */
111
+ readonly usage?: UsageAggregate | undefined;
104
112
  readonly completedAt: number;
105
113
  };
106
114
 
@@ -142,6 +150,8 @@ export type WorkflowRun = {
142
150
  ReadonlyArray<StepExecutionAttempt> | undefined;
143
151
  /** Older current-step attempts compacted from the checkpoint. */
144
152
  readonly currentStepOmittedAttempts?: number | undefined;
153
+ /** Aggregate usage for the current step, including compacted attempts. */
154
+ readonly currentStepUsage?: UsageAggregate | undefined;
145
155
  readonly startedAt: number;
146
156
  readonly updatedAt: number;
147
157
  /**
@@ -11,6 +11,15 @@ export {
11
11
  MAX_WORKFLOW_TRACE_CHARS,
12
12
  RUN_STATE_VERSION,
13
13
  } from './state-types.ts';
14
+ export {
15
+ addUsage,
16
+ emptyUsage,
17
+ emptyUsageAggregate,
18
+ isUsageAggregate,
19
+ isUsageTotals,
20
+ mergeUsage,
21
+ normalizeUsage,
22
+ } from './usage.ts';
14
23
  export type {
15
24
  GateApproval,
16
25
  GateResolution,
@@ -23,3 +32,4 @@ export type {
23
32
  WorkflowRun,
24
33
  WorkflowRunStatus,
25
34
  } from './state-types.ts';
35
+ export type { ModelUsage, UsageAggregate, UsageTotals } from './usage.ts';
@@ -1,6 +1,11 @@
1
1
  import { isAbsolute, relative, resolve, sep } from 'node:path';
2
2
  import type { WorkflowStepResult } from '../runtime/step-result.ts';
3
3
  import { redactStepLogText } from '../step-log.ts';
4
+ import {
5
+ emptyUsageAggregate,
6
+ mergeUsage,
7
+ type UsageAggregate,
8
+ } from './usage.ts';
4
9
  import {
5
10
  MAX_STEP_TRACE_ARTIFACT_CHARS,
6
11
  MAX_STEP_TRACE_ATTEMPTS,
@@ -52,7 +57,8 @@ function attemptSize(attempt: StepExecutionAttempt): number {
52
57
  : 0) +
53
58
  (attempt.gateDecision?.requestId.length ?? 0) +
54
59
  (attempt.gateDecision?.feedback.length ?? 0) +
55
- (attempt.gateDecision?.reviewId?.length ?? 0)
60
+ (attempt.gateDecision?.reviewId?.length ?? 0) +
61
+ (attempt.usage ? JSON.stringify(attempt.usage).length : 0)
56
62
  );
57
63
  }
58
64
 
@@ -122,10 +128,18 @@ function compactAttempt(attempt: StepExecutionAttempt): StepExecutionAttempt {
122
128
 
123
129
  /** Returns the bounded checkpoint payload attributable to step traces. */
124
130
  export function workflowTraceChars(run: WorkflowRun): number {
125
- return [
131
+ const attempts = [
126
132
  ...run.history.flatMap((entry) => entry.attempts ?? []),
127
133
  ...(run.currentStepAttempts ?? []),
128
134
  ].reduce((total, attempt) => total + attemptSize(attempt), 0);
135
+ const aggregates = [
136
+ ...run.history.map((entry) => entry.usage),
137
+ run.currentStepUsage,
138
+ ].reduce(
139
+ (total, usage) => total + (usage ? JSON.stringify(usage).length : 0),
140
+ 0,
141
+ );
142
+ return attempts + aggregates;
129
143
  }
130
144
 
131
145
  function compactRunTraceBudget(run: WorkflowRun): WorkflowRun {
@@ -441,6 +455,45 @@ function attemptResult(
441
455
  }
442
456
 
443
457
  /** Stores the submitted result on the latest attempt before transition. */
458
+ /**
459
+ * Attaches finalized usage to its exact attempt. The separate current
460
+ * aggregate deliberately survives trace eviction and is copied to history.
461
+ */
462
+ export function recordCurrentStepUsage(
463
+ run: WorkflowRun,
464
+ requestId: string,
465
+ usage: UsageAggregate,
466
+ now: number,
467
+ ): WorkflowRun {
468
+ const attempts = run.currentStepAttempts;
469
+ const index = attempts?.findIndex(
470
+ (attempt) => attempt.requestId === requestId,
471
+ );
472
+ if (index === undefined || index < 0 || !attempts) return run;
473
+ const attempt = attempts[index];
474
+ if (!attempt) return run;
475
+ const currentStepAttempts = [...attempts];
476
+ currentStepAttempts[index] = {
477
+ ...attempt,
478
+ usage: mergeUsage(attempt.usage ?? emptyUsageAggregate(), usage.models),
479
+ };
480
+ return compactRunTraceBudget({
481
+ ...run,
482
+ currentStepAttempts,
483
+ currentStepUsage: mergeUsage(
484
+ run.currentStepUsage ?? emptyUsageAggregate(),
485
+ usage.models,
486
+ ),
487
+ updatedAt: now,
488
+ });
489
+ }
490
+
491
+ export function usageAggregateFromModels(
492
+ entries: UsageAggregate['models'],
493
+ ): UsageAggregate {
494
+ return mergeUsage(emptyUsageAggregate(), entries);
495
+ }
496
+
444
497
  export function recordCurrentStepResult(
445
498
  run: WorkflowRun,
446
499
  result: WorkflowStepResult,
@@ -0,0 +1,312 @@
1
+ export type UsageTotals = {
2
+ readonly inputTokens: number;
3
+ readonly outputTokens: number;
4
+ readonly cacheReadTokens: number;
5
+ readonly cacheWriteTokens: number;
6
+ readonly inputCostUsd: number;
7
+ readonly outputCostUsd: number;
8
+ readonly cacheReadCostUsd: number;
9
+ readonly cacheWriteCostUsd: number;
10
+ /** Provider-reported cost without a token-category breakdown. */
11
+ readonly otherCostUsd: number;
12
+ readonly totalCostUsd: number;
13
+ };
14
+
15
+ export type ModelUsage = {
16
+ readonly provider: string;
17
+ readonly model: string;
18
+ readonly usage: UsageTotals;
19
+ };
20
+
21
+ export type UsageAggregate = {
22
+ readonly usage: UsageTotals;
23
+ readonly models: ReadonlyArray<ModelUsage>;
24
+ };
25
+
26
+ type UnknownRecord = Readonly<Record<string, unknown>>;
27
+ const isRecord = (value: unknown): value is UnknownRecord =>
28
+ value !== null && typeof value === 'object' && !Array.isArray(value);
29
+
30
+ export const emptyUsage = (): UsageTotals => ({
31
+ inputTokens: 0,
32
+ outputTokens: 0,
33
+ cacheReadTokens: 0,
34
+ cacheWriteTokens: 0,
35
+ inputCostUsd: 0,
36
+ outputCostUsd: 0,
37
+ cacheReadCostUsd: 0,
38
+ cacheWriteCostUsd: 0,
39
+ otherCostUsd: 0,
40
+ totalCostUsd: 0,
41
+ });
42
+
43
+ export const emptyUsageAggregate = (): UsageAggregate => ({
44
+ usage: emptyUsage(),
45
+ models: [],
46
+ });
47
+
48
+ const fields = [
49
+ 'inputTokens',
50
+ 'outputTokens',
51
+ 'cacheReadTokens',
52
+ 'cacheWriteTokens',
53
+ 'inputCostUsd',
54
+ 'outputCostUsd',
55
+ 'cacheReadCostUsd',
56
+ 'cacheWriteCostUsd',
57
+ 'otherCostUsd',
58
+ 'totalCostUsd',
59
+ ] as const satisfies ReadonlyArray<keyof UsageTotals>;
60
+
61
+ function finiteNonNegative(value: unknown): number | undefined {
62
+ return typeof value === 'number' && Number.isFinite(value) && value >= 0
63
+ ? value
64
+ : undefined;
65
+ }
66
+
67
+ function numberAt(
68
+ value: UnknownRecord,
69
+ ...keys: ReadonlyArray<string>
70
+ ): number {
71
+ for (const key of keys) {
72
+ const candidate = finiteNonNegative(value[key]);
73
+ if (candidate !== undefined) return candidate;
74
+ }
75
+ return 0;
76
+ }
77
+
78
+ function costNumber(
79
+ value: UnknownRecord,
80
+ costObj: UnknownRecord | undefined,
81
+ flatKeys: ReadonlyArray<string>,
82
+ nestedKey: string,
83
+ ): number {
84
+ for (const key of flatKeys) {
85
+ const candidate = finiteNonNegative(value[key]);
86
+ if (candidate !== undefined) return candidate;
87
+ }
88
+ if (costObj) {
89
+ const candidate = finiteNonNegative(costObj[nestedKey]);
90
+ if (candidate !== undefined) return candidate;
91
+ }
92
+ return 0;
93
+ }
94
+
95
+ function hasMalformedNumber(value: UnknownRecord): boolean {
96
+ for (const candidate of Object.values(value)) {
97
+ if (
98
+ typeof candidate === 'number' &&
99
+ (!Number.isFinite(candidate) || candidate < 0)
100
+ ) {
101
+ return true;
102
+ }
103
+ if (isRecord(candidate)) {
104
+ for (const nested of Object.values(candidate)) {
105
+ if (
106
+ typeof nested === 'number' &&
107
+ (!Number.isFinite(nested) || nested < 0)
108
+ ) {
109
+ return true;
110
+ }
111
+ }
112
+ }
113
+ }
114
+ return false;
115
+ }
116
+
117
+ /** Normalizes Pi usage shapes while rejecting malformed or inconsistent totals. */
118
+ export function normalizeUsage(value: unknown): UsageTotals | undefined {
119
+ if (!isRecord(value)) return undefined;
120
+ if (hasMalformedNumber(value)) return undefined;
121
+
122
+ const costObj = isRecord(value.cost) ? value.cost : undefined;
123
+ const flatTotal =
124
+ finiteNonNegative(value.totalCostUsd) ??
125
+ (typeof value.cost === 'number'
126
+ ? finiteNonNegative(value.cost)
127
+ : undefined) ??
128
+ finiteNonNegative(value.costUsd);
129
+ const nestedTotal = costObj ? finiteNonNegative(costObj.total) : undefined;
130
+ const hasTotal = flatTotal !== undefined || nestedTotal !== undefined;
131
+ const totalCostUsd = flatTotal ?? nestedTotal ?? 0;
132
+
133
+ let usage: UsageTotals = {
134
+ inputTokens: numberAt(value, 'inputTokens', 'input'),
135
+ outputTokens: numberAt(value, 'outputTokens', 'output'),
136
+ cacheReadTokens: numberAt(value, 'cacheReadTokens', 'cacheRead'),
137
+ cacheWriteTokens: numberAt(value, 'cacheWriteTokens', 'cacheWrite'),
138
+ inputCostUsd: costNumber(
139
+ value,
140
+ costObj,
141
+ ['inputCostUsd', 'inputCost'],
142
+ 'input',
143
+ ),
144
+ outputCostUsd: costNumber(
145
+ value,
146
+ costObj,
147
+ ['outputCostUsd', 'outputCost'],
148
+ 'output',
149
+ ),
150
+ cacheReadCostUsd: costNumber(
151
+ value,
152
+ costObj,
153
+ ['cacheReadCostUsd', 'cacheReadCost'],
154
+ 'cacheRead',
155
+ ),
156
+ cacheWriteCostUsd: costNumber(
157
+ value,
158
+ costObj,
159
+ ['cacheWriteCostUsd', 'cacheWriteCost'],
160
+ 'cacheWrite',
161
+ ),
162
+ otherCostUsd: 0,
163
+ totalCostUsd,
164
+ };
165
+ const recognized = [
166
+ 'inputTokens',
167
+ 'input',
168
+ 'outputTokens',
169
+ 'output',
170
+ 'cacheReadTokens',
171
+ 'cacheRead',
172
+ 'cacheWriteTokens',
173
+ 'cacheWrite',
174
+ 'inputCostUsd',
175
+ 'inputCost',
176
+ 'outputCostUsd',
177
+ 'outputCost',
178
+ 'cacheReadCostUsd',
179
+ 'cacheReadCost',
180
+ 'cacheWriteCostUsd',
181
+ 'cacheWriteCost',
182
+ 'totalCostUsd',
183
+ 'cost',
184
+ 'costUsd',
185
+ ].some((key) => key in value);
186
+ if (!recognized) return undefined;
187
+ const componentCost =
188
+ usage.inputCostUsd +
189
+ usage.outputCostUsd +
190
+ usage.cacheReadCostUsd +
191
+ usage.cacheWriteCostUsd;
192
+ if (hasTotal) {
193
+ if (componentCost > usage.totalCostUsd + 1e-9) return undefined;
194
+ usage = { ...usage, otherCostUsd: usage.totalCostUsd - componentCost };
195
+ } else {
196
+ usage = { ...usage, totalCostUsd: componentCost };
197
+ }
198
+ return usage;
199
+ }
200
+
201
+ export function isUsageTotals(value: unknown): value is UsageTotals {
202
+ if (
203
+ !isRecord(value) ||
204
+ !fields.every((field) => finiteNonNegative(value[field]) !== undefined)
205
+ )
206
+ return false;
207
+ return (
208
+ Math.abs(
209
+ (value.inputCostUsd as number) +
210
+ (value.outputCostUsd as number) +
211
+ (value.cacheReadCostUsd as number) +
212
+ (value.cacheWriteCostUsd as number) +
213
+ (value.otherCostUsd as number) -
214
+ (value.totalCostUsd as number),
215
+ ) <= 1e-9
216
+ );
217
+ }
218
+
219
+ export function addUsage(left: UsageTotals, right: UsageTotals): UsageTotals {
220
+ return Object.fromEntries(
221
+ fields.map((field) => [field, left[field] + right[field]]),
222
+ ) as UsageTotals;
223
+ }
224
+
225
+ export function mergeUsage(
226
+ aggregate: UsageAggregate,
227
+ entries: ReadonlyArray<ModelUsage>,
228
+ ): UsageAggregate {
229
+ const byModel = new Map(
230
+ aggregate.models.map((entry) => [
231
+ `${entry.provider}\0${entry.model}`,
232
+ entry.usage,
233
+ ]),
234
+ );
235
+ for (const entry of entries) {
236
+ if (!entry.provider || !entry.model || !isUsageTotals(entry.usage))
237
+ continue;
238
+ const key = `${entry.provider}\0${entry.model}`;
239
+ byModel.set(key, addUsage(byModel.get(key) ?? emptyUsage(), entry.usage));
240
+ }
241
+ const models = [...byModel.entries()]
242
+ .map(([key, usage]) => {
243
+ const [provider, model] = key.split('\0');
244
+ return { provider: provider ?? '', model: model ?? '', usage };
245
+ })
246
+ .sort((left, right) =>
247
+ `${left.provider}/${left.model}`.localeCompare(
248
+ `${right.provider}/${right.model}`,
249
+ ),
250
+ );
251
+ return {
252
+ usage: models.reduce(
253
+ (total, entry) => addUsage(total, entry.usage),
254
+ emptyUsage(),
255
+ ),
256
+ models,
257
+ };
258
+ }
259
+
260
+ export function isUsageAggregate(value: unknown): value is UsageAggregate {
261
+ const usage =
262
+ value && typeof value === 'object'
263
+ ? (value as UnknownRecord).usage
264
+ : undefined;
265
+ const models =
266
+ value && typeof value === 'object'
267
+ ? (value as UnknownRecord).models
268
+ : undefined;
269
+ if (!isRecord(value) || !isUsageTotals(usage) || !Array.isArray(models))
270
+ return false;
271
+ if (
272
+ !models.every(
273
+ (entry: unknown) =>
274
+ isRecord(entry) &&
275
+ typeof entry.provider === 'string' &&
276
+ entry.provider.length > 0 &&
277
+ typeof entry.model === 'string' &&
278
+ entry.model.length > 0 &&
279
+ isUsageTotals(entry.usage),
280
+ )
281
+ )
282
+ return false;
283
+ if (
284
+ new Set(
285
+ models.map(
286
+ (entry) =>
287
+ `${(entry as ModelUsage).provider}\0${(entry as ModelUsage).model}`,
288
+ ),
289
+ ).size !== models.length
290
+ )
291
+ return false;
292
+ const typedModels = models as ReadonlyArray<ModelUsage>;
293
+ const total = typedModels.reduce(
294
+ (sum, entry) => addUsage(sum, entry.usage),
295
+ emptyUsage(),
296
+ );
297
+ return fields.every((field) => Math.abs(total[field] - usage[field]) <= 1e-9);
298
+ }
299
+
300
+ /** Extracts one model-keyed usage record from a Pi message-like payload. */
301
+ export function modelUsageFromMessage(
302
+ value: unknown,
303
+ fallbackProvider?: string,
304
+ fallbackModel?: string,
305
+ ): ModelUsage | undefined {
306
+ if (!isRecord(value)) return undefined;
307
+ const usage = normalizeUsage(value.usage);
308
+ const provider =
309
+ typeof value.provider === 'string' ? value.provider : fallbackProvider;
310
+ const model = typeof value.model === 'string' ? value.model : fallbackModel;
311
+ return usage && provider && model ? { provider, model, usage } : undefined;
312
+ }
@@ -19,6 +19,7 @@ import type {
19
19
  import type { MainStepRuntimeController } from '../runtime/main-step-runtime.ts';
20
20
  import type { SerialTaskQueueController } from '../runtime/serial-task-queue.ts';
21
21
  import type { WorkflowStepResult } from '../runtime/step-result.ts';
22
+ import type { ModelUsage } from '../engine/usage.ts';
22
23
  import type { WorkflowStatusSnapshot } from '../workflow-status.ts';
23
24
  import type { WorkflowHarnessDependencies } from './dependencies.ts';
24
25
  import type { SettledStepReport } from './step-reporting.ts';
@@ -104,11 +105,13 @@ export type HarnessActionContext = {
104
105
  identity: MainStepIdentity,
105
106
  lines: ReadonlyArray<string>,
106
107
  context: ExtensionContext,
108
+ usage?: ReadonlyArray<ModelUsage>,
107
109
  ) => Promise<void>;
108
110
  recordMainStepLog: (
109
111
  identity: MainStepIdentity,
110
112
  lines: ReadonlyArray<string>,
111
113
  context: ExtensionContext,
114
+ usage?: ReadonlyArray<ModelUsage>,
112
115
  ) => Promise<void>;
113
116
  queueMainStepResult: (
114
117
  identity: MainStepIdentity,
@@ -4,7 +4,11 @@ import {
4
4
  type SubagentDelegationUpdate,
5
5
  } from '../integrations/subagents/protocol.ts';
6
6
  import { advanceRun } from '../engine/transitions.ts';
7
- import { recordCurrentStepResult } from '../engine/step-trace.ts';
7
+ import {
8
+ recordCurrentStepResult,
9
+ recordCurrentStepUsage,
10
+ usageAggregateFromModels,
11
+ } from '../engine/step-trace.ts';
8
12
  import type { WorkflowStepResult } from '../runtime/step-result.ts';
9
13
  import type { HarnessActionContext as FullHarnessActionContext } from './action-context.ts';
10
14
  import type { ActiveDelegation } from './types.ts';
@@ -23,6 +27,7 @@ type HarnessActionContext = Pick<
23
27
  | 'launchCurrentStep'
24
28
  | 'mutationQueue'
25
29
  | 'pauseForDelegationFailure'
30
+ | 'persist'
26
31
  | 'releaseMainAfterCancellation'
27
32
  | 'retainUnconfirmedDelegation'
28
33
  | 'run'
@@ -150,7 +155,22 @@ async function finishDelegation(
150
155
  ) {
151
156
  return;
152
157
  }
158
+ if (response.requestId !== active.requestId) {
159
+ throw new Error(
160
+ 'Workflow worker returned an uncorrelated terminal response',
161
+ );
162
+ }
153
163
  const terminalAt = this.dependencies.now();
164
+ if (response.usage && response.usage.length > 0) {
165
+ this.run = recordCurrentStepUsage(
166
+ this.run,
167
+ active.requestId,
168
+ usageAggregateFromModels(response.usage),
169
+ terminalAt,
170
+ );
171
+ this.persist();
172
+ this.updateStatus();
173
+ }
154
174
  const workflow = this.catalog.workflows.get(this.run.workflowId);
155
175
  const step = workflow?.definition.steps[this.run.currentStepId];
156
176
  if (!workflow || !step) {