test(subagent): await the image capability read before draining

packages/subagent/subagent/tests/continuation.spec.ts

@@ -659,17 +659,26 @@ describe('continuable image Queue prompts', () => {
     const started = await ctx.subagents.startContinuable(startSpec(parent))
     await vi.waitFor(() => { expect(adapter.requests).toHaveLength(1) })
     const capability = Promise.withResolvers<{ inputModalities: string[] }>()
-    const resolve = vi.spyOn(ctx.llm, 'resolveModelInfo').mockReturnValue(capability.promise as never)
+    const readingCapability = Promise.withResolvers<undefined>()
+    vi.spyOn(ctx.llm, 'resolveModelInfo').mockImplementation(() => {
+      readingCapability.resolve(undefined)
+      return capability.promise as never
+    })
 
     const delivery = queuePrompt(ctx, parent, started.childId, [imageBlock])
-    delivery.catch(() => undefined)
-    await vi.waitFor(() => { expect(resolve).toHaveBeenCalled() })
-    releaseFirst.resolve(undefined)
-    const draining = drainManager(ctx)
-    capability.resolve({ inputModalities: ['text', 'image'] })
+    try {
+      await Promise.race([readingCapability.promise, delivery])
+      releaseFirst.resolve(undefined)
+      const draining = drainManager(ctx)
+      capability.resolve({ inputModalities: ['text', 'image'] })
 
-    await expect(delivery).rejects.toMatchObject({ code: 'DRAINING' })
-    await draining
+      await expect(delivery).rejects.toMatchObject({ code: 'DRAINING' })
+      await draining
+    } finally {
+      releaseFirst.resolve(undefined)
+      capability.resolve({ inputModalities: ['text', 'image'] })
+      await Promise.allSettled([delivery, drainManager(ctx)])
+    }
   })
 
   it('rejects a materialized image follow-up whose capability read raced a drain', async () => {
@@ -677,16 +686,24 @@ describe('continuable image Queue prompts', () => {
     const started = await ctx.subagents.startContinuable(startSpec(parent))
     await waitNoActivation(ctx, started.childId)
     const capability = Promise.withResolvers<{ inputModalities: string[] }>()
-    const resolve = vi.spyOn(ctx.llm, 'resolveModelInfo').mockReturnValue(capability.promise as never)
+    const readingCapability = Promise.withResolvers<undefined>()
+    vi.spyOn(ctx.llm, 'resolveModelInfo').mockImplementation(() => {
+      readingCapability.resolve(undefined)
+      return capability.promise as never
+    })
 
     const delivery = queuePrompt(ctx, parent, started.childId, [imageBlock])
-    delivery.catch(() => undefined)
-    await vi.waitFor(() => { expect(resolve).toHaveBeenCalled() })
-    const draining = drainManager(ctx)
-    capability.resolve({ inputModalities: ['text', 'image'] })
+    try {
+      await Promise.race([readingCapability.promise, delivery])
+      const draining = drainManager(ctx)
+      capability.resolve({ inputModalities: ['text', 'image'] })
 
-    await expect(delivery).rejects.toMatchObject({ code: 'ACTIVATION_CLOSING' })
-    await draining
+      await expect(delivery).rejects.toMatchObject({ code: 'ACTIVATION_CLOSING' })
+      await draining
+    } finally {
+      capability.resolve({ inputModalities: ['text', 'image'] })
+      await Promise.allSettled([delivery, drainManager(ctx)])
+    }
     const loaded = await loadStoredSession(ctx.sessionPersistence, started.childId)
     expect(loaded.events.some(event => event.type === 'user/message'
       && event.data.content.some(block => block.type === 'image'))).toBe(false)