fix(search): align activity headers and preserve complete drafts

Author: waleedlatif1

apps/sim/app/o/[organizationId]/home/organization-home.test.tsx

@@ -170,6 +170,15 @@ it('waits for the user before first rendering their saved mode and level', async
   expect(composerProps()).not.toHaveProperty('assistantSearchLevel')
 })
 
+async function attachDraftImage() {
+  const files = [new File(['image'], 'draft.png', { type: 'image/png' })]
+  await act(async () =>
+    composerProps().files.processFiles(
+      Object.assign(files, { item: (index: number) => files[index] ?? null })
+    )
+  )
+}
+
 function hasCompletedMcpStep() {
   const link = container.querySelector('a[href="/o/organization-a/settings/search-mcp"]')
   expect(link).not.toBeNull()
@@ -914,8 +923,11 @@ it.each(['agent', 'plan', 'assistant'] as const)(
   async (mode) => {
     mocks.plan = true
     await act(async () => renderHome(<OrganizationHome requestMode={mode} />))
+    await attachDraftImage()
     await act(async () => composerProps().onChange('Unsent launch follow-up'))
     expect(localStorage.getItem('mothership-drafts:v1')).toContain('Unsent launch follow-up')
+    expect(localStorage.getItem('mothership-drafts:v1')).toContain('image-key')
+    expect(localStorage.getItem('mothership-drafts:v1')).not.toContain('blob:')
     await act(async () =>
       composerProps().onModeChange?.(mode === 'assistant' ? 'agent' : 'assistant')
     )
@@ -928,20 +940,30 @@ it.each(['agent', 'plan', 'assistant'] as const)(
     root = createRoot(container)
     await act(async () => renderHome(<OrganizationHome requestMode={mode} />))
     expect(composerProps().value).toBe('Unsent launch follow-up')
+    expect(composerProps().files.attachedFiles).toEqual([
+      expect.objectContaining({
+        key: 'image-key',
+        uploading: false,
+        previewUrl: '/api/files/serve/image-key?context=mothership&preview=1',
+      }),
+    ])
     await act(async () => composerProps().onSubmit(composerProps().value))
     expect(composerProps().value).toBe('')
     expect(localStorage.getItem('mothership-drafts:v1')).not.toContain('Unsent launch follow-up')
+    expect(localStorage.getItem('mothership-drafts:v1')).not.toContain('image-key')
   }
 )
 
 it('isolates saved drafts by user, organization, and conversation', async () => {
   mocks.renderer.mockImplementation(({ composer }: { composer: ReactNode }) => composer)
   await act(async () => renderHome(<OrganizationHome chatId='chat-a' requestMode='assistant' />))
+  await attachDraftImage()
   await act(async () => composerProps().onChange('Chat A follow-up'))
   await act(async () => renderHome(<OrganizationHome chatId='chat-b' requestMode='assistant' />))
   expect(composerProps().value).toBe('')
   await act(async () => renderHome(<OrganizationHome chatId='chat-a' requestMode='assistant' />))
   expect(composerProps().value).toBe('Chat A follow-up')
+  expect(composerProps().files.attachedFiles[0]?.key).toBe('image-key')
   mocks.session.mockReturnValue({ data: { user: { id: 'other-user' } } })
   await act(async () => renderHome(<OrganizationHome chatId='chat-a' requestMode='assistant' />))
   expect(composerProps().value).toBe('')
@@ -958,15 +980,36 @@ it('moves a follow-up draft to the newly resolved conversation before it is reop
   mocks.renderer.mockImplementation(({ composer }: { composer: ReactNode }) => composer)
   await act(async () => renderHome(<OrganizationHome requestMode='assistant' />))
   await act(async () => composerProps().onSubmit('First question'))
+  await attachDraftImage()
   await act(async () => composerProps().onChange('Follow-up typed before admission'))
   mocks.chat.mockReturnValue({ ...mocks.chat(), resolvedChatId: 'adopted-chat' })
   await act(async () => renderHome(<OrganizationHome requestMode='assistant' />))
   expect(composerProps().value).toBe('Follow-up typed before admission')
+  expect(composerProps().files.attachedFiles[0]?.key).toBe('image-key')
   expect(
     useMothershipDraftsStore.getState().drafts['reader:organization:organization-a:new']
   ).toBeUndefined()
   await act(async () =>
     renderHome(<OrganizationHome chatId='adopted-chat' requestMode='assistant' />)
   )
   expect(composerProps().value).toBe('Follow-up typed before admission')
+  expect(composerProps().files.attachedFiles[0]?.key).toBe('image-key')
+})
+
+it('keeps an attachment-only draft across remounts and clears it when its last file is removed', async () => {
+  await act(async () => renderHome(<OrganizationHome requestMode='assistant' />))
+  await attachDraftImage()
+  await act(async () => composerProps().onChange('temporary text'))
+  await act(async () => composerProps().onChange(''))
+  await act(async () => root.unmount())
+  root = createRoot(container)
+  await act(async () => renderHome(<OrganizationHome requestMode='assistant' />))
+  expect(composerProps().value).toBe('')
+  const attached = composerProps().files.attachedFiles[0]
+  expect(attached.key).toBe('image-key')
+  await act(async () => composerProps().files.removeFile(attached.id))
+  expect(
+    useMothershipDraftsStore.getState().drafts['reader:organization:organization-a:new']
+  ).toBeUndefined()
+  expect(composerProps().files.attachedFiles).toEqual([])
 })

apps/sim/app/o/[organizationId]/home/organization-home.tsx

@@ -109,17 +109,20 @@ function OrganizationHomeContent({
   })
   const initialDraftKey = `${userId}:organization:${organization.id}:${chatId ?? 'new'}`
   const draftKey = `${userId}:organization:${organization.id}:${chat.resolvedChatId ?? chatId ?? 'new'}`
+  const savedDraft = useMothershipDraftsStore.getState().drafts
+  const initialDraft = savedDraft[draftKey] ?? savedDraft[initialDraftKey]
   const draft = useMothershipDraftsStore(
     (state) => (state.drafts[draftKey] ?? state.drafts[initialDraftKey])?.text ?? ''
   )
   const setDraft = useCallback(
     (text: string, contexts?: ChatContext[]) => {
-      useMothershipDraftsStore.getState().setDraft(draftKey, { text, contexts })
+      const store = useMothershipDraftsStore.getState()
+      store.setDraft(draftKey, { ...store.drafts[draftKey], text, contexts })
     },
     [draftKey]
   )
   const [restoredContexts, setRestoredContexts] = useState<ChatContext[]>(
-    () => useMothershipDraftsStore.getState().drafts[draftKey]?.contexts ?? []
+    () => initialDraft?.contexts ?? []
   )
   useEffect(() => {
     useMothershipDraftsStore.getState().migrateDraft(initialDraftKey, draftKey)
@@ -189,7 +192,25 @@ function OrganizationHomeContent({
     userId: session?.user?.id,
     organizationId: organization.id,
     requestMode,
+    initialAttachments: initialDraft?.fileAttachments,
   })
+  useEffect(() => {
+    const store = useMothershipDraftsStore.getState()
+    const fileAttachments = files.attachedFiles
+      .filter((file) => !file.uploading && file.key)
+      .map((file) => ({
+        id: file.id,
+        key: file.key!,
+        filename: file.name,
+        media_type: file.type,
+        size: file.size,
+        path: file.path,
+      }))
+    store.setDraft(draftKey, {
+      ...(store.drafts[draftKey] ?? { text: '' }),
+      fileAttachments,
+    })
+  }, [draftKey, files.attachedFiles])
   useEffect(() => {
     if (chat.error) toast.error(chat.error)
   }, [chat.error])
@@ -262,7 +283,7 @@ function OrganizationHomeContent({
         path: file.path,
       }))
     if (!message && !attachments.length) return
-    setDraft('')
+    useMothershipDraftsStore.getState().clearDraft(draftKey)
     send(message, attachments.length ? attachments : undefined, contexts)
     setRestoredContexts([])
     files.clearAttachedFiles()
@@ -320,7 +341,7 @@ function OrganizationHomeContent({
             if (queued) {
               const queuedMode = queued.requestMode ?? requestMode
               setSelectedMode(queuedMode)
-              setDraft(queued.content)
+              setDraft(queued.content, queued.contexts)
               setRestoredContexts(queued.contexts ?? [])
               files.restoreAttachedFiles(
                 (queued.fileAttachments ?? []).map((file) => ({

apps/sim/app/o/[organizationId]/integrations/live-member-integrations.test.tsx

@@ -23,6 +23,9 @@ vi.mock('@/hooks/queries/organization-secrets', () => ({
   useRemoveOrganizationSecretSource: () => ({ mutate: vi.fn() }),
 }))
 vi.mock('@/hooks/queries/search-integrations', () => ({ useSearchIntegrations: mocks.policies }))
+vi.mock('@/app/o/[organizationId]/providers/organization-provider', () => ({
+  useOrganizationContext: () => ({ organization: { id: 'org', name: 'Example Organization' } }),
+}))
 vi.mock('@/app/o/[organizationId]/integrations/disconnect-account-menu', () => ({
   DisconnectAccountMenu: ({ accounts }: { accounts: { displayName: string }[] }) => (
     <span>{accounts.map((account) => account.displayName).join(', ')}</span>

apps/sim/app/workspace/[workspaceId]/home/components/knowledge-search-results/knowledge-search-results.tsx

@@ -255,14 +255,18 @@ function SearchResults({
       </ChipLink>
     </div>
   ) : (
-    <div className='flex flex-col'>
+    <div aria-busy={!awaitingRange && fetching} className='flex flex-col'>
       <div className='flex items-center gap-2 px-2 py-2'>
         <div className='min-w-0 flex-1'>
           {awaitingRange ? (
             <p role='status' className='text-[var(--text-muted)] text-caption'>
               Choose the days to search.
             </p>
-          ) : fetching || (pending && !failed) ? null : (
+          ) : fetching ? (
+            <p role='status' className='sr-only'>
+              {pending ? 'Searching…' : 'Updating results…'}
+            </p>
+          ) : pending && !failed ? null : (
             <p role='status' className='text-[var(--text-muted)] text-caption'>
               {failed
                 ? 'Search couldn’t run.'
@@ -442,7 +446,12 @@ function LiveSearchResults({
     ]),
   ]
   return (
-    <div className='flex flex-col'>
+    <div aria-busy={!awaitingRange && isFetching} className='flex flex-col'>
+      {!awaitingRange && isFetching && (
+        <p role='status' className='sr-only'>
+          {isPending ? 'Searching…' : 'Updating results…'}
+        </p>
+      )}
       {(awaitingRange || isError) && (
         <div className='flex items-center gap-2 px-2 py-2'>
           {awaitingRange ? (

apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/search-activity.tsx

@@ -1,10 +1,11 @@
 'use client'
 
 import { useState } from 'react'
-import { cn, OverflowText } from '@sim/emcn'
+import { cn } from '@sim/emcn'
 import { Search } from '@sim/emcn/icons'
 import { toStringOrNull } from '@sim/utils/coerce'
 import { toArray, toRecord } from '@sim/utils/object'
+import { ActivityStatus } from '@/components/ui/activity-status'
 import {
   collectRetrievalCitationEvidence,
   parseCitationRecord,
@@ -91,10 +92,7 @@ function SearchQueryActivity({ tool }: SearchQueryActivityProps) {
   return (
     <ActivityDisclosure
       header={
-        <span className='flex min-w-0 items-center gap-2 text-[var(--text-muted)] text-small'>
-          <Search aria-hidden className='size-[14px] shrink-0 text-[var(--text-icon)]' />
-          <OverflowText label={label} focusTarget='nearest-interactive' />
-        </span>
+        <ActivityStatus label={label} isActive={false} icon={<Search className='size-[14px]' />} />
       }
       expanded={expanded}
       onToggle={() => setExpanded(!expanded)}

apps/sim/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-content/components/search-resource-content.test.tsx

@@ -118,7 +118,7 @@ async function render(query: string, nativeQueries?: [{ provider: 'github'; quer
   await vi.waitFor(async () => {
     await act(async () => {})
     expect(client.isFetching()).toBe(0)
-    expect(container.textContent).not.toContain('Searching…')
+    expect(container.querySelector('[aria-busy="true"]')).toBeNull()
   })
 }
 describe('shared Search resource content', () => {

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/copilot/components/user-input/hooks/use-file-attachments.ts

@@ -17,7 +17,10 @@ import {
 } from '@/lib/uploads/shared/assistant-images'
 import { MAX_WORKSPACE_FILE_SIZE } from '@/lib/uploads/shared/types'
 import { resolveFileType } from '@/lib/uploads/utils/file-utils'
-import type { ChatRequestMode } from '@/app/workspace/[workspaceId]/home/types'
+import type {
+  ChatRequestMode,
+  FileAttachmentForApi,
+} from '@/app/workspace/[workspaceId]/home/types'
 
 const logger = createLogger('useFileAttachments')
 
@@ -89,6 +92,7 @@ interface UseFileAttachmentsProps {
   requestMode?: ChatRequestMode
   disabled?: boolean
   isLoading?: boolean
+  initialAttachments?: FileAttachmentForApi[]
 }
 
 /**
@@ -102,11 +106,22 @@ export function useFileAttachments(props: UseFileAttachmentsProps) {
   const { userId, workspaceId, organizationId, requestMode, disabled, isLoading } = props
   const imagesOnly = Boolean(organizationId) && requestMode !== 'agent' && requestMode !== 'plan'
 
-  const [attachedFiles, setAttachedFiles] = useState<AttachedFile[]>([])
+  const [attachedFiles, setAttachedFiles] = useState<AttachedFile[]>(() =>
+    (props.initialAttachments ?? []).map((file) => ({
+      id: file.id,
+      name: file.filename,
+      size: file.size,
+      type: file.media_type,
+      key: file.key,
+      path: file.path || getMothershipAttachmentPreviewUrl(file) || '',
+      previewUrl: getMothershipAttachmentPreviewUrl(file),
+      uploading: false,
+    }))
+  )
   const [dragCounter, setDragCounter] = useState(0)
   const isDragging = dragCounter > 0
   const fileInputRef = useRef<HTMLInputElement>(null)
-  const attachedFilesRef = useRef<AttachedFile[]>([])
+  const attachedFilesRef = useRef<AttachedFile[]>(attachedFiles)
   const uploadControllersRef = useRef(new Map<string, AbortController>())
 
   const updateAttachedFiles = useCallback((update: (files: AttachedFile[]) => AttachedFile[]) => {

apps/sim/stores/mothership-drafts/store.ts

@@ -22,6 +22,7 @@ interface MothershipDraftsState {
   setDraft: (key: string, payload: DraftPayload) => void
   clearDraft: (key: string) => void
   migrateDraft: (fromKey: string, toKey: string) => void
+  reset: () => void
 }
 
 const LEGACY_WORKFLOW_COPILOT_KEY = /^[^:]+:workflow-copilot:[^:]+$/
@@ -58,6 +59,7 @@ export const useMothershipDraftsStore = create<MothershipDraftsState>()(
     persist(
       (set) => ({
         drafts: {},
+        reset: () => set({ drafts: {} }),
         setDraft: (key, payload) =>
           set((s) => {
             if (isEmpty(payload)) {

apps/sim/stores/reset-all-stores.test.ts

@@ -10,6 +10,7 @@ const {
   mockClearAllExecutionPointers,
   mockGetQueryClient,
   mockMothershipQueueReset,
+  mockMothershipDraftsReset,
   mockOperationQueueReset,
   mockResetRegisteredUserData,
   mockRegistrySetState,
@@ -22,6 +23,7 @@ const {
   mockConsoleReset: vi.fn(),
   mockGetQueryClient: vi.fn(),
   mockMothershipQueueReset: vi.fn(),
+  mockMothershipDraftsReset: vi.fn(),
   mockOperationQueueReset: vi.fn(),
   mockResetRegisteredUserData: vi.fn(),
   mockRegistrySetState: vi.fn(),
@@ -40,7 +42,7 @@ vi.mock('@/stores/execution', () => ({
   useExecutionStore: { getState: () => ({ reset: vi.fn() }) },
 }))
 vi.mock('@/stores/mothership-drafts/store', () => ({
-  useMothershipDraftsStore: { setState: vi.fn() },
+  useMothershipDraftsStore: { getState: () => ({ reset: mockMothershipDraftsReset }) },
 }))
 vi.mock('@/stores/mothership-queue/store', () => ({
   useMothershipQueueStore: { getState: () => ({ reset: mockMothershipQueueReset }) },
@@ -104,6 +106,7 @@ describe('resetAllStores', () => {
     expect(mockConsoleReset).toHaveBeenCalledOnce()
     expect(mockClearAllExecutionPointers).toHaveBeenCalledOnce()
     expect(mockMothershipQueueReset).toHaveBeenCalledOnce()
+    expect(mockMothershipDraftsReset).toHaveBeenCalledOnce()
     expect(mockConsolePersist).toHaveBeenCalledWith({ merge: false })
   })
 

apps/sim/stores/reset-all-stores.ts

@@ -53,7 +53,7 @@ export async function resetAllStores(): Promise<void> {
   })
   consolePersistence.reset()
   clearAllExecutionPointers()
-  useMothershipDraftsStore.setState({ drafts: {} })
+  useMothershipDraftsStore.getState().reset()
   useMothershipQueueStore.getState().reset()
   await consolePersistence.persist({ merge: false })
 }