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`
 
@@ -130,7 +130,7 @@ const storage = restStorage({ baseUrl: "/api/threads" }); // was threadApiUrl
 <AgentInterface llm={llm} storage={storage} />;
 ```
 
-`restStorage` hits the exact endpoints the old `threadApiUrl` prop did — `GET {baseUrl}/get`, `POST {baseUrl}/create`, `GET {baseUrl}/get/{threadId}`, `PATCH {baseUrl}/update/{id}`, `DELETE {baseUrl}/delete/{id}` — so an existing backend keeps working. Pass `messageFormat` here too if your backend stores messages in a provider shape.
+`restStorage` hits the exact endpoints the old `threadApiUrl` prop did — `GET {baseUrl}/get`, `POST {baseUrl}/create`, `GET {baseUrl}/get/{threadId}`, `PATCH {baseUrl}/update/{id}`, `DELETE {baseUrl}/delete/{id}` — so an existing backend keeps working. It also PATCHes `{baseUrl}/messages/{threadId}/{messageId}` on in-message edits (`{ messages: messageFormat.toApi([message]) }`); a missing route is swallowed and does not break the UI. Pass `messageFormat` here too if your backend stores messages in a provider shape.
 
 <Callout type="info">If you pass no `storage` at all, `AgentInterface` uses an internal in-memory store — fine for prototyping, but wiped on reload.</Callout>
 

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

@@ -158,7 +158,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 {
@@ -167,6 +167,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>;
 }
 ```
 
@@ -177,6 +178,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:
 
@@ -270,7 +272,7 @@ const storage = restStorage({
 
 ### Endpoint contract
 
-Each `ThreadStorage` operation maps to exactly one HTTP call under `baseUrl`. The paths are literal — these are the exact five endpoints you implement.
+Each `ThreadStorage` operation maps to exactly one HTTP call under `baseUrl`. The paths are literal — these are the exact endpoints you implement.
 
 | Operation | Method | Path | Request body | Response |
 |-----------|--------|------|--------------|----------|
@@ -279,14 +281,18 @@ Each `ThreadStorage` operation maps to exactly one HTTP call under `baseUrl`. Th
 | Get messages | `GET` | `{baseUrl}/get/{threadId}` | — | `Message[]` (run through `messageFormat.fromApi`) |
 | Update thread | `PATCH` | `{baseUrl}/update/{thread.id}` | the `Thread` | the updated `Thread` |
 | Delete thread | `DELETE` | `{baseUrl}/delete/{id}` | — | — |
+| Update message | `PATCH` | `{baseUrl}/messages/{threadId}/{message.id}` | `{ messages: messageFormat.toApi([message]) }` | — |
 
-With `baseUrl: "/api/threads"` the concrete paths are `/api/threads/get`, `/api/threads/create`, `/api/threads/get/{threadId}`, `/api/threads/update/{thread.id}`, and `/api/threads/delete/{id}`.
+With `baseUrl: "/api/threads"` the concrete paths are `/api/threads/get`, `/api/threads/create`, `/api/threads/get/{threadId}`, `/api/threads/update/{thread.id}`, `/api/threads/delete/{id}`, and `/api/threads/messages/{threadId}/{message.id}`.
+
+`updateMessage` is fire-and-forget on the client: a 404 from a backend that hasn't added this route yet is swallowed and does not surface as a thread error.
 
 ### Message format application
 
-`messageFormat` applies at exactly two points:
+`messageFormat` applies at these points:
 
 - **Create** — the request body's `messages` is `messageFormat.toApi([firstMessage])`.
+- **Update message** — the request body's `messages` is `messageFormat.toApi([message])`. One AG-UI message can expand to several wire items; the array is the full conversion, not a single item.
 - **Get messages** — the `{baseUrl}/get/{threadId}` response is run through `messageFormat.fromApi` on the way back.
 
 With the default `identityMessageFormat`, messages cross the wire as the canonical `Message` type unchanged. Pass a provider-specific format (e.g. `openAIMessageFormat`) when your backend stores messages in that provider's shape.

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

@@ -203,7 +203,7 @@ export default function App() {
 }
 ```
 
-That's the whole frontend change — one prop. You implement the five endpoints `restStorage` calls; the sidebar's thread lifecycle comes for free. With `baseUrl: "/api/threads"`:
+That's the whole frontend change — one prop. You implement the endpoints `restStorage` calls; the sidebar's thread lifecycle comes for free. With `baseUrl: "/api/threads"`:
 
 | Operation | Method | Path | Request body | Returns |
 |---|---|---|---|---|
@@ -212,8 +212,9 @@ That's the whole frontend change — one prop. You implement the five endpoints
 | Get messages | `GET` | `/api/threads/get/{threadId}` | — | the thread's `Message[]` |
 | Update thread | `PATCH` | `/api/threads/update/{threadId}` | the full `Thread` | the updated `Thread` |
 | Delete thread | `DELETE` | `/api/threads/delete/{threadId}` | — | nothing |
+| Update message | `PATCH` | `/api/threads/messages/{threadId}/{messageId}` | `{ messages: [...] }` (the edited message via the message format; may be several wire items) | nothing |
 
-`restStorage` throws a descriptive error on any non-`ok` response. The `Thread` shape at the boundary is `{ id, title, createdAt: string | number, isPending? }`. By default `restStorage` uses the identity message format; pass `messageFormat` (and optional `headers` / `fetch`) if your backend stores a provider-specific shape — it is applied to the `create` body (`toApi`) and the `get/{threadId}` response (`fromApi`).
+`restStorage` throws a descriptive error on any non-`ok` response. The client swallows `updateMessage` failures so a backend that hasn't added this route yet doesn't surface a thread error. The `Thread` shape at the boundary is `{ id, title, createdAt: string | number, isPending? }`. By default `restStorage` uses the identity message format; pass `messageFormat` (and optional `headers` / `fetch`) if your backend stores a provider-specific shape — it is applied to the `create` and `update message` bodies (`toApi`) and the `get/{threadId}` response (`fromApi`).
 
 ```tsx
 import { openAIMessageFormat } from "@openuidev/react-ui";
@@ -262,13 +263,20 @@ export async function DELETE(_req: NextRequest, { params }: { params: { threadId
   await db.deleteThread(params.threadId);
   return new NextResponse(null, { status: 204 });
 }
+
+// app/api/threads/messages/[threadId]/[messageId]/route.ts — persist an in-message edit
+export async function PATCH(req: NextRequest, { params }: { params: { threadId: string; messageId: string } }) {
+  const { messages } = await req.json(); // messageFormat.toApi([editedMessage]) — one AG-UI message may be several wire items
+  await db.updateMessage(params.threadId, params.messageId, messages);
+  return new NextResponse(null, { status: 204 });
+}
 ```
 
 `db` is a stand-in for your persistence layer — Postgres, SQLite, Redis, a cloud KV store, anything. The endpoints are thin: read/write threads and messages, return the shapes the table describes.
 
 ### 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";
@@ -295,11 +303,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 |
 |---|---|
@@ -308,6 +320,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/__tests__/restStorage.test.ts

@@ -83,6 +83,24 @@ describe("restStorage", () => {
     expect(result).toEqual(updated);
   });
 
+  it("updateMessage PATCHes {base}/messages/:threadId/:messageId with the full toApi result", async () => {
+    fetchSpy.mockResolvedValue(json({}, true));
+    const toApi = vi.fn((msgs) => msgs.flatMap((m: any) => [{ custom: m.id }, { extra: true }]));
+    const messageFormat: MessageFormat = { toApi, fromApi: (d) => d as any };
+    const message = { id: "m1", role: "assistant" as const, content: "hi" };
+
+    await make({ messageFormat }).updateMessage("t1", message);
+
+    const [url, opts] = fetchSpy.mock.calls[0];
+    expect(url).toBe("/api/threads/messages/t1/m1");
+    expect(opts.method).toBe("PATCH");
+    expect(opts.headers["Content-Type"]).toBe("application/json");
+    expect(toApi).toHaveBeenCalledWith([message]);
+    expect(JSON.parse(opts.body)).toEqual({
+      messages: [{ custom: "m1" }, { extra: true }],
+    });
+  });
+
   it("deleteThread DELETEs {base}/delete/:id", async () => {
     fetchSpy.mockResolvedValue(json({}, true));
     await make().deleteThread("t1");

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

@@ -7,11 +7,12 @@ export interface RestStorageOptions {
   /**
    * Base URL for thread endpoints (the old `threadApiUrl`). The factory hits
    * the same conventions the legacy default used:
-   *   - list:   GET    {baseUrl}/get  (·  ?cursor={cursor})
-   *   - create: POST   {baseUrl}/create
-   *   - get:    GET    {baseUrl}/get/{threadId}
-   *   - update: PATCH  {baseUrl}/update/{threadId}
-   *   - delete: DELETE {baseUrl}/delete/{threadId}
+   *   - list:    GET    {baseUrl}/get  (·  ?cursor={cursor})
+   *   - create:  POST   {baseUrl}/create
+   *   - get:     GET    {baseUrl}/get/{threadId}
+   *   - update:  PATCH  {baseUrl}/update/{threadId}
+   *   - delete:  DELETE {baseUrl}/delete/{threadId}
+   *   - message: PATCH  {baseUrl}/messages/{threadId}/{messageId}
    */
   baseUrl: string;
   /** Wire-format conversion. Defaults to identity (canonical Message). */
@@ -84,6 +85,14 @@ export function restStorage({
       async deleteThread(id: string): Promise<void> {
         await request(`${baseUrl}/delete/${id}`, { method: "DELETE" });
       },
+      async updateMessage(threadId: string, message: Message): Promise<void> {
+        // Send the full toApi conversion. One AG-UI message can map to several
+        // wire items (e.g. OpenAI Responses flattens text + tool calls).
+        await request(`${baseUrl}/messages/${threadId}/${message.id}`, {
+          method: "PATCH",
+          body: JSON.stringify({ messages: messageFormat.toApi([message]) }),
+        });
+      },
     },
   };
 }

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) {
+          // Fire-and-forget: a backend that hasn't implemented updateMessage yet
+          // (or a transient failure) shouldn't surface a thread-level error.
+          threadStorage.updateMessage?.(threadId, message).catch(() => {});
+        }
       },
 
       setMessages: (messages: Message[]) => {

packages/react-headless/src/stream/adapters/langgraph.ts

@@ -66,7 +66,11 @@ export const langGraphAdapter = (options?: LangGraphAdapterOptions): StreamProto
     if (!reader) throw new Error("No response body");
 
     const decoder = new TextDecoder();
-    const messageId = crypto.randomUUID();
+    // Prefer the LangChain message id (`msg.id`, stable across all chunks of
+    // the same AI message and known to the backend) so a persisted edit keys
+    // on an id the backend agrees on. Fall back to a client uuid only if the
+    // stream omits it. Set from the first "messages" chunk below.
+    let messageId = "";
     const toolCallIds: Record<number, string> = {};
     let messageStarted = false;
     let buffer = "";
@@ -122,6 +126,8 @@ export const langGraphAdapter = (options?: LangGraphAdapterOptions): StreamProto
               break;
             }
 
+            if (!messageId) messageId = msg.id || crypto.randomUUID();
+
             // Emit TEXT_MESSAGE_START on first AI message chunk
             if (!messageStarted) {
               yield {

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;