From 503df29195aebb7437985c66ae9ef6ef48e19f7d Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Mon, 4 May 2026 12:51:18 +0900 Subject: [PATCH] fix(team-mode): register spawned sessions early --- src/features/background-agent/manager.ts | 1 + .../session-created-callback.test.ts | 65 ++++++++++++++++++ src/features/background-agent/spawner.ts | 1 + src/features/background-agent/types.ts | 1 + src/features/team-mode/integration.test.ts | 4 +- .../team-mode/team-runtime/create.test.ts | 68 ++++++++++++++++--- src/features/team-mode/team-runtime/create.ts | 15 ++-- .../delete-team-bg-cancel.test.ts | 20 +++--- .../team-mode/team-runtime/delete-team.ts | 2 +- .../team-mode/team-runtime/shutdown.test.ts | 4 +- 10 files changed, 151 insertions(+), 30 deletions(-) create mode 100644 src/features/background-agent/session-created-callback.test.ts diff --git a/src/features/background-agent/manager.ts b/src/features/background-agent/manager.ts index b6a894c19..688ad0ec7 100644 --- a/src/features/background-agent/manager.ts +++ b/src/features/background-agent/manager.ts @@ -580,6 +580,7 @@ export class BackgroundManager { return } + await input.onSessionCreated?.(sessionID) this.settlePreStartDescendantReservation(task) subagentSessions.add(sessionID) diff --git a/src/features/background-agent/session-created-callback.test.ts b/src/features/background-agent/session-created-callback.test.ts new file mode 100644 index 000000000..b08297701 --- /dev/null +++ b/src/features/background-agent/session-created-callback.test.ts @@ -0,0 +1,65 @@ +/// + +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 { + 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() + }) +}) diff --git a/src/features/background-agent/spawner.ts b/src/features/background-agent/spawner.ts index 2cb3edc35..aefe15829 100644 --- a/src/features/background-agent/spawner.ts +++ b/src/features/background-agent/spawner.ts @@ -112,6 +112,7 @@ export async function startTask( } const sessionID = createResult.data.id + await input.onSessionCreated?.(sessionID) subagentSessions.add(sessionID) task.status = "running" diff --git a/src/features/background-agent/types.ts b/src/features/background-agent/types.ts index 0851c179b..d72ab4a5e 100644 --- a/src/features/background-agent/types.ts +++ b/src/features/background-agent/types.ts @@ -117,6 +117,7 @@ export interface LaunchInput { skillContent?: string category?: string sessionPermission?: SessionPermissionRule[] + onSessionCreated?: (sessionId: string) => void | Promise } export interface ResumeInput { diff --git a/src/features/team-mode/integration.test.ts b/src/features/team-mode/integration.test.ts index 2e3b48a11..ffd33f653 100644 --- a/src/features/team-mode/integration.test.ts +++ b/src/features/team-mode/integration.test.ts @@ -77,7 +77,7 @@ function createManager(launchImpl?: (input: LaunchInput) => Promise 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 diff --git a/src/features/team-mode/team-runtime/create.test.ts b/src/features/team-mode/team-runtime/create.test.ts index 4f19c2578..a45d8da2e 100644 --- a/src/features/team-mode/team-runtime/create.test.ts +++ b/src/features/team-mode/team-runtime/create.test.ts @@ -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, getTaskImpl: (taskId: string) => BackgroundTask | undefined = () => undefined, ): { manager: BackgroundManager; launchMock: ReturnType; cancelTaskMock: ReturnType } { - const manager = new BackgroundManager({ client: {} as ExecutorContext["client"], directory: baseDir } as ConstructorParameters[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() + 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 = { diff --git a/src/features/team-mode/team-runtime/create.ts b/src/features/team-mode/team-runtime/create.ts index 590425105..5bfdf9e01 100644 --- a/src/features/team-mode/team-runtime/create.ts +++ b/src/features/team-mode/team-runtime/create.ts @@ -82,14 +82,14 @@ async function createMemberWorktree(memberWorktreePath: string, projectRoot: str } async function waitForTaskSessionId(bgMgr: BackgroundManager, task: BackgroundTask, deadlineAt: number): Promise { - 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) diff --git a/src/features/team-mode/team-runtime/delete-team-bg-cancel.test.ts b/src/features/team-mode/team-runtime/delete-team-bg-cancel.test.ts index d29bd0088..7433dd989 100644 --- a/src/features/team-mode/team-runtime/delete-team-bg-cancel.test.ts +++ b/src/features/team-mode/team-runtime/delete-team-bg-cancel.test.ts @@ -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) diff --git a/src/features/team-mode/team-runtime/delete-team.ts b/src/features/team-mode/team-runtime/delete-team.ts index 385e90a3e..c385d281b 100644 --- a/src/features/team-mode/team-runtime/delete-team.ts +++ b/src/features/team-mode/team-runtime/delete-team.ts @@ -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}`, diff --git a/src/features/team-mode/team-runtime/shutdown.test.ts b/src/features/team-mode/team-runtime/shutdown.test.ts index 14b93f6fe..5975d637e 100644 --- a/src/features/team-mode/team-runtime/shutdown.test.ts +++ b/src/features/team-mode/team-runtime/shutdown.test.ts @@ -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, }