From c8898034f9756afd084b188c8c6466a99e1acc7b Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Fri, 13 Mar 2026 21:36:42 +0000
Subject: [PATCH] refactor: share agent wait dedupe cleanup

---
 .../server-methods/agent-wait-dedupe.ts       | 33 ++++++++-----------
 1 file changed, 13 insertions(+), 20 deletions(-)

diff --git a/src/gateway/server-methods/agent-wait-dedupe.ts b/src/gateway/server-methods/agent-wait-dedupe.ts
index 98d0df72fa3..50629beb3eb 100644
--- a/src/gateway/server-methods/agent-wait-dedupe.ts
+++ b/src/gateway/server-methods/agent-wait-dedupe.ts
@@ -23,6 +23,17 @@ function asFiniteNumber(value: unknown): number | undefined {
   return typeof value === "number" && Number.isFinite(value) ? value : undefined;
 }
 
+function removeWaiter(runId: string, waiter: () => void): void {
+  const waiters = AGENT_WAITERS_BY_RUN_ID.get(runId);
+  if (!waiters) {
+    return;
+  }
+  waiters.delete(waiter);
+  if (waiters.size === 0) {
+    AGENT_WAITERS_BY_RUN_ID.delete(runId);
+  }
+}
+
 function addWaiter(runId: string, waiter: () => void): () => void {
   const normalizedRunId = runId.trim();
   if (!normalizedRunId) {
@@ -31,28 +42,10 @@ function addWaiter(runId: string, waiter: () => void): () => void {
   const existing = AGENT_WAITERS_BY_RUN_ID.get(normalizedRunId);
   if (existing) {
     existing.add(waiter);
-    return () => {
-      const waiters = AGENT_WAITERS_BY_RUN_ID.get(normalizedRunId);
-      if (!waiters) {
-        return;
-      }
-      waiters.delete(waiter);
-      if (waiters.size === 0) {
-        AGENT_WAITERS_BY_RUN_ID.delete(normalizedRunId);
-      }
-    };
+    return () => removeWaiter(normalizedRunId, waiter);
   }
   AGENT_WAITERS_BY_RUN_ID.set(normalizedRunId, new Set([waiter]));
-  return () => {
-    const waiters = AGENT_WAITERS_BY_RUN_ID.get(normalizedRunId);
-    if (!waiters) {
-      return;
-    }
-    waiters.delete(waiter);
-    if (waiters.size === 0) {
-      AGENT_WAITERS_BY_RUN_ID.delete(normalizedRunId);
-    }
-  };
+  return () => removeWaiter(normalizedRunId, waiter);
 }
 
 function notifyWaiters(runId: string): void {