🤖 fix: keep ask_user_question waiting across restart (#1152)

Author: ThomasK33

src/browser/App.tsx

@@ -416,14 +416,21 @@ function AppInner() {
 
       // Compute streaming models here (only when command palette opens)
       const allStates = workspaceStore.getAllStates();
+      const selectedWorkspaceState = params.selectedWorkspace
+        ? (allStates.get(params.selectedWorkspace.workspaceId) ?? null)
+        : null;
       const streamingModels = new Map<string, string>();
       for (const [workspaceId, state] of allStates) {
         if (state.canInterrupt && state.currentModel) {
           streamingModels.set(workspaceId, state.currentModel);
         }
       }
 
-      const factories = buildCoreSources({ ...params, streamingModels });
+      const factories = buildCoreSources({
+        ...params,
+        streamingModels,
+        selectedWorkspaceState,
+      });
       const actions: CommandAction[] = [];
       for (const factory of factories) {
         actions.push(...factory());

src/browser/components/tools/AskUserQuestionToolCall.tsx

@@ -2,6 +2,7 @@ import assert from "@/common/utils/assert";
 
 import { useMemo, useState } from "react";
 
+import { CUSTOM_EVENTS, createCustomEvent } from "@/common/constants/events";
 import { useAPI } from "@/browser/contexts/API";
 import { Checkbox } from "@/browser/components/ui/checkbox";
 import { Input } from "@/browser/components/ui/input";
@@ -222,6 +223,7 @@ export function AskUserQuestionToolCall(props: {
     setSubmitError(null);
 
     let answers: Record<string, string>;
+    let workspaceId: string;
 
     try {
       answers = {};
@@ -237,6 +239,7 @@ export function AskUserQuestionToolCall(props: {
 
       assert(api, "API not connected");
       assert(props.workspaceId, "workspaceId is required");
+      workspaceId = props.workspaceId;
     } catch (error) {
       const errorMessage = error instanceof Error ? error.message : String(error);
       setSubmitError(errorMessage);
@@ -246,14 +249,24 @@ export function AskUserQuestionToolCall(props: {
 
     api.workspace
       .answerAskUserQuestion({
-        workspaceId: props.workspaceId,
+        workspaceId,
         toolCallId: props.toolCallId,
         answers,
       })
       .then((result) => {
         if (!result.success) {
           setSubmitError(result.error);
+          return;
         }
+
+        // If the stream was interrupted (e.g. app restart) we need to explicitly
+        // kick the resume manager so the assistant continues after answers.
+        window.dispatchEvent(
+          createCustomEvent(CUSTOM_EVENTS.RESUME_CHECK_REQUESTED, {
+            workspaceId,
+            isManual: true,
+          })
+        );
       })
       .catch((error) => {
         const errorMessage = error instanceof Error ? error.message : String(error);

src/browser/hooks/useAIViewKeybinds.ts

@@ -67,6 +67,12 @@ export function useAIViewKeybinds({
       // Interrupt stream: Ctrl+C in vim mode, Esc in normal mode
       // Only intercept if actively compacting (otherwise allow browser default for copy in vim mode)
       if (matchesKeybind(e, interruptKeybind)) {
+        // ask_user_question is a special waiting state: don't interrupt it with Esc/Ctrl+C.
+        // Users can still respond via the questions UI, or type in chat to cancel.
+        if (aggregator?.hasAwaitingUserQuestion()) {
+          return;
+        }
+
         if (canInterrupt && aggregator && isCompactingStream(aggregator)) {
           // Ctrl+C during compaction: restore original state and enter edit mode
           // Stores cancellation marker in localStorage (persists across reloads)

src/browser/utils/commands/sources.ts

@@ -9,6 +9,7 @@ import { CommandIds } from "@/browser/utils/commandIds";
 import type { ProjectConfig } from "@/node/config";
 import type { FrontendWorkspaceMetadata } from "@/common/types/workspace";
 import type { BranchListResult } from "@/common/orpc/types";
+import type { WorkspaceState } from "@/browser/stores/WorkspaceStore";
 import type { RuntimeConfig } from "@/common/types/runtime";
 
 export interface BuildSourcesParams {
@@ -17,6 +18,7 @@ export interface BuildSourcesParams {
   /** Map of workspace ID to workspace metadata (keyed by metadata.id, not path) */
   workspaceMetadata: Map<string, FrontendWorkspaceMetadata>;
   theme: ThemeMode;
+  selectedWorkspaceState?: WorkspaceState | null;
   selectedWorkspace: {
     projectPath: string;
     projectName: string;
@@ -393,6 +395,9 @@ export function buildCoreSources(p: BuildSourcesParams): Array<() => CommandActi
         title: "Interrupt Streaming",
         section: section.chat,
         run: async () => {
+          if (p.selectedWorkspaceState?.awaitingUserQuestion) {
+            return;
+          }
           await p.api?.workspace.interruptStream({ workspaceId: id });
         },
       });

src/browser/utils/messages/StreamingMessageAggregator.status.test.ts

@@ -46,6 +46,55 @@ afterAll(() => {
   }
 });
 
+describe("ask_user_question waiting state", () => {
+  it("treats partial ask_user_question as executing (waiting) not interrupted", () => {
+    const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z");
+
+    aggregator.loadHistoricalMessages([
+      {
+        id: "assistant-1",
+        role: "assistant" as const,
+        parts: [
+          {
+            type: "dynamic-tool" as const,
+            toolCallId: "call-ask-1",
+            toolName: "ask_user_question",
+            state: "input-available" as const,
+            input: {
+              questions: [
+                {
+                  header: "Approach",
+                  question: "Which approach should we take?",
+                  options: [
+                    { label: "A", description: "Approach A" },
+                    { label: "B", description: "Approach B" },
+                  ],
+                  multiSelect: false,
+                },
+              ],
+            },
+          },
+        ],
+        metadata: {
+          timestamp: 1000,
+          historySequence: 1,
+          partial: true,
+        },
+      },
+    ]);
+
+    const displayed = aggregator.getDisplayedMessages();
+    const toolMsg = displayed.find((m) => m.type === "tool" && m.toolName === "ask_user_question");
+    expect(toolMsg).toBeDefined();
+    if (toolMsg?.type === "tool") {
+      expect(toolMsg.status).toBe("executing");
+      expect(toolMsg.isPartial).toBe(true);
+    }
+
+    expect(aggregator.hasAwaitingUserQuestion()).toBe(true);
+  });
+});
+
 describe("StreamingMessageAggregator - Agent Status", () => {
   it("should start with undefined agent status", () => {
     const aggregator = new StreamingMessageAggregator("2024-01-01T00:00:00.000Z");

src/browser/utils/messages/StreamingMessageAggregator.ts

@@ -299,18 +299,19 @@ export class StreamingMessageAggregator {
    * Used to show "Awaiting your input" instead of "streaming..." in the UI.
    */
   hasAwaitingUserQuestion(): boolean {
-    // Scan displayed messages for an ask_user_question tool in "executing" state
+    // Only treat the workspace as "awaiting input" when the *latest* displayed
+    // message is an executing ask_user_question tool.
+    //
+    // This avoids false positives from stale historical partials if the user
+    // continued the chat after skipping/canceling the questions.
     const displayed = this.getDisplayedMessages();
-    for (const msg of displayed) {
-      if (
-        msg.type === "tool" &&
-        msg.toolName === "ask_user_question" &&
-        msg.status === "executing"
-      ) {
-        return true;
-      }
+    const last = displayed[displayed.length - 1];
+
+    if (last?.type !== "tool") {
+      return false;
     }
-    return false;
+
+    return last.toolName === "ask_user_question" && last.status === "executing";
   }
 
   /**
@@ -1095,10 +1096,18 @@ export class StreamingMessageAggregator {
               if (part.state === "output-available") {
                 // Check if result indicates failure (for tools that return { success: boolean })
                 status = hasFailureResult(part.output) ? "failed" : "completed";
-              } else if (part.state === "input-available" && message.metadata?.partial) {
-                status = "interrupted";
               } else if (part.state === "input-available") {
-                status = "executing";
+                // Most unfinished tool calls in partial messages represent an interruption.
+                // ask_user_question is different: it's intentionally waiting on user input,
+                // so after restart we should keep it answerable ("executing") instead of
+                // showing retry/auto-resume UX.
+                if (part.toolName === "ask_user_question") {
+                  status = "executing";
+                } else if (message.metadata?.partial) {
+                  status = "interrupted";
+                } else {
+                  status = "executing";
+                }
               } else {
                 status = "pending";
               }

src/browser/utils/messages/messageUtils.test.ts

@@ -1,7 +1,46 @@
 import { describe, it, expect } from "@jest/globals";
+
+import { shouldShowInterruptedBarrier } from "./messageUtils";
 import { mergeConsecutiveStreamErrors, computeBashOutputGroupInfo } from "./messageUtils";
 import type { DisplayedMessage } from "@/common/types/message";
 
+describe("shouldShowInterruptedBarrier", () => {
+  it("returns false for executing ask_user_question", () => {
+    const msg: DisplayedMessage = {
+      type: "tool",
+      id: "tool-1",
+      historyId: "assistant-1",
+      toolName: "ask_user_question",
+      toolCallId: "call-1",
+      args: { questions: [] },
+      status: "executing",
+      isPartial: true,
+      historySequence: 2,
+      streamSequence: 0,
+      isLastPartOfMessage: true,
+    };
+
+    expect(shouldShowInterruptedBarrier(msg)).toBe(false);
+  });
+
+  it("returns true for interrupted tool (non ask_user_question)", () => {
+    const msg: DisplayedMessage = {
+      type: "tool",
+      id: "tool-1",
+      historyId: "assistant-1",
+      toolName: "bash",
+      toolCallId: "call-1",
+      args: { script: "echo hi", timeout_secs: 1, display_name: "test" },
+      status: "interrupted",
+      isPartial: true,
+      historySequence: 2,
+      streamSequence: 0,
+      isLastPartOfMessage: true,
+    };
+
+    expect(shouldShowInterruptedBarrier(msg)).toBe(true);
+  });
+});
 describe("mergeConsecutiveStreamErrors", () => {
   it("returns empty array for empty input", () => {
     const result = mergeConsecutiveStreamErrors([]);

src/browser/utils/messages/messageUtils.ts

@@ -54,6 +54,13 @@ export function shouldShowInterruptedBarrier(msg: DisplayedMessage): boolean {
   )
     return false;
 
+  // ask_user_question is intentionally a "waiting for input" state. Even if the
+  // underlying message is a persisted partial (e.g. after app restart), we keep
+  // it answerable instead of showing "Interrupted".
+  if (msg.type === "tool" && msg.toolName === "ask_user_question" && msg.status === "executing") {
+    return false;
+  }
+
   // Only show on the last part of multi-part messages
   if (!msg.isLastPartOfMessage) return false;
 

src/browser/utils/messages/retryEligibility.test.ts

@@ -59,6 +59,32 @@ describe("hasInterruptedStream", () => {
     expect(hasInterruptedStream(messages)).toBe(true);
   });
 
+  it("returns false for executing ask_user_question (waiting state)", () => {
+    const messages: DisplayedMessage[] = [
+      {
+        type: "user",
+        id: "user-1",
+        historyId: "user-1",
+        content: "Hello",
+        historySequence: 1,
+      },
+      {
+        type: "tool",
+        id: "tool-1",
+        historyId: "assistant-1",
+        toolName: "ask_user_question",
+        toolCallId: "call-1",
+        args: { questions: [] },
+        status: "executing",
+        isPartial: true,
+        historySequence: 2,
+        streamSequence: 0,
+        isLastPartOfMessage: true,
+      },
+    ];
+
+    expect(hasInterruptedStream(messages)).toBe(false);
+  });
   it("returns true for partial tool message", () => {
     const messages: DisplayedMessage[] = [
       {

src/browser/utils/messages/retryEligibility.ts

@@ -90,6 +90,19 @@ export function hasInterruptedStream(
 
   const lastMessage = messages[messages.length - 1];
 
+  // ask_user_question is a special case: an unfinished tool call represents an
+  // intentional "waiting for user input" state, not a stream interruption.
+  //
+  // Treating it as interrupted causes RetryBarrier + auto-resume to fire on app
+  // restart, which re-runs the LLM call and re-asks the questions.
+  if (
+    lastMessage.type === "tool" &&
+    lastMessage.toolName === "ask_user_question" &&
+    lastMessage.status === "executing"
+  ) {
+    return false;
+  }
+
   return (
     lastMessage.type === "stream-error" || // Stream errored out (show UI for ALL error types)
     lastMessage.type === "user" || // No response received yet (app restart during slow model)