diff --git a/src/features/team-mode/tools/messaging-missing-session.test.ts b/src/features/team-mode/tools/messaging-missing-session.test.ts new file mode 100644 index 000000000..3dceab444 --- /dev/null +++ b/src/features/team-mode/tools/messaging-missing-session.test.ts @@ -0,0 +1,92 @@ +/// + +import { describe, expect, mock, test } from "bun:test" +import { mkdtemp, readdir } from "node:fs/promises" +import { randomUUID } from "node:crypto" +import { tmpdir } from "node:os" +import path from "node:path" + +import type { ToolContext } from "@opencode-ai/plugin/tool" + +import { TeamModeConfigSchema } from "../../../config/schema/team-mode" +import { getInboxDir, resolveBaseDir } from "../team-registry/paths" + +function createToolContext(sessionID: string, directory: string): ToolContext { + return { + sessionID, + messageID: randomUUID(), + agent: "test-agent", + directory, + worktree: directory, + abort: new AbortController().signal, + metadata: () => {}, + ask: async () => undefined, + } +} + +describe("createTeamSendMessageTool missing recipient session fallback", () => { + test("releases the .delivering reservation when the recipient session disappears before live delivery", async () => { + // given + const baseDir = await mkdtemp(path.join(tmpdir(), "team-send-message-missing-session-")) + const config = TeamModeConfigSchema.parse({ base_dir: baseDir }) + const teamRunId = randomUUID() + const leadSessionId = randomUUID() + const memberOneSessionId = randomUUID() + const memberTwoSessionId = randomUUID() + + const runtimeStateWithRecipientSession = { + teamRunId, + leadSessionId, + status: "active", + members: [ + { name: "team-lead", agentType: "leader", sessionId: leadSessionId }, + { name: "m1", agentType: "member", sessionId: memberOneSessionId }, + { name: "m2", agentType: "member", sessionId: memberTwoSessionId }, + ], + } + const runtimeStateWithoutRecipientSession = { + ...runtimeStateWithRecipientSession, + members: runtimeStateWithRecipientSession.members.map((member) => ( + member.name === "m2" + ? { ...member, sessionId: undefined } + : member + )), + } + + let loadRuntimeStateCalls = 0 + mock.module("../team-state-store/store", () => ({ + listActiveTeams: async () => [{ teamRunId }], + loadRuntimeState: async () => { + loadRuntimeStateCalls += 1 + return loadRuntimeStateCalls >= 3 + ? runtimeStateWithoutRecipientSession + : runtimeStateWithRecipientSession + }, + })) + + const { createTeamSendMessageTool } = await import("./messaging") + type LiveDeliveryClient = Parameters[1] + const client = { + session: { + promptAsync: async () => { + throw new Error("promptAsync should not run when the recipient session is missing") + }, + }, + } satisfies LiveDeliveryClient + const tool = createTeamSendMessageTool(config, client) + + // when + const result = await tool.execute({ + teamRunId, + to: "m2", + body: "ping", + }, createToolContext(memberOneSessionId, baseDir)) + const parsedResult = JSON.parse(result) as { deliveredTo: string[]; messageId: string } + const inboxDir = getInboxDir(resolveBaseDir(config), teamRunId, "m2") + const inboxEntries = (await readdir(inboxDir)).filter((entry) => entry.endsWith(".json")) + + // then + expect(parsedResult.deliveredTo).toEqual(["m2"]) + expect(inboxEntries).toEqual([`${parsedResult.messageId}.json`]) + }) +}) diff --git a/src/features/team-mode/tools/messaging.test.ts b/src/features/team-mode/tools/messaging.test.ts new file mode 100644 index 000000000..0faf67771 --- /dev/null +++ b/src/features/team-mode/tools/messaging.test.ts @@ -0,0 +1,623 @@ +/// + +import { afterEach, describe, expect, test } from "bun:test" +import { mkdtemp, readdir, readFile } from "node:fs/promises" +import { randomUUID } from "node:crypto" +import { tmpdir } from "node:os" +import path from "node:path" + +import { type ToolContext } from "@opencode-ai/plugin/tool" +import { TeamModeConfigSchema } from "../../../config/schema/team-mode" +import { _resetForTesting, registerAgentName } from "../../claude-code-session-state" +import { SessionCategoryRegistry } from "../../../shared/session-category-registry" +import { + clearAllSessionPromptParams, + getSessionPromptParams, +} from "../../../shared/session-prompt-params-state" +import { listUnreadMessages } from "../team-mailbox/inbox" +import { BroadcastNotPermittedError } from "../team-mailbox/send" +import { getInboxDir, resolveBaseDir } from "../team-registry/paths" +import { createRuntimeState, saveRuntimeState } from "../team-state-store/store" +import { clearTeamSessionRegistry, registerTeamSession } from "../team-session-registry" +import type { Message } from "../types" +import { MessageSchema } from "../types" +import { createTeamSendMessageTool } from "./messaging" + +type PromptAsyncCall = { + sessionId: string + parts: Array<{ type: string; text?: string }> + agent?: string + model?: { providerID: string; modelID: string } + variant?: string + directory?: string +} + +type LiveDeliveryClient = { + session: { + promptAsync(input: { + path: { id: string } + body: { + parts: Array<{ type: "text"; text: string }> + agent?: string + model?: { providerID: string; modelID: string } + variant?: string + } + query?: { directory: string } + }): Promise + } +} + +function createRecordingClient(): { client: LiveDeliveryClient; calls: PromptAsyncCall[] } { + const calls: PromptAsyncCall[] = [] + const client = { + session: { + promptAsync: async (input: { + path: { id: string } + body: { + parts: Array<{ type: "text"; text: string }> + agent?: string + model?: { providerID: string; modelID: string } + variant?: string + } + query?: { directory: string } + }) => { + calls.push({ + sessionId: input.path.id, + parts: input.body.parts, + agent: input.body.agent, + model: input.body.model, + variant: input.body.variant, + directory: input.query?.directory, + }) + return undefined + }, + }, + } + return { client, calls } +} + +const mockClient: LiveDeliveryClient = { + session: { + promptAsync: async () => { throw new Error("live delivery disabled in fixture") }, + }, +} + +afterEach(() => { + clearTeamSessionRegistry() + SessionCategoryRegistry.clear() + clearAllSessionPromptParams() + _resetForTesting() +}) + +async function createFixtureBaseDir(): Promise { + return await mkdtemp(path.join(tmpdir(), "team-send-message-")) +} + +function createConfig(baseDir: string) { + return TeamModeConfigSchema.parse({ base_dir: baseDir }) +} + +function createToolContext(sessionID: string, directory: string): ToolContext { + return { + sessionID, + messageID: randomUUID(), + agent: "test-agent", + directory, + worktree: directory, + abort: new AbortController().signal, + metadata: () => {}, + ask: async () => undefined, + } +} + +async function createTeamFixture() { + const baseDir = await createFixtureBaseDir() + const config = createConfig(baseDir) + const leadSessionId = randomUUID() + const memberOneSessionId = randomUUID() + const memberTwoSessionId = randomUUID() + + const runtimeState = await createRuntimeState( + { + version: 1, + name: "team-alpha", + createdAt: Date.now(), + leadAgentId: "team-lead", + members: [ + { kind: "subagent_type", name: "team-lead", subagent_type: "sisyphus-junior", backendType: "in-process", isActive: true }, + { kind: "subagent_type", name: "m1", subagent_type: "sisyphus-junior", backendType: "in-process", isActive: true }, + { kind: "subagent_type", name: "m2", subagent_type: "sisyphus-junior", backendType: "in-process", isActive: true }, + ], + }, + leadSessionId, + "project", + config, + ) + + runtimeState.leadSessionId = leadSessionId + runtimeState.members[0].sessionId = leadSessionId + runtimeState.members[1].sessionId = memberOneSessionId + runtimeState.members[2].sessionId = memberTwoSessionId + runtimeState.members[0].status = "idle" + runtimeState.members[1].status = "idle" + runtimeState.members[2].status = "idle" + await saveRuntimeState(runtimeState, config) + + return { + config, + teamRunId: runtimeState.teamRunId, + leadSessionId, + memberOneSessionId, + memberTwoSessionId, + tool: createTeamSendMessageTool(config, mockClient), + toolContext: (sessionID: string) => createToolContext(sessionID, baseDir), + } +} + +describe("createTeamSendMessageTool", () => { + test("routes a member message to one recipient", async () => { + // given + const fixture = await createTeamFixture() + + // when + const result = await fixture.tool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "hello", + }, fixture.toolContext(fixture.memberOneSessionId)) + const parsedResult = JSON.parse(result) + + // then + expect(parsedResult.deliveredTo).toEqual(["m2"]) + const inboxDir = getInboxDir(resolveBaseDir(fixture.config), fixture.teamRunId, "m2") + const [messageFile] = (await readdir(inboxDir)).filter((entry) => entry.endsWith(".json")) + const message = MessageSchema.parse(JSON.parse(await readFile(path.join(inboxDir, messageFile), "utf8"))) + expect(message.from).toBe("m1") + }) + + test("gates broadcast to the lead and fans out to active members", async () => { + // given + const fixture = await createTeamFixture() + + // when + const nonLeadResult = fixture.tool.execute({ + teamRunId: fixture.teamRunId, + to: "*", + body: "hello everyone", + }, fixture.toolContext(fixture.memberOneSessionId)) + + // then + expect(nonLeadResult).rejects.toBeInstanceOf(BroadcastNotPermittedError) + + // when + const leadResult = await fixture.tool.execute({ + teamRunId: fixture.teamRunId, + to: "*", + body: "team announcement", + kind: "announcement", + }, fixture.toolContext(fixture.leadSessionId)) + const parsedLeadResult = JSON.parse(leadResult) + + // then + expect(parsedLeadResult.deliveredTo).toEqual(["m1", "m2"]) + const memberOneInbox = await readdir(getInboxDir(resolveBaseDir(fixture.config), fixture.teamRunId, "m1")) + const memberTwoInbox = await readdir(getInboxDir(resolveBaseDir(fixture.config), fixture.teamRunId, "m2")) + expect(memberOneInbox.filter((entry) => entry.endsWith(".json") && !entry.startsWith("."))).toHaveLength(1) + expect(memberTwoInbox.filter((entry) => entry.endsWith(".json") && !entry.startsWith("."))).toHaveLength(1) + }) + + test("live-delivers the envelope via promptAsync to the recipient session", async () => { + // given + const fixture = await createTeamFixture() + const { client, calls } = createRecordingClient() + const liveTool = createTeamSendMessageTool(fixture.config, client) + + // when + await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "ping", + }, fixture.toolContext(fixture.memberOneSessionId)) + + // then + expect(calls).toHaveLength(1) + expect(calls[0].sessionId).toBe(fixture.memberTwoSessionId) + expect(calls[0].directory).toBe(resolveBaseDir(fixture.config)) + const envelopeText = calls[0].parts[0]?.text ?? "" + expect(envelopeText).toContain(" { + // given + const fixture = await createTeamFixture() + const { loadRuntimeState: loadState, saveRuntimeState: saveState } = await import("../team-state-store/store") + const state = await loadState(fixture.teamRunId, fixture.config) + const memberTwo = state.members.find((member) => member.name === "m2") + if (!memberTwo) throw new Error("m2 runtime member missing") + memberTwo.worktreePath = "/tmp/team-worker-m2" + await saveState(state, fixture.config) + + const { client, calls } = createRecordingClient() + const liveTool = createTeamSendMessageTool(fixture.config, client) + + // when + await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "ping", + }, fixture.toolContext(fixture.memberOneSessionId)) + + // then + expect(calls).toHaveLength(1) + expect(calls[0]?.directory).toBe("/tmp/team-worker-m2") + }) + + test("live-delivers to running recipients so active teammates receive messages immediately", async () => { + // given + const fixture = await createTeamFixture() + const { loadRuntimeState: loadState, saveRuntimeState: saveState } = await import("../team-state-store/store") + const state = await loadState(fixture.teamRunId, fixture.config) + const memberTwo = state.members.find((member) => member.name === "m2") + if (!memberTwo) throw new Error("m2 runtime member missing") + memberTwo.status = "running" + await saveState(state, fixture.config) + + const { client, calls } = createRecordingClient() + const liveTool = createTeamSendMessageTool(fixture.config, client) + + // when + const result = await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "ping", + }, fixture.toolContext(fixture.memberOneSessionId)) + const parsedResult = JSON.parse(result) + + // then + expect(parsedResult.deliveredTo).toEqual(["m2"]) + expect(calls).toHaveLength(1) + expect(calls[0]?.sessionId).toBe(fixture.memberTwoSessionId) + expect(calls[0]?.directory).toBe(resolveBaseDir(fixture.config)) + }) + + test("live delivery pins the recipient's resolved subagent_type and model on promptAsync", async () => { + // given + const fixture = await createTeamFixture() + const { loadRuntimeState: loadState, saveRuntimeState: saveState } = await import("../team-state-store/store") + const state = await loadState(fixture.teamRunId, fixture.config) + const memberTwo = state.members.find((member) => member.name === "m2") + if (!memberTwo) throw new Error("m2 runtime member missing") + memberTwo.subagent_type = "atlas" + memberTwo.model = { providerID: "anthropic", modelID: "claude-opus-4-7", variant: "high" } + await saveState(state, fixture.config) + + const { client, calls } = createRecordingClient() + const liveTool = createTeamSendMessageTool(fixture.config, client) + + // when + await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "ping", + }, fixture.toolContext(fixture.memberOneSessionId)) + + // then + expect(calls).toHaveLength(1) + expect(calls[0].sessionId).toBe(fixture.memberTwoSessionId) + expect(calls[0].agent).toBe("atlas") + expect(calls[0].model).toEqual({ providerID: "anthropic", modelID: "claude-opus-4-7" }) + expect(calls[0].variant).toBe("high") + }) + + test("live delivery uses the registered agent alias when the runtime stores a config-key agent name", async () => { + // given + registerAgentName("\u200B\u200B\u200B\u200BAtlas - Plan Executor") + const fixture = await createTeamFixture() + const { loadRuntimeState: loadState, saveRuntimeState: saveState } = await import("../team-state-store/store") + const state = await loadState(fixture.teamRunId, fixture.config) + const memberTwo = state.members.find((member) => member.name === "m2") + if (!memberTwo) throw new Error("m2 runtime member missing") + memberTwo.subagent_type = "atlas" + await saveState(state, fixture.config) + + const { client, calls } = createRecordingClient() + const liveTool = createTeamSendMessageTool(fixture.config, client) + + // when + await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "ping", + }, fixture.toolContext(fixture.memberOneSessionId)) + + // then + expect(calls).toHaveLength(1) + expect(calls[0]?.agent).toBe("\u200B\u200B\u200B\u200BAtlas - Plan Executor") + }) + + test("live delivery reapplies category routing and advanced model params for category members", async () => { + // given + const fixture = await createTeamFixture() + const { loadRuntimeState: loadState, saveRuntimeState: saveState } = await import("../team-state-store/store") + const state = await loadState(fixture.teamRunId, fixture.config) + const memberTwo = state.members.find((member) => member.name === "m2") + if (!memberTwo) throw new Error("m2 runtime member missing") + memberTwo.subagent_type = "Sisyphus-Junior" + memberTwo.category = "quick" + memberTwo.model = { + providerID: "openai", + modelID: "gpt-5.4", + variant: "medium", + reasoningEffort: "high", + temperature: 0.2, + top_p: 0.8, + maxTokens: 4096, + thinking: { type: "enabled", budgetTokens: 2048 }, + } + await saveState(state, fixture.config) + + const { client, calls } = createRecordingClient() + const liveTool = createTeamSendMessageTool(fixture.config, client) + + // when + await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "ping", + }, fixture.toolContext(fixture.memberOneSessionId)) + + // then + expect(calls).toHaveLength(1) + expect(calls[0].agent).toBe("Sisyphus-Junior") + expect(calls[0].model).toEqual({ providerID: "openai", modelID: "gpt-5.4" }) + expect(calls[0].variant).toBe("medium") + expect(SessionCategoryRegistry.get(fixture.memberTwoSessionId)).toBe("quick") + expect(getSessionPromptParams(fixture.memberTwoSessionId)).toEqual({ + temperature: 0.2, + topP: 0.8, + maxOutputTokens: 4096, + options: { + reasoningEffort: "high", + thinking: { type: "enabled", budgetTokens: 2048 }, + }, + }) + }) + + test("live delivery omits agent and model on promptAsync when the runtime member has none recorded", async () => { + // given + const fixture = await createTeamFixture() + const { client, calls } = createRecordingClient() + const liveTool = createTeamSendMessageTool(fixture.config, client) + + // when + await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "ping", + }, fixture.toolContext(fixture.memberOneSessionId)) + + // then + expect(calls).toHaveLength(1) + expect(calls[0].agent).toBeUndefined() + expect(calls[0].model).toBeUndefined() + expect(calls[0].variant).toBeUndefined() + }) + + test("prefers the team session registry when the runtime member session has not been persisted yet", async () => { + // given + const fixture = await createTeamFixture() + registerTeamSession(fixture.memberOneSessionId, { + teamRunId: fixture.teamRunId, + memberName: "m1", + role: "member", + }) + + const { loadRuntimeState: loadState, saveRuntimeState: saveState } = await import("../team-state-store/store") + const runtimeState = await loadState(fixture.teamRunId, fixture.config) + const memberOne = runtimeState.members.find((member) => member.name === "m1") + if (!memberOne) throw new Error("m1 runtime member missing") + memberOne.sessionId = undefined + await saveState(runtimeState, fixture.config) + + // when + const result = await fixture.tool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "hello", + }, fixture.toolContext(fixture.memberOneSessionId)) + const parsedResult = JSON.parse(result) + + // then + expect(parsedResult.deliveredTo).toEqual(["m2"]) + const inboxDir = getInboxDir(resolveBaseDir(fixture.config), fixture.teamRunId, "m2") + const [messageFile] = (await readdir(inboxDir)).filter((entry) => entry.endsWith(".json")) + const message = MessageSchema.parse(JSON.parse(await readFile(path.join(inboxDir, messageFile), "utf8"))) + expect(message.from).toBe("m1") + }) + + test("acks the message after live delivery so the transform hook does not redeliver", async () => { + // given + const fixture = await createTeamFixture() + const { client } = createRecordingClient() + const liveTool = createTeamSendMessageTool(fixture.config, client) + + // when + await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "ping", + }, fixture.toolContext(fixture.memberOneSessionId)) + + // then + const inboxDir = getInboxDir(resolveBaseDir(fixture.config), fixture.teamRunId, "m2") + const inboxEntries = (await readdir(inboxDir)).filter((entry) => entry.endsWith(".json")) + const processedEntries = (await readdir(path.join(inboxDir, "processed"))).filter((entry) => entry.endsWith(".json")) + expect(inboxEntries).toHaveLength(0) + expect(processedEntries).toHaveLength(1) + }) + + test("broadcast fans out live delivery to every member except the sender", async () => { + // given + const fixture = await createTeamFixture() + const { client, calls } = createRecordingClient() + const liveTool = createTeamSendMessageTool(fixture.config, client) + + // when + await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "*", + body: "broadcast ping", + kind: "announcement", + }, fixture.toolContext(fixture.leadSessionId)) + + // then + const targetedSessionIds = calls.map((entry) => entry.sessionId).sort() + expect(targetedSessionIds).toEqual([ + fixture.memberOneSessionId, + fixture.memberTwoSessionId, + ].sort()) + }) + + test("broadcast still queues for members whose session has not spawned yet", async () => { + // given + const fixture = await createTeamFixture() + const { loadRuntimeState: loadState } = await import("../team-state-store/store") + const stateBefore = await loadState(fixture.teamRunId, fixture.config) + const pendingMember = stateBefore.members.find((member) => member.name === "m2") + if (!pendingMember) throw new Error("m2 runtime member missing") + pendingMember.sessionId = undefined + await saveRuntimeState(stateBefore, fixture.config) + + const { client, calls } = createRecordingClient() + const liveTool = createTeamSendMessageTool(fixture.config, client) + + // when + const result = await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "*", + body: "broadcast ping", + kind: "announcement", + }, fixture.toolContext(fixture.leadSessionId)) + const parsedResult = JSON.parse(result) + + // then + expect(parsedResult.deliveredTo).toEqual(["m1", "m2"]) + const targetedSessionIds = calls.map((entry) => entry.sessionId) + expect(targetedSessionIds).toEqual([fixture.memberOneSessionId]) + const memberTwoInbox = await readdir(getInboxDir(resolveBaseDir(fixture.config), fixture.teamRunId, "m2")) + expect(memberTwoInbox.filter((entry) => entry.endsWith(".json") && !entry.startsWith("."))).toHaveLength(1) + }) + + test("inbox stays intact when live delivery fails so the fallback path still works", async () => { + // given + const fixture = await createTeamFixture() + const failingClient = { + session: { + promptAsync: async () => { throw new Error("network down") }, + }, + } satisfies LiveDeliveryClient + const liveTool = createTeamSendMessageTool(fixture.config, failingClient) + + // when + await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "ping", + }, fixture.toolContext(fixture.memberOneSessionId)) + + // then + const inboxDir = getInboxDir(resolveBaseDir(fixture.config), fixture.teamRunId, "m2") + const inboxEntries = (await readdir(inboxDir)).filter((entry) => entry.endsWith(".json") && !entry.startsWith(".")) + expect(inboxEntries).toHaveLength(1) + }) + + test("reserves the message during live delivery so concurrent listings cannot surface it", async () => { + // given + const fixture = await createTeamFixture() + let unreadDuringDelivery: Message[] = [] + const reservingClient = { + session: { + promptAsync: async () => { + unreadDuringDelivery = await listUnreadMessages(fixture.teamRunId, "m2", fixture.config) + return undefined + }, + }, + } satisfies LiveDeliveryClient + const liveTool = createTeamSendMessageTool(fixture.config, reservingClient) + + // when + await liveTool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "ping", + }, fixture.toolContext(fixture.memberOneSessionId)) + + // then + expect(unreadDuringDelivery).toHaveLength(0) + }) + + test("hides the message from the inbox from the moment it is written for a live recipient", async () => { + // given + const fixture = await createTeamFixture() + const { sendMessage } = await import("../team-mailbox/send") + const messageId = randomUUID() + + // when + await sendMessage({ + version: 1, + messageId, + from: "m1", + to: "m2", + kind: "message", + body: "ping", + timestamp: Date.now(), + }, fixture.teamRunId, fixture.config, { + isLead: false, + activeMembers: ["m2"], + reservedRecipients: new Set(["m2"]), + }) + const unreadImmediately = await listUnreadMessages(fixture.teamRunId, "m2", fixture.config) + const inboxDir = getInboxDir(resolveBaseDir(fixture.config), fixture.teamRunId, "m2") + const rawEntries = (await readdir(inboxDir)) + .filter((entry) => entry.endsWith(".json")) + + // then + expect(unreadImmediately).toHaveLength(0) + expect(rawEntries).toEqual([`.delivering-${messageId}.json`]) + }) + + test("rejects shutdown_request kind", async () => { + // given + const fixture = await createTeamFixture() + + // when + const result = fixture.tool.execute({ + teamRunId: fixture.teamRunId, + to: "m1", + body: "stop", + kind: "shutdown_request", + }, fixture.toolContext(fixture.leadSessionId)) + + // then + expect(result).rejects.toBeInstanceOf(Error) + }) + + test("rejects a non-UUID correlationId before writing the message", async () => { + // given + const fixture = await createTeamFixture() + + // when + const result = fixture.tool.execute({ + teamRunId: fixture.teamRunId, + to: "m2", + body: "hello", + correlationId: "task-1", + }, fixture.toolContext(fixture.memberOneSessionId)) + + // then + await expect(result).rejects.toThrow("correlationId") + await expect(readdir(getInboxDir(resolveBaseDir(fixture.config), fixture.teamRunId, "m2"))).rejects.toThrow() + }) +}) diff --git a/src/features/team-mode/tools/messaging.ts b/src/features/team-mode/tools/messaging.ts new file mode 100644 index 000000000..bee322ed1 --- /dev/null +++ b/src/features/team-mode/tools/messaging.ts @@ -0,0 +1,271 @@ +import { randomUUID } from "node:crypto" + +import { tool, type ToolDefinition } from "@opencode-ai/plugin/tool" +import { z } from "zod" + +import type { TeamModeConfig } from "../../../config/schema/team-mode" +import { log } from "../../../shared/logger" +import { applyMemberSessionRouting, buildMemberPromptBody } from "../member-session-routing" +import { lookupTeamSession } from "../team-session-registry" +import { loadRuntimeState } from "../team-state-store/store" +import { buildEnvelope } from "../team-mailbox/poll" +import { + commitDeliveryReservation, + releaseDeliveryReservation, + reserveMessageForDelivery, +} from "../team-mailbox/reservation" +import { BroadcastNotPermittedError, sendMessage } from "../team-mailbox/send" + +import type { Message } from "../types" +import { MessageSchema } from "../types" + +const MESSAGE_TOOL_KINDS = ["message", "announcement"] as const + +export type LiveDeliveryClient = { + session: { + promptAsync(input: { + path: { id: string } + body: { + parts: Array<{ type: "text"; text: string }> + agent?: string + model?: { providerID: string; modelID: string } + variant?: string + } + query?: { directory: string } + }): Promise + } +} + +type TeamRuntimeDetails = { + teamRunId: string + isLead: boolean + senderName: string + activeMembers: string[] +} + +const TeamReferenceArgsSchema = z.object({ + path: z.string().min(1), + description: z.string().optional(), +}) + +const TeamSendMessageArgsSchema = z.object({ + teamRunId: z.string().min(1), + to: z.string().min(1), + body: z.string(), + kind: z.enum(MESSAGE_TOOL_KINDS).optional(), + correlationId: z.string().uuid().optional(), + summary: z.string().optional(), + references: z.array(TeamReferenceArgsSchema).optional(), +}) + +type DeliveryReservation = Awaited> + +async function resolveTeamRuntimeDetails(teamRunId: string, sessionID: string, config: TeamModeConfig): Promise { + const registryEntry = lookupTeamSession(sessionID) + if (registryEntry?.teamRunId === teamRunId) { + const runtimeState = await loadRuntimeState(teamRunId, config) + + return { + teamRunId: runtimeState.teamRunId, + isLead: registryEntry.role === "lead", + senderName: registryEntry.memberName, + activeMembers: runtimeState.members + .map((entry) => entry.name) + .filter((name) => name !== registryEntry.memberName), + } + } + + try { + const runtimeState = await loadRuntimeState(teamRunId, config) + const isLead = runtimeState.leadSessionId === sessionID + const leadMember = isLead + ? runtimeState.members.find((member) => member.agentType === "leader") + : undefined + const member = runtimeState.members.find((entry) => entry.sessionId === sessionID) + const senderName = leadMember?.name ?? member?.name ?? "unknown" + + return { + teamRunId: runtimeState.teamRunId, + isLead, + senderName, + activeMembers: runtimeState.members + .map((entry) => entry.name) + .filter((name) => name !== senderName), + } + } catch { + return { + teamRunId, + isLead: false, + senderName: "unknown", + activeMembers: [], + } + } +} + +async function releaseReservationSafely( + reservation: DeliveryReservation, + input: { teamRunId: string; recipient: string; messageId: string }, +): Promise { + if (reservation === null) return + + try { + await releaseDeliveryReservation(reservation) + } catch (releaseError) { + log("[team-mailbox] failed to release delivery reservation", { + error: releaseError instanceof Error ? releaseError.message : String(releaseError), + teamRunId: input.teamRunId, + recipient: input.recipient, + messageId: input.messageId, + }) + } +} + +async function deliverLive( + client: LiveDeliveryClient, + message: Message, + teamRunId: string, + deliveredTo: readonly string[], + config: TeamModeConfig, + directory: string, +): Promise { + const runtimeState = await loadRuntimeState(teamRunId, config) + const envelope = buildEnvelope(message) + + for (const recipientName of deliveredTo) { + // Reserve the inbox file before delivering so the transform-hook fallback + // cannot re-read the same message while promptAsync is in flight. + const reservation = await reserveMessageForDelivery(teamRunId, recipientName, message.messageId, config) + if (reservation === null) continue + + const recipientMember = runtimeState.members.find((entry) => entry.name === recipientName) + if (!recipientMember) { + await releaseReservationSafely(reservation, { + teamRunId, + recipient: recipientName, + messageId: message.messageId, + }) + continue + } + + const recipientSessionId = recipientMember.sessionId + if (!recipientSessionId) { + log("[team-mailbox] live delivery unavailable, falling back to inbox injection", { + reason: "missing-session-id", + teamRunId, + recipient: recipientName, + messageId: message.messageId, + }) + await releaseReservationSafely(reservation, { + teamRunId, + recipient: recipientName, + messageId: message.messageId, + }) + continue + } + + applyMemberSessionRouting(recipientSessionId, recipientMember) + + try { + await client.session.promptAsync({ + path: { id: recipientSessionId }, + body: buildMemberPromptBody(recipientMember, envelope), + query: { directory: recipientMember.worktreePath ?? directory }, + }) + await commitDeliveryReservation(reservation) + log("[team-mailbox] live delivery committed", { + teamRunId, + recipient: recipientName, + recipientSessionId, + messageId: message.messageId, + }) + } catch (error) { + log("[team-mailbox] live delivery failed, falling back to inbox injection", { + error: error instanceof Error ? error.message : String(error), + teamRunId, + recipient: recipientName, + messageId: message.messageId, + }) + await releaseReservationSafely(reservation, { + teamRunId, + recipient: recipientName, + messageId: message.messageId, + }) + } + } +} + +export function createTeamSendMessageTool(config: TeamModeConfig, client: LiveDeliveryClient): ToolDefinition { + return tool({ + description: "Send a message to a team member or broadcast to the team.", + args: { + teamRunId: tool.schema.string().describe("Team run ID"), + to: tool.schema.string().describe("Recipient name or * for broadcast"), + body: tool.schema.string().describe("Message body"), + kind: tool.schema.enum(MESSAGE_TOOL_KINDS).optional().default("message").describe("Message kind"), + correlationId: tool.schema.string().optional().describe("Optional UUID correlation ID. Do not use task IDs like 'task-1'."), + summary: tool.schema.string().optional().describe("Optional summary"), + references: tool.schema.array(tool.schema.object({ + path: tool.schema.string(), + description: tool.schema.string().optional(), + })).optional().describe("Optional references as [{ path, description? }]"), + }, + execute: async (rawArgs, context) => { + const args = TeamSendMessageArgsSchema.parse(rawArgs) + const runtimeContext = context as { sessionID?: string; directory?: string } + const sessionID = runtimeContext.sessionID + + if (!sessionID) { + throw new Error("session ID is required") + } + + const targetDirectory = typeof runtimeContext.directory === "string" ? runtimeContext.directory : process.cwd() + + const teamRuntime = await resolveTeamRuntimeDetails(args.teamRunId, sessionID, config) + const message = MessageSchema.parse({ + version: 1, + messageId: randomUUID(), + from: teamRuntime.senderName, + to: args.to, + body: args.body, + kind: args.kind ?? "message", + timestamp: Date.now(), + correlationId: args.correlationId, + summary: args.summary, + references: args.references, + }) + + if (message.kind === "shutdown_request" || message.kind === "shutdown_approved" || message.kind === "shutdown_rejected") { + throw new Error("must use lifecycle tools for shutdown kinds") + } + + if (message.to === "*" && !teamRuntime.isLead) { + throw new BroadcastNotPermittedError() + } + + const runtimeState = await loadRuntimeState(teamRuntime.teamRunId, config) + const reservedRecipients = new Set( + runtimeState.members + .filter((member) => member.sessionId !== undefined && member.name !== teamRuntime.senderName) + .map((member) => member.name), + ) + + const result = await sendMessage(message, teamRuntime.teamRunId, config, { + isLead: teamRuntime.isLead, + activeMembers: teamRuntime.activeMembers, + reservedRecipients, + }) + + try { + await deliverLive(client, message, teamRuntime.teamRunId, result.deliveredTo, config, targetDirectory) + } catch (liveError) { + log("[team-mailbox] deliverLive top-level error (message already in inbox, safe to ignore)", { + error: liveError instanceof Error ? liveError.message : String(liveError), + teamRunId: teamRuntime.teamRunId, + messageId: message.messageId, + }) + } + + return JSON.stringify(result) + }, + }) +}