fix(team-mode): close peer message delivery races

This commit is contained in:
YeonGyu-Kim
2026-05-19 18:23:57 +09:00
parent bcea4a9d28
commit b2918fd4db
13 changed files with 492 additions and 47 deletions
@@ -196,7 +196,6 @@ export class ParentWakeNotifier {
sessionID,
source: "background-agent-parent-wake",
settleMs: 0,
postDispatchHoldMs: 250,
queueBehavior: "defer",
checkToolState: !toolWaitDecision.skipPromptGateToolStateCheck,
input: {
@@ -79,6 +79,30 @@ describe("pollAndBuildInjection", () => {
})
})
test("#given concurrent transforms for one turn #when mailbox injection is claimed #then only one call injects the peer message", async () => {
// given
const { teamRunId, config } = await setupRuntime(["m1"])
await sendMessage({
version: 1,
messageId: randomUUID(),
from: "lead",
to: "m1",
kind: "message",
body: "race",
timestamp: 100,
}, teamRunId, config, { isLead: true, activeMembers: ["m1"] })
// when
const results = await Promise.all(Array.from({ length: 8 }, () =>
pollAndBuildInjection("session-1", "m1", teamRunId, config, "turn-race")
))
// then
expect(results.filter((result) => result.injected)).toHaveLength(1)
expect(results.filter((result) => !result.injected)).toHaveLength(7)
})
test("wraps hostile message bodies in a literal peer_message envelope", async () => {
// given
const { teamRunId, config } = await setupRuntime(["m1"])
+45 -38
View File
@@ -48,47 +48,54 @@ export async function pollAndBuildInjection(
config: TeamModeConfig,
turnMarker: string,
): Promise<InjectionResult> {
const runtimeState = await loadRuntimeState(teamRunId, config)
const runtimeMember = runtimeState.members.find((member) => member.name === memberName)
if (runtimeMember === undefined) {
throw new Error(`runtime member not found for session ${sessionID}: ${memberName}`)
}
const unreadMessages = await listUnreadMessages(teamRunId, memberName, config)
let result: InjectionResult | undefined
if (runtimeMember.lastInjectedTurnMarker === turnMarker) {
return { injected: false, messageIds: [], reason: "already injected this turn" }
}
const pendingMessageIds = new Set(runtimeMember.pendingInjectedMessageIds)
const unreadMessages = (await listUnreadMessages(teamRunId, memberName, config))
.filter((message) => !pendingMessageIds.has(message.messageId))
if (unreadMessages.length === 0) {
if (pendingMessageIds.size > 0) {
return { injected: false, messageIds: [], reason: "pending ack" }
await transitionRuntimeState(teamRunId, (currentRuntimeState) => {
const runtimeMember = currentRuntimeState.members.find((member) => member.name === memberName)
if (runtimeMember === undefined) {
throw new Error(`runtime member not found for session ${sessionID}: ${memberName}`)
}
return { injected: false, messageIds: [], reason: "no unread" }
if (runtimeMember.lastInjectedTurnMarker === turnMarker) {
result = { injected: false, messageIds: [], reason: "already injected this turn" }
return currentRuntimeState
}
const pendingMessageIds = new Set(runtimeMember.pendingInjectedMessageIds)
const injectableMessages = unreadMessages.filter((message) => !pendingMessageIds.has(message.messageId))
if (injectableMessages.length === 0) {
result = pendingMessageIds.size > 0
? { injected: false, messageIds: [], reason: "pending ack" }
: { injected: false, messageIds: [], reason: "no unread" }
return currentRuntimeState
}
const messageIds: string[] = []
const envelopes: string[] = []
for (const unreadMessage of injectableMessages) {
messageIds.push(unreadMessage.messageId)
envelopes.push(buildEnvelope(unreadMessage))
}
result = { injected: true, content: envelopes.join("\n"), messageIds }
return {
...currentRuntimeState,
members: currentRuntimeState.members.map((member) => (
member.name === memberName
? {
...member,
lastInjectedTurnMarker: turnMarker,
pendingInjectedMessageIds: Array.from(new Set([...member.pendingInjectedMessageIds, ...messageIds])),
}
: member
)),
}
}, config)
if (result === undefined) {
throw new Error(`mailbox injection claim failed for session ${sessionID}: ${memberName}`)
}
const messageIds: string[] = []
const envelopes: string[] = []
for (const unreadMessage of unreadMessages) {
messageIds.push(unreadMessage.messageId)
envelopes.push(buildEnvelope(unreadMessage))
}
const content = envelopes.join("\n")
await transitionRuntimeState(teamRunId, (currentRuntimeState) => ({
...currentRuntimeState,
members: currentRuntimeState.members.map((member) => (
member.name === memberName
? {
...member,
lastInjectedTurnMarker: turnMarker,
pendingInjectedMessageIds: Array.from(new Set([...member.pendingInjectedMessageIds, ...messageIds])),
}
: member
)),
}), config)
return { injected: true, content, messageIds }
return result
}
+69 -1
View File
@@ -1,7 +1,7 @@
/// <reference types="bun-types" />
import { afterEach, describe, expect, test } from "bun:test"
import { mkdtemp, readdir, readFile } from "node:fs/promises"
import { mkdtemp, readdir, readFile, rm } from "node:fs/promises"
import { randomUUID } from "node:crypto"
import { tmpdir } from "node:os"
import path from "node:path"
@@ -673,6 +673,74 @@ describe("createTeamSendMessageTool", () => {
expect(recipient?.pendingInjectedMessageIds).toHaveLength(1)
})
test("#given live delivery prompt dispatches but pending mark fails #when delivery finishes #then the message is not re-exposed as unread", async () => {
// given
const fixture = await createTeamFixture()
let promptCalls = 0
const client = {
session: {
promptAsync: async () => {
promptCalls += 1
await rm(path.join(resolveBaseDir(fixture.config), "runtime", fixture.teamRunId, "state.json"))
},
},
} satisfies LiveDeliveryClient
const liveTool = createTeamSendMessageTool(fixture.config, client)
// when
await liveTool.execute({
teamRunId: fixture.teamRunId,
to: "m2",
body: "accepted before state vanished",
}, fixture.toolContext(fixture.memberOneSessionId))
// then
expect(promptCalls).toBe(1)
const unread = await listUnreadMessages(fixture.teamRunId, "m2", fixture.config)
expect(unread).toHaveLength(0)
const inboxDir = getInboxDir(resolveBaseDir(fixture.config), fixture.teamRunId, "m2")
const inboxEntries = (await readdir(inboxDir)).filter((entry) => entry.endsWith(".json"))
expect(inboxEntries).toHaveLength(1)
expect(inboxEntries[0]?.startsWith(".delivering-")).toBe(true)
})
test("#given live delivery cannot reload runtime after pre-reserve #when delivery aborts #then the message is released for mailbox injection", async () => {
// given
const fixture = await createTeamFixture()
const { loadRuntimeState: loadState } = await import("../team-state-store/store")
const runtimeState = await loadState(fixture.teamRunId, fixture.config)
let loadCount = 0
const deps = {
loadRuntimeState: async () => {
loadCount += 1
if (loadCount === 3) {
throw new Error("runtime reload failed")
}
return runtimeState
},
}
const { client, calls } = createRecordingClient()
const liveTool = createTeamSendMessageTool(fixture.config, client, deps)
// when
await liveTool.execute({
teamRunId: fixture.teamRunId,
to: "m2",
body: "fallback unread",
}, fixture.toolContext(fixture.memberOneSessionId))
// then
expect(calls).toHaveLength(0)
const unread = await listUnreadMessages(fixture.teamRunId, "m2", fixture.config)
expect(unread).toHaveLength(1)
expect(unread[0]?.body).toBe("fallback unread")
const inboxEntries = await readdir(getInboxDir(resolveBaseDir(fixture.config), fixture.teamRunId, "m2"))
expect(inboxEntries.filter((entry) => entry.endsWith(".json") && !entry.startsWith("."))).toHaveLength(1)
expect(inboxEntries.some((entry) => entry.startsWith(".delivering-"))).toBe(false)
})
test("reserves the message during live delivery so concurrent listings cannot surface it", async () => {
// given
const fixture = await createTeamFixture()
+53 -3
View File
@@ -161,6 +161,22 @@ async function markLiveDeliveryPending(
}), config)
}
async function releaseReservationsForRecipients(
teamRunId: string,
recipientNames: readonly string[],
messageId: string,
config: TeamModeConfig,
): Promise<void> {
for (const recipientName of recipientNames) {
const reservation = await reserveMessageForDelivery(teamRunId, recipientName, messageId, config)
await releaseReservationSafely(reservation, {
teamRunId,
recipient: recipientName,
messageId,
})
}
}
async function deliverLive(
client: LiveDeliveryClient,
message: Message,
@@ -170,7 +186,19 @@ async function deliverLive(
directory: string,
deps: TeamSendMessageToolDeps,
): Promise<void> {
const runtimeState = await deps.loadRuntimeState(teamRunId, config)
let runtimeState: RuntimeState
try {
runtimeState = await deps.loadRuntimeState(teamRunId, config)
} catch (error) {
await releaseReservationsForRecipients(teamRunId, deliveredTo, message.messageId, config)
log("[team-mailbox] live delivery unavailable after pre-reserve, released recipients to inbox", {
teamRunId,
messageId: message.messageId,
deliveredTo,
error: error instanceof Error ? error.message : String(error),
})
return
}
const envelope = buildEnvelope(message)
for (const recipientName of deliveredTo) {
@@ -246,7 +274,18 @@ async function deliverLive(
},
})
if (promptResult.status === "failed" && isAmbiguousPromptDispatchFailure(promptResult.error)) {
await markLiveDeliveryPending(teamRunId, recipientName, message.messageId, config)
try {
await markLiveDeliveryPending(teamRunId, recipientName, message.messageId, config)
} catch (markError) {
log("[team-mailbox] live delivery prompt may be accepted but pending mark failed, keeping reservation hidden", {
teamRunId,
recipient: recipientName,
recipientSessionId,
messageId: message.messageId,
error: markError instanceof Error ? markError.message : String(markError),
})
continue
}
log("[team-mailbox] live delivery prompt failed after dispatch attempt, keeping reservation pending", {
teamRunId,
recipient: recipientName,
@@ -271,7 +310,18 @@ async function deliverLive(
})
continue
}
await markLiveDeliveryPending(teamRunId, recipientName, message.messageId, config)
try {
await markLiveDeliveryPending(teamRunId, recipientName, message.messageId, config)
} catch (markError) {
log("[team-mailbox] live delivery prompt dispatched but pending mark failed, keeping reservation hidden", {
teamRunId,
recipient: recipientName,
recipientSessionId,
messageId: message.messageId,
error: markError instanceof Error ? markError.message : String(markError),
})
continue
}
log("[team-mailbox] live delivery reserved until recipient idle", {
teamRunId,
recipient: recipientName,