fix(ai): onStepFinish type for AgentUIStream methods

[rovo89]

Background

See #12841. Recently, onStepFinish callbacks where added in a few places, but the type for AgentUIStreams focussed on the "agent" part instead of the "UI stream" part.

Summary

This changes the type of the onStepFinish callback from ToolLoopAgentOnStepFinishCallback<TOOLS> to
UIMessageStreamOnStepFinishCallback<UI_MESSAGE>
for methods which create a UIMessageChunk stream for an Agent, giving the implementation access to the generatedUIMessage so far instead of just the raw parts.

This partly reverts #11980, which explicitly added the onStepFinish callback to the method signatures. Instead, onStepFinish is now added to the UIMessageStreamOptions type in the same way as onFinish already was, and is passed through via options.

Affected methods are:

• createAgentUIStream()
• createAgentUIStreamResponse()
• pipeAgentUIStreamToResponse()

For even more consistency, the onStepFinish callback is also added to:

• DefaultStreamTextResult.toUIMessageStream() (used by the methods above)
• DefaultStreamTextResult.toUIMessageStreamResponse()
• DefaultStreamTextResult.pipeUIMessageStreamToResponse()

Callers which want to receive the ToolLoopAgentOnStepFinishCallback can easily do that with the onStepFinish callback to the constructor of the ToolLoopAgent that they're passing to the methods above.

Fixes #12841.

Manual Verification

Used createAgentUIStream() in a project, added a onStepFinish callback, logged the responseMessage and verified that it looks as expected.

Checklist

• Tests have been added / updated (for bug fixes / features)
• Documentation has been added / updated (for bug fixes / features)
• A patch changeset for relevant packages has been added (for bug fixes / features - run pnpm changeset in the project root)
• I have reviewed this pull request (self-review)

Related Issues

Issues #12841 and #12383 describe the motivation to have the callback.
PRs #11980 and #12448 added implementations recently.
PR #12654 also touched callbacks afterwards.

[rovo89]

@lgrammel @aayush-kapoor @sullyo Could you please have a look? I think this makes the onStepFinish callback more consistent. Thanks!

[rovo89]

I don't really have an idea why onFinish is omitted here, but I found it while searching for references to UIMessageStreamOptions. Since onStepFinish is now added to that type and is similar to onFinish, I added it here. If that's not correct, I'll remove it, but this way we can at least talk about it.

[rovo89]

@lgrammel @aayush-kapoor @sullyo Friendly ping. I think this PR would make the types for onStepFinish much more consistent, support it in more places, remove special handling and documentation is fully adjusted. Lots of wins. 😉 And with v7 betas coming in, even the tiny (easy to resolve) backwards-incompatibility isn't a problem. Would really appreciate getting this merged.

[lgrammel]

@rovo89 can you update the relevant examples / add examples to show how things would look after the change? when designing this feature, I discussed with @nicoalbanese where one would most likely want to listen for the agent finish, and i want to make sure this does not break the intended agent definition / agent execution separation.

[rovo89]

@lgrammel Sure, I'll have a look. You're referring to https://ai-sdk.dev/docs/reference/ai-sdk-core/tool-loop-agent, https://ai-sdk.dev/docs/reference/ai-sdk-core/create-agent-ui-stream etc.?

Currently I don't see any examples for neither onFinish nor onStepFinish there, but I can add one showing the differences between adding them to either the ToolLoopAgent constructor or createAgentUIStreamResponse (e.g. for useChat from AI Elements).

But maybe these tables is even more helpful to understand my intention. This is the current situation:

Callback	ToolLoopAgent constructor	createUIMessageStream	createAgentUIStream etc.
onFinish	OnFinishEvent (ModelMessage-based)	UIMessageStreamOnFinishCallback (UiMessage-based)	UIMessageStreamOnFinishCallback (UiMessage-based)
onStepFinish	OnStepFinishEvent (ModelMessage-based)	UIMessageStreamOnStepFinishCallback (UiMessage-based)	OnStepFinishEvent (ModelMessage-based) ❌

And I think it should be:

Callback	ToolLoopAgent constructor	createUIMessageStream	createAgentUIStream etc.
onFinish	OnFinishEvent (ModelMessage-based)	UIMessageStreamOnFinishCallback (UiMessage-based)	UIMessageStreamOnFinishCallback (UiMessage-based)
onStepFinish	OnStepFinishEvent (ModelMessage-based)	UIMessageStreamOnStepFinishCallback (UiMessage-based)	UIMessageStreamOnStepFinishCallback (UiMessage-based) ✅

That gives devs the capability to hook into either the low-level part with ModelMessage or the higher-level with UIMessage. For onFinish, that is already implemented, but onStepFinish isn't consistent with that.