diff --git a/docs/plan-mode.mdx b/docs/plan-mode.mdx index 35ee7f2e47..a395fe0123 100644 --- a/docs/plan-mode.mdx +++ b/docs/plan-mode.mdx @@ -29,13 +29,34 @@ This means you can make edits in your preferred editor, return to mux, send a me ## Plan File Location -Plans are stored in a dedicated directory: +Plans are stored in a dedicated directory under your Mux home: ``` -~/.mux/plans/.md +~/.mux/plans//.md ``` -The file is created when the agent first writes a plan and persists across sessions. +Notes: + +- `` includes the random suffix (e.g. `feature-x7k2`), so it’s globally unique with high probability. + +## ask_user_question (Plan Mode Only) + +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. + +What you’ll see: + +- An inline “tool call card” in the chat with a small form (single-select or multi-select). +- An always-available **Other** option for free-form answers. + +How to respond: + +- **Recommended:** answer in the form and click **Submit answers**. +- **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. + +Availability: + +- `ask_user_question` is only registered for the agent in **Plan Mode**. +- In Exec Mode, the agent cannot call `ask_user_question`. ## UI Features @@ -55,7 +76,7 @@ User: "Add user authentication to the app" ▼ ┌─────────────────────────────────────┐ │ Agent reads codebase, writes plan │ -│ to ~/.mux/plans/.md │ +│ to ~/.mux/plans//.md│ └─────────────────────────────────────┘ │ ▼ diff --git a/src/browser/components/AIView.tsx b/src/browser/components/AIView.tsx index c7078bd27b..f37d518111 100644 --- a/src/browser/components/AIView.tsx +++ b/src/browser/components/AIView.tsx @@ -160,7 +160,8 @@ const AIViewInner: React.FC = ({ useEffect(() => { workspaceStateRef.current = workspaceState; }, [workspaceState]); - const { messages, canInterrupt, isCompacting, loading, currentModel } = workspaceState; + const { messages, canInterrupt, isCompacting, awaitingUserQuestion, loading, currentModel } = + workspaceState; // Apply message transformations: // 1. Merge consecutive identical stream errors @@ -673,24 +674,34 @@ const AIViewInner: React.FC = ({ {canInterrupt && ( )} diff --git a/src/browser/components/Messages/ToolMessage.tsx b/src/browser/components/Messages/ToolMessage.tsx index 2d4e68948f..e581e91844 100644 --- a/src/browser/components/Messages/ToolMessage.tsx +++ b/src/browser/components/Messages/ToolMessage.tsx @@ -5,6 +5,7 @@ import { GenericToolCall } from "../tools/GenericToolCall"; import { BashToolCall } from "../tools/BashToolCall"; import { FileEditToolCall } from "../tools/FileEditToolCall"; import { FileReadToolCall } from "../tools/FileReadToolCall"; +import { AskUserQuestionToolCall } from "../tools/AskUserQuestionToolCall"; import { ProposePlanToolCall } from "../tools/ProposePlanToolCall"; import { TodoToolCall } from "../tools/TodoToolCall"; import { StatusSetToolCall } from "../tools/StatusSetToolCall"; @@ -29,6 +30,8 @@ import type { FileEditReplaceStringToolResult, FileEditReplaceLinesToolArgs, FileEditReplaceLinesToolResult, + AskUserQuestionToolArgs, + AskUserQuestionToolResult, ProposePlanToolArgs, ProposePlanToolResult, TodoWriteToolArgs, @@ -90,6 +93,11 @@ function isFileEditInsertTool(toolName: string, args: unknown): args is FileEdit return TOOL_DEFINITIONS.file_edit_insert.schema.safeParse(args).success; } +function isAskUserQuestionTool(toolName: string, args: unknown): args is AskUserQuestionToolArgs { + if (toolName !== "ask_user_question") return false; + return TOOL_DEFINITIONS.ask_user_question.schema.safeParse(args).success; +} + function isProposePlanTool(toolName: string, args: unknown): args is ProposePlanToolArgs { if (toolName !== "propose_plan") return false; return TOOL_DEFINITIONS.propose_plan.schema.safeParse(args).success; @@ -213,6 +221,20 @@ export const ToolMessage: React.FC = ({ ); } + if (isAskUserQuestionTool(message.toolName, message.args)) { + return ( +
+ +
+ ); + } + if (isProposePlanTool(message.toolName, message.args)) { return (
diff --git a/src/browser/components/WorkspaceListItem.tsx b/src/browser/components/WorkspaceListItem.tsx index 4ec1557a24..9c465a4484 100644 --- a/src/browser/components/WorkspaceListItem.tsx +++ b/src/browser/components/WorkspaceListItem.tsx @@ -99,7 +99,8 @@ const WorkspaceListItemInner: React.FC = ({ } }; - const { canInterrupt } = useWorkspaceSidebarState(workspaceId); + const { canInterrupt, awaitingUserQuestion } = useWorkspaceSidebarState(workspaceId); + const isWorking = canInterrupt && !awaitingUserQuestion; return ( @@ -167,7 +168,7 @@ const WorkspaceListItemInner: React.FC = ({ Remove workspace )} - + {isEditing ? ( = ({ }} title={isDisabled ? undefined : "Double-click to edit title"} > - {canInterrupt || isCreating ? ( + {isWorking || isCreating ? ( {displayTitle} @@ -213,7 +214,7 @@ const WorkspaceListItemInner: React.FC = ({ gitStatus={gitStatus} workspaceId={workspaceId} tooltipPosition="right" - isWorking={canInterrupt} + isWorking={isWorking} /> )}
diff --git a/src/browser/components/WorkspaceStatusDot.tsx b/src/browser/components/WorkspaceStatusDot.tsx index eae4d9bc1c..15f14501ea 100644 --- a/src/browser/components/WorkspaceStatusDot.tsx +++ b/src/browser/components/WorkspaceStatusDot.tsx @@ -11,10 +11,10 @@ export const WorkspaceStatusDot = memo<{ size?: number; }>( ({ workspaceId, lastReadTimestamp, onClick, size = 8 }) => { - const { canInterrupt, currentModel, agentStatus, recencyTimestamp } = + const { canInterrupt, awaitingUserQuestion, currentModel, agentStatus, recencyTimestamp } = useWorkspaceSidebarState(workspaceId); - const streaming = canInterrupt; + const streaming = canInterrupt && !awaitingUserQuestion; // Compute unread status if lastReadTimestamp provided (sidebar only) const unread = useMemo(() => { @@ -27,12 +27,13 @@ export const WorkspaceStatusDot = memo<{ () => getStatusTooltip({ isStreaming: streaming, + isAwaitingInput: awaitingUserQuestion, streamingModel: currentModel, agentStatus, isUnread: unread, recencyTimestamp, }), - [streaming, currentModel, agentStatus, unread, recencyTimestamp] + [streaming, awaitingUserQuestion, currentModel, agentStatus, unread, recencyTimestamp] ); const bgColor = canInterrupt ? "bg-blue-400" : unread ? "bg-gray-300" : "bg-muted-dark"; diff --git a/src/browser/components/WorkspaceStatusIndicator.tsx b/src/browser/components/WorkspaceStatusIndicator.tsx index fc03d7899d..db2c3798bf 100644 --- a/src/browser/components/WorkspaceStatusIndicator.tsx +++ b/src/browser/components/WorkspaceStatusIndicator.tsx @@ -5,7 +5,17 @@ import { Tooltip, TooltipTrigger, TooltipContent } from "./ui/tooltip"; import { Button } from "./ui/button"; export const WorkspaceStatusIndicator = memo<{ workspaceId: string }>(({ workspaceId }) => { - const { agentStatus } = useWorkspaceSidebarState(workspaceId); + const { agentStatus, awaitingUserQuestion } = useWorkspaceSidebarState(workspaceId); + + // Show prompt when ask_user_question is pending - make it prominent + if (awaitingUserQuestion) { + return ( +
+ + Mux has a few questions +
+ ); + } if (!agentStatus) { return null; diff --git a/src/browser/components/tools/AskUserQuestionToolCall.tsx b/src/browser/components/tools/AskUserQuestionToolCall.tsx new file mode 100644 index 0000000000..2ffd24c20a --- /dev/null +++ b/src/browser/components/tools/AskUserQuestionToolCall.tsx @@ -0,0 +1,561 @@ +import assert from "@/common/utils/assert"; + +import { useMemo, useState } from "react"; + +import { useAPI } from "@/browser/contexts/API"; +import { Checkbox } from "@/browser/components/ui/checkbox"; +import { Input } from "@/browser/components/ui/input"; +import { Button } from "@/browser/components/ui/button"; +import { + ErrorBox, + ExpandIcon, + StatusIndicator, + ToolContainer, + ToolDetails, + ToolHeader, + ToolName, +} from "@/browser/components/tools/shared/ToolPrimitives"; +import { + getStatusDisplay, + useToolExpansion, + type ToolStatus, +} from "@/browser/components/tools/shared/toolUtils"; +import type { + AskUserQuestionQuestion, + AskUserQuestionToolArgs, + AskUserQuestionToolResult, + AskUserQuestionToolSuccessResult, + ToolErrorResult, +} from "@/common/types/tools"; + +const OTHER_VALUE = "__other__"; + +interface DraftAnswer { + selected: string[]; + otherText: string; +} + +function unwrapJsonContainer(value: unknown): unknown { + if (!value || typeof value !== "object") { + return value; + } + + const record = value as Record; + if (record.type === "json" && "value" in record) { + return record.value; + } + + return value; +} + +function isAskUserQuestionToolSuccessResult(val: unknown): val is AskUserQuestionToolSuccessResult { + if (!val || typeof val !== "object") { + return false; + } + + const record = val as Record; + if (!Array.isArray(record.questions)) { + return false; + } + + if (!record.answers || typeof record.answers !== "object") { + return false; + } + + for (const [, v] of Object.entries(record.answers as Record)) { + if (typeof v !== "string") { + return false; + } + } + + return true; +} + +function isToolErrorResult(val: unknown): val is ToolErrorResult { + if (!val || typeof val !== "object") { + return false; + } + + const record = val as Record; + return record.success === false && typeof record.error === "string"; +} + +function parsePrefilledAnswer(question: AskUserQuestionQuestion, answer: string): DraftAnswer { + const trimmed = answer.trim(); + if (trimmed.length === 0) { + return { selected: [], otherText: "" }; + } + + const optionLabels = new Set(question.options.map((o) => o.label)); + + if (!question.multiSelect) { + if (optionLabels.has(trimmed)) { + return { selected: [trimmed], otherText: "" }; + } + + return { selected: [OTHER_VALUE], otherText: trimmed }; + } + + const tokens = trimmed + .split(",") + .map((t) => t.trim()) + .filter((t) => t.length > 0); + + const selected: string[] = []; + const otherParts: string[] = []; + + for (const token of tokens) { + if (optionLabels.has(token)) { + selected.push(token); + } else { + otherParts.push(token); + } + } + + if (otherParts.length > 0) { + selected.push(OTHER_VALUE); + } + + return { selected, otherText: otherParts.join(", ") }; +} + +function isQuestionAnswered(question: AskUserQuestionQuestion, draft: DraftAnswer): boolean { + if (draft.selected.length === 0) { + return false; + } + + if (draft.selected.includes(OTHER_VALUE)) { + return draft.otherText.trim().length > 0; + } + + return true; +} + +function draftToAnswerString(question: AskUserQuestionQuestion, draft: DraftAnswer): string { + assert(isQuestionAnswered(question, draft), "draftToAnswerString requires a complete answer"); + + const parts: string[] = []; + for (const label of draft.selected) { + if (label === OTHER_VALUE) { + parts.push(draft.otherText.trim()); + } else { + parts.push(label); + } + } + + if (!question.multiSelect) { + assert(parts.length === 1, "Single-select questions must have exactly one answer"); + return parts[0]; + } + + return parts.join(", "); +} + +export function AskUserQuestionToolCall(props: { + args: AskUserQuestionToolArgs; + result: AskUserQuestionToolResult | null; + status: ToolStatus; + toolCallId: string; + workspaceId?: string; +}): JSX.Element { + const { api } = useAPI(); + + const { expanded, toggleExpanded } = useToolExpansion(props.status === "executing"); + const statusDisplay = getStatusDisplay(props.status); + + const [activeIndex, setActiveIndex] = useState(0); + const [isSubmitting, setIsSubmitting] = useState(false); + const [submitError, setSubmitError] = useState(null); + + const argsAnswers = props.args.answers ?? {}; + + const [draftAnswers, setDraftAnswers] = useState>(() => { + const initial: Record = {}; + for (const q of props.args.questions) { + const prefilled = argsAnswers[q.question]; + if (typeof prefilled === "string") { + initial[q.question] = parsePrefilledAnswer(q, prefilled); + } else { + initial[q.question] = { selected: [], otherText: "" }; + } + } + return initial; + }); + + const resultUnwrapped = useMemo(() => { + if (!props.result) { + return null; + } + + return unwrapJsonContainer(props.result); + }, [props.result]); + + const successResult = + resultUnwrapped && isAskUserQuestionToolSuccessResult(resultUnwrapped) ? resultUnwrapped : null; + + const errorResult = + resultUnwrapped && isToolErrorResult(resultUnwrapped) ? resultUnwrapped : null; + + const isComplete = useMemo(() => { + return props.args.questions.every((q) => { + const draft = draftAnswers[q.question]; + return draft ? isQuestionAnswered(q, draft) : false; + }); + }, [draftAnswers, props.args.questions]); + + const summaryIndex = props.args.questions.length; + const isOnSummary = activeIndex === summaryIndex; + const currentQuestion = isOnSummary + ? null + : props.args.questions[Math.min(activeIndex, props.args.questions.length - 1)]; + const currentDraft = currentQuestion ? draftAnswers[currentQuestion.question] : undefined; + + const unansweredCount = useMemo(() => { + return props.args.questions.filter((q) => { + const draft = draftAnswers[q.question]; + return !draft || !isQuestionAnswered(q, draft); + }).length; + }, [draftAnswers, props.args.questions]); + + const handleSubmit = (): void => { + setIsSubmitting(true); + setSubmitError(null); + + let answers: Record; + + try { + answers = {}; + for (const q of props.args.questions) { + const draft = draftAnswers[q.question]; + if (draft && isQuestionAnswered(q, draft)) { + answers[q.question] = draftToAnswerString(q, draft); + } else { + // Unanswered questions get empty string + answers[q.question] = ""; + } + } + + assert(api, "API not connected"); + assert(props.workspaceId, "workspaceId is required"); + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + setSubmitError(errorMessage); + setIsSubmitting(false); + return; + } + + api.workspace + .answerAskUserQuestion({ + workspaceId: props.workspaceId, + toolCallId: props.toolCallId, + answers, + }) + .then((result) => { + if (!result.success) { + setSubmitError(result.error); + } + }) + .catch((error) => { + const errorMessage = error instanceof Error ? error.message : String(error); + setSubmitError(errorMessage); + }) + .finally(() => { + setIsSubmitting(false); + }); + }; + const title = "ask_user_question"; + + return ( + + + +
+ {title} +
+ Answer below, or type in chat to cancel. +
+
+ {statusDisplay} +
+ + {expanded && ( + +
+ {props.status === "executing" && ( +
+
+ {props.args.questions.map((q, idx) => { + const draft = draftAnswers[q.question]; + const answered = draft ? isQuestionAnswered(q, draft) : false; + const isActive = idx === activeIndex; + return ( + + ); + })} + +
+ + {!isOnSummary && currentQuestion && currentDraft && ( + <> +
+
{currentQuestion.question}
+
+ +
+ {currentQuestion.options.map((opt) => { + const checked = currentDraft.selected.includes(opt.label); + + const toggle = () => { + setDraftAnswers((prev) => { + const next = { ...prev }; + const draft = next[currentQuestion.question] ?? { + selected: [], + otherText: "", + }; + + if (currentQuestion.multiSelect) { + const selected = new Set(draft.selected); + if (selected.has(opt.label)) { + selected.delete(opt.label); + } else { + selected.add(opt.label); + } + next[currentQuestion.question] = { + ...draft, + selected: Array.from(selected), + }; + } else { + next[currentQuestion.question] = { + selected: checked ? [] : [opt.label], + otherText: "", + }; + } + + return next; + }); + }; + + return ( +
{ + if (e.key === "Enter" || e.key === " ") { + e.preventDefault(); + toggle(); + } + }} + > + e.stopPropagation()} + /> +
+
{opt.label}
+
{opt.description}
+
+
+ ); + })} + + {(() => { + const checked = currentDraft.selected.includes(OTHER_VALUE); + const toggle = () => { + setDraftAnswers((prev) => { + const next = { ...prev }; + const draft = next[currentQuestion.question] ?? { + selected: [], + otherText: "", + }; + const selected = new Set(draft.selected); + if (selected.has(OTHER_VALUE)) { + selected.delete(OTHER_VALUE); + next[currentQuestion.question] = { + ...draft, + selected: Array.from(selected), + }; + } else { + if (!currentQuestion.multiSelect) { + selected.clear(); + } + selected.add(OTHER_VALUE); + next[currentQuestion.question] = { + ...draft, + selected: Array.from(selected), + }; + } + return next; + }); + }; + + return ( +
{ + if (e.key === "Enter" || e.key === " ") { + e.preventDefault(); + toggle(); + } + }} + > + e.stopPropagation()} + /> +
+
Other
+
+ Provide a custom answer. +
+
+
+ ); + })()} + + {currentDraft.selected.includes(OTHER_VALUE) && ( + { + const value = e.target.value; + setDraftAnswers((prev) => ({ + ...prev, + [currentQuestion.question]: { + ...(prev[currentQuestion.question] ?? { + selected: [], + otherText: "", + }), + otherText: value, + }, + })); + }} + /> + )} +
+ + )} + + {isOnSummary && ( +
+
Review your answers
+ {unansweredCount > 0 && ( +
+ ⚠️ {unansweredCount} question{unansweredCount > 1 ? "s" : ""} not answered +
+ )} +
+ {props.args.questions.map((q, idx) => { + const draft = draftAnswers[q.question]; + const answered = draft ? isQuestionAnswered(q, draft) : false; + const answerText = answered ? draftToAnswerString(q, draft) : null; + return ( +
setActiveIndex(idx)} + onKeyDown={(e) => { + if (e.key === "Enter" || e.key === " ") { + e.preventDefault(); + setActiveIndex(idx); + } + }} + > + {answered ? ( + + ) : ( + ⚠️ + )}{" "} + {q.header}:{" "} + {answered ? ( + {answerText} + ) : ( + Not answered + )} +
+ ); + })} +
+
+ )} + +
+ Tip: you can also just type a message to respond in chat (this will cancel these + questions). +
+ + {submitError && {submitError}} +
+ )} + + {props.status !== "executing" && ( +
+ {successResult && ( +
+
User answered:
+ {Object.entries(successResult.answers).map(([question, answer]) => ( +
+ • {question}: {answer} +
+ ))} +
+ )} + + {errorResult && {errorResult.error}} +
+ )} + + {props.status === "executing" && ( +
+ {isOnSummary ? ( + + ) : ( + + )} +
+ )} +
+
+ )} +
+ ); +} diff --git a/src/browser/stores/WorkspaceStore.ts b/src/browser/stores/WorkspaceStore.ts index 134440d523..fe0d15ebf2 100644 --- a/src/browser/stores/WorkspaceStore.ts +++ b/src/browser/stores/WorkspaceStore.ts @@ -38,6 +38,7 @@ export interface WorkspaceState { queuedMessage: QueuedMessage | null; canInterrupt: boolean; isCompacting: boolean; + awaitingUserQuestion: boolean; loading: boolean; muxMessages: MuxMessage[]; currentModel: string | null; @@ -53,6 +54,7 @@ export interface WorkspaceState { */ export interface WorkspaceSidebarState { canInterrupt: boolean; + awaitingUserQuestion: boolean; currentModel: string | null; recencyTimestamp: number | null; agentStatus: { emoji: string; message: string; url?: string } | undefined; @@ -480,6 +482,7 @@ export class WorkspaceStore { queuedMessage: this.queuedMessages.get(workspaceId) ?? null, canInterrupt: activeStreams.length > 0, isCompacting: aggregator.isCompacting(), + awaitingUserQuestion: aggregator.hasAwaitingUserQuestion(), loading: !hasMessages && !isCaughtUp, muxMessages: messages, currentModel: aggregator.getCurrentModel() ?? null, @@ -507,6 +510,7 @@ export class WorkspaceStore { if ( cached && cached.canInterrupt === fullState.canInterrupt && + cached.awaitingUserQuestion === fullState.awaitingUserQuestion && cached.currentModel === fullState.currentModel && cached.recencyTimestamp === fullState.recencyTimestamp && cached.agentStatus === fullState.agentStatus @@ -517,6 +521,7 @@ export class WorkspaceStore { // Create and cache new state const newState: WorkspaceSidebarState = { canInterrupt: fullState.canInterrupt, + awaitingUserQuestion: fullState.awaitingUserQuestion, currentModel: fullState.currentModel, recencyTimestamp: fullState.recencyTimestamp, agentStatus: fullState.agentStatus, diff --git a/src/browser/stories/App.chat.stories.tsx b/src/browser/stories/App.chat.stories.tsx index 0fac0723e4..b956d85a35 100644 --- a/src/browser/stories/App.chat.stories.tsx +++ b/src/browser/stories/App.chat.stories.tsx @@ -270,6 +270,112 @@ export const Streaming: AppStory = { ), }; +/** Streaming/working state with ask_user_question pending */ +export const AskUserQuestionPending: AppStory = { + render: () => ( + + setupStreamingChatStory({ + messages: [ + createUserMessage("msg-1", "Please implement the feature", { + historySequence: 1, + timestamp: STABLE_TIMESTAMP - 3000, + }), + ], + streamingMessageId: "msg-2", + historySequence: 2, + streamText: "I have a few clarifying questions.", + pendingTool: { + toolCallId: "call-ask-1", + toolName: "ask_user_question", + args: { + questions: [ + { + question: "Which approach should we take?", + header: "Approach", + options: [ + { label: "A", description: "Approach A" }, + { label: "B", description: "Approach B" }, + ], + multiSelect: false, + }, + { + question: "Which platforms do we need to support?", + header: "Platforms", + options: [ + { label: "macOS", description: "Apple macOS" }, + { label: "Windows", description: "Microsoft Windows" }, + { label: "Linux", description: "Linux desktops" }, + ], + multiSelect: true, + }, + ], + }, + }, + gitStatus: { dirty: 1 }, + }) + } + /> + ), +}; + +/** Completed ask_user_question tool call */ +export const AskUserQuestionCompleted: AppStory = { + render: () => ( + + setupSimpleChatStory({ + messages: [ + createUserMessage("msg-1", "Please implement the feature", { + historySequence: 1, + timestamp: STABLE_TIMESTAMP - 60000, + }), + createAssistantMessage("msg-2", "I asked some questions.", { + historySequence: 2, + timestamp: STABLE_TIMESTAMP - 55000, + toolCalls: [ + createGenericTool( + "call-ask-1", + "ask_user_question", + { + questions: [ + { + question: "Which approach should we take?", + header: "Approach", + options: [ + { label: "A", description: "Approach A" }, + { label: "B", description: "Approach B" }, + ], + multiSelect: false, + }, + ], + }, + { + questions: [ + { + question: "Which approach should we take?", + header: "Approach", + options: [ + { label: "A", description: "Approach A" }, + { label: "B", description: "Approach B" }, + ], + multiSelect: false, + }, + ], + answers: { + "Which approach should we take?": "A", + }, + } + ), + ], + }), + ], + }) + } + /> + ), +}; + /** Generic tool call with JSON-highlighted arguments and results */ export const GenericTool: AppStory = { render: () => ( @@ -319,21 +425,6 @@ export const GenericTool: AppStory = { }, }, }, - play: async ({ canvasElement }: { canvasElement: HTMLElement }) => { - const canvas = within(canvasElement); - - // Wait for workspace metadata to load and main content to render - await waitFor( - async () => { - const toolHeader = canvas.getByText("fetch_data"); - await userEvent.click(toolHeader); - }, - { timeout: 5000 } - ); - // Wait for any auto-focus timers (ChatInput has 100ms delay), then blur - await new Promise((resolve) => setTimeout(resolve, 150)); - (document.activeElement as HTMLElement)?.blur(); - }, }; /** Streaming compaction with shimmer effect - tests GPU-accelerated animation */ @@ -467,7 +558,8 @@ export const ModeHelpTooltip: AppStory = { /> ), play: async ({ canvasElement }) => { - const canvas = within(canvasElement); + const storyRoot = document.getElementById("storybook-root") ?? canvasElement; + const canvas = within(storyRoot); // Wait for app to fully load - the chat input with mode selector should be present await canvas.findAllByText("Exec", {}, { timeout: 10000 }); diff --git a/src/browser/utils/messages/StreamingMessageAggregator.ts b/src/browser/utils/messages/StreamingMessageAggregator.ts index 0c8e0909dc..81ec063dd8 100644 --- a/src/browser/utils/messages/StreamingMessageAggregator.ts +++ b/src/browser/utils/messages/StreamingMessageAggregator.ts @@ -234,6 +234,25 @@ export class StreamingMessageAggregator { return this.agentStatus; } + /** + * Check if there's an executing ask_user_question tool awaiting user input. + * 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 + const displayed = this.getDisplayedMessages(); + for (const msg of displayed) { + if ( + msg.type === "tool" && + msg.toolName === "ask_user_question" && + msg.status === "executing" + ) { + return true; + } + } + return false; + } + /** * Extract compaction summary text from a completed assistant message. * Used when a compaction stream completes to get the summary for history replacement. diff --git a/src/browser/utils/messages/toolOutputRedaction.ts b/src/browser/utils/messages/toolOutputRedaction.ts index d578f65e98..a75be97b27 100644 --- a/src/browser/utils/messages/toolOutputRedaction.ts +++ b/src/browser/utils/messages/toolOutputRedaction.ts @@ -11,6 +11,7 @@ */ import type { + AskUserQuestionToolSuccessResult, FileEditInsertToolResult, FileEditReplaceStringToolResult, FileEditReplaceLinesToolResult, @@ -103,9 +104,46 @@ function redactFileEditInsert(output: unknown): unknown { return output; } +function isAskUserQuestionToolSuccessResult(val: unknown): val is AskUserQuestionToolSuccessResult { + if (!val || typeof val !== "object") return false; + const record = val as Record; + + if (!Array.isArray(record.questions)) return false; + if (!record.answers || typeof record.answers !== "object") return false; + + // answers is Record + for (const [k, v] of Object.entries(record.answers as Record)) { + if (typeof k !== "string" || typeof v !== "string") return false; + } + + return true; +} + +function redactAskUserQuestion(output: unknown): unknown { + const unwrapped = unwrapJsonContainer(output); + const val = unwrapped.value; + + if (!isAskUserQuestionToolSuccessResult(val)) { + return output; + } + + const pairs = Object.entries(val.answers) + .map(([question, answer]) => `"${question}"="${answer}"`) + .join(", "); + + const summary = + pairs.length > 0 + ? `User has answered your questions: ${pairs}. You can now continue with the user's answers in mind.` + : "User has answered your questions. You can now continue with the user's answers in mind."; + + return rewrapJsonContainer(unwrapped.wrapped, summary); +} + // Public API - registry entrypoint. Add new tools here as needed. export function redactToolOutput(toolName: string, output: unknown): unknown { switch (toolName) { + case "ask_user_question": + return redactAskUserQuestion(output); case "file_edit_replace_string": case "file_edit_replace_lines": return redactFileEditReplace(output); diff --git a/src/browser/utils/ui/statusTooltip.tsx b/src/browser/utils/ui/statusTooltip.tsx index 7f67bc66a6..e3d8f37a90 100644 --- a/src/browser/utils/ui/statusTooltip.tsx +++ b/src/browser/utils/ui/statusTooltip.tsx @@ -8,12 +8,14 @@ import { formatRelativeTime } from "@/browser/utils/ui/dateTime"; */ export function getStatusTooltip(options: { isStreaming: boolean; + isAwaitingInput?: boolean; streamingModel: string | null; agentStatus?: { emoji: string; message: string; url?: string }; isUnread?: boolean; recencyTimestamp?: number | null; }): React.ReactNode { - const { isStreaming, streamingModel, agentStatus, isUnread, recencyTimestamp } = options; + const { isStreaming, isAwaitingInput, streamingModel, agentStatus, isUnread, recencyTimestamp } = + options; // If agent status is set, show message and URL (if available) if (agentStatus) { @@ -29,6 +31,11 @@ export function getStatusTooltip(options: { return agentStatus.message; } + // Show awaiting input status + if (isAwaitingInput) { + return "Awaiting your input"; + } + // Otherwise show streaming/idle status if (isStreaming && streamingModel) { return ( diff --git a/src/common/orpc/schemas/api.ts b/src/common/orpc/schemas/api.ts index 46c44da95e..8de78fbe3f 100644 --- a/src/common/orpc/schemas/api.ts +++ b/src/common/orpc/schemas/api.ts @@ -237,6 +237,16 @@ export const workspace = { }), output: ResultSchema(z.object({}), SendMessageErrorSchema), }, + answerAskUserQuestion: { + input: z + .object({ + workspaceId: z.string(), + toolCallId: z.string(), + answers: z.record(z.string(), z.string()), + }) + .strict(), + output: ResultSchema(z.void(), z.string()), + }, resumeStream: { input: z.object({ workspaceId: z.string(), diff --git a/src/common/types/tools.ts b/src/common/types/tools.ts index ac86a7acbc..d0341dc584 100644 --- a/src/common/types/tools.ts +++ b/src/common/types/tools.ts @@ -4,7 +4,12 @@ */ import type { z } from "zod"; -import type { TOOL_DEFINITIONS } from "@/common/utils/tools/toolDefinitions"; +import type { + AskUserQuestionOptionSchema, + AskUserQuestionQuestionSchema, + AskUserQuestionToolResultSchema, + TOOL_DEFINITIONS, +} from "@/common/utils/tools/toolDefinitions"; // Bash Tool Types export interface BashToolArgs { @@ -161,11 +166,27 @@ export const NOTE_READ_FILE_AGAIN_RETRY = "Read the file again and retry."; export const TOOL_EDIT_WARNING = "Always check the tool result before proceeding with other operations."; +// Generic tool error shape emitted via streamManager on tool-error parts. +export interface ToolErrorResult { + success: false; + error: string; +} export type FileEditToolArgs = | FileEditReplaceStringToolArgs | FileEditReplaceLinesToolArgs | FileEditInsertToolArgs; +// Ask User Question Tool Types +// Args derived from schema (avoid drift) +export type AskUserQuestionToolArgs = z.infer; + +export type AskUserQuestionOption = z.infer; +export type AskUserQuestionQuestion = z.infer; + +export type AskUserQuestionToolSuccessResult = z.infer; + +export type AskUserQuestionToolResult = AskUserQuestionToolSuccessResult | ToolErrorResult; + // Propose Plan Tool Types // Args derived from schema export type ProposePlanToolArgs = z.infer; diff --git a/src/common/utils/tools/toolDefinitions.ts b/src/common/utils/tools/toolDefinitions.ts index 522eee72f5..06be704621 100644 --- a/src/common/utils/tools/toolDefinitions.ts +++ b/src/common/utils/tools/toolDefinitions.ts @@ -17,6 +17,71 @@ import { TOOL_EDIT_WARNING } from "@/common/types/tools"; import { zodToJsonSchema } from "zod-to-json-schema"; +// ----------------------------------------------------------------------------- +// ask_user_question (plan-mode interactive questions) +// ----------------------------------------------------------------------------- + +export const AskUserQuestionOptionSchema = z + .object({ + label: z.string().min(1), + description: z.string().min(1), + }) + .strict(); + +export const AskUserQuestionQuestionSchema = z + .object({ + question: z.string().min(1), + header: z.string().min(1).max(12), + options: z.array(AskUserQuestionOptionSchema).min(2).max(4), + multiSelect: z.boolean(), + }) + .strict() + .superRefine((question, ctx) => { + const labels = question.options.map((o) => o.label); + const labelSet = new Set(labels); + if (labelSet.size !== labels.length) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "Option labels must be unique within a question", + path: ["options"], + }); + } + + // Claude Code provides "Other" automatically; do not include it explicitly. + if (labels.some((label) => label.trim().toLowerCase() === "other")) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "Do not include an 'Other' option; it is provided automatically", + path: ["options"], + }); + } + }); + +export const AskUserQuestionToolArgsSchema = z + .object({ + questions: z.array(AskUserQuestionQuestionSchema).min(1).max(4), + // Optional prefilled answers (Claude Code supports this, though Mux typically won't use it) + answers: z.record(z.string(), z.string()).optional(), + }) + .strict() + .superRefine((args, ctx) => { + const questionTexts = args.questions.map((q) => q.question); + const questionTextSet = new Set(questionTexts); + if (questionTextSet.size !== questionTexts.length) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "Question text must be unique across questions", + path: ["questions"], + }); + } + }); + +export const AskUserQuestionToolResultSchema = z + .object({ + questions: z.array(AskUserQuestionQuestionSchema), + answers: z.record(z.string(), z.string()), + }) + .strict(); const FILE_EDIT_FILE_PATH = z .string() .describe("Path to the file to edit (absolute or relative to the current workspace)"); @@ -160,6 +225,13 @@ export const TOOL_DEFINITIONS = { path: ["before"], }), }, + ask_user_question: { + description: + "Ask 1–4 multiple-choice questions (with optional multi-select) and wait for the user's answers. " + + "This tool is intended for plan mode and should be used when proceeding requires clarification. " + + "Each question must include 2–4 options; an 'Other' choice is provided automatically.", + schema: AskUserQuestionToolArgsSchema, + }, propose_plan: { description: "Signal that your plan is complete and ready for user approval. " + @@ -329,9 +401,10 @@ export function getToolSchemas(): Record { /** * Get which tools are available for a given model * @param modelString The model string (e.g., "anthropic:claude-opus-4-1") + * @param mode Optional mode ("plan" | "exec") - ask_user_question only available in plan mode * @returns Array of tool names available for the model */ -export function getAvailableTools(modelString: string): string[] { +export function getAvailableTools(modelString: string, mode?: "plan" | "exec"): string[] { const [provider] = modelString.split(":"); // Base tools available for all models @@ -344,6 +417,8 @@ export function getAvailableTools(modelString: string): string[] { "file_edit_replace_string", // "file_edit_replace_lines", // DISABLED: causes models to break repo state "file_edit_insert", + // ask_user_question only available in plan mode + ...(mode === "plan" ? ["ask_user_question"] : []), "propose_plan", "todo_write", "todo_read", diff --git a/src/common/utils/tools/tools.ts b/src/common/utils/tools/tools.ts index 94c19a2ec0..e3fb536adb 100644 --- a/src/common/utils/tools/tools.ts +++ b/src/common/utils/tools/tools.ts @@ -7,6 +7,7 @@ import { createBashBackgroundTerminateTool } from "@/node/services/tools/bash_ba import { createFileEditReplaceStringTool } from "@/node/services/tools/file_edit_replace_string"; // DISABLED: import { createFileEditReplaceLinesTool } from "@/node/services/tools/file_edit_replace_lines"; import { createFileEditInsertTool } from "@/node/services/tools/file_edit_insert"; +import { createAskUserQuestionTool } from "@/node/services/tools/ask_user_question"; import { createProposePlanTool } from "@/node/services/tools/propose_plan"; import { createTodoWriteTool, createTodoReadTool } from "@/node/services/tools/todo"; import { createStatusSetTool } from "@/node/services/tools/status_set"; @@ -126,6 +127,7 @@ export async function getToolsForModel( // Non-runtime tools execute immediately (no init wait needed) const nonRuntimeTools: Record = { + ...(config.mode === "plan" ? { ask_user_question: createAskUserQuestionTool(config) } : {}), propose_plan: createProposePlanTool(config), todo_write: createTodoWriteTool(config), todo_read: createTodoReadTool(config), diff --git a/src/node/orpc/router.ts b/src/node/orpc/router.ts index 9a2cf18db3..f2e2044283 100644 --- a/src/node/orpc/router.ts +++ b/src/node/orpc/router.ts @@ -350,6 +350,22 @@ export const router = (authToken?: string) => { return { success: true, data: {} }; }), + answerAskUserQuestion: t + .input(schemas.workspace.answerAskUserQuestion.input) + .output(schemas.workspace.answerAskUserQuestion.output) + .handler(({ context, input }) => { + const result = context.workspaceService.answerAskUserQuestion( + input.workspaceId, + input.toolCallId, + input.answers + ); + + if (!result.success) { + return { success: false, error: result.error }; + } + + return { success: true, data: undefined }; + }), resumeStream: t .input(schemas.workspace.resumeStream.input) .output(schemas.workspace.resumeStream.output) diff --git a/src/node/services/aiService.ts b/src/node/services/aiService.ts index 31571929dc..5c230b9dc8 100644 --- a/src/node/services/aiService.ts +++ b/src/node/services/aiService.ts @@ -941,6 +941,8 @@ export class AIService extends EventEmitter { // This is idempotent - won't double-commit if already in chat.jsonl await this.partialService.commitToHistory(workspaceId); + const uiMode: UIMode | undefined = + mode === "plan" ? "plan" : mode === "exec" ? "exec" : undefined; const effectiveMuxProviderOptions: MuxProviderOptions = muxProviderOptions ?? {}; // For xAI models, swap between reasoning and non-reasoning variants based on thinkingLevel @@ -982,6 +984,7 @@ export class AIService extends EventEmitter { runtime: earlyRuntime, runtimeTempDir: os.tmpdir(), secrets: {}, + mode: uiMode, }, "", // Empty workspace ID for early stub config this.initStateManager, @@ -1170,11 +1173,13 @@ export class AIService extends EventEmitter { const runtimeTempDir = await this.streamManager.createTempDirForStream(streamToken, runtime); // Extract tool-specific instructions from AGENTS.md files + // Pass uiMode so ask_user_question is included in plan mode const toolInstructions = await readToolInstructions( metadata, runtime, workspacePath, - modelString + modelString, + uiMode === "plan" ? "plan" : "exec" ); // Get model-specific tools with workspace path (correct for local or remote) @@ -1198,7 +1203,7 @@ export class AIService extends EventEmitter { // Plan/exec mode configuration for plan file access. // - read: plan file is readable in all modes (useful context) // - write: enforced by file_edit_* tools (plan file is read-only outside plan mode) - mode: mode as UIMode | undefined, + mode: uiMode, planFilePath, workspaceId, // External edit detection callback diff --git a/src/node/services/askUserQuestionManager.test.ts b/src/node/services/askUserQuestionManager.test.ts new file mode 100644 index 0000000000..5151ce0c3a --- /dev/null +++ b/src/node/services/askUserQuestionManager.test.ts @@ -0,0 +1,67 @@ +import { describe, expect, it } from "bun:test"; + +import { AskUserQuestionManager } from "@/node/services/askUserQuestionManager"; + +const QUESTIONS = [ + { + question: "What should we do?", + header: "Next", + options: [ + { label: "A", description: "Option A" }, + { label: "B", description: "Option B" }, + ], + multiSelect: false, + }, +]; + +describe("AskUserQuestionManager", () => { + it("resolves when answered", async () => { + const manager = new AskUserQuestionManager(); + + const promise = manager.registerPending("ws", "tool-1", [...QUESTIONS]); + manager.answer("ws", "tool-1", { "What should we do?": "A" }); + + const answers = await promise; + expect(answers).toEqual({ "What should we do?": "A" }); + expect(manager.getLatestPending("ws")).toBeNull(); + }); + + it("rejects when canceled", async () => { + const manager = new AskUserQuestionManager(); + + const promise = manager.registerPending("ws", "tool-1", [...QUESTIONS]); + + // Attach handler *before* cancel to avoid Bun treating the rejection as unhandled. + const caught = promise.catch((err: unknown) => err); + + manager.cancel("ws", "tool-1", "User canceled"); + + const error = await caught; + expect(error).toBeInstanceOf(Error); + expect((error as Error).message).toContain("User canceled"); + expect(manager.getLatestPending("ws")).toBeNull(); + }); + + it("tracks latest pending per workspace", async () => { + const manager = new AskUserQuestionManager(); + + const promise1 = manager.registerPending("ws", "tool-1", [...QUESTIONS]); + await new Promise((r) => setTimeout(r, 5)); + const promise2 = manager.registerPending("ws", "tool-2", [...QUESTIONS]); + + expect(manager.getLatestPending("ws")?.toolCallId).toEqual("tool-2"); + + // Attach handlers *before* cancel to avoid Bun treating the rejection as unhandled. + const caught1 = promise1.catch((err: unknown) => err); + const caught2 = promise2.catch((err: unknown) => err); + + manager.cancel("ws", "tool-1", "cleanup"); + manager.cancel("ws", "tool-2", "cleanup"); + + const error1 = await caught1; + const error2 = await caught2; + + expect(error1).toBeInstanceOf(Error); + expect(error2).toBeInstanceOf(Error); + }); +}); diff --git a/src/node/services/askUserQuestionManager.ts b/src/node/services/askUserQuestionManager.ts new file mode 100644 index 0000000000..a2352b8b5f --- /dev/null +++ b/src/node/services/askUserQuestionManager.ts @@ -0,0 +1,140 @@ +import assert from "node:assert/strict"; + +import type { AskUserQuestionQuestion } from "@/common/types/tools"; + +export interface PendingAskUserQuestion { + toolCallId: string; + questions: AskUserQuestionQuestion[]; +} + +interface PendingAskUserQuestionInternal extends PendingAskUserQuestion { + createdAt: number; + resolve: (answers: Record) => void; + reject: (error: Error) => void; +} + +export class AskUserQuestionManager { + private pendingByWorkspace = new Map>(); + + registerPending( + workspaceId: string, + toolCallId: string, + questions: AskUserQuestionQuestion[] + ): Promise> { + assert(workspaceId.length > 0, "workspaceId must be non-empty"); + assert(toolCallId.length > 0, "toolCallId must be non-empty"); + assert(Array.isArray(questions) && questions.length > 0, "questions must be a non-empty array"); + + const workspaceMap = this.getOrCreateWorkspaceMap(workspaceId); + assert( + !workspaceMap.has(toolCallId), + `ask_user_question already pending for toolCallId=${toolCallId}` + ); + + return new Promise>((resolve, reject) => { + const entry: PendingAskUserQuestionInternal = { + toolCallId, + questions, + createdAt: Date.now(), + resolve, + reject, + }; + + workspaceMap.set(toolCallId, entry); + }).finally(() => { + // Ensure cleanup no matter how the promise resolves. + this.deletePending(workspaceId, toolCallId); + }); + } + + answer(workspaceId: string, toolCallId: string, answers: Record): void { + assert(workspaceId.length > 0, "workspaceId must be non-empty"); + assert(toolCallId.length > 0, "toolCallId must be non-empty"); + assert(answers && typeof answers === "object", "answers must be an object"); + + const entry = this.getPending(workspaceId, toolCallId); + entry.resolve(answers); + } + + cancel(workspaceId: string, toolCallId: string, reason: string): void { + assert(workspaceId.length > 0, "workspaceId must be non-empty"); + assert(toolCallId.length > 0, "toolCallId must be non-empty"); + assert(reason.length > 0, "reason must be non-empty"); + + const entry = this.getPending(workspaceId, toolCallId); + entry.reject(new Error(reason)); + } + + cancelAll(workspaceId: string, reason: string): void { + assert(workspaceId.length > 0, "workspaceId must be non-empty"); + assert(reason.length > 0, "reason must be non-empty"); + + const workspaceMap = this.pendingByWorkspace.get(workspaceId); + if (!workspaceMap) { + return; + } + + for (const toolCallId of workspaceMap.keys()) { + // cancel() will delete from map via finally cleanup + this.cancel(workspaceId, toolCallId, reason); + } + } + + getLatestPending(workspaceId: string): PendingAskUserQuestion | null { + assert(workspaceId.length > 0, "workspaceId must be non-empty"); + + const workspaceMap = this.pendingByWorkspace.get(workspaceId); + if (!workspaceMap || workspaceMap.size === 0) { + return null; + } + + let latest: PendingAskUserQuestionInternal | null = null; + for (const entry of workspaceMap.values()) { + if (!latest || entry.createdAt > latest.createdAt) { + latest = entry; + } + } + + assert(latest !== null, "Expected latest pending entry to be non-null"); + + return { + toolCallId: latest.toolCallId, + questions: latest.questions, + }; + } + + private getOrCreateWorkspaceMap( + workspaceId: string + ): Map { + let workspaceMap = this.pendingByWorkspace.get(workspaceId); + if (!workspaceMap) { + workspaceMap = new Map(); + this.pendingByWorkspace.set(workspaceId, workspaceMap); + } + return workspaceMap; + } + + private getPending(workspaceId: string, toolCallId: string): PendingAskUserQuestionInternal { + const workspaceMap = this.pendingByWorkspace.get(workspaceId); + assert(workspaceMap, `No pending ask_user_question entries for workspaceId=${workspaceId}`); + + const entry = workspaceMap.get(toolCallId); + assert(entry, `No pending ask_user_question entry for toolCallId=${toolCallId}`); + + return entry; + } + + private deletePending(workspaceId: string, toolCallId: string): void { + const workspaceMap = this.pendingByWorkspace.get(workspaceId); + if (!workspaceMap) { + return; + } + + workspaceMap.delete(toolCallId); + if (workspaceMap.size === 0) { + this.pendingByWorkspace.delete(workspaceId); + } + } +} + +export const askUserQuestionManager = new AskUserQuestionManager(); diff --git a/src/node/services/systemMessage.ts b/src/node/services/systemMessage.ts index 4d051bb59f..9f05024b1c 100644 --- a/src/node/services/systemMessage.ts +++ b/src/node/services/systemMessage.ts @@ -148,14 +148,16 @@ function getSystemDirectory(): string { * @param globalInstructions Global instructions from ~/.mux/AGENTS.md * @param contextInstructions Context instructions from workspace/project AGENTS.md * @param modelString Active model identifier to determine available tools + * @param mode Optional mode ("plan" | "exec") - affects which tools are available * @returns Map of tool names to their additional instructions */ export function extractToolInstructions( globalInstructions: string | null, contextInstructions: string | null, - modelString: string + modelString: string, + mode?: "plan" | "exec" ): Record { - const availableTools = getAvailableTools(modelString); + const availableTools = getAvailableTools(modelString, mode); const toolInstructions: Record = {}; for (const toolName of availableTools) { @@ -181,13 +183,15 @@ export function extractToolInstructions( * @param runtime - Runtime for reading workspace files (supports SSH) * @param workspacePath - Workspace directory path * @param modelString - Active model identifier to determine available tools + * @param mode Optional mode ("plan" | "exec") - affects which tools are available * @returns Map of tool names to their additional instructions */ export async function readToolInstructions( metadata: WorkspaceMetadata, runtime: Runtime, workspacePath: string, - modelString: string + modelString: string, + mode?: "plan" | "exec" ): Promise> { const [globalInstructions, contextInstructions] = await readInstructionSources( metadata, @@ -195,7 +199,7 @@ export async function readToolInstructions( workspacePath ); - return extractToolInstructions(globalInstructions, contextInstructions, modelString); + return extractToolInstructions(globalInstructions, contextInstructions, modelString, mode); } /** diff --git a/src/node/services/tools/ask_user_question.ts b/src/node/services/tools/ask_user_question.ts new file mode 100644 index 0000000000..3507a20332 --- /dev/null +++ b/src/node/services/tools/ask_user_question.ts @@ -0,0 +1,66 @@ +import assert from "node:assert/strict"; + +import { tool } from "ai"; + +import type { AskUserQuestionToolResult } from "@/common/types/tools"; +import type { ToolConfiguration, ToolFactory } from "@/common/utils/tools/tools"; +import { TOOL_DEFINITIONS } from "@/common/utils/tools/toolDefinitions"; +import { askUserQuestionManager } from "@/node/services/askUserQuestionManager"; + +export const createAskUserQuestionTool: ToolFactory = (config: ToolConfiguration) => { + return tool({ + description: TOOL_DEFINITIONS.ask_user_question.description, + inputSchema: TOOL_DEFINITIONS.ask_user_question.schema, + execute: async (args, { abortSignal, toolCallId }): Promise => { + // Claude Code allows passing pre-filled answers directly. If provided, we can short-circuit + // and return immediately without prompting. + if (args.answers && Object.keys(args.answers).length > 0) { + return { questions: args.questions, answers: args.answers }; + } + + assert(config.workspaceId, "ask_user_question requires a workspaceId"); + assert(toolCallId, "ask_user_question requires toolCallId"); + + const pendingPromise = askUserQuestionManager.registerPending( + config.workspaceId, + toolCallId, + args.questions + ); + + if (!abortSignal) { + const answers = await pendingPromise; + return { questions: args.questions, answers }; + } + + if (abortSignal.aborted) { + // Ensure we don't leak a pending prompt entry. + try { + askUserQuestionManager.cancel(config.workspaceId, toolCallId, "Interrupted"); + } catch { + // ignore + } + throw new Error("Interrupted"); + } + + const abortPromise = new Promise>((_, reject) => { + abortSignal.addEventListener( + "abort", + () => { + try { + askUserQuestionManager.cancel(config.workspaceId!, toolCallId, "Interrupted"); + } catch { + // ignore + } + reject(new Error("Interrupted")); + }, + { once: true } + ); + }); + + const answers = await Promise.race([pendingPromise, abortPromise]); + assert(answers && typeof answers === "object", "Expected answers to be an object"); + + return { questions: args.questions, answers }; + }, + }); +}; diff --git a/src/node/services/workspaceService.ts b/src/node/services/workspaceService.ts index 64cf56fa9b..2eb9977635 100644 --- a/src/node/services/workspaceService.ts +++ b/src/node/services/workspaceService.ts @@ -5,6 +5,7 @@ import assert from "@/common/utils/assert"; import type { Config } from "@/node/config"; import type { Result } from "@/common/types/result"; import { Ok, Err } from "@/common/types/result"; +import { askUserQuestionManager } from "@/node/services/askUserQuestionManager"; import { log } from "@/node/services/log"; import { AgentSession } from "@/node/services/agentSession"; import type { HistoryService } from "@/node/services/historyService"; @@ -977,6 +978,23 @@ export class WorkspaceService extends EventEmitter { void this.updateRecencyTimestamp(workspaceId); if (this.aiService.isStreaming(workspaceId) && !options?.editMessageId) { + const pendingAskUserQuestion = askUserQuestionManager.getLatestPending(workspaceId); + if (pendingAskUserQuestion) { + try { + askUserQuestionManager.cancel( + workspaceId, + pendingAskUserQuestion.toolCallId, + "User responded in chat; questions canceled" + ); + } catch (error) { + log.debug("Failed to cancel pending ask_user_question", { + workspaceId, + toolCallId: pendingAskUserQuestion.toolCallId, + error: error instanceof Error ? error.message : String(error), + }); + } + } + session.queueMessage(message, options); return Ok(undefined); } @@ -1090,6 +1108,20 @@ export class WorkspaceService extends EventEmitter { } } + answerAskUserQuestion( + workspaceId: string, + toolCallId: string, + answers: Record + ): Result { + try { + askUserQuestionManager.answer(workspaceId, toolCallId, answers); + return Ok(undefined); + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + return Err(errorMessage); + } + } + clearQueue(workspaceId: string): Result { try { const session = this.getOrCreateSession(workspaceId); @@ -1124,23 +1156,36 @@ export class WorkspaceService extends EventEmitter { ? expandTildeForSSH(legacyPlanPath) : shellQuote(expandTilde(legacyPlanPath)); - // Delete plan files through runtime (supports both local and SSH) - const runtime = createRuntime(metadata.runtimeConfig, { - projectPath: metadata.projectPath, - }); - - try { - // Use exec to delete files since runtime doesn't have a deleteFile method. - // Delete both paths in one command for efficiency. - const execStream = await runtime.exec(`rm -f ${quotedPlanPath} ${quotedLegacyPlanPath}`, { - cwd: metadata.projectPath, - timeout: 10, + // SSH runtime: delete via remote shell so $HOME expands on the remote. + if (isSSHRuntime(metadata.runtimeConfig)) { + const runtime = createRuntime(metadata.runtimeConfig, { + projectPath: metadata.projectPath, }); - // Wait for completion so callers can rely on the plan file actually being removed. - await execStream.exitCode; - } catch { - // Plan files don't exist or can't be deleted - ignore + + try { + // Use exec to delete files since runtime doesn't have a deleteFile method. + // Delete both paths in one command for efficiency. + const execStream = await runtime.exec(`rm -f ${quotedPlanPath} ${quotedLegacyPlanPath}`, { + cwd: metadata.projectPath, + timeout: 10, + }); + // Wait for completion so callers can rely on the plan file actually being removed. + await execStream.exitCode; + } catch { + // Plan files don't exist or can't be deleted - ignore + } + + return; } + + // Local runtimes: delete directly on the local filesystem. + const planPathAbs = expandTilde(planPath); + const legacyPlanPathAbs = expandTilde(legacyPlanPath); + + await Promise.allSettled([ + fsPromises.rm(planPathAbs, { force: true }), + fsPromises.rm(legacyPlanPathAbs, { force: true }), + ]); } async truncateHistory(workspaceId: string, percentage?: number): Promise> {