elias/oh-my-opencode
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(agent-teams): harden lead auth and require teammate categories

Browse Source
This commit is contained in:
Nguyen Khac Trung Kien 2026-02-08 11:10:02 +07:00 committed by YeonGyu-Kim
parent e984ce7493
commit fe05a1f254
9 changed files with 186 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
src/tools/agent-teams/inbox-store.ts
View File

@ -84,10 +84,7 @@ export function appendInboxMessage(teamName: string, agentName: string, message:
assertValidInboxAgentName(agentName)
withInboxLock(teamName, () => {
const path = getTeamInboxPath(teamName, agentName)
if (!existsSync(path)) {
writeJsonAtomic(path, [])
}
const messages = readInboxMessages(teamName, agentName)
const messages = existsSync(path) ? parseInboxFile(readFileSync(path, "utf-8")) : []
messages.push(TeamInboxMessageSchema.parse(message))
writeInboxMessages(teamName, agentName, messages)
})
@ -153,18 +150,27 @@ export function readInbox(
return selected
}
const selectedSet = unreadOnly ? new Set(selected) : null
let changed = false
const updated = messages.map((message) => {
if (!unreadOnly) {
if (!message.read) {
changed = true
}
return { ...message, read: true }
}
if (selected.some((selectedMessage) => selectedMessage.timestamp === message.timestamp && selectedMessage.from === message.from && selectedMessage.text === message.text)) {
if (selectedSet?.has(message)) {
changed = true
return { ...message, read: true }
}
return message
})
writeInboxMessages(teamName, agentName, updated)
if (changed) {
writeInboxMessages(teamName, agentName, updated)
}
return selected
})
}

25
src/tools/agent-teams/messaging-tools.ts
View File

@ -1,7 +1,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, updateTeamConfig } from "./team-config-store"
import { getTeamMember, listTeammates, readTeamConfigOrThrow } from "./team-config-store"
import { validateAgentNameOrLead, validateTeamName } from "./name-validation"
import { resumeTeammateWithMessage } from "./teammate-runtime"
import {
@ -25,13 +25,6 @@ function resolveSenderFromContext(config: TeamConfig, context: TeamToolContext):
return matchedMember?.name ?? null
}
function claimLeadSession(teamName: string, nextLeadSessionId: string): TeamConfig {
return updateTeamConfig(teamName, (current) => ({
...current,
leadSessionId: nextLeadSessionId,
}))
}
export function createSendMessageTool(manager: BackgroundManager): ToolDefinition {
return tool({
description: "Send direct or broadcast team messages and protocol responses.",
@ -57,12 +50,8 @@ export function createSendMessageTool(manager: BackgroundManager): ToolDefinitio
if (senderError) {
return JSON.stringify({ error: senderError })
}
let config = readTeamConfigOrThrow(input.team_name)
let actor = resolveSenderFromContext(config, context)
if (!actor && requestedSender === "team-lead") {
config = claimLeadSession(input.team_name, context.sessionID)
actor = "team-lead"
}
const config = readTeamConfigOrThrow(input.team_name)
const actor = resolveSenderFromContext(config, context)
if (!actor) {
return JSON.stringify({ error: "unauthorized_sender_session" })
}
@ -234,12 +223,8 @@ export function createReadInboxTool(): ToolDefinition {
if (agentError) {
return JSON.stringify({ error: agentError })
}
let config = readTeamConfigOrThrow(input.team_name)
let actor = resolveSenderFromContext(config, context)
if (!actor && input.agent_name === "team-lead") {
config = claimLeadSession(input.team_name, context.sessionID)
actor = "team-lead"
}
const config = readTeamConfigOrThrow(input.team_name)
const actor = resolveSenderFromContext(config, context)
if (!actor) {
return JSON.stringify({ error: "unauthorized_reader_session" })
}

8
src/tools/agent-teams/team-config-store.ts
View File

@ -51,12 +51,12 @@ function withTeamLock<T>(teamName: string, operation: () => T): T {
}
}
function createLeadMember(teamName: string, cwd: string, model: string): TeamLeadMember {
function createLeadMember(teamName: string, cwd: string, leadModel: string): TeamLeadMember {
return {
agentId: `team-lead@${teamName}`,
name: "team-lead",
agentType: "team-lead",
model,
model: leadModel,
joinedAt: nowMs(),
cwd,
subscriptions: [],
@ -82,7 +82,7 @@ export function createTeamConfig(
description: string,
leadSessionId: string,
cwd: string,
model: string,
leadModel: string,
): TeamConfig {
ensureTeamStorageDirs(teamName)
@ -93,7 +93,7 @@ export function createTeamConfig(
createdAt: nowMs(),
leadAgentId,
leadSessionId,
members: [createLeadMember(teamName, cwd, model)],
members: [createLeadMember(teamName, cwd, leadModel)],
}
return withTeamLock(teamName, () => {

2
src/tools/agent-teams/team-lifecycle-tools.ts
View File

@ -30,7 +30,7 @@ export function createTeamCreateTool(): ToolDefinition {
input.description ?? "",
context.sessionID,
process.cwd(),
context.agent ?? "native",
"native/team-lead",
)
ensureInbox(config.name, "team-lead")

13
src/tools/agent-teams/teammate-runtime.ts
View File

@ -70,7 +70,7 @@ export interface SpawnTeammateParams {
teamName: string
name: string
prompt: string
category?: string
category: string
subagentType: string
model?: string
planModeRequired: boolean
@ -100,17 +100,20 @@ async function getSystemDefaultModel(client: PluginInput["client"]): Promise<str
}
async function resolveSpawnExecution(params: SpawnTeammateParams): Promise<SpawnExecution> {
if (!params.category) {
if (params.model) {
const launchModel = parseModel(params.model)
return {
agentType: params.subagentType,
teammateModel: params.model ?? "native",
teammateModel: params.model,
...(launchModel ? { launchModel } : {}),
}
}
if (!params.categoryContext?.client) {
throw new Error("category_requires_client_context")
return {
agentType: params.subagentType,
teammateModel: "native",
}
}
const parentContext = resolveParentContext({
@ -179,7 +182,7 @@ export async function spawnTeammate(params: SpawnTeammateParams): Promise<TeamTe
agentId: `${params.name}@${params.teamName}`,
name: params.name,
agentType: execution.agentType,
...(params.category ? { category: params.category } : {}),
category: params.category,
model: execution.teammateModel,
prompt: params.prompt,
color: assignNextColor(current),

26
src/tools/agent-teams/teammate-tools.ts
View File

@ -28,7 +28,7 @@ export function createSpawnTeammateTool(manager: BackgroundManager, options?: Ag
team_name: tool.schema.string().describe("Team name"),
name: tool.schema.string().describe("Teammate name"),
prompt: tool.schema.string().describe("Initial teammate prompt"),
category: tool.schema.string().optional().describe("Optional category (spawns sisyphus-junior with category model/prompt)") ,
category: tool.schema.string().optional().describe("Required category for teammate metadata and routing"),
subagent_type: tool.schema.string().optional().describe("Agent name to run (default: sisyphus-junior)"),
model: tool.schema.string().optional().describe("Optional model override in provider/model format"),
plan_mode_required: tool.schema.boolean().optional().describe("Enable plan mode flag in teammate metadata"),
@ -46,19 +46,15 @@ export function createSpawnTeammateTool(manager: BackgroundManager, options?: Ag
return JSON.stringify({ error: agentError })
}
if (!input.category || !input.category.trim()) {
return JSON.stringify({ error: "category_required" })
}
if (input.category && input.subagent_type && input.subagent_type !== "sisyphus-junior") {
return JSON.stringify({ error: "category_conflicts_with_subagent_type" })
}
if (input.category && input.model) {
return JSON.stringify({ error: "category_conflicts_with_model_override" })
}
if (input.category && !options?.client) {
return JSON.stringify({ error: "category_requires_client_context" })
}
const resolvedSubagentType = input.category ? "sisyphus-junior" : input.subagent_type ?? "sisyphus-junior"
const resolvedSubagentType = input.subagent_type ?? "sisyphus-junior"
const teammate = await spawnTeammate({
teamName: input.team_name,
@ -100,7 +96,7 @@ export function createForceKillTeammateTool(manager: BackgroundManager): ToolDef
team_name: tool.schema.string().describe("Team name"),
agent_name: tool.schema.string().describe("Teammate name"),
},
execute: async (args: Record<string, unknown>): Promise<string> => {
execute: async (args: Record<string, unknown>, context: TeamToolContext): Promise<string> => {
try {
const input = TeamForceKillInputSchema.parse(args)
const teamError = validateTeamName(input.team_name)
@ -112,6 +108,9 @@ export function createForceKillTeammateTool(manager: BackgroundManager): ToolDef
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" })
@ -149,7 +148,7 @@ export function createProcessShutdownTool(manager: BackgroundManager): ToolDefin
team_name: tool.schema.string().describe("Team name"),
agent_name: tool.schema.string().describe("Teammate name"),
},
execute: async (args: Record<string, unknown>): Promise<string> => {
execute: async (args: Record<string, unknown>, context: TeamToolContext): Promise<string> => {
try {
const input = TeamProcessShutdownInputSchema.parse(args)
const teamError = validateTeamName(input.team_name)
@ -165,6 +164,9 @@ export function createProcessShutdownTool(manager: BackgroundManager): ToolDefin
}
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" })

151
src/tools/agent-teams/tools.functional.test.ts
View File

@ -181,12 +181,13 @@ describe("agent-teams tools functional", () => {
//#when
const config = await executeJsonTool(tools, "read_config", { team_name: "core" }, context) as {
name: string
members: Array<{ name: string }>
members: Array<{ name: string; model?: string }>
}
//#then
expect(config.name).toBe("core")
expect(config.members.map((member) => member.name)).toEqual(["team-lead"])
expect(config.members[0]?.model).toBe("native/team-lead")
//#when
const deleted = await executeJsonTool(tools, "team_delete", { team_name: "core" }, context) as {
@ -213,6 +214,7 @@ describe("agent-teams tools functional", () => {
team_name: "core",
name: "worker_1",
prompt: "Handle release prep",
category: "quick",
},
context,
)
@ -340,6 +342,30 @@ describe("agent-teams tools functional", () => {
expect(teammate?.model).toBe(`${resolvedModel.providerID}/${resolvedModel.modelID}`)
})
test("spawn_teammate requires a category", 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,
"spawn_teammate",
{
team_name: "core",
name: "worker_1",
prompt: "Handle release prep",
},
context,
) as { error?: string }
//#then
expect(result.error).toBe("category_required")
})
test("rejects category with incompatible subagent_type", async () => {
//#given
const { manager } = createMockManager()
@ -483,6 +509,7 @@ describe("agent-teams tools functional", () => {
team_name: "core",
name: "worker_1",
prompt: "Handle release prep",
category: "quick",
},
context,
) as { name: string; session_id: string; task_id: string }
@ -617,6 +644,7 @@ describe("agent-teams tools functional", () => {
team_name: "core",
name: "worker_1",
prompt: "Handle release prep",
category: "quick",
},
leadContext,
)
@ -653,6 +681,78 @@ describe("agent-teams tools functional", () => {
expect(configAfterShutdown.members.some((member) => member.name === "worker_1")).toBe(false)
})
test("force_kill_teammate requires lead session authorization", 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",
category: "quick",
},
leadContext,
)
const teammateContext = createContext("ses-worker-1")
//#when
const unauthorized = await executeJsonTool(
tools,
"force_kill_teammate",
{
team_name: "core",
agent_name: "worker_1",
},
teammateContext,
) as { error?: string }
//#then
expect(unauthorized.error).toBe("unauthorized_lead_session")
})
test("process_shutdown_approved requires lead session authorization", 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",
category: "quick",
},
leadContext,
)
const teammateContext = createContext("ses-worker-1")
//#when
const unauthorized = await executeJsonTool(
tools,
"process_shutdown_approved",
{
team_name: "core",
agent_name: "worker_1",
},
teammateContext,
) as { error?: string }
//#then
expect(unauthorized.error).toBe("unauthorized_lead_session")
})
test("rolls back teammate and cancels background task when launch fails", async () => {
//#given
const { manager, cancelCalls } = createFailingLaunchManager()
@ -668,6 +768,7 @@ describe("agent-teams tools functional", () => {
team_name: "core",
name: "worker_1",
prompt: "Handle release prep",
category: "quick",
},
context,
) as { error?: string }
@ -709,6 +810,7 @@ describe("agent-teams tools functional", () => {
team_name: "core",
name: "worker_1",
prompt: "Handle release prep",
category: "quick",
model: "invalid-format",
},
context,
@ -742,6 +844,7 @@ describe("agent-teams tools functional", () => {
team_name: "core",
name: "worker_1",
prompt: "Handle release prep",
category: "quick",
model: "openai/gpt-5.3-codex/reasoning",
},
context,
@ -792,6 +895,7 @@ describe("agent-teams tools functional", () => {
team_name: "core",
name: "worker_1",
prompt: "Handle release prep",
category: "quick",
},
leadContext,
)
@ -827,14 +931,13 @@ describe("agent-teams tools functional", () => {
expect(Array.isArray(ownInbox)).toBe(true)
})
test("allows lead session to rotate after restart using team-lead identity", async () => {
test("rejects unknown session claiming team-lead identity", async () => {
//#given
const { manager } = createMockManager()
const tools = createAgentTeamsTools(manager)
const originalLead = createContext("ses-main")
await executeJsonTool(tools, "team_create", { team_name: "core" }, originalLead)
const restartedLead = createContext("ses-restarted")
const leadContext = createContext("ses-main")
await executeJsonTool(tools, "team_create", { team_name: "core" }, leadContext)
const unknownContext = createContext("ses-unknown")
//#when
const sendResult = await executeJsonTool(
@ -848,20 +951,43 @@ describe("agent-teams tools functional", () => {
summary: "restart",
content: "Lead session migrated",
},
restartedLead,
unknownContext,
) as { success?: boolean; error?: string }
//#then
expect(sendResult.error).toBeUndefined()
expect(sendResult.success).toBe(true)
expect(sendResult.success).toBeUndefined()
expect(sendResult.error).toBe("unauthorized_sender_session")
//#when
const config = await executeJsonTool(tools, "read_config", { team_name: "core" }, restartedLead) as {
const config = await executeJsonTool(tools, "read_config", { team_name: "core" }, leadContext) as {
leadSessionId: string
}
//#then
expect(config.leadSessionId).toBe("ses-restarted")
expect(config.leadSessionId).toBe("ses-main")
})
test("rejects unknown session claiming team-lead inbox", async () => {
//#given
const { manager } = createMockManager()
const tools = createAgentTeamsTools(manager)
const leadContext = createContext("ses-main")
await executeJsonTool(tools, "team_create", { team_name: "core" }, leadContext)
const unknownContext = createContext("ses-unknown")
//#when
const result = await executeJsonTool(
tools,
"read_inbox",
{
team_name: "core",
agent_name: "team-lead",
},
unknownContext,
) as { error?: string }
//#then
expect(result.error).toBe("unauthorized_reader_session")
})
test("clears old inbox when teammate is removed then re-spawned", async () => {
@ -878,6 +1004,7 @@ describe("agent-teams tools functional", () => {
team_name: "core",
name: "worker_1",
prompt: "First run",
category: "quick",
},
context,
)
@ -913,6 +1040,7 @@ describe("agent-teams tools functional", () => {
team_name: "core",
name: "worker_1",
prompt: "Second run",
category: "quick",
},
context,
)
@ -1004,6 +1132,7 @@ describe("agent-teams tools functional", () => {
team_name: "core",
name: "worker_1",
prompt: "Handle release prep",
category: "quick",
},
leadContext,
)

1
src/tools/agent-teams/types.test.ts
View File

@ -9,6 +9,7 @@ describe("agent-teams types", () => {
agentId: "worker@team",
name: "worker",
agentType: "team-lead",
category: "quick",
model: "native",
prompt: "do work",
color: "blue",

2
src/tools/agent-teams/types.ts
View File

@ -30,7 +30,7 @@ export const TeamTeammateMemberSchema = z.object({
agentType: z.string().refine((value) => value !== "team-lead", {
message: "agent_type_reserved",
}),
category: z.string().optional(),
category: z.string(),
model: z.string(),
prompt: z.string(),
color: z.string(),