From dfbffed030c39e786cd2f80919c37a2c34427ce2 Mon Sep 17 00:00:00 2001 From: Declan Brady Date: Thu, 16 Jul 2026 11:11:14 -0400 Subject: [PATCH] feat(ui): Stop button + queue-at-boundary for task interrupt (AGX1-391) Dual-mode composer button (Stop when a turn is streaming and the input is empty; Send that enqueues when there is text), Esc-to-interrupt, a queued-message chip flushed on the running->idle edge, and INTERRUPTED treated as a non-terminal status (composer stays enabled). Exposes the sync message/send AbortController so Stop tears down the stream, and settles the rpcStatus meta cache to idle so no thinking indicator hangs. The interrupt call goes through a thin adapter (POST /tasks/{id}/interrupt) with a TODO to swap to the typed SDK method once regenerated. Part of AGX1-391; depends on the backend PR scale-agentex#365. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../primary-content/prompt-input.tsx | 229 +++++++++++++----- agentex-ui/hooks/use-task-messages.ts | 157 +++++++++--- agentex-ui/lib/types.ts | 16 ++ 3 files changed, 302 insertions(+), 100 deletions(-) diff --git a/agentex-ui/components/primary-content/prompt-input.tsx b/agentex-ui/components/primary-content/prompt-input.tsx index 81738fd8..29d68628 100644 --- a/agentex-ui/components/primary-content/prompt-input.tsx +++ b/agentex-ui/components/primary-content/prompt-input.tsx @@ -8,7 +8,7 @@ import { Prec } from '@codemirror/state'; import { EditorView, keymap } from '@codemirror/view'; import CodeMirror from '@uiw/react-codemirror'; import { DataContent, TextContent } from 'agentex/resources'; -import { ArrowUp } from 'lucide-react'; +import { ArrowUp, Square, X } from 'lucide-react'; import { useAgentexClient } from '@/components/providers'; import { IconButton } from '@/components/ui/icon-button'; @@ -19,9 +19,13 @@ import { SearchParamKey, useSafeSearchParams, } from '@/hooks/use-safe-search-params'; -import { useSendMessage } from '@/hooks/use-task-messages'; +import { + useInterruptTurn, + useSendMessage, + useTaskMessages, +} from '@/hooks/use-task-messages'; import { useTask } from '@/hooks/use-tasks'; -import { TaskStatusEnum } from '@/lib/types'; +import { isTaskTerminalStatus } from '@/lib/types'; type PromptInputProps = { prompt: string; @@ -51,20 +55,32 @@ export function PromptInput({ prompt, setPrompt }: PromptInputProps) { const [isSendingJSON, setIsSendingJSON] = useState(false); const [isTaskParamsOpen, setIsTaskParamsOpen] = useState(false); const [taskParams, setTaskParams] = useState(''); + const [queuedMessage, setQueuedMessage] = useState<{ + text: string; + isJson: boolean; + } | null>(null); const taskParamsViewRef = useRef(null); const { agentexClient } = useAgentexClient(); const createTaskMutation = useCreateTask({ agentexClient }); const sendMessageMutation = useSendMessage({ agentexClient }); + const interruptTurnMutation = useInterruptTurn({ agentexClient }); const { data: task } = useTask({ agentexClient, taskId: taskID ?? '' }); + const { rpcStatus } = useTaskMessages({ + agentexClient, + taskId: taskID ?? '', + }); + + const isStreaming = rpcStatus === 'pending'; const textInputRef = useRef(null); const codeMirrorViewRef = useRef(null); + const wasStreamingRef = useRef(false); const isTaskTerminal = useMemo(() => { if (!taskID || !task) return false; - return task.status != null && task.status !== TaskStatusEnum.RUNNING; + return isTaskTerminalStatus(task.status); }, [taskID, task]); const handleSetJson = useCallback( @@ -101,15 +117,80 @@ export function PromptInput({ prompt, setPrompt }: PromptInputProps) { [agentName, isClient, isTaskTerminal] ); + const performSend = useCallback( + async (rawPrompt: string, asJson: boolean) => { + if (!agentName || !rawPrompt.trim()) return; + + if (asJson) { + try { + JSON.parse(rawPrompt); + } catch { + toast.error('Invalid JSON'); + return; + } + } + + let currentTaskId = taskID; + + if (!currentTaskId) { + let extraTaskParams: Record = {}; + if (taskParams.trim()) { + try { + extraTaskParams = JSON.parse(taskParams); + } catch { + toast.error('Invalid Task Parameters JSON'); + return; + } + } + + const task = await createTaskMutation.mutateAsync({ + agentName: agentName, + params: { + ...extraTaskParams, + description: rawPrompt, + content: rawPrompt, + }, + }); + currentTaskId = task.id; + updateParams({ [SearchParamKey.TASK_ID]: currentTaskId }); + } + + const content: TextContent | DataContent = asJson + ? { + type: 'data', + author: 'user', + data: JSON.parse(rawPrompt) as Record, + } + : { + type: 'text', + author: 'user', + format: 'plain', + attachments: [], + content: rawPrompt, + }; + + await sendMessageMutation.mutateAsync({ + taskId: currentTaskId, + agentName, + content, + }); + }, + [ + taskID, + agentName, + createTaskMutation, + updateParams, + sendMessageMutation, + taskParams, + ] + ); + const handleSendPrompt = useCallback(async () => { if (isDisabled || !prompt.trim() || !agentName) { toast.error('Please select an agent and enter a prompt'); return; } - let currentTaskId = taskID; - const currentPrompt = prompt; - if (isSendingJSON) { try { JSON.parse(prompt); @@ -119,62 +200,71 @@ export function PromptInput({ prompt, setPrompt }: PromptInputProps) { } } + const text = prompt; + const asJson = isSendingJSON; setPrompt(''); + await performSend(text, asJson); + }, [isDisabled, prompt, agentName, isSendingJSON, setPrompt, performSend]); - if (!currentTaskId) { - let extraTaskParams: Record = {}; - if (taskParams.trim()) { - try { - extraTaskParams = JSON.parse(taskParams); - } catch { - toast.error('Invalid Task Parameters JSON'); - return; - } - } + const handleEnqueue = useCallback(() => { + if (!prompt.trim()) return; - const task = await createTaskMutation.mutateAsync({ - agentName: agentName, - params: { - ...extraTaskParams, - description: prompt, - content: currentPrompt, - }, - }); - currentTaskId = task.id; - updateParams({ [SearchParamKey.TASK_ID]: currentTaskId }); + if (isSendingJSON) { + try { + JSON.parse(prompt); + } catch { + toast.error('Invalid JSON'); + return; + } } - const content: TextContent | DataContent = isSendingJSON - ? { - type: 'data', - author: 'user', - data: JSON.parse(prompt) as Record, - } - : { - type: 'text', - author: 'user', - format: 'plain', - attachments: [], - content: prompt as string, - }; - - await sendMessageMutation.mutateAsync({ - taskId: currentTaskId, - agentName: agentName!, - content, + setQueuedMessage({ text: prompt, isJson: isSendingJSON }); + setPrompt(''); + }, [prompt, isSendingJSON, setPrompt]); + + const handleInterrupt = useCallback(() => { + if (!taskID || !agentName) return; + sendMessageMutation.abortStream(taskID); + interruptTurnMutation.mutate({ + taskId: taskID, + agentName, + reason: 'user_interrupt', }); - }, [ - isDisabled, - prompt, - taskID, - agentName, - createTaskMutation, - updateParams, - sendMessageMutation, - setPrompt, - isSendingJSON, - taskParams, - ]); + }, [taskID, agentName, sendMessageMutation, interruptTurnMutation]); + + const handleComposerSubmit = useCallback(() => { + if (isStreaming) { + handleEnqueue(); + } else { + handleSendPrompt(); + } + }, [isStreaming, handleEnqueue, handleSendPrompt]); + + useEffect(() => { + if (wasStreamingRef.current && !isStreaming && queuedMessage) { + const queued = queuedMessage; + setQueuedMessage(null); + performSend(queued.text, queued.isJson); + } + wasStreamingRef.current = isStreaming; + }, [isStreaming, queuedMessage, performSend]); + + useEffect(() => { + if (!isStreaming) return; + const onKeyDown = (e: KeyboardEvent) => { + if (e.key === 'Escape') { + e.preventDefault(); + handleInterrupt(); + } + }; + window.addEventListener('keydown', onKeyDown); + return () => window.removeEventListener('keydown', onKeyDown); + }, [isStreaming, handleInterrupt]); + + const isStopMode = isStreaming && !prompt.trim(); + const isPrimaryDisabled = isStopMode + ? isDisabled + : isDisabled || !prompt.trim(); return (
@@ -199,6 +289,19 @@ export function PromptInput({ prompt, setPrompt }: PromptInputProps) { )}
)} + {queuedMessage && ( +
+ Queued: {queuedMessage.text} + setQueuedMessage(null)} + icon={X} + aria-label="Remove queued message" + /> +
+ )}
@@ -207,7 +310,7 @@ export function PromptInput({ prompt, setPrompt }: PromptInputProps) { prompt={prompt} setPrompt={setPrompt} isDisabled={isDisabled} - handleSendPrompt={handleSendPrompt} + handleSendPrompt={handleComposerSubmit} codeMirrorViewRef={codeMirrorViewRef} /> ) : ( @@ -217,16 +320,16 @@ export function PromptInput({ prompt, setPrompt }: PromptInputProps) { isDisabled={isDisabled} isTaskTerminal={isTaskTerminal} taskStatus={task?.status} - handleSendPrompt={handleSendPrompt} + handleSendPrompt={handleComposerSubmit} inputRef={textInputRef} /> )}
>(new Map()); + + const abortStream = useCallback((taskId: string) => { + const controller = abortControllersRef.current.get(taskId); + if (controller) { + controller.abort(); + abortControllersRef.current.delete(taskId); + } + }, []); + + const mutation = useMutation({ mutationFn: async ({ taskId, agentName, content }: SendMessageParams) => { const queryKey = taskMessagesKeys.byTaskId(taskId); const metaKey = taskMessagesMetaKeys.byTaskId(taskId); @@ -217,44 +230,55 @@ export function useSendMessage({ }); const controller = new AbortController(); - - for await (const response of agentRPCWithStreaming( - agentexClient, - { agentName }, - 'message/send', - { task_id: taskId, content }, - { signal: controller.signal } - )) { - if (response.error != null) { + abortControllersRef.current.set(taskId, controller); + + try { + for await (const response of agentRPCWithStreaming( + agentexClient, + { agentName }, + 'message/send', + { task_id: taskId, content }, + { signal: controller.signal } + )) { + if (response.error != null) { + queryClient.setQueryData(metaKey, { + deltaAccumulator: null, + rpcStatus: 'error', + }); + throw new Error(response.error.message); + } + + const result = aggregateMessageEvents( + latestMessages, + latestDeltaAccumulator, + [response.result] + ); + + latestMessages = result.messages; + latestDeltaAccumulator = result.deltaAccumulator; + + // Write streaming state to first page + queryClient.setQueryData>( + queryKey, + old => updateFirstPage(old, () => latestMessages) + ); queryClient.setQueryData(metaKey, { - deltaAccumulator: null, - rpcStatus: 'error', + deltaAccumulator: latestDeltaAccumulator, + rpcStatus: 'pending', }); - throw new Error(response.error.message); - } - - const result = aggregateMessageEvents( - latestMessages, - latestDeltaAccumulator, - [response.result] - ); - latestMessages = result.messages; - latestDeltaAccumulator = result.deltaAccumulator; - - // Write streaming state to first page - queryClient.setQueryData>( - queryKey, - old => updateFirstPage(old, () => latestMessages) - ); - queryClient.setQueryData(metaKey, { - deltaAccumulator: latestDeltaAccumulator, - rpcStatus: 'pending', - }); - - if (response.result.type === 'done') { - queryClient.invalidateQueries({ queryKey: ['spans', taskId] }); + if (response.result.type === 'done') { + queryClient.invalidateQueries({ queryKey: ['spans', taskId] }); + } + } + } catch (streamError) { + // A user Stop aborts the controller; that is an expected end of the + // stream, not a failure. Any other error propagates to onError. + if (!controller.signal.aborted) { + throw streamError; } + } finally { + abortControllersRef.current.delete(taskId); } // Final reconciliation: refetch all loaded pages so the cache stays @@ -264,9 +288,11 @@ export function useSendMessage({ // when there are no active observers (e.g. component unmounted mid-stream). await queryClient.refetchQueries({ queryKey }); + // On a Stop the shimmer must clear immediately (idle); a normal + // completion settles to success. queryClient.setQueryData(metaKey, { deltaAccumulator: null, - rpcStatus: 'success', + rpcStatus: controller.signal.aborted ? 'idle' : 'success', }); const latestData = @@ -288,4 +314,61 @@ export function useSendMessage({ }); }, }); + + return Object.assign(mutation, { abortStream }); +} + +type InterruptTurnParams = { + taskId: string; + agentName: string; + reason?: string | null; +}; + +/** + * Thin adapter for the assumed `task/interrupt` shape. The typed SDK method is + * not generated yet, so this hits the REST route directly. + * + * TODO(AGX1-391): swap to the typed SDK method once the OpenAPI `task/interrupt` + * change lands, e.g. `agentexClient.tasks.interrupt(taskId, { reason })`, or the + * RPC form `agentRPCNonStreaming(agentexClient, { agentName }, 'task/interrupt', + * { task_id: taskId })`. + */ +async function interruptTaskAdapter( + agentexClient: AgentexSDK, + { taskId, reason }: InterruptTurnParams +): Promise { + await agentexClient.post(`/tasks/${taskId}/interrupt`, { + body: { reason: reason ?? null }, + }); +} + +/** + * Interrupts the current turn without terminating the task. Forces the message + * meta cache to idle on settle so the thinking indicator clears. Does NOT abort + * the SSE task subscription — that stays live so later turns keep streaming. + */ +export function useInterruptTurn({ + agentexClient, +}: { + agentexClient: AgentexSDK; +}) { + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async (params: InterruptTurnParams) => { + return interruptTaskAdapter(agentexClient, params); + }, + onSettled: (_data, _error, variables) => { + queryClient.setQueryData( + taskMessagesMetaKeys.byTaskId(variables.taskId), + { deltaAccumulator: null, rpcStatus: 'idle' } + ); + }, + onError: error => { + toast.error({ + title: 'Failed to stop the current turn', + message: error instanceof Error ? error.message : 'Please try again.', + }); + }, + }); } diff --git a/agentex-ui/lib/types.ts b/agentex-ui/lib/types.ts index 38345e8a..ea6c815d 100644 --- a/agentex-ui/lib/types.ts +++ b/agentex-ui/lib/types.ts @@ -15,11 +15,27 @@ export enum TaskStatusEnum { COMPLETED = 'COMPLETED', FAILED = 'FAILED', RUNNING = 'RUNNING', + INTERRUPTED = 'INTERRUPTED', TERMINATED = 'TERMINATED', TIMED_OUT = 'TIMED_OUT', DELETED = 'DELETED', } +// Non-terminal statuses: the task is still continuable, so the composer stays +// enabled. INTERRUPTED is a resting state ("stopped by user, waiting for the +// next message") and must not be treated as terminal. +const NON_TERMINAL_TASK_STATUSES: ReadonlySet = new Set([ + TaskStatusEnum.RUNNING, + TaskStatusEnum.INTERRUPTED, +]); + +export function isTaskTerminalStatus( + status: string | null | undefined +): boolean { + if (status == null) return false; + return !NON_TERMINAL_TASK_STATUSES.has(status); +} + // API to frontend type checks export const _taskStatusEnumCheck: Record = { [TaskStatusEnum.CANCELED]: true,