Skip to content

fix(ai): change createAgentUIStream onStepFinish to UIMessageStreamOnStepFinishCallback#12978

Open
sleitor wants to merge 1 commit intovercel:mainfrom
sleitor:fix-12841
Open

fix(ai): change createAgentUIStream onStepFinish to UIMessageStreamOnStepFinishCallback#12978
sleitor wants to merge 1 commit intovercel:mainfrom
sleitor:fix-12841

Conversation

@sleitor
Copy link
Copy Markdown
Contributor

@sleitor sleitor commented Mar 1, 2026

Summary

Fixes #12841

The onStepFinish callback in createAgentUIStream, createAgentUIStreamResponse, and pipeAgentUIStreamToResponse was typed as ToolLoopAgentOnStepFinishCallback (LLM-level step result with token usage, finish reason, etc.), which is inconsistent with the UIMessageStreamOnStepFinishCallback (UI-level messages, isContinuation, responseMessage) used in createUIMessageStream.

This makes it impossible to use createAgentUIStream for the primary use-case of persisting intermediate messages during multi-step agent runs (as described in #12383).

Changes

  • create-agent-ui-stream.ts: Changed onStepFinish type from ToolLoopAgentOnStepFinishCallback to UIMessageStreamOnStepFinishCallback; moved it from agent.stream() to result.toUIMessageStream()
  • create-agent-ui-stream-response.ts: Same type change
  • pipe-agent-ui-stream-to-response.ts: Same type change
  • stream-text-result.ts: Added onStepFinish to UIMessageStreamOptions so it threads through toUIMessageStream(), pipeUIMessageStreamToResponse(), and toUIMessageStreamResponse()
  • stream-text.ts: Destructured and passed onStepFinish to handleUIMessageStreamFinish() in toUIMessageStream()

Before / After

Before: onStepFinish in createAgentUIStream received OnStepFinishEvent with raw LLM details (usage tokens, finish reason, model response, etc.)

After: onStepFinish receives { messages, isContinuation, responseMessage } — the same shape as createUIMessageStream's onStepFinish, making it easy to persist intermediate UI messages:

await createAgentUIStream({
  agent,
  uiMessages,
  onStepFinish: async ({ messages, responseMessage }) => {
    await db.saveMessages(messages); // persist intermediate state
  },
});

…StepFinishCallback

The onStepFinish callback in createAgentUIStream, createAgentUIStreamResponse,
and pipeAgentUIStreamToResponse was typed as ToolLoopAgentOnStepFinishCallback
(LLM-level step result), which differs from the UIMessageStreamOnStepFinishCallback
(UI-level messages, isContinuation, responseMessage) used in createUIMessageStream.

This commit:
- Changes the onStepFinish type in all three agent UI stream functions to
  UIMessageStreamOnStepFinishCallback for API consistency
- Moves onStepFinish from agent.stream() to result.toUIMessageStream() so it
  receives the processed UI messages instead of raw LLM step results
- Adds onStepFinish to UIMessageStreamOptions so it can be passed through
  toUIMessageStream(), pipeUIMessageStreamToResponse(), and toUIMessageStreamResponse()
- Adds onStepFinish to handleUIMessageStreamFinish() call in toUIMessageStream()

Fixes vercel#12841
@tigent tigent Bot added ai/ui anything UI related bug Something isn't working as documented labels Mar 1, 2026
@rovo89
Copy link
Copy Markdown

rovo89 commented Mar 4, 2026

@sleitor I believe my approach is cleaner: #12843. Did you miss anything there that made you come up with your own PR?

@sleitor
Copy link
Copy Markdown
Contributor Author

sleitor commented Apr 7, 2026

👋 Gentle ping — just checking if this is still on the radar for review. Happy to address any feedback or rebase if needed!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ai/ui anything UI related bug Something isn't working as documented

Projects

None yet

Development

Successfully merging this pull request may close these issues.

onStepFinish callback in createAgentUIStream() doesn't contain responseMessage

2 participants