fix(sessions): close the delegation dialog without waiting for sessions to start (#1631)

* fix(sessions): close the delegation dialog without waiting for sessions to start

Start and immediate model changes answered only after every delegated
session had its worktree created and its CLI started (or stopped and
restarted), and the dialog ignored Close until then. The UI routes now
answer once the batch or change is accepted; the launch continues in the
background and reports through the batch and change state. The model's
own CreateSessions call still waits, since it needs the new session ids.
Closing the dialog is never blocked by an in-flight request.

* test(sessions): wait for both first messages before checking a launched batch

Author: NanmiCoder

desktop/src/components/sessionDelegation/SessionDelegationDialog.test.tsx

@@ -82,6 +82,31 @@ describe('SessionDelegationDialog', () => {
     await waitFor(() => expect(screen.queryByRole('dialog')).not.toBeInTheDocument())
   })
 
+  it('closes while a confirmation is still in flight and stays closed when it answers', async () => {
+    let answer!: (batch: SessionDelegationSnapshot['batches'][number]) => void
+    vi.mocked(sessionCollaborationApi.act).mockReturnValueOnce(new Promise(resolve => { answer = resolve }))
+    openDialog(fixture())
+    fireEvent.click(screen.getByRole('button', { name: 'Start sessions' }))
+    await waitFor(() => expect(sessionCollaborationApi.act).toHaveBeenCalledWith('parent', expect.objectContaining({ id: 'batch' }), 'confirm'))
+    expect(screen.getByRole('button', { name: 'Start sessions' })).toBeDisabled()
+    fireEvent.click(screen.getByRole('button', { name: 'Close dialog' }))
+    expect(screen.queryByRole('dialog')).not.toBeInTheDocument()
+    expect(useSessionDelegationStore.getState().bySession.parent?.manageOpen).toBe(false)
+    await act(async () => { answer({ ...fixture().batches[0]!, state: 'launching', revision: 2 }) })
+    await waitFor(() => expect(sessionCollaborationApi.get).toHaveBeenCalled())
+    expect(screen.queryByRole('dialog')).not.toBeInTheDocument()
+  })
+
+  it('closes with Escape while a model change is still in flight', async () => {
+    vi.mocked(sessionCollaborationApi.change).mockReturnValueOnce(new Promise(() => {}))
+    openDialog(launchedFixture())
+    fireEvent.click(screen.getByLabelText('Builder'))
+    fireEvent.click(screen.getByRole('button', { name: 'Apply to 1 session' }))
+    await waitFor(() => expect(screen.getByRole('button', { name: 'Apply to 1 session' })).toBeDisabled())
+    fireEvent.keyDown(document, { key: 'Escape' })
+    expect(screen.queryByRole('dialog')).not.toBeInTheDocument()
+  })
+
   it('blocks unresolved models and can cancel without starting', async () => {
     const data = fixture()
     data.batches[0]!.sessions[0]!.runtime = { providerId: 'deleted', modelId: 'same' }

desktop/src/components/sessionDelegation/SessionDelegationDialog.tsx

@@ -78,7 +78,8 @@ function SessionDelegationDialogContent({ sessionId }: { sessionId: string }) {
   const effective = (row: DelegationEntryView) => drafts[row.key] ?? row.entry.runtime
   const conflicted = (batch: SessionDelegationBatch) => draftRevisions.current.has(batch.id) && draftRevisions.current.get(batch.id) !== batch.revision
 
-  const close = () => { if (!busy) setManageOpen(sessionId, false) }
+  // Closing never cancels a request; its outcome shows in the review bar and the Activity rail.
+  const close = () => setManageOpen(sessionId, false)
   /**
    * Forget what a successful action submitted, and close only when nothing else
    * is left. A selection can mix started sessions, which Apply changes, with

scripts/quality-gate/mock-llm/scenarios.ts

@@ -103,8 +103,16 @@ export function createMockOnlyRunners(
         const response = await fetch(`${base}/batches/${encodeURIComponent(batch.id)}/confirm`, {
           method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ expectedRevision: batch.revision }),
         })
-        const launched = await response.json() as any
-        if (!response.ok || launched.state !== 'launched' || !launched.sessions[0]?.sessionId) throw new Error(`confirm did not launch the child: ${JSON.stringify(launched)}`)
+        const accepted = await response.json() as any
+        if (!response.ok || !['launching', 'launched'].includes(accepted.state)) throw new Error(`confirm was not accepted: ${JSON.stringify(accepted)}`)
+        // The UI request answers once the batch is launching; the child starts after it.
+        const deadline = Date.now() + ctx.stepTimeoutMs
+        let launched = accepted
+        while (launched.state !== 'launched' && launched.state !== 'failed' && Date.now() < deadline) {
+          await Bun.sleep(100)
+          launched = (await read()).batches.find((item: any) => item.id === batch.id) ?? launched
+        }
+        if (launched.state !== 'launched' || !launched.sessions[0]?.sessionId) throw new Error(`confirm did not launch the child: ${JSON.stringify(launched)}`)
         const runtime = launched.sessions[0].runtime
         if (runtime.providerId !== ctx.target.providerId || runtime.modelId !== ctx.target.modelId) throw new Error('Child runtime did not preserve the selected provider and model')
         await fetch(`${base}/stop`, { method: 'POST' })

src/server/__tests__/e2e/session-collaboration.test.ts

@@ -63,9 +63,10 @@ test('desktop host dispatches independent sessions, exchanges messages, exposes
     expect(proposals.every(batch => batch.state === 'review' && batch.sessions.every((item: any) => !item.sessionId))).toBe(true)
     const children = []
     for (const batch of proposals) {
-      const confirmed = await api(`/api/session-collaboration/${root.sessionId}/batches/${encodeURIComponent(batch.id)}/confirm`, { expectedRevision: batch.revision })
-      expect(confirmed.state).toBe('launched')
-      children.push(confirmed.sessions[0])
+      const accepted = await api(`/api/session-collaboration/${root.sessionId}/batches/${encodeURIComponent(batch.id)}/confirm`, { expectedRevision: batch.revision })
+      expect(['launching', 'launched']).toContain(accepted.state)
+      const snapshot = await eventually(() => api(`/api/session-collaboration/${root.sessionId}/delegation`), value => value.batches.some((item: any) => item.id === batch.id && item.state === 'launched'), 'batch launched')
+      children.push(snapshot.batches.find((item: any) => item.id === batch.id).sessions[0])
     }
     expect(new Set([root.sessionId, ...children.map(child => child.sessionId)]).size).toBe(3)
     const busyDelivery = await tool('send', { targetSessionId: children[0].sessionId, content: 'Progress while running', messageId: 'busy-message' })

src/server/api/sessionCollaboration.test.ts

@@ -66,7 +66,7 @@ test('UI model changes validate input and bind the route session', async () => {
   const input = { requestId: 'change', mode: 'after_turn', selections: [{ sessionId: 'child', runtime: { providerId: 'b', modelId: 'shared' } }] }
   const response = await handleSessionCollaborationUiApi(new Request(url, { method: 'POST', body: JSON.stringify(input) }), url, service)
   expect(response.status).toBe(200)
-  expect(calls).toEqual([['parent', input]])
+  expect(calls).toEqual([['parent', input, { wait: false }]])
   const invalid = await handleSessionCollaborationUiApi(new Request(url, { method: 'POST', body: '{' }), url, service)
   expect(invalid.status).toBe(400)
 })

src/server/api/sessionCollaboration.ts

@@ -99,11 +99,12 @@ export async function handleSessionCollaborationUiApi(req: Request, url: URL, se
         return Response.json(await service.updateBatch(sessionId, batchId, input.expectedRevision, input.selections))
       }
       const input = parse(revisionSchema, body)
-      if (parts[3] === 'confirm') return Response.json(await service.confirmBatch(sessionId, batchId, input.expectedRevision))
+      // The dialog closes on this response; the launch reports through the batch's state.
+      if (parts[3] === 'confirm') return Response.json(await service.confirmBatch(sessionId, batchId, input.expectedRevision, { wait: false }))
       if (parts[3] === 'cancel') return Response.json(await service.cancelBatch(sessionId, batchId, input.expectedRevision))
     }
     if (parts.length === 2 && parts[1] === 'model-changes' && req.method === 'POST') {
-      return Response.json(await service.changeModels(sessionId, parse(changeSchema, await jsonBody(req))))
+      return Response.json(await service.changeModels(sessionId, parse(changeSchema, await jsonBody(req)), { wait: false }))
     }
     if (parts.length === 4 && parts[1] === 'model-changes' && parts[3] === 'cancel' && req.method === 'POST') {
       const input = parse(revisionSchema, await jsonBody(req))

src/server/services/sessionCollaborationService.delegation.test.ts

@@ -4,6 +4,7 @@ import { join } from 'node:path'
 import { tmpdir } from 'node:os'
 import { SessionCollaborationService, type SessionCollaborationDependencies, type CollaborationMessage } from './sessionCollaborationService.js'
 import type { SessionRuntimeSelection } from '../../shared/sessionDelegation.js'
+import { handleSessionCollaborationUiApi } from '../api/sessionCollaboration.js'
 
 let directory: string
 let service: SessionCollaborationService
@@ -344,3 +345,66 @@ test('manual child Stop fences a model change still validating before persistenc
   expect(await service.listModelChanges('root')).toEqual([])
   expect((await service.status()).members.find(item => item.sessionId === child.sessionId)?.stopped).toBe(true)
 })
+
+async function until<T>(read: () => Promise<T>, ready: (value: T) => boolean): Promise<T> {
+  for (let attempt = 0; attempt < 200; attempt++) {
+    const value = await read()
+    if (ready(value)) return value
+    await Bun.sleep(5)
+  }
+  throw new Error('condition was not reached')
+}
+const uiPost = (path: string, body: unknown) => {
+  const url = new URL(`http://localhost/api/session-collaboration/root/${path}`)
+  return handleSessionCollaborationUiApi(new Request(url, { method: 'POST', body: JSON.stringify(body) }), url, service)
+}
+
+// The dialog closes on the response. It used to arrive only after every
+// workspace was created and every CLI had started, with Close ignored meanwhile.
+test('UI Start answers once the batch is launching and the launch finishes on its own', async () => {
+  let release!: () => void
+  const gate = new Promise<void>(resolve => { release = resolve })
+  const create = deps.sessions.create
+  deps.sessions.create = async (...args) => { await gate; return create(...args) }
+  // The batch reads launched before each first message is delivered; a slow
+  // start keeps that window open on every machine.
+  deps.runtime.start = async (_, message) => { await Bun.sleep(20); started.push(message) }
+  const proposal = await service.propose('root', input())
+  const response = await uiPost(`batches/${encodeURIComponent(proposal.id)}/confirm`, { expectedRevision: proposal.revision })
+  expect(response.status).toBe(200)
+  expect(await response.json()).toMatchObject({ id: proposal.id, state: 'launching' })
+  expect(created).toBe(0)
+  release()
+  await until(async () => started.length, count => count === 2)
+  const [launched] = await service.listBatches('root')
+  expect(launched?.state).toBe('launched')
+  expect(launched?.sessions.every(entry => entry.launched && entry.sessionId)).toBe(true)
+})
+
+test('UI Start still refuses a stale revision, and a launch failing later is recorded on the batch', async () => {
+  let release!: () => void
+  const gate = new Promise<void>(resolve => { release = resolve })
+  deps.sessions.create = async () => { await gate; throw new Error('workspace failure') }
+  const proposal = await service.propose('root', input())
+  expect((await uiPost(`batches/${encodeURIComponent(proposal.id)}/confirm`, { expectedRevision: proposal.revision + 1 })).status).toBe(409)
+  expect((await uiPost(`batches/${encodeURIComponent(proposal.id)}/confirm`, { expectedRevision: proposal.revision })).status).toBe(200)
+  release()
+  const [failed] = await until(() => service.listBatches('root'), batches => batches[0]?.state === 'failed')
+  expect(failed?.error).toContain('workspace failure')
+  expect(started).toHaveLength(0)
+})
+
+test('UI immediate change answers before running sessions are stopped and restarted', async () => {
+  const child = await service.create('root', { prompt: 'work', runtime })
+  await service.onMessageConsumed(child.messageId)
+  let release!: () => void
+  const gate = new Promise<void>(resolve => { release = resolve })
+  deps.runtime.stopAndWait = async id => { await gate; await service.onStopped(id) }
+  const response = await uiPost('model-changes', { requestId: 'ui', mode: 'immediate', selections: [{ sessionId: child.sessionId, runtime: { ...runtime, modelId: 'better' } }] })
+  expect(response.status).toBe(200)
+  expect(['pending', 'applying']).toContain((await response.json()).state)
+  expect(applied).toEqual([])
+  release()
+  await until(() => service.listModelChanges('root'), changes => changes[0]?.state === 'applied')
+  expect(applied).toEqual([child.sessionId])
+})

src/server/services/sessionCollaborationService.ts

@@ -72,7 +72,8 @@ export type SessionCollaborationDependencies = {
   }
   now?: () => Date
 }
-type Store = { batches?: Record<string, SessionDelegationBatch & { input: string; stopEpoch: number }>; modelChanges?: Record<string, SessionModelChange & { input: string; continuations?: string[]; stopEpoch?: number }>; version: 1; revision: number; members: Record<string, CollaborationMember>; messages: CollaborationMessage[]; stopEpochs?: Record<string, number>; creations?: Record<string, { input: string; rootSessionId?: string; stopEpoch?: number; result?: { sessionId: string; workDir?: string; messageId: string; title?: string }; failure?: { message: string; code: string; status: number } }> }
+type StoredBatch = SessionDelegationBatch & { input: string; stopEpoch: number }
+type Store = { batches?: Record<string, StoredBatch>; modelChanges?: Record<string, SessionModelChange & { input: string; continuations?: string[]; stopEpoch?: number }>; version: 1; revision: number; members: Record<string, CollaborationMember>; messages: CollaborationMessage[]; stopEpochs?: Record<string, number>; creations?: Record<string, { input: string; rootSessionId?: string; stopEpoch?: number; result?: { sessionId: string; workDir?: string; messageId: string; title?: string }; failure?: { message: string; code: string; status: number } }> }
 /** Pages one ReadSession cursor chain may serve before it stops instead of walking the whole transcript. */
 export const COLLABORATION_READ_MAX_PAGES = 8
 export const COLLABORATION_WAIT_MIN_MS = 10_000
@@ -334,7 +335,13 @@ export class SessionCollaborationService {
     return (await this.listBatches(callerSessionId)).find(item => item.id === id)!
   }
 
-  async confirmBatch(callerSessionId: string, id: string, revision: number): Promise<SessionDelegationBatch> {
+  /**
+   * The model's tool call waits for the launch: it needs the new session ids.
+   * The UI passes `wait: false` and returns once the batch is launching.
+   * Creating each workspace and starting its CLI can take far longer than a
+   * dialog should stay open, and the batch's state already reports the outcome.
+   */
+  async confirmBatch(callerSessionId: string, id: string, revision: number, options: { wait?: boolean } = {}): Promise<SessionDelegationBatch> {
     if (!this.deps.models) throw ApiError.conflict('Model selection is unavailable')
     await this.ready; await this.tail
     const snapshot = structuredClone(this.batch(callerSessionId, id))
@@ -352,6 +359,13 @@ export class SessionCollaborationService {
       batch.state = 'launching'; batch.revision++; delete batch.error
       batch.sessions.forEach((item, index) => { item.runtime = runtimes[index]! })
     })
+    const launch = this.launchBatch(callerSessionId, id, snapshot, runtimes)
+    if (options.wait === false) launch.catch(error => console.error('[SessionCollaboration] Batch launch failed', error))
+    else await launch
+    return (await this.listBatches(callerSessionId)).find(item => item.id === id)!
+  }
+
+  private async launchBatch(callerSessionId: string, id: string, snapshot: StoredBatch, runtimes: SessionRuntimeSelection[]): Promise<void> {
     try {
       for (const entry of snapshot.sessions) {
         const runtime = runtimes[snapshot.sessions.indexOf(entry)]!
@@ -388,10 +402,10 @@ export class SessionCollaborationService {
     } catch (cause) {
       await this.mutate(() => { const batch = this.batch(callerSessionId, id); batch.state = 'failed'; batch.revision++; batch.error = cause instanceof Error ? cause.message : String(cause) })
     }
-    return (await this.listBatches(callerSessionId)).find(item => item.id === id)!
   }
 
-  async changeModels(callerSessionId: string, input: SessionModelChangeInput): Promise<SessionModelChange> {
+  /** As with confirmBatch, the UI passes `wait: false`: stopping and restarting sessions reports through the change's state. */
+  async changeModels(callerSessionId: string, input: SessionModelChangeInput, options: { wait?: boolean } = {}): Promise<SessionModelChange> {
     if (!input.requestId?.trim() || !Array.isArray(input.selections) || !input.selections.length || input.selections.length > 8 || new Set(input.selections.map(item => item.sessionId)).size !== input.selections.length) throw ApiError.badRequest('Choose between one and eight distinct sessions')
     if (!this.deps.models || !this.deps.runtime.apply || !this.deps.runtime.stopAndWait) throw ApiError.conflict('Runtime changes are unavailable')
     if (input.mode !== undefined && input.mode !== 'after_turn' && input.mode !== 'immediate') throw ApiError.badRequest('Invalid change mode')
@@ -430,7 +444,9 @@ export class SessionCollaborationService {
       this.store.modelChanges ??= {}
       this.store.modelChanges[id] = { id, callerSessionId, requestId: input.requestId, revision: 1, mode: input.mode ?? 'after_turn', resume: input.resume !== false, selections, state: 'pending', input: JSON.stringify(input), stopEpoch: this.store.stopEpochs?.[caller.rootSessionId] ?? 0 }
     })
-    await this.processPendingModelChanges()
+    const processing = this.processPendingModelChanges()
+    if (options.wait === false) processing.catch(error => console.error('[SessionCollaboration] Model change failed', error))
+    else await processing
     return (await this.listModelChanges(callerSessionId)).find(item => item.id === id)!
   }