diff --git a/src/tools/agent-teams/teammate-tools.test.ts b/src/tools/agent-teams/teammate-tools.test.ts new file mode 100644 index 000000000..aced47da9 --- /dev/null +++ b/src/tools/agent-teams/teammate-tools.test.ts @@ -0,0 +1,97 @@ +/// +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 MockManagerHandles { + manager: BackgroundManager + launchCalls: Array> +} + +function createMockManager(): MockManagerHandles { + const launchCalls: Array> = [] + + const manager = { + launch: async (args: Record) => { + launchCalls.push(args) + return { id: `bg-${launchCalls.length}`, sessionID: `ses-worker-${launchCalls.length}` } + }, + getTask: () => undefined, + resume: async () => ({ id: "resume-1" }), + cancelTask: async () => true, + } as unknown as BackgroundManager + + return { manager, launchCalls } +} + +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 teammate tools", () => { + let originalCwd: string + let tempProjectDir: string + + beforeEach(() => { + originalCwd = process.cwd() + tempProjectDir = mkdtempSync(join(tmpdir(), "agent-teams-teammate-tools-")) + process.chdir(tempProjectDir) + }) + + afterEach(() => { + process.chdir(originalCwd) + rmSync(tempProjectDir, { recursive: true, force: true }) + }) + + test("spawn_teammate requires lead session authorization", async () => { + //#given + const { manager, launchCalls } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext("ses-lead") + const teammateContext = createContext("ses-worker") + + await executeJsonTool(tools, "team_create", { team_name: "core" }, leadContext) + + //#when + const unauthorized = await executeJsonTool( + tools, + "spawn_teammate", + { + team_name: "core", + name: "worker_1", + prompt: "Handle release prep", + category: "quick", + }, + teammateContext, + ) as { error?: string } + + //#then + expect(unauthorized.error).toBe("unauthorized_lead_session") + expect(launchCalls).toHaveLength(0) + }) +}) diff --git a/src/tools/agent-teams/teammate-tools.ts b/src/tools/agent-teams/teammate-tools.ts index 8b60f3c66..9d3541af4 100644 --- a/src/tools/agent-teams/teammate-tools.ts +++ b/src/tools/agent-teams/teammate-tools.ts @@ -21,6 +21,42 @@ export interface AgentTeamsSpawnOptions { sisyphusJuniorModel?: string } +async function shutdownTeammateWithCleanup( + manager: BackgroundManager, + context: TeamToolContext, + teamName: string, + agentName: string, +): Promise { + const config = readTeamConfigOrThrow(teamName) + if (context.sessionID !== config.leadSessionId) { + return "unauthorized_lead_session" + } + + const member = getTeamMember(config, agentName) + if (!member || !isTeammateMember(member)) { + return "teammate_not_found" + } + + await cancelTeammateRun(manager, member) + let removed = false + + updateTeamConfig(teamName, (current) => { + const refreshedMember = getTeamMember(current, agentName) + if (!refreshedMember || !isTeammateMember(refreshedMember)) { + return current + } + removed = true + return removeTeammate(current, agentName) + }) + + if (removed) { + clearInbox(teamName, agentName) + } + + resetOwnerTasks(teamName, agentName) + return null +} + export function createSpawnTeammateTool(manager: BackgroundManager, options?: AgentTeamsSpawnOptions): ToolDefinition { return tool({ description: "Spawn a teammate using native internal agent execution.", @@ -54,6 +90,11 @@ export function createSpawnTeammateTool(manager: BackgroundManager, options?: Ag return JSON.stringify({ error: "category_conflicts_with_subagent_type" }) } + const config = readTeamConfigOrThrow(input.team_name) + if (context.sessionID !== config.leadSessionId) { + return JSON.stringify({ error: "unauthorized_lead_session" }) + } + const resolvedSubagentType = input.subagent_type ?? "sisyphus-junior" const teammate = await spawnTeammate({ @@ -107,32 +148,12 @@ export function createForceKillTeammateTool(manager: BackgroundManager): ToolDef if (agentError) { return JSON.stringify({ error: agentError }) } - const config = readTeamConfigOrThrow(input.team_name) - if (context.sessionID !== config.leadSessionId) { - return JSON.stringify({ error: "unauthorized_lead_session" }) - } - const member = getTeamMember(config, input.agent_name) - if (!member || !isTeammateMember(member)) { - return JSON.stringify({ error: "teammate_not_found" }) - } - await cancelTeammateRun(manager, member) - let removed = false - updateTeamConfig(input.team_name, (current) => { - const refreshedMember = getTeamMember(current, input.agent_name) - if (!refreshedMember || !isTeammateMember(refreshedMember)) { - return current - } - removed = true - return removeTeammate(current, input.agent_name) - }) - - if (removed) { - clearInbox(input.team_name, input.agent_name) + const shutdownError = await shutdownTeammateWithCleanup(manager, context, input.team_name, input.agent_name) + if (shutdownError) { + return JSON.stringify({ error: shutdownError }) } - resetOwnerTasks(input.team_name, input.agent_name) - return JSON.stringify({ success: true, message: `${input.agent_name} stopped` }) } catch (error) { return JSON.stringify({ error: error instanceof Error ? error.message : "force_kill_teammate_failed" }) @@ -163,32 +184,10 @@ export function createProcessShutdownTool(manager: BackgroundManager): ToolDefin return JSON.stringify({ error: agentError }) } - const config = readTeamConfigOrThrow(input.team_name) - if (context.sessionID !== config.leadSessionId) { - return JSON.stringify({ error: "unauthorized_lead_session" }) + const shutdownError = await shutdownTeammateWithCleanup(manager, context, input.team_name, input.agent_name) + if (shutdownError) { + return JSON.stringify({ error: shutdownError }) } - const member = getTeamMember(config, input.agent_name) - if (!member || !isTeammateMember(member)) { - return JSON.stringify({ error: "teammate_not_found" }) - } - - await cancelTeammateRun(manager, member) - let removed = false - - updateTeamConfig(input.team_name, (current) => { - const refreshedMember = getTeamMember(current, input.agent_name) - if (!refreshedMember || !isTeammateMember(refreshedMember)) { - return current - } - removed = true - return removeTeammate(current, input.agent_name) - }) - - if (removed) { - clearInbox(input.team_name, input.agent_name) - } - - resetOwnerTasks(input.team_name, input.agent_name) return JSON.stringify({ success: true, message: `${input.agent_name} removed` }) } catch (error) {