feat: customizable tool render behavior in ai-core - #561

packages/ai-core/src/agents/AgentUtils.ts (1)

444-455: ⚠️ Potential issue | 🟠 Major

Make the approval wait abortable too.

The new final sweep at Lines 520-532 is not reached if abortSignal fires while this Promise is awaiting user approval. That can still leave an approval-requested tool call and approval UI stuck until the user responds. Route this path through the same cancel/snapshot handling and clear the pending approval on abort.

🐛 Proposed direction

+  const markPendingToolCallsCancelled = () => {
+    const now = Date.now();
+    for (const [id, tc] of toolCallMap) {
+      if (tc.state === 'pending' || tc.state === 'approval-requested') {
+        toolCallMap.set(id, {
+          ...tc,
+          state: 'error',
+          errorText: TOOL_CALL_CANCELLED,
+          completedAt: now,
+        });
+      }
+    }
+  };
+
+  const throwAbortWithSnapshot = (): never => {
+    markPendingToolCallsCancelled();
+    const snapshot = buildAbortSnapshot(
+      parentToolCallId,
+      toolCallMap,
+      finalText,
+      store,
+    );
+    store.getState().ai.writeAbortSnapshot?.(parentToolCallId, snapshot);
+    pushProgress();
+    throw new ToolAbortError(TOOL_CALL_CANCELLED, snapshot);
+  };
+
-    const approved = await new Promise<boolean>((resolve) => {
-      throwIfAborted();
+    const approved = await new Promise<boolean>((resolve, reject) => {
+      if (abortSignal?.aborted) {
+        reject(new ToolAbortError(TOOL_CALL_CANCELLED));
+        return;
+      }
+
+      const onAbort = () => reject(new ToolAbortError(TOOL_CALL_CANCELLED));
+      abortSignal?.addEventListener('abort', onAbort, {once: true});
+
       store.getState().ai.requestSubAgentApproval!({
         toolCallId: pendingApproval!.toolCallId,
         approvalId: pendingApproval!.approvalId,
         toolName: approvedEntry.toolName,
         input: approvedEntry.input,
-        resolve,
+        resolve: (approved) => {
+          abortSignal?.removeEventListener('abort', onAbort);
+          resolve(approved);
+        },
       });
+    }).catch((err) => {
+      if (abortSignal?.aborted) {
+        store.getState().ai.clearSubAgentApproval?.(
+          pendingApproval!.approvalId,
+        );
+        throwAbortWithSnapshot();
+      }
+      throw err;
     });

Also applies to: 520-532

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/ai-core/src/agents/AgentUtils.ts` around lines 444 - 455, The
approval-await block using new Promise<boolean> is not abortable: if abortSignal
fires while waiting the pendingApproval and the `approval-requested` tool UI
remain stuck. Make this wait listen to the same `abortSignal`/cancel + snapshot
handling used later (the path around the code at 520-532) by registering an
abort handler before calling `store.getState().ai.requestSubAgentApproval` that
clears `pendingApproval` and invokes the same cancel/snapshot cleanup (or
rejects/resolves the promise) and then removes the abort listener; ensure
`throwIfAborted()` is still called before and after and that the handler clears
the pending approval and resolves/rejects the awaiting Promise so the flow
continues into the existing abort cleanup.

packages/ai-core/src/agents/AgentUtils.ts (1)

456-467: ⚠️ Potential issue | 🟠 Major

Handle aborts while waiting for sub-agent approval.

Line 456 can wait forever if abortSignal fires after requestSubAgentApproval is surfaced, so the new cancellation normalization is never reached and the tool can remain approval-requested. Add an abort listener around the approval promise and reuse the same snapshot/cancel path.

🐛 Proposed fix

-    const approved = await new Promise<boolean>((resolve) => {
-      throwIfAborted();
-      store.getState().ai.requestSubAgentApproval!({
-        toolCallId: pendingApproval!.toolCallId,
-        approvalId: pendingApproval!.approvalId,
-        toolName: approvedEntry.toolName,
-        input: approvedEntry.input,
-        resolve,
-      });
-    });
-
-    throwIfAborted();
+    let approved: boolean;
+    try {
+      approved = await new Promise<boolean>((resolve, reject) => {
+        const onAbort = () => reject(new ToolAbortError(TOOL_CALL_CANCELLED));
+
+        if (abortSignal?.aborted) {
+          onAbort();
+          return;
+        }
+
+        abortSignal?.addEventListener('abort', onAbort, {once: true});
+        store.getState().ai.requestSubAgentApproval!({
+          toolCallId: pendingApproval!.toolCallId,
+          approvalId: pendingApproval!.approvalId,
+          toolName: approvedEntry.toolName,
+          input: approvedEntry.input,
+          resolve: (value) => {
+            abortSignal?.removeEventListener('abort', onAbort);
+            resolve(value);
+          },
+        });
+      });
+
+      throwIfAborted();
+    } catch (err) {
+      if (abortSignal?.aborted) {
+        store.getState().ai.clearSubAgentApproval?.(pendingApproval.approvalId);
+        markPendingToolCallsAsCancelled(toolCallMap);
+        const snapshot = buildAbortSnapshot(
+          parentToolCallId,
+          toolCallMap,
+          finalText,
+          store,
+        );
+        store.getState().ai.writeAbortSnapshot?.(parentToolCallId, snapshot);
+        pushProgress();
+        throw new ToolAbortError(TOOL_CALL_CANCELLED, snapshot);
+      }
+      throw err;
+    }

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/ai-core/src/agents/AgentUtils.ts` around lines 456 - 467, The
approval Promise can hang if abortSignal fires after requestSubAgentApproval is
emitted; wrap the new Promise in an abort-aware pattern: register an abort
listener on abortSignal before calling
store.getState().ai.requestSubAgentApproval and, if aborted, call the same
snapshot/cancel path used by throwIfAborted (reuse throwIfAborted behavior) to
resolve/reject the Promise and clean up the request (e.g., call the original
cancel logic for pendingApproval), and ensure you remove the abort listener when
the Promise resolves; update the code surrounding pendingApproval,
requestSubAgentApproval, and throwIfAborted to use this abort-listener + cleanup
pattern.

packages/ai-core/src/components/FlatAgentRenderer.tsx (1)

100-109: ⚠️ Potential issue | 🔴 Critical

Add runtime type-check for reasoning before rendering.

toolCall.input is unknown-shaped; the as string | undefined type assertion only silences TypeScript and does not validate runtime values. If a tool sends {reasoning: {...}} or {reasoning: [...]}, the assertion allows the non-string value through, causing it to be assigned to label and rendered in JSX. React will crash when rendering non-string primitives as children.

🛡️ Proposed fix

-  const reasoning = inputObj?.reasoning as string | undefined;
+  const rawReasoning = inputObj?.reasoning;
+  const reasoning =
+    typeof rawReasoning === 'string' ? rawReasoning : undefined;

Apply to both ActivityLogLine (lines 100–109) and OrchestratorLogLineInner (lines 580–589).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/ai-core/src/components/FlatAgentRenderer.tsx` around lines 100 -
109, The code assigns reasoning with a TypeScript-only assertion and may pass
non-string values into label causing React to crash; update the runtime check in
ActivityLogLine (and similarly in OrchestratorLogLineInner) so that after
deriving inputObj and reasoning you validate with typeof reasoning === 'string'
before using it (e.g., const reasoning = typeof inputObj?.reasoning === 'string'
? inputObj.reasoning : undefined) and then build label using that validated
string along with getActivityLabel and isPending as before.