diff --git a/src/tools/agent-teams/teammate-runtime.ts b/src/tools/agent-teams/teammate-runtime.ts index 052a58c6..0a531412 100644 --- a/src/tools/agent-teams/teammate-runtime.ts +++ b/src/tools/agent-teams/teammate-runtime.ts @@ -18,6 +18,18 @@ function delay(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)) } +function resolveLaunchFailureMessage(status: string | undefined, error: string | undefined): string { + if (status === "error") { + return error ? `teammate_launch_failed:${error}` : "teammate_launch_failed" + } + + if (status === "cancelled") { + return "teammate_launch_cancelled" + } + + return "teammate_launch_timeout" +} + function buildLaunchPrompt(teamName: string, teammateName: string, userPrompt: string): string { return [ `You are teammate "${teammateName}" in team "${teamName}".`, @@ -87,17 +99,28 @@ export async function spawnTeammate(params: SpawnTeammateParams): Promise): Promise => { try { const input = TeamForceKillInputSchema.parse(args) + const teamError = validateTeamName(input.team_name) + if (teamError) { + return JSON.stringify({ error: teamError }) + } + const agentError = validateAgentName(input.agent_name) + if (agentError) { + return JSON.stringify({ error: agentError }) + } const config = readTeamConfigOrThrow(input.team_name) const member = getTeamMember(config, input.agent_name) if (!member || !isTeammateMember(member)) { @@ -78,7 +86,11 @@ export function createForceKillTeammateTool(manager: BackgroundManager): ToolDef } await cancelTeammateRun(manager, member) - writeTeamConfig(input.team_name, removeTeammate(config, input.agent_name)) + const refreshedConfig = readTeamConfigOrThrow(input.team_name) + const refreshedMember = getTeamMember(refreshedConfig, input.agent_name) + if (refreshedMember && isTeammateMember(refreshedMember)) { + writeTeamConfig(input.team_name, removeTeammate(refreshedConfig, input.agent_name)) + } resetOwnerTasks(input.team_name, input.agent_name) return JSON.stringify({ success: true, message: `${input.agent_name} stopped` }) @@ -99,9 +111,17 @@ export function createProcessShutdownTool(): ToolDefinition { execute: async (args: Record): Promise => { try { const input = TeamProcessShutdownInputSchema.parse(args) + const teamError = validateTeamName(input.team_name) + if (teamError) { + return JSON.stringify({ error: teamError }) + } if (input.agent_name === "team-lead") { return JSON.stringify({ error: "cannot_shutdown_team_lead" }) } + const agentError = validateAgentName(input.agent_name) + if (agentError) { + return JSON.stringify({ error: agentError }) + } const config = readTeamConfigOrThrow(input.team_name) const member = getTeamMember(config, input.agent_name) diff --git a/src/tools/agent-teams/tools.functional.test.ts b/src/tools/agent-teams/tools.functional.test.ts index ab3cbbf9..00fd680b 100644 --- a/src/tools/agent-teams/tools.functional.test.ts +++ b/src/tools/agent-teams/tools.functional.test.ts @@ -76,6 +76,24 @@ function createMockManager(): MockManagerHandles { return { manager, launchCalls, resumeCalls, cancelCalls } } +function createFailingLaunchManager(): BackgroundManager { + return { + launch: async () => ({ id: "bg-fail" }), + getTask: () => ({ + id: "bg-fail", + parentSessionID: "ses-main", + parentMessageID: "msg-main", + description: "failed launch", + prompt: "prompt", + agent: "sisyphus-junior", + status: "error", + error: "launch failed", + }), + resume: async () => ({ id: "resume-unused" }), + cancelTask: async () => true, + } as unknown as BackgroundManager +} + function createContext(): TestToolContext { return { sessionID: "ses-main", @@ -225,6 +243,26 @@ describe("agent-teams tools functional", () => { expect(payload.taskId).toBe(createdTask.id) }) + test("rejects invalid task id input for task_get", async () => { + //#given + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const context = createContext() + + await executeJsonTool(tools, "team_create", { team_name: "core" }, context) + + //#when + const result = await executeJsonTool( + tools, + "team_task_get", + { team_name: "core", task_id: "../../etc/passwd" }, + context, + ) as { error?: string } + + //#then + expect(result.error).toBe("task_id_invalid") + }) + test("spawn_teammate + send_message + force_kill_teammate execute end-to-end", async () => { //#given const { manager, launchCalls, resumeCalls, cancelCalls } = createMockManager() @@ -342,4 +380,35 @@ describe("agent-teams tools functional", () => { expect(taskAfterKill.owner).toBeUndefined() expect(taskAfterKill.status).toBe("pending") }) + + test("rolls back teammate when launch fails", async () => { + //#given + const manager = createFailingLaunchManager() + 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", + }, + context, + ) as { error?: string } + + //#then + expect(spawnResult.error).toBe("teammate_launch_failed:launch failed") + + //#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"]) + }) })