fix(agent): stabilize stream token output

src/routes/chat.ts

@@ -764,6 +764,15 @@ export const buildChatRouter = (
           textPart.state = "streaming";
         }
         writeUiChunk({ type: "text-delta", id: textPartId, delta });
+        writeUiChunk({
+          type: "data-stream_token",
+          data: {
+            session_id: clientSessionId,
+            message_id: assistantUiMessageId,
+            content: delta,
+          },
+          transient: true,
+        } as UIMessageChunk);
         syncLatestUiMessages();
       };
       const writeDataPart = (event: string, data: Record<string, unknown>) => {
@@ -879,6 +888,16 @@ export const buildChatRouter = (
       res.on("close", handleClientClose);
 
       const publish = (event: string, data: Record<string, unknown>) => {
+        const subscriberData =
+          event === "token"
+            ? {
+                ...data,
+                message_id:
+                  typeof data.message_id === "string"
+                    ? data.message_id
+                    : assistantUiMessageId,
+              }
+            : data;
         if (event === "token") {
           activeRun.messages = updateLastAssistantMessage(activeRun.messages, (message) => ({
             ...message,
@@ -1011,7 +1030,7 @@ export const buildChatRouter = (
         }
 
         for (const subscriber of activeRun.subscribers) {
-          subscriber.write(event, data);
+          subscriber.write(event, subscriberData);
         }
         void queueSessionUiStatePersist().catch((error) => {
           logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");