From f422cfc7af6cd8d73b640f66aea8f0163b8e475a Mon Sep 17 00:00:00 2001 From: Nguyen Khac Trung Kien Date: Sun, 8 Feb 2026 13:35:03 +0700 Subject: [PATCH] fix(agent-teams): harden deletion and messaging safety --- src/tools/agent-teams/messaging-tools.test.ts | 179 ++++++++++++++++++ src/tools/agent-teams/messaging-tools.ts | 32 +++- .../agent-teams/team-config-store.test.ts | 20 +- src/tools/agent-teams/team-config-store.ts | 8 +- .../agent-teams/team-lifecycle-tools.test.ts | 69 +++++++ src/tools/agent-teams/team-lifecycle-tools.ts | 5 +- 6 files changed, 305 insertions(+), 8 deletions(-) create mode 100644 src/tools/agent-teams/messaging-tools.test.ts create mode 100644 src/tools/agent-teams/team-lifecycle-tools.test.ts diff --git a/src/tools/agent-teams/messaging-tools.test.ts b/src/tools/agent-teams/messaging-tools.test.ts new file mode 100644 index 000000000..3d4e72fb1 --- /dev/null +++ b/src/tools/agent-teams/messaging-tools.test.ts @@ -0,0 +1,179 @@ +/// +import { afterEach, beforeEach, describe, expect, test } from "bun:test" +import { mkdtempSync, rmSync } from "node:fs" +import { tmpdir } from "node:os" +import { join } from "node:path" +import type { BackgroundManager } from "../../features/background-agent" +import { createAgentTeamsTools } from "./tools" + +interface TestToolContext { + sessionID: string + messageID: string + agent: string + abort: AbortSignal +} + +interface ResumeCall { + sessionId: string + prompt: string +} + +function createContext(sessionID = "ses-main"): TestToolContext { + return { + sessionID, + messageID: "msg-main", + agent: "sisyphus", + abort: new AbortController().signal, + } +} + +async function executeJsonTool( + tools: ReturnType, + toolName: keyof ReturnType, + args: Record, + context: TestToolContext, +): Promise { + const output = await tools[toolName].execute(args, context) + return JSON.parse(output) +} + +function createManagerWithImmediateResume(): { manager: BackgroundManager; resumeCalls: ResumeCall[] } { + const resumeCalls: ResumeCall[] = [] + let launchCount = 0 + + const manager = { + launch: async () => { + launchCount += 1 + return { id: `bg-${launchCount}`, sessionID: `ses-worker-${launchCount}` } + }, + getTask: () => undefined, + resume: async (args: ResumeCall) => { + resumeCalls.push(args) + return { id: `resume-${resumeCalls.length}` } + }, + } as unknown as BackgroundManager + + 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?.() + } + }, + } +} + +describe("agent-teams messaging tools", () => { + let originalCwd: string + let tempProjectDir: string + + beforeEach(() => { + originalCwd = process.cwd() + tempProjectDir = mkdtempSync(join(tmpdir(), "agent-teams-messaging-")) + process.chdir(tempProjectDir) + }) + + afterEach(() => { + process.chdir(originalCwd) + 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, + ) + + //#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("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( + tools, + "spawn_teammate", + { team_name: "core", name, prompt: "Handle release prep", category: "quick" }, + leadContext, + ) + } + + //#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 }> + + await Promise.resolve() + await Promise.resolve() + + //#then + expect(resumeCalls).toHaveLength(3) + + //#when + resolveAllResumes() + const broadcastResult = await broadcastPromise + + //#then + expect(broadcastResult.success).toBe(true) + expect(broadcastResult.message).toBe("broadcast_sent:3") + }) +}) diff --git a/src/tools/agent-teams/messaging-tools.ts b/src/tools/agent-teams/messaging-tools.ts index 8010b613d..43add6461 100644 --- a/src/tools/agent-teams/messaging-tools.ts +++ b/src/tools/agent-teams/messaging-tools.ts @@ -16,6 +16,25 @@ function nowIso(): string { return new Date().toISOString() } +function validateRecipientTeam(recipient: unknown, teamName: string): string | null { + if (typeof recipient !== "string") { + return null + } + + const trimmed = recipient.trim() + const atIndex = trimmed.indexOf("@") + if (atIndex <= 0) { + return null + } + + const specifiedTeam = trimmed.slice(atIndex + 1).trim() + if (!specifiedTeam || specifiedTeam === teamName) { + return null + } + + return "recipient_team_mismatch" +} + function resolveSenderFromContext(config: TeamConfig, context: TeamToolContext): string | null { if (context.sessionID === config.leadSessionId) { return "team-lead" @@ -45,6 +64,10 @@ export function createSendMessageTool(manager: BackgroundManager): ToolDefinitio if (teamError) { return JSON.stringify({ error: teamError }) } + const recipientTeamError = validateRecipientTeam(args.recipient, input.team_name) + if (recipientTeamError) { + return JSON.stringify({ error: recipientTeamError }) + } const requestedSender = input.sender const senderError = requestedSender ? validateAgentNameOrLead(requestedSender) : null if (senderError) { @@ -88,11 +111,16 @@ export function createSendMessageTool(manager: BackgroundManager): ToolDefinitio if (!input.summary) { return JSON.stringify({ error: "broadcast_requires_summary" }) } + const broadcastSummary = input.summary const teammates = listTeammates(config) for (const teammate of teammates) { - sendPlainInboxMessage(input.team_name, sender, teammate.name, input.content ?? "", input.summary) - await resumeTeammateWithMessage(manager, context, input.team_name, teammate, input.summary, input.content ?? "") + sendPlainInboxMessage(input.team_name, sender, teammate.name, input.content ?? "", broadcastSummary) } + await Promise.allSettled( + teammates.map((teammate) => + resumeTeammateWithMessage(manager, context, input.team_name, teammate, broadcastSummary, input.content ?? ""), + ), + ) return JSON.stringify({ success: true, message: `broadcast_sent:${teammates.length}` }) } diff --git a/src/tools/agent-teams/team-config-store.test.ts b/src/tools/agent-teams/team-config-store.test.ts index d1f9c1f45..254056ab6 100644 --- a/src/tools/agent-teams/team-config-store.test.ts +++ b/src/tools/agent-teams/team-config-store.test.ts @@ -1,6 +1,6 @@ /// import { afterEach, beforeEach, describe, expect, test } from "bun:test" -import { mkdtempSync, rmSync } from "node:fs" +import { mkdtempSync, readFileSync, rmSync } from "node:fs" import { tmpdir } from "node:os" import { join } from "node:path" import { acquireLock } from "../../features/claude-tasks/storage" @@ -68,4 +68,22 @@ describe("agent-teams team config store", () => { //#then expect(teamExists("core")).toBe(false) }) + + test("deleteTeamData removes task files before team files", () => { + //#given + const sourceUrl = new URL("./team-config-store.ts", import.meta.url) + const source = readFileSync(sourceUrl, "utf-8") + const deleteFnStart = source.indexOf("export function deleteTeamData") + const deleteFnSlice = deleteFnStart >= 0 ? source.slice(deleteFnStart, deleteFnStart + 700) : "" + + //#when + const taskDeleteIndex = deleteFnSlice.indexOf("rmSync(taskDir") + const teamDeleteIndex = deleteFnSlice.indexOf("rmSync(teamDir") + + //#then + expect(deleteFnStart).toBeGreaterThanOrEqual(0) + expect(taskDeleteIndex).toBeGreaterThanOrEqual(0) + expect(teamDeleteIndex).toBeGreaterThanOrEqual(0) + expect(taskDeleteIndex).toBeLessThan(teamDeleteIndex) + }) }) diff --git a/src/tools/agent-teams/team-config-store.ts b/src/tools/agent-teams/team-config-store.ts index b916f24f1..7c733e539 100644 --- a/src/tools/agent-teams/team-config-store.ts +++ b/src/tools/agent-teams/team-config-store.ts @@ -179,13 +179,13 @@ export function deleteTeamData(teamName: string): void { const teamDir = getTeamDir(teamName) const taskDir = getTeamTaskDir(teamName) - if (existsSync(teamDir)) { - rmSync(teamDir, { recursive: true, force: true }) - } - if (existsSync(taskDir)) { rmSync(taskDir, { recursive: true, force: true }) } + + if (existsSync(teamDir)) { + rmSync(teamDir, { recursive: true, force: true }) + } }) }) } diff --git a/src/tools/agent-teams/team-lifecycle-tools.test.ts b/src/tools/agent-teams/team-lifecycle-tools.test.ts new file mode 100644 index 000000000..cffbb1e69 --- /dev/null +++ b/src/tools/agent-teams/team-lifecycle-tools.test.ts @@ -0,0 +1,69 @@ +/// +import { afterEach, beforeEach, describe, expect, test } from "bun:test" +import { existsSync, mkdtempSync, rmSync } from "node:fs" +import { tmpdir } from "node:os" +import { join } from "node:path" +import type { BackgroundManager } from "../../features/background-agent" +import { getTeamDir } from "./paths" +import { createAgentTeamsTools } from "./tools" + +interface TestToolContext { + sessionID: string + messageID: string + agent: string + abort: AbortSignal +} + +function createContext(sessionID = "ses-main"): TestToolContext { + return { + sessionID, + messageID: "msg-main", + agent: "sisyphus", + abort: new AbortController().signal, + } +} + +async function executeJsonTool( + tools: ReturnType, + toolName: keyof ReturnType, + args: Record, + context: TestToolContext, +): Promise { + const output = await tools[toolName].execute(args, context) + return JSON.parse(output) +} + +describe("agent-teams team lifecycle tools", () => { + let originalCwd: string + let tempProjectDir: string + + beforeEach(() => { + originalCwd = process.cwd() + tempProjectDir = mkdtempSync(join(tmpdir(), "agent-teams-lifecycle-")) + process.chdir(tempProjectDir) + }) + + afterEach(() => { + process.chdir(originalCwd) + rmSync(tempProjectDir, { recursive: true, force: true }) + }) + + test("team_delete requires lead session authorization", async () => { + //#given + const tools = createAgentTeamsTools({} as BackgroundManager) + const leadContext = createContext("ses-main") + await executeJsonTool(tools, "team_create", { team_name: "core" }, leadContext) + + //#when + const unauthorized = await executeJsonTool( + tools, + "team_delete", + { team_name: "core" }, + createContext("ses-intruder"), + ) as { error?: string } + + //#then + expect(unauthorized.error).toBe("unauthorized_lead_session") + expect(existsSync(getTeamDir("core"))).toBe(true) + }) +}) diff --git a/src/tools/agent-teams/team-lifecycle-tools.ts b/src/tools/agent-teams/team-lifecycle-tools.ts index bd252d039..4fb53b2f6 100644 --- a/src/tools/agent-teams/team-lifecycle-tools.ts +++ b/src/tools/agent-teams/team-lifecycle-tools.ts @@ -52,10 +52,13 @@ export function createTeamDeleteTool(): ToolDefinition { args: { team_name: tool.schema.string().describe("Team name"), }, - execute: async (args: Record): Promise => { + execute: async (args: Record, context: TeamToolContext): Promise => { try { const input = TeamDeleteInputSchema.parse(args) const config = readTeamConfigOrThrow(input.team_name) + if (context.sessionID !== config.leadSessionId) { + return JSON.stringify({ error: "unauthorized_lead_session" }) + } const teammates = listTeammates(config) if (teammates.length > 0) { return JSON.stringify({