Show tool use before it completes (#24)

* Show tool use before it completes

* Review fixes

* Review fixes

* Review fixes

* Update storybook

Author: ScriptSmith

ui/src/components/MultiModelResponse/ContentRound.tsx

@@ -1,4 +1,4 @@
-import { memo, useState, useCallback, useMemo } from "react";
+import { memo, useState, useCallback, useMemo, useEffect, useRef } from "react";
 import type { ToolExecutionRound, Artifact, DisplaySelectionData } from "@/components/chat-types";
 import { Artifact as ArtifactComponent } from "@/components/Artifact";
 import { ReasoningSection } from "@/components/ReasoningSection/ReasoningSection";
@@ -47,8 +47,33 @@ function ContentRoundComponent({
   displaySelection,
   allOutputArtifacts,
 }: ContentRoundProps) {
-  const [toolsExpanded, setToolsExpanded] = useState(false);
-  const handleToggleTools = useCallback(() => setToolsExpanded((p) => !p), []);
+  const [isManuallyExpanded, setIsManuallyExpanded] = useState(false);
+  const [userOverride, setUserOverride] = useState(false);
+  const [wasAutoExpanded, setWasAutoExpanded] = useState(false);
+
+  // Auto-expand when tools are streaming, stay open after stream ends until user collapses
+  const toolsExpanded = useMemo(() => {
+    if (userOverride) return isManuallyExpanded;
+    if (isToolsStreaming) return true;
+    if (wasAutoExpanded) return true;
+    return isManuallyExpanded;
+  }, [isToolsStreaming, isManuallyExpanded, userOverride, wasAutoExpanded]);
+
+  // Reset user override only on false→true transition (new streaming session)
+  const prevStreamingRef = useRef(false);
+  useEffect(() => {
+    if (isToolsStreaming && !prevStreamingRef.current) {
+      setUserOverride(false);
+      setWasAutoExpanded(true);
+    }
+    prevStreamingRef.current = isToolsStreaming;
+  }, [isToolsStreaming]);
+
+  const handleToggleTools = useCallback(() => {
+    setIsManuallyExpanded(!toolsExpanded);
+    setUserOverride(true);
+    setWasAutoExpanded(false);
+  }, [toolsExpanded]);
   const compactMode = useCompactMode();
 
   // Resolve display selection to actual artifacts

ui/src/components/MultiModelResponse/MultiModelResponse.stories.tsx

@@ -562,7 +562,7 @@ export const WithReasoningContent: Story = {
 };
 
 /**
- * Test: Streaming with empty content shows Thinking indicator
+ * Test: Streaming with running tool shows tool execution UI (not Thinking indicator)
  */
 export const WithToolCallSearching: Story = {
   args: {
@@ -596,8 +596,9 @@ export const WithToolCallSearching: Story = {
   play: async ({ canvasElement }) => {
     const canvas = within(canvasElement);
 
-    // While streaming with no content, shows Thinking indicator
-    await expect(canvas.getByText("Thinking")).toBeInTheDocument();
+    // While tools are visibly running, shows tool status instead of Thinking indicator
+    const runningElements = canvas.getAllByText("running");
+    await expect(runningElements.length).toBeGreaterThan(0);
   },
 };
 

ui/src/components/MultiModelResponse/MultiModelResponse.tsx

@@ -740,10 +740,36 @@ const ModelResponseCard = memo(function ModelResponseCard({
     response.toolExecutionRounds,
   ]);
 
+  // Bridge live tool execution data into the last completed round so that
+  // ContentRound can render the execution timeline while tools are still running.
+  // Gated on isStreaming so committed (DB-loaded) messages are unaffected.
+  const completedRoundsWithLiveTools = useMemo(() => {
+    if (!response.isStreaming || !toolExecutionRounds.length) {
+      return completedRounds;
+    }
+    const liveRound = toolExecutionRounds[toolExecutionRounds.length - 1];
+    // No completed rounds yet (model called tool immediately) — synthesize one
+    if (!completedRounds.length) {
+      return [{ toolExecution: liveRound }];
+    }
+    const last = completedRounds[completedRounds.length - 1];
+    // Last round already has tool execution (back-to-back tool calls) — append new round
+    if (last.toolExecution) {
+      return [...completedRounds, { toolExecution: liveRound }];
+    }
+    // Last round is text-only — inject live tools into it
+    const merged = [...completedRounds];
+    merged[merged.length - 1] = {
+      ...last,
+      toolExecution: liveRound,
+    };
+    return merged;
+  }, [completedRounds, toolExecutionRounds, response.isStreaming]);
+
   // All output artifacts across all rounds (for resolving display_artifacts selections)
   const allOutputArtifacts = useMemo(() => {
     const result: ArtifactType[] = [];
-    for (const round of completedRounds) {
+    for (const round of completedRoundsWithLiveTools) {
       if (round.toolExecution) {
         for (const execution of round.toolExecution.executions) {
           for (const a of execution.outputArtifacts) {
@@ -753,7 +779,7 @@ const ModelResponseCard = memo(function ModelResponseCard({
       }
     }
     return result;
-  }, [completedRounds]);
+  }, [completedRoundsWithLiveTools]);
 
   // Extract display selection for a specific tool execution round
   const getDisplaySelectionForRound = useCallback(
@@ -959,31 +985,44 @@ const ModelResponseCard = memo(function ModelResponseCard({
             {/* Content: unified rendering via ContentRound for all responses */}
             {(() => {
               // Detect in-flight content that hasn't been captured in a completed round yet.
-              // completedRounds is always populated (even for single-round responses),
+              // completedRoundsWithLiveTools is always populated (even for single-round responses),
               // so this only shows content actively streaming in the current round.
               const currentReasoning =
                 response.isStreaming &&
                 response.reasoningContent &&
-                !completedRounds.some((r) => r.reasoning === response.reasoningContent)
+                !completedRoundsWithLiveTools.some((r) => r.reasoning === response.reasoningContent)
                   ? response.reasoningContent
                   : null;
               const currentContent =
                 response.isStreaming && response.content?.trim() ? response.content : null;
               const showInFlight = currentReasoning || currentContent;
+
+              // Suppress the streaming status indicator when running tools are
+              // already visible in expanded ContentRounds (non-compact mode only)
+              const hasVisibleRunningTools =
+                !compactMode &&
+                completedRoundsWithLiveTools.some((r) =>
+                  r.toolExecution?.executions.some(
+                    (e) => e.status === "running" || e.status === "pending"
+                  )
+                );
+
               return (
                 <div className="space-y-3">
-                  {completedRounds.map((round, i) => (
+                  {completedRoundsWithLiveTools.map((round, i) => (
                     <ContentRound
                       key={i}
                       reasoning={round.reasoning}
                       content={round.content}
                       reasoningTokenCount={
-                        completedRounds.length === 1 ? response.usage?.reasoningTokens : undefined
+                        completedRoundsWithLiveTools.length === 1
+                          ? response.usage?.reasoningTokens
+                          : undefined
                       }
                       toolExecutionRound={round.toolExecution}
                       isToolsStreaming={
                         response.isStreaming &&
-                        i === completedRounds.length - 1 &&
+                        i === completedRoundsWithLiveTools.length - 1 &&
                         !!round.toolExecution?.executions.some(
                           (e) => e.status === "pending" || e.status === "running"
                         )
@@ -1005,10 +1044,12 @@ const ModelResponseCard = memo(function ModelResponseCard({
                       isReasoningStreaming={response.isStreaming && !currentContent}
                     />
                   )}
-                  <StreamingStatusIndicator
-                    phase={streamingPhase}
-                    toolStatusMessage={toolStatusMessage}
-                  />
+                  {!hasVisibleRunningTools && (
+                    <StreamingStatusIndicator
+                      phase={streamingPhase}
+                      toolStatusMessage={toolStatusMessage}
+                    />
+                  )}
                 </div>
               );
             })()}

ui/src/components/ToolExecution/ToolExecutionStep.tsx

@@ -71,6 +71,29 @@ const TOOL_LANGUAGES: Record<string, string> = {
   chart_render: "json",
 };
 
+/** Extract a displayable code string from raw tool input (before artifacts are populated) */
+function extractCodeFromInput(input: unknown): string | null {
+  if (!input || typeof input !== "object") return null;
+  const obj = input as Record<string, unknown>;
+  if (typeof obj.code === "string") return obj.code;
+  if (typeof obj.sql === "string") return obj.sql;
+  if (typeof obj.query === "string") return obj.query;
+  if (obj.spec && typeof obj.spec === "object") {
+    try {
+      return JSON.stringify(obj.spec, null, 2);
+    } catch {
+      return null;
+    }
+  }
+  // Generic fallback: show all args as JSON
+  try {
+    const json = JSON.stringify(obj, null, 2);
+    return json === "{}" ? null : json;
+  } catch {
+    return null;
+  }
+}
+
 /** Format duration in human-readable form */
 function formatDuration(ms: number): string {
   if (ms < 1000) return `${ms}ms`;
@@ -127,17 +150,25 @@ function ToolExecutionStepComponent({
   );
   const hasOutputArtifacts = visibleOutputArtifacts.length > 0;
 
-  // Extract inline code preview from first code input artifact
+  // Extract inline code preview from first code input artifact, falling back to raw input
   const inlineCode = useMemo(() => {
     const codeArtifact = execution.inputArtifacts.find((a) => a.type === "code");
-    if (!codeArtifact) return null;
-    const code = getCodeFromArtifact(codeArtifact);
-    if (!code) return null;
-    const { preview, isTruncated } = getCodePreview(code, 4);
-    const language =
-      (codeArtifact.data as CodeArtifactData)?.language || TOOL_LANGUAGES[execution.toolName];
-    return { code, preview, isTruncated, artifact: codeArtifact, language };
-  }, [execution.inputArtifacts, execution.toolName]);
+    if (codeArtifact) {
+      const code = getCodeFromArtifact(codeArtifact);
+      if (code) {
+        const { preview, isTruncated } = getCodePreview(code, 4);
+        const language =
+          (codeArtifact.data as CodeArtifactData)?.language || TOOL_LANGUAGES[execution.toolName];
+        return { code, preview, isTruncated, artifact: codeArtifact, language };
+      }
+    }
+    // Fallback: extract from raw input (available immediately, before artifacts populate)
+    const rawCode = extractCodeFromInput(execution.input);
+    if (!rawCode) return null;
+    const { preview, isTruncated } = getCodePreview(rawCode, 4);
+    const language = TOOL_LANGUAGES[execution.toolName] || "json";
+    return { code: rawCode, preview, isTruncated, artifact: null, language };
+  }, [execution.inputArtifacts, execution.toolName, execution.input]);
 
   // Other input artifacts (non-code)
   const otherInputArtifacts = useMemo(
@@ -225,18 +256,20 @@ function ToolExecutionStepComponent({
                   {inlineCode.language || displayName}
                 </span>
                 <div className="flex-1" />
-                <button
-                  type="button"
-                  onClick={() => onArtifactClick?.(inlineCode.artifact)}
-                  className={cn(
-                    "p-0.5 rounded",
-                    "text-zinc-400 hover:text-zinc-700 dark:text-zinc-500 dark:hover:text-zinc-200",
-                    "opacity-0 group-hover/code:opacity-100 transition-opacity"
-                  )}
-                  aria-label="Expand"
-                >
-                  <Maximize2 className="h-3.5 w-3.5" />
-                </button>
+                {inlineCode.artifact && (
+                  <button
+                    type="button"
+                    onClick={() => onArtifactClick?.(inlineCode.artifact!)}
+                    className={cn(
+                      "p-0.5 rounded",
+                      "text-zinc-400 hover:text-zinc-700 dark:text-zinc-500 dark:hover:text-zinc-200",
+                      "opacity-0 group-hover/code:opacity-100 transition-opacity"
+                    )}
+                    aria-label="Expand"
+                  >
+                    <Maximize2 className="h-3.5 w-3.5" />
+                  </button>
+                )}
               </div>
 
               {/* Code content */}