diff --git a/packages/subagent/subagent/tests/continuation.spec.ts b/packages/subagent/subagent/tests/continuation.spec.ts index a99ac93e71..4ca915b77e 100644 --- a/packages/subagent/subagent/tests/continuation.spec.ts +++ b/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() + 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() + 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)