Add UI Message Support to DurableAgent onFinish callback #739
Description
The onFinish callback in DurableAgent returns ModelMessage[], but it's recommended to use UIMessage[] to persist messages to database for display on the frontend in a format that is consistent with useChat.
Currently messages with the UIMessage seem to only be available on the client via useChat, I would prefer to save messages on the server without needing to first send them to the client.
Workflow steps might take some time to complete and it would be nice to allow the user to navigate away, persist their messages on the server and have them ready for when they return back to the chat.
Current Implementation: (Based on workflow-examples repo here)
import { convertToModelMessages, UIMessageChunk, type UIMessage } from "ai";
import { DurableAgent } from "@workflow/ai/agent";
import { getWritable } from "workflow";
export async function chat(messages: UIMessage[]) {
"use workflow";
const writable = getWritable<UIMessageChunk>();
const agent = new DurableAgent({
model: "bedrock/claude-4-sonnet-20250514-v1",
system: FLIGHT_ASSISTANT_PROMPT,
tools: flightBookingTools,
});
await agent.stream({
messages: convertToModelMessages(messages),
writable,
onFinish: async ({ messages, steps }) => {
// messages is ModelMessage[] - can't save to DB directly
// Want: saveMessages(messages) with UIMessage[] format
}
});
}We need either:
1. A convertToUIMessages utility
The AI SDK has convertToModelMessages but no reverse conversion, making persistence difficult. This would be better fixed in the ai package.
Related: #vercel/ai#7180.
2. Return uiMessages in onFinish:
Perhaps the durable agent onFinish messages could directly return the UIMessage type, or an additional uiMessages could be made available, eg.
onFinish: async ({ messages, uiMessages, steps }) => {
// Save UI Messages to DB in a "use step" function
saveMessages(uiMessages)
}I was able to persist the current ModelMessage messages returned to the database and it sort of works, but ran into issues like missing tool result data, misordered messages in the db etc. It took some manual extra code I had to write to convert, felt error prone and overall wasn't a great experience. I think it could be a lot easier with direct access to the UI messages, or perhaps someone has other suggestions on how to approach this.