feat: extend agent-driven map workbench

features/workbench/hooks/agent-session-message-state.ts

@@ -1,5 +1,5 @@
 import type { UIMessage } from "ai";
-import type { Dispatch, MutableRefObject, SetStateAction } from "react";
+import type { Dispatch, SetStateAction } from "react";
 import type {
   AgentChatMessage,
   AgentPermissionReply,
@@ -48,13 +48,11 @@ export type QuestionOverride = {
 };
 
 type StreamRenderStateSetter = Dispatch<SetStateAction<AgentStreamRenderState>>;
-type StreamRenderChunkIdRef = MutableRefObject<number>;
 
-export function appendStreamRenderToken(
+export function markStreamRenderPending(
   data: Record<string, unknown>,
   messages: AgentUiMessage[],
-  setStreamRenderState: StreamRenderStateSetter,
-  chunkIdRef: StreamRenderChunkIdRef
+  setStreamRenderState: StreamRenderStateSetter
 ) {
   const text = getDataString(data, "content");
   if (!text) {
@@ -66,17 +64,14 @@ export function appendStreamRenderToken(
     return;
   }
 
-  const chunkId = chunkIdRef.current;
-  chunkIdRef.current += 1;
-
   setStreamRenderState((current) => {
-    const previous = current[messageId] ?? { chunks: [], done: false };
+    if (current[messageId]?.done === false) {
+      return current;
+    }
+
     return {
       ...current,
-      [messageId]: {
-        chunks: [...previous.chunks, { id: chunkId, text }],
-        done: false
-      }
+      [messageId]: { done: false }
     };
   });
 }
@@ -84,10 +79,7 @@ export function appendStreamRenderToken(
 export function createCompletedStreamRenderState(messages: AgentUiMessage[]): AgentStreamRenderState {
   return messages.reduce<AgentStreamRenderState>((next, message) => {
     if (message.role === "assistant") {
-      next[message.id] = {
-        chunks: [],
-        done: true
-      };
+      next[message.id] = { done: true };
     }
     return next;
   }, {});
@@ -107,10 +99,7 @@ export function markLastAssistantStreamDone(
 
     return {
       ...current,
-      [messageId]: {
-        chunks: current[messageId]?.chunks ?? [],
-        done: true
-      }
+      [messageId]: { done: true }
     };
   });
 }