Keep desktop AskUserQuestion flows from stalling in plan mode

Desktop chat was rendering AskUserQuestion twice: once as an inline question card and again as a generic permission request. That split also dropped the structured answers on the floor because the websocket permission_response shape only carried allow/deny state. This change keeps AskUserQuestion on the permission pipeline end-to-end. The desktop websocket contract now carries toolUseId and updatedInput, AskUserQuestion submits answers through permission_response, and the generic permission card is suppressed for that tool so the user sees a single question flow. Constraint: AskUserQuestion answers must round-trip through updatedInput.answers for the CLI tool contract to complete Rejected: Leave AskUserQuestion as a plain chat reply in desktop | the tool never receives structured answers and the pending approval UI remains stuck Confidence: high Scope-risk: narrow Reversibility: clean Directive: Keep AskUserQuestion bound to the permission-response path unless the desktop protocol grows a separate structured elicitation channel Tested: desktop lint; vitest src/components/chat/AskUserQuestion.test.tsx src/stores/chatStore.test.ts; bun test src/server/__tests__/conversations.test.ts Not-tested: Manual desktop click-through against a live plan-mode session after bundling
2026-07-17 13:13:35 +08:00 · 2026-04-20 22:27:00 +08:00 · 2026-04-20 22:27:00 +08:00 · 0d3933c2c9
commit 0d3933c2c9
parent 870cfe8b74
10 changed files with 327 additions and 28 deletions
--- a/desktop/src/components/chat/AskUserQuestion.test.tsx
+++ b/desktop/src/components/chat/AskUserQuestion.test.tsx
@ -0,0 +1,109 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { fireEvent, render, screen } from '@testing-library/react'
+
+const { sendMock } = vi.hoisted(() => ({
+  sendMock: vi.fn(),
+}))
+
+vi.mock('../../api/websocket', () => ({
+  wsManager: {
+    connect: vi.fn(),
+    disconnect: vi.fn(),
+    onMessage: vi.fn(() => () => {}),
+    clearHandlers: vi.fn(),
+    send: sendMock,
+  },
+}))
+
+vi.mock('../../api/sessions', () => ({
+  sessionsApi: {
+    getMessages: vi.fn(async () => ({ messages: [] })),
+    getSlashCommands: vi.fn(async () => ({ commands: [] })),
+  },
+}))
+
+import { AskUserQuestion } from './AskUserQuestion'
+import { useChatStore } from '../../stores/chatStore'
+import { useTabStore } from '../../stores/tabStore'
+
+const ACTIVE_TAB = 'active-tab'
+
+describe('AskUserQuestion', () => {
+  beforeEach(() => {
+    sendMock.mockReset()
+    useTabStore.setState({
+      activeTabId: ACTIVE_TAB,
+      tabs: [{ sessionId: ACTIVE_TAB, title: 'Test', type: 'session', status: 'idle' }],
+    })
+    useChatStore.setState({
+      sessions: {
+        [ACTIVE_TAB]: {
+          messages: [],
+          chatState: 'permission_pending',
+          connectionState: 'connected',
+          streamingText: '',
+          streamingToolInput: '',
+          activeToolUseId: null,
+          activeToolName: null,
+          activeThinkingId: null,
+          pendingPermission: {
+            requestId: 'perm-1',
+            toolName: 'AskUserQuestion',
+            toolUseId: 'tool-1',
+            input: {
+              questions: [
+                {
+                  question: 'Should we persist data?',
+                  options: [{ label: 'No' }, { label: 'Yes' }],
+                },
+              ],
+            },
+          },
+          pendingComputerUsePermission: null,
+          tokenUsage: { input_tokens: 0, output_tokens: 0 },
+          elapsedSeconds: 0,
+          statusVerb: '',
+          slashCommands: [],
+          agentTaskNotifications: {},
+          elapsedTimer: null,
+        },
+      },
+    })
+  })
+
+  it('submits answers through permission_response updatedInput instead of sending a chat message', () => {
+    render(
+      <AskUserQuestion
+        toolUseId="tool-1"
+        input={{
+          questions: [
+            {
+              question: 'Should we persist data?',
+              options: [{ label: 'No' }, { label: 'Yes' }],
+            },
+          ],
+        }}
+      />,
+    )
+
+    fireEvent.click(screen.getByRole('button', { name: /^No$/ }))
+    fireEvent.click(screen.getByRole('button', { name: /submit/i }))
+
+    expect(sendMock).toHaveBeenCalledWith(ACTIVE_TAB, {
+      type: 'permission_response',
+      requestId: 'perm-1',
+      allowed: true,
+      updatedInput: {
+        questions: [
+          {
+            question: 'Should we persist data?',
+            options: [{ label: 'No' }, { label: 'Yes' }],
+          },
+        ],
+        answers: {
+          'Should we persist data?': 'No',
+        },
+      },
+    })
+  })
+})
--- a/desktop/src/components/chat/AskUserQuestion.tsx
+++ b/desktop/src/components/chat/AskUserQuestion.tsx
@ -1,4 +1,4 @@
-import { useState, useRef } from 'react'
+import { useMemo, useRef, useState } from 'react'
 import { useChatStore } from '../../stores/chatStore'
 import { useTabStore } from '../../stores/tabStore'
 import { useTranslation } from '../../i18n'
@ -24,6 +24,7 @@ type AskUserInput = {
 type Props = {
  toolUseId: string
  input: unknown
+  result?: unknown
 }

 /**
@ -46,19 +47,41 @@ function parseInput(input: unknown): Question[] {
  return []
 }

-export function AskUserQuestion({ toolUseId: _toolUseId, input }: Props) {
-  const { sendMessage } = useChatStore()
+export function AskUserQuestion({ toolUseId, input, result }: Props) {
+  const { respondToPermission } = useChatStore()
  const activeTabId = useTabStore((s) => s.activeTabId)
+  const pendingPermission = useChatStore((s) => activeTabId ? s.sessions[activeTabId]?.pendingPermission : undefined)
  const t = useTranslation()
  const questions = parseInput(input)
+  const inputObject = (input && typeof input === 'object') ? input as Record<string, unknown> : {}
  const [activeTab, setActiveTab] = useState(0)
  const [selections, setSelections] = useState<Record<number, string>>({})
  const [freeText, setFreeText] = useState('')
-  const [submitted, setSubmitted] = useState(false)
+  const [hasSubmitted, setHasSubmitted] = useState(false)
  const composingRef = useRef(false)

  if (questions.length === 0) return null

+  const resultAnswers = useMemo(() => {
+    if (!result || typeof result !== 'object') return {}
+    const answers = (result as { answers?: unknown }).answers
+    return answers && typeof answers === 'object'
+      ? answers as Record<string, string>
+      : {}
+  }, [result])
+
+  const pendingRequest = pendingPermission?.toolUseId === toolUseId ? pendingPermission : null
+  const answeredText = useMemo(() => {
+    if (Object.keys(resultAnswers).length > 0) {
+      return questions
+        .map((question) => resultAnswers[question.question])
+        .filter((answer): answer is string => typeof answer === 'string' && answer.trim().length > 0)
+        .join(', ')
+    }
+    return freeText.trim() || Object.values(selections).join(', ')
+  }, [freeText, questions, resultAnswers, selections])
+  const submitted = Object.keys(resultAnswers).length > 0 || hasSubmitted
+
  const handleSelect = (qIndex: number, label: string) => {
    if (submitted) return
    setSelections((prev) => {
@ -84,9 +107,24 @@ export function AskUserQuestion({ toolUseId: _toolUseId, input }: Props) {
    const response = freeText.trim() || parts.join('; ') || ''
    if (!response) return

-    setSubmitted(true)
-    if (!activeTabId) return
-    sendMessage(activeTabId, response)
+    if (!activeTabId || !pendingRequest) return
+
+    const answers = questions.reduce<Record<string, string>>((acc, question, index) => {
+      if (freeText.trim()) {
+        acc[question.question] = freeText.trim()
+      } else if (selections[index]) {
+        acc[question.question] = selections[index]!
+      }
+      return acc
+    }, {})
+
+    setHasSubmitted(true)
+    respondToPermission(activeTabId, pendingRequest.requestId, true, {
+      updatedInput: {
+        ...inputObject,
+        answers,
+      },
+    })
  }

  // All questions must be answered (via selection or free text) to enable submit
@ -241,7 +279,7 @@ export function AskUserQuestion({ toolUseId: _toolUseId, input }: Props) {
          <div className="flex items-center gap-2 text-xs text-[var(--color-text-secondary)]">
            <span className="material-symbols-outlined text-[14px] text-[var(--color-success)]">check_circle</span>
            <span>
-              {t('question.answeredPrefix')}<strong>{freeText.trim() || Object.values(selections).join(', ')}</strong>
+              {t('question.answeredPrefix')}<strong>{answeredText}</strong>
            </span>
          </div>
        )}
@ -253,7 +291,7 @@ export function AskUserQuestion({ toolUseId: _toolUseId, input }: Props) {
          <Button
            variant="primary"
            size="sm"
-            disabled={!allAnswered}
+            disabled={!allAnswered || !pendingRequest}
            onClick={handleSubmit}
            icon={
              <span className="material-symbols-outlined text-[14px]">send</span>
--- a/desktop/src/components/chat/MessageList.tsx
+++ b/desktop/src/components/chat/MessageList.tsx
@ -210,6 +210,7 @@ export const MessageBlock = memo(function MessageBlock({
          <AskUserQuestion
            toolUseId={message.toolUseId}
            input={message.input}
+            result={toolResult?.content}
          />
        )
      }
--- a/desktop/src/components/chat/PermissionDialog.tsx
+++ b/desktop/src/components/chat/PermissionDialog.tsx
@ -237,7 +237,7 @@ export function PermissionDialog({ requestId, toolName, input, description }: Pr
          <Button
            variant="ghost"
            size="sm"
-            onClick={() => activeTabId && respondToPermission(activeTabId, requestId, true, 'always')}
+            onClick={() => activeTabId && respondToPermission(activeTabId, requestId, true, { rule: 'always' })}
            icon={
              <span className="material-symbols-outlined text-[14px]">verified</span>
            }
--- a/desktop/src/stores/chatStore.test.ts
+++ b/desktop/src/stores/chatStore.test.ts
@ -223,6 +223,74 @@ describe('chatStore history mapping', () => {
    ])
  })

+  it('keeps AskUserQuestion permission requests out of the message list while tracking the pending request', () => {
+    useChatStore.setState({
+      sessions: {
+        [TEST_SESSION_ID]: {
+          messages: [
+            {
+              id: 'ask-1',
+              type: 'tool_use',
+              toolName: 'AskUserQuestion',
+              toolUseId: 'tool-ask-1',
+              input: {
+                questions: [
+                  {
+                    question: 'Should we persist data?',
+                    options: [{ label: 'No' }, { label: 'Yes' }],
+                  },
+                ],
+              },
+              timestamp: 1,
+            },
+          ],
+          chatState: 'idle',
+          connectionState: 'connected',
+          streamingText: '',
+          streamingToolInput: '',
+          activeToolUseId: null,
+          activeToolName: null,
+          activeThinkingId: null,
+          pendingPermission: null,
+          pendingComputerUsePermission: null,
+          tokenUsage: { input_tokens: 0, output_tokens: 0 },
+          elapsedSeconds: 0,
+          statusVerb: '',
+          slashCommands: [],
+          agentTaskNotifications: {},
+          elapsedTimer: null,
+        },
+      },
+    })
+
+    useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
+      type: 'permission_request',
+      requestId: 'perm-ask-1',
+      toolName: 'AskUserQuestion',
+      toolUseId: 'tool-ask-1',
+      input: {
+        questions: [
+          {
+            question: 'Should we persist data?',
+            options: [{ label: 'No' }, { label: 'Yes' }],
+          },
+        ],
+      },
+    })
+
+    const session = useChatStore.getState().sessions[TEST_SESSION_ID]
+    expect(session?.pendingPermission).toMatchObject({
+      requestId: 'perm-ask-1',
+      toolName: 'AskUserQuestion',
+      toolUseId: 'tool-ask-1',
+    })
+    expect(session?.messages).toHaveLength(1)
+    expect(session?.messages[0]).toMatchObject({
+      type: 'tool_use',
+      toolUseId: 'tool-ask-1',
+    })
+  })
+
  it('sends permission mode updates to the active session only', () => {
    useChatStore.getState().setSessionPermissionMode('nonexistent-session', 'acceptEdits')
    expect(sendMock).not.toHaveBeenCalled()
--- a/desktop/src/stores/chatStore.ts
+++ b/desktop/src/stores/chatStore.ts
@ -35,6 +35,7 @@ export type PerSessionState = {
  pendingPermission: {
    requestId: string
    toolName: string
+    toolUseId?: string
    input: unknown
    description?: string
  } | null
@ -80,7 +81,15 @@ type ChatStore = {
  connectToSession: (sessionId: string) => void
  disconnectSession: (sessionId: string) => void
  sendMessage: (sessionId: string, content: string, attachments?: AttachmentRef[]) => void
-  respondToPermission: (sessionId: string, requestId: string, allowed: boolean, rule?: string) => void
+  respondToPermission: (
+    sessionId: string,
+    requestId: string,
+    allowed: boolean,
+    options?: {
+      rule?: string
+      updatedInput?: Record<string, unknown>
+    },
+  ) => void
  respondToComputerUsePermission: (
    sessionId: string,
    requestId: string,
@ -280,8 +289,14 @@ export const useChatStore = create<ChatStore>((set, get) => ({
    wsManager.send(sessionId, { type: 'user_message', content, attachments })
  },

-  respondToPermission: (sessionId, requestId, allowed, rule?) => {
-    wsManager.send(sessionId, { type: 'permission_response', requestId, allowed, ...(rule ? { rule } : {}) })
+  respondToPermission: (sessionId, requestId, allowed, options) => {
+    wsManager.send(sessionId, {
+      type: 'permission_response',
+      requestId,
+      allowed,
+      ...(options?.rule ? { rule: options.rule } : {}),
+      ...(options?.updatedInput ? { updatedInput: options.updatedInput } : {}),
+    })
    set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ pendingPermission: null, chatState: allowed ? 'tool_executing' : 'idle' })) }))
  },

@ -493,14 +508,29 @@ export const useChatStore = create<ChatStore>((set, get) => ({

      case 'permission_request':
        update((s) => ({
-          pendingPermission: { requestId: msg.requestId, toolName: msg.toolName, input: msg.input, description: msg.description },
+          pendingPermission: {
+            requestId: msg.requestId,
+            toolName: msg.toolName,
+            toolUseId: msg.toolUseId,
+            input: msg.input,
+            description: msg.description,
+          },
          pendingComputerUsePermission: null,
          chatState: 'permission_pending',
          activeThinkingId: null,
-          messages: [...s.messages, {
-            id: nextId(), type: 'permission_request', requestId: msg.requestId,
-            toolName: msg.toolName, input: msg.input, description: msg.description, timestamp: Date.now(),
-          }],
+          messages:
+            msg.toolName === 'AskUserQuestion'
+              ? s.messages
+              : [...s.messages, {
+                  id: nextId(),
+                  type: 'permission_request',
+                  requestId: msg.requestId,
+                  toolName: msg.toolName,
+                  toolUseId: msg.toolUseId,
+                  input: msg.input,
+                  description: msg.description,
+                  timestamp: Date.now(),
+                }],
        }))
        break

--- a/desktop/src/types/chat.ts
+++ b/desktop/src/types/chat.ts
@ -6,7 +6,13 @@ import type { PermissionMode } from './settings'

 export type ClientMessage =
  | { type: 'user_message'; content: string; attachments?: AttachmentRef[] }
-  | { type: 'permission_response'; requestId: string; allowed: boolean; rule?: string }
+  | {
+      type: 'permission_response'
+      requestId: string
+      allowed: boolean
+      rule?: string
+      updatedInput?: Record<string, unknown>
+    }
  | {
      type: 'computer_use_permission_response'
      requestId: string
@ -39,7 +45,14 @@ export type ServerMessage =
  | { type: 'content_delta'; text?: string; toolInput?: string }
  | { type: 'tool_use_complete'; toolName: string; toolUseId: string; input: unknown; parentToolUseId?: string }
  | { type: 'tool_result'; toolUseId: string; content: unknown; isError: boolean; parentToolUseId?: string }
-  | { type: 'permission_request'; requestId: string; toolName: string; input: unknown; description?: string }
+  | {
+      type: 'permission_request'
+      requestId: string
+      toolName: string
+      toolUseId?: string
+      input: unknown
+      description?: string
+    }
  | {
      type: 'computer_use_permission_request'
      requestId: string
@ -147,6 +160,15 @@ export type UIMessage =
  | { id: string; type: 'tool_use'; toolName: string; toolUseId: string; input: unknown; timestamp: number; parentToolUseId?: string }
  | { id: string; type: 'tool_result'; toolUseId: string; content: unknown; isError: boolean; timestamp: number; parentToolUseId?: string }
  | { id: string; type: 'system'; content: string; timestamp: number }
-  | { id: string; type: 'permission_request'; requestId: string; toolName: string; input: unknown; description?: string; timestamp: number }
+  | {
+      id: string
+      type: 'permission_request'
+      requestId: string
+      toolName: string
+      toolUseId?: string
+      input: unknown
+      description?: string
+      timestamp: number
+    }
  | { id: string; type: 'error'; message: string; code: string; timestamp: number }
  | { id: string; type: 'task_summary'; tasks: TaskSummaryItem[]; timestamp: number }
--- a/src/server/services/conversationService.ts
+++ b/src/server/services/conversationService.ts
@ -156,10 +156,7 @@ export class ConversationService {
    this.readErrorStream(sessionId, proc)

    proc.exited.then((code) => {
-      console.log(
-        `[ConversationService] CLI process for ${sessionId} exited with code ${code}`,
-      )
-      this.sessions.delete(sessionId)
+      this.handleProcessExit(sessionId, proc, code)
    })

    const STARTUP_GRACE_MS = 3000
@ -232,6 +229,7 @@ export class ConversationService {
    requestId: string,
    allowed: boolean,
    rule?: string,
+    updatedInput?: Record<string, unknown>,
  ): boolean {
    const session = this.sessions.get(sessionId)
    const pendingRequest = session?.pendingPermissionRequests.get(requestId)
@ -247,7 +245,7 @@ export class ConversationService {
        response: allowed
          ? {
              behavior: 'allow',
-              updatedInput: {},
+              updatedInput: updatedInput ?? {},
              ...(rule === 'always' && pendingRequest
                ? {
                    updatedPermissions: [
@ -440,6 +438,21 @@ export class ConversationService {
    return true
  }

+  private handleProcessExit(
+    sessionId: string,
+    proc: SessionProcess['proc'],
+    code: number,
+  ): void {
+    console.log(
+      `[ConversationService] CLI process for ${sessionId} exited with code ${code}`,
+    )
+
+    const activeSession = this.sessions.get(sessionId)
+    if (activeSession?.proc === proc) {
+      this.sessions.delete(sessionId)
+    }
+  }
+
  private getPermissionArgs(
    mode: string | undefined,
    dangerousMode: boolean,
--- a/src/server/ws/events.ts
+++ b/src/server/ws/events.ts
@ -10,7 +10,13 @@

 export type ClientMessage =
  | { type: 'user_message'; content: string; attachments?: AttachmentRef[] }
-  | { type: 'permission_response'; requestId: string; allowed: boolean; rule?: string }
+  | {
+      type: 'permission_response'
+      requestId: string
+      allowed: boolean
+      rule?: string
+      updatedInput?: Record<string, unknown>
+    }
  | {
      type: 'computer_use_permission_response'
      requestId: string
@ -38,7 +44,14 @@ export type ServerMessage =
  | { type: 'content_delta'; text?: string; toolInput?: string }
  | { type: 'tool_use_complete'; toolName: string; toolUseId: string; input: unknown; parentToolUseId?: string }
  | { type: 'tool_result'; toolUseId: string; content: unknown; isError: boolean; parentToolUseId?: string }
-  | { type: 'permission_request'; requestId: string; toolName: string; input: unknown; description?: string }
+  | {
+      type: 'permission_request'
+      requestId: string
+      toolName: string
+      toolUseId?: string
+      input: unknown
+      description?: string
+    }
  | {
      type: 'computer_use_permission_request'
      requestId: string
--- a/src/server/ws/handler.ts
+++ b/src/server/ws/handler.ts
@ -302,6 +302,7 @@ function handlePermissionResponse(
    message.requestId,
    message.allowed,
    message.rule,
+    message.updatedInput,
  )
  console.log(`[WS] Permission response for ${message.requestId}: ${message.allowed}`)
 }
@ -693,6 +694,10 @@ function translateCliMessage(cliMsg: any, sessionId: string): ServerMessage[] {
          type: 'permission_request',
          requestId: cliMsg.request_id,
          toolName: cliMsg.request.tool_name || 'Unknown',
+          toolUseId:
+            typeof cliMsg.request.tool_use_id === 'string'
+              ? cliMsg.request.tool_use_id
+              : undefined,
          input: cliMsg.request.input || {},
          description: cliMsg.request.description,
        }]