diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md index 11ce285c0f..c0ca09f838 100644 --- a/.claude/CLAUDE.md +++ b/.claude/CLAUDE.md @@ -245,8 +245,9 @@ put to a person still runs one at a time, and results return in the order asked) `maxRoundRetries` (default 2, for a `RetryableJobError`, waiting as the provider's retry-after says within 1–60 s), `roundTimeoutMs` (a provider can accept a request and never answer; past it the round is abandoned as a retryable failure, so the retries cover it; a round's text and -reasoning are forwarded only once the attempt settles, since a failed attempt's partial cannot be -taken back from the accumulated port; reasoning goes out on the turn's own `reasoning` port, a +reasoning are forwarded live, and an attempt that fails takes back exactly what it forwarded with a +`text-retract` per port, so the folded ports read as if it never ran and earlier rounds are untouched; +reasoning goes out on the turn's own `reasoning` port, a blank line between rounds, and is absent from the output when no round returned any), and budgets — `maxInputTokens`, `maxCostUsd` (refused without a price card, since a budget it cannot measure never stops anything), `maxDurationMs` — each ending the turn @@ -306,7 +307,9 @@ switch. DeepSeek's `replay_reasoning` is a different mechanism: it controls repl (they draw on a screen, ask a person, read state only that process holds) but whose model is reachable only through a backend binds `AGENT_ROUND_RUNNER` on the run's registry: every round is then handed to it as a `ToolCallingTaskInput` and answered as a `ToolCallingTaskOutput`, while the -loop and every tool stay in-process. A serialized graph cannot carry the closures, so without it a +loop and every tool stay in-process. The runner streams back through its `AgentRoundContext`: +text, reasoning, and `onTextRetract` for text its own call withdrew, which the turn caps at what +that attempt forwarded. A serialized graph cannot carry the closures, so without it a host chooses between tools that work and a model it can reach. A registry binding, not an input, for the same reason as `AGENT_APPROVAL_OPT_OUT`: where a turn's model calls go is the host's call. @@ -318,7 +321,7 @@ task: `normalizeHistoryForModel` / `trimHistoryForModel` (`ChatHistory.ts`), and unit-tested in `packages/ai/src/task/__tests__/`: `AgentTurnBudget` (tokens, cost, deadline — `exhausted()` answers the answer path, `mustStopLoop()` the loop, so changing what an answered turn reports cannot move where the loop stops), `AgentRoundAttempt` (one round's retries, -timeout and held deltas), `AgentTurnPolicy` (submit tool, reminders, review, time-left notice) +timeout, and the live deltas a failed attempt takes back), `AgentTurnPolicy` (submit tool, reminders, review, time-left notice) and `AgentToolRound` (dispatching a round's calls and building their results). `validateModelAuthoredSchema` (`ModelAuthoredSchema.ts`) bounds a JSON Schema a model wrote diff --git a/docs/technical/04-dataflow-and-streaming.md b/docs/technical/04-dataflow-and-streaming.md index 5ee31fe5b2..c90fb736cc 100644 --- a/docs/technical/04-dataflow-and-streaming.md +++ b/docs/technical/04-dataflow-and-streaming.md @@ -348,12 +348,16 @@ All streaming data flows through the `StreamEvent` discriminated union type: ```typescript type StreamEvent> = | StreamTextDelta + | StreamTextRetract | StreamObjectDelta | StreamSnapshot | StreamFinish | StreamError; ``` +(Abridged: the union also carries `binary-delta`, `refusal`, `usage`, `phase` and `tool-call` +events; see `StreamTypes.ts`.) + ### StreamTextDelta ```typescript @@ -366,6 +370,29 @@ interface StreamTextDelta { Used with `x-stream: "append"`. Each event carries a fragment of text that should be appended to the accumulated result for the named port. +### StreamTextRetract + +```typescript +interface StreamTextRetract { + type: "text-retract"; + port: string; // Output port name + length: number; // UTF-16 code units to take back +} +``` + +Append mode only. Removes the last `length` characters streamed on `port` in this run, clamped at +the start; later `text-delta`s append after what remains. Every consumer folding a port's text +truncates with `retractText(acc, length)`, which removes `retractionLength(length)` characters (a +fraction rounds down, and anything not positive removes nothing). A +producer that streams text it may yet abandon uses it: `AgentTask` forwards each round live and, +when an attempt fails and is retried, takes back exactly what that attempt sent, so the folded port +reads as if the attempt never ran. Dataflow edges carrying the port carry the retraction too. + +A port whose text was already teed to a cache sink cannot be truncated there: the processor fails +that port's sink (discarding the partial write) and settles the port inline from the accumulator +for the rest of the run. With no accumulator either, the run fails at `finish` unless the finish +payload states the port's value itself. + ### StreamObjectDelta ```typescript @@ -451,6 +478,11 @@ text-delta: - Emit "stream_chunk" on the task - Update progress (asymptotic curve: 1 - e^(-0.05*chunkCount)) +text-retract: + - Truncate the port's accumulated text (if shouldAccumulate) + - Fail that port's cache sink, if one was teed (see StreamTextRetract) + - Emit "stream_chunk" on the task + object-delta: - Accumulate per-port (if shouldAccumulate) - Update runOutputData progressively diff --git a/examples/cli/src/agent/runAgentChat.ts b/examples/cli/src/agent/runAgentChat.ts index 498c57c87b..c2459711c4 100644 --- a/examples/cli/src/agent/runAgentChat.ts +++ b/examples/cli/src/agent/runAgentChat.ts @@ -7,6 +7,7 @@ import type { AgentApprovalMode, AgentTaskOutput, ChatMessage, ToolDefinition } from "@workglow/ai"; import { AGENT_APPROVAL_OPT_OUT, AgentTask } from "@workglow/ai"; import type { StreamEvent } from "@workglow/task-graph"; +import { retractionLength } from "@workglow/task-graph"; import { globalServiceRegistry, HUMAN_CONNECTOR, @@ -202,14 +203,26 @@ async function runTurn( // here it would read as part of the answer. if (event.type === "text-delta" && (event.port ?? "text") === "text") { transcript.delta(event.textDelta); + } else if ( + event.type === "text-retract" && + (event.port ?? "text") === "text" && + retractionLength(event.length) > 0 && + !controller.signal.aborted + ) { + // An attempt that failed partway takes its text back. What it printed is + // already in the scrollback and cannot be erased, so say it no longer + // counts; the retry's reply follows on a line of its own. An interrupted + // turn retracts too, but it did not fail: the "interrupted" note says so. + transcript.note(" · that attempt failed; the partial reply above is discarded"); } }); const offProgress = task.subscribe("progress", (_progress, message) => { - // Only the tool lines are worth a row: "Thinking" is what the blank space - // between the prompt and the first token already says. + // Only the tool and retry lines are worth a row: "Thinking" is what the + // blank space between the prompt and the first token already says. if (!message || message === phase) return; phase = message; if (message.startsWith("Running ")) transcript.note(` · ${message.slice("Running ".length)}`); + else if (message.startsWith("Retrying ")) transcript.note(` · ${message.toLowerCase()}`); }); try { diff --git a/examples/cli/src/run-events/RunEventTypes.ts b/examples/cli/src/run-events/RunEventTypes.ts index 2daba2d195..5152b3e7f7 100644 --- a/examples/cli/src/run-events/RunEventTypes.ts +++ b/examples/cli/src/run-events/RunEventTypes.ts @@ -61,6 +61,16 @@ export type RunEvent = | { readonly k: "text"; readonly id: string; readonly delta: string } /** A model's thinking, streamed beside its answer; never folded into `text`. */ | { readonly k: "reasoning"; readonly id: string; readonly delta: string } + /** + * Takes back the last `length` characters of a row's `text` or `reasoning`: + * a model call that streamed and then failed, retracted before its retry. + */ + | { + readonly k: "retract"; + readonly id: string; + readonly port: "text" | "reasoning"; + readonly length: number; + } | { readonly k: "messages"; readonly id: string; readonly messages: unknown } | { readonly k: "human_request"; diff --git a/examples/cli/src/run-events/projectRunEvents.test.ts b/examples/cli/src/run-events/projectRunEvents.test.ts index 71d07142f0..207a3cc55d 100644 --- a/examples/cli/src/run-events/projectRunEvents.test.ts +++ b/examples/cli/src/run-events/projectRunEvents.test.ts @@ -4,7 +4,8 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { Task, Workflow, type IExecuteContext } from "@workglow/task-graph"; +import type { IExecuteContext, StreamEvent } from "@workglow/task-graph"; +import { Task, Workflow } from "@workglow/task-graph"; import type { DataPortSchema } from "@workglow/util/schema"; import { describe, expect, it } from "vitest"; import type { RunEvent } from "./RunEventTypes"; @@ -94,7 +95,41 @@ class ThinkingTask extends Task, { text: string; reasoning } } +/** Streams a partial, takes it back, then streams the reply, as a retried `AgentTask` round does. */ +class RetractingTask extends Task, { text: string; reasoning: string }> { + static override readonly type = "RetractingTask"; + static override readonly title = "Retracting task"; + static override readonly category = "Test"; + static override readonly cacheable = false; + static override inputSchema(): DataPortSchema { + return EMPTY; + } + static override outputSchema(): DataPortSchema { + return STREAMED_PORTS; + } + async *executeStream(): AsyncIterable> { + yield { type: "text-delta" as const, port: "text", textDelta: "Oops" }; + yield { type: "text-retract" as const, port: "text", length: 4 }; + yield { type: "text-delta" as const, port: "text", textDelta: "Hello." }; + yield { type: "finish" as const, data: {} as { text: string; reasoning: string } }; + } +} + describe("projectRunEvents streamed text", () => { + it("reports a retraction as its own kind, naming the port", async () => { + const { events, sink } = collector(); + const workflow = new Workflow(); + workflow.pipe(new RetractingTask() as never); + const stop = projectRunEvents(workflow.graph, sink); + await workflow.run({}); + stop(); + + const text = events.flatMap((e) => + e.k === "text" ? [e.delta] : e.k === "retract" ? [`-${e.port}:${e.length}`] : [] + ); + expect(text).toEqual(["Oops", "-text:4", "Hello."]); + }); + it("reports a task's text and reasoning as separate kinds", async () => { const { events, sink } = collector(); const workflow = new Workflow(); diff --git a/examples/cli/src/run-events/projectRunEvents.ts b/examples/cli/src/run-events/projectRunEvents.ts index 0ec1111a6d..5035085322 100644 --- a/examples/cli/src/run-events/projectRunEvents.ts +++ b/examples/cli/src/run-events/projectRunEvents.ts @@ -5,7 +5,7 @@ */ import type { ITask, StreamEvent, TaskGraph, TaskIdType, Usage } from "@workglow/task-graph"; -import { classifyTextDelta } from "../ui/model/streamedText"; +import { classifyTextDelta, classifyTextRetract } from "../ui/model/streamedText"; import { cliTaskLabel } from "../ui/taskGraphCliSubscriptions"; import { projectTaskSubtree } from "./projectSubtrees"; import type { RunEventSink } from "./runEventChannel"; @@ -130,6 +130,16 @@ export function projectRunEvents( delta: classified.delta, }); } + } else if (event.type === "text-retract") { + const classified = classifyTextRetract(event); + if (classified) { + sink.emit({ + k: "retract", + id: String(taskId), + port: classified.port, + length: classified.length, + }); + } } else if (event.type === "object-delta" && (event as { port?: string }).port === "messages") { const task = graph.getTask(taskId); sink.emit({ diff --git a/examples/cli/src/ui/model/streamedText.test.ts b/examples/cli/src/ui/model/streamedText.test.ts index 1c7bcddced..5c53eb4da6 100644 --- a/examples/cli/src/ui/model/streamedText.test.ts +++ b/examples/cli/src/ui/model/streamedText.test.ts @@ -48,6 +48,42 @@ describe("foldTextDelta", () => { }); }); +describe("foldTextDelta on a text-retract", () => { + it("takes back the tail of its own port, so the retry appends after what remains", () => { + const out = fold([ + { type: "text-delta", port: "reasoning", textDelta: "Thinking" }, + { type: "text-delta", port: "text", textDelta: "First. " }, + { type: "text-delta", port: "text", textDelta: "Oops" }, + { type: "text-retract", port: "text", length: 4 }, + { type: "text-delta", port: "text", textDelta: "Second." }, + ]); + expect(out).toEqual({ text: "First. Second.", reasoning: "Thinking" }); + }); + + it("retracts reasoning on its own port, and clamps at the start", () => { + const out = fold([ + { type: "text-delta", port: "reasoning", textDelta: "abc" }, + { type: "text-delta", port: "text", textDelta: "kept" }, + { type: "text-retract", port: "reasoning", length: 10 }, + ]); + expect(out).toEqual({ text: "kept", reasoning: "" }); + }); + + it("returns the same object when there is nothing to take back", () => { + const prev: StreamedText = { text: "a", reasoning: "" }; + expect(foldTextDelta(prev, { type: "text-retract", port: "text", length: 0 })).toBe(prev); + expect(foldTextDelta(prev, { type: "text-retract", port: "reasoning", length: 3 })).toBe(prev); + }); + + it("removes the whole number of characters retractText removes", () => { + const prev: StreamedText = { text: "Pabc", reasoning: "" }; + expect(foldTextDelta(prev, { type: "text-retract", port: "text", length: 1.5 }).text).toBe( + "Pab" + ); + expect(foldTextDelta(prev, { type: "text-retract", port: "text", length: 0.5 })).toBe(prev); + }); +}); + describe("visibleReasoning", () => { it("shows thinking until the answer starts", () => { expect(visibleReasoning({ text: "", reasoning: "hmm" })).toBe("hmm"); diff --git a/examples/cli/src/ui/model/streamedText.ts b/examples/cli/src/ui/model/streamedText.ts index 755e568671..6b2df9d4e9 100644 --- a/examples/cli/src/ui/model/streamedText.ts +++ b/examples/cli/src/ui/model/streamedText.ts @@ -24,6 +24,8 @@ export interface TextDeltaLike { readonly type: string; readonly port?: string; readonly textDelta?: string; + /** On a `text-retract`: how many trailing characters to take back. */ + readonly length?: number; } /** Lines of a streaming panel shown at once; the oldest scroll off the top. */ @@ -48,11 +50,46 @@ export function classifyTextDelta(event: TextDeltaLike): ClassifiedDelta | undef return { port: event.port === "reasoning" ? "reasoning" : "text", delta }; } +/** A text retraction on a port worth showing: which one, and how much it takes back. */ +export interface ClassifiedRetract { + readonly port: "text" | "reasoning"; + readonly length: number; +} + +/** + * Reads a stream event as a `text-retract`, or `undefined` when it is not one + * that takes anything back. Ports are read as {@link classifyTextDelta} reads + * them, so a retraction lands on the text its deltas went to. + */ +export function classifyTextRetract(event: TextDeltaLike): ClassifiedRetract | undefined { + if (event.type !== "text-retract") return undefined; + // Rounded as `retractText` rounds it, so every fold downstream of this one, + // the web console's included, removes the same whole number of characters. + const length = Math.floor(event.length ?? 0); + if (!(length > 0)) return undefined; + return { port: event.port === "reasoning" ? "reasoning" : "text", length }; +} + +/** `text` without its last `length` characters, clamped at the start. */ +function dropTail(text: string, length: number): string { + return text.slice(0, Math.max(0, text.length - length)); +} + /** - * Folds one stream event into the accumulated text. Returns `prev` itself when - * nothing changed so a state setter can bail out of a re-render. + * Folds one stream event into the accumulated text: a `text-delta` appends and + * a `text-retract` takes its tail back. Returns `prev` itself when nothing + * changed so a state setter can bail out of a re-render. */ export function foldTextDelta(prev: StreamedText, event: TextDeltaLike): StreamedText { + const retract = classifyTextRetract(event); + if (retract) { + if (retract.port === "text") { + if (prev.text.length === 0) return prev; + return { text: dropTail(prev.text, retract.length), reasoning: prev.reasoning }; + } + if (prev.reasoning.length === 0) return prev; + return { text: prev.text, reasoning: dropTail(prev.reasoning, retract.length) }; + } const classified = classifyTextDelta(event); if (!classified) return prev; return classified.port === "text" diff --git a/examples/cli/src/ui/rows/ChatTaskRow.tsx b/examples/cli/src/ui/rows/ChatTaskRow.tsx index 6e8ce603b6..89936b259f 100644 --- a/examples/cli/src/ui/rows/ChatTaskRow.tsx +++ b/examples/cli/src/ui/rows/ChatTaskRow.tsx @@ -83,7 +83,8 @@ export function ChatTaskRow({ task, line }: TaskRowProps): React.ReactElement { () => messagesFromTask(task) ); // `streamed` is the in-progress assistant text (and thinking, kept apart by - // port) for the *current* turn. Appends on every text-delta; resets whenever + // port) for the *current* turn. Appends on every text-delta, drops the tail a + // text-retract takes back (a failed attempt's partial); resets whenever // history grows (meaning the previous turn's assistant message just landed, // or a new user turn has started), and on stream_end. const [streamed, setStreamed] = useState(EMPTY_STREAMED_TEXT); @@ -91,7 +92,7 @@ export function ChatTaskRow({ task, line }: TaskRowProps): React.ReactElement { useEffect(() => { const onChunk = (event: StreamEvent): void => { - if (event.type === "text-delta") { + if (event.type === "text-delta" || event.type === "text-retract") { setStreamed((prev) => foldTextDelta(prev, event)); } else if (event.type === "object-delta" && event.port === "messages") { // Each object-delta marks a turn boundary — clear the per-turn diff --git a/examples/cli/src/web/client/state.test.ts b/examples/cli/src/web/client/state.test.ts index ff485bf397..f35172a1a2 100644 --- a/examples/cli/src/web/client/state.test.ts +++ b/examples/cli/src/web/client/state.test.ts @@ -72,6 +72,19 @@ describe("reduceRunEvent", () => { expect(state.rows.get("t1")!.streamText).toBe("Hello"); }); + it("takes back the tail a retraction names, on its own port", () => { + const state = apply([ + added("t1"), + { k: "reasoning", id: "t1", delta: "Hmm" }, + { k: "text", id: "t1", delta: "First. Oops" }, + { k: "retract", id: "t1", port: "text", length: 4 }, + { k: "text", id: "t1", delta: "Second." }, + { k: "retract", id: "t1", port: "reasoning", length: 99 }, + ]); + expect(state.rows.get("t1")!.streamText).toBe("First. Second."); + expect(state.rows.get("t1")!.reasoningText).toBe(""); + }); + it("drops a row the run released", () => { const state = apply([added("t1"), { k: "task_removed", id: "t1" }]); expect(state.rows.has("t1")).toBe(false); diff --git a/examples/cli/src/web/client/state.ts b/examples/cli/src/web/client/state.ts index abac520685..286c97585f 100644 --- a/examples/cli/src/web/client/state.ts +++ b/examples/cli/src/web/client/state.ts @@ -229,6 +229,18 @@ export function reduceRunEvent( if (!row) return state; return withRow(state, event.id, { reasoningText: row.reasoningText + event.delta }); } + case "retract": { + const row = state.rows.get(event.id); + if (!row) return state; + const held = event.port === "text" ? row.streamText : row.reasoningText; + const kept = held.slice(0, Math.max(0, held.length - event.length)); + if (kept === held) return state; + return withRow( + state, + event.id, + event.port === "text" ? { streamText: kept } : { reasoningText: kept } + ); + } case "messages": return withRow(state, event.id, { messages: event.messages }); case "usage": diff --git a/examples/eval/src/commands/run.ts b/examples/eval/src/commands/run.ts index 7b0e6248f3..94d28a1845 100644 --- a/examples/eval/src/commands/run.ts +++ b/examples/eval/src/commands/run.ts @@ -6,7 +6,7 @@ import { withCli } from "@workglow/cli"; import type { StreamEvent } from "@workglow/task-graph"; -import { Workflow } from "@workglow/task-graph"; +import { retractionLength, Workflow } from "@workglow/task-graph"; import type { Command } from "commander"; import { EvalSweepTask, type EvalSweepTaskOutput } from "../evals/EvalSweepTask"; import type { ColumnOptions, DatasetContext } from "../evals/types"; @@ -167,6 +167,12 @@ async function runEval(kind: EvalKind, flags: RunFlags, stores: EvalStores): Pro const onStreamChunk: ((event: StreamEvent) => void) | undefined = flags.stream ? (event) => { if (event.type === "text-delta") process.stderr.write(event.textDelta); + else if (event.type === "text-retract" && retractionLength(event.length) > 0) { + // Already written, so it cannot be erased: say it no longer counts. + process.stderr.write( + `\n[attempt failed; the last ${retractionLength(event.length)} characters are withdrawn]\n` + ); + } } : undefined; // Through `withCli`, the seam that decides what a run does with itself. It diff --git a/examples/web/src/graph/TaskNode.tsx b/examples/web/src/graph/TaskNode.tsx index 228be6a064..e80e02e6ea 100644 --- a/examples/web/src/graph/TaskNode.tsx +++ b/examples/web/src/graph/TaskNode.tsx @@ -6,7 +6,7 @@ import { formatUsage } from "@workglow/ai"; import type { ITask, StreamEvent, Usage } from "@workglow/task-graph"; -import { TaskStatus } from "@workglow/task-graph"; +import { retractText, TaskStatus } from "@workglow/task-graph"; import { ArrayTask } from "@workglow/tasks"; import type { Node, NodeProps } from "@xyflow/react"; import { useEffect, useRef, useState } from "react"; @@ -119,6 +119,16 @@ export function TaskNode(props: NodeProps>) { streamingTextRef.current += event.textDelta; setStreamingText(streamingTextRef.current); } + } else if (event.type === "text-retract") { + // A failed attempt takes back what it streamed, from the same port. + const keep = (text: string): string => retractText(text, event.length); + if (event.port === "reasoning") { + reasoningTextRef.current = keep(reasoningTextRef.current); + setReasoningText(reasoningTextRef.current); + } else { + streamingTextRef.current = keep(streamingTextRef.current); + setStreamingText(streamingTextRef.current); + } } else if (event.type === "snapshot") { const text = typeof event.data === "string" ? event.data : JSON.stringify(event.data); streamingTextRef.current = text; diff --git a/packages/ai/src/capability/StreamEventAccumulator.ts b/packages/ai/src/capability/StreamEventAccumulator.ts index 31825f9975..f1db704fea 100644 --- a/packages/ai/src/capability/StreamEventAccumulator.ts +++ b/packages/ai/src/capability/StreamEventAccumulator.ts @@ -9,6 +9,7 @@ import { mergeUsage, REFUSAL_CATEGORY_OUTPUT_KEY, REFUSAL_OUTPUT_KEY, + retractText, USAGE_OUTPUT_KEY, } from "@workglow/task-graph"; @@ -20,6 +21,7 @@ const isNonEmptyObject = (v: unknown): v is Record => * of {@link StreamEvent}s. * * - `text-delta` per port → concatenated per port (`Map`). + * - `text-retract` per port → the port's text loses its last `length` chars. * - `object-delta` non-array → replace semantics (each delta is a more * complete partial snapshot — last wins per port). * - `object-delta` array → upsert by `id` per port. @@ -83,6 +85,14 @@ export class StreamEventAccumulator { this.textAccumulator.set(e.port, (this.textAccumulator.get(e.port) ?? "") + e.textDelta); return; } + case "text-retract": { + this.hasTextDeltas = true; + this.lastEventType = "text-retract"; + const e = event as Extract, { type: "text-retract" }>; + const text = this.textAccumulator.get(e.port); + if (text !== undefined) this.textAccumulator.set(e.port, retractText(text, e.length)); + return; + } case "object-delta": { this.hasObjectDeltas = true; this.lastEventType = "object-delta"; diff --git a/packages/ai/src/capability/__tests__/StreamEventAccumulator.test.ts b/packages/ai/src/capability/__tests__/StreamEventAccumulator.test.ts index a6444593b8..6cdbf96bcd 100644 --- a/packages/ai/src/capability/__tests__/StreamEventAccumulator.test.ts +++ b/packages/ai/src/capability/__tests__/StreamEventAccumulator.test.ts @@ -45,6 +45,24 @@ describe("StreamEventAccumulator", () => { expect(acc.materialize()).toEqual({ text: "Hello" }); }); + it("takes back a port's tail on text-retract, so later deltas append after what remains", () => { + const acc = new StreamEventAccumulator(); + acc.observe({ type: "text-delta", port: "text", textDelta: "First. Oops" } as StreamEvent); + acc.observe({ type: "text-delta", port: "notes", textDelta: "kept" } as StreamEvent); + acc.observe({ type: "text-retract", port: "text", length: 4 } as StreamEvent); + acc.observe({ type: "text-delta", port: "text", textDelta: "Second." } as StreamEvent); + acc.observeFinish({ type: "finish", data: {} }); + expect(acc.materialize()).toEqual({ text: "First. Second.", notes: "kept" }); + }); + + it("falls back to the finish payload for a port retracted to nothing", () => { + const acc = new StreamEventAccumulator(); + acc.observe({ type: "text-delta", port: "text", textDelta: "partial" } as StreamEvent); + acc.observe({ type: "text-retract", port: "text", length: 99 } as StreamEvent); + acc.observeFinish({ type: "finish", data: { text: "stated" } }); + expect(acc.materialize()).toEqual({ text: "stated" }); + }); + it("replaces object-delta non-array payloads (latest wins)", () => { const acc = new StreamEventAccumulator(); acc.observe({ type: "object-delta", port: "obj", objectDelta: { a: 1 } } as StreamEvent); diff --git a/packages/ai/src/task/AgentRoundAttempt.ts b/packages/ai/src/task/AgentRoundAttempt.ts index 1ea8f72afc..668b8fdfa2 100644 --- a/packages/ai/src/task/AgentRoundAttempt.ts +++ b/packages/ai/src/task/AgentRoundAttempt.ts @@ -5,7 +5,13 @@ */ import { RetryableJobError } from "@workglow/job-queue"; -import type { IExecuteContext, StreamEvent, StreamTextDelta } from "@workglow/task-graph"; +import type { + IExecuteContext, + StreamEvent, + StreamTextDelta, + StreamTextRetract, +} from "@workglow/task-graph"; +import { retractionLength } from "@workglow/task-graph"; import { createEmitQueue } from "../capability/emitQueue"; import type { AgentRoundRunner } from "./AgentRoundRunner"; import { AGENT_ROUND_RUNNER } from "./AgentRoundRunner"; @@ -16,6 +22,9 @@ import type { ChatMessage } from "./ChatMessage"; import type { ToolCallingTaskInput, ToolCallingTaskOutput } from "./ToolCallingTask"; import { ToolCallingTask } from "./ToolCallingTask"; +/** What a round puts on the turn's wire: its text and reasoning, and the retraction of a failed attempt's. */ +type RoundStreamEvent = StreamTextDelta | StreamTextRetract; + /** Between two rounds' reasoning, which would otherwise run together into one thought. */ export const REASONING_ROUND_SEPARATOR = "\n\n"; @@ -68,13 +77,13 @@ function roundRunnerOf(context: IExecuteContext): AgentRoundRunner | undefined { * failure (a rate limit, an overload, a call abandoned after `timeoutMs`) up to * `maxRetries` times, waiting as the provider asks. * - * Yields the round's text and reasoning as this task's own, but only from the - * attempt that settled: a failed attempt's partial cannot be taken back from - * the accumulated port, so it is never forwarded. + * Yields the round's text and reasoning as this task's own, as they arrive. A + * failed attempt takes back exactly what it yielded with a `text-retract` per + * port before the retry starts, so the folded ports read as if it never ran. */ export async function* settleRound( options: AgentRoundOptions -): AsyncGenerator { +): AsyncGenerator { const { context } = options; const captured: { output: ToolCallingTaskOutput | undefined } = { output: undefined }; let attempts = 0; @@ -116,7 +125,7 @@ export async function* settleRound( async function* attemptRound( options: AgentRoundOptions, captured: { output: ToolCallingTaskOutput | undefined } -): AsyncGenerator { +): AsyncGenerator { const { round, turn, context, timeoutMs } = options; const roundInput: ToolCallingTaskInput = { model: turn.model, @@ -134,7 +143,7 @@ async function* attemptRound( maxTokens: turn.maxTokens, }; const hostRunner = roundRunnerOf(context); - const queue = createEmitQueue(); + const queue = createEmitQueue(); let off = (): void => {}; let runRound: () => Promise; // A round that never answers holds the turn forever: a provider can accept @@ -157,6 +166,7 @@ async function* attemptRound( onTextDelta: (delta) => queue.push({ type: "text-delta", port: "text", textDelta: delta }), onReasoningDelta: (delta) => queue.push({ type: "text-delta", port: "reasoning", textDelta: delta }), + onTextRetract: (port, length) => queue.push({ type: "text-retract", port, length }), // Progress is advisory: a failed update must not surface as an // unhandled rejection, nor fail a round that is otherwise fine. onProgress: (progress, message) => { @@ -168,10 +178,10 @@ async function* attemptRound( context.own(task); abandon = () => task.abort(); off = task.subscribe("stream_chunk", (event: StreamEvent) => { - if (event.type !== "text-delta") return; + if (event.type !== "text-delta" && event.type !== "text-retract") return; const port = event.port ?? "text"; if (port !== "text" && port !== "reasoning") return; - queue.push({ type: "text-delta", port, textDelta: event.textDelta }); + queue.push({ ...event, port }); }); runRound = () => task.run(roundInput); } @@ -182,6 +192,7 @@ async function* attemptRound( timedOut = true; abandon(); }, timeoutMs); + let settled = false; const run = (async () => { try { captured.output = await runRound(); @@ -191,6 +202,7 @@ async function* attemptRound( } throw err; } finally { + settled = true; if (timer !== undefined) clearTimeout(timer); off(); queue.close(); @@ -200,24 +212,67 @@ async function* attemptRound( // unhandled; it is re-thrown below, once the events already produced have // reached the caller. run.catch(() => {}); - // Held until the attempt settles: the stream accumulates every delta it is - // handed into the task's text and reasoning ports and offers no way to take - // one back, so either forwarded from an attempt that then fails and is - // retried would be joined to the retry's. A failed attempt throws here and - // its text and reasoning are dropped together. - const held: StreamTextDelta[] = []; - for await (const event of queue.iterable) held.push(event); - await run; - let separated = !options.reasonedBefore; - for (const event of held) { - // Rounds are told apart on the wire exactly as in the settled output, and - // only once this round has something to say: a separator ahead of nothing - // would leave the streamed reasoning ending differently than the final one. - if (!separated && event.port === "reasoning") { + // Forwarded as they arrive, so a host sees the reply as it is written. What + // this attempt forwarded is counted per port, because the stream folds every + // delta into the task's text and reasoning ports: an attempt that fails takes + // exactly that much back with a `text-retract`, and a retry or the next round + // appends after the earlier rounds as if it never ran. + let textSent = 0; + let reasoningSent = 0; + // Rounds are told apart on the wire exactly as in the settled output, and + // only once this round has something to say: a separator ahead of nothing + // would leave the streamed reasoning ending differently than the final one. + let separatorSent = false; + try { + for await (const event of queue.iterable) { + if (event.type === "text-retract") { + // Counted as the fold removes it: a NaN would poison the counters and a + // fraction would drift them from the text, so a failed attempt's own + // retraction would take back too little, or an earlier round's text. + const requested = retractionLength(event.length); + if (requested === 0) continue; + if (event.port === "text") { + const length = Math.min(requested, textSent); + textSent -= length; + if (length > 0) yield { type: "text-retract", port: "text", length }; + } else { + let length = Math.min(requested, reasoningSent); + reasoningSent -= length; + // Everything this round reasoned was taken back, so the separator + // ahead of it goes too. + if (reasoningSent === 0 && separatorSent) { + length += REASONING_ROUND_SEPARATOR.length; + separatorSent = false; + } + if (length > 0) yield { type: "text-retract", port: "reasoning", length }; + } + continue; + } if (event.textDelta.length === 0) continue; - yield { type: "text-delta", port: "reasoning", textDelta: REASONING_ROUND_SEPARATOR }; - separated = true; + if (event.port === "text") { + textSent += event.textDelta.length; + } else { + if (options.reasonedBefore && !separatorSent) { + yield { type: "text-delta", port: "reasoning", textDelta: REASONING_ROUND_SEPARATOR }; + separatorSent = true; + } + reasoningSent += event.textDelta.length; + } + yield event; + } + await run; + } catch (err) { + if (textSent > 0) yield { type: "text-retract", port: "text", length: textSent }; + const reasoningLength = reasoningSent + (separatorSent ? REASONING_ROUND_SEPARATOR.length : 0); + if (reasoningLength > 0) { + yield { type: "text-retract", port: "reasoning", length: reasoningLength }; } - yield event; + throw err; + } finally { + // Deltas are yielded while the call is still running, so a consumer that + // stops reading leaves this generator at a `yield` with the call in flight. + // Abandoning it ends the request; the run's own `finally` then clears the + // timer and the listener. + if (!settled) abandon(); } } diff --git a/packages/ai/src/task/AgentRoundRunner.ts b/packages/ai/src/task/AgentRoundRunner.ts index 5d64e482d7..fe8474c349 100644 --- a/packages/ai/src/task/AgentRoundRunner.ts +++ b/packages/ai/src/task/AgentRoundRunner.ts @@ -15,10 +15,20 @@ export interface AgentRoundContext { readonly onTextDelta: (delta: string) => void; /** * The round's reasoning as it streams, re-emitted as the turn's own on its - * `reasoning` port. Held with the round's text until the attempt settles, so - * a round that fails and is retried contributes neither. + * `reasoning` port. Forwarded live like the round's text, and both are taken + * back with a `text-retract` if the attempt fails, so a round that is retried + * contributes only its retry's. */ readonly onReasoningDelta: (delta: string) => void; + /** + * Takes back the last `length` characters the round streamed on `port`, + * counted as `retractionLength` counts them: for a host whose own call + * retracts text it already forwarded, such as a `ToolCallingTask` whose + * provider withdrew part of a reply. The turn passes it on as its own + * `text-retract`, capped at what this attempt forwarded on that port, so a + * host can never take back an earlier round's text. + */ + readonly onTextRetract: (port: "text" | "reasoning", length: number) => void; /** * Progress while the round has nothing to say yet — a model being downloaded * or loaded. `progress` is 0..100, or undefined when there is no denominator. diff --git a/packages/ai/src/task/TextGenerationTask.ts b/packages/ai/src/task/TextGenerationTask.ts index 37fd3571b0..64479a99a5 100644 --- a/packages/ai/src/task/TextGenerationTask.ts +++ b/packages/ai/src/task/TextGenerationTask.ts @@ -11,7 +11,7 @@ import type { StreamEvent, TaskConfig, } from "@workglow/task-graph"; -import { CreateWorkflow, Workflow } from "@workglow/task-graph"; +import { CreateWorkflow, retractText, Workflow } from "@workglow/task-graph"; import type { DataPortSchema } from "@workglow/util/schema"; import type { Capability } from "../capability/Capabilities"; import type { AiJobInput } from "../job/AiJob"; @@ -282,6 +282,8 @@ export class TextGenerationTask extends StreamingAiTask< for await (const event of super.executeStream(input, context)) { if (event.type === "text-delta" && (event.port ?? "text") === "text") { text += event.textDelta; + } else if (event.type === "text-retract" && (event.port ?? "text") === "text") { + text = retractText(text, event.length); } if (event.type === "finish") { await this.finalizeCheckpoint(input, text); diff --git a/packages/ai/src/task/ToolCallingTask.ts b/packages/ai/src/task/ToolCallingTask.ts index c10180f70b..c70a2d3699 100644 --- a/packages/ai/src/task/ToolCallingTask.ts +++ b/packages/ai/src/task/ToolCallingTask.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { CreateWorkflow, getTaskConstructors, Workflow } from "@workglow/task-graph"; +import { CreateWorkflow, getTaskConstructors, retractText, Workflow } from "@workglow/task-graph"; import type { CachePolicy, @@ -615,6 +615,10 @@ export class ToolCallingTask extends StreamingAiTask< text += event.textDelta; } else if (event.type === "text-delta" && event.port === "reasoning") { reasoning += event.textDelta; + } else if (event.type === "text-retract" && (event.port ?? "text") === "text") { + text = retractText(text, event.length); + } else if (event.type === "text-retract" && event.port === "reasoning") { + reasoning = retractText(reasoning, event.length); } else if (event.type === "object-delta" && event.port === "nativeTurn") { nativeTurn = event.objectDelta as unknown as ToolCallingTaskOutput["nativeTurn"]; } else if (event.type === "object-delta" && event.port === "toolCalls") { diff --git a/packages/ai/src/task/__tests__/AgentRoundAttempt.test.ts b/packages/ai/src/task/__tests__/AgentRoundAttempt.test.ts index 64c29f5f9e..ae59972634 100644 --- a/packages/ai/src/task/__tests__/AgentRoundAttempt.test.ts +++ b/packages/ai/src/task/__tests__/AgentRoundAttempt.test.ts @@ -5,7 +5,8 @@ */ import { PermanentJobError, RetryableJobError } from "@workglow/job-queue"; -import type { IExecuteContext, StreamTextDelta } from "@workglow/task-graph"; +import type { IExecuteContext, StreamTextDelta, StreamTextRetract } from "@workglow/task-graph"; +import { retractText } from "@workglow/task-graph"; import { Container, ServiceRegistry } from "@workglow/util"; import { describe, expect, it } from "vitest"; import type { AgentRoundOptions, SettledRound } from "../AgentRoundAttempt"; @@ -68,11 +69,23 @@ function options( }; } +type RoundEvent = StreamTextDelta | StreamTextRetract; + +/** Folds one port of a round's events the way the stream does, retractions included. */ +function folded(events: readonly RoundEvent[], port: string): string { + let text = ""; + for (const event of events) { + if (event.port !== port) continue; + text = event.type === "text-retract" ? retractText(text, event.length) : text + event.textDelta; + } + return text; +} + /** Drains a round, returning what it forwarded and how it settled. */ async function drain( - run: AsyncGenerator -): Promise<{ readonly events: StreamTextDelta[]; readonly settled: SettledRound }> { - const events: StreamTextDelta[] = []; + run: AsyncGenerator +): Promise<{ readonly events: RoundEvent[]; readonly settled: SettledRound }> { + const events: RoundEvent[] = []; for (;;) { const next = await run.next(); if (next.done) return { events, settled: next.value }; @@ -91,7 +104,10 @@ describe("settleRound", () => { ]); const { events, settled } = await drain(settleRound(options(context))); expect(called()).toBe(1); - expect(events.map((event) => event.textDelta)).toEqual(["do", "ne"]); + expect(events).toEqual([ + { type: "text-delta", port: "text", textDelta: "do" }, + { type: "text-delta", port: "text", textDelta: "ne" }, + ]); expect(settled).toEqual({ status: "settled", output: ANSWER, attempts: 1 }); }); @@ -110,7 +126,7 @@ describe("settleRound", () => { expect(sent?.messages).toHaveLength(1); }); - it("retries a retryable failure and drops the failed attempt's partial text", async () => { + it("retries a retryable failure and takes back the failed attempt's partial text", async () => { const { context, called, progress } = contextFor([ async (_input, round) => { round.onTextDelta("half an ans"); @@ -124,11 +140,41 @@ describe("settleRound", () => { ]); const { events, settled } = await drain(settleRound(options(context))); expect(called()).toBe(2); - expect(events.map((event) => event.textDelta)).toEqual(["done"]); + // The partial is forwarded as it streams, then retracted before the retry, + // so the folded port holds only the retry's text. + expect(events).toEqual([ + { type: "text-delta", port: "text", textDelta: "half an ans" }, + { type: "text-retract", port: "text", length: "half an ans".length }, + { type: "text-delta", port: "text", textDelta: "done" }, + ]); + expect(folded(events, "text")).toBe("done"); expect(settled).toMatchObject({ status: "settled", attempts: 2 }); expect(progress).toEqual(["Retrying in 1s", "Thinking (attempt 2)"]); }); + it("abandons the call in flight when the consumer stops reading", async () => { + let signal: AbortSignal | undefined; + const { context } = contextFor([ + async (_input, round) => { + signal = round.signal; + round.onTextDelta("partial"); + await new Promise((resolve) => + round.signal.addEventListener("abort", () => resolve()) + ); + throw new Error("aborted"); + }, + ]); + const run = settleRound(options(context)); + expect((await run.next()).value).toEqual({ + type: "text-delta", + port: "text", + textDelta: "partial", + }); + expect(signal?.aborted).toBe(false); + await run.return({ status: "out-of-time" }); + expect(signal?.aborted).toBe(true); + }); + it("ends out of time rather than sleeping through a wait that passes the deadline", async () => { const { context, called, progress } = contextFor([ async () => { @@ -179,6 +225,31 @@ describe("settleRound", () => { ).rejects.toThrow(/Round 1 had no answer within 10 ms/); }); + it("passes on a host runner's retraction, capped at what this attempt forwarded", async () => { + const { context } = contextFor([ + async (_input, round) => { + round.onTextDelta("draft"); + round.onTextRetract("text", 2); + round.onTextDelta("y"); + // More than the attempt sent: only its own text can be taken back. + round.onTextRetract("text", 100); + round.onTextDelta("done"); + round.onTextRetract("reasoning", 3); + return ANSWER; + }, + ]); + const { events, settled } = await drain(settleRound(options(context))); + expect(events).toEqual([ + { type: "text-delta", port: "text", textDelta: "draft" }, + { type: "text-retract", port: "text", length: 2 }, + { type: "text-delta", port: "text", textDelta: "y" }, + { type: "text-retract", port: "text", length: 4 }, + { type: "text-delta", port: "text", textDelta: "done" }, + ]); + expect(folded(events, "text")).toBe("done"); + expect(settled).toMatchObject({ status: "settled", attempts: 1 }); + }); + it("sets a round's reasoning apart from an earlier round's, once it has some", async () => { const { context } = contextFor([ async (_input, round) => { diff --git a/packages/ai/src/task/base/StreamingAiTask.ts b/packages/ai/src/task/base/StreamingAiTask.ts index 0d2764f4d3..affa725154 100644 --- a/packages/ai/src/task/base/StreamingAiTask.ts +++ b/packages/ai/src/task/base/StreamingAiTask.ts @@ -165,7 +165,7 @@ export class StreamingAiTask< usage = mergeUsage(usage, event.usage); } - if (event.type === "text-delta") { + if (event.type === "text-delta" || event.type === "text-retract") { yield { ...event, port: event.port ?? defaultPort } as StreamEvent; } else if (event.type === "object-delta") { yield { ...event, port: event.port ?? defaultPort } as StreamEvent; diff --git a/packages/ai/src/task/base/chatTurn.ts b/packages/ai/src/task/base/chatTurn.ts index 52290b4a89..3848e6ab07 100644 --- a/packages/ai/src/task/base/chatTurn.ts +++ b/packages/ai/src/task/base/chatTurn.ts @@ -11,7 +11,7 @@ import type { TaskOutput, Usage, } from "@workglow/task-graph"; -import { mergeUsage } from "@workglow/task-graph"; +import { mergeUsage, retractText } from "@workglow/task-graph"; import type { AiEmit } from "../../capability/AiEmit"; import type { IAiExecutionStrategy } from "../../execution/IAiExecutionStrategy"; import type { AiJobInput } from "../../job/AiJob"; @@ -23,7 +23,7 @@ export interface ChatTurnArgs { readonly jobInput: AiJobInput; readonly context: IExecuteContext; readonly runnerId: string | undefined; - /** Port stamped onto port-less `text-delta` events from the provider. */ + /** Port stamped onto port-less `text-delta` / `text-retract` events from the provider. */ readonly textPort: string; } @@ -67,6 +67,12 @@ export function runChatTurn(args: ChatTurnArgs): Chat ...event, port: event.port ?? args.textPort, } as StreamEvent); + } else if (event.type === "text-retract") { + text = retractText(text, event.length); + queue.push({ + ...event, + port: event.port ?? args.textPort, + } as StreamEvent); } else if (event.type === "finish") { // Invariant: inner turn run-fns must be text.generation only; // finish.data from inner turns is intentionally discarded. If diff --git a/packages/task-graph/src/EXECUTION_MODEL.md b/packages/task-graph/src/EXECUTION_MODEL.md index 273a9237a3..e61b5761a4 100644 --- a/packages/task-graph/src/EXECUTION_MODEL.md +++ b/packages/task-graph/src/EXECUTION_MODEL.md @@ -340,7 +340,7 @@ The **SQL** backings share one implementation. `TabularBlobChunkStore` persists ### Live cross-process stream transport (`onStream`) -`JobHandle.outputStream` (above) streams the _final_ binary result out of the cache after completion. Live **stream events** (`text-delta` / `object-delta` / `binary-delta` / `snapshot` / `finish` / `phase`) are a separate channel with two serialization boundaries: +`JobHandle.outputStream` (above) streams the _final_ binary result out of the cache after completion. Live **stream events** (`text-delta` / `text-retract` / `object-delta` / `binary-delta` / `snapshot` / `finish` / `phase`) are a separate channel with two serialization boundaries: - **Worker-thread boundary** (`@workglow/util`): provider run-fns execute inside a worker; `WorkerServerBase.postStreamChunk` posts each event to the main thread. It extracts transferables from the event (mirroring `postResult`), so binary payloads — `binary-delta` buffers and `snapshot` image bytes — transfer zero-copy instead of being structure-cloned; non-binary events yield an empty transfer list and clone byte-for-byte as before. (The same-process `job_stream` fan-out inside `@workglow/job-queue` has no structured-clone boundary — events pass by reference — so transferables are a no-op there.) diff --git a/packages/task-graph/src/task-graph/StreamPump.ts b/packages/task-graph/src/task-graph/StreamPump.ts index 4728f4fbb5..b982e18f17 100644 --- a/packages/task-graph/src/task-graph/StreamPump.ts +++ b/packages/task-graph/src/task-graph/StreamPump.ts @@ -774,11 +774,16 @@ export class StreamPump { /** * Returns true if an event carries a port-specific delta (text-delta, - * object-delta, or binary-delta). + * text-retract, object-delta, or binary-delta). A retraction belongs to its + * port like the text it takes back: an edge carrying that port's deltas must + * carry it too, and an edge for another port must not. */ private static isPortDelta(event: StreamEvent): event is StreamEvent & { port: string } { return ( - event.type === "text-delta" || event.type === "object-delta" || event.type === "binary-delta" + event.type === "text-delta" || + event.type === "text-retract" || + event.type === "object-delta" || + event.type === "binary-delta" ); } @@ -839,7 +844,7 @@ export class StreamPump { } // A task's own reporting is not content on a port, so it never // reaches an edge. The port filter above cannot decide this: it - // recognises only the three delta types, so an event carrying no + // recognises only the per-port delta types, so an event carrying no // `port` passes it by construction. if (isDataflowExcluded(event)) { return; diff --git a/packages/task-graph/src/task-graph/__tests__/StreamTextRetract.test.ts b/packages/task-graph/src/task-graph/__tests__/StreamTextRetract.test.ts new file mode 100644 index 0000000000..3d4dace979 --- /dev/null +++ b/packages/task-graph/src/task-graph/__tests__/StreamTextRetract.test.ts @@ -0,0 +1,339 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * `text-retract` takes back the tail of a port's streamed text. The processor + * truncates its accumulator, gives up a cache sink it can no longer keep + * honest, and the pump carries the retraction down the edge of its own port + * only, so a consumer folding that port truncates too. + */ + +import type { IExecuteContext, StreamEvent, StreamSink, TaskOutput } from "@workglow/task-graph"; +import { + Dataflow, + isCacheRef, + makeCacheRef, + retractionLength, + retractText, + Task, + TaskError, + TaskGraph, + TaskGraphRunner, + TaskRegistry, +} from "@workglow/task-graph"; +import type { DataPortSchema } from "@workglow/util/schema"; +import { beforeAll, describe, expect, it } from "vitest"; + +type TextOut = { text: string }; + +const TEXT_SCHEMA = { + type: "object", + properties: { text: { type: "string", "x-stream": "append" } }, + additionalProperties: false, +} as const satisfies DataPortSchema; + +/** Streams the events it was constructed with, verbatim. */ +class ScriptedTextTask extends Task, TextOut> { + public static override type = "StreamTextRetractTest_Scripted"; + public static override category = "Test"; + public static override cacheable = false; + public script: readonly StreamEvent[] = []; + + public static override outputSchema(): DataPortSchema { + return TEXT_SCHEMA; + } + + async *executeStream(): AsyncIterable> { + for (const event of this.script) yield event; + } +} + +beforeAll(() => { + TaskRegistry.registerTask(ScriptedTextTask as any); +}); + +const delta = (textDelta: string, port = "text"): StreamEvent => ({ + type: "text-delta", + port, + textDelta, +}); +const retract = (length: number, port = "text"): StreamEvent => ({ + type: "text-retract", + port, + length, +}); +const FINISH: StreamEvent = { type: "finish", data: {} }; + +/** Folds a port's text the way every consumer of the event does. */ +function fold(events: readonly StreamEvent[], port: string): string { + let text = ""; + for (const event of events) { + if (event.type === "text-delta" && event.port === port) text += event.textDelta; + else if (event.type === "text-retract" && event.port === port) { + text = retractText(text, event.length); + } + } + return text; +} + +function makeSink(): { sink: StreamSink; settled: Promise<"written" | Error> } { + let settle: (v: "written" | Error) => void = () => {}; + const settled = new Promise<"written" | Error>((resolve) => { + settle = resolve; + }); + const sink: StreamSink = { + mode: "append", + write: async (chunks) => { + try { + for await (const _ of chunks) { + // drain + } + } catch (err) { + settle(err as Error); + throw err; + } + settle("written"); + return makeCacheRef({ $ref: "inmem://retract/text", port: "text", mode: "append", size: 0 }); + }, + }; + return { sink, settled }; +} + +async function runProcessor( + task: ScriptedTextTask, + options: { shouldAccumulate: boolean; refSinks?: ReadonlyMap } +): Promise { + const processor = (task as any).runner.streamProcessor as { + run(input: any, ctx: any, deps: any): Promise; + }; + const ctx = { + abortController: new AbortController(), + shouldAccumulate: options.shouldAccumulate, + registry: undefined, + runOutputData: {}, + dispose: () => {}, + } as any; + return (await processor.run({}, ctx, { + registry: undefined as any, + resourceScope: undefined, + inputStreams: undefined, + onProgress: async () => {}, + own: (t: T) => t, + refSinks: options.refSinks, + })) as TextOut; +} + +describe("retractText", () => { + it("drops the tail, clamps at the start, and ignores a non-positive length", () => { + expect(retractText("Hello there", 5)).toBe("Hello "); + expect(retractText("Hi", 10)).toBe(""); + expect(retractText("Hi", 0)).toBe("Hi"); + expect(retractText("Hi", -3)).toBe("Hi"); + }); + + it("removes a whole number of characters: a fraction rounds down and NaN removes nothing", () => { + expect(retractionLength(1.5)).toBe(1); + expect(retractText("Pabc", 1.5)).toBe("Pab"); + expect(retractionLength(Number.NaN)).toBe(0); + expect(retractText("Hi", Number.NaN)).toBe("Hi"); + expect(retractText("Hi", Number.POSITIVE_INFINITY)).toBe(""); + }); +}); + +describe("StreamProcessor — text-retract", () => { + it("truncates the accumulated port, so later deltas append after what remains", async () => { + const task = new ScriptedTextTask(); + task.script = [delta("Hello "), delta("there"), retract(5), delta("world"), FINISH]; + const chunks: StreamEvent[] = []; + task.on("stream_chunk", (event: StreamEvent) => chunks.push(event)); + + const output = await task.run({}); + + expect(output.text).toBe("Hello world"); + // The retraction reaches subscribers, and the enriched finish carries the + // folded value rather than every delta ever sent. + expect(chunks.filter((event) => event.type === "text-retract")).toEqual([retract(5)]); + const finish = chunks.find( + (event): event is Extract => event.type === "finish" + ); + expect(finish?.data.text).toBe("Hello world"); + expect(fold(chunks, "text")).toBe("Hello world"); + }); + + it("clamps a retraction longer than the port at the start", async () => { + const task = new ScriptedTextTask(); + task.script = [delta("abc"), retract(100), delta("x"), FINISH]; + expect((await task.run({})).text).toBe("x"); + }); + + it("settles a port retracted to nothing on the finish payload", async () => { + const task = new ScriptedTextTask(); + task.script = [delta("partial"), retract(7), { type: "finish", data: { text: "stated" } }]; + expect((await task.run({})).text).toBe("stated"); + }); + + it("fails a port's cache sink once text it holds is retracted, and settles the port inline", async () => { + const task = new ScriptedTextTask(); + task.script = [delta("Hello "), delta("there"), retract(5), delta("world"), FINISH]; + const { sink, settled } = makeSink(); + + const output = await runProcessor(task, { + shouldAccumulate: true, + refSinks: new Map([["text", sink]]), + }); + + // The bytes already teed were "Hello there": the sink must not commit them. + const outcome = await settled; + expect(outcome).toBeInstanceOf(TaskError); + expect(String(outcome)).toMatch(/retracted/); + // No ref minted from withdrawn text; the accumulator holds the real value. + expect(isCacheRef(output.text)).toBe(false); + expect(output.text).toBe("Hello world"); + }); + + it("keeps a sink that was never written to before the retraction", async () => { + const task = new ScriptedTextTask(); + task.script = [retract(3), delta("fresh"), FINISH]; + const { sink, settled } = makeSink(); + + const output = await runProcessor(task, { + shouldAccumulate: true, + refSinks: new Map([["text", sink]]), + }); + + expect(await settled).toBe("written"); + expect(isCacheRef(output.text)).toBe(true); + }); + + it("keeps a written sink when the retraction removes nothing, even with no accumulator", async () => { + const task = new ScriptedTextTask(); + task.script = [delta("Hello"), retract(0.5), FINISH]; + const { sink, settled } = makeSink(); + + const output = await runProcessor(task, { + shouldAccumulate: false, + refSinks: new Map([["text", sink]]), + }); + + expect(await settled).toBe("written"); + expect(isCacheRef(output.text)).toBe(true); + }); + + it("refuses to settle a sunk port it kept no copy of, unless the finish states it", async () => { + const lost = new ScriptedTextTask(); + lost.script = [delta("Hello there"), retract(5), delta("world"), FINISH]; + const lostSink = makeSink(); + await expect( + runProcessor(lost, { shouldAccumulate: false, refSinks: new Map([["text", lostSink.sink]]) }) + ).rejects.toThrow(/retracted text on port "text"/); + expect(await lostSink.settled).toBeInstanceOf(TaskError); + + const stated = new ScriptedTextTask(); + stated.script = [ + delta("Hello there"), + retract(5), + delta("world"), + { type: "finish", data: { text: "Hello world" } }, + ]; + const statedSink = makeSink(); + const output = await runProcessor(stated, { + shouldAccumulate: false, + refSinks: new Map([["text", statedSink.sink]]), + }); + expect(output.text).toBe("Hello world"); + }); +}); + +// ========================================================================== +// Edges: a retraction travels with its port's deltas, and only with them. +// ========================================================================== + +type TwoPorts = { text: string; notes: string }; + +class TwoPortSource extends Task, TwoPorts> { + public static override type = "StreamTextRetractTest_TwoPortSource"; + public static override category = "Test"; + public static override cacheable = false; + + public static override inputSchema(): DataPortSchema { + return { type: "object", properties: {}, additionalProperties: false } as const; + } + public static override outputSchema(): DataPortSchema { + return { + type: "object", + properties: { + text: { type: "string", "x-stream": "append" }, + notes: { type: "string", "x-stream": "append" }, + }, + additionalProperties: false, + } as const satisfies DataPortSchema; + } + async *executeStream(): AsyncIterable> { + yield delta("Hello there") as StreamEvent; + yield delta("note", "notes") as StreamEvent; + yield retract(5) as StreamEvent; + yield delta("world") as StreamEvent; + yield { type: "finish", data: {} as TwoPorts }; + } +} + +/** Append passthrough that records what reached it and re-streams it as its own. */ +class RecordingConsumer extends Task<{ text: string }, TextOut> { + public static override type = "StreamTextRetractTest_RecordingConsumer"; + public static override category = "Test"; + public static override cacheable = false; + public received: StreamEvent[] = []; + + public static override inputSchema(): DataPortSchema { + return TEXT_SCHEMA; + } + public static override outputSchema(): DataPortSchema { + return TEXT_SCHEMA; + } + async *executeStream( + _input: { text: string }, + ctx: IExecuteContext + ): AsyncIterable> { + const stream = ctx.inputStreams?.get("text"); + if (stream) { + const reader = stream.getReader(); + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + if (value.type === "text-delta" || value.type === "text-retract") { + this.received.push(value); + yield { ...value, port: "text" } as StreamEvent; + } + } + } finally { + reader.releaseLock(); + } + } + yield { type: "finish", data: {} as TextOut }; + } +} + +describe("StreamPump — text-retract on a port-filtered edge", () => { + it("forwards a retraction to its own port's consumer and to no other", async () => { + const graph = new TaskGraph(); + const source = new TwoPortSource({ id: "source" }); + const textConsumer = new RecordingConsumer({ id: "textConsumer" }); + const notesConsumer = new RecordingConsumer({ id: "notesConsumer" }); + graph.addTasks([source, textConsumer, notesConsumer]); + graph.addDataflow(new Dataflow("source", "text", "textConsumer", "text")); + graph.addDataflow(new Dataflow("source", "notes", "notesConsumer", "text")); + + const results = await new TaskGraphRunner(graph).runGraph({}); + + expect(textConsumer.received).toEqual([delta("Hello there"), retract(5), delta("world")]); + expect(notesConsumer.received).toEqual([delta("note", "notes")]); + const text = results.find((r) => r.id === "textConsumer")!.data as TaskOutput; + expect(text.text).toBe("Hello world"); + const notes = results.find((r) => r.id === "notesConsumer")!.data as TaskOutput; + expect(notes.text).toBe("note"); + }); +}); diff --git a/packages/task-graph/src/task/StreamProcessor.ts b/packages/task-graph/src/task/StreamProcessor.ts index 5a4c943ccf..0126e5481d 100644 --- a/packages/task-graph/src/task/StreamProcessor.ts +++ b/packages/task-graph/src/task/StreamProcessor.ts @@ -24,6 +24,8 @@ import { mergeUsage, REFUSAL_CATEGORY_OUTPUT_KEY, REFUSAL_OUTPUT_KEY, + retractionLength, + retractText, USAGE_OUTPUT_KEY, } from "./StreamTypes"; import { TaskAbortedError, TaskError } from "./TaskError"; @@ -180,11 +182,14 @@ export class StreamProcessor // output at the port slot. The codec (chosen by the sink's mode) turns each // text/object/binary delta into the ordered bytes the sink persists. const routers = new Map(); + // Ports whose router was failed by a `text-retract` (see that case below); + // they never get a router again in this run. + const detachedPorts = new Set(); const ensureRouter = ( port: string ): { router: BinaryStreamRouter; codec: StreamPortCodec } | undefined => { const sink = sinks?.get(port); - if (!sink) return undefined; + if (!sink || detachedPorts.has(port)) return undefined; let r = routers.get(port); if (!r) { r = { @@ -294,6 +299,40 @@ export class StreamProcessor if (deps.edgeBackpressure) await deps.edgeBackpressure(event.port); break; } + case "text-retract": { + if (!streamingStarted) { + streamingStarted = true; + this.task.status = TaskStatus.STREAMING; + this.task.emit("status", this.task.status); + } + if (accumulated) { + const text = accumulated.get(event.port); + if (text !== undefined) accumulated.set(event.port, retractText(text, event.length)); + } + // Bytes already teed to a cache sink cannot be taken back, and a + // ref minted from them would settle the port on text the producer + // withdrew. So the port's router is failed (its sink discards the + // partial write) and the port gets no router again this run: later + // deltas go to the accumulator alone, and the port settles inline. + // With no accumulator either, nothing here can say what the port + // holds now; the finish handler refuses that unless the producer + // states the port's value itself. + const routed = routers.get(event.port); + // Decided on what the fold removes, not on the raw length: a + // retraction that removes nothing (0.5, NaN) leaves the teed bytes + // exactly what the port holds. + if (routed && retractionLength(event.length) > 0) { + routed.router.fail( + new TaskError( + `Task ${this.task.type} retracted text already written to the cache for port "${event.port}"; discarding that write.` + ) + ); + routers.delete(event.port); + detachedPorts.add(event.port); + } + this.task.emit("stream_chunk", event as StreamEvent); + break; + } case "object-delta": { if (!streamingStarted) { streamingStarted = true; @@ -382,6 +421,17 @@ export class StreamProcessor // provider reported. Spread conditionally: an unreported usage must // leave no key at all rather than an explicit `usage: undefined`. const usageOnEvent = event.usage ? { usage: event.usage } : {}; + if (!accumulated) { + const stated = (event.data ?? {}) as Record; + for (const port of detachedPorts) { + if (!(port in stated)) { + throw new TaskError( + `Task ${this.task.type} retracted text on port "${port}" after it was ` + + `written to the cache, and this run keeps no other copy of the port.` + ); + } + } + } // `accumulated` witnesses all three accumulator maps — they are // created together under one `ctx.shouldAccumulate` ternary. const hasEnrichment = accumulated !== undefined || routers.size > 0; diff --git a/packages/task-graph/src/task/StreamTypes.ts b/packages/task-graph/src/task/StreamTypes.ts index e7bb91421f..7bad863c1f 100644 --- a/packages/task-graph/src/task/StreamTypes.ts +++ b/packages/task-graph/src/task/StreamTypes.ts @@ -30,6 +30,43 @@ export type StreamTextDelta = { textDelta: string; }; +/** + * Append mode only: takes back the last `length` characters (UTF-16 code + * units) streamed on `port` in this run, clamped at the start. Later + * `text-delta`s on that port append after what remains. Consumers fold it with + * {@link retractText}, which removes {@link retractionLength} characters: a + * fraction rounds down and anything not positive removes nothing. + * + * For a producer that streams text it may yet abandon, such as a model call + * that fails partway and is retried: the folded port then reads as if the + * abandoned text had never been sent. + */ +export type StreamTextRetract = { + type: "text-retract"; + port: string; + length: number; +}; + +/** + * How many characters a {@link StreamTextRetract} of `length` takes back: a + * fraction rounds down, and anything not positive (`NaN` included) takes back + * nothing. A producer counting what it may retract counts through this, so its + * count moves by exactly what {@link retractText} removes. + */ +export function retractionLength(length: number): number { + return length > 0 ? Math.floor(length) : 0; +} + +/** + * Removes the last {@link retractionLength} characters of `text`, clamped at + * the start: the fold every consumer of a {@link StreamTextRetract} applies. + */ +export function retractText(text: string, length: number): string { + const n = retractionLength(length); + if (n === 0) return text; + return text.slice(0, Math.max(0, text.length - n)); +} + /** * Object delta for structured/object streaming. * `port` identifies which output port this delta belongs to. @@ -442,6 +479,7 @@ export type StreamToolCall = */ export type StreamEvent> = | StreamTextDelta + | StreamTextRetract | StreamObjectDelta | StreamBinaryDelta | StreamSnapshot @@ -726,8 +764,8 @@ export const DEFAULT_STREAM_GATE_WATCHDOG_MS = 60_000; * passes the backpressure gate uncounted and a slow consumer accumulates one * per call without the gate ever closing. * - * The port filter cannot do this job: it only recognises the three delta types, - * so anything without a `port` passes it by construction. + * The port filter cannot do this job: it only recognises the per-port delta + * types, so anything without a `port` passes it by construction. */ export function isDataflowExcluded(event: StreamEvent): boolean { return event.type === "phase" || event.type === "tool-call"; diff --git a/packages/test/src/test/ai/AgentTaskReasoningAndImages.test.ts b/packages/test/src/test/ai/AgentTaskReasoningAndImages.test.ts index d1b7c93755..c64e5f64ad 100644 --- a/packages/test/src/test/ai/AgentTaskReasoningAndImages.test.ts +++ b/packages/test/src/test/ai/AgentTaskReasoningAndImages.test.ts @@ -294,6 +294,34 @@ describe("AgentTask reasoning and tool images", () => { return events; } + type WireEvent = { port: string; delta: string } | { port: string; retract: number }; + + /** Every delta and retraction, in the order the turn sent them. */ + function listenAll(task: AgentTask): WireEvent[] { + const events: WireEvent[] = []; + task.subscribe("stream_chunk", (event) => { + if (event.type === "text-delta") { + events.push({ port: event.port ?? "text", delta: event.textDelta }); + } else if (event.type === "text-retract") { + events.push({ port: event.port ?? "text", retract: event.length }); + } + }); + return events; + } + + /** What a host folding `port` reads, retractions applied. */ + function fold(events: readonly WireEvent[], port: string): string { + let text = ""; + for (const event of events) { + if (event.port !== port) continue; + text = + "delta" in event + ? text + event.delta + : text.slice(0, Math.max(0, text.length - event.retract)); + } + return text; + } + it("streams a round's reasoning on its own port, apart from the text", async () => { script([ { reasoning: "Check the file first.", calls: [{ id: "c1", name: "look", input: {} }] }, @@ -348,7 +376,7 @@ describe("AgentTask reasoning and tool images", () => { expect(output.reasoning).toBe("Only now."); }); - it("drops a failed attempt's reasoning and text, so the retry's are not joined to them", async () => { + it("takes back a failed attempt's reasoning and text, so the retry's are not joined to them", async () => { script([ { failWith: { @@ -360,7 +388,7 @@ describe("AgentTask reasoning and tool images", () => { { reasoning: "A whole thought.", text: "A whole answer." }, ]); const task = new AgentTask(); - const events = listen(task); + const events = listenAll(task); const output = await task.run( { model: MODEL, prompt: "?", tools: [], approval: "never" }, @@ -370,10 +398,53 @@ describe("AgentTask reasoning and tool images", () => { expect(output.steps[0]!.attempts).toBe(2); expect(output.reasoning).toBe("A whole thought."); expect(output.text).toBe("A whole answer."); + // Live, then retracted by exactly what the failed attempt sent on each port. expect(events).toEqual([ + { port: "reasoning", delta: "Half a thought" }, + { port: "text", delta: "Half an ans" }, + { port: "text", retract: "Half an ans".length }, + { port: "reasoning", retract: "Half a thought".length }, { port: "reasoning", delta: "A whole thought." }, { port: "text", delta: "A whole answer." }, ]); + expect(fold(events, "reasoning")).toBe(output.reasoning); + expect(fold(events, "text")).toBe(output.text); + }); + + it("keeps earlier rounds' reasoning when a later round's attempt fails, separator included", async () => { + script([ + { + reasoning: "Check the file first.", + text: "Looking. ", + calls: [{ id: "c1", name: "look", input: {} }], + }, + { + failWith: { + error: new RetryableJobError("overloaded", new Date(Date.now())), + reasoning: "Half", + text: "Par", + }, + }, + { reasoning: "It looks fine.", text: "All good." }, + ]); + const task = new AgentTask(); + const events = listenAll(task); + + const output = await task.run( + { model: MODEL, prompt: "Check it", tools: [LOOK], approval: "never" }, + { registry } + ); + + expect(output.reasoning).toBe("Check the file first.\n\nIt looks fine."); + expect(output.text).toBe("Looking. All good."); + // The failed attempt sent its own separator ahead of its reasoning, so it + // takes that back too; the retry sends one again. + expect(events.filter((event) => "retract" in event)).toEqual([ + { port: "text", retract: 3 }, + { port: "reasoning", retract: "\n\nHalf".length }, + ]); + expect(fold(events, "reasoning")).toBe(output.reasoning); + expect(fold(events, "text")).toBe(output.text); }); }); }); diff --git a/packages/test/src/test/ai/AgentTaskTurnControls.test.ts b/packages/test/src/test/ai/AgentTaskTurnControls.test.ts index 5d95320c98..cb71288676 100644 --- a/packages/test/src/test/ai/AgentTaskTurnControls.test.ts +++ b/packages/test/src/test/ai/AgentTaskTurnControls.test.ts @@ -45,10 +45,12 @@ interface Round { readonly usage?: Usage; /** Thrown instead of answering. */ readonly error?: Error; - /** Streamed before `error` is thrown, as a provider failing mid-stream does. */ + /** Streamed before `error` is thrown or the hang, as a provider failing mid-stream does. */ readonly partial?: string; /** Never answers: waits until the call is abandoned. */ readonly hang?: boolean; + /** Awaited after `text` is streamed and before the round answers. */ + readonly afterText?: () => Promise; } function usage(input: number, cached: number, output: number): Usage { @@ -69,10 +71,8 @@ function script(rounds: readonly Round[]): () => number { const runFn: AiProviderRunFn = async (_input, _model, signal, emit) => { const round = rounds[Math.min(called, rounds.length - 1)]!; called++; - if (round.error) { - if (round.partial) emit({ type: "text-delta", port: "text", textDelta: round.partial }); - throw round.error; - } + if (round.partial) emit({ type: "text-delta", port: "text", textDelta: round.partial }); + if (round.error) throw round.error; if (round.hang) { await new Promise((_resolve, reject) => { if (signal.aborted) reject(signal.reason); @@ -80,6 +80,7 @@ function script(rounds: readonly Round[]): () => number { }); } if (round.text) emit({ type: "text-delta", port: "text", textDelta: round.text }); + await round.afterText?.(); if (round.calls?.length) { emit({ type: "object-delta", port: "toolCalls", objectDelta: [...round.calls] }); } @@ -125,6 +126,33 @@ const ECHO: ToolDefinition = { execute: async (input) => String((input as { text?: string }).text ?? ""), }; +type TextEvent = { delta: string } | { retract: number }; + +/** The `text` port's deltas and retractions, in the order the turn sent them. */ +function listenText(task: AgentTask): TextEvent[] { + const events: TextEvent[] = []; + task.subscribe("stream_chunk", (event) => { + if (event.type === "text-delta" && event.port === "text") { + events.push({ delta: event.textDelta }); + } else if (event.type === "text-retract" && event.port === "text") { + events.push({ retract: event.length }); + } + }); + return events; +} + +/** What a host folding the wire reads, retractions applied. */ +function foldText(events: readonly TextEvent[]): string { + let text = ""; + for (const event of events) { + text = + "delta" in event + ? text + event.delta + : text.slice(0, Math.max(0, text.length - event.retract)); + } + return text; +} + function toolResults(messages: readonly { role: string; content: readonly unknown[] }[]) { return messages .filter((message) => message.role === "tool") @@ -708,16 +736,41 @@ describe("AgentTask turn controls", () => { expect(output.steps[0]!.attempts).toBe(2); }); - it("does not duplicate a failed attempt's streamed text on the retry", async () => { + it("streams a round's text before the round settles", async () => { + const task = new AgentTask(); + let streamed = ""; + const sawText = Promise.withResolvers(); + task.subscribe("stream_chunk", (event) => { + if (event.type !== "text-delta") return; + streamed += event.textDelta; + sawText.resolve(); + }); + let streamedBeforeAnswer: string | undefined; + script([ + { + text: "Hello", + afterText: async () => { + // Bounded, so a turn that held its text fails here rather than hanging. + await Promise.race([sawText.promise, new Promise((r) => setTimeout(r, 2_000))]); + streamedBeforeAnswer = streamed; + }, + }, + ]); + const output = await task.run( + { model: MODEL, prompt: "?", tools: [], approval: "never" }, + { registry } + ); + expect(streamedBeforeAnswer).toBe("Hello"); + expect(output.text).toBe("Hello"); + }); + + it("streams a failed attempt's text live, then takes it back before the retry", async () => { script([ { partial: "Hello ", error: new RetryableJobError("overloaded", new Date(Date.now())) }, { text: "Hello world", usage: usage(10, 0, 2) }, ]); const task = new AgentTask(); - let streamed = ""; - task.subscribe("stream_chunk", (event) => { - if (event.type === "text-delta") streamed += event.textDelta; - }); + const events = listenText(task); const output = await task.run( { model: MODEL, prompt: "?", tools: [], approval: "never" }, { registry } @@ -725,7 +778,69 @@ describe("AgentTask turn controls", () => { expect(output.steps[0]!.attempts).toBe(2); expect(output.text).toBe("Hello world"); expect(output.steps[0]!.text).toBe("Hello world"); - expect(streamed).toBe(output.text); + // The partial reached the wire as it was written, and exactly it was retracted. + expect(events).toEqual([{ delta: "Hello " }, { retract: 6 }, { delta: "Hello world" }]); + expect(foldText(events)).toBe(output.text); + }); + + it("retracts only the failed attempt's tail, leaving earlier rounds' text", async () => { + script([ + { text: "Checking. ", calls: [{ id: "c1", name: "echo", input: { text: "hi" } }] }, + { partial: "Oops", error: new RetryableJobError("overloaded", new Date(Date.now())) }, + { text: "Done." }, + ]); + const task = new AgentTask(); + const events = listenText(task); + const output = await task.run( + { model: MODEL, prompt: "?", tools: [ECHO], approval: "never" }, + { registry } + ); + expect(output.steps.map((step) => step.attempts)).toEqual([1, 2]); + expect(output.text).toBe("Checking. Done."); + expect(events).toEqual([ + { delta: "Checking. " }, + { delta: "Oops" }, + { retract: 4 }, + { delta: "Done." }, + ]); + expect(foldText(events)).toBe(output.text); + }); + + it("takes back a timed-out attempt's text before the retry", async () => { + script([{ partial: "Hel", hang: true }, { text: "Hello" }]); + const task = new AgentTask(); + const events = listenText(task); + const output = await task.run( + { model: MODEL, prompt: "?", tools: [], roundTimeoutMs: 50, approval: "never" }, + { registry } + ); + expect(output.text).toBe("Hello"); + expect(events).toEqual([{ delta: "Hel" }, { retract: 3 }, { delta: "Hello" }]); + }); + + it("takes back a failed attempt's text when the turn ends on it", async () => { + // The retry wait would pass the deadline, so the turn ends "budget" instead. + script([ + { partial: "Half", error: new RetryableJobError("x", new Date(Date.now() + 30_000)) }, + ]); + const budget = new AgentTask(); + const budgetEvents = listenText(budget); + const output = await budget.run( + { model: MODEL, prompt: "?", tools: [], maxDurationMs: 5_000, approval: "never" }, + { registry } + ); + expect(output.stopReason).toBe("budget"); + expect(output.text).toBe(""); + expect(budgetEvents).toEqual([{ delta: "Half" }, { retract: 4 }]); + + getAiProviderRegistry().unregisterProvider(PROVIDER); + script([{ partial: "Half", error: new PermanentJobError("401 bad key") }]); + const failed = new AgentTask(); + const failedEvents = listenText(failed); + await expect( + failed.run({ model: MODEL, prompt: "?", tools: [], approval: "never" }, { registry }) + ).rejects.toThrow(/401/); + expect(failedEvents).toEqual([{ delta: "Half" }, { retract: 4 }]); }); it("fails at once on a permanent failure, or when retries are off", async () => {