From dc3d81a0b81c80def632c7edc5e6fa3cd3ee3dbe Mon Sep 17 00:00:00 2001 From: Nguyen Khac Trung Kien Date: Sun, 8 Feb 2026 09:11:51 +0700 Subject: [PATCH] fix(agent-teams): tighten reviewer-raised runtime and messaging guards Validate sender/owner/team flows more strictly, fail fast on invalid model overrides, and cancel failed launches to prevent orphaned background tasks while expanding functional coverage for these paths. --- src/tools/agent-teams/inbox-store.ts | 2 +- src/tools/agent-teams/messaging-tools.ts | 22 ++ src/tools/agent-teams/team-config-store.ts | 4 +- .../agent-teams/team-task-update-tool.ts | 19 +- src/tools/agent-teams/team-task-update.ts | 2 +- src/tools/agent-teams/teammate-runtime.ts | 14 +- .../agent-teams/tools.functional.test.ts | 191 +++++++++++++++++- 7 files changed, 238 insertions(+), 16 deletions(-) diff --git a/src/tools/agent-teams/inbox-store.ts b/src/tools/agent-teams/inbox-store.ts index 6d9dcb1ef..fe8f40a06 100644 --- a/src/tools/agent-teams/inbox-store.ts +++ b/src/tools/agent-teams/inbox-store.ts @@ -1,4 +1,4 @@ -import { existsSync, readFileSync, writeFileSync } from "node:fs" +import { existsSync, readFileSync } from "node:fs" import { z } from "zod" import { acquireLock, ensureDir, writeJsonAtomic } from "../../features/claude-tasks/storage" import { getTeamInboxDir, getTeamInboxPath } from "./paths" diff --git a/src/tools/agent-teams/messaging-tools.ts b/src/tools/agent-teams/messaging-tools.ts index 6608a0978..aa1b3db68 100644 --- a/src/tools/agent-teams/messaging-tools.ts +++ b/src/tools/agent-teams/messaging-tools.ts @@ -2,6 +2,7 @@ import { tool, type ToolDefinition } from "@opencode-ai/plugin/tool" import type { BackgroundManager } from "../../features/background-agent" import { buildShutdownRequestId, readInbox, sendPlainInboxMessage, sendStructuredInboxMessage } from "./inbox-store" import { getTeamMember, listTeammates, readTeamConfigOrThrow } from "./team-config-store" +import { validateAgentNameOrLead, validateTeamName } from "./name-validation" import { resumeTeammateWithMessage } from "./teammate-runtime" import { TeamReadInboxInputSchema, @@ -30,9 +31,20 @@ export function createSendMessageTool(manager: BackgroundManager): ToolDefinitio execute: async (args: Record, context: TeamToolContext): Promise => { try { const input = TeamSendMessageInputSchema.parse(args) + const teamError = validateTeamName(input.team_name) + if (teamError) { + return JSON.stringify({ error: teamError }) + } const sender = input.sender ?? "team-lead" + const senderError = validateAgentNameOrLead(sender) + if (senderError) { + return JSON.stringify({ error: senderError }) + } const config = readTeamConfigOrThrow(input.team_name) const memberNames = new Set(config.members.map((member) => member.name)) + if (sender !== "team-lead" && !memberNames.has(sender)) { + return JSON.stringify({ error: "invalid_sender" }) + } if (input.type === "message") { if (!input.recipient || !input.summary || !input.content) { @@ -184,6 +196,16 @@ export function createReadInboxTool(): ToolDefinition { execute: async (args: Record): Promise => { try { const input = TeamReadInboxInputSchema.parse(args) + const teamError = validateTeamName(input.team_name) + if (teamError) { + return JSON.stringify({ error: teamError }) + } + const agentError = validateAgentNameOrLead(input.agent_name) + if (agentError) { + return JSON.stringify({ error: agentError }) + } + readTeamConfigOrThrow(input.team_name) + const messages = readInbox( input.team_name, input.agent_name, diff --git a/src/tools/agent-teams/team-config-store.ts b/src/tools/agent-teams/team-config-store.ts index 72e50cf6b..2243dc7ce 100644 --- a/src/tools/agent-teams/team-config-store.ts +++ b/src/tools/agent-teams/team-config-store.ts @@ -51,7 +51,7 @@ function withTeamLock(teamName: string, operation: () => T): T { } } -function createLeadMember(teamName: string, leadSessionId: string, cwd: string, model: string): TeamLeadMember { +function createLeadMember(teamName: string, cwd: string, model: string): TeamLeadMember { return { agentId: `team-lead@${teamName}`, name: "team-lead", @@ -93,7 +93,7 @@ export function createTeamConfig( createdAt: nowMs(), leadAgentId, leadSessionId, - members: [createLeadMember(teamName, leadSessionId, cwd, model)], + members: [createLeadMember(teamName, cwd, model)], } return withTeamLock(teamName, () => { diff --git a/src/tools/agent-teams/team-task-update-tool.ts b/src/tools/agent-teams/team-task-update-tool.ts index 3518b0a49..367dce52b 100644 --- a/src/tools/agent-teams/team-task-update-tool.ts +++ b/src/tools/agent-teams/team-task-update-tool.ts @@ -1,6 +1,6 @@ import { tool, type ToolDefinition } from "@opencode-ai/plugin/tool" import { readTeamConfigOrThrow } from "./team-config-store" -import { validateAgentName, validateTaskId, validateTeamName } from "./name-validation" +import { validateAgentNameOrLead, validateTaskId, validateTeamName } from "./name-validation" import { TeamTaskUpdateInputSchema } from "./types" import { updateTeamTask } from "./team-task-update" import { notifyOwnerAssignment } from "./team-task-tools" @@ -31,10 +31,19 @@ export function createTeamTaskUpdateTool(): ToolDefinition { if (taskIdError) { return JSON.stringify({ error: taskIdError }) } + + const config = readTeamConfigOrThrow(input.team_name) + const memberNames = new Set(config.members.map((member) => member.name)) if (input.owner !== undefined) { - const ownerError = validateAgentName(input.owner) - if (ownerError) { - return JSON.stringify({ error: ownerError }) + if (input.owner !== "") { + const ownerError = validateAgentNameOrLead(input.owner) + if (ownerError) { + return JSON.stringify({ error: ownerError }) + } + + if (!memberNames.has(input.owner)) { + return JSON.stringify({ error: "owner_not_in_team" }) + } } } if (input.add_blocks) { @@ -53,8 +62,6 @@ export function createTeamTaskUpdateTool(): ToolDefinition { } } } - readTeamConfigOrThrow(input.team_name) - const task = updateTeamTask(input.team_name, input.task_id, { status: input.status, owner: input.owner, diff --git a/src/tools/agent-teams/team-task-update.ts b/src/tools/agent-teams/team-task-update.ts index 2298b496f..9916db7a2 100644 --- a/src/tools/agent-teams/team-task-update.ts +++ b/src/tools/agent-teams/team-task-update.ts @@ -134,7 +134,7 @@ export function updateTeamTask(teamName: string, taskId: string, patch: TeamTask nextTask.activeForm = patch.activeForm } if (patch.owner !== undefined) { - nextTask.owner = patch.owner + nextTask.owner = patch.owner === "" ? undefined : patch.owner } const pendingWrites = new Map() diff --git a/src/tools/agent-teams/teammate-runtime.ts b/src/tools/agent-teams/teammate-runtime.ts index 0a531412e..09e0c6c3d 100644 --- a/src/tools/agent-teams/teammate-runtime.ts +++ b/src/tools/agent-teams/teammate-runtime.ts @@ -9,7 +9,7 @@ function parseModel(model: string | undefined): { providerID: string; modelID: s } const [providerID, modelID] = model.split("/", 2) if (!providerID || !modelID) { - return undefined + throw new Error("invalid_model_override_format") } return { providerID, modelID } } @@ -61,6 +61,7 @@ export interface SpawnTeammateParams { export async function spawnTeammate(params: SpawnTeammateParams): Promise { const config = readTeamConfigOrThrow(params.teamName) + let launchedTaskID: string | undefined if (getTeamMember(config, params.name)) { throw new Error("teammate_already_exists") @@ -96,6 +97,7 @@ export async function spawnTeammate(params: SpawnTeammateParams): Promise undefined) + } + const rollback = readTeamConfigOrThrow(params.teamName) writeTeamConfig(params.teamName, removeTeammate(rollback, params.name)) throw error diff --git a/src/tools/agent-teams/tools.functional.test.ts b/src/tools/agent-teams/tools.functional.test.ts index 00fd680b0..09aec2772 100644 --- a/src/tools/agent-teams/tools.functional.test.ts +++ b/src/tools/agent-teams/tools.functional.test.ts @@ -76,8 +76,10 @@ function createMockManager(): MockManagerHandles { return { manager, launchCalls, resumeCalls, cancelCalls } } -function createFailingLaunchManager(): BackgroundManager { - return { +function createFailingLaunchManager(): { manager: BackgroundManager; cancelCalls: CancelCall[] } { + const cancelCalls: CancelCall[] = [] + + const manager = { launch: async () => ({ id: "bg-fail" }), getTask: () => ({ id: "bg-fail", @@ -90,8 +92,13 @@ function createFailingLaunchManager(): BackgroundManager { error: "launch failed", }), resume: async () => ({ id: "resume-unused" }), - cancelTask: async () => true, + cancelTask: async (taskId: string, options?: unknown) => { + cancelCalls.push({ taskId, options }) + return true + }, } as unknown as BackgroundManager + + return { manager, cancelCalls } } function createContext(): TestToolContext { @@ -177,6 +184,16 @@ describe("agent-teams tools functional", () => { const context = createContext() await executeJsonTool(tools, "team_create", { team_name: "core" }, context) + await executeJsonTool( + tools, + "spawn_teammate", + { + team_name: "core", + name: "worker_1", + prompt: "Handle release prep", + }, + context, + ) //#when const createdTask = await executeJsonTool( @@ -241,6 +258,21 @@ describe("agent-teams tools functional", () => { const payload = JSON.parse(assignment!.text) as { type: string; taskId: string } expect(payload.type).toBe("task_assignment") expect(payload.taskId).toBe(createdTask.id) + + //#when + const clearedOwnerTask = await executeJsonTool( + tools, + "team_task_update", + { + team_name: "core", + task_id: createdTask.id, + owner: "", + }, + context, + ) as { owner?: string } + + //#then + expect(clearedOwnerTask.owner).toBeUndefined() }) test("rejects invalid task id input for task_get", async () => { @@ -263,6 +295,74 @@ describe("agent-teams tools functional", () => { expect(result.error).toBe("task_id_invalid") }) + test("requires owner to be a team member when setting task owner", async () => { + //#given + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const context = createContext() + + await executeJsonTool(tools, "team_create", { team_name: "core" }, context) + const createdTask = await executeJsonTool( + tools, + "team_task_create", + { + team_name: "core", + subject: "Investigate bug", + description: "Investigate and report root cause", + }, + context, + ) as { id: string } + + //#when + const result = await executeJsonTool( + tools, + "team_task_update", + { + team_name: "core", + task_id: createdTask.id, + owner: "ghost_user", + }, + context, + ) as { error?: string } + + //#then + expect(result.error).toBe("owner_not_in_team") + }) + + test("allows assigning team-lead as task owner", async () => { + //#given + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const context = createContext() + + await executeJsonTool(tools, "team_create", { team_name: "core" }, context) + const createdTask = await executeJsonTool( + tools, + "team_task_create", + { + team_name: "core", + subject: "Prepare checklist", + description: "Prepare release checklist", + }, + context, + ) as { id: string } + + //#when + const updated = await executeJsonTool( + tools, + "team_task_update", + { + team_name: "core", + task_id: createdTask.id, + owner: "team-lead", + }, + context, + ) as { owner?: string } + + //#then + expect(updated.owner).toBe("team-lead") + }) + test("spawn_teammate + send_message + force_kill_teammate execute end-to-end", async () => { //#given const { manager, launchCalls, resumeCalls, cancelCalls } = createMockManager() @@ -315,6 +415,24 @@ describe("agent-teams tools functional", () => { expect(resumeCalls).toHaveLength(1) expect(resumeCalls[0].sessionId).toBe("ses-worker-1") + //#when + const invalidSender = await executeJsonTool( + tools, + "send_message", + { + team_name: "core", + type: "message", + sender: "ghost_user", + recipient: "worker_1", + summary: "sync", + content: "Please update status.", + }, + context, + ) as { error?: string } + + //#then + expect(invalidSender.error).toBe("invalid_sender") + //#given const createdTask = await executeJsonTool( tools, @@ -381,9 +499,9 @@ describe("agent-teams tools functional", () => { expect(taskAfterKill.status).toBe("pending") }) - test("rolls back teammate when launch fails", async () => { + test("rolls back teammate and cancels background task when launch fails", async () => { //#given - const manager = createFailingLaunchManager() + const { manager, cancelCalls } = createFailingLaunchManager() const tools = createAgentTeamsTools(manager) const context = createContext() await executeJsonTool(tools, "team_create", { team_name: "core" }, context) @@ -410,5 +528,68 @@ describe("agent-teams tools functional", () => { //#then expect(config.members.map((member) => member.name)).toEqual(["team-lead"]) + expect(cancelCalls).toHaveLength(1) + expect(cancelCalls[0].taskId).toBe("bg-fail") + expect(cancelCalls[0].options).toEqual( + expect.objectContaining({ + source: "team_launch_failed", + abortSession: true, + skipNotification: true, + }), + ) + }) + + test("returns explicit error on invalid model override format", async () => { + //#given + const { manager, launchCalls } = createMockManager() + const tools = createAgentTeamsTools(manager) + const context = createContext() + await executeJsonTool(tools, "team_create", { team_name: "core" }, context) + + //#when + const spawnResult = await executeJsonTool( + tools, + "spawn_teammate", + { + team_name: "core", + name: "worker_1", + prompt: "Handle release prep", + model: "invalid-format", + }, + context, + ) as { error?: string } + + //#then + expect(spawnResult.error).toBe("invalid_model_override_format") + expect(launchCalls).toHaveLength(0) + + //#when + const config = await executeJsonTool(tools, "read_config", { team_name: "core" }, context) as { + members: Array<{ name: string }> + } + + //#then + expect(config.members.map((member) => member.name)).toEqual(["team-lead"]) + }) + + test("read_inbox returns team_not_found for unknown team", async () => { + //#given + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const context = createContext() + + //#when + const result = await executeJsonTool( + tools, + "read_inbox", + { + team_name: "missing_team", + agent_name: "team-lead", + }, + context, + ) as { error?: string } + + //#then + expect(result.error).toBe("team_not_found") }) })