feat(chat): tool approvals support — ID-matched message replacement, sendEmail example, approval UI

2026-04-14 13:07:51 +01:00
parent cc2ad1f565
commit 3241e1cd71
4 changed files with 133 additions and 14 deletions
@@ -3025,11 +3025,32 @@ function chatAgent<
                    accumulatedUIMessages = [...cleanedUIMessages];
                    // No new user messages for regenerate — just the response (added below)
                  } else {
-                    // Submit: frontend sent only the new user message(s). Append to accumulator.
-                    accumulatedMessages.push(...incomingModelMessages);
-                    accumulatedUIMessages.push(...cleanedUIMessages);
-                    turnNewModelMessages.push(...incomingModelMessages);
-                    turnNewUIMessages.push(...cleanedUIMessages);
+                    // Submit: check if any incoming message updates an existing one (by ID).
+                    // This handles tool approval responses, where the frontend resends the
+                    // assistant message with updated tool parts (approval-responded).
+                    // IDs match because we always pass generateMessageId + originalMessages
+                    // to toUIMessageStream, so the backend's start chunk carries the same
+                    // messageId that the frontend uses.
+                    let replaced = false;
+                    for (const incoming of cleanedUIMessages) {
+                      const idx = accumulatedUIMessages.findIndex((m) => m.id === incoming.id);
+                      if (idx !== -1) {
+                        accumulatedUIMessages[idx] = incoming as TUIMessage;
+                        replaced = true;
+                      } else {
+                        accumulatedUIMessages.push(incoming as TUIMessage);
+                        turnNewUIMessages.push(incoming as TUIMessage);
+                      }
+                    }
+                    if (replaced) {
+                      // Reconvert all model messages since a replacement changes the structure
+                      accumulatedMessages = await toModelMessages(accumulatedUIMessages);
+                    } else {
+                      accumulatedMessages.push(...incomingModelMessages);
+                    }
+                    if (turnNewUIMessages.length > 0) {
+                      turnNewModelMessages.push(...(await toModelMessages(turnNewUIMessages)));
+                    }
                  }

                  // Mint a scoped public access token once per turn, reused for
@@ -3157,15 +3178,20 @@ function chatAgent<
                  let runResult: unknown;

                  try {
-                    // Drain any messages injected by background work (e.g. self-review from previous turn)
+                    // Drain any messages injected by background work (e.g. self-review from previous turn).
+                    // Skip if the last message is a tool message — appending after it would
+                    // prevent streamText from finding pending tool approvals (it checks
+                    // the last message). The queued messages will be picked up by prepareStep
+                    // at the next step boundary instead.
+                    const lastAccumulated = accumulatedMessages[accumulatedMessages.length - 1];
                    const bgQueue = locals.get(chatBackgroundQueueKey);
-                    if (bgQueue && bgQueue.length > 0) {
+                    if (bgQueue && bgQueue.length > 0 && lastAccumulated?.role !== "tool") {
                      accumulatedMessages.push(...bgQueue.splice(0));
                    }

                    runResult = await userRun({
                      ...restWire,
-                      messages: await applyPrepareMessages(accumulatedMessages, "run"),
+                      messages: messagesForRun,
                      clientData,
                      continuation,
                      previousRunId,
@@ -3185,9 +3211,16 @@ function chatAgent<
                    // (e.g. for tool approval continuations / HITL flows).
                    if ((locals.get(chatPipeCountKey) ?? 0) === 0 && isUIMessageStreamable(runResult)) {
                      onFinishAttached = true;
+                      const resolvedOptions = resolveUIMessageStreamOptions();
                      const uiStream = runResult.toUIMessageStream({
-                        ...resolveUIMessageStreamOptions(),
+                        ...resolvedOptions,
+                        // Pass originalMessages so the AI SDK reuses message IDs across
+                        // turns (e.g. for tool approval continuations / HITL flows).
                        originalMessages: accumulatedUIMessages,
+                        // Always provide generateMessageId so the start chunk carries a
+                        // messageId. Without this, the frontend and backend generate IDs
+                        // independently and they won't match for ID-based dedup.
+                        generateMessageId: resolvedOptions.generateMessageId ?? generateMessageId,
                        onFinish: ({ responseMessage }: { responseMessage: UIMessage }) => {
                          capturedResponseMessage = responseMessage as TUIMessage;
                          resolveOnFinish!();
@@ -463,9 +463,10 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
    // If we have an existing run, send the message via input stream
    // to resume the conversation in the same run.
    if (session?.runId) {
+      const slicedMessages = trigger === "submit-message" ? messages.slice(-1) : messages;
      const minimalPayload = {
        ...payload,
-        messages: trigger === "submit-message" ? messages.slice(-1) : messages,
+        messages: slicedMessages,
      };

      const sendChatMessages = async (token: string) => {
@@ -1,6 +1,7 @@
 "use client";

 import { useChat } from "@ai-sdk/react";
+import { lastAssistantMessageIsCompleteWithApprovalResponses } from "ai";
 import type { ChatUiMessage } from "@/lib/chat-tools";
 import type { TriggerChatTransport } from "@trigger.dev/sdk/chat";
 import type { CompactionChunkData } from "@trigger.dev/sdk/ai";
@@ -9,7 +10,15 @@ import { useCallback, useEffect, useRef, useState } from "react";
 import { Streamdown } from "streamdown";
 import { MODEL_OPTIONS } from "@/lib/models";

-function ToolInvocation({ part }: { part: any }) {
+function ToolInvocation({
+  part,
+  onApprove,
+  onDeny,
+}: {
+  part: any;
+  onApprove?: (approvalId: string) => void;
+  onDeny?: (approvalId: string) => void;
+}) {
  const [expanded, setExpanded] = useState(false);
  const toolName = part.type.startsWith("tool-") ? part.type.slice(5) : "tool";
  const state = part.state ?? "input-available";
@@ -18,6 +27,9 @@ function ToolInvocation({ part }: { part: any }) {

  const isLoading = state === "input-streaming" || state === "input-available";
  const isError = state === "output-error";
+  const needsApproval = state === "approval-requested";
+  const wasApproved = state === "approval-responded" && part.approval?.approved === true;
+  const wasDenied = state === "approval-responded" && part.approval?.approved === false;

  return (
    <div className="my-1 rounded border border-gray-200 bg-gray-50 text-xs">
@@ -29,12 +41,37 @@ function ToolInvocation({ part }: { part: any }) {
        {isLoading && (
          <span className="inline-block h-3 w-3 animate-spin rounded-full border-2 border-gray-300 border-t-gray-600" />
        )}
-        {!isLoading && !isError && <span className="text-green-600">&#10003;</span>}
+        {needsApproval && <span className="text-amber-500">&#9888;</span>}
+        {wasApproved && <span className="text-green-600">&#10003;</span>}
+        {wasDenied && <span className="text-red-600">&#10007;</span>}
+        {!isLoading && !needsApproval && !wasApproved && !wasDenied && !isError && (
+          <span className="text-green-600">&#10003;</span>
+        )}
        {isError && <span className="text-red-600">&#10007;</span>}
        <span>{toolName}</span>
+        {needsApproval && <span className="text-amber-500 text-[10px]">needs approval</span>}
        <span className="ml-auto text-gray-400">{expanded ? "▲" : "▼"}</span>
      </button>

+      {needsApproval && (
+        <div className="flex gap-2 border-t border-gray-200 px-3 py-2">
+          <button
+            type="button"
+            onClick={() => onApprove?.(part.approval.id)}
+            className="rounded bg-green-600 px-3 py-1 text-xs font-medium text-white hover:bg-green-700"
+          >
+            Approve
+          </button>
+          <button
+            type="button"
+            onClick={() => onDeny?.(part.approval.id)}
+            className="rounded bg-red-600 px-3 py-1 text-xs font-medium text-white hover:bg-red-700"
+          >
+            Deny
+          </button>
+        </div>
+      )}
+
      {expanded && (
        <div className="border-t border-gray-200 px-3 py-2 space-y-2">
          {args && Object.keys(args).length > 0 && (
@@ -267,11 +304,20 @@ export function Chat({
  const turnCounter = useRef(0);
  const [ttfbHistory, setTtfbHistory] = useState<TtfbEntry[]>([]);

-  const { messages, setMessages, sendMessage, stop: aiStop, status, error } = useChat({
+  const {
+    messages,
+    setMessages,
+    sendMessage,
+    stop: aiStop,
+    addToolApprovalResponse,
+    status,
+    error,
+  } = useChat({
    id: chatId,
    messages: initialMessages,
    transport,
    resume: resumeProp,
+    sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses,
  });

  // Use transport.stopGeneration for reliable stop after reconnect.
@@ -282,6 +328,21 @@ export function Chat({
    aiStop();
  }, [transport, chatId, aiStop]);

+  // Tool approval callbacks
+  const handleApprove = useCallback(
+    (approvalId: string) => {
+      addToolApprovalResponse({ id: approvalId, approved: true });
+    },
+    [addToolApprovalResponse, chatId, messages, status]
+  );
+
+  const handleDeny = useCallback(
+    (approvalId: string) => {
+      addToolApprovalResponse({ id: approvalId, approved: false, reason: "User denied" });
+    },
+    [addToolApprovalResponse, chatId]
+  );
+
  // Notify parent of first user message (for chat metadata creation)
  useEffect(() => {
    if (hasCalledFirstMessage.current) return;
@@ -549,7 +610,14 @@ export function Chat({
                  }

                  if (part.type.startsWith("tool-")) {
-                    return <ToolInvocation key={i} part={part} />;
+                    return (
+                      <ToolInvocation
+                        key={i}
+                        part={part}
+                        onApprove={handleApprove}
+                        onDeny={handleDeny}
+                      />
+                    );
                  }

                  if (pending.isInjectionPoint(part)) {
@@ -308,6 +308,22 @@ export const executeJs = tool({
  },
 });

+export const sendEmail = tool({
+  description:
+    "Send an email to a recipient. Requires human approval before sending. " +
+    "Use when the user asks you to send, draft, or compose an email.",
+  inputSchema: z.object({
+    to: z.string().describe("Recipient email address"),
+    subject: z.string().describe("Email subject line"),
+    body: z.string().describe("Email body text"),
+  }),
+  needsApproval: true,
+  execute: async ({ to, subject, body }) => {
+    // Simulated — in a real app this would call an email API
+    return { sent: true, to, subject, preview: body.slice(0, 100) };
+  },
+});
+
 /** Tool set passed to `streamText` for the main `chat.agent` run (includes PostHog). */
 export const chatTools = {
  inspectEnvironment,
@@ -316,6 +332,7 @@ export const chatTools = {
  posthogQuery,
  executeCode,
  executeJs,
+  sendEmail,
 };

 type ChatToolSet = typeof chatTools;