diff --git a/src/tools/agent-teams/inbox-message-sender.ts b/src/tools/agent-teams/inbox-message-sender.ts new file mode 100644 index 00000000..d4801170 --- /dev/null +++ b/src/tools/agent-teams/inbox-message-sender.ts @@ -0,0 +1,61 @@ +import { randomUUID } from "node:crypto" +import { InboxMessageSchema } from "./types" +import { appendInboxMessage } from "./inbox-store" + +function nowIso(): string { + return new Date().toISOString() +} + +const STRUCTURED_TYPE_MAP: Record = { + shutdown_request: "shutdown_request", + shutdown_approved: "shutdown_response", + shutdown_rejected: "shutdown_response", + plan_approved: "plan_approval_response", + plan_rejected: "plan_approval_response", +} + +export function buildShutdownRequestId(recipient: string): string { + return `shutdown-${recipient}-${randomUUID().slice(0, 8)}` +} + +export function sendPlainInboxMessage( + teamName: string, + sender: string, + recipient: string, + content: string, + summary: string, + _color?: string, +): void { + const message = InboxMessageSchema.parse({ + id: randomUUID(), + type: "message", + sender, + recipient, + content, + summary, + timestamp: nowIso(), + read: false, + }) + appendInboxMessage(teamName, recipient, message) +} + +export function sendStructuredInboxMessage( + teamName: string, + sender: string, + recipient: string, + data: Record, + summaryType: string, +): void { + const messageType = STRUCTURED_TYPE_MAP[summaryType] ?? "message" + const message = InboxMessageSchema.parse({ + id: randomUUID(), + type: messageType, + sender, + recipient, + content: JSON.stringify(data), + summary: summaryType, + timestamp: nowIso(), + read: false, + }) + appendInboxMessage(teamName, recipient, message) +} diff --git a/src/tools/agent-teams/inbox-store.ts b/src/tools/agent-teams/inbox-store.ts index 5c464dc6..a22dfb15 100644 --- a/src/tools/agent-teams/inbox-store.ts +++ b/src/tools/agent-teams/inbox-store.ts @@ -2,7 +2,8 @@ import { existsSync, readFileSync, unlinkSync } from "node:fs" import { z } from "zod" import { acquireLock, ensureDir, writeJsonAtomic } from "../../features/claude-tasks/storage" import { getTeamInboxPath } from "./paths" -import { InboxMessage, InboxMessageSchema } from "./types" +import type { InboxMessage } from "./types" +import { InboxMessageSchema } from "./types" const InboxMessageListSchema = z.array(InboxMessageSchema) @@ -110,16 +111,9 @@ export function appendInboxMessage(teamName: string, agentName: string, message: }) } -export interface ReadInboxOptions { - unreadOnly?: boolean - markAsRead?: boolean -} - -export function readInbox(teamName: string, agentName: string, options?: ReadInboxOptions): InboxMessage[] { +export function readInbox(teamName: string, agentName: string, unreadOnly = false, markAsRead = false): InboxMessage[] { return withInboxLock(teamName, () => { const messages = readInboxMessages(teamName, agentName) - const unreadOnly = options?.unreadOnly ?? false - const markAsRead = options?.markAsRead ?? false const selectedIndexes = new Set() @@ -197,3 +191,7 @@ export function deleteInbox(teamName: string, agentName: string): void { } }) } + +export const clearInbox = deleteInbox + +export { buildShutdownRequestId, sendPlainInboxMessage, sendStructuredInboxMessage } from "./inbox-message-sender" diff --git a/src/tools/agent-teams/messaging-tools.test.ts b/src/tools/agent-teams/messaging-tools.test.ts index 18c2ccd0..32ee20fc 100644 --- a/src/tools/agent-teams/messaging-tools.test.ts +++ b/src/tools/agent-teams/messaging-tools.test.ts @@ -1,9 +1,11 @@ /// import { afterEach, beforeEach, describe, expect, test } from "bun:test" +import { randomUUID } from "node:crypto" import { mkdtempSync, rmSync } from "node:fs" import { tmpdir } from "node:os" import { join } from "node:path" import type { BackgroundManager } from "../../features/background-agent" +import { readInbox } from "./inbox-store" import { createAgentTeamsTools } from "./tools" interface TestToolContext { @@ -37,7 +39,11 @@ async function executeJsonTool( return JSON.parse(output) } -function createManagerWithImmediateResume(): { manager: BackgroundManager; resumeCalls: ResumeCall[] } { +function uniqueTeam(): string { + return `msg-${randomUUID().slice(0, 8)}` +} + +function createMockManager(): { manager: BackgroundManager; resumeCalls: ResumeCall[] } { const resumeCalls: ResumeCall[] = [] let launchCount = 0 @@ -56,39 +62,19 @@ function createManagerWithImmediateResume(): { manager: BackgroundManager; resum return { manager, resumeCalls } } -function createManagerWithDeferredResume(): { - manager: BackgroundManager - resumeCalls: ResumeCall[] - resolveAllResumes: () => void -} { - const resumeCalls: ResumeCall[] = [] - const pendingResolves: Array<() => void> = [] - let launchCount = 0 - - const manager = { - launch: async () => { - launchCount += 1 - return { id: `bg-${launchCount}`, sessionID: `ses-worker-${launchCount}` } - }, - getTask: () => undefined, - resume: (args: ResumeCall) => { - resumeCalls.push(args) - return new Promise<{ id: string }>((resolve) => { - pendingResolves.push(() => resolve({ id: `resume-${resumeCalls.length}` })) - }) - }, - } as unknown as BackgroundManager - - return { - manager, - resumeCalls, - resolveAllResumes: () => { - while (pendingResolves.length > 0) { - const next = pendingResolves.shift() - next?.() - } - }, - } +async function setupTeamWithWorker( + tools: ReturnType, + context: TestToolContext, + teamName = "core", + workerName = "worker_1", +): Promise { + await executeJsonTool(tools, "team_create", { team_name: teamName }, context) + await executeJsonTool( + tools, + "spawn_teammate", + { team_name: teamName, name: workerName, prompt: "Handle tasks", category: "quick" }, + context, + ) } describe("agent-teams messaging tools", () => { @@ -106,106 +92,342 @@ describe("agent-teams messaging tools", () => { rmSync(tempProjectDir, { recursive: true, force: true }) }) - test("send_message rejects recipient team suffix mismatch", async () => { - //#given - const { manager, resumeCalls } = createManagerWithImmediateResume() - const tools = createAgentTeamsTools(manager) - const leadContext = createContext() - await executeJsonTool(tools, "team_create", { team_name: "core" }, leadContext) - await executeJsonTool( - tools, - "spawn_teammate", - { team_name: "core", name: "worker_1", prompt: "Handle release prep", category: "quick" }, - leadContext, - ) + describe("message type", () => { + test("delivers message to recipient inbox", async () => { + //#given + const tn = uniqueTeam() + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) - //#when - const mismatchedRecipient = await executeJsonTool( - tools, - "send_message", - { - team_name: "core", - type: "message", - recipient: "worker_1@other-team", - summary: "sync", - content: "Please update status.", - }, - leadContext, - ) as { error?: string } - - //#then - expect(mismatchedRecipient.error).toBe("recipient_team_mismatch") - expect(resumeCalls).toHaveLength(0) - }) - - test("send_message rejects recipient with empty team suffix", async () => { - //#given - const { manager, resumeCalls } = createManagerWithImmediateResume() - const tools = createAgentTeamsTools(manager) - const leadContext = createContext() - await executeJsonTool(tools, "team_create", { team_name: "core" }, leadContext) - await executeJsonTool( - tools, - "spawn_teammate", - { team_name: "core", name: "worker_1", prompt: "Handle release prep", category: "quick" }, - leadContext, - ) - - //#when - const invalidRecipient = await executeJsonTool( - tools, - "send_message", - { - team_name: "core", - type: "message", - recipient: "worker_1@", - summary: "sync", - content: "Please update status.", - }, - leadContext, - ) as { error?: string } - - //#then - expect(invalidRecipient.error).toBe("recipient_team_invalid") - expect(resumeCalls).toHaveLength(0) - }) - - test("broadcast schedules teammate resumes without serial await", async () => { - //#given - const { manager, resumeCalls, resolveAllResumes } = createManagerWithDeferredResume() - const tools = createAgentTeamsTools(manager) - const leadContext = createContext() - await executeJsonTool(tools, "team_create", { team_name: "core" }, leadContext) - - for (const name of ["worker_1", "worker_2", "worker_3"]) { - await executeJsonTool( + //#when + const result = await executeJsonTool( tools, - "spawn_teammate", - { team_name: "core", name, prompt: "Handle release prep", category: "quick" }, + "send_message", + { + team_name: tn, + type: "message", + recipient: "worker_1", + content: "Please update status.", + summary: "status_request", + }, leadContext, - ) - } + ) as { success?: boolean; message?: string } - //#when - const broadcastPromise = executeJsonTool( - tools, - "send_message", - { team_name: "core", type: "broadcast", summary: "sync", content: "Please update status." }, - leadContext, - ) as Promise<{ success?: boolean; message?: string }> + //#then + expect(result.success).toBe(true) + expect(result.message).toBe("message_sent:worker_1") + const inbox = readInbox(tn, "worker_1") + const delivered = inbox.filter((m) => m.summary === "status_request") + expect(delivered.length).toBeGreaterThanOrEqual(1) + expect(delivered[0]?.sender).toBe("team-lead") + expect(delivered[0]?.content).toBe("Please update status.") + }) - await Promise.resolve() - await Promise.resolve() + test("rejects message to nonexistent recipient", async () => { + //#given + const tn = uniqueTeam() + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) - //#then - expect(resumeCalls).toHaveLength(3) + //#when + const result = await executeJsonTool( + tools, + "send_message", + { team_name: tn, type: "message", recipient: "nonexistent", content: "hello", summary: "test" }, + leadContext, + ) as { error?: string } - //#when - resolveAllResumes() - const broadcastResult = await broadcastPromise + //#then + expect(result.error).toBe("message_recipient_not_found") + }) - //#then - expect(broadcastResult.success).toBe(true) - expect(broadcastResult.message).toBe("broadcast_sent:3") + test("rejects recipient with team suffix mismatch", async () => { + //#given + const tn = uniqueTeam() + const { manager, resumeCalls } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { team_name: tn, type: "message", recipient: "worker_1@other-team", summary: "sync", content: "hi" }, + leadContext, + ) as { error?: string } + + //#then + expect(result.error).toBe("recipient_team_mismatch") + expect(resumeCalls).toHaveLength(0) + }) + + test("rejects recipient with empty team suffix", async () => { + //#given + const tn = uniqueTeam() + const { manager, resumeCalls } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { team_name: tn, type: "message", recipient: "worker_1@", summary: "sync", content: "hi" }, + leadContext, + ) as { error?: string } + + //#then + expect(result.error).toBe("recipient_team_invalid") + expect(resumeCalls).toHaveLength(0) + }) + }) + + describe("broadcast type", () => { + test("writes to all teammate inboxes", async () => { + //#given + const tn = uniqueTeam() + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await executeJsonTool(tools, "team_create", { team_name: tn }, leadContext) + for (const name of ["worker_1", "worker_2"]) { + await executeJsonTool( + tools, + "spawn_teammate", + { team_name: tn, name, prompt: "Handle tasks", category: "quick" }, + leadContext, + ) + } + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { team_name: tn, type: "broadcast", summary: "sync", content: "Status update needed" }, + leadContext, + ) as { success?: boolean; message?: string } + + //#then + expect(result.success).toBe(true) + expect(result.message).toBe("broadcast_sent:2") + for (const name of ["worker_1", "worker_2"]) { + const inbox = readInbox(tn, name) + const broadcastMessages = inbox.filter((m) => m.summary === "sync") + expect(broadcastMessages.length).toBeGreaterThanOrEqual(1) + } + }) + + test("rejects broadcast without summary", async () => { + //#given + const tn = uniqueTeam() + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { team_name: tn, type: "broadcast", content: "hello" }, + leadContext, + ) as { error?: string } + + //#then + expect(result.error).toBe("broadcast_requires_summary") + }) + }) + + describe("shutdown_request type", () => { + test("sends shutdown request and returns request_id", async () => { + //#given + const tn = uniqueTeam() + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { team_name: tn, type: "shutdown_request", recipient: "worker_1", content: "Work completed" }, + leadContext, + ) as { success?: boolean; request_id?: string; target?: string } + + //#then + expect(result.success).toBe(true) + expect(result.request_id).toMatch(/^shutdown-worker_1-/) + expect(result.target).toBe("worker_1") + }) + + test("rejects shutdown targeting team-lead", async () => { + //#given + const tn = uniqueTeam() + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { team_name: tn, type: "shutdown_request", recipient: "team-lead" }, + leadContext, + ) as { error?: string } + + //#then + expect(result.error).toBe("cannot_shutdown_team_lead") + }) + }) + + describe("shutdown_response type", () => { + test("sends approved shutdown response to team-lead inbox", async () => { + //#given + const tn = uniqueTeam() + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { team_name: tn, type: "shutdown_response", request_id: "shutdown-worker_1-abc12345", approve: true }, + leadContext, + ) as { success?: boolean; message?: string } + + //#then + expect(result.success).toBe(true) + expect(result.message).toBe("shutdown_approved:shutdown-worker_1-abc12345") + const leadInbox = readInbox(tn, "team-lead") + const shutdownMessages = leadInbox.filter((m) => m.summary === "shutdown_approved") + expect(shutdownMessages.length).toBeGreaterThanOrEqual(1) + }) + + test("sends rejected shutdown response", async () => { + //#given + const tn = uniqueTeam() + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { + team_name: tn, + type: "shutdown_response", + request_id: "shutdown-worker_1-abc12345", + approve: false, + content: "Still working on it", + }, + leadContext, + ) as { success?: boolean; message?: string } + + //#then + expect(result.success).toBe(true) + expect(result.message).toBe("shutdown_rejected:shutdown-worker_1-abc12345") + }) + }) + + describe("plan_approval_response type", () => { + test("sends approved plan response", async () => { + //#given + const tn = uniqueTeam() + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { team_name: tn, type: "plan_approval_response", request_id: "plan-req-001", approve: true, recipient: "worker_1" }, + leadContext, + ) as { success?: boolean; message?: string } + + //#then + expect(result.success).toBe(true) + expect(result.message).toBe("plan_approved:worker_1") + }) + + test("sends rejected plan response with content", async () => { + //#given + const tn = uniqueTeam() + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { + team_name: tn, + type: "plan_approval_response", + request_id: "plan-req-002", + approve: false, + recipient: "worker_1", + content: "Need more details", + }, + leadContext, + ) as { success?: boolean; message?: string } + + //#then + expect(result.success).toBe(true) + expect(result.message).toBe("plan_rejected:worker_1") + }) + }) + + describe("authorization", () => { + test("rejects message from unauthorized session", async () => { + //#given + const tn = uniqueTeam() + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext() + await setupTeamWithWorker(tools, leadContext, tn) + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { team_name: tn, type: "message", recipient: "worker_1", content: "hello", summary: "test" }, + createContext("ses-intruder"), + ) as { error?: string } + + //#then + expect(result.error).toBe("unauthorized_sender_session") + }) + + test("rejects message to nonexistent team", async () => { + //#given + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + + //#when + const result = await executeJsonTool( + tools, + "send_message", + { team_name: "nonexistent-xyz", type: "message", recipient: "w", content: "hello", summary: "test" }, + createContext(), + ) as { error?: string } + + //#then + expect(result.error).toBe("team_not_found") + }) }) })