diff --git a/apps/web/src/app/(sandbox)/task/[taskId]/WakeTaskInput.client.test.tsx b/apps/web/src/app/(sandbox)/task/[taskId]/WakeTaskInput.client.test.tsx index 1dce9673e..7dd35aa38 100644 --- a/apps/web/src/app/(sandbox)/task/[taskId]/WakeTaskInput.client.test.tsx +++ b/apps/web/src/app/(sandbox)/task/[taskId]/WakeTaskInput.client.test.tsx @@ -9,6 +9,7 @@ const { removeOptimisticQueuedMessageMock, restoreMutateAsyncMock, toastErrorMock, + useAuthorizedUserMock, useSandboxCurrentUserInfoMock, } = vi.hoisted(() => ({ appendOptimisticAcpEventMock: vi.fn(), @@ -18,6 +19,7 @@ const { removeOptimisticQueuedMessageMock: vi.fn(), restoreMutateAsyncMock: vi.fn(), toastErrorMock: vi.fn(), + useAuthorizedUserMock: vi.fn(), useSandboxCurrentUserInfoMock: vi.fn(), })); @@ -26,6 +28,7 @@ let capturedSuggestion: unknown; let capturedSubmitWithMetaKey: boolean | undefined; let capturedSubmitIcon: unknown; let capturedSurface: string | undefined; +let capturedModelSwitcherDisabled: boolean | undefined; const submittedFilesRef: { current: Array<{ url?: string; @@ -54,16 +57,7 @@ vi.mock('@/hooks/snapshots', () => ({ })); vi.mock('@/hooks/useUser', () => ({ - useAuthorizedUser: () => ({ - managedAccess: { - state: 'active', - reason: null, - revision: 1, - effectiveAt: '2026-01-01T00:00:00.000Z', - restrictionStartsAt: null, - remediationUrl: null, - }, - }), + useAuthorizedUser: useAuthorizedUserMock, })); vi.mock('@/trpc/client', () => ({ @@ -99,6 +93,7 @@ vi.mock('@/components/tasks', () => ({ suggestion, submitWithMetaKey, submitIcon, + tools, surface, submitDisabledReason, }: { @@ -117,6 +112,7 @@ vi.mock('@/components/tasks', () => ({ suggestion?: unknown; submitWithMetaKey?: boolean; submitIcon?: unknown; + tools?: React.ReactNode; surface?: string; submitDisabledReason?: string; }) => { @@ -145,11 +141,51 @@ vi.mock('@/components/tasks', () => ({ > Send + {tools} ); }, })); +vi.mock('./prompt-input/TaskModelSwitcher', () => ({ + TaskModelSwitcher: ({ + disabled, + onPendingChange, + }: { + disabled?: boolean; + onPendingChange?: (pending: boolean) => void; + }) => { + capturedModelSwitcherDisabled = disabled; + + return ( + + ); + }, +})); + +vi.mock('./sidebar-actions/TaskToolsButton', () => ({ + TaskToolsMenu: ({ + disabled, + onSelect, + }: { + disabled?: boolean; + onSelect: (actionId: 'simplify') => void; + }) => ( +
+ +
+ ), +})); + import { WakeTaskInput } from './WakeTaskInput'; function renderWithQueryClient(ui: React.ReactNode, queryClient: QueryClient) { @@ -166,6 +202,7 @@ describe('WakeTaskInput', () => { capturedSubmitWithMetaKey = undefined; capturedSubmitIcon = undefined; capturedSurface = undefined; + capturedModelSwitcherDisabled = undefined; submittedFilesRef.current = []; preparePromptAttachmentsMock.mockResolvedValue({ text: 'Wake up and keep going', @@ -175,6 +212,16 @@ describe('WakeTaskInput', () => { runId: 84, taskId: 'task-42', }); + useAuthorizedUserMock.mockReturnValue({ + managedAccess: { + state: 'active', + reason: null, + revision: 1, + effectiveAt: '2026-01-01T00:00:00.000Z', + restrictionStartsAt: null, + remediationUrl: null, + }, + }); useSandboxCurrentUserInfoMock.mockReturnValue(null); }); @@ -219,6 +266,121 @@ describe('WakeTaskInput', () => { expect(capturedSurface).toBe('embedded'); }); + it('allows model selection before waking an OpenCode task', async () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + }, + }); + + renderWithQueryClient( + , + queryClient, + ); + + expect( + screen.getByRole('button', { name: 'Model selector' }), + ).toBeVisible(); + expect(capturedModelSwitcherDisabled).toBe(false); + + fireEvent.click(screen.getByRole('button', { name: 'Model selector' })); + + await waitFor(() => { + expect(screen.getByRole('button', { name: 'Send' })).toBeDisabled(); + expect(capturedModelSwitcherDisabled).toBe(true); + }); + + fireEvent.click(screen.getByRole('button', { name: 'Send' })); + expect(restoreMutateAsyncMock).not.toHaveBeenCalled(); + }); + + it('runs a task tool as the wake-up prompt', async () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + }, + }); + preparePromptAttachmentsMock.mockImplementation( + ({ text }: { text: string }) => Promise.resolve({ text }), + ); + + renderWithQueryClient( + , + queryClient, + ); + + fireEvent.click( + screen.getByRole('button', { name: 'Simplify changed code' }), + ); + + await waitFor(() => { + expect(restoreMutateAsyncMock).toHaveBeenCalledWith({ + sourceSnapshotId: 'snap-42', + sourceRunId: 42, + clientMessageId: expect.any(String), + resumePrompt: '$simplify', + }); + }); + + expect(appendOptimisticAcpEventMock).toHaveBeenCalledWith( + expect.objectContaining({ + text: '$simplify', + }), + ); + }); + + it('disables task tools when managed access blocks waking the task', () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + }, + }); + useAuthorizedUserMock.mockReturnValue({ + managedAccess: { + state: 'read_only', + reason: 'billing_required', + revision: 2, + effectiveAt: '2026-01-02T00:00:00.000Z', + restrictionStartsAt: null, + remediationUrl: null, + }, + }); + + renderWithQueryClient( + , + queryClient, + ); + + expect(screen.getByRole('button', { name: 'Task Tools' })).toBeDisabled(); + fireEvent.click( + screen.getByRole('button', { name: 'Simplify changed code' }), + ); + expect(restoreMutateAsyncMock).not.toHaveBeenCalled(); + }); + it('prefills the sleeping draft, appends an optimistic transcript row, and resumes the task with a deferred prompt', async () => { const queryClient = new QueryClient({ defaultOptions: { diff --git a/apps/web/src/app/(sandbox)/task/[taskId]/WakeTaskInput.tsx b/apps/web/src/app/(sandbox)/task/[taskId]/WakeTaskInput.tsx index c15e8a651..b160792a3 100644 --- a/apps/web/src/app/(sandbox)/task/[taskId]/WakeTaskInput.tsx +++ b/apps/web/src/app/(sandbox)/task/[taskId]/WakeTaskInput.tsx @@ -3,7 +3,11 @@ import { useEffect, useState } from 'react'; import { toast } from 'sonner'; -import { DEFAULT_MANAGED_DEPLOYMENT_ACCESS } from '@roomote/types'; +import { + DEFAULT_MANAGED_DEPLOYMENT_ACCESS, + getTaskToolInvocation, + type TaskToolActionId, +} from '@roomote/types'; import type { PromptInputMessage } from '@/components/ai-elements'; import { TaskPromptInput } from '@/components/tasks'; @@ -16,9 +20,13 @@ import type { TaskRunDetail } from '@/lib/server'; import { cn } from '@/lib/utils'; import { useOptimisticPromptSubmission } from './prompt-input/useOptimisticPromptSubmission'; +import { TaskModelSwitcher } from './prompt-input/TaskModelSwitcher'; +import { TaskToolsMenu } from './sidebar-actions/TaskToolsButton'; +import { shouldShowTaskToolsActions } from './sidebar-actions/utils'; interface WakeTaskInputProps { - taskRun: Pick; + taskRun: Pick & + Partial>; initialPrompt?: string; embedded?: boolean; } @@ -34,6 +42,7 @@ export function WakeTaskInput({ } = useOptimisticPromptSubmission(); const [promptText, setPromptText] = useState(initialPrompt); const [sending, setSending] = useState(false); + const [modelSettingsPending, setModelSettingsPending] = useState(false); const { managedAccess = DEFAULT_MANAGED_DEPLOYMENT_ACCESS } = useAuthorizedUser(); const taskLaunchDisabledReason = getTaskLaunchDisabledReason(managedAccess); @@ -41,7 +50,7 @@ export function WakeTaskInput({ const restore = useRestoreTaskRunSnapshot({ onSuccess: () => setPromptText(''), }); - const isBusy = sending || restore.isPending; + const isBusy = sending || restore.isPending || modelSettingsPending; useEffect(() => { setPromptText(initialPrompt); @@ -126,6 +135,13 @@ export function WakeTaskInput({ } }; + const handleTaskToolSelect = (actionId: TaskToolActionId) => { + void handleSubmit({ + text: getTaskToolInvocation(actionId, taskRun.harness), + files: [], + }); + }; + const input = (
: undefined} + tools={ + <> + {shouldShowTaskToolsActions(taskRun.payloadKind) && ( + + )} + {taskRun.harness === 'opencode-server' && ( + + )} + + } surface={embedded ? 'embedded' : 'default'} submitDisabledReason={taskLaunchDisabledReason} /> diff --git a/apps/web/src/app/(sandbox)/task/[taskId]/prompt-input/TaskModelSwitcher.tsx b/apps/web/src/app/(sandbox)/task/[taskId]/prompt-input/TaskModelSwitcher.tsx index 9502fc317..26a615475 100644 --- a/apps/web/src/app/(sandbox)/task/[taskId]/prompt-input/TaskModelSwitcher.tsx +++ b/apps/web/src/app/(sandbox)/task/[taskId]/prompt-input/TaskModelSwitcher.tsx @@ -1,6 +1,6 @@ 'use client'; -import { useEffect, useMemo, useState } from 'react'; +import { useEffect, useMemo, useRef, useState } from 'react'; import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; import { toast } from 'sonner'; @@ -60,9 +60,12 @@ type RoleSelection = { export function TaskModelSwitcher({ taskRun, disabled = false, + onPendingChange, }: { - taskRun: TaskRunDetail; + taskRun: Pick & + Partial>; disabled?: boolean; + onPendingChange?: (pending: boolean) => void; }) { const trpc = useTRPC(); const queryClient = useQueryClient(); @@ -73,6 +76,20 @@ export function TaskModelSwitcher({ const [localSelections, setLocalSelections] = useState< Partial> >({}); + const pendingOperationsRef = useRef(0); + + const startPendingOperation = () => { + pendingOperationsRef.current += 1; + onPendingChange?.(true); + }; + + const finishPendingOperation = () => { + pendingOperationsRef.current = Math.max( + 0, + pendingOperationsRef.current - 1, + ); + onPendingChange?.(pendingOperationsRef.current > 0); + }; const { data: roleDefaults } = useQuery( trpc.taskModels.roleDefaults.queryOptions(undefined, { enabled: open }), @@ -137,6 +154,7 @@ export function TaskModelSwitcher({ }); toast.error(error.message || 'Failed to update the model settings'); }, + onSettled: finishPendingOperation, }), ); @@ -148,6 +166,7 @@ export function TaskModelSwitcher({ const applyRoleSelection = (role: SwitcherRole, selection: RoleSelection) => { setLocalSelections((current) => ({ ...current, [role]: selection })); + startPendingOperation(); updateModelSelection.mutate({ taskId: taskRun.taskId, role, @@ -215,6 +234,7 @@ export function TaskModelSwitcher({ ]; const cleared: RoleSelection = { model: null, reasoningEffort: null }; + startPendingOperation(); setResetting(true); setLocalSelections( Object.fromEntries(rolesToClear.map((role) => [role, cleared])), @@ -238,6 +258,7 @@ export function TaskModelSwitcher({ ); } finally { setResetting(false); + finishPendingOperation(); void invalidateSession(); } }; @@ -401,7 +422,7 @@ export function TaskModelSwitcher({ size="sm" className="text-xs font-medium" onClick={() => void handleReset()} - disabled={disabled || resetting} + disabled={disabled || updateModelSelection.isPending || resetting} > Defaults diff --git a/apps/web/src/app/(sandbox)/task/[taskId]/sidebar-actions/TaskToolsButton.tsx b/apps/web/src/app/(sandbox)/task/[taskId]/sidebar-actions/TaskToolsButton.tsx index ca0df70be..debd60615 100644 --- a/apps/web/src/app/(sandbox)/task/[taskId]/sidebar-actions/TaskToolsButton.tsx +++ b/apps/web/src/app/(sandbox)/task/[taskId]/sidebar-actions/TaskToolsButton.tsx @@ -3,6 +3,8 @@ import { Fragment, memo } from 'react'; import { toast } from 'sonner'; +import type { TaskToolActionId } from '@roomote/types'; + import { useUser } from '@/hooks/useUser'; import { generateClientUuid } from '@/lib/client-uuid'; import { useTRPCClient } from '@/trpc/client'; @@ -31,6 +33,44 @@ import { TASK_TOOL_CATALOG } from '../task-tools'; import { type SidebarActionBaseProps } from './types'; import { isTaskRunAsleep } from './utils'; +export function TaskToolsMenu({ + onSelect, + disabled = false, +}: { + onSelect: (actionId: TaskToolActionId) => void | Promise; + disabled?: boolean; +}) { + return ( + + + + + + + + Task Tools + {TASK_TOOL_CATALOG.map(({ actionId, label, separator, icon: Icon }) => ( + + {separator && } + void onSelect(actionId)} + className="flex items-center gap-2 cursor-pointer min-w-64" + > + + {label} + + + ))} + + + ); +} + function TaskToolsButtonBase({ taskRun, }: Pick) { @@ -52,54 +92,28 @@ function TaskToolsButtonBase({ return null; } - const content = ( - <> - Task Tools - {TASK_TOOL_CATALOG.map(({ actionId, label, separator, icon: Icon }) => ( - - {separator && } - { - const clientMessageId = generateClientUuid(); - - try { - await trpcClient.sandboxSession.sendPrompt.mutate({ - taskId: taskRun.taskId, - taskTool: { actionId }, - source: 'web', - clientMessageId, - userImageUrl, - }); - } catch (error) { - toast.error( - error instanceof Error - ? error.message - : 'Failed to send task tool.', - ); - } - }} - className="flex items-center gap-2 cursor-pointer min-w-64" - > - - {label} - - - ))} - - ); - return ( - - - - - - - {content} - + { + const clientMessageId = generateClientUuid(); + + try { + await trpcClient.sandboxSession.sendPrompt.mutate({ + taskId: taskRun.taskId, + taskTool: { actionId }, + source: 'web', + clientMessageId, + userImageUrl, + }); + } catch (error) { + toast.error( + error instanceof Error + ? error.message + : 'Failed to send task tool.', + ); + } + }} + /> ); } diff --git a/apps/web/src/components/tasks/TaskPromptInput.client.test.tsx b/apps/web/src/components/tasks/TaskPromptInput.client.test.tsx index a361a166a..741717e9e 100644 --- a/apps/web/src/components/tasks/TaskPromptInput.client.test.tsx +++ b/apps/web/src/components/tasks/TaskPromptInput.client.test.tsx @@ -37,4 +37,19 @@ describe('TaskPromptInput', () => { expect(screen.getByRole('button', { name: 'Submit' })).toBeEnabled(); }); + + it('renders additional prompt tools', () => { + render( + {}} + onSubmit={() => {}} + placeholder="Describe a task" + tools={} + />, + ); + + expect(screen.getByRole('button', { name: 'Choose model' })).toBeVisible(); + }); }); diff --git a/apps/web/src/components/tasks/TaskPromptInput.tsx b/apps/web/src/components/tasks/TaskPromptInput.tsx index 069224adb..b53c9f6ba 100644 --- a/apps/web/src/components/tasks/TaskPromptInput.tsx +++ b/apps/web/src/components/tasks/TaskPromptInput.tsx @@ -112,6 +112,8 @@ type TaskPromptInputProps = { animateContainer?: boolean; /** Optional content rendered inside the prompt box, below the input. */ suggestion?: ReactNode; + /** Optional actions rendered alongside the attachment menu. */ + tools?: ReactNode; /** Optional reason that disables the submit button and explains why. */ submitDisabledReason?: string; /** When true, submit on Cmd/Ctrl+Enter instead of plain Enter. */ @@ -131,6 +133,7 @@ export function TaskPromptInput({ textareaMaxHeight, animateContainer = true, suggestion, + tools, submitDisabledReason, submitWithMetaKey = true, submitIcon, @@ -187,6 +190,7 @@ export function TaskPromptInput({ + {tools}