test(team): await mailbox acknowledgement flushes before teardown

packages/experimental/agent-team/tests/persistence.spec.ts

@@ -4,7 +4,7 @@ import { rm } from 'node:fs/promises'
 import { tmpdir } from 'node:os'
 import { join } from 'node:path'
 import { Context } from '@deepseek-ai/cordis'
-import type { Agent } from '@deepseek-ai/dsh-agent'
+import type { Agent, AgentHandle } from '@deepseek-ai/dsh-agent'
 import AgentLoop from '@deepseek-ai/dsh-agent-loop'
 import { mountAgentLoopTestDependencies } from '@deepseek-ai/dsh-agent-loop-testkit'
 import { createUserMessage } from '@deepseek-ai/dsh-llm'
@@ -14,6 +14,7 @@ import SubagentService, { snapshotSubagentDescriptor } from '@deepseek-ai/dsh-su
 import * as SubagentSpawn from '@deepseek-ai/dsh-subagent-spawn-in-process'
 import { MockAdapter, textResponse } from '../../../core/agent-loop/tests/mock-adapter.ts'
 import TeamService, { TeamId, TeamMessageId } from '../src/index.ts'
+import type { TeamMailbox } from '../src/mailbox.ts'
 import { teamProjectionDefinition } from '../src/projection.ts'
 import type { TeamMemberSnapshot, TeamMessageSnapshot, TeamTaskSnapshot } from '../src/index.ts'
 import { TestSessionQuery } from './test-session-query.ts'
@@ -50,6 +51,12 @@ async function storedEvents(ctx: Context, id: SessionId): Promise<readonly Sessi
   }
 }
 
+/** Await mailbox acknowledgements through their flush and dispatch completion. */
+async function settleMailbox(ctx: Context): Promise<void> {
+  const { mailbox } = ctx.agentTeams as unknown as { readonly mailbox: TeamMailbox }
+  await Promise.all(mailbox.pendingDispatches())
+}
+
 async function disposeContext(ctx: Context): Promise<void> {
   try {
     await ctx.fiber.dispose()
@@ -360,9 +367,8 @@ for (const backend of backends) {
         },
       }), { surfaceOp: 'append' })
       await first.ctx.sessions.flush(targetHandle.agent.session)
-      // Let the pre-queue acknowledgement observer prove there is no mailbox
-      // row yet before authoring the simulated crash prefix below.
-      await new Promise<void>((resolve) => { setTimeout(resolve, 0) })
+      // Finish the target observer before writing the crash-only queued prefix.
+      await settleMailbox(first.ctx)
       await targetHandle.dispose()
 
       const queued: TeamMessageSnapshot = {
@@ -382,11 +388,41 @@ for (const backend of backends) {
       await first.dispose()
 
       const second = await stack(backend, storageRoot, [])
-      const rootHandle = await second.ctx.agents.resume({
-        resumeSessionId: rootId,
-        agentOptions: { provider: 'mock', model: 'mock' },
+      const { mailbox } = second.ctx.agentTeams as unknown as { readonly mailbox: TeamMailbox }
+      const flush = second.ctx.sessions.flush.bind(second.ctx.sessions)
+      const checkpointEntered = Promise.withResolvers<undefined>()
+      const releaseCheckpoint = Promise.withResolvers<undefined>()
+      const delayedCheckpoint = vi.spyOn(second.ctx.sessions, 'flush').mockImplementation(async (session) => {
+        if (session.id === rootId && session.snapshotEvents().some(event =>
+          event.type === 'team/message/delivered' && event.data.messageId === messageId)) {
+          checkpointEntered.resolve(undefined)
+          await releaseCheckpoint.promise
+        }
+        return await flush(session)
       })
-      await vi.waitFor(() => { expect(durable(rootHandle.agent).pendingMessages).toEqual([]) })
+      let rootHandle: AgentHandle
+      try {
+        rootHandle = await second.ctx.agents.resume({
+          resumeSessionId: rootId,
+          agentOptions: { provider: 'mock', model: 'mock' },
+        })
+        await checkpointEntered.promise
+        expect(durable(rootHandle.agent).pendingMessages).toEqual([])
+        expect(mailbox.pendingDispatches().length).toBeGreaterThan(0)
+        let settled = false
+        const settlement = settleMailbox(second.ctx).then(() => { settled = true })
+        await Promise.resolve()
+        expect(settled).toBe(false)
+        releaseCheckpoint.resolve(undefined)
+        await settlement
+      } finally {
+        releaseCheckpoint.resolve(undefined)
+        try {
+          await settleMailbox(second.ctx)
+        } finally {
+          delayedCheckpoint.mockRestore()
+        }
+      }
       expect(second.ctx.agents.get(started.member.id)).toBeUndefined()
       expect(second.adapter.requests).toEqual([])
 
@@ -461,6 +497,7 @@ for (const backend of backends) {
       await vi.waitFor(() => {
         expect(durable(rootHandle.agent).pendingMessages).toEqual([])
       })
+      await settleMailbox(second.ctx)
       expect(second.adapter.requests).toEqual([])
       expect(second.ctx.agents.get(childId)).toBeUndefined()
       const stored = await storedEvents(second.ctx, childId)