From 2a57feb810414a10a184bd26e00930d26c806a9e Mon Sep 17 00:00:00 2001 From: Nguyen Khac Trung Kien Date: Sun, 8 Feb 2026 13:49:59 +0700 Subject: [PATCH] fix(agent-teams): tighten config access and context propagation --- src/tools/agent-teams/messaging-tools.test.ts | 32 ++++++++++ src/tools/agent-teams/messaging-tools.ts | 10 ++- src/tools/agent-teams/team-lifecycle-tools.ts | 37 ++++++++++- .../teammate-parent-context.test.ts | 14 +++++ .../agent-teams/teammate-parent-context.ts | 2 +- .../agent-teams/tools.functional.test.ts | 61 +++++++++++++++++++ src/tools/agent-teams/types.ts | 1 + 7 files changed, 153 insertions(+), 4 deletions(-) create mode 100644 src/tools/agent-teams/teammate-parent-context.test.ts diff --git a/src/tools/agent-teams/messaging-tools.test.ts b/src/tools/agent-teams/messaging-tools.test.ts index 3d4e72fb1..18c2ccd0f 100644 --- a/src/tools/agent-teams/messaging-tools.test.ts +++ b/src/tools/agent-teams/messaging-tools.test.ts @@ -138,6 +138,38 @@ describe("agent-teams messaging tools", () => { expect(resumeCalls).toHaveLength(0) }) + test("send_message rejects recipient with empty team suffix", async () => { + //#given + const { manager, resumeCalls } = createManagerWithImmediateResume() + 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", category: "quick" }, + leadContext, + ) + + //#when + const invalidRecipient = await executeJsonTool( + tools, + "send_message", + { + team_name: "core", + type: "message", + recipient: "worker_1@", + summary: "sync", + content: "Please update status.", + }, + leadContext, + ) as { error?: string } + + //#then + expect(invalidRecipient.error).toBe("recipient_team_invalid") + expect(resumeCalls).toHaveLength(0) + }) + test("broadcast schedules teammate resumes without serial await", async () => { //#given const { manager, resumeCalls, resolveAllResumes } = createManagerWithDeferredResume() diff --git a/src/tools/agent-teams/messaging-tools.ts b/src/tools/agent-teams/messaging-tools.ts index 43add6461..0554f1c4b 100644 --- a/src/tools/agent-teams/messaging-tools.ts +++ b/src/tools/agent-teams/messaging-tools.ts @@ -28,7 +28,10 @@ function validateRecipientTeam(recipient: unknown, teamName: string): string | n } const specifiedTeam = trimmed.slice(atIndex + 1).trim() - if (!specifiedTeam || specifiedTeam === teamName) { + if (!specifiedTeam) { + return "recipient_team_invalid" + } + if (specifiedTeam === teamName) { return null } @@ -55,7 +58,10 @@ export function createSendMessageTool(manager: BackgroundManager): ToolDefinitio summary: tool.schema.string().optional().describe("Short summary"), request_id: tool.schema.string().optional().describe("Protocol request id"), approve: tool.schema.boolean().optional().describe("Approval flag"), - sender: tool.schema.string().optional().describe("Sender name (default: team-lead)"), + sender: tool.schema + .string() + .optional() + .describe("Sender name inferred from calling session; explicit value must match resolved sender."), }, execute: async (args: Record, context: TeamToolContext): Promise => { try { diff --git a/src/tools/agent-teams/team-lifecycle-tools.ts b/src/tools/agent-teams/team-lifecycle-tools.ts index 4fb53b2f6..73adae13f 100644 --- a/src/tools/agent-teams/team-lifecycle-tools.ts +++ b/src/tools/agent-teams/team-lifecycle-tools.ts @@ -3,13 +3,38 @@ import { getTeamConfigPath } from "./paths" import { validateTeamName } from "./name-validation" import { ensureInbox } from "./inbox-store" import { + TeamConfig, TeamCreateInputSchema, TeamDeleteInputSchema, TeamReadConfigInputSchema, TeamToolContext, + isTeammateMember, } from "./types" import { createTeamConfig, deleteTeamData, listTeammates, readTeamConfig, readTeamConfigOrThrow } from "./team-config-store" +function resolveReaderFromContext(config: TeamConfig, context: TeamToolContext): "team-lead" | 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 +} + +function toPublicTeamConfig(config: TeamConfig): { + team_name: string + description: string + lead_agent_id: string + teammates: Array<{ name: string }> +} { + return { + team_name: config.name, + description: config.description, + lead_agent_id: config.leadAgentId, + teammates: listTeammates(config).map((member) => ({ name: member.name })), + } +} + export function createTeamCreateTool(): ToolDefinition { return tool({ description: "Create a team workspace with config, inboxes, and task storage.", @@ -82,13 +107,23 @@ export function createTeamReadConfigTool(): ToolDefinition { args: { team_name: tool.schema.string().describe("Team name"), }, - execute: async (args: Record): Promise => { + execute: async (args: Record, context: TeamToolContext): Promise => { try { const input = TeamReadConfigInputSchema.parse(args) const config = readTeamConfig(input.team_name) if (!config) { return JSON.stringify({ error: "team_not_found" }) } + + const actor = resolveReaderFromContext(config, context) + if (!actor) { + return JSON.stringify({ error: "unauthorized_reader_session" }) + } + + if (actor !== "team-lead") { + return JSON.stringify(toPublicTeamConfig(config)) + } + return JSON.stringify(config) } catch (error) { return JSON.stringify({ error: error instanceof Error ? error.message : "team_read_config_failed" }) diff --git a/src/tools/agent-teams/teammate-parent-context.test.ts b/src/tools/agent-teams/teammate-parent-context.test.ts new file mode 100644 index 000000000..8ad11e286 --- /dev/null +++ b/src/tools/agent-teams/teammate-parent-context.test.ts @@ -0,0 +1,14 @@ +/// +import { describe, expect, test } from "bun:test" +import { readFileSync } from "node:fs" + +describe("agent-teams teammate parent context", () => { + test("forwards incoming abort signal to parent context resolver", () => { + //#given + const sourceUrl = new URL("./teammate-parent-context.ts", import.meta.url) + const source = readFileSync(sourceUrl, "utf-8") + + //#then + expect(source.includes("abort: context.abort ?? new AbortController().signal")).toBe(true) + }) +}) diff --git a/src/tools/agent-teams/teammate-parent-context.ts b/src/tools/agent-teams/teammate-parent-context.ts index caeace6c9..59e5f63e3 100644 --- a/src/tools/agent-teams/teammate-parent-context.ts +++ b/src/tools/agent-teams/teammate-parent-context.ts @@ -8,6 +8,6 @@ export function resolveTeamParentContext(context: TeamToolContext): ParentContex sessionID: context.sessionID, messageID: context.messageID, agent: context.agent ?? "sisyphus", - abort: new AbortController().signal, + abort: context.abort ?? new AbortController().signal, } as ToolContextWithMetadata) } diff --git a/src/tools/agent-teams/tools.functional.test.ts b/src/tools/agent-teams/tools.functional.test.ts index c34dc1ac5..9a1b7d37a 100644 --- a/src/tools/agent-teams/tools.functional.test.ts +++ b/src/tools/agent-teams/tools.functional.test.ts @@ -1032,6 +1032,67 @@ describe("agent-teams tools functional", () => { expect(config.leadSessionId).toBe("ses-main") }) + test("read_config rejects sessions outside the team", async () => { + //#given + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext("ses-main") + await executeJsonTool(tools, "team_create", { team_name: "core" }, leadContext) + + //#when + const result = await executeJsonTool( + tools, + "read_config", + { team_name: "core" }, + createContext("ses-unknown"), + ) as { error?: string } + + //#then + expect(result.error).toBe("unauthorized_reader_session") + }) + + test("read_config returns sanitized config for teammate sessions", async () => { + //#given + const { manager } = createMockManager() + const tools = createAgentTeamsTools(manager) + const leadContext = createContext("ses-main") + await executeJsonTool(tools, "team_create", { team_name: "core" }, leadContext) + await executeJsonTool( + tools, + "spawn_teammate", + { + team_name: "core", + name: "worker_1", + prompt: "Handle release prep", + category: "quick", + }, + leadContext, + ) + + //#when + const teammateView = await executeJsonTool( + tools, + "read_config", + { team_name: "core" }, + createContext("ses-worker-1"), + ) as { + team_name?: string + description?: string + lead_agent_id?: string + teammates?: Array<{ name: string }> + leadSessionId?: string + members?: unknown + } + + //#then + expect(teammateView.team_name).toBe("core") + expect(teammateView.description).toBe("") + expect(teammateView.lead_agent_id).toBe("team-lead@core") + expect(teammateView.teammates).toEqual(expect.arrayContaining([expect.objectContaining({ name: "worker_1" })])) + expect("leadSessionId" in teammateView).toBe(false) + expect("members" in teammateView).toBe(false) + }) + test("rejects unknown session claiming team-lead inbox", async () => { //#given const { manager } = createMockManager() diff --git a/src/tools/agent-teams/types.ts b/src/tools/agent-teams/types.ts index 099dcb856..dad9126d4 100644 --- a/src/tools/agent-teams/types.ts +++ b/src/tools/agent-teams/types.ts @@ -192,6 +192,7 @@ export const TeamProcessShutdownInputSchema = z.object({ export interface TeamToolContext { sessionID: string messageID: string + abort: AbortSignal agent?: string }