diff --git a/src/tools/agent-teams/inbox-store.ts b/src/tools/agent-teams/inbox-store.ts index fe8f40a06..d9b4327ed 100644 --- a/src/tools/agent-teams/inbox-store.ts +++ b/src/tools/agent-teams/inbox-store.ts @@ -1,4 +1,4 @@ -import { existsSync, readFileSync } from "node:fs" +import { existsSync, readFileSync, unlinkSync } from "node:fs" import { z } from "zod" import { acquireLock, ensureDir, writeJsonAtomic } from "../../features/claude-tasks/storage" import { getTeamInboxDir, getTeamInboxPath } from "./paths" @@ -93,6 +93,17 @@ export function appendInboxMessage(teamName: string, agentName: string, message: }) } +export function clearInbox(teamName: string, agentName: string): void { + assertValidTeamName(teamName) + assertValidInboxAgentName(agentName) + withInboxLock(teamName, () => { + const path = getTeamInboxPath(teamName, agentName) + if (existsSync(path)) { + unlinkSync(path) + } + }) +} + export function sendPlainInboxMessage( teamName: string, from: string, diff --git a/src/tools/agent-teams/messaging-tools.ts b/src/tools/agent-teams/messaging-tools.ts index aa1b3db68..5d382eed9 100644 --- a/src/tools/agent-teams/messaging-tools.ts +++ b/src/tools/agent-teams/messaging-tools.ts @@ -5,6 +5,7 @@ import { getTeamMember, listTeammates, readTeamConfigOrThrow } from "./team-conf import { validateAgentNameOrLead, validateTeamName } from "./name-validation" import { resumeTeammateWithMessage } from "./teammate-runtime" import { + TeamConfig, TeamReadInboxInputSchema, TeamSendMessageInputSchema, TeamToolContext, @@ -15,6 +16,15 @@ function nowIso(): string { return new Date().toISOString() } +function resolveSenderFromContext(config: TeamConfig, context: TeamToolContext): string | null { + if (context.sessionID === config.leadSessionId) { + return "team-lead" + } + + const matchedMember = config.members.find((member) => isTeammateMember(member) && member.sessionID === context.sessionID) + return matchedMember?.name ?? null +} + export function createSendMessageTool(manager: BackgroundManager): ToolDefinition { return tool({ description: "Send direct or broadcast team messages and protocol responses.", @@ -35,12 +45,21 @@ export function createSendMessageTool(manager: BackgroundManager): ToolDefinitio if (teamError) { return JSON.stringify({ error: teamError }) } - const sender = input.sender ?? "team-lead" - const senderError = validateAgentNameOrLead(sender) + const requestedSender = input.sender + const senderError = requestedSender ? validateAgentNameOrLead(requestedSender) : null if (senderError) { return JSON.stringify({ error: senderError }) } const config = readTeamConfigOrThrow(input.team_name) + const actor = resolveSenderFromContext(config, context) + if (!actor) { + return JSON.stringify({ error: "unauthorized_sender_session" }) + } + if (requestedSender && requestedSender !== actor) { + return JSON.stringify({ error: "sender_context_mismatch" }) + } + const sender = requestedSender ?? actor + const memberNames = new Set(config.members.map((member) => member.name)) if (sender !== "team-lead" && !memberNames.has(sender)) { return JSON.stringify({ error: "invalid_sender" }) @@ -92,12 +111,12 @@ export function createSendMessageTool(manager: BackgroundManager): ToolDefinitio const requestId = buildShutdownRequestId(input.recipient) sendStructuredInboxMessage( input.team_name, - "team-lead", + sender, input.recipient, { type: "shutdown_request", requestId, - from: "team-lead", + from: sender, reason: input.content ?? "", timestamp: nowIso(), }, diff --git a/src/tools/agent-teams/team-config-store.ts b/src/tools/agent-teams/team-config-store.ts index 2243dc7ce..45572966b 100644 --- a/src/tools/agent-teams/team-config-store.ts +++ b/src/tools/agent-teams/team-config-store.ts @@ -127,6 +127,20 @@ export function writeTeamConfig(teamName: string, config: TeamConfig): TeamConfi }) } +export function updateTeamConfig(teamName: string, updater: (config: TeamConfig) => TeamConfig): TeamConfig { + assertValidTeamName(teamName) + return withTeamLock(teamName, () => { + const current = readJsonSafe(getTeamConfigPath(teamName), TeamConfigSchema) + if (!current) { + throw new Error("team_not_found") + } + + const next = TeamConfigSchema.parse(updater(current)) + writeJsonAtomic(getTeamConfigPath(teamName), next) + return next + }) +} + export function listTeammates(config: TeamConfig): TeamTeammateMember[] { return config.members.filter(isTeammateMember) } diff --git a/src/tools/agent-teams/team-task-tools.ts b/src/tools/agent-teams/team-task-tools.ts index d2aff70c8..2febdeed1 100644 --- a/src/tools/agent-teams/team-task-tools.ts +++ b/src/tools/agent-teams/team-task-tools.ts @@ -1,7 +1,7 @@ import { tool, type ToolDefinition } from "@opencode-ai/plugin/tool" import { sendStructuredInboxMessage } from "./inbox-store" import { readTeamConfigOrThrow } from "./team-config-store" -import { validateAgentName, validateTaskId, validateTeamName } from "./name-validation" +import { validateAgentNameOrLead, validateTaskId, validateTeamName } from "./name-validation" import { TeamTaskCreateInputSchema, TeamTaskGetInputSchema, @@ -117,7 +117,7 @@ export function notifyOwnerAssignment(teamName: string, task: TeamTask): void { return } - if (validateAgentName(task.owner)) { + if (validateAgentNameOrLead(task.owner)) { return } diff --git a/src/tools/agent-teams/teammate-runtime.ts b/src/tools/agent-teams/teammate-runtime.ts index 09e0c6c3d..90063c6b3 100644 --- a/src/tools/agent-teams/teammate-runtime.ts +++ b/src/tools/agent-teams/teammate-runtime.ts @@ -1,6 +1,6 @@ import type { BackgroundManager } from "../../features/background-agent" -import { ensureInbox, sendPlainInboxMessage } from "./inbox-store" -import { assignNextColor, getTeamMember, readTeamConfigOrThrow, removeTeammate, upsertTeammate, writeTeamConfig } from "./team-config-store" +import { clearInbox, ensureInbox, sendPlainInboxMessage } from "./inbox-store" +import { assignNextColor, getTeamMember, removeTeammate, updateTeamConfig, upsertTeammate } from "./team-config-store" import type { TeamTeammateMember, TeamToolContext } from "./types" function parseModel(model: string | undefined): { providerID: string; modelID: string } | undefined { @@ -60,29 +60,36 @@ export interface SpawnTeammateParams { } export async function spawnTeammate(params: SpawnTeammateParams): Promise { - const config = readTeamConfigOrThrow(params.teamName) + let teammate: TeamTeammateMember | undefined let launchedTaskID: string | undefined - if (getTeamMember(config, params.name)) { - throw new Error("teammate_already_exists") + updateTeamConfig(params.teamName, (current) => { + if (getTeamMember(current, params.name)) { + throw new Error("teammate_already_exists") + } + + teammate = { + agentId: `${params.name}@${params.teamName}`, + name: params.name, + agentType: params.subagentType, + model: params.model ?? "native", + prompt: params.prompt, + color: assignNextColor(current), + planModeRequired: params.planModeRequired, + joinedAt: Date.now(), + cwd: process.cwd(), + subscriptions: [], + backendType: "native", + isActive: false, + } + + return upsertTeammate(current, teammate) + }) + + if (!teammate) { + throw new Error("teammate_create_failed") } - const teammate: TeamTeammateMember = { - agentId: `${params.name}@${params.teamName}`, - name: params.name, - agentType: params.subagentType, - model: params.model ?? "native", - prompt: params.prompt, - color: assignNextColor(config), - planModeRequired: params.planModeRequired, - joinedAt: Date.now(), - cwd: process.cwd(), - subscriptions: [], - backendType: "native", - isActive: false, - } - - writeTeamConfig(params.teamName, upsertTeammate(config, teammate)) ensureInbox(params.teamName, params.name) sendPlainInboxMessage(params.teamName, "team-lead", params.name, params.prompt, "initial_prompt", teammate.color) @@ -125,8 +132,7 @@ export async function spawnTeammate(params: SpawnTeammateParams): Promise upsertTeammate(current, nextMember)) return nextMember } catch (error) { if (launchedTaskID) { @@ -138,9 +144,8 @@ export async function spawnTeammate(params: SpawnTeammateParams): Promise undefined) } - - const rollback = readTeamConfigOrThrow(params.teamName) - writeTeamConfig(params.teamName, removeTeammate(rollback, params.name)) + updateTeamConfig(params.teamName, (current) => removeTeammate(current, params.name)) + clearInbox(params.teamName, params.name) throw error } } diff --git a/src/tools/agent-teams/teammate-tools.ts b/src/tools/agent-teams/teammate-tools.ts index f5c47f552..154f0e841 100644 --- a/src/tools/agent-teams/teammate-tools.ts +++ b/src/tools/agent-teams/teammate-tools.ts @@ -101,7 +101,7 @@ export function createForceKillTeammateTool(manager: BackgroundManager): ToolDef }) } -export function createProcessShutdownTool(): ToolDefinition { +export function createProcessShutdownTool(manager: BackgroundManager): ToolDefinition { return tool({ description: "Finalize an approved shutdown by removing teammate and resetting owned tasks.", args: { @@ -129,6 +129,7 @@ export function createProcessShutdownTool(): ToolDefinition { return JSON.stringify({ error: "teammate_not_found" }) } + await cancelTeammateRun(manager, member) writeTeamConfig(input.team_name, removeTeammate(config, input.agent_name)) resetOwnerTasks(input.team_name, input.agent_name) diff --git a/src/tools/agent-teams/tools.functional.test.ts b/src/tools/agent-teams/tools.functional.test.ts index 09aec2772..0620f534b 100644 --- a/src/tools/agent-teams/tools.functional.test.ts +++ b/src/tools/agent-teams/tools.functional.test.ts @@ -101,9 +101,9 @@ function createFailingLaunchManager(): { manager: BackgroundManager; cancelCalls return { manager, cancelCalls } } -function createContext(): TestToolContext { +function createContext(sessionID = "ses-main"): TestToolContext { return { - sessionID: "ses-main", + sessionID, messageID: "msg-main", agent: "sisyphus", abort: new AbortController().signal, @@ -359,8 +359,21 @@ describe("agent-teams tools functional", () => { context, ) as { owner?: string } + const leadInbox = await executeJsonTool( + tools, + "read_inbox", + { + team_name: "core", + agent_name: "team-lead", + unread_only: true, + mark_as_read: false, + }, + context, + ) as Array<{ summary?: string; text: string }> + //#then expect(updated.owner).toBe("team-lead") + expect(leadInbox.some((message) => message.summary === "task_assignment")).toBe(true) }) test("spawn_teammate + send_message + force_kill_teammate execute end-to-end", async () => { @@ -431,7 +444,7 @@ describe("agent-teams tools functional", () => { ) as { error?: string } //#then - expect(invalidSender.error).toBe("invalid_sender") + expect(invalidSender.error).toBe("sender_context_mismatch") //#given const createdTask = await executeJsonTool( @@ -499,6 +512,56 @@ describe("agent-teams tools functional", () => { expect(taskAfterKill.status).toBe("pending") }) + test("process_shutdown_approved cancels and removes teammate", async () => { + //#given + const { manager, cancelCalls } = createMockManager() + 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", + }, + leadContext, + ) + + //#when + const shutdownResult = await executeJsonTool( + tools, + "process_shutdown_approved", + { + team_name: "core", + agent_name: "worker_1", + }, + leadContext, + ) as { success: boolean } + + //#then + expect(shutdownResult.success).toBe(true) + expect(cancelCalls).toHaveLength(1) + expect(cancelCalls[0].taskId).toBe("bg-1") + expect(cancelCalls[0].options).toEqual( + expect.objectContaining({ + source: "team_force_kill", + abortSession: true, + skipNotification: true, + }), + ) + + //#when + const configAfterShutdown = await executeJsonTool(tools, "read_config", { team_name: "core" }, leadContext) as { + members: Array<{ name: string }> + } + + //#then + expect(configAfterShutdown.members.some((member) => member.name === "worker_1")).toBe(false) + }) + test("rolls back teammate and cancels background task when launch fails", async () => { //#given const { manager, cancelCalls } = createFailingLaunchManager() @@ -537,6 +600,7 @@ describe("agent-teams tools functional", () => { skipNotification: true, }), ) + expect(existsSync(getTeamInboxPath("core", "worker_1"))).toBe(false) }) test("returns explicit error on invalid model override format", async () => { @@ -592,4 +656,59 @@ describe("agent-teams tools functional", () => { //#then expect(result.error).toBe("team_not_found") }) + + test("binds sender to calling context and rejects sender spoofing", async () => { + //#given + const { manager } = createMockManager() + 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", + }, + leadContext, + ) + + const teammateContext = createContext("ses-worker-1") + + //#when + const spoofed = await executeJsonTool( + tools, + "send_message", + { + team_name: "core", + type: "message", + sender: "team-lead", + recipient: "worker_1", + summary: "spoof", + content: "I am lead", + }, + teammateContext, + ) as { error?: string } + + //#then + expect(spoofed.error).toBe("sender_context_mismatch") + + //#when + const validFromContext = await executeJsonTool( + tools, + "send_message", + { + team_name: "core", + type: "message", + recipient: "team-lead", + summary: "update", + content: "status from worker", + }, + teammateContext, + ) as { success?: boolean } + + //#then + expect(validFromContext.success).toBe(true) + }) }) diff --git a/src/tools/agent-teams/tools.ts b/src/tools/agent-teams/tools.ts index dd4a26e44..904f636f1 100644 --- a/src/tools/agent-teams/tools.ts +++ b/src/tools/agent-teams/tools.ts @@ -19,6 +19,6 @@ export function createAgentTeamsTools(manager: BackgroundManager): Record