fix(team-mode): register spawned sessions early

2026-05-04 12:51:18 +09:00
parent cb819307ff
commit 503df29195
10 changed files with 151 additions and 30 deletions
@@ -580,6 +580,7 @@ export class BackgroundManager {
      return
    }

+    await input.onSessionCreated?.(sessionID)
    this.settlePreStartDescendantReservation(task)
    subagentSessions.add(sessionID)

@@ -0,0 +1,65 @@
+/// <reference types="bun-types" />
+
+import { describe, expect, test } from "bun:test"
+import { tmpdir } from "node:os"
+
+import type { PluginInput } from "@opencode-ai/plugin"
+
+import { BackgroundManager } from "./manager"
+
+async function waitForEvent(events: readonly string[], eventName: string): Promise<void> {
+  const deadlineAt = Date.now() + 1_000
+  while (!events.includes(eventName)) {
+    if (Date.now() > deadlineAt) {
+      throw new Error(`timed out waiting for ${eventName}`)
+    }
+    await new Promise((resolve) => setTimeout(resolve, 10))
+  }
+}
+
+describe("BackgroundManager session created callback", () => {
+  test("fires onSessionCreated before the launch prompt is sent", async () => {
+    //#given
+    const events: string[] = []
+    const client = {
+      session: {
+        get: async ({ path }: { path: { id: string } }) => ({
+          data: { id: path.id, directory: tmpdir() },
+        }),
+        create: async () => {
+          events.push("session.create")
+          return { data: { id: "child-session" } }
+        },
+        promptAsync: async () => {
+          events.push("promptAsync")
+          return { data: {} }
+        },
+      },
+    }
+    const manager = new BackgroundManager({
+      pluginContext: { client, directory: tmpdir() } as PluginInput,
+    })
+
+    //#when
+    await manager.launch({
+      description: "Create child",
+      prompt: "Do work",
+      agent: "general",
+      parentSessionId: "parent-session",
+      parentMessageId: "parent-message",
+      onSessionCreated: (sessionId) => {
+        events.push(`onSessionCreated:${sessionId}`)
+      },
+    })
+    await waitForEvent(events, "promptAsync")
+
+    //#then
+    expect(events).toEqual([
+      "session.create",
+      "onSessionCreated:child-session",
+      "promptAsync",
+    ])
+
+    manager.shutdown()
+  })
+})
@@ -112,6 +112,7 @@ export async function startTask(
  }

  const sessionID = createResult.data.id
+  await input.onSessionCreated?.(sessionID)
  subagentSessions.add(sessionID)

  task.status = "running"
@@ -117,6 +117,7 @@ export interface LaunchInput {
  skillContent?: string
  category?: string
  sessionPermission?: SessionPermissionRule[]
+  onSessionCreated?: (sessionId: string) => void | Promise<void>
 }

 export interface ResumeInput {
@@ -77,7 +77,7 @@ function createManager(launchImpl?: (input: LaunchInput) => Promise<BackgroundTa
  let launchCount = 0
  manager.launch = mock((input: LaunchInput) => launchImpl?.(input) ?? Promise.resolve({
    id: `task-${++launchCount}`,
-    sessionID: `ses_mock_${randomUUID()}`,
+    sessionId: `ses_mock_${randomUUID()}`,
    status: "running",
  } as BackgroundTask))
  manager.getTask = mock(() => undefined)
@@ -301,7 +301,7 @@ describe("team-mode integration", () => {
      maxInFlight = Math.max(maxInFlight, inFlight)
      await new Promise((resolve) => setTimeout(resolve, 10))
      inFlight -= 1
-      return { id: `task-${randomUUID()}`, sessionID: `ses_mock_${randomUUID()}`, status: "running" } as BackgroundTask
+      return { id: `task-${randomUUID()}`, sessionId: `ses_mock_${randomUUID()}`, status: "running" } as BackgroundTask
    })

    // when
@@ -5,11 +5,14 @@ import { access, mkdtemp, readdir, rm } from "node:fs/promises"
 import { tmpdir } from "node:os"
 import path from "node:path"

+import type { PluginInput } from "@opencode-ai/plugin"
+
 import { TeamModeConfigSchema } from "../../../config/schema/team-mode"
 import type { ExecutorContext } from "../../../tools/delegate-task/executor-types"
 import type { BackgroundTask, LaunchInput } from "../../background-agent/types"
 import { BackgroundManager } from "../../background-agent/manager"
 import { loadRuntimeState } from "../team-state-store/store"
+import { clearTeamSessionRegistry, lookupTeamSession } from "../team-session-registry"
 import type { TeamSpec } from "../types"

 const resolveMemberMock = mock(async (member: TeamSpec["members"][number]) => ({
@@ -59,7 +62,7 @@ function createManager(
  launchImpl: (input: LaunchInput) => Promise<BackgroundTask>,
  getTaskImpl: (taskId: string) => BackgroundTask | undefined = () => undefined,
 ): { manager: BackgroundManager; launchMock: ReturnType<typeof mock>; cancelTaskMock: ReturnType<typeof mock> } {
-  const manager = new BackgroundManager({ client: {} as ExecutorContext["client"], directory: baseDir } as ConstructorParameters<typeof BackgroundManager>[0])
+  const manager = new BackgroundManager({ pluginContext: { client: {} as ExecutorContext["client"], directory: baseDir } as PluginInput })
  const launchMock = mock((input: LaunchInput) => launchImpl(input))
  const getTaskMock = mock((taskId: string) => getTaskImpl(taskId))
  const cancelTaskMock = mock(async () => true)
@@ -88,6 +91,7 @@ describe("createTeamRun", () => {

  beforeEach(() => {
    resolveMemberMock.mockClear()
+    clearTeamSessionRegistry()
  })

  afterAll(async () => {
@@ -99,7 +103,7 @@ describe("createTeamRun", () => {
    const baseDir = await mkdtemp(path.join(tmpdir(), "team-runtime-create-"))
    temporaryDirectories.push(baseDir)
    let launchCount = 0
-    const { manager, launchMock } = createManager(baseDir, async () => ({ id: `task-${++launchCount}`, sessionID: `session-${launchCount}`, status: "running" } as BackgroundTask))
+    const { manager, launchMock } = createManager(baseDir, async () => ({ id: `task-${++launchCount}`, sessionId: `session-${launchCount}`, status: "running" } as BackgroundTask))
    const context = createContext(baseDir, manager)

    // when
@@ -113,12 +117,54 @@ describe("createTeamRun", () => {
    expect((launchMock.mock.calls as Array<[LaunchInput]>).every(([input]) => input.suppressTmuxSpawn === true)).toBe(true)
  })

+  test("registers a member session as soon as launch reports the real sessionId", async () => {
+    // given
+    const baseDir = await mkdtemp(path.join(tmpdir(), "team-runtime-session-lineage-"))
+    temporaryDirectories.push(baseDir)
+    const tasks = new Map<string, BackgroundTask>()
+    const { manager } = createManager(
+      baseDir,
+      async (input) => {
+        const task = {
+          id: "task-lineage",
+          status: "pending",
+          parentSessionId: input.parentSessionId,
+          parentMessageId: input.parentMessageId,
+          description: input.description,
+          prompt: input.prompt,
+          agent: input.agent,
+        } satisfies BackgroundTask
+        tasks.set(task.id, task)
+        input.onSessionCreated?.("session-lineage")
+        tasks.set(task.id, { ...task, sessionId: "session-lineage", status: "running" })
+        expect(lookupTeamSession("session-lineage")).toEqual({
+          teamRunId: expect.any(String),
+          memberName: "member-1",
+          role: "lead",
+        })
+        return task
+      },
+      (taskId) => tasks.get(taskId),
+    )
+
+    // when
+    const runtimeState = await createTeamRun(createSpec(1), "lead-session", createContext(baseDir, manager), createConfig(baseDir), manager)
+
+    // then
+    expect(runtimeState.members[0]?.sessionId).toBe("session-lineage")
+    expect(lookupTeamSession("session-lineage")).toEqual({
+      teamRunId: runtimeState.teamRunId,
+      memberName: "member-1",
+      role: "lead",
+    })
+  })
+
  test("persists the resolved subagent_type and model on each spawned runtime member", async () => {
    // given
    const baseDir = await mkdtemp(path.join(tmpdir(), "team-runtime-subagent-type-"))
    temporaryDirectories.push(baseDir)
    let launchCount = 0
-    const { manager } = createManager(baseDir, async () => ({ id: `task-${++launchCount}`, sessionID: `session-${launchCount}`, status: "running" } as BackgroundTask))
+    const { manager } = createManager(baseDir, async () => ({ id: `task-${++launchCount}`, sessionId: `session-${launchCount}`, status: "running" } as BackgroundTask))

    // when
    const runtimeState = await createTeamRun(createSpec(3), "lead-session", createContext(baseDir, manager), createConfig(baseDir), manager)
@@ -141,7 +187,7 @@ describe("createTeamRun", () => {
    temporaryDirectories.push(baseDir)
    const { manager, launchMock } = createManager(baseDir, async () => ({
      id: "task-1",
-      sessionID: "session-1",
+      sessionId: "session-1",
      status: "running",
    } as BackgroundTask))

@@ -169,7 +215,7 @@ describe("createTeamRun", () => {
    const { manager, cancelTaskMock } = createManager(baseDir, async () => {
      launchCount += 1
      if (launchCount === 4) throw new Error("launch-4 failed")
-      return { id: `task-${launchCount}`, sessionID: `session-${launchCount}`, status: "running" } as BackgroundTask
+      return { id: `task-${launchCount}`, sessionId: `session-${launchCount}`, status: "running" } as BackgroundTask
    })

    // when
@@ -194,7 +240,7 @@ describe("createTeamRun", () => {
    const { manager } = createManager(baseDir, async () => {
      launchCount += 1
      if (launchCount === 2) throw new Error("launch-2 failed")
-      return { id: `task-${launchCount}`, sessionID: `session-${launchCount}`, status: "running" } as BackgroundTask
+      return { id: `task-${launchCount}`, sessionId: `session-${launchCount}`, status: "running" } as BackgroundTask
    })
    const spec = createSpec(2, true)

@@ -216,7 +262,7 @@ describe("createTeamRun", () => {
    const baseDir = await mkdtemp(path.join(tmpdir(), "team-runtime-idempotent-"))
    temporaryDirectories.push(baseDir)
    let launchCount = 0
-    const { manager, launchMock } = createManager(baseDir, async () => ({ id: `task-${++launchCount}`, sessionID: `session-${launchCount}`, status: "running" } as BackgroundTask))
+    const { manager, launchMock } = createManager(baseDir, async () => ({ id: `task-${++launchCount}`, sessionId: `session-${launchCount}`, status: "running" } as BackgroundTask))
    const spec = createSpec(2)
    const context = createContext(baseDir, manager)

@@ -242,7 +288,7 @@ describe("createTeamRun", () => {
      maxInFlight = Math.max(maxInFlight, inFlight)
      await new Promise((resolve) => setTimeout(resolve, 10))
      inFlight -= 1
-      return { id: `task-${launchCount}`, sessionID: `session-${launchCount}`, status: "running" } as BackgroundTask
+      return { id: `task-${launchCount}`, sessionId: `session-${launchCount}`, status: "running" } as BackgroundTask
    })

    // when
@@ -259,7 +305,7 @@ describe("createTeamRun", () => {
    let launchCount = 0
    const { manager, launchMock } = createManager(baseDir, async (input) => ({
      id: `task-${++launchCount}`,
-      sessionID: `${input.agent}-session-${launchCount}`,
+      sessionId: `${input.agent}-session-${launchCount}`,
      status: "running",
    } as BackgroundTask))
    const spec: TeamSpec = {
@@ -301,7 +347,7 @@ describe("createTeamRun", () => {
    temporaryDirectories.push(baseDir)
    const { manager } = createManager(baseDir, async (input) => ({
      id: `task-${input.agent}`,
-      sessionID: `${input.agent}-session`,
+      sessionId: `${input.agent}-session`,
      status: "running",
    } as BackgroundTask))
    const spec: TeamSpec = {
@@ -340,7 +386,7 @@ describe("createTeamRun", () => {
    let launchCount = 0
    const { manager, launchMock } = createManager(baseDir, async (input) => ({
      id: `task-${++launchCount}`,
-      sessionID: `${input.agent}-session-${launchCount}`,
+      sessionId: `${input.agent}-session-${launchCount}`,
      status: "running",
    } as BackgroundTask))
    const spec: TeamSpec = {
@@ -82,14 +82,14 @@ async function createMemberWorktree(memberWorktreePath: string, projectRoot: str
 }

 async function waitForTaskSessionId(bgMgr: BackgroundManager, task: BackgroundTask, deadlineAt: number): Promise<string> {
-  let sessionId = task.sessionID
+  let sessionId = task.sessionId
  while (!sessionId) {
    if (Date.now() > deadlineAt) throw new Error(`timed out waiting for child session for task ${task.id}`)
    const updatedTask = bgMgr.getTask(task.id)
    if (updatedTask?.status === "error" || updatedTask?.status === "cancelled" || updatedTask?.status === "interrupt") {
      throw new Error(updatedTask.error ?? `task ${task.id} failed before session creation`)
    }
-    sessionId = updatedTask?.sessionID
+    sessionId = updatedTask?.sessionId
    if (!sessionId) await new Promise((resolve) => setTimeout(resolve, SESSION_ID_POLL_MS))
  }
  return sessionId
@@ -190,8 +190,8 @@ export async function createTeamRun(
            description: `Create team member ${spec.name}/${member.name}`,
            prompt: buildMemberPrompt(spec, member, runtimeState.teamRunId, config, resource.worktreePath),
            agent: resolvedMember.agentToUse,
-            parentSessionID: leadSessionId,
-            parentMessageID: options?.parentMessageID ?? `team-create:${runtimeState.teamRunId}:${member.name}`,
+            parentSessionId: leadSessionId,
+            parentMessageId: options?.parentMessageID ?? `team-create:${runtimeState.teamRunId}:${member.name}`,
            teamRunId: runtimeState.teamRunId,
            suppressTmuxSpawn: true,
            model: resolvedMember.model,
@@ -199,6 +199,13 @@ export async function createTeamRun(
            skillContent: resolvedMember.systemContent,
            category: member.kind === "category" ? member.category : undefined,
            sessionPermission: QUESTION_DENIED_SESSION_PERMISSION,
+            onSessionCreated: (sessionId) => {
+              registerTeamSession(sessionId, {
+                teamRunId: runtimeState.teamRunId,
+                memberName: member.name,
+                role: member.name === spec.leadAgentId ? "lead" : "member",
+              })
+            },
          })
          resource.taskId = task.id
          const sessionId = await waitForTaskSessionId(bgMgr, task, deadlineAt)
@@ -26,18 +26,18 @@ describe("deleteTeam cancels only this team's background tasks", () => {
      "member-b": "shutdown_approved",
    })

-    const getTasksByParentSessionMock = mock((sessionID: string) => {
-      if (sessionID !== "lead-session") return []
+    const getTasksByParentSessionMock = mock((sessionId: string) => {
+      if (sessionId !== "lead-session") return []
      return [
-        { id: "team-task-a", sessionID: "session-a", parentMessageID: `team-create:${fixture.teamRunId}:member-a` },
-        { id: "team-task-b", sessionID: "session-b", parentMessageID: `team-create:${fixture.teamRunId}:member-b` },
+        { id: "team-task-a", sessionId: "session-a", parentMessageId: `team-create:${fixture.teamRunId}:member-a` },
+        { id: "team-task-b", sessionId: "session-b", parentMessageId: `team-create:${fixture.teamRunId}:member-b` },
      ]
    })
    const cancelTaskMock = mock(async () => true)
    const bgMgr = {
      getTasksByParentSession: getTasksByParentSessionMock,
      cancelTask: cancelTaskMock,
-    } as unknown as BackgroundManager
+    } as BackgroundManager

    // when
    await deleteTeam(fixture.teamRunId, fixture.config, undefined, bgMgr)
@@ -62,16 +62,16 @@ describe("deleteTeam cancels only this team's background tasks", () => {
    })

    const getTasksByParentSessionMock = mock(() => [
-      { id: "team-task-a", sessionID: "session-a", parentMessageID: `team-create:${fixture.teamRunId}:member-a` },
-      { id: "delegate-task-x", sessionID: "session-x", parentMessageID: "delegate-task:plan-refactor" },
-      { id: "background-task-y", sessionID: "session-y", parentMessageID: undefined },
-      { id: "team-task-other", sessionID: "session-other", parentMessageID: "team-create:other-team-id:member-a" },
+      { id: "team-task-a", sessionId: "session-a", parentMessageId: `team-create:${fixture.teamRunId}:member-a` },
+      { id: "delegate-task-x", sessionId: "session-x", parentMessageId: "delegate-task:plan-refactor" },
+      { id: "background-task-y", sessionId: "session-y", parentMessageId: undefined },
+      { id: "team-task-other", sessionId: "session-other", parentMessageId: "team-create:other-team-id:member-a" },
    ])
    const cancelTaskMock = mock(async () => true)
    const bgMgr = {
      getTasksByParentSession: getTasksByParentSessionMock,
      cancelTask: cancelTaskMock,
-    } as unknown as BackgroundManager
+    } as BackgroundManager

    // when
    await deleteTeam(fixture.teamRunId, fixture.config, undefined, bgMgr)
@@ -44,7 +44,7 @@ export async function deleteTeam(
  if (bgMgr && runtimeState.leadSessionId) {
    const teamMessageMarkerPrefix = `team-create:${teamRunId}:`
    const teamTasks = bgMgr.getTasksByParentSession(runtimeState.leadSessionId)
-      .filter((task) => task.teamRunId === teamRunId || task.parentMessageID?.startsWith(teamMessageMarkerPrefix))
+      .filter((task) => task.teamRunId === teamRunId || task.parentMessageId?.startsWith(teamMessageMarkerPrefix))
    await Promise.all(teamTasks.map((task) => bgMgr.cancelTask(task.id, {
      source: "team-mode-delete",
      reason: `delete team ${teamRunId}`,
@@ -346,8 +346,8 @@ describe("team-runtime shutdown", () => {
    })
    const bgMgr = {
      getTasksByParentSession: () => [
-        { id: "team-task-a", sessionID: "session-a", parentMessageID: `team-create:${fixture.teamRunId}:member-a` },
-        { id: "team-task-b", sessionID: "session-b", parentMessageID: `team-create:${fixture.teamRunId}:member-b` },
+        { id: "team-task-a", sessionId: "session-a", parentMessageId: `team-create:${fixture.teamRunId}:member-a` },
+        { id: "team-task-b", sessionId: "session-b", parentMessageId: `team-create:${fixture.teamRunId}:member-b` },
      ],
      cancelTask: cancelTaskMock,
    }