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 the interface.
 
 ### The REST case → `restStorage`
 

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

@@ -147,7 +147,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 {
@@ -156,6 +156,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>;
 }
 ```
 
@@ -166,6 +167,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";
@@ -308,6 +308,7 @@ The five methods map one-to-one onto the sidebar:
 | `updateThread(thread)` | A thread changes (e.g. rename). Returns the updated `Thread`. |
 | `deleteThread(id)` | User deletes a thread. |
 
+
 ## 4. Store artifacts
 
 Adding an optional `artifact` channel to your storage makes every dashboard, report, and presentation the agent produces a durable, searchable, cross-thread record — and `<AgentInterface>` renders the entire artifact browser (sidebar entry → searchable list → full-page view) for you. `ArtifactStorage` is three methods:

docs/public/AGENTS.md

@@ -282,6 +282,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>; // optional: persists client-side message rewrites (form state); omit → edits are in-memory only
 }
 
 interface ArtifactStorage {

packages/react-headless/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@openuidev/react-headless",
-  "version": "0.9.2",
+  "version": "0.9.3",
   "description": "Headless React primitives for AI chat — state management, streaming adapters for OpenAI and AG-UI, message format converters, and thread management for OpenUI generative UI apps",
   "license": "MIT",
   "type": "module",

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

@@ -184,6 +184,10 @@ export const createChatStore = (config: 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) =>
@@ -226,6 +230,12 @@ export const createChatStore = (config: 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/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(),
@@ -42,6 +45,9 @@ export const processStreamedMessage = async ({
 
   let isFirst = true;
 
+  // First TEXT_MESSAGE_START wins — see that case below.
+  let serverIdAdopted = false;
+
   // Tool messages by toolCallId, so repeated TOOL_CALL_RESULTs for the same
   // call UPDATE one message in place instead of duplicating it.
   const toolMessagesByCallId = new Map<string, ToolMessage>();
@@ -154,11 +160,23 @@ export const processStreamedMessage = async ({
       }
 
       case EventType.TEXT_MESSAGE_START:
-        // The optimistic id is kept regardless of `event.messageId` — swapping
-        // ids mid-stream by deleting + re-creating the assistant message
-        // breaks ordering when tool messages have already been appended
-        // between the original create and this event (e.g. from
-        // TOOL_CALL_RESULT). Persistence layers should map ids on save.
+        // Use the server id when available (usually available from
+        // the TEXT_MESSAGE_START event. The swap is done IN PLACE
+        // — deleting + re-creating the assistant message would break
+        // ordering when tool messages were already appended between the
+        // If the message isn't in the store yet (this is the first event), the
+        // trailing isFirst createMessage below already carries the server id.
+        //
+        // First TEXT_MESSAGE_START wins: a multi-text-item run emits one per
+        // item, and re-relabeling would remount the message component (keyed
+        // by id) for each.
+        if (!serverIdAdopted && event.messageId) {
+          serverIdAdopted = true;
+          if (event.messageId !== currentMessage.id) {
+            replaceMessageId?.(currentMessage.id, event.messageId);
+            currentMessage = { ...currentMessage, id: event.messageId };
+          }
+        }
         break;
 
       case EventType.TOOL_CALL_RESULT: {

packages/react-ui/package.json

@@ -2,7 +2,7 @@
   "type": "module",
   "name": "@openuidev/react-ui",
   "license": "MIT",
-  "version": "0.13.0",
+  "version": "0.13.1",
   "description": "Component library for Generative UI SDK",
   "main": "dist/index.cjs",
   "module": "dist/index.mjs",

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

@@ -9,7 +9,7 @@ import {
 } 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 {
   separateContentAndContext,
   wrapContent,
@@ -81,6 +81,8 @@ 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 code = openuiCode ?? "";
@@ -89,6 +91,10 @@ export const GenUIAssistantMessage = ({
       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, openuiCode, contentHeader],