Skip to content

Commit 6fcca8b

Browse files
authored
🤖 feat: add ask_user_question interactive plan-mode tool (#1121)
Adds a new interactive plan-mode tool `ask_user_question` (inspired by Claude Code's `AskUserQuestion`) and a corresponding inline UI for answering multiple-choice questions. Key behaviors: - Tool is only registered/available in **plan mode**. - While the tool is pending, the UI shows an inline multi-question form. - If the user types a normal chat message instead, the tool is canceled (recorded as a failed tool call via `tool-error`) and the message is sent as a real user message. - Tool args/results are inferred from Zod schemas to avoid drift. Validation: - `make static-check` --- <details> <summary>📋 Implementation Plan</summary> # 🤖 Plan: Add an `AskUserQuestion`-style tool to Mux (Plan Mode UX) ## Goal Implement a tool compatible with Claude Code’s `AskUserQuestion` semantics (per `cc-plan-mode.md`) so that, **while staying in plan mode**, the assistant can ask the user 1–4 multiple-choice questions and receive structured answers. Key UX requirement: the question UI should be **non-invasive** and **optional**—the user is encouraged to answer via the form, but they should also be able to just type a normal chat reply without having to explicitly “close/deny” the UI. ## Source-of-truth behavior to mirror (from `cc-plan-mode.md`) We should match the Claude Code *behavior + schema* as closely as practical, while using Mux’s snake_case naming convention: - **Tool name (CC):** `AskUserQuestion` - **Tool name (Mux):** `ask_user_question` - **Input:** `{ questions: Question[], answers?: Record<string,string> }` - **Output:** `{ questions: Question[], answers: Record<string,string> }` - answers map: **question text → answer string** - multi-select answers are **comma-separated** - Constraints: - questions: 1–4 - options per question: 2–4 - question texts unique; option labels unique within a question - “Other” is always available (but **not** included in `options` input) ## What already exists in Mux (relevant implementation hooks) - Tool call rendering is centralized in `src/browser/components/Messages/ToolMessage.tsx`, which routes known tools to custom React components; otherwise it falls back to `GenericToolCall`. - Tool UI primitives are in `src/browser/components/tools/shared/ToolPrimitives.tsx` + `toolUtils.tsx`. - We already have reusable UI building blocks: - `src/browser/components/ui/toggle-group.tsx` (Radix ToggleGroup; supports `type="single"|"multiple"`) - `src/browser/components/ui/checkbox.tsx` - `src/browser/components/ui/input.tsx` - `src/browser/components/ui/button.tsx` - Tool definitions live in `src/common/utils/tools/toolDefinitions.ts` (Zod schema + description); backend tools reference this. - Backend tools are in `src/node/services/tools/*` and are registered through `src/common/utils/tools/tools.ts` (`getToolsForModel`). - Streaming tool calls are represented as `dynamic-tool` parts; while a tool is running it’s `state: "input-available"` and UI status becomes `"executing"`. ## Recommended approach (single approach; net new product LoC ~500–900) ### Summary Add a new **interactive tool** `ask_user_question` (CC-compatible `AskUserQuestion`) that: 1. **Pauses** the current assistant stream until answers are provided. 2. Renders an **inline, collapsible** multiple-choice form in the chat (like other tool calls). 3. Lets the user **either**: - submit answers via the form (structured), or - type a normal chat message, which **cancels** the question tool call and is sent as a real user message (no “deny/close” step). ### Backend design #### 1) Tool definition (shared) - Add `ask_user_question` to `TOOL_DEFINITIONS` with a Zod schema matching `cc-plan-mode.md`: - Prefer defining and exporting named sub-schemas (e.g. `AskUserQuestionOptionSchema`, `AskUserQuestionQuestionSchema`) so the result schema + UI types can reuse them. - `Option`: `{ label: string; description: string }` - `Question`: `{ question: string; header: string; options: Option[]; multiSelect: boolean }` - Input: `{ questions: Question[]; answers?: Record<string,string> }` + refine uniqueness - Output: `{ questions: Question[]; answers: Record<string,string> }` TypeScript types: **infer from Zod** (avoid handwritten interfaces) to prevent drift. - Args: - `export type AskUserQuestionToolArgs = z.infer<typeof TOOL_DEFINITIONS.ask_user_question.schema>;` - Result: - Define a shared `AskUserQuestionToolResultSchema` (Zod) alongside the input schema (reusing the same `Question`/`Option` schemas), then: - `export type AskUserQuestionToolResult = z.infer<typeof AskUserQuestionToolResultSchema>;` - (Optional) also export `AskUserQuestionQuestion` / `AskUserQuestionOption` via `z.infer<...>` for the UI reducer. #### 2) Pending-question manager (new node service) Create an in-memory manager owned by the backend process (NOT persisted): - `AskUserQuestionManager` - `registerPending(workspaceId, toolCallId, questions): Promise<Record<string,string>>` - `answer(workspaceId, toolCallId, answers): void` - `cancel(workspaceId, toolCallId, reason): void` - `getLatestPending(workspaceId): { toolCallId, questions } | null` Important defensive behavior: - Assert toolCallIds are unique per workspace. - Always delete pending entries on resolve/reject. - Add a long timeout (e.g., 30 minutes) to avoid leaks if the UI never responds. - Support abort: if the stream is interrupted, tool execution must unblock (reject) so interrupt actually works. #### 3) Tool implementation (`src/node/services/tools/ask_user_question.ts`) Implement `createAskUserQuestionTool(config)` using `tool()` from `ai`. Execution flow: - Validate args using the shared Zod schema. - Determine `workspaceId` from `config.workspaceId` (assert present). - Determine `toolCallId` from `ToolCallOptions.toolCallId`. - Register the pending question with `AskUserQuestionManager` and **await** answers. - Return `{ questions, answers }` (matching Claude Code’s output schema). Abort/cancel handling: - Race the awaited promise against an abort signal if available. - On abort/interrupt/**cancel** (including “user responded in chat” cancellation), **throw** an Error so the AI SDK emits a `tool-error` part (Mux will record this as a failed tool result). Model-context parity (optional but recommended): - Mux providers will see tool output as structured JSON by default. - To be more CC-like (and more token-efficient), add a small transform in `applyToolOutputRedaction(...)` to replace the provider-facing output for `ask_user_question` with a short string summary like: - `User has answered your questions: "<Q>"="<A>", ... You can now continue with the user's answers in mind.` - while keeping the full structured `{ questions, answers }` persisted for UI. #### 4) ORPC endpoint to submit answers Add a workspace-scoped ORPC endpoint: - `workspace.answerAskUserQuestion` - input: `{ workspaceId: string; toolCallId: string; answers: Record<string,string> }` - output: `Result<void, string>` Implementation: - Route → `WorkspaceService` → `AskUserQuestionManager.answer(...)`. #### 5) “Chat message instead of clicking” behavior (non-invasive requirement) Requirement (confirmed): if the user just types a normal chat reply while `ask_user_question` is pending, we should **not** treat it as tool input. Instead we should: 1. **Cancel the tool call** (so it shows as failed with a message like “User responded in chat; questions canceled”). 2. Send the user’s text as a **real user message** (normal send/queue semantics). Implementation sketch: - Add `AskUserQuestionManager.getLatestPending(workspaceId)` (or `getPendingByToolCallId`) that returns the active pending prompt. - Modify `WorkspaceService.sendMessage()` behavior: - if `aiService.isStreaming(workspaceId)` and a pending `ask_user_question` exists: - call `AskUserQuestionManager.cancel(workspaceId, toolCallId, "User responded in chat; questions canceled")` - then proceed with existing behavior (queue message via `session.queueMessage(message, options)`) - return `Ok()` Effect: - The stream unblocks because the tool throws → AI SDK emits `tool-error` → Mux records a failed tool result. - The user’s message remains a normal user message and will be processed next (queued during streaming, then auto-sent on tool completion / stream end, per existing queue behavior). ### Frontend design #### 1) Add a custom tool renderer Add `AskUserQuestionToolCall.tsx` under `src/browser/components/tools/` and route it in `ToolMessage.tsx` similarly to `ProposePlanToolCall`. Props needed: - `args`, `result`, `status`, `workspaceId` #### 2) UI behavior (inline form) Render as an inline tool card, expanded by default while awaiting answers: - Header: `ask_user_question` + status - Body: - Tab bar for 1–4 questions (mirroring CC’s chips): show each `header`, with a checkmark for answered. - Per-question view: - question text - options: - single-select → `ToggleGroup type="single"` - multi-select → `Checkbox` list or `ToggleGroup type="multiple"` - “Other”: - option always available - selecting “Other” reveals an `Input` for custom text - Submit view: - show all answers in a review list - `Submit answers` button → calls `api.workspace.answerAskUserQuestion({ workspaceId, toolCallId, answers })` State management: - Use a reducer similar to CC’s (`currentQuestionIndex`, `answers`, `questionStates`, `isInTextInput`). Accessibility/keyboard (follow-up, but planned): - Basic: click-based + tab focus. - Optional parity: left/right arrow to navigate question tabs; enter to select. #### 3) Display after completion When `result` is present: - Show a compact “User answered …” summary (like CC) - Keep details collapsible. #### 4) Non-invasive encouragement While the tool is pending: - Add a small helper line in the tool body: “Tip: you can also just type a message to respond in chat (this will cancel these questions).” (We can later refine this into a better “waiting for input” streaming banner, but it’s not required for v1.) ### Testing / validation (net new test LoC ~250–450) - Unit tests for `AskUserQuestionManager`: - registers + answers - cancel/timeout cleanup - rejects on abort - Backend test for ORPC endpoint (or WorkspaceService method): answering resolves the tool. - Storybook: - Add a chat story with a pending `ask_user_question` tool call - Add a story with completed answers ## Alternatives considered ### A) Non-blocking “suggestion card” tool (net new product LoC ~300–600) Have `ask_user_question` return immediately and not pause the stream; user answers later and the answer is posted as a normal user message. Downside: this diverges from Claude Code semantics and won’t behave like the tool LLMs were trained on (tool answers won’t be available in the same turn). ## Decisions / requirements confirmed - If the user types a chat message while `ask_user_question` is pending: - mark the tool call as **failed** with an output indicating the user canceled by responding in chat - then treat the typed text as a **real user message** (normal queue/send behavior) - Tool name in Mux should be `ask_user_question` (snake_case). - No explicit “Skip” button required; a small helper label is enough. - No feature flags: include this tool + UI by default in **plan mode**. ## Availability / gating - Register `ask_user_question` only when `SendMessageOptions.mode === "plan"` (so it is available by default in plan mode, and not callable in exec mode). - Ensure `AIService.streamMessage()` passes `mode` into the *early* `getToolsForModel(...)` call used for the mode-transition “Available tools: …” sentinel. Otherwise plan mode would omit `ask_user_question` from the advertised tool list even though it’s registered. - UI: always render custom tool UI for `ask_user_question` tool parts (no toggle). --- **Net new product code estimate (recommended approach):** ~500–900 LoC </details> --- _Generated with `mux` • Model: `anthropic:claude-opus-4-5` • Thinking: `high`_ --------- Signed-off-by: Thomas Kosiewski <tk@coder.com>
1 parent 3e9740e commit 6fcca8b

23 files changed

+1306
-67
lines changed

docs/plan-mode.mdx

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,34 @@ This means you can make edits in your preferred editor, return to mux, send a me
2929

3030
## Plan File Location
3131

32-
Plans are stored in a dedicated directory:
32+
Plans are stored in a dedicated directory under your Mux home:
3333

3434
```
35-
~/.mux/plans/<workspace-id>.md
35+
~/.mux/plans/<project>/<workspace-name>.md
3636
```
3737

38-
The file is created when the agent first writes a plan and persists across sessions.
38+
Notes:
39+
40+
- `<workspace-name>` includes the random suffix (e.g. `feature-x7k2`), so it’s globally unique with high probability.
41+
42+
## ask_user_question (Plan Mode Only)
43+
44+
In plan mode, the agent may call `ask_user_question` to ask up to 4 structured multiple-choice questions when it needs clarification before finalizing a plan.
45+
46+
What you’ll see:
47+
48+
- An inline “tool call card” in the chat with a small form (single-select or multi-select).
49+
- An always-available **Other** option for free-form answers.
50+
51+
How to respond:
52+
53+
- **Recommended:** answer in the form and click **Submit answers**.
54+
- **Optional:** you can also just type a normal chat message. This will **cancel** the pending `ask_user_question` tool call and your message will be sent as a regular chat message.
55+
56+
Availability:
57+
58+
- `ask_user_question` is only registered for the agent in **Plan Mode**.
59+
- In Exec Mode, the agent cannot call `ask_user_question`.
3960

4061
## UI Features
4162

@@ -55,7 +76,7 @@ User: "Add user authentication to the app"
5576
5677
┌─────────────────────────────────────┐
5778
│ Agent reads codebase, writes plan │
58-
│ to ~/.mux/plans/<id>.md
79+
│ to ~/.mux/plans/<project>/<name>.md│
5980
└─────────────────────────────────────┘
6081
6182

src/browser/components/AIView.tsx

Lines changed: 26 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,8 @@ const AIViewInner: React.FC<AIViewProps> = ({
160160
useEffect(() => {
161161
workspaceStateRef.current = workspaceState;
162162
}, [workspaceState]);
163-
const { messages, canInterrupt, isCompacting, loading, currentModel } = workspaceState;
163+
const { messages, canInterrupt, isCompacting, awaitingUserQuestion, loading, currentModel } =
164+
workspaceState;
164165

165166
// Apply message transformations:
166167
// 1. Merge consecutive identical stream errors
@@ -673,24 +674,34 @@ const AIViewInner: React.FC<AIViewProps> = ({
673674
{canInterrupt && (
674675
<StreamingBarrier
675676
statusText={
676-
isCompacting
677-
? currentModel
678-
? `${getModelName(currentModel)} compacting...`
679-
: "compacting..."
680-
: currentModel
681-
? `${getModelName(currentModel)} streaming...`
682-
: "streaming..."
677+
awaitingUserQuestion
678+
? "Awaiting your input..."
679+
: isCompacting
680+
? currentModel
681+
? `${getModelName(currentModel)} compacting...`
682+
: "compacting..."
683+
: currentModel
684+
? `${getModelName(currentModel)} streaming...`
685+
: "streaming..."
686+
}
687+
cancelText={
688+
awaitingUserQuestion
689+
? "type a message to respond"
690+
: `hit ${formatKeybind(vimEnabled ? KEYBINDS.INTERRUPT_STREAM_VIM : KEYBINDS.INTERRUPT_STREAM_NORMAL)} to cancel`
683691
}
684-
cancelText={`hit ${formatKeybind(vimEnabled ? KEYBINDS.INTERRUPT_STREAM_VIM : KEYBINDS.INTERRUPT_STREAM_NORMAL)} to cancel`}
685692
tokenCount={
686-
activeStreamMessageId
687-
? aggregator?.getStreamingTokenCount(activeStreamMessageId)
688-
: undefined
693+
awaitingUserQuestion
694+
? undefined
695+
: activeStreamMessageId
696+
? aggregator?.getStreamingTokenCount(activeStreamMessageId)
697+
: undefined
689698
}
690699
tps={
691-
activeStreamMessageId
692-
? aggregator?.getStreamingTPS(activeStreamMessageId)
693-
: undefined
700+
awaitingUserQuestion
701+
? undefined
702+
: activeStreamMessageId
703+
? aggregator?.getStreamingTPS(activeStreamMessageId)
704+
: undefined
694705
}
695706
/>
696707
)}

src/browser/components/Messages/ToolMessage.tsx

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { GenericToolCall } from "../tools/GenericToolCall";
55
import { BashToolCall } from "../tools/BashToolCall";
66
import { FileEditToolCall } from "../tools/FileEditToolCall";
77
import { FileReadToolCall } from "../tools/FileReadToolCall";
8+
import { AskUserQuestionToolCall } from "../tools/AskUserQuestionToolCall";
89
import { ProposePlanToolCall } from "../tools/ProposePlanToolCall";
910
import { TodoToolCall } from "../tools/TodoToolCall";
1011
import { StatusSetToolCall } from "../tools/StatusSetToolCall";
@@ -29,6 +30,8 @@ import type {
2930
FileEditReplaceStringToolResult,
3031
FileEditReplaceLinesToolArgs,
3132
FileEditReplaceLinesToolResult,
33+
AskUserQuestionToolArgs,
34+
AskUserQuestionToolResult,
3235
ProposePlanToolArgs,
3336
ProposePlanToolResult,
3437
TodoWriteToolArgs,
@@ -90,6 +93,11 @@ function isFileEditInsertTool(toolName: string, args: unknown): args is FileEdit
9093
return TOOL_DEFINITIONS.file_edit_insert.schema.safeParse(args).success;
9194
}
9295

96+
function isAskUserQuestionTool(toolName: string, args: unknown): args is AskUserQuestionToolArgs {
97+
if (toolName !== "ask_user_question") return false;
98+
return TOOL_DEFINITIONS.ask_user_question.schema.safeParse(args).success;
99+
}
100+
93101
function isProposePlanTool(toolName: string, args: unknown): args is ProposePlanToolArgs {
94102
if (toolName !== "propose_plan") return false;
95103
return TOOL_DEFINITIONS.propose_plan.schema.safeParse(args).success;
@@ -213,6 +221,20 @@ export const ToolMessage: React.FC<ToolMessageProps> = ({
213221
);
214222
}
215223

224+
if (isAskUserQuestionTool(message.toolName, message.args)) {
225+
return (
226+
<div className={className}>
227+
<AskUserQuestionToolCall
228+
args={message.args}
229+
result={(message.result as AskUserQuestionToolResult | undefined) ?? null}
230+
status={message.status}
231+
toolCallId={message.toolCallId}
232+
workspaceId={workspaceId}
233+
/>
234+
</div>
235+
);
236+
}
237+
216238
if (isProposePlanTool(message.toolName, message.args)) {
217239
return (
218240
<div className={className}>

src/browser/components/WorkspaceListItem.tsx

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,8 @@ const WorkspaceListItemInner: React.FC<WorkspaceListItemProps> = ({
9999
}
100100
};
101101

102-
const { canInterrupt } = useWorkspaceSidebarState(workspaceId);
102+
const { canInterrupt, awaitingUserQuestion } = useWorkspaceSidebarState(workspaceId);
103+
const isWorking = canInterrupt && !awaitingUserQuestion;
103104

104105
return (
105106
<React.Fragment>
@@ -167,7 +168,7 @@ const WorkspaceListItemInner: React.FC<WorkspaceListItemProps> = ({
167168
<TooltipContent align="start">Remove workspace</TooltipContent>
168169
</Tooltip>
169170
)}
170-
<RuntimeBadge runtimeConfig={metadata.runtimeConfig} isWorking={canInterrupt} />
171+
<RuntimeBadge runtimeConfig={metadata.runtimeConfig} isWorking={isWorking} />
171172
{isEditing ? (
172173
<input
173174
className="bg-input-bg text-input-text border-input-border font-inherit focus:border-input-border-focus col-span-2 min-w-0 flex-1 rounded-sm border px-1 text-left text-[13px] outline-none"
@@ -195,7 +196,7 @@ const WorkspaceListItemInner: React.FC<WorkspaceListItemProps> = ({
195196
}}
196197
title={isDisabled ? undefined : "Double-click to edit title"}
197198
>
198-
{canInterrupt || isCreating ? (
199+
{isWorking || isCreating ? (
199200
<Shimmer className="w-full truncate" colorClass="var(--color-foreground)">
200201
{displayTitle}
201202
</Shimmer>
@@ -213,7 +214,7 @@ const WorkspaceListItemInner: React.FC<WorkspaceListItemProps> = ({
213214
gitStatus={gitStatus}
214215
workspaceId={workspaceId}
215216
tooltipPosition="right"
216-
isWorking={canInterrupt}
217+
isWorking={isWorking}
217218
/>
218219
)}
219220
</div>

src/browser/components/WorkspaceStatusDot.tsx

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,10 @@ export const WorkspaceStatusDot = memo<{
1111
size?: number;
1212
}>(
1313
({ workspaceId, lastReadTimestamp, onClick, size = 8 }) => {
14-
const { canInterrupt, currentModel, agentStatus, recencyTimestamp } =
14+
const { canInterrupt, awaitingUserQuestion, currentModel, agentStatus, recencyTimestamp } =
1515
useWorkspaceSidebarState(workspaceId);
1616

17-
const streaming = canInterrupt;
17+
const streaming = canInterrupt && !awaitingUserQuestion;
1818

1919
// Compute unread status if lastReadTimestamp provided (sidebar only)
2020
const unread = useMemo(() => {
@@ -27,12 +27,13 @@ export const WorkspaceStatusDot = memo<{
2727
() =>
2828
getStatusTooltip({
2929
isStreaming: streaming,
30+
isAwaitingInput: awaitingUserQuestion,
3031
streamingModel: currentModel,
3132
agentStatus,
3233
isUnread: unread,
3334
recencyTimestamp,
3435
}),
35-
[streaming, currentModel, agentStatus, unread, recencyTimestamp]
36+
[streaming, awaitingUserQuestion, currentModel, agentStatus, unread, recencyTimestamp]
3637
);
3738

3839
const bgColor = canInterrupt ? "bg-blue-400" : unread ? "bg-gray-300" : "bg-muted-dark";

src/browser/components/WorkspaceStatusIndicator.tsx

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,17 @@ import { Tooltip, TooltipTrigger, TooltipContent } from "./ui/tooltip";
55
import { Button } from "./ui/button";
66

77
export const WorkspaceStatusIndicator = memo<{ workspaceId: string }>(({ workspaceId }) => {
8-
const { agentStatus } = useWorkspaceSidebarState(workspaceId);
8+
const { agentStatus, awaitingUserQuestion } = useWorkspaceSidebarState(workspaceId);
9+
10+
// Show prompt when ask_user_question is pending - make it prominent
11+
if (awaitingUserQuestion) {
12+
return (
13+
<div className="flex min-w-0 items-center gap-1.5 rounded bg-yellow-500/20 px-1.5 py-0.5 text-xs text-yellow-400">
14+
<span className="-mt-0.5 shrink-0 text-[10px]"></span>
15+
<span className="min-w-0 truncate font-medium">Mux has a few questions</span>
16+
</div>
17+
);
18+
}
919

1020
if (!agentStatus) {
1121
return null;

0 commit comments

Comments
 (0)