From a7078eb0cbfa71f41e17d4beb3ee8a4a59c7973f Mon Sep 17 00:00:00 2001
From: Eric Allam <eallam@icloud.com>
Date: Fri, 17 Apr 2026 09:59:18 +0100
Subject: [PATCH] feat(ai-chat): add askUser tool for HITL testing, verify
 TRI-8556 fix

---
 references/ai-chat/src/components/chat.tsx | 40 ++++++++++++++++++++--
 references/ai-chat/src/lib/chat-tools.ts   | 22 ++++++++++++
 references/ai-chat/src/trigger/chat.ts     | 17 ++++-----
 3 files changed, 69 insertions(+), 10 deletions(-)

diff --git a/references/ai-chat/src/components/chat.tsx b/references/ai-chat/src/components/chat.tsx
index b5f3d6d27..93990c0b3 100644
--- a/references/ai-chat/src/components/chat.tsx
+++ b/references/ai-chat/src/components/chat.tsx
@@ -1,7 +1,10 @@
 "use client";
 
 import { useChat } from "@ai-sdk/react";
-import { lastAssistantMessageIsCompleteWithApprovalResponses } from "ai";
+import {
+  lastAssistantMessageIsCompleteWithApprovalResponses,
+  lastAssistantMessageIsCompleteWithToolCalls,
+} 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";
@@ -14,10 +17,12 @@ function ToolInvocation({
   part,
   onApprove,
   onDeny,
+  onToolOutput,
 }: {
   part: any;
   onApprove?: (approvalId: string) => void;
   onDeny?: (approvalId: string) => void;
+  onToolOutput?: (toolCallId: string, output: unknown) => void;
 }) {
   const [expanded, setExpanded] = useState(false);
   const toolName = part.type.startsWith("tool-") ? part.type.slice(5) : "tool";
@@ -72,6 +77,31 @@ function ToolInvocation({
         </div>
       )}
 
+      {/* askUser tool: show question + option buttons when input-available */}
+      {toolName === "askUser" && state === "input-available" && args?.question && (
+        <div className="border-t border-gray-200 px-3 py-2 space-y-2">
+          <div className="font-medium text-gray-700">{args.question}</div>
+          <div className="flex flex-wrap gap-2">
+            {(args.options ?? []).map((opt: any) => (
+              <button
+                key={opt.id}
+                type="button"
+                onClick={() =>
+                  onToolOutput?.(part.toolCallId, {
+                    skipped: false,
+                    answers: [{ questionId: args.question, optionId: opt.id, text: opt.label }],
+                  })
+                }
+                className="rounded border border-blue-300 bg-blue-50 px-3 py-1.5 text-xs text-blue-700 hover:bg-blue-100"
+                title={opt.description}
+              >
+                {opt.label}
+              </button>
+            ))}
+          </div>
+        </div>
+      )}
+
       {expanded && (
         <div className="border-t border-gray-200 px-3 py-2 space-y-2">
           {args && Object.keys(args).length > 0 && (
@@ -310,6 +340,7 @@ export function Chat({
     sendMessage,
     stop: aiStop,
     addToolApprovalResponse,
+    addToolOutput,
     status,
     error,
   } = useChat({
@@ -317,7 +348,9 @@ export function Chat({
     messages: initialMessages,
     transport,
     resume: resumeProp,
-    sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses,
+    sendAutomaticallyWhen: (opts) =>
+      lastAssistantMessageIsCompleteWithApprovalResponses(opts) ||
+      lastAssistantMessageIsCompleteWithToolCalls(opts),
   });
 
   // Use transport.stopGeneration for reliable stop after reconnect.
@@ -616,6 +649,9 @@ export function Chat({
                         part={part}
                         onApprove={handleApprove}
                         onDeny={handleDeny}
+                        onToolOutput={(toolCallId, output) =>
+                          addToolOutput({ toolCallId, output })
+                        }
                       />
                     );
                   }
diff --git a/references/ai-chat/src/lib/chat-tools.ts b/references/ai-chat/src/lib/chat-tools.ts
index a2330529e..9057e4c3b 100644
--- a/references/ai-chat/src/lib/chat-tools.ts
+++ b/references/ai-chat/src/lib/chat-tools.ts
@@ -324,6 +324,27 @@ export const sendEmail = tool({
   },
 });
 
+export const askUser = tool({
+  description:
+    "Ask the user a question when you need clarification or input before proceeding. " +
+    "Present 2-4 options for the user to choose from. Use when uncertain about the user's intent.",
+  inputSchema: z.object({
+    question: z.string().describe("The question to ask the user"),
+    options: z
+      .array(
+        z.object({
+          id: z.string().describe("Unique option identifier"),
+          label: z.string().describe("Short option title"),
+          description: z.string().optional().describe("Longer explanation"),
+        })
+      )
+      .min(2)
+      .max(4),
+  }),
+  // No execute function — streamText ends, turn completes,
+  // frontend sends the answer via addToolOutput
+});
+
 /** Tool set passed to `streamText` for the main `chat.agent` run (includes PostHog). */
 export const chatTools = {
   inspectEnvironment,
@@ -333,6 +354,7 @@ export const chatTools = {
   executeCode,
   executeJs,
   sendEmail,
+  askUser,
 };
 
 type ChatToolSet = typeof chatTools;
diff --git a/references/ai-chat/src/trigger/chat.ts b/references/ai-chat/src/trigger/chat.ts
index 1a5a7667b..5f298fb10 100644
--- a/references/ai-chat/src/trigger/chat.ts
+++ b/references/ai-chat/src/trigger/chat.ts
@@ -176,7 +176,7 @@ export const aiChat = chat
   .agent({
     id: "ai-chat",
     idleTimeoutInSeconds: 60,
-    chatAccessTokenTTL: "1m",
+    chatAccessTokenTTL: "1m", // TRI-8556 test
 
     // #region Compaction — automatic context window management
     compaction: {
@@ -392,15 +392,16 @@ export const aiChat = chat
       chatAccessToken,
       lastEventId,
     }) => {
-      // Log whether data-turn-metadata persisted to the response
-      const metadataParts = responseMessage?.parts?.filter(
-        (p: any) => p.type === "data-turn-metadata"
-      );
-      logger.info("onTurnComplete response parts check", {
+      // Log responseMessage parts for debugging TRI-8556
+      const partTypes = responseMessage?.parts?.map((p: any) => p.type) ?? [];
+      const toolParts = responseMessage?.parts?.filter((p: any) => p.type?.startsWith("tool-")) ?? [];
+      logger.info("onTurnComplete responseMessage", {
         hasResponseMessage: !!responseMessage,
+        responseMessageId: responseMessage?.id,
         totalParts: responseMessage?.parts?.length ?? 0,
-        metadataPartsCount: metadataParts?.length ?? 0,
-        metadataParts,
+        partTypes,
+        toolPartsCount: toolParts.length,
+        toolParts: toolParts.map((p: any) => ({ type: p.type, state: p.state, toolCallId: p.toolCallId })),
       });
       await prisma.chat.update({
         where: { id: chatId },