From 0db3d4b75092055ce7f07a40ce73680419cf021b Mon Sep 17 00:00:00 2001 From: zhangweijian Date: Thu, 3 Sep 2026 02:01:50 +0800 Subject: [PATCH 1/3] fix(zcode): retry JSON-less role turns in-session before failing A successful turn whose final response contains no JSON object used to terminalize the whole attempt as invalid_structured_output. High-effort models periodically reply with a prose report instead, so the harness now sends a schema-repeating correction message into the same session and waits for the next completion, up to two times, before failing. The retry counter is ephemeral: a restart already abandons the in-flight turn, so no cursor migration is needed. --- src/agents/zcode/harness.ts | 52 +++++++++++++++++++++- src/agents/zcode/prompts.ts | 20 +++++++++ test/agents/zcode/harness.test.ts | 72 ++++++++++++++++++++++++++++++- 3 files changed, 142 insertions(+), 2 deletions(-) diff --git a/src/agents/zcode/harness.ts b/src/agents/zcode/harness.ts index fb5fff6..53f2717 100644 --- a/src/agents/zcode/harness.ts +++ b/src/agents/zcode/harness.ts @@ -17,7 +17,12 @@ import { restoreApprovedSourceCommit, } from "../source-commit"; import type { ZcodeClientApi } from "./client"; -import { implementPrompt, reviewPrompt, scoutPrompt } from "./prompts"; +import { + implementPrompt, + reviewPrompt, + scoutPrompt, + structuredOutputRetryPrompt, +} from "./prompts"; import { classifyZcodeTurnFailure, mapZcodeUsage, @@ -61,8 +66,12 @@ type ActiveAttempt = { reviewStatusBefore?: string; reconciledCompletion?: boolean; pendingDeliveries: HarnessDelivery[]; + structuredRetries: number; }; +/** In-session corrections allowed before a JSON-less turn terminalizes the attempt. */ +const maxStructuredOutputRetries = 2; + const zeroUsage: Usage = { inputTokens: 0, cachedInputTokens: 0, @@ -93,6 +102,16 @@ function extractJsonObject(text: string): unknown { throw new Error("Response does not contain a JSON object"); } +/** Reports whether a final model response contains an extractable JSON object. */ +function responseContainsJsonObject(text: string): boolean { + try { + extractJsonObject(text); + return true; + } catch { + return false; + } +} + /** Creates a task-scoped protocol error with optional session context and cause. */ function protocolError( code: string, @@ -259,6 +278,7 @@ export function createZcodeHarness(input: { outputDelivered: false, startedAt: now(), pendingDeliveries: [], + structuredRetries: 0, }; const failureCursor: BackendCursor = cursor ?? { version: 1, @@ -526,6 +546,7 @@ export function createZcodeHarness(input: { startedAt, reviewStatusBefore, pendingDeliveries: [], + structuredRetries: 0, }; activeAttempts.set(active.attemptId, active); const cursor: BackendCursor = { @@ -809,6 +830,34 @@ export function createZcodeHarness(input: { active.sessionId, ); } + // A successful turn whose response carries no JSON object gets an + // in-session correction round instead of terminalizing the attempt; + // the retry counter is ephemeral because a restart already abandons + // the in-flight turn (see reconcile). + if ( + active.structuredRetries < maxStructuredOutputRetries && + !responseContainsJsonObject(event.data.response) + ) { + active.structuredRetries += 1; + try { + await input.client.request("session/send", { + sessionId: active.sessionId, + content: structuredOutputRetryPrompt(active.role), + }); + } catch (error) { + throw normalizeError(error, { + code: "zcode_retry_send_failed", + category: "infra", + retryable: true, + component: "zcode-harness", + message: "ZCode did not accept the structured-output retry", + taskId: request.attempt.taskId, + attemptId: request.attempt.attemptId, + threadId: active.sessionId, + }); + } + continue; + } const deliveries = await completeFromTurn( request, cursor, @@ -872,6 +921,7 @@ export function createZcodeHarness(input: { startedAt: now(), reviewStatusBefore: cursor.reviewStatusBefore, pendingDeliveries: [], + structuredRetries: 0, }; if (request.attempt.role === "review") { try { diff --git a/src/agents/zcode/prompts.ts b/src/agents/zcode/prompts.ts index 3ff73fb..2e0eac7 100644 --- a/src/agents/zcode/prompts.ts +++ b/src/agents/zcode/prompts.ts @@ -94,3 +94,23 @@ export function reviewPrompt( JSON.stringify(validated.implementation, null, 2), ].join("\n"); } + +/** Builds the in-session correction sent when a completed turn carries no JSON object. */ +export function structuredOutputRetryPrompt( + role: "scout" | "implement" | "review", +): string { + const schema = + role === "scout" + ? ScoutOutputJsonSchema + : role === "implement" + ? ImplementDraftOutputJsonSchema + : ReviewOutputJsonSchema; + return [ + "Your previous final message contained no JSON object, so it could not be parsed.", + "Reply again with your final answer for the same task.", + "Your final message must be exactly one JSON object and nothing else.", + "The JSON object must match this exact schema:", + JSON.stringify(schema), + "Output the JSON object without markdown fences or any surrounding prose.", + ].join("\n"); +} diff --git a/test/agents/zcode/harness.test.ts b/test/agents/zcode/harness.test.ts index 9c561ed..2090226 100644 --- a/test/agents/zcode/harness.test.ts +++ b/test/agents/zcode/harness.test.ts @@ -210,9 +210,10 @@ function turnCompleted( fence?: boolean; resultType?: string; usage?: Record; + rawResponse?: string; } = {}, ): ServerMessage { - let response = JSON.stringify(output); + let response = options.rawResponse ?? JSON.stringify(output); if (options.fence) { response = "```json\n" + response + "\n```"; } @@ -583,3 +584,72 @@ test("provider rpc rejection text never reaches the durable failure event", asyn // stderr renders the same fixed message field. expect(JSON.stringify(failed.event)).not.toContain("zcode-secret-sentinel"); }); + +test("a JSON-less final response gets an in-session retry before completing", async () => { + const client = new RecordedZcodeClient(); + const harness = createZcodeHarness({ + client, + branches: memoryBranches(), + now: () => "2026-08-27T00:00:00.000Z", + }); + + const started = await harness.step(makeScoutRequest()); + if (started.kind !== "event") throw new Error("unreachable"); + + client.enqueue( + turnCompleted("sess-1", undefined, { + rawResponse: + "The repository inspection went well.\nNo structured payload in this reply.", + }), + ); + client.enqueue(turnCompleted("sess-1", scoutOutput)); + const { events } = await collect(harness, { + ...makeScoutRequest(), + backendCursor: started.nextCursor, + }); + expect(events.map((event) => event.type)).toEqual([ + "attempt.usage_delta", + "attempt.output", + "attempt.completed", + ]); + + const sends = client.requests.filter((r) => r.method === "session/send"); + expect(sends).toHaveLength(2); + const correction = sends[1]?.params as { sessionId?: string; content?: string }; + expect(correction.sessionId).toBe("sess-1"); + expect(correction.content).toContain("no JSON object"); + expect(correction.content).toContain("exactly one JSON object"); +}); + +test("persistently JSON-less responses exhaust retries and fail the attempt", async () => { + const client = new RecordedZcodeClient(); + const harness = createZcodeHarness({ + client, + branches: memoryBranches(), + now: () => "2026-08-27T00:00:00.000Z", + }); + + const started = await harness.step(makeScoutRequest()); + if (started.kind !== "event") throw new Error("unreachable"); + + for (let index = 0; index < 3; index += 1) { + client.enqueue( + turnCompleted("sess-1", undefined, { + rawResponse: `Prose report ${index}: findings in paragraph form only.`, + }), + ); + } + const { events } = await collect(harness, { + ...makeScoutRequest(), + backendCursor: started.nextCursor, + }); + expect(events.map((event) => event.type)).toEqual(["attempt.failed_infra"]); + const failure = events[0] as Extract< + HarnessEvent, + { type: "attempt.failed_infra" } + >; + expect(failure.code).toBe("invalid_structured_output"); + + const sends = client.requests.filter((r) => r.method === "session/send"); + expect(sends).toHaveLength(3); // prompt + two corrections, then give up +}); From 43270f9a6a1dd3dddaeba48a057f9aaaab57909f Mon Sep 17 00:00:00 2001 From: zhangweijian Date: Thu, 3 Sep 2026 15:03:25 +0800 Subject: [PATCH 2/3] fix(zcode): retry schema-violating JSON turns before failing Extend the in-session structured-output retry to responses whose JSON object violates the role's strict schema, not only JSON-less prose. Live runs of the ZCode backend showed the dominant failure shape is a JSON reply that embeds the prompt's schema metadata as a "\$schema" key or omits required fields, which strict validation rejects; the correction prompt already restates the schema, so both shapes share one retry budget (two corrections total) before the attempt fails through the existing invalid_structured_output path. Implement commit resolution stays outside the retry: a schema-valid draft whose commit cannot be resolved is a workspace problem restating the schema cannot correct. --- src/agents/zcode/harness.ts | 53 +++++++++++++++---- src/agents/zcode/prompts.ts | 9 +++- test/agents/conformance.ts | 15 ++++-- test/agents/zcode/harness.test.ts | 84 +++++++++++++++++++++++++++++++ 4 files changed, 145 insertions(+), 16 deletions(-) diff --git a/src/agents/zcode/harness.ts b/src/agents/zcode/harness.ts index 53f2717..2ed382e 100644 --- a/src/agents/zcode/harness.ts +++ b/src/agents/zcode/harness.ts @@ -102,14 +102,41 @@ function extractJsonObject(text: string): unknown { throw new Error("Response does not contain a JSON object"); } -/** Reports whether a final model response contains an extractable JSON object. */ -function responseContainsJsonObject(text: string): boolean { +/** One decoded, schema-valid role output; commit resolution is excluded. */ +type DecodedRoleOutput = + | z.infer + | Omit, "commitSha"> + | z.infer; + +/** + * Decodes a final model response into a schema-valid role output, or returns + * undefined when the response carries no extractable JSON object or the JSON + * object does not satisfy the role's strict schema. Implement commit + * resolution is deliberately excluded: a schema-valid draft whose commit + * cannot be resolved is a workspace problem that restating the schema cannot + * correct. + */ +function decodeStructuredOutput( + role: "scout" | "implement" | "review", + text: string, +): DecodedRoleOutput | undefined { + let decoded: unknown; try { - extractJsonObject(text); - return true; + decoded = extractJsonObject(text); } catch { - return false; + return undefined; + } + if (role === "implement") { + const draft = ImplementOutputSchema.omit({ + commitSha: true, + }).safeParse(decoded); + return draft.success ? draft.data : undefined; } + const parsed = + role === "scout" + ? ScoutOutputSchema.safeParse(decoded) + : ReviewOutputSchema.safeParse(decoded); + return parsed.success ? parsed.data : undefined; } /** Creates a task-scoped protocol error with optional session context and cause. */ @@ -830,13 +857,19 @@ export function createZcodeHarness(input: { active.sessionId, ); } - // A successful turn whose response carries no JSON object gets an - // in-session correction round instead of terminalizing the attempt; - // the retry counter is ephemeral because a restart already abandons - // the in-flight turn (see reconcile). + // A successful turn whose response carries no schema-valid JSON + // object gets an in-session correction round instead of + // terminalizing the attempt; the correction restates the schema, so + // it covers both a JSON-less prose reply and a JSON reply that + // violates the schema (e.g. an injected "$schema" key). The retry + // counter is ephemeral because a restart already abandons the + // in-flight turn (see reconcile). if ( active.structuredRetries < maxStructuredOutputRetries && - !responseContainsJsonObject(event.data.response) + decodeStructuredOutput( + active.role, + event.data.response, + ) === undefined ) { active.structuredRetries += 1; try { diff --git a/src/agents/zcode/prompts.ts b/src/agents/zcode/prompts.ts index 2e0eac7..128f0ad 100644 --- a/src/agents/zcode/prompts.ts +++ b/src/agents/zcode/prompts.ts @@ -95,7 +95,12 @@ export function reviewPrompt( ].join("\n"); } -/** Builds the in-session correction sent when a completed turn carries no JSON object. */ +/** + * Builds the in-session correction sent when a completed turn carries no + * schema-valid structured output: either no extractable JSON object at all, + * or a JSON object that violates the role's strict schema (for example an + * injected "$schema" key or a missing required field). + */ export function structuredOutputRetryPrompt( role: "scout" | "implement" | "review", ): string { @@ -106,7 +111,7 @@ export function structuredOutputRetryPrompt( ? ImplementDraftOutputJsonSchema : ReviewOutputJsonSchema; return [ - "Your previous final message contained no JSON object, so it could not be parsed.", + "Your previous final message contained no JSON object matching the required schema, so it could not be accepted.", "Reply again with your final answer for the same task.", "Your final message must be exactly one JSON object and nothing else.", "The JSON object must match this exact schema:", diff --git a/test/agents/conformance.ts b/test/agents/conformance.ts index ea66fb6..c3a41c2 100644 --- a/test/agents/conformance.ts +++ b/test/agents/conformance.ts @@ -526,10 +526,17 @@ export function defineNormalizedConformance( "output before completion", "invalid structured output fails closed without output or completion", async (fixture) => { - fixture.driver.scriptSuccessfulTurn("scout", { - usage: scriptedUsage, - invalidOutput: true, - }); + // Backends that retry invalid structured output in-session (the + // ZCode backend sends up to two schema-restating corrections) consume + // one scripted invalid turn per correction round before failing; + // backends without the retry fail on the first turn and ignore the + // extra scripted turns. + for (let index = 0; index < 3; index += 1) { + fixture.driver.scriptSuccessfulTurn("scout", { + usage: scriptedUsage, + invalidOutput: true, + }); + } const { events } = await collect( fixture.harness, conformanceRequest("scout", { attemptId: "attempt-invalid" }), diff --git a/test/agents/zcode/harness.test.ts b/test/agents/zcode/harness.test.ts index 2090226..87c7be0 100644 --- a/test/agents/zcode/harness.test.ts +++ b/test/agents/zcode/harness.test.ts @@ -653,3 +653,87 @@ test("persistently JSON-less responses exhaust retries and fail the attempt", as const sends = client.requests.filter((r) => r.method === "session/send"); expect(sends).toHaveLength(3); // prompt + two corrections, then give up }); + +test("a schema-violating JSON final response gets an in-session retry before completing", async () => { + const client = new RecordedZcodeClient(); + const harness = createZcodeHarness({ + client, + branches: memoryBranches(), + now: () => "2026-08-27T00:00:00.000Z", + }); + + const started = await harness.step(makeScoutRequest()); + if (started.kind !== "event") throw new Error("unreachable"); + + // Observed live failure shape: the model embeds the prompt's JSON-schema + // metadata as a "$schema" key, which the strict role schema rejects. + const polluted = { + $schema: "https://json-schema.org/draft/2020-12/schema", + ...scoutOutput, + }; + client.enqueue(turnCompleted("sess-1", undefined, { + rawResponse: JSON.stringify(polluted), + })); + client.enqueue(turnCompleted("sess-1", scoutOutput)); + const { events } = await collect(harness, { + ...makeScoutRequest(), + backendCursor: started.nextCursor, + }); + expect(events.map((event) => event.type)).toEqual([ + "attempt.usage_delta", + "attempt.output", + "attempt.completed", + ]); + + const sends = client.requests.filter((r) => r.method === "session/send"); + expect(sends).toHaveLength(2); + const correction = sends[1]?.params as { sessionId?: string; content?: string }; + expect(correction.sessionId).toBe("sess-1"); + expect(correction.content).toContain("no JSON object"); + expect(correction.content).toContain('"$schema"'); +}); + +test("mixed JSON-less and schema-violating responses exhaust retries and fail the attempt", async () => { + const client = new RecordedZcodeClient(); + const harness = createZcodeHarness({ + client, + branches: memoryBranches(), + now: () => "2026-08-27T00:00:00.000Z", + }); + + const started = await harness.step(makeScoutRequest()); + if (started.kind !== "event") throw new Error("unreachable"); + + client.enqueue( + turnCompleted("sess-1", undefined, { + rawResponse: "Prose report: findings in paragraph form only.", + }), + ); + client.enqueue( + turnCompleted("sess-1", undefined, { + rawResponse: JSON.stringify({ + $schema: "https://json-schema.org/draft/2020-12/schema", + ...scoutOutput, + }), + }), + ); + client.enqueue( + turnCompleted("sess-1", undefined, { + rawResponse: JSON.stringify({ kind: "scout" }), // missing required fields + }), + ); + const { events } = await collect(harness, { + ...makeScoutRequest(), + backendCursor: started.nextCursor, + }); + expect(events.map((event) => event.type)).toEqual(["attempt.failed_infra"]); + const failure = events[0] as Extract< + HarnessEvent, + { type: "attempt.failed_infra" } + >; + expect(failure.code).toBe("invalid_structured_output"); + + // Both failure shapes share one retry budget: prompt + two corrections. + const sends = client.requests.filter((r) => r.method === "session/send"); + expect(sends).toHaveLength(3); +}); From 76ca96231adf4bdb3c4d278fb1f0b8bb929035f8 Mon Sep 17 00:00:00 2001 From: zhangweijian Date: Thu, 3 Sep 2026 17:23:10 +0800 Subject: [PATCH 3/3] fix(zcode): accumulate correction-round token usage per native semantics The native protocol reports usage per turn, but the correction branch dropped each JSON-less turn's usage before sending the retry, so a corrected attempt reported only the final turn's tokens and a retry-exhausted attempt reported zero. Every successful completion now adds to an attempt-scoped accumulator seeded from the persisted cursor baseline: correction success, correction-send failure, and retry exhaustion all publish the consumed tokens through a usage_delta delivery whose cursor carries the total, and a replay from that cursor never re-counts already-published usage. Also adds boundary tests for JSON scalars and arrays (null, [], string), which earn a correction round under the retry extension instead of bypassing it, and expands the harness.test.ts correction type declaration per Biome format. --- src/agents/zcode/harness.ts | 187 +++++++++++++----- test/agents/zcode/harness.test.ts | 303 +++++++++++++++++++++++++++++- 2 files changed, 429 insertions(+), 61 deletions(-) diff --git a/src/agents/zcode/harness.ts b/src/agents/zcode/harness.ts index 2ed382e..24c1ea0 100644 --- a/src/agents/zcode/harness.ts +++ b/src/agents/zcode/harness.ts @@ -66,7 +66,15 @@ type ActiveAttempt = { reviewStatusBefore?: string; reconciledCompletion?: boolean; pendingDeliveries: HarnessDelivery[]; + /** Deferred terminal delivery handed out after queued deliveries drain. */ + pendingTerminal?: () => HarnessDelivery; structuredRetries: number; + /** + * Sum of every per-turn usage the native protocol has reported for this + * attempt, seeded from the persisted cursor baseline so a resumed attempt + * never re-counts usage that a prior delivery already published. + */ + accumulatedUsage: Usage; }; /** In-session corrections allowed before a JSON-less turn terminalizes the attempt. */ @@ -79,6 +87,29 @@ const zeroUsage: Usage = { reasoningOutputTokens: 0, }; +/** Adds one native per-turn usage report onto a running attempt total. */ +function addUsage(total: Usage, turn: Usage): Usage { + return { + inputTokens: total.inputTokens + turn.inputTokens, + cachedInputTokens: total.cachedInputTokens + turn.cachedInputTokens, + outputTokens: total.outputTokens + turn.outputTokens, + reasoningOutputTokens: + total.reasoningOutputTokens + turn.reasoningOutputTokens, + }; +} + +/** Returns the not-yet-published usage delta against a cursor baseline. */ +function usageDelta(accumulated: Usage, baseline: Usage): Usage { + return { + inputTokens: accumulated.inputTokens - baseline.inputTokens, + cachedInputTokens: + accumulated.cachedInputTokens - baseline.cachedInputTokens, + outputTokens: accumulated.outputTokens - baseline.outputTokens, + reasoningOutputTokens: + accumulated.reasoningOutputTokens - baseline.reasoningOutputTokens, + }; +} + /** Extracts a JSON object from a final model response, tolerating fences and prose. */ function extractJsonObject(text: string): unknown { const trimmed = text.trim(); @@ -268,6 +299,38 @@ export function createZcodeHarness(input: { }); } + /** + * Publishes any turn usage the attempt consumed but the persisted cursor + * baseline does not carry yet, then hands out the first delivery of the + * terminal chain. The terminal delivery itself is deferred: building it + * eagerly would terminalize the attempt while the queued usage delivery + * has not been consumed yet, so it is armed on the attempt and built when + * the delivery queue drains. + */ + function terminalDeliveryChain( + cursor: BackendCursor, + active: ActiveAttempt, + build: (cursor: BackendCursor) => HarnessDelivery, + ): HarnessDelivery { + const delta = usageDelta(active.accumulatedUsage, cursor.usage); + if (Object.values(delta).every((value) => value === 0)) { + return build(cursor); + } + const running: BackendCursor = { + ...cursor, + usage: active.accumulatedUsage, + }; + active.pendingTerminal = () => + build({ ...running, nextSequence: running.nextSequence + 1 }); + return delivery(running, { + type: "attempt.usage_delta", + eventId: `${active.attemptId}:${active.sessionId}:usage:${cumulativeHash(running.usage)}`, + attemptId: active.attemptId, + occurredAt: now(), + ...delta, + }); + } + /** Converts task-scoped operational failures into policy or infrastructure deliveries. */ function taskFailure( request: SupportedRequest, @@ -306,6 +369,7 @@ export function createZcodeHarness(input: { startedAt: now(), pendingDeliveries: [], structuredRetries: 0, + accumulatedUsage: cursor?.usage ?? zeroUsage, }; const failureCursor: BackendCursor = cursor ?? { version: 1, @@ -314,22 +378,26 @@ export function createZcodeHarness(input: { usage: zeroUsage, }; if (normalized.category === "policy") { - terminalize(active); - return delivery(withoutTerminalMarkers(failureCursor), { - type: "attempt.blocked_policy", - eventId: `${active.attemptId}:${active.sessionId}:blocked_policy:${normalized.code}`, - attemptId: active.attemptId, - occurredAt: now(), - code: normalized.code, - message: normalized.message, + return terminalDeliveryChain(failureCursor, active, (terminalCursor) => { + terminalize(active); + return delivery(withoutTerminalMarkers(terminalCursor), { + type: "attempt.blocked_policy", + eventId: `${active.attemptId}:${active.sessionId}:blocked_policy:${normalized.code}`, + attemptId: active.attemptId, + occurredAt: now(), + code: normalized.code, + message: normalized.message, + }); }); } - return failedDelivery( - failureCursor, - active, - normalized.code, - normalized.message, - normalized.retryable, + return terminalDeliveryChain(failureCursor, active, (terminalCursor) => + failedDelivery( + terminalCursor, + active, + normalized.code, + normalized.message, + normalized.retryable, + ), ); } @@ -356,22 +424,26 @@ export function createZcodeHarness(input: { ): HarnessDelivery { const failure = classifyZcodeTurnFailure(resultType, errorText); if (failure.category === "policy") { - terminalize(active); - return delivery(withoutTerminalMarkers(cursor), { - type: "attempt.blocked_policy", - eventId: `${active.attemptId}:${active.sessionId}:blocked_policy:${failure.code}`, - attemptId: active.attemptId, - occurredAt: now(), - code: failure.code, - message: failure.message, + return terminalDeliveryChain(cursor, active, (terminalCursor) => { + terminalize(active); + return delivery(withoutTerminalMarkers(terminalCursor), { + type: "attempt.blocked_policy", + eventId: `${active.attemptId}:${active.sessionId}:blocked_policy:${failure.code}`, + attemptId: active.attemptId, + occurredAt: now(), + code: failure.code, + message: failure.message, + }); }); } - return failedDelivery( - cursor, - active, - failure.code, - failure.message, - failure.retryable, + return terminalDeliveryChain(cursor, active, (terminalCursor) => + failedDelivery( + terminalCursor, + active, + failure.code, + failure.message, + failure.retryable, + ), ); } @@ -574,6 +646,7 @@ export function createZcodeHarness(input: { reviewStatusBefore, pendingDeliveries: [], structuredRetries: 0, + accumulatedUsage: zeroUsage, }; activeAttempts.set(active.attemptId, active); const cursor: BackendCursor = { @@ -672,11 +745,13 @@ export function createZcodeHarness(input: { ); } catch { return [ - failedDelivery( - cursor, - active, - "review_status_snapshot_failed", - "Could not verify the Review workspace", + terminalDeliveryChain(cursor, active, (terminalCursor) => + failedDelivery( + terminalCursor, + active, + "review_status_snapshot_failed", + "Could not verify the Review workspace", + ), ), ]; } @@ -685,26 +760,24 @@ export function createZcodeHarness(input: { statusAfter !== active.reviewStatusBefore ) { return [ - failedDelivery( - cursor, - active, - "review_mutated_workspace", - "Review changed the task checkout", + terminalDeliveryChain(cursor, active, (terminalCursor) => + failedDelivery( + terminalCursor, + active, + "review_mutated_workspace", + "Review changed the task checkout", + ), ), ]; } } const deliveries: HarnessDelivery[] = []; - const cumulative = mapZcodeUsage(payload.usage); - const delta: Usage = { - inputTokens: cumulative.inputTokens - cursor.usage.inputTokens, - cachedInputTokens: - cumulative.cachedInputTokens - cursor.usage.cachedInputTokens, - outputTokens: cumulative.outputTokens - cursor.usage.outputTokens, - reasoningOutputTokens: - cumulative.reasoningOutputTokens - cursor.usage.reasoningOutputTokens, - }; + // The attempt total already includes this turn's usage — it accumulates + // when the completion event arrives — so correction rounds stay + // accounted instead of being replaced by the final turn alone. + const cumulative = active.accumulatedUsage; + const delta: Usage = usageDelta(cumulative, cursor.usage); if (Object.values(delta).some((value) => value < 0)) { throw protocolError( "non_monotonic_token_usage", @@ -759,6 +832,11 @@ export function createZcodeHarness(input: { const queued = active.pendingDeliveries.shift(); if (queued !== undefined) return queued; } + if (active.pendingTerminal !== undefined) { + const terminal = active.pendingTerminal; + active.pendingTerminal = undefined; + return terminal(); + } if (active.outputDelivered) { // The turn's structured output was already delivered; hand out the // terminal completion delivery now. @@ -857,6 +935,14 @@ export function createZcodeHarness(input: { active.sessionId, ); } + // The native protocol reports usage per turn, so every completion + // — including one that only earns a correction round — adds to the + // attempt total. Dropping it here would under-report corrected + // attempts and zero out retry-exhausted ones. + active.accumulatedUsage = addUsage( + active.accumulatedUsage, + mapZcodeUsage(event.data.usage), + ); // A successful turn whose response carries no schema-valid JSON // object gets an in-session correction round instead of // terminalizing the attempt; the correction restates the schema, so @@ -866,10 +952,8 @@ export function createZcodeHarness(input: { // in-flight turn (see reconcile). if ( active.structuredRetries < maxStructuredOutputRetries && - decodeStructuredOutput( - active.role, - event.data.response, - ) === undefined + decodeStructuredOutput(active.role, event.data.response) === + undefined ) { active.structuredRetries += 1; try { @@ -955,6 +1039,7 @@ export function createZcodeHarness(input: { reviewStatusBefore: cursor.reviewStatusBefore, pendingDeliveries: [], structuredRetries: 0, + accumulatedUsage: cursor.usage, }; if (request.attempt.role === "review") { try { diff --git a/test/agents/zcode/harness.test.ts b/test/agents/zcode/harness.test.ts index 87c7be0..a4895a6 100644 --- a/test/agents/zcode/harness.test.ts +++ b/test/agents/zcode/harness.test.ts @@ -615,7 +615,10 @@ test("a JSON-less final response gets an in-session retry before completing", as const sends = client.requests.filter((r) => r.method === "session/send"); expect(sends).toHaveLength(2); - const correction = sends[1]?.params as { sessionId?: string; content?: string }; + const correction = sends[1]?.params as { + sessionId?: string; + content?: string; + }; expect(correction.sessionId).toBe("sess-1"); expect(correction.content).toContain("no JSON object"); expect(correction.content).toContain("exactly one JSON object"); @@ -632,23 +635,64 @@ test("persistently JSON-less responses exhaust retries and fail the attempt", as const started = await harness.step(makeScoutRequest()); if (started.kind !== "event") throw new Error("unreachable"); + const perTurnUsage = [ + { + inputTokens: 1000, + outputTokens: 50, + reasoningTokens: 20, + cacheReadTokens: 400, + }, + { + inputTokens: 200, + outputTokens: 10, + reasoningTokens: 0, + cacheReadTokens: 0, + }, + { + inputTokens: 30, + outputTokens: 2, + reasoningTokens: 0, + cacheReadTokens: 0, + }, + ]; for (let index = 0; index < 3; index += 1) { client.enqueue( turnCompleted("sess-1", undefined, { rawResponse: `Prose report ${index}: findings in paragraph form only.`, + usage: perTurnUsage[index], }), ); } - const { events } = await collect(harness, { + const { events, cursors } = await collect(harness, { ...makeScoutRequest(), backendCursor: started.nextCursor, }); - expect(events.map((event) => event.type)).toEqual(["attempt.failed_infra"]); - const failure = events[0] as Extract< + expect(events.map((event) => event.type)).toEqual([ + "attempt.usage_delta", + "attempt.failed_infra", + ]); + const usage = events[0] as Extract< + HarnessEvent, + { type: "attempt.usage_delta" } + >; + expect(usage).toMatchObject({ + inputTokens: 1230, + cachedInputTokens: 400, + outputTokens: 62, + reasoningOutputTokens: 20, + }); + const failure = events[1] as Extract< HarnessEvent, { type: "attempt.failed_infra" } >; expect(failure.code).toBe("invalid_structured_output"); + // The terminal cursor keeps the consumed usage instead of zeroing it. + expect(JSON.parse(cursors.at(-1) ?? "{}").usage).toEqual({ + inputTokens: 1230, + cachedInputTokens: 400, + outputTokens: 62, + reasoningOutputTokens: 20, + }); const sends = client.requests.filter((r) => r.method === "session/send"); expect(sends).toHaveLength(3); // prompt + two corrections, then give up @@ -671,9 +715,11 @@ test("a schema-violating JSON final response gets an in-session retry before com $schema: "https://json-schema.org/draft/2020-12/schema", ...scoutOutput, }; - client.enqueue(turnCompleted("sess-1", undefined, { - rawResponse: JSON.stringify(polluted), - })); + client.enqueue( + turnCompleted("sess-1", undefined, { + rawResponse: JSON.stringify(polluted), + }), + ); client.enqueue(turnCompleted("sess-1", scoutOutput)); const { events } = await collect(harness, { ...makeScoutRequest(), @@ -687,7 +733,10 @@ test("a schema-violating JSON final response gets an in-session retry before com const sends = client.requests.filter((r) => r.method === "session/send"); expect(sends).toHaveLength(2); - const correction = sends[1]?.params as { sessionId?: string; content?: string }; + const correction = sends[1]?.params as { + sessionId?: string; + content?: string; + }; expect(correction.sessionId).toBe("sess-1"); expect(correction.content).toContain("no JSON object"); expect(correction.content).toContain('"$schema"'); @@ -726,8 +775,11 @@ test("mixed JSON-less and schema-violating responses exhaust retries and fail th ...makeScoutRequest(), backendCursor: started.nextCursor, }); - expect(events.map((event) => event.type)).toEqual(["attempt.failed_infra"]); - const failure = events[0] as Extract< + expect(events.map((event) => event.type)).toEqual([ + "attempt.usage_delta", + "attempt.failed_infra", + ]); + const failure = events[1] as Extract< HarnessEvent, { type: "attempt.failed_infra" } >; @@ -737,3 +789,234 @@ test("mixed JSON-less and schema-violating responses exhaust retries and fail th const sends = client.requests.filter((r) => r.method === "session/send"); expect(sends).toHaveLength(3); }); + +test("correction rounds accumulate per-turn usage into the final delta and cursor", async () => { + const client = new RecordedZcodeClient(); + const harness = createZcodeHarness({ + client, + branches: memoryBranches(), + now: () => "2026-08-27T00:00:00.000Z", + }); + + const started = await harness.step(makeScoutRequest()); + if (started.kind !== "event") throw new Error("unreachable"); + + client.enqueue( + turnCompleted("sess-1", undefined, { + rawResponse: "Scout prose reply without any JSON payload.", + usage: { + inputTokens: 1000, + outputTokens: 50, + reasoningTokens: 20, + cacheReadTokens: 400, + }, + }), + ); + client.enqueue( + turnCompleted("sess-1", scoutOutput, { + usage: { + inputTokens: 200, + outputTokens: 10, + reasoningTokens: 0, + cacheReadTokens: 0, + }, + }), + ); + const { events, cursors } = await collect(harness, { + ...makeScoutRequest(), + backendCursor: started.nextCursor, + }); + expect(events.map((event) => event.type)).toEqual([ + "attempt.usage_delta", + "attempt.output", + "attempt.completed", + ]); + const usage = events[0] as Extract< + HarnessEvent, + { type: "attempt.usage_delta" } + >; + expect(usage).toMatchObject({ + inputTokens: 1200, + cachedInputTokens: 400, + outputTokens: 60, + reasoningOutputTokens: 20, + }); + expect(JSON.parse(cursors.at(-1) ?? "{}").usage).toEqual({ + inputTokens: 1200, + cachedInputTokens: 400, + outputTokens: 60, + reasoningOutputTokens: 20, + }); +}); + +test("a failed correction send still preserves the consumed turn usage", async () => { + class FailingCorrectionClient extends RecordedZcodeClient { + private sendCount = 0; + + override async request(method: string, params: unknown): Promise { + if (method === "session/send") { + this.sendCount += 1; + if (this.sendCount > 1) throw new Error("correction send rejected"); + } + return super.request(method, params); + } + } + const client = new FailingCorrectionClient(); + const harness = createZcodeHarness({ + client, + branches: memoryBranches(), + now: () => "2026-08-27T00:00:00.000Z", + }); + + const started = await harness.step(makeScoutRequest()); + if (started.kind !== "event") throw new Error("unreachable"); + + client.enqueue( + turnCompleted("sess-1", undefined, { + rawResponse: "Scout prose reply without any JSON payload.", + usage: { + inputTokens: 1000, + outputTokens: 50, + reasoningTokens: 0, + cacheReadTokens: 0, + }, + }), + ); + const { events, cursors } = await collect(harness, { + ...makeScoutRequest(), + backendCursor: started.nextCursor, + }); + expect(events.map((event) => event.type)).toEqual([ + "attempt.usage_delta", + "attempt.failed_infra", + ]); + const usage = events[0] as Extract< + HarnessEvent, + { type: "attempt.usage_delta" } + >; + expect(usage).toMatchObject({ + inputTokens: 1000, + cachedInputTokens: 0, + outputTokens: 50, + reasoningOutputTokens: 0, + }); + const failure = events[1] as Extract< + HarnessEvent, + { type: "attempt.failed_infra" } + >; + expect(failure.code).toBe("zcode_retry_send_failed"); + expect(JSON.parse(cursors.at(-1) ?? "{}").usage).toEqual({ + inputTokens: 1000, + cachedInputTokens: 0, + outputTokens: 50, + reasoningOutputTokens: 0, + }); +}); + +test("JSON scalars and arrays earn a correction round before completing", async () => { + for (const malformed of ["null", "[]", '"just a string"']) { + const client = new RecordedZcodeClient(); + const harness = createZcodeHarness({ + client, + branches: memoryBranches(), + now: () => "2026-08-27T00:00:00.000Z", + }); + + const started = await harness.step(makeScoutRequest()); + if (started.kind !== "event") throw new Error("unreachable"); + + // Valid JSON that is not an object must not bypass the correction path. + client.enqueue( + turnCompleted("sess-1", undefined, { rawResponse: malformed }), + ); + client.enqueue(turnCompleted("sess-1", scoutOutput)); + const { events } = await collect(harness, { + ...makeScoutRequest(), + backendCursor: started.nextCursor, + }); + expect(events.map((event) => event.type)).toEqual([ + "attempt.usage_delta", + "attempt.output", + "attempt.completed", + ]); + const sends = client.requests.filter((r) => r.method === "session/send"); + expect(sends).toHaveLength(2); // initial prompt + one correction + } +}); + +test("replay after a mid-chain crash keeps usage without re-counting it", async () => { + const client = new RecordedZcodeClient(); + const harness = createZcodeHarness({ + client, + branches: memoryBranches(), + now: () => "2026-08-27T00:00:00.000Z", + }); + + const started = await harness.step(makeScoutRequest()); + if (started.kind !== "event") throw new Error("unreachable"); + + const perTurnUsage = [ + { + inputTokens: 1000, + outputTokens: 50, + reasoningTokens: 0, + cacheReadTokens: 0, + }, + { + inputTokens: 200, + outputTokens: 10, + reasoningTokens: 0, + cacheReadTokens: 0, + }, + { + inputTokens: 30, + outputTokens: 2, + reasoningTokens: 0, + cacheReadTokens: 0, + }, + ]; + for (let index = 0; index < 3; index += 1) { + client.enqueue( + turnCompleted("sess-1", undefined, { + rawResponse: `Prose report ${index}: findings in paragraph form only.`, + usage: perTurnUsage[index], + }), + ); + } + // The retry-exhausted attempt hands out its usage_delta first; the process + // then "crashes" before the terminal delivery drains. + const usageDelivery = await harness.step({ + ...makeScoutRequest(), + backendCursor: started.nextCursor, + }); + expect(usageDelivery).toMatchObject({ + kind: "event", + event: { type: "attempt.usage_delta", inputTokens: 1230 }, + }); + if (usageDelivery.kind !== "event") throw new Error("unreachable"); + + // A restarted process replays from the persisted usage cursor: the + // orphaned-turn failure keeps the consumed usage and emits no second + // usage_delta for tokens the cursor already carries. + const replayed = createZcodeHarness({ + client: new RecordedZcodeClient(), + branches: memoryBranches(), + now: () => "2026-08-27T00:00:00.000Z", + }); + const orphaned = await replayed.step({ + ...makeScoutRequest(), + mode: "reconcile", + backendCursor: usageDelivery.nextCursor, + }); + expect(orphaned).toMatchObject({ + kind: "event", + event: { type: "attempt.failed_infra", code: "orphaned_turn" }, + }); + if (orphaned.kind !== "event") throw new Error("unreachable"); + expect(JSON.parse(orphaned.nextCursor).usage).toEqual({ + inputTokens: 1230, + cachedInputTokens: 0, + outputTokens: 62, + reasoningOutputTokens: 0, + }); +});