fix(wework): preserve composer drafts across tasks

wework/e2e/desktop/task-flow.e2e.mjs

@@ -23,6 +23,9 @@ const REQUEST_USER_INPUT_PROMPT =
 const REQUEST_USER_INPUT_QUESTION = 'Which implementation direction should be used?'
 const REQUEST_USER_INPUT_COMPLETION_TEXT = 'WEWORK_DESKTOP_E2E_REQUEST_INPUT_COMPLETE'
 const SEND_MODE_DRAFT = 'WEWORK_DESKTOP_E2E_SEND_MODE_DRAFT'
+const UNSENT_BLANK_TASK_DRAFT = 'WEWORK_DESKTOP_E2E_UNSENT_BLANK_TASK_DRAFT'
+const UNSENT_FIRST_TASK_DRAFT = 'WEWORK_DESKTOP_E2E_UNSENT_FIRST_TASK_DRAFT'
+const UNSENT_SECOND_TASK_DRAFT = 'WEWORK_DESKTOP_E2E_UNSENT_SECOND_TASK_DRAFT'
 const WINDOW_LIFECYCLE_PROMPT =
   'WEWORK_DESKTOP_E2E_WINDOW_LIFECYCLE: keep this response running until released.'
 const WINDOW_LIFECYCLE_COMPLETION_TEXT = 'WEWORK_DESKTOP_E2E_WINDOW_LIFECYCLE_COMPLETE'
@@ -4033,6 +4036,27 @@ async function main() {
       testId.startsWith('runtime-local-task-row-')
     )
     assert.ok(taskRowTestId, 'The completed task row was not found')
+
+    phase = 'blank-task-draft-restoration'
+    await control.command('click', '[data-testid="new-chat-button"]')
+    await control.command('waitFor', composerSelector, {
+      timeoutMs: WORKBENCH_READY_TIMEOUT_MS,
+    })
+    await control.command('fill', composerSelector, { value: UNSENT_BLANK_TASK_DRAFT })
+    await control.command('click', `[data-testid="${taskRowTestId}"]`)
+    await control.command('waitFor', '[data-testid="message-assistant"]', {
+      text: COMPLETION_TEXT,
+      timeoutMs: UI_TIMEOUT_MS,
+    })
+    await control.command('click', '[data-testid="new-chat-button"]')
+    await waitForControlValue(
+      control,
+      composerSelector,
+      UNSENT_BLANK_TASK_DRAFT,
+      'The blank task lost its unsent composer draft after switching tasks'
+    )
+    await control.command('fill', composerSelector, { value: '' })
+
     if (!REQUEST_INPUT_ONLY) {
       await control.command('click', '[data-testid="new-chat-button"]')
       await control.command('waitFor', composerSelector, {
@@ -4228,6 +4252,11 @@ async function main() {
 
     phase = 'fresh-chat'
     control.setScenario('fresh_chat')
+    const taskRowsBeforeFreshChat = new Set(
+      JSON.parse(await control.command('snapshot', 'body')).testIds.filter(testId =>
+        testId.startsWith('runtime-local-task-row-')
+      )
+    )
     await control.command('click', '[data-testid="new-chat-button"]')
     await control.command('waitFor', composerSelector, {
       timeoutMs: WORKBENCH_READY_TIMEOUT_MS,
@@ -4245,6 +4274,41 @@ async function main() {
       timeoutMs: UI_TIMEOUT_MS,
     })
 
+    phase = 'task-draft-isolation'
+    const secondTaskSnapshot = await waitForSnapshot(
+      control,
+      snapshot =>
+        snapshot.testIds.some(
+          testId =>
+            testId.startsWith('runtime-local-task-row-') && !taskRowsBeforeFreshChat.has(testId)
+        ),
+      'The second task was not available for task draft isolation'
+    )
+    const secondTaskRowTestId = secondTaskSnapshot.testIds.find(
+      testId => testId.startsWith('runtime-local-task-row-') && !taskRowsBeforeFreshChat.has(testId)
+    )
+    assert.ok(secondTaskRowTestId, 'The second task row was not found')
+    await control.command('fill', composerSelector, { value: UNSENT_SECOND_TASK_DRAFT })
+    await control.command('click', `[data-testid="${taskRowTestId}"]`)
+    await control.command('fill', composerSelector, { value: UNSENT_FIRST_TASK_DRAFT })
+    await control.command('click', `[data-testid="${secondTaskRowTestId}"]`)
+    await waitForControlValue(
+      control,
+      composerSelector,
+      UNSENT_SECOND_TASK_DRAFT,
+      'The second task lost its unsent composer draft after switching tasks'
+    )
+    await control.command('click', `[data-testid="${taskRowTestId}"]`)
+    await waitForControlValue(
+      control,
+      composerSelector,
+      UNSENT_FIRST_TASK_DRAFT,
+      'The first task lost its unsent composer draft after switching tasks'
+    )
+    await control.command('fill', composerSelector, { value: '' })
+    await control.command('click', `[data-testid="${secondTaskRowTestId}"]`)
+    await control.command('fill', composerSelector, { value: '' })
+
     phase = 'standalone-new-task-state'
     await control.command('click', '[data-testid="runtime-chat-section-new-chat-button"]')
     const standaloneTaskSnapshot = await waitForSnapshot(

wework/src/components/chat/ChatInput.tsx

@@ -30,6 +30,7 @@ import { TaskPlanProgress } from './composer/TaskPlanProgress'
 export type ProjectCreateMode = 'scratch' | 'existing' | 'git'
 
 export interface ProjectChatControls {
+  scopeKey?: string
   models: UnifiedModel[]
   skills: UnifiedSkill[]
   selectedModel: UnifiedModel | null

wework/src/components/layout/BufferedChatInput.test.tsx

@@ -41,4 +41,64 @@ describe('BufferedChatInput', () => {
       )
     })
   })
+
+  test('restores a buffered draft after switching chat scopes', async () => {
+    const drafts = new Map<string, string>()
+    const onBlankChange = vi.fn((value: string) => drafts.set('blank', value))
+    const onTaskChange = vi.fn((value: string) => drafts.set('task', value))
+    const baseProjectChat = {
+      models: [],
+      skills: [],
+      selectedModel: null,
+      selectedModelOptions: {},
+      selectedSkills: [],
+      attachments: [],
+      uploadingFiles: new Map(),
+      errors: new Map(),
+      isOptionsLocked: false,
+      setSelectedModel: vi.fn(),
+      setSelectedModelOption: vi.fn(),
+      toggleSkill: vi.fn(),
+      handleFileSelect: vi.fn(),
+      removeAttachment: vi.fn(),
+      listLocalSkills: vi.fn(),
+    }
+    const { rerender } = render(
+      <BufferedChatInput
+        value=""
+        onChange={onBlankChange}
+        onSubmit={vi.fn()}
+        disabled={false}
+        projectChat={{ ...baseProjectChat, scopeKey: 'blank' }}
+      />
+    )
+    await userEvent.type(screen.getByTestId('chat-message-input'), 'unfinished draft')
+
+    rerender(
+      <BufferedChatInput
+        value=""
+        onChange={onTaskChange}
+        onSubmit={vi.fn()}
+        disabled={false}
+        projectChat={{ ...baseProjectChat, scopeKey: 'task' }}
+      />
+    )
+
+    await waitFor(() => expect(onBlankChange).toHaveBeenCalledWith('unfinished draft'))
+    expect(screen.getByTestId('chat-message-input')).toHaveValue('')
+
+    rerender(
+      <BufferedChatInput
+        value={drafts.get('blank') ?? ''}
+        onChange={onBlankChange}
+        onSubmit={vi.fn()}
+        disabled={false}
+        projectChat={{ ...baseProjectChat, scopeKey: 'blank' }}
+      />
+    )
+
+    await waitFor(() => {
+      expect(screen.getByTestId('chat-message-input')).toHaveValue('unfinished draft')
+    })
+  })
 })

wework/src/components/layout/BufferedChatInput.tsx

@@ -18,32 +18,51 @@ export const BufferedChatInput = memo(function BufferedChatInput({
   insertion,
   ...props
 }: BufferedChatInputProps) {
+  const scopeKey = props.projectChat?.scopeKey
   const [draftState, setDraftState] = useState(() => ({
+    scopeKey,
     sourceValue: value,
     draft: value,
   }))
-  const draft = draftState.sourceValue === value ? draftState.draft : value
+  const draft =
+    draftState.scopeKey === scopeKey && draftState.sourceValue === value ? draftState.draft : value
+  const draftRef = useRef(draft)
   const appliedInsertionIdRef = useRef<number | null>(null)
 
+  useEffect(() => {
+    draftRef.current = value
+    return () => {
+      const pendingDraft = draftRef.current
+      if (pendingDraft !== value) {
+        onChange(pendingDraft)
+      }
+    }
+  }, [onChange, scopeKey, value])
+
   useEffect(() => {
     if (!insertion || appliedInsertionIdRef.current === insertion.id) return
     appliedInsertionIdRef.current = insertion.id
     setDraftState(current => {
-      const currentDraft = current.sourceValue === value ? current.draft : value
+      const currentDraft =
+        current.scopeKey === scopeKey && current.sourceValue === value ? current.draft : value
+      const nextDraft = currentDraft ? `${currentDraft}\n${insertion.text}` : insertion.text
+      draftRef.current = nextDraft
       return {
+        scopeKey,
         sourceValue: value,
-        draft: currentDraft ? `${currentDraft}\n${insertion.text}` : insertion.text,
+        draft: nextDraft,
       }
     })
-  }, [insertion, value])
+  }, [insertion, scopeKey, value])
   const setDraft = useCallback(
     (nextDraft: string) => {
-      setDraftState({ sourceValue: value, draft: nextDraft })
+      draftRef.current = nextDraft
+      setDraftState({ scopeKey, sourceValue: value, draft: nextDraft })
       if (parseComposerMentions(nextDraft).length > 0) {
         onChange(nextDraft)
       }
     },
-    [onChange, value]
+    [onChange, scopeKey, value]
   )
 
   const handleSubmit = useCallback(
@@ -55,10 +74,11 @@ export const BufferedChatInput = memo(function BufferedChatInput({
         void onSubmit(submittedDraft, options)
       }
       if (submittedDraft.trim()) {
-        setDraftState({ sourceValue: value, draft: '' })
+        draftRef.current = ''
+        setDraftState({ scopeKey, sourceValue: value, draft: '' })
       }
     },
-    [draft, onSubmit, value]
+    [draft, onSubmit, scopeKey, value]
   )
 
   return <ChatInput {...props} value={draft} onChange={setDraft} onSubmit={handleSubmit} />

wework/src/features/workbench/WorkbenchProvider.tsx

@@ -1306,6 +1306,7 @@ export function WorkbenchProvider({
   )
   const projectChatValue = useMemo(
     () => ({
+      scopeKey: projectChatScopeKey,
       models: modelSelection.models,
       skills: skillSelection.skills,
       selectedModel: modelSelection.selectedModel,
@@ -1345,6 +1346,7 @@ export function WorkbenchProvider({
       attachmentSelection.removeAttachment,
       attachmentSelection.resetAttachments,
       attachmentSelection.uploadingFiles,
+      projectChatScopeKey,
       draftInput,
       trialTemplates,
       handleBlockedModelSelect,
@@ -1370,6 +1372,7 @@ export function WorkbenchProvider({
   )
   const paneProjectChatValue = useMemo(
     () => ({
+      scopeKey: projectChatScopeKey,
       models: modelSelection.models,
       skills: skillSelection.skills,
       selectedModel: modelSelection.selectedModel,
@@ -1409,6 +1412,7 @@ export function WorkbenchProvider({
       attachmentSelection.removeAttachment,
       attachmentSelection.resetAttachments,
       attachmentSelection.uploadingFiles,
+      projectChatScopeKey,
       draftInput,
       trialTemplates,
       handleBlockedModelSelect,