Need streaming callbacks in chat for persistence

[nauvalazhar]

We already have callbacks like onFinish and onChunk in the useChat hook, but it would be great to have the same callbacks available in the chat function as well.

Currently, I’m using StreamProcessor to capture the final message so I can persist it to the database:

export const collectChatStream = createServerOnlyFn(async function* (
  stream: AsyncIterable<StreamChunk>,
  options?: CollectChatStreamOptions
) {
  const processor = new StreamProcessor({
    events: {
      onMessagesChange: (message) => {
        console.log(message);
      },
      onStreamEnd: (message: UIMessage) => {
        options?.onFinish?.(message);
        console.log("onstreamend");
      },
    },
  });

  processor.startAssistantMessage();

  for await (const event of stream) {
    processor.processChunk(event);

    yield event;
  }

  processor.finalizeStream();
});

The problem is that onStreamEnd is not called when the chat is aborted by the client using stop() and I'm not be able to capture the final message.

Thank you.