From 129388387bddb66a655d75f41b5eab2bb04fb971 Mon Sep 17 00:00:00 2001
From: Gladdonilli <tianyi35@illinois.edu>
Date: Tue, 13 Jan 2026 14:40:50 +0800
Subject: [PATCH] fix: address PR review feedback

- Add pendingByParent cleanup to ALL completion paths (session.idle, polling, stability)
- Add null guard for task.parentSessionID before Map access
- Add consistency guard in prune function (set concurrencyKey = undefined)
- Remove redundant setTimeout release (already released at completion)
---
 src/features/background-agent/manager.ts | 44 +++++++++++++++++++++---
 1 file changed, 39 insertions(+), 5 deletions(-)

diff --git a/src/features/background-agent/manager.ts b/src/features/background-agent/manager.ts
index b91b0d7c1..c2ad078f8 100644
--- a/src/features/background-agent/manager.ts
+++ b/src/features/background-agent/manager.ts
@@ -426,6 +426,16 @@ export class BackgroundManager {
           this.concurrencyManager.release(task.concurrencyKey)
           task.concurrencyKey = undefined  // Prevent double-release
         }
+        // Clean up pendingByParent to prevent stale entries
+        if (task.parentSessionID) {
+          const pending = this.pendingByParent.get(task.parentSessionID)
+          if (pending) {
+            pending.delete(task.id)
+            if (pending.size === 0) {
+              this.pendingByParent.delete(task.parentSessionID)
+            }
+          }
+        }
         this.markForNotification(task)
         await this.notifyParentSession(task)
         log("[background-agent] Task completed via session.idle event:", task.id)
@@ -453,11 +463,13 @@ export class BackgroundManager {
         task.concurrencyKey = undefined  // Prevent double-release
       }
       // Clean up pendingByParent to prevent stale entries
-      const pending = this.pendingByParent.get(task.parentSessionID)
-      if (pending) {
-        pending.delete(task.id)
-        if (pending.size === 0) {
-          this.pendingByParent.delete(task.parentSessionID)
+      if (task.parentSessionID) {
+        const pending = this.pendingByParent.get(task.parentSessionID)
+        if (pending) {
+          pending.delete(task.id)
+          if (pending.size === 0) {
+            this.pendingByParent.delete(task.parentSessionID)
+          }
         }
       }
       this.tasks.delete(task.id)
@@ -678,6 +690,7 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea
 
     const taskId = task.id
     setTimeout(() => {
+      // Concurrency already released at completion - just cleanup notifications and task
       this.clearNotificationsForTask(taskId)
       this.tasks.delete(taskId)
       log("[background-agent] Removed completed task from memory:", taskId)
@@ -717,6 +730,7 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea
         task.completedAt = new Date()
         if (task.concurrencyKey) {
           this.concurrencyManager.release(task.concurrencyKey)
+          task.concurrencyKey = undefined  // Prevent double-release
         }
         this.clearNotificationsForTask(taskId)
         this.tasks.delete(taskId)
@@ -775,6 +789,16 @@ try {
             this.concurrencyManager.release(task.concurrencyKey)
             task.concurrencyKey = undefined  // Prevent double-release
           }
+          // Clean up pendingByParent to prevent stale entries
+          if (task.parentSessionID) {
+            const pending = this.pendingByParent.get(task.parentSessionID)
+            if (pending) {
+              pending.delete(task.id)
+              if (pending.size === 0) {
+                this.pendingByParent.delete(task.parentSessionID)
+              }
+            }
+          }
           this.markForNotification(task)
           await this.notifyParentSession(task)
           log("[background-agent] Task completed via polling:", task.id)
@@ -846,6 +870,16 @@ if (lastMessage) {
                     this.concurrencyManager.release(task.concurrencyKey)
                     task.concurrencyKey = undefined  // Prevent double-release
                   }
+                  // Clean up pendingByParent to prevent stale entries
+                  if (task.parentSessionID) {
+                    const pending = this.pendingByParent.get(task.parentSessionID)
+                    if (pending) {
+                      pending.delete(task.id)
+                      if (pending.size === 0) {
+                        this.pendingByParent.delete(task.parentSessionID)
+                      }
+                    }
+                  }
                   this.markForNotification(task)
                   await this.notifyParentSession(task)
                   log("[background-agent] Task completed via stability detection:", task.id)