diff --git a/src/native-subagents.ts b/src/native-subagents.ts index bef41528..ef83b2b3 100644 --- a/src/native-subagents.ts +++ b/src/native-subagents.ts @@ -162,9 +162,17 @@ export class NativeSubagentRuntime { const previous = this.children.get(task.taskId); if (previous && previous.terminalState === undefined) return; - const knownParentSessionId = task.toolUseId - ? this.parentByToolUse.get(task.toolUseId) - : undefined; + // A SendMessage resume reuses the finished generation's tool id, whose + // control state is already cleaned up. Keep its parent while that session + // is still live; a finished parent can no longer host the new generation. + const resumedParentSessionId = + previous && + (this.isLiveSession(previous.parentSessionId) + ? previous.parentSessionId + : this.rootSessionId); + const knownParentSessionId = + (task.toolUseId ? this.parentByToolUse.get(task.toolUseId) : undefined) ?? + resumedParentSessionId; const identity = task.toolUseId ? this.identityByToolUse.get(task.toolUseId) : undefined; const child: NativeSubagent = { sessionId: this.nextChildSessionId(task.taskId, previous), @@ -303,6 +311,14 @@ export class NativeSubagentRuntime { this.parentByToolUse.delete(toolUseId); } + private isLiveSession(sessionId: string): boolean { + if (sessionId === this.rootSessionId) return true; + for (const child of this.children.values()) { + if (child.sessionId === sessionId) return child.terminalState === undefined; + } + return false; + } + private nextChildSessionId(taskId: string, previous: NativeSubagent | undefined): string { if (!previous) { this.generationByTaskId.set(taskId, 1); diff --git a/src/tests/native-subagents.test.ts b/src/tests/native-subagents.test.ts index cdd07a09..ab7dc8a3 100644 --- a/src/tests/native-subagents.test.ts +++ b/src/tests/native-subagents.test.ts @@ -322,4 +322,107 @@ describe("NativeSubagentRuntime lifecycle", () => { expect(spawnedIds).toEqual(["worker-1", "worker-1:generation:2"]); expect(terminalIds).toEqual(spawnedIds); }); + + it("announces resumed subagent generation immediately when SendMessage restarts the task (#1158)", async () => { + const published: AcpSessionNotification[] = []; + const runtime = new NativeSubagentRuntime( + true, + "root", + {}, + async (notification) => { + published.push(notification); + }, + { log: () => {} }, + ); + const launch = (toolCallId: string) => + ({ + ...control("tool_call", "pending"), + update: { ...control("tool_call", "pending").update, toolCallId }, + }) as AcpSessionNotification; + + // First generation starts from Agent tool call + await runtime.route(launch("launch-1"), async () => {}); + await runtime.taskStarted( + { + taskId: "worker-1", + toolUseId: "launch-1", + subagentType: "Explore", + description: "First run", + }, + async () => {}, + ); + // Generation 1 finishes, which cleans up parentByToolUse for launch-1 + await runtime.finishTask("worker-1", "completed", async () => {}, "launch-1"); + + expect(published.map(({ update }) => update.sessionUpdate)).toEqual([ + "subagent_spawned", + "subagent_state_update", + ]); + + // SendMessage resumes the subagent with the same toolUseId and taskId without a new Agent control frame + await runtime.taskStarted( + { + taskId: "worker-1", + toolUseId: "launch-1", + subagentType: "Explore", + description: "Resumed run", + }, + async () => {}, + ); + + // Generation 2 must be announced immediately at taskStarted, not held until finishTask + const spawnedIds = published.flatMap(({ update }) => + update.sessionUpdate === "subagent_spawned" ? [update.subagentSessionId] : [], + ); + expect(spawnedIds).toEqual(["worker-1", "worker-1:generation:2"]); + + // Terminal state then ends generation 2 + await runtime.finishTask("worker-1", "completed", async () => {}, "launch-1"); + const terminalIds = published.flatMap(({ update }) => + update.sessionUpdate === "subagent_state_update" ? [update.subagentSessionId] : [], + ); + expect(terminalIds).toEqual(["worker-1", "worker-1:generation:2"]); + }); + + it("announces a resumed nested subagent under a live ancestor when its parent has finished", async () => { + const published: AcpSessionNotification[] = []; + const runtime = new NativeSubagentRuntime( + true, + "root", + {}, + async (notification) => { + published.push(notification); + }, + { log: () => {} }, + ); + const launch = (toolCallId: string, parentToolUseId?: string) => + ({ + ...control("tool_call", "pending", parentToolUseId), + update: { ...control("tool_call", "pending", parentToolUseId).update, toolCallId }, + }) as AcpSessionNotification; + + await runtime.route(launch("launch-a"), async () => {}); + await runtime.taskStarted( + { taskId: "agent-a", toolUseId: "launch-a", subagentType: "Explore" }, + async () => {}, + ); + await runtime.route(launch("launch-b", "launch-a"), async () => {}); + await runtime.taskStarted( + { taskId: "agent-b", toolUseId: "launch-b", subagentType: "Explore" }, + async () => {}, + ); + await runtime.finishTask("agent-b", "completed", async () => {}, "launch-b"); + await runtime.finishTask("agent-a", "completed", async () => {}, "launch-a"); + + // The root resumes B with SendMessage after A, B's original parent, is gone. + await runtime.taskStarted( + { taskId: "agent-b", toolUseId: "launch-b", subagentType: "Explore" }, + async () => {}, + ); + + expect(published.at(-1)).toMatchObject({ + sessionId: "root", + update: { sessionUpdate: "subagent_spawned", subagentSessionId: "agent-b:generation:2" }, + }); + }); });