Merge pull request #2725 from cphoward/fix/spawn-budget-lifetime-semantics-clean

fix(background-agent): decrement spawn budget on task completion, cancellation, error, and interrupt

src/features/background-agent/manager.test.ts

@@ -2484,6 +2484,133 @@ describe("BackgroundManager - Non-blocking Queue Integration", () => {
       expect(abortCalls).toEqual([createdSessionID])
       expect(getConcurrencyManager(manager).getCount("test-agent")).toBe(0)
     })
+
+    test("should release descendant quota when task completes", async () => {
+      manager.shutdown()
+      manager = new BackgroundManager(
+        {
+          client: createMockClientWithSessionChain({
+            "session-root": { directory: "/test/dir" },
+          }),
+          directory: tmpdir(),
+        } as unknown as PluginInput,
+        { maxDescendants: 1 },
+      )
+      stubNotifyParentSession(manager)
+
+      const input = {
+        description: "Test task",
+        prompt: "Do something",
+        agent: "test-agent",
+        parentSessionID: "session-root",
+        parentMessageID: "parent-message",
+      }
+
+      const task = await manager.launch(input)
+      const internalTask = getTaskMap(manager).get(task.id)!
+      internalTask.status = "running"
+      internalTask.sessionID = "child-session-complete"
+      internalTask.rootSessionID = "session-root"
+
+      // Complete via internal method (session.status events go through the poller, not handleEvent)
+      await tryCompleteTaskForTest(manager, internalTask)
+
+      await expect(manager.launch(input)).resolves.toBeDefined()
+    })
+
+    test("should release descendant quota when running task is cancelled", async () => {
+      manager.shutdown()
+      manager = new BackgroundManager(
+        {
+          client: createMockClientWithSessionChain({
+            "session-root": { directory: "/test/dir" },
+          }),
+          directory: tmpdir(),
+        } as unknown as PluginInput,
+        { maxDescendants: 1 },
+      )
+
+      const input = {
+        description: "Test task",
+        prompt: "Do something",
+        agent: "test-agent",
+        parentSessionID: "session-root",
+        parentMessageID: "parent-message",
+      }
+
+      const task = await manager.launch(input)
+      const internalTask = getTaskMap(manager).get(task.id)!
+      internalTask.status = "running"
+      internalTask.sessionID = "child-session-cancel"
+
+      await manager.cancelTask(task.id)
+
+      await expect(manager.launch(input)).resolves.toBeDefined()
+    })
+
+    test("should release descendant quota when task errors", async () => {
+      manager.shutdown()
+      manager = new BackgroundManager(
+        {
+          client: createMockClientWithSessionChain({
+            "session-root": { directory: "/test/dir" },
+          }),
+          directory: tmpdir(),
+        } as unknown as PluginInput,
+        { maxDescendants: 1 },
+      )
+
+      const input = {
+        description: "Test task",
+        prompt: "Do something",
+        agent: "test-agent",
+        parentSessionID: "session-root",
+        parentMessageID: "parent-message",
+      }
+
+      const task = await manager.launch(input)
+      const internalTask = getTaskMap(manager).get(task.id)!
+      internalTask.status = "running"
+      internalTask.sessionID = "child-session-error"
+
+      manager.handleEvent({
+        type: "session.error",
+        properties: { sessionID: internalTask.sessionID, info: { id: internalTask.sessionID } },
+      })
+      await new Promise((resolve) => setTimeout(resolve, 100))
+
+      await expect(manager.launch(input)).resolves.toBeDefined()
+    })
+
+    test("should not double-decrement quota when pending task is cancelled", async () => {
+      manager.shutdown()
+      manager = new BackgroundManager(
+        {
+          client: createMockClientWithSessionChain({
+            "session-root": { directory: "/test/dir" },
+          }),
+          directory: tmpdir(),
+        } as unknown as PluginInput,
+        { maxDescendants: 2 },
+      )
+
+      const input = {
+        description: "Test task",
+        prompt: "Do something",
+        agent: "test-agent",
+        parentSessionID: "session-root",
+        parentMessageID: "parent-message",
+      }
+
+      const task1 = await manager.launch(input)
+      const task2 = await manager.launch(input)
+
+      await manager.cancelTask(task1.id)
+      await manager.cancelTask(task2.id)
+
+      await expect(manager.launch(input)).resolves.toBeDefined()
+      await expect(manager.launch(input)).resolves.toBeDefined()
+    })
   })
 
   describe("pending task can be cancelled", () => {

src/features/background-agent/manager.ts

@@ -550,6 +550,9 @@ export class BackgroundManager {
           existingTask.error = errorMessage
         }
         existingTask.completedAt = new Date()
+        if (existingTask.rootSessionID) {
+          this.unregisterRootDescendant(existingTask.rootSessionID)
+        }
         if (existingTask.concurrencyKey) {
           this.concurrencyManager.release(existingTask.concurrencyKey)
           existingTask.concurrencyKey = undefined
@@ -826,6 +829,9 @@ export class BackgroundManager {
       const errorMessage = error instanceof Error ? error.message : String(error)
       existingTask.error = errorMessage
       existingTask.completedAt = new Date()
+      if (existingTask.rootSessionID) {
+        this.unregisterRootDescendant(existingTask.rootSessionID)
+      }
 
       // Release concurrency on error to prevent slot leaks
       if (existingTask.concurrencyKey) {
@@ -1022,6 +1028,9 @@ export class BackgroundManager {
       task.status = "error"
       task.error = errorMsg
       task.completedAt = new Date()
+      if (task.rootSessionID) {
+        this.unregisterRootDescendant(task.rootSessionID)
+      }
       this.taskHistory.record(task.parentSessionID, { id: task.id, sessionID: task.sessionID, agent: task.agent, description: task.description, status: "error", category: task.category, startedAt: task.startedAt, completedAt: task.completedAt })
 
       if (task.concurrencyKey) {
@@ -1354,8 +1363,12 @@ export class BackgroundManager {
       log("[background-agent] Cancelled pending task:", { taskId, key })
     }
 
+    const wasRunning = task.status === "running"
     task.status = "cancelled"
     task.completedAt = new Date()
+    if (wasRunning && task.rootSessionID) {
+      this.unregisterRootDescendant(task.rootSessionID)
+    }
     if (reason) {
       task.error = reason
     }
@@ -1476,6 +1489,10 @@ export class BackgroundManager {
     task.completedAt = new Date()
     this.taskHistory.record(task.parentSessionID, { id: task.id, sessionID: task.sessionID, agent: task.agent, description: task.description, status: "completed", category: task.category, startedAt: task.startedAt, completedAt: task.completedAt })
 
+    if (task.rootSessionID) {
+      this.unregisterRootDescendant(task.rootSessionID)
+    }
+
     removeTaskToastTracking(task.id)
 
     // Release concurrency BEFORE any async operations to prevent slot leaks
@@ -1714,6 +1731,9 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea
         task.status = "error"
         task.error = errorMessage
         task.completedAt = new Date()
+        if (!wasPending && task.rootSessionID) {
+          this.unregisterRootDescendant(task.rootSessionID)
+        }
         this.taskHistory.record(task.parentSessionID, { id: task.id, sessionID: task.sessionID, agent: task.agent, description: task.description, status: "error", category: task.category, startedAt: task.startedAt, completedAt: task.completedAt })
         if (task.concurrencyKey) {
           this.concurrencyManager.release(task.concurrencyKey)