fix(agent-teams): enforce lead spawn auth and dedupe shutdown

src/tools/agent-teams/teammate-tools.test.ts

@@ -0,0 +1,97 @@
+/// <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 MockManagerHandles {
+  manager: BackgroundManager
+  launchCalls: Array<Record<string, unknown>>
+}
+
+function createMockManager(): MockManagerHandles {
+  const launchCalls: Array<Record<string, unknown>> = []
+
+  const manager = {
+    launch: async (args: Record<string, unknown>) => {
+      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<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 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)
+  })
+})

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<string | null> {
+  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) {