feat(hooks): add team tool gating hook with tests

src/hooks/team-tool-gating/hook.test.ts

@@ -0,0 +1,276 @@
+import { afterEach, beforeEach, describe, expect, test } from "bun:test"
+import { mkdir, mkdtemp, rm } from "node:fs/promises"
+import { tmpdir } from "node:os"
+import path from "node:path"
+
+import type { PluginInput } from "@opencode-ai/plugin"
+import type { TeamModeConfig } from "../../config/schema/team-mode"
+import { TeamModeConfigSchema } from "../../config/schema/team-mode"
+import {
+  clearTeamSessionRegistry,
+  registerTeamSession,
+} from "../../features/team-mode/team-session-registry"
+import type { RuntimeState } from "../../features/team-mode/types"
+import { saveRuntimeState } from "../../features/team-mode/team-state-store/store"
+import { createTeamToolGating } from "./hook"
+
+function createConfig(overrides?: Partial<TeamModeConfig>, baseDir = "/tmp/team-mode"): TeamModeConfig {
+  return {
+    enabled: true,
+    tmux_visualization: false,
+    max_parallel_members: 4,
+    max_members: 8,
+    max_messages_per_run: 10_000,
+    max_wall_clock_minutes: 120,
+    max_member_turns: 500,
+    base_dir: baseDir,
+    message_payload_max_bytes: 32_768,
+    recipient_unread_max_bytes: 262_144,
+    mailbox_poll_interval_ms: 3_000,
+    ...overrides,
+  }
+}
+
+function createRuntimeState(): RuntimeState {
+  return {
+    version: 1,
+    teamRunId: "11111111-1111-4111-8111-111111111111",
+    teamName: "team-alpha",
+    specSource: "project",
+    createdAt: 1,
+    status: "active",
+    leadSessionId: "lead-session",
+    members: [
+      { name: "m1", sessionId: "member-session-1", agentType: "general-purpose", status: "running", pendingInjectedMessageIds: [] },
+      { name: "m2", sessionId: "member-session-2", agentType: "general-purpose", status: "running", pendingInjectedMessageIds: [] },
+    ],
+    shutdownRequests: [],
+    bounds: { maxMembers: 8, maxParallelMembers: 4, maxMessagesPerRun: 10_000, maxWallClockMinutes: 120, maxMemberTurns: 500 },
+  }
+}
+
+async function seedTeams(baseDir: string, ...runtimeStates: RuntimeState[]): Promise<void> {
+  const config = TeamModeConfigSchema.parse({ base_dir: baseDir, enabled: true })
+  await Promise.all(runtimeStates.map(async (runtimeState) => {
+    await mkdir(path.join(baseDir, "runtime", runtimeState.teamRunId), { recursive: true })
+    await saveRuntimeState(runtimeState, config)
+  }))
+}
+
+async function runHook(tool: string, sessionID: string, args: Record<string, unknown>, config?: Partial<TeamModeConfig>, baseDir = "/tmp/team-mode"): Promise<void> {
+  const hook = createTeamToolGating({ directory: baseDir } as PluginInput, createConfig(config, baseDir))
+  await hook["tool.execute.before"]?.({ tool, sessionID, callID: "call-1" }, { args })
+}
+
+describe("createTeamToolGating", () => {
+  const temporaryDirectories: string[] = []
+
+  beforeEach(() => {
+    temporaryDirectories.length = 0
+    clearTeamSessionRegistry()
+  })
+
+  afterEach(async () => {
+    clearTeamSessionRegistry()
+    await Promise.all(temporaryDirectories.splice(0).map(async (directoryPath) => rm(directoryPath, { recursive: true, force: true })))
+  })
+
+  test("allows a fresh session to call team_create", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    await seedTeams(baseDir, createRuntimeState())
+
+    // when
+    const result = runHook("team_create", "fresh-session", {}, undefined, baseDir)
+
+    // then
+    await expect(result).resolves.toBeUndefined()
+  })
+
+  test("allows team_list from a fresh session", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    await seedTeams(baseDir, createRuntimeState())
+
+    // when
+    const result = runHook("team_list", "fresh-session", {}, undefined, baseDir)
+
+    // then
+    await expect(result).resolves.toBeUndefined()
+  })
+
+  test("rejects team_create when the caller is already a team member", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    await seedTeams(baseDir, createRuntimeState())
+
+    // when
+    const result = runHook("team_create", "member-session-1", {}, undefined, baseDir)
+
+    // then
+    await expect(result).rejects.toThrow("team_create denied: session is already a participant of team 11111111-1111-4111-8111-111111111111")
+  })
+
+  test("allows the target member to self-approve shutdown", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    await seedTeams(baseDir, createRuntimeState())
+
+    // when
+    const result = runHook("team_approve_shutdown", "member-session-1", { teamRunId: "11111111-1111-4111-8111-111111111111", memberName: "m1" }, undefined, baseDir)
+
+    // then
+    await expect(result).resolves.toBeUndefined()
+  })
+
+  test("allows the lead to force-approve shutdown", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    await seedTeams(baseDir, createRuntimeState())
+
+    // when
+    const result = runHook("team_approve_shutdown", "lead-session", { teamRunId: "11111111-1111-4111-8111-111111111111", memberName: "m1" }, undefined, baseDir)
+
+    // then
+    await expect(result).resolves.toBeUndefined()
+  })
+
+  test("rejects a non-target member from approving shutdown", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    await seedTeams(baseDir, createRuntimeState())
+
+    // when
+    const result = runHook("team_approve_shutdown", "member-session-2", { teamRunId: "11111111-1111-4111-8111-111111111111", memberName: "m1" }, undefined, baseDir)
+
+    // then
+    await expect(result).rejects.toThrow("team_approve_shutdown: caller must be target member or team lead")
+  })
+
+  test("allows delegate-task for team members without a run-wide budget", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    await seedTeams(baseDir, createRuntimeState())
+
+    // when
+    const result = runHook("delegate-task", "member-session-1", {}, undefined, baseDir)
+
+    // then
+    await expect(result).resolves.toBeUndefined()
+  })
+
+  test("allows team_delete for the lead of the target team", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    await seedTeams(baseDir, createRuntimeState())
+
+    // when
+    const result = runHook("team_delete", "lead-session", { teamRunId: "11111111-1111-4111-8111-111111111111" }, undefined, baseDir)
+
+    // then
+    await expect(result).resolves.toBeUndefined()
+  })
+
+  test("no-ops for unrelated tools without querying team state", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    await seedTeams(baseDir, createRuntimeState())
+
+    // when
+    const result = runHook("write", "fresh-session", {}, undefined, baseDir)
+
+    // then
+    await expect(result).resolves.toBeUndefined()
+  })
+
+  test("allows team_send_message during the spawn race when runtime state lacks the member's sessionId but the registry already has it", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    const staleRuntimeState: RuntimeState = {
+      ...createRuntimeState(),
+      members: [
+        { name: "m1", agentType: "general-purpose", status: "pending", pendingInjectedMessageIds: [] },
+        { name: "m2", agentType: "general-purpose", status: "pending", pendingInjectedMessageIds: [] },
+      ],
+    }
+    await seedTeams(baseDir, staleRuntimeState)
+    registerTeamSession("just-spawned-session", {
+      teamRunId: "11111111-1111-4111-8111-111111111111",
+      memberName: "m1",
+      role: "member",
+    })
+
+    // when
+    const result = runHook("team_send_message", "just-spawned-session", { teamRunId: "11111111-1111-4111-8111-111111111111" }, undefined, baseDir)
+
+    // then
+    await expect(result).resolves.toBeUndefined()
+  })
+
+  test("allows team_send_message from a lead whose session is tracked only in the registry", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    const staleRuntimeState: RuntimeState = {
+      ...createRuntimeState(),
+      leadSessionId: undefined,
+      members: [
+        { name: "lead", agentType: "leader", status: "pending", pendingInjectedMessageIds: [] },
+      ],
+    }
+    await seedTeams(baseDir, staleRuntimeState)
+    registerTeamSession("caller-lead-session", {
+      teamRunId: "11111111-1111-4111-8111-111111111111",
+      memberName: "lead",
+      role: "lead",
+    })
+
+    // when
+    const result = runHook("team_send_message", "caller-lead-session", { teamRunId: "11111111-1111-4111-8111-111111111111" }, undefined, baseDir)
+
+    // then
+    await expect(result).resolves.toBeUndefined()
+  })
+
+  test("rejects team_send_message when the session is not in the registry and not in runtime state", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    await seedTeams(baseDir, createRuntimeState())
+
+    // when
+    const result = runHook("team_send_message", "unknown-session", { teamRunId: "11111111-1111-4111-8111-111111111111" }, undefined, baseDir)
+
+    // then
+    await expect(result).rejects.toThrow("team-mode tool team_send_message denied: not a participant of team 11111111-1111-4111-8111-111111111111")
+  })
+
+  test("rejects team_send_message when the registry only has the caller for a different team than the requested teamRunId", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-tool-gating-"))
+    temporaryDirectories.push(baseDir)
+    const emptyState: RuntimeState = { ...createRuntimeState(), members: [] }
+    await seedTeams(baseDir, emptyState)
+    registerTeamSession("cross-team-session", {
+      teamRunId: "22222222-2222-4222-8222-222222222222",
+      memberName: "other-team-member",
+      role: "member",
+    })
+
+    // when
+    const result = runHook("team_send_message", "cross-team-session", { teamRunId: "11111111-1111-4111-8111-111111111111" }, undefined, baseDir)
+
+    // then
+    await expect(result).rejects.toThrow("denied: not a participant of team 11111111-1111-4111-8111-111111111111")
+  })
+})

src/hooks/team-tool-gating/hook.ts

@@ -0,0 +1,149 @@
+import type { Hooks, PluginInput } from "@opencode-ai/plugin"
+
+import type { TeamModeConfig } from "../../config/schema/team-mode"
+import { lookupTeamSession } from "../../features/team-mode/team-session-registry"
+import type { RuntimeState } from "../../features/team-mode/types"
+import {
+  listActiveTeams,
+  loadRuntimeState,
+} from "../../features/team-mode/team-state-store"
+
+const ACTIVE_RUNTIME_STATUSES = new Set<RuntimeState["status"]>(["creating", "active", "shutdown_requested"])
+const UNIVERSAL_TOOL_NAMES = new Set([
+  "team_send_message",
+  "team_task_create",
+  "team_task_list",
+  "team_task_update",
+  "team_task_get",
+  "team_status",
+])
+
+type TeamParticipant =
+  | { role: "neither" }
+  | { role: "lead"; teamRunId: string }
+  | { role: "member"; teamRunId: string; memberName: string }
+
+function getStringArg(args: Record<string, unknown>, key: string): string | undefined {
+  const value = args[key]
+  return typeof value === "string" ? value : undefined
+}
+
+function resolveParticipantFromRegistry(sessionID: string): TeamParticipant | undefined {
+  const entry = lookupTeamSession(sessionID)
+  if (!entry) return undefined
+  if (entry.role === "lead") {
+    return { role: "lead", teamRunId: entry.teamRunId }
+  }
+  return { role: "member", teamRunId: entry.teamRunId, memberName: entry.memberName }
+}
+
+async function resolveParticipant(sessionID: string, config: TeamModeConfig): Promise<TeamParticipant> {
+  const fromRegistry = resolveParticipantFromRegistry(sessionID)
+  if (fromRegistry) {
+    return fromRegistry
+  }
+
+  const activeTeams = await listActiveTeams(config)
+
+  for (const activeTeam of activeTeams) {
+    const runtimeState = await loadRuntimeState(activeTeam.teamRunId, config)
+    if (!ACTIVE_RUNTIME_STATUSES.has(runtimeState.status)) {
+      continue
+    }
+
+    if (runtimeState.leadSessionId === sessionID) {
+      return { role: "lead", teamRunId: runtimeState.teamRunId }
+    }
+
+    const matchedMember = runtimeState.members.find((member) => member.sessionId === sessionID)
+    if (matchedMember) {
+      return {
+        role: "member",
+        teamRunId: runtimeState.teamRunId,
+        memberName: matchedMember.name,
+      }
+    }
+  }
+
+  return { role: "neither" }
+}
+
+function isLeadOfTargetTeam(participant: TeamParticipant, teamRunId: string | undefined): boolean {
+  return participant.role === "lead" && participant.teamRunId === teamRunId
+}
+
+function isTargetMember(participant: TeamParticipant, teamRunId: string | undefined, memberName: string | undefined): boolean {
+  return participant.role === "member"
+    && participant.teamRunId === teamRunId
+    && participant.memberName === memberName
+}
+
+export function createTeamToolGating(_ctx: PluginInput, config: TeamModeConfig | undefined): Hooks {
+  return {
+    "tool.execute.before": async (
+      input: { tool: string; sessionID: string; callID: string },
+      output: { args: Record<string, unknown> },
+    ): Promise<void> => {
+      if (!config?.enabled) {
+        return
+      }
+
+      const toolName = input.tool
+      if (!toolName.startsWith("team_") && toolName !== "delegate-task") {
+        return
+      }
+
+      const participant = await resolveParticipant(input.sessionID, config)
+
+      if (toolName === "delegate-task") {
+        return
+      }
+
+      if (toolName === "team_create") {
+        if (participant.role !== "neither") {
+          throw new Error(`team_create denied: session is already a participant of team ${participant.teamRunId}`)
+        }
+
+        return
+      }
+
+      const teamRunId = getStringArg(output.args, "teamRunId")
+      const memberName = getStringArg(output.args, "memberName")
+
+      if (toolName === "team_delete" || toolName === "team_shutdown_request") {
+        if (!isLeadOfTargetTeam(participant, teamRunId)) {
+          throw new Error(`${toolName} is lead-only`)
+        }
+
+        return
+      }
+
+      if (toolName === "team_approve_shutdown" || toolName === "team_reject_shutdown") {
+        if (!isLeadOfTargetTeam(participant, teamRunId) && !isTargetMember(participant, teamRunId, memberName)) {
+          throw new Error(`${toolName}: caller must be target member or team lead`)
+        }
+
+        return
+      }
+
+      if (toolName === "team_list") {
+        return
+      }
+
+      if (UNIVERSAL_TOOL_NAMES.has(toolName)) {
+        if (
+          (participant.role === "lead" || participant.role === "member")
+          && participant.teamRunId === teamRunId
+        ) {
+          return
+        }
+
+        throw new Error(
+          teamRunId === undefined
+            ? `team-mode tool ${toolName} requires teamRunId argument`
+            : `team-mode tool ${toolName} denied: not a participant of team ${teamRunId}`,
+        )
+      }
+    },
+  }
+}

src/hooks/team-tool-gating/index.ts

@@ -0,0 +1 @@
+export { createTeamToolGating } from "./hook"