Add updateMessage handler on ThreadStorage

docs/content/docs/agent/guides/migrating.mdx

@@ -116,7 +116,7 @@ The behavior change to internalize: **you no longer create or own an `AbortContr
 
 ## Thread props → `storage`
 
-`threadApiUrl` plus the per-operation callbacks become a single `ChatStorage` whose `thread` member holds five methods.
+`threadApiUrl` plus the per-operation callbacks become a single `ChatStorage` whose `thread` member implements `ThreadStorage`.
 
 ### The REST case → `restStorage`
 

docs/content/docs/agent/reference/adapters-and-formats.mdx

@@ -156,7 +156,7 @@ Omit `storage` entirely and `AgentInterface` uses an internal in-memory store 
 
 ### `ThreadStorage`
 
-The five methods that back thread management. Implement these and the default sidebar's thread list, "New chat" button, thread switching, and deletion all operate against your backend.
+Implements methods that back thread and message management. Implement the method and the default sidebar's thread list, "New chat" button, thread switching, deletion and message interactions all operate against your backend.
 
 ```ts
 interface ThreadStorage {
@@ -165,6 +165,7 @@ interface ThreadStorage {
   getMessages(threadId: string): Promise<Message[]>;
   updateThread(thread: Thread): Promise<Thread>;
   deleteThread(id: string): Promise<void>;
+  updateMessage?(threadId: string, message: Message): Promise<void>;
 }
 ```
 
@@ -175,6 +176,7 @@ interface ThreadStorage {
 | `getMessages(threadId)` | `Message[]` | The user opens a thread. |
 | `updateThread(thread)` | the updated `Thread` | A thread changes (e.g. a rename). |
 | `deleteThread(id)` | `void` | The user deletes a thread. |
+| `updateMessage?(threadId, message)` | `void` | *Optional.* The user edits form state in a rendered message; called fire-and-forget. |
 
 Implement these directly when your storage doesn't fit the REST shape `restStorage` expects — a different route layout, GraphQL, or a client-side store like IndexedDB:
 

docs/content/docs/agent/reference/self-hosting.mdx

@@ -267,7 +267,7 @@ export async function DELETE(_req: NextRequest, { params }: { params: { threadId
 
 ### Custom `ChatStorage` instead
 
-If the REST endpoint shape doesn't fit your backend — GraphQL, a client-side store like IndexedDB, a SaaS SDK, or just a different URL layout — implement `ChatStorage` directly. It's an object with a `thread` member satisfying `ThreadStorage` (five methods) plus an optional `artifact` member. `restStorage` is itself just a `ChatStorage` built this way for the common REST case.
+If the REST endpoint shape doesn't fit your backend — GraphQL, a client-side store like IndexedDB, a SaaS SDK, or just a different URL layout — implement `ChatStorage` directly. It's an object with a `thread` member satisfying `ThreadStorage` plus an optional `artifact` member. `restStorage` is itself just a `ChatStorage` built this way for the common REST case.
 
 ```ts
 import type { ChatStorage } from "@openuidev/react-ui";
@@ -294,11 +294,15 @@ export const storage: ChatStorage = {
     async deleteThread(id) {
       await gql(DELETE_THREAD, { id });
     },
+    // Optional — persists in-message edits
+    async updateMessage(threadId, message) {
+      await gql(UPDATE_MESSAGE, { threadId, message });
+    },
   },
 };
 ```
 
-The five methods map one-to-one onto the sidebar:
+The required methods map onto the sidebar; `updateMessage` is optional and backs in-message edits:
 
 | Method | When it runs |
 |---|---|
@@ -307,6 +311,7 @@ The five methods map one-to-one onto the sidebar:
 | `getMessages(threadId)` | User opens a thread. Returns its `Message[]`. |
 | `updateThread(thread)` | A thread changes (e.g. rename). Returns the updated `Thread`. |
 | `deleteThread(id)` | User deletes a thread. |
+| `updateMessage(threadId, message)` | *Optional.* User edits form state in a rendered message; called fire-and-forget. |
 
 ## 4. Store artifacts
 

packages/react-headless/src/adapters/restStorage.ts

@@ -84,6 +84,13 @@ export function restStorage({
       async deleteThread(id: string): Promise<void> {
         await request(`${baseUrl}/delete/${id}`, { method: "DELETE" });
       },
+      async updateMessage(threadId: string, message: Message): Promise<void> {
+        const [wire] = messageFormat.toApi([message]) as unknown[];
+        await request(`${baseUrl}/messages/${threadId}/${message.id}`, {
+          method: "PATCH",
+          body: JSON.stringify(wire),
+        });
+      },
     },
   };
 }

packages/react-headless/src/adapters/types.ts

@@ -11,6 +11,7 @@ export interface ThreadStorage {
   getMessages(threadId: string): Promise<Message[]>;
   updateThread(thread: Thread): Promise<Thread>;
   deleteThread(id: string): Promise<void>;
+  updateMessage?(threadId: string, message: Message): Promise<void>;
 }
 
 // ── Artifact storage (global, cross-thread) ──

packages/react-headless/src/store/createChatStore.ts

@@ -189,6 +189,10 @@ export const createChatStore = (configRef: React.RefObject<CreateChatStoreConfig
               set((s) => ({
                 messages: s.messages.map((m) => (m.id === msg.id ? msg : m)),
               })),
+            replaceMessageId: (previousId, serverId) =>
+              set((s) => ({
+                messages: s.messages.map((m) => (m.id === previousId ? { ...m, id: serverId } : m)),
+              })),
             // A tool's args have closed (TOOL_CALL_END) → it is now executing.
             markToolExecuting: (id) =>
               set((s) =>
@@ -231,6 +235,12 @@ export const createChatStore = (configRef: React.RefObject<CreateChatStoreConfig
         set((s) => ({
           messages: s.messages.map((m) => (m.id === message.id ? message : m)),
         }));
+        const threadId = get().selectedThreadId;
+        if (threadId !== null) {
+          threadStorage
+            .updateMessage?.(threadId, message)
+            .catch((e) => set(() => ({ threadError: e })));
+        }
       },
 
       setMessages: (messages: Message[]) => {

packages/react-headless/src/stream/adapters/openai-completions.ts

@@ -4,7 +4,11 @@ import { sseLineIterator } from "./_shared/sseLines";
 
 export const openAIAdapter = (): StreamProtocolAdapter => ({
   async *parse(response: Response): AsyncIterable<AGUIEvent> {
-    const messageId = crypto.randomUUID();
+    // Prefer the completion id (`json.id`, stable across the whole response and
+    // seen by both client and backend) as the message id, so an edit persisted
+    // later keys on an id the backend agrees on. Fall back to a client uuid only
+    // if the stream omits it. Set from the first chunk below.
+    let messageId = "";
     const toolCallIds: Record<number, string> = {};
     let messageStarted = false;
 
@@ -15,6 +19,7 @@ export const openAIAdapter = (): StreamProtocolAdapter => ({
 
       try {
         const json = JSON.parse(data) as ChatCompletionChunk;
+        if (!messageId) messageId = json.id || crypto.randomUUID();
         const choice = json.choices?.[0];
         const delta = choice?.delta;
 

packages/react-headless/src/stream/processStreamedMessage.ts

@@ -10,6 +10,8 @@ interface Parameters {
   createMessage: (message: Message) => void;
   /** A function that updates an existing message in the thread (matched by id). */
   updateMessage: (message: Message) => void;
+  /** Relabels an existing message in place (same position, new id) */
+  replaceMessageId?: (previousId: string, serverId: string) => void;
   /**
    * Marks a tool call as executing (args closed, awaiting result). Wired to the
    * store's `executingToolCallIds` set so `pairToolActivity` can report the
@@ -29,6 +31,7 @@ export const processStreamedMessage = async ({
   response,
   createMessage,
   updateMessage,
+  replaceMessageId,
   markToolExecuting = () => {},
   clearToolExecuting = () => {},
   adapter = agUIAdapter(),
@@ -159,9 +162,7 @@ export const processStreamedMessage = async ({
       case EventType.TEXT_MESSAGE_START: {
         // A DIFFERENT item id after content/tool calls have accumulated means
         // the model opened a new output message item — interleaving prose with
-        // tool calls (several sections in one run). Split into a fresh assistant
-        // message so the live structure matches what reload reconstructs from
-        // storage
+        // tool calls (several sections in one run).
         const startId = (event as { messageId?: string }).messageId ?? null;
         const hasBody =
           (currentMessage.content?.length ?? 0) > 0 || (currentMessage.toolCalls?.length ?? 0) > 0;
@@ -173,13 +174,27 @@ export const processStreamedMessage = async ({
             rafId = null;
             if (!isFirst) updateMessage(currentMessage);
           }
+          // Key the new segment by the server id when present (else a uuid) so
+          // it is created already carrying the persistable id.
           currentMessage = {
-            id: crypto.randomUUID(),
+            id: startId ?? crypto.randomUUID(),
             role: "assistant",
             content: "",
             toolCalls: [],
           };
           isFirst = true;
+        } else if (startId && startId !== currentMessage.id) {
+          // First (or same) item: adopt the server id in place of the optimistic
+          // uuid. Swap IN PLACE via replaceMessageId — deleting + re-creating
+          // would break ordering when tool messages were appended between
+          // the create and this event. Without replaceMessageId we keep the
+          // optimistic id rather  than desync currentMessage from the store.
+          if (isFirst) {
+            currentMessage = { ...currentMessage, id: startId };
+          } else if (replaceMessageId) {
+            replaceMessageId(currentMessage.id, startId);
+            currentMessage = { ...currentMessage, id: startId };
+          }
         }
         currentTextItemId = startId;
         break;

packages/react-ui/src/components/OpenUIChat/GenUIAssistantMessage.tsx

@@ -4,7 +4,7 @@ import type { AssistantMessage } from "@openuidev/react-headless";
 import { useThread } from "@openuidev/react-headless";
 import type { ActionEvent, Library } from "@openuidev/react-lang";
 import { BuiltinActionType, Renderer } from "@openuidev/react-lang";
-import { useCallback, useMemo } from "react";
+import { useCallback, useMemo, useRef } from "react";
 import { getLastAssistantMessageId } from "../../utils/messages";
 import {
   separateContentAndContext,
@@ -54,13 +54,19 @@ export const GenUIAssistantMessage = ({
   // Persist form state into the inline-wrapped message content. The original
   // header line (which may include `libraryVersion` and telemetry tags emitted
   // by the backend) is reused so attrs survive the persist round-trip.
+
+  const lastPersistedContentRef = useRef<string | null>(null);
   const handleStateUpdate = useCallback(
     (state: Record<string, any>) => {
       const hasState = Object.keys(state).length > 0;
       const contentPart = wrapContentWithHeader(content ?? "", contentHeader);
       const fullMessage = hasState
         ? contentPart + wrapContext(JSON.stringify([state]))
         : contentPart;
+      if (fullMessage === lastPersistedContentRef.current || fullMessage === message.content) {
+        return;
+      }
+      lastPersistedContentRef.current = fullMessage;
       updateMessage({ ...message, content: fullMessage });
     },
     [updateMessage, message, content, contentHeader],