From 080bace3f0737888d051758a62e88bc193e38d28 Mon Sep 17 00:00:00 2001
From: Eric Allam <eallam@icloud.com>
Date: Wed, 6 May 2026 10:23:24 +0100
Subject: [PATCH] feat(ai-chat-ref): add undo action + button to demo agent
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Wire a typed  action on the main `aiChat` agent in
references/ai-chat — `actionSchema` accepts `{ type: "undo" }` and
`onAction` calls `chat.history.slice(0, -2)` to drop the last
user/assistant exchange. Adds an Undo button to the chat input row
that calls `transport.sendAction(chatId, { type: "undo" })` and
optimistically updates local `useChat` state via `setMessages`.

Exercises the new TRI-9118 action semantics end-to-end through the
demo UI: clicking Undo emits a `chat action` span (not `chat turn`),
fires only `onAction()`, no `run()` / `streamText` / turn lifecycle
hooks. The next message turn sees the truncated server-side history.
---
 references/ai-chat/src/components/chat.tsx | 18 ++++++++++++++++++
 references/ai-chat/src/trigger/chat.ts     | 15 +++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/references/ai-chat/src/components/chat.tsx b/references/ai-chat/src/components/chat.tsx
index ff73688ae..12b0c78ad 100644
--- a/references/ai-chat/src/components/chat.tsx
+++ b/references/ai-chat/src/components/chat.tsx
@@ -937,6 +937,24 @@ export function Chat({
               Stop
             </button>
           )}
+          {/* Undo — server-side `chat.history.slice(0, -2)` via the
+              `undo` action, optimistically reflected in the local
+              `useChat` state. Drops the last user / assistant exchange.
+              Only meaningful when there's at least one full exchange
+              and the chat isn't currently streaming. */}
+          {status !== "streaming" && messages.length >= 2 && (
+            <button
+              type="button"
+              disabled={isReadOnly}
+              onClick={() => {
+                void transport.sendAction(chatId, { type: "undo" });
+                setMessages((prev) => prev.slice(0, -2));
+              }}
+              className="rounded-lg bg-amber-500 px-4 py-2 text-sm font-medium text-white hover:bg-amber-600 disabled:opacity-50"
+            >
+              Undo
+            </button>
+          )}
         </div>
       </form>
     </div>
diff --git a/references/ai-chat/src/trigger/chat.ts b/references/ai-chat/src/trigger/chat.ts
index b8d674828..2d90ca6df 100644
--- a/references/ai-chat/src/trigger/chat.ts
+++ b/references/ai-chat/src/trigger/chat.ts
@@ -401,6 +401,21 @@ export const aiChat = chat
     },
     // #endregion
 
+    // #region actionSchema + onAction — typed actions for state-only mutations
+    // Actions are not turns: only `hydrateMessages` and `onAction` fire,
+    // no `run()` invocation, no model call. The `undo` action drops the
+    // last user/assistant exchange so the next message turn sees a
+    // truncated history.
+    actionSchema: z.discriminatedUnion("type", [
+      z.object({ type: z.literal("undo") }),
+    ]),
+    onAction: async ({ action }) => {
+      if (action.type === "undo") {
+        chat.history.slice(0, -2);
+      }
+    },
+    // #endregion
+
     // #region onTurnComplete — persist + background self-review via chat.inject()
     onTurnComplete: async ({
       chatId,