fix(agent-teams): harden deletion and messaging safety
This commit is contained in:
parent
0f0ba0f71b
commit
f422cfc7af
179
src/tools/agent-teams/messaging-tools.test.ts
Normal file
179
src/tools/agent-teams/messaging-tools.test.ts
Normal file
@ -0,0 +1,179 @@
|
|||||||
|
/// <reference types="bun-types" />
|
||||||
|
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<typeof createAgentTeamsTools>,
|
||||||
|
toolName: keyof ReturnType<typeof createAgentTeamsTools>,
|
||||||
|
args: Record<string, unknown>,
|
||||||
|
context: TestToolContext,
|
||||||
|
): Promise<unknown> {
|
||||||
|
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")
|
||||||
|
})
|
||||||
|
})
|
||||||
@ -16,6 +16,25 @@ function nowIso(): string {
|
|||||||
return new Date().toISOString()
|
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 {
|
function resolveSenderFromContext(config: TeamConfig, context: TeamToolContext): string | null {
|
||||||
if (context.sessionID === config.leadSessionId) {
|
if (context.sessionID === config.leadSessionId) {
|
||||||
return "team-lead"
|
return "team-lead"
|
||||||
@ -45,6 +64,10 @@ export function createSendMessageTool(manager: BackgroundManager): ToolDefinitio
|
|||||||
if (teamError) {
|
if (teamError) {
|
||||||
return JSON.stringify({ error: 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 requestedSender = input.sender
|
||||||
const senderError = requestedSender ? validateAgentNameOrLead(requestedSender) : null
|
const senderError = requestedSender ? validateAgentNameOrLead(requestedSender) : null
|
||||||
if (senderError) {
|
if (senderError) {
|
||||||
@ -88,11 +111,16 @@ export function createSendMessageTool(manager: BackgroundManager): ToolDefinitio
|
|||||||
if (!input.summary) {
|
if (!input.summary) {
|
||||||
return JSON.stringify({ error: "broadcast_requires_summary" })
|
return JSON.stringify({ error: "broadcast_requires_summary" })
|
||||||
}
|
}
|
||||||
|
const broadcastSummary = input.summary
|
||||||
const teammates = listTeammates(config)
|
const teammates = listTeammates(config)
|
||||||
for (const teammate of teammates) {
|
for (const teammate of teammates) {
|
||||||
sendPlainInboxMessage(input.team_name, sender, teammate.name, input.content ?? "", input.summary)
|
sendPlainInboxMessage(input.team_name, sender, teammate.name, input.content ?? "", broadcastSummary)
|
||||||
await resumeTeammateWithMessage(manager, context, input.team_name, teammate, input.summary, input.content ?? "")
|
|
||||||
}
|
}
|
||||||
|
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}` })
|
return JSON.stringify({ success: true, message: `broadcast_sent:${teammates.length}` })
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -1,6 +1,6 @@
|
|||||||
/// <reference types="bun-types" />
|
/// <reference types="bun-types" />
|
||||||
import { afterEach, beforeEach, describe, expect, test } from "bun:test"
|
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 { tmpdir } from "node:os"
|
||||||
import { join } from "node:path"
|
import { join } from "node:path"
|
||||||
import { acquireLock } from "../../features/claude-tasks/storage"
|
import { acquireLock } from "../../features/claude-tasks/storage"
|
||||||
@ -68,4 +68,22 @@ describe("agent-teams team config store", () => {
|
|||||||
//#then
|
//#then
|
||||||
expect(teamExists("core")).toBe(false)
|
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)
|
||||||
|
})
|
||||||
})
|
})
|
||||||
|
|||||||
@ -179,13 +179,13 @@ export function deleteTeamData(teamName: string): void {
|
|||||||
const teamDir = getTeamDir(teamName)
|
const teamDir = getTeamDir(teamName)
|
||||||
const taskDir = getTeamTaskDir(teamName)
|
const taskDir = getTeamTaskDir(teamName)
|
||||||
|
|
||||||
if (existsSync(teamDir)) {
|
|
||||||
rmSync(teamDir, { recursive: true, force: true })
|
|
||||||
}
|
|
||||||
|
|
||||||
if (existsSync(taskDir)) {
|
if (existsSync(taskDir)) {
|
||||||
rmSync(taskDir, { recursive: true, force: true })
|
rmSync(taskDir, { recursive: true, force: true })
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (existsSync(teamDir)) {
|
||||||
|
rmSync(teamDir, { recursive: true, force: true })
|
||||||
|
}
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
69
src/tools/agent-teams/team-lifecycle-tools.test.ts
Normal file
69
src/tools/agent-teams/team-lifecycle-tools.test.ts
Normal file
@ -0,0 +1,69 @@
|
|||||||
|
/// <reference types="bun-types" />
|
||||||
|
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<typeof createAgentTeamsTools>,
|
||||||
|
toolName: keyof ReturnType<typeof createAgentTeamsTools>,
|
||||||
|
args: Record<string, unknown>,
|
||||||
|
context: TestToolContext,
|
||||||
|
): Promise<unknown> {
|
||||||
|
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)
|
||||||
|
})
|
||||||
|
})
|
||||||
@ -52,10 +52,13 @@ export function createTeamDeleteTool(): ToolDefinition {
|
|||||||
args: {
|
args: {
|
||||||
team_name: tool.schema.string().describe("Team name"),
|
team_name: tool.schema.string().describe("Team name"),
|
||||||
},
|
},
|
||||||
execute: async (args: Record<string, unknown>): Promise<string> => {
|
execute: async (args: Record<string, unknown>, context: TeamToolContext): Promise<string> => {
|
||||||
try {
|
try {
|
||||||
const input = TeamDeleteInputSchema.parse(args)
|
const input = TeamDeleteInputSchema.parse(args)
|
||||||
const config = readTeamConfigOrThrow(input.team_name)
|
const config = readTeamConfigOrThrow(input.team_name)
|
||||||
|
if (context.sessionID !== config.leadSessionId) {
|
||||||
|
return JSON.stringify({ error: "unauthorized_lead_session" })
|
||||||
|
}
|
||||||
const teammates = listTeammates(config)
|
const teammates = listTeammates(config)
|
||||||
if (teammates.length > 0) {
|
if (teammates.length > 0) {
|
||||||
return JSON.stringify({
|
return JSON.stringify({
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user