Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/quiet-runs-finish-once.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'@tanstack/ai': patch
---

Emit one lifecycle pair for the full agent loop instead of one pair per provider iteration.
29 changes: 13 additions & 16 deletions docs/resumable-streams/custom-adapter.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,23 +39,21 @@ Get these wrong and resume breaks in subtle ways:
containing `NUL`/CR/LF, one with leading or trailing whitespace, or a
duplicate.
- **`read` replays strictly *after* the offset**, oldest first, and ends when the
log is **closed** never when it sees a terminal chunk. This is the rule most
log is **closed**, never when it sees a terminal chunk. This is the rule most
likely to be "simplified" back, so here is the reason, quoted from the
invariant `memoryStream`'s own `read` loop carries in
`packages/ai/src/stream-durability.ts` (a test pins it):

> A terminal chunk (`RUN_FINISHED` / `RUN_ERROR`) does NOT end the read: an
> agent-loop run emits one per iteration (`finishReason` `"tool_calls"` then
> `"stop"`), so stopping on the first would truncate a tool-calling run at its
> first tool call. The producer signals true completion by calling `close()`
> (it does so on every exit — see `StreamDurability.close`), which sets
> `log.complete`. Read tails until then, or until the caller aborts.

So an adapter that returns at the first terminal chunk truncates **every**
resumed tool-calling run — the resumed client sees the first tool call and
then a clean end, which it reads as "the run is over". `close()` is your only
end-of-log signal, and core awaits it on every producer exit (completion,
cancellation, and failure), so tailing until then always terminates.
> A terminal chunk (`RUN_FINISHED` / `RUN_ERROR`) does NOT end the read.
> `StreamDurability` accepts lower-level streams that append more chunks after
> a terminal, so only the producer's `close()` marks the log complete. Read
> tails until then, or until the caller aborts.

An adapter that returns at the first terminal chunk can truncate a valid
persisted stream. `close()` is your only end-of-log signal. Core awaits it
when a producer completes, is cancelled, or fails. During a detached handoff,
core leaves the log open so the successor producer can continue it and close
it later.
- **`read` must never end the response empty while the run is still producing.**
Park (wait for the next append) instead. A clean end with no new data tells
the client the run is over; if it isn't, the client fails with
Expand Down Expand Up @@ -131,9 +129,8 @@ export function customDurability(
const entries = await log.readAfter(cursor)
for (const entry of entries) {
cursor = entry.cursor
// Yield terminal chunks like any other. An agent-loop run emits a
// RUN_FINISHED per iteration, so returning on one would truncate a
// resumed tool-calling run at its first tool call.
// Yield terminal chunks like any other. A lower-level producer may
// append more chunks before it closes.
yield { offset: entry.cursor, chunk: entry.chunk }
}
// The ONLY end-of-log condition: the producer called `close()`.
Expand Down
11 changes: 4 additions & 7 deletions packages/ai-client/src/connection-adapters.ts
Original file line number Diff line number Diff line change
Expand Up @@ -749,13 +749,10 @@ async function* resumableStream(
sawTerminal = true
}
yield chunk
// Do NOT stop on a terminal mid-source: an agent loop emits one
// RUN_STARTED/RUN_FINISHED pair PER turn, so a tool-calling run carries
// several RUN_FINISHED events before the run is truly done. Returning on
// the first one would drop every subsequent turn (the tool result and
// the final answer). Instead, drain the event source to its natural end
// — the server closes the response only when the run is actually
// complete — and use `sawTerminal` below to decide done-vs-reconnect.
// Do NOT stop on a terminal mid-source. A lower-level stream may contain
// multiple terminal events before its source ends, so returning on the
// first would drop later chunks. Drain the event source to its natural
// end, then use `sawTerminal` below to decide done-vs-reconnect.
}
} catch (error) {
if (abortSignal?.aborted) return
Expand Down
10 changes: 3 additions & 7 deletions packages/ai-client/tests/connection-adapters.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -211,13 +211,9 @@ describe('connection-adapters', () => {
})
})

it('does not truncate an agentic run at the first RUN_FINISHED (multiple terminals per run)', async () => {
// An agent loop emits one RUN_STARTED/RUN_FINISHED pair PER turn, so a
// tool-calling run streams several RUN_FINISHED events in a single
// non-durable (untagged) response: turn 1 ends with RUN_FINISHED, then the
// tool result + turn 2 (the final answer) follow. Regression guard: the
// client must forward EVERY event and stop only when the source closes,
// never returning early on the first terminal.
it('does not truncate a lower-level stream at its first terminal', async () => {
// The transport must forward every event and stop only when the source
// closes, even when a lower-level producer emits multiple terminals.
const body =
'data: {"type":"RUN_STARTED","timestamp":1}\n\n' +
'data: {"type":"TOOL_CALL_END","timestamp":2}\n\n' +
Expand Down
163 changes: 85 additions & 78 deletions packages/ai/src/activities/chat/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -779,8 +779,11 @@ class TextEngine<
private readonly loopStrategy: AgentLoopStrategy
private toolCallManager: ToolCallManager<ReadonlyArray<AnyTool>, TContext>
private readonly lazyToolManager: LazyToolManager
/** A public interruption terminal must always have this run's start event. */
private hasPublicRunStarted = false
/** The first public start owns the outer lifecycle identity. */
private publicRunIdentity: Pick<
RunFinishedEvent,
'runId' | 'threadId'
> | null = null
private readonly initialMessageCount: number
private readonly requestId: string
private readonly streamId: string
Expand Down Expand Up @@ -810,9 +813,8 @@ class TextEngine<
private eventToolNames?: Array<string>
private finishedEvent: RunFinishedEvent | null = null
private readonly streamedToolErrorResults = new Map<string, ToolResult>()
private deferredToolCallRunFinishedChunks: Array<StreamChunk> = []
/** The model terminal is held until afterModel can choose an interrupt. */
private deferredModelRunFinishedChunks: Array<StreamChunk> = []
/** The provider terminal is held until the agent loop reaches its public boundary. */
private deferredRunFinishedChunk: RunFinishedEvent | null = null

private earlyTermination = false
private toolPhase: ToolPhaseResult = 'continue'
Expand Down Expand Up @@ -1226,14 +1228,6 @@ class TextEngine<
this.setToolPhase('wait')
return
}
if (this.shouldExecuteToolPhase()) {
this.deferredToolCallRunFinishedChunks.push(
...this.deferredModelRunFinishedChunks,
)
this.deferredModelRunFinishedChunks = []
} else {
yield* this.flushDeferredModelRunFinishedChunks()
}
} else {
yield* this.processToolCalls()
}
Expand Down Expand Up @@ -1270,6 +1264,24 @@ class TextEngine<
}
}

if (
this.deferredRunFinishedChunk &&
this.toolPhase !== 'wait' &&
!this.isCancelled() &&
!this.finalizationError &&
!this.earlyTermination
) {
const providerFinishedChunk = this.deferredRunFinishedChunk
this.deferredRunFinishedChunk = null
yield* this.emitSyntheticRunStarted(providerFinishedChunk)
const finishedChunk = this.withPublicRunIdentity(providerFinishedChunk)
this.logger.output(`type=${finishedChunk.type}`, {
chunk: finishedChunk,
})
yield finishedChunk
this.middlewareCtx.chunkIndex++
}

// Call terminal hook (skip when waiting for client — stream is paused, not finished).
// Priority: finalizationError → onError; otherwise normal onFinish.
// Skip on cancellation — the finally block routes aborts to onAbort.
Expand Down Expand Up @@ -1601,13 +1613,13 @@ class TextEngine<
outboundChunk,
)
// When a streaming structured-output finalization step will run after
// the agent loop, suppress the agent-loop's RUN_STARTED/RUN_FINISHED
// here — the finalization step emits the single outer lifecycle pair
// that reaches the consumer.
// the agent loop, suppress the agent-loop's RUN_FINISHED here. Its first
// RUN_STARTED owns the outer lifecycle, and the finalization step emits
// the terminal that closes it.
//
// Native combined mode does NOT issue a second adapter stream — the
// agent loop's lifecycle IS the outer pair the consumer sees.
const suppressAgentLifecycle =
const suppressAgentRunFinished =
!!this.finalStructuredOutput &&
this.finalStructuredOutput.yieldChunks &&
this.finalStructuredOutput.nativeCombined !== true
Expand All @@ -1617,25 +1629,25 @@ class TextEngine<
)) {
restorePublicUsage(spec)
if (
suppressAgentLifecycle &&
(spec.type === EventType.RUN_STARTED ||
spec.type === EventType.RUN_FINISHED)
suppressAgentRunFinished &&
spec.type === EventType.RUN_FINISHED
) {
continue
}
if (spec.type === EventType.RUN_FINISHED) {
this.deferredModelRunFinishedChunks.push(spec)
continue
}
if (this.shouldDeferToolCallRunFinished(spec)) {
this.deferredToolCallRunFinishedChunks.push(spec)
this.deferredRunFinishedChunk = spec
continue
}
if (spec.type === EventType.RUN_STARTED) {
this.hasPublicRunStarted = true
if (this.publicRunIdentity) continue
this.publicRunIdentity = {
runId: spec.runId,
threadId: spec.threadId,
}
}
this.logger.output(`type=${spec.type}`, { chunk: spec })
yield spec
const publicChunk = this.withPublicRunIdentity(spec)
this.logger.output(`type=${publicChunk.type}`, { chunk: publicChunk })
yield publicChunk
this.middlewareCtx.chunkIndex++
}
}
Expand Down Expand Up @@ -2058,8 +2070,6 @@ class TextEngine<
executionResult.needsApproval.length > 0 ||
executionResult.needsClientExecution.length > 0
) {
this.discardDeferredToolCallRunFinishedChunks()

if (allResults.length > 0) {
for (const chunk of this.buildToolResultChunks(
allResults,
Expand Down Expand Up @@ -2140,7 +2150,6 @@ class TextEngine<
]

if (executableToolCalls.length === 0) {
yield* this.flushDeferredToolCallRunFinishedChunks()
// All tool calls already have error results — emit them, then continue
// the loop (strategy / onShouldContinue may stop).
if (deferredErrorResults.length > 0) {
Expand Down Expand Up @@ -2278,8 +2287,6 @@ class TextEngine<
return
}

yield* this.flushDeferredToolCallRunFinishedChunks()

const toolResultChunks = afterToolBoundaryChunks

for (const chunk of toolResultChunks) {
Expand All @@ -2302,37 +2309,10 @@ class TextEngine<
this.setToolPhase('continue')
}

private shouldDeferToolCallRunFinished(chunk: StreamChunk): boolean {
return (
chunk.type === EventType.RUN_FINISHED &&
this.lastFinishReason === 'tool_calls' &&
this.tools.length > 0 &&
this.toolCallManager.hasToolCalls()
)
}

private *flushDeferredToolCallRunFinishedChunks(): Generator<StreamChunk> {
for (const chunk of this.deferredToolCallRunFinishedChunks) {
this.logger.output(`type=${chunk.type}`, { chunk })
yield chunk
this.middlewareCtx.chunkIndex++
}
this.deferredToolCallRunFinishedChunks = []
}

private *flushDeferredModelRunFinishedChunks(): Generator<StreamChunk> {
for (const chunk of this.deferredModelRunFinishedChunks) {
this.logger.output(`type=${chunk.type}`, { chunk })
yield chunk
this.middlewareCtx.chunkIndex++
}
this.deferredModelRunFinishedChunks = []
}

private async *emitSyntheticRunStarted(
finishEvent: RunFinishedEvent,
): AsyncGenerator<StreamChunk, void, void> {
if (this.hasPublicRunStarted) return
if (this.publicRunIdentity) return
yield* this.pipeThroughMiddleware({
type: EventType.RUN_STARTED,
runId: finishEvent.runId,
Expand All @@ -2358,10 +2338,6 @@ class TextEngine<
})
}

private discardDeferredToolCallRunFinishedChunks(): void {
this.deferredToolCallRunFinishedChunks = []
}

private shouldExecuteToolPhase(): boolean {
return (
this.lastFinishReason === 'tool_calls' &&
Expand Down Expand Up @@ -3678,12 +3654,16 @@ class TextEngine<
)

// 7c. Decide consumer visibility — only yieldChunks=true callers get them.
// We do NOT strip the finalization stream's RUN_STARTED/RUN_FINISHED:
// they are the single outer lifecycle pair the consumer sees (the
// agent-loop's pair was suppressed in streamModelResponse when
// finalStructuredOutput.yieldChunks is true).
// Keep the finalization's full lifecycle when the no-tools path skipped
// the agent loop. Otherwise, the agent start already owns the public
// lifecycle, so only the finalization terminal remains public.
if (this.finalStructuredOutput.yieldChunks) {
for (const spec of this.emitPublicChunks(outputChunks)) {
const publicOutputChunks = this.publicRunIdentity
? outputChunks.filter(
(output) => output.type !== EventType.RUN_STARTED,
)
: outputChunks
for (const spec of this.emitPublicChunks(publicOutputChunks)) {
if (spec.type === EventType.RUN_ERROR) {
runErrorYielded = true
}
Expand Down Expand Up @@ -3899,11 +3879,7 @@ class TextEngine<

// On success, emit the synthetic `structured-output.complete` carrying
// the parsed object + raw text. Pin the messageId so the client-side
// handler can target the right UIMessage even when the agent loop's
// terminal RUN_FINISHED has already cleared `activeMessageIds` (the
// complete event yields AFTER the loop ends, by which point
// `getActiveAssistantMessageId()` returns null and would otherwise drop
// the event silently).
// handler targets the same UIMessage as `structured-output.start`.
if (
this.structuredOutputResult &&
!this.finalizationError &&
Expand Down Expand Up @@ -4398,15 +4374,46 @@ class TextEngine<
for (const output of outputs) {
for (const spec of normalizeStreamChunk(output as AdapterYieldChunk)) {
restorePublicUsage(spec)
if (spec.type === EventType.RUN_STARTED) {
this.hasPublicRunStarted = true
if (
spec.type === EventType.RUN_STARTED &&
this.publicRunIdentity === null
) {
this.publicRunIdentity = {
runId: spec.runId,
threadId: spec.threadId,
}
}
yield spec
yield this.withPublicRunIdentity(spec)
this.middlewareCtx.chunkIndex++
}
}
}

private withPublicRunIdentity(chunk: StreamChunk): StreamChunk {
if (
this.publicRunIdentity === null ||
(chunk.type !== EventType.RUN_FINISHED &&
chunk.type !== EventType.RUN_ERROR)
) {
return chunk
}
if (chunk.type === EventType.RUN_ERROR) {
return withTanstackMetadata(
{ ...chunk, ...this.publicRunIdentity },
this.publicRunIdentity,
)
}
const publicChunk = { ...chunk, ...this.publicRunIdentity }
const metadataIdentity = tanstackMetadata(chunk)
if (
metadataIdentity?.runId !== undefined ||
metadataIdentity?.threadId !== undefined
) {
return withTanstackMetadata(publicChunk, this.publicRunIdentity)
}
return publicChunk
}

/**
* Pipe a single internal chunk through middleware, then spec-normalize
* before the public `for await` stream.
Expand Down
10 changes: 4 additions & 6 deletions packages/ai/src/stream-durability.ts
Original file line number Diff line number Diff line change
Expand Up @@ -501,12 +501,10 @@ export function memoryStream(
yield { offset: entry.offset, chunk: entry.chunk }
}
}
// A terminal chunk (RUN_FINISHED / RUN_ERROR) does NOT end the read: an
// agent-loop run emits one per iteration (finishReason "tool_calls" then
// "stop"), so stopping on the first would truncate a tool-calling run at
// its first tool call. The producer signals true completion by calling
// `close()` (it does so on every exit — see StreamDurability.close), which
// sets `log.complete`. Read tails until then, or until the caller aborts.
// A terminal chunk (RUN_FINISHED / RUN_ERROR) does NOT end the read.
// StreamDurability accepts lower-level streams that append more chunks
// after a terminal, so only the producer's `close()` marks the log
// complete. Read tails until then, or until the caller aborts.
if (log.complete || signal?.aborted) return

// Bound only the wait for the very first chunk: once a run has produced
Expand Down
Loading