fix(sse): propagate AbortSignal to pre-fetch semaphore and rate-limit awaits

open-sse/handlers/chatCore.ts

@@ -2162,67 +2162,78 @@ export async function handleChatCore({
         accountSemaphoreKey && accountSemaphoreMaxConcurrency != null
           ? await acquireAccountSemaphore(accountSemaphoreKey, {
               maxConcurrency: accountSemaphoreMaxConcurrency,
+              signal: streamController.signal,
             })
           : () => {};
 
       try {
-        const rawResult = await withRateLimit(provider, connectionId, modelToCall, async () => {
-          let attempts = 0;
-          const maxAttempts = provider === "qwen" ? 3 : 1;
-
-          while (attempts < maxAttempts) {
-            const res = await executor.execute({
-              model: modelToCall,
-              body: bodyToSend,
-              stream: upstreamStream,
-              credentials: executionCredentials,
-              signal: streamController.signal,
-              log,
-              extendedContext,
-              upstreamExtraHeaders: buildUpstreamHeadersForExecute(modelToCall),
-              clientHeaders: buildExecutorClientHeaders(clientRawRequest?.headers, userAgent),
-              onCredentialsRefreshed,
-            });
-
-            // Qwen 429 strict quota backoff (wait 1.5s, 3s and retry)
-            if (provider === "qwen" && res.response.status === 429 && attempts < maxAttempts - 1) {
-              const bodyPeek = await res.response
-                .clone()
-                .text()
-                .catch(() => "");
-              if (bodyPeek.toLowerCase().includes("exceeded your current quota")) {
-                const delay = 1500 * (attempts + 1);
-                log?.warn?.("QWEN_RETRY", `Quota 429 hit. Retrying in ${delay}ms...`);
-                await new Promise((r) => setTimeout(r, delay));
-                attempts++;
-                continue;
+        const rawResult = await withRateLimit(
+          provider,
+          connectionId,
+          modelToCall,
+          async () => {
+            let attempts = 0;
+            const maxAttempts = provider === "qwen" ? 3 : 1;
+
+            while (attempts < maxAttempts) {
+              const res = await executor.execute({
+                model: modelToCall,
+                body: bodyToSend,
+                stream: upstreamStream,
+                credentials: executionCredentials,
+                signal: streamController.signal,
+                log,
+                extendedContext,
+                upstreamExtraHeaders: buildUpstreamHeadersForExecute(modelToCall),
+                clientHeaders: buildExecutorClientHeaders(clientRawRequest?.headers, userAgent),
+                onCredentialsRefreshed,
+              });
+
+              // Qwen 429 strict quota backoff (wait 1.5s, 3s and retry)
+              if (
+                provider === "qwen" &&
+                res.response.status === 429 &&
+                attempts < maxAttempts - 1
+              ) {
+                const bodyPeek = await res.response
+                  .clone()
+                  .text()
+                  .catch(() => "");
+                if (bodyPeek.toLowerCase().includes("exceeded your current quota")) {
+                  const delay = 1500 * (attempts + 1);
+                  log?.warn?.("QWEN_RETRY", `Quota 429 hit. Retrying in ${delay}ms...`);
+                  await new Promise((r) => setTimeout(r, delay));
+                  attempts++;
+                  continue;
+                }
               }
-            }
 
-            // For streaming: release the semaphore when the client drains or cancels the stream.
-            if (stream) {
-              const originalBody = res.response.body;
-              if (!originalBody) {
-                acquireAccountSemaphoreRelease();
-                return res;
+              // For streaming: release the semaphore when the client drains or cancels the stream.
+              if (stream) {
+                const originalBody = res.response.body;
+                if (!originalBody) {
+                  acquireAccountSemaphoreRelease();
+                  return res;
+                }
+
+                return {
+                  ...res,
+                  response: new Response(
+                    wrapReadableStreamWithFinalize(originalBody, acquireAccountSemaphoreRelease),
+                    {
+                      status: res.response.status,
+                      statusText: res.response.statusText,
+                      headers: res.response.headers,
+                    }
+                  ),
+                };
               }
 
-              return {
-                ...res,
-                response: new Response(
-                  wrapReadableStreamWithFinalize(originalBody, acquireAccountSemaphoreRelease),
-                  {
-                    status: res.response.status,
-                    statusText: res.response.statusText,
-                    headers: res.response.headers,
-                  }
-                ),
-              };
+              return res;
             }
-
-            return res;
-          }
-        });
+          },
+          streamController.signal
+        );
 
         if (stream) {
           return rawResult;

open-sse/services/accountSemaphore.ts

@@ -28,6 +28,7 @@ interface AccountGate {
 export interface AcquireAccountSemaphoreOptions {
   maxConcurrency?: number | null;
   timeoutMs?: number;
+  signal?: AbortSignal | null;
 }
 
 export interface AccountSemaphoreStatsEntry {
@@ -168,18 +169,34 @@ function createSemaphoreTimeoutError(
   return error;
 }
 
+function makeAbortError(signal: AbortSignal): Error {
+  const reason = signal.reason;
+  if (reason instanceof Error) return reason;
+  const err = new Error(typeof reason === "string" ? reason : "The operation was aborted");
+  err.name = "AbortError";
+  return err;
+}
+
 /**
  * Acquire a slot for a provider/model/account tuple.
  * Returns an idempotent release function that is safe to call in finally blocks.
  */
 export function acquire(
   semaphoreKey: string,
-  { maxConcurrency = null, timeoutMs = DEFAULT_TIMEOUT_MS }: AcquireAccountSemaphoreOptions = {}
+  {
+    maxConcurrency = null,
+    timeoutMs = DEFAULT_TIMEOUT_MS,
+    signal = null,
+  }: AcquireAccountSemaphoreOptions = {}
 ): Promise<() => void> {
   if (isBypassed(maxConcurrency)) {
     return Promise.resolve(createNoopReleaseFn());
   }
 
+  if (signal?.aborted) {
+    return Promise.reject(makeAbortError(signal));
+  }
+
   const gate = ensureGate(semaphoreKey, maxConcurrency);
   clearCleanupTimer(gate);
 
@@ -189,7 +206,16 @@ export function acquire(
   }
 
   return new Promise((resolve, reject) => {
+    let abortListener: (() => void) | null = null;
+
+    const cleanup = () => {
+      if (abortListener && signal) {
+        signal.removeEventListener("abort", abortListener);
+      }
+    };
+
     const timer = setTimeout(() => {
+      cleanup();
       const nextGate = gates.get(semaphoreKey);
       if (!nextGate) {
         reject(createSemaphoreTimeoutError(semaphoreKey, timeoutMs));
@@ -209,7 +235,49 @@ export function acquire(
     }, timeoutMs);
 
     timer.unref?.();
-    gate.queue.push({ resolve, reject, timer });
+
+    const queueItem: QueuedAcquire = {
+      resolve: (release) => {
+        cleanup();
+        resolve(release);
+      },
+      reject: (error) => {
+        cleanup();
+        reject(error);
+      },
+      timer,
+    };
+
+    gate.queue.push(queueItem);
+
+    if (signal) {
+      abortListener = () => {
+        cleanup();
+        clearTimeout(timer);
+
+        const nextGate = gates.get(semaphoreKey);
+        if (!nextGate) {
+          reject(makeAbortError(signal));
+          return;
+        }
+
+        const queueIndex = nextGate.queue.findIndex((item) => item.timer === timer);
+        if (queueIndex !== -1) {
+          nextGate.queue.splice(queueIndex, 1);
+        }
+
+        if (nextGate.running === 0 && nextGate.queue.length === 0) {
+          scheduleCleanup(semaphoreKey);
+        }
+
+        reject(makeAbortError(signal));
+      };
+      if (signal.aborted) {
+        abortListener();
+      } else {
+        signal.addEventListener("abort", abortListener);
+      }
+    }
   });
 }
 

open-sse/services/rateLimitManager.ts

@@ -274,18 +274,57 @@ function getLimiter(provider, connectionId, model = null) {
  * @param {string} connectionId - Connection ID
  * @param {string} model - Model name (optional, for per-model limits)
  * @param {Function} fn - The async function to execute (e.g., executor.execute)
+ * @param {AbortSignal} signal - Optional abort signal to cancel waiting
  * @returns {Promise<unknown>} Result of fn()
  */
-export async function withRateLimit(provider, connectionId, model, fn) {
+export async function withRateLimit(provider, connectionId, model, fn, signal = null) {
   if (!enabledConnections.has(connectionId)) {
     return fn();
   }
 
+  if (signal?.aborted) {
+    const reason = signal.reason;
+    if (reason instanceof Error) throw reason;
+    const err = new Error(typeof reason === "string" ? reason : "The operation was aborted");
+    err.name = "AbortError";
+    throw err;
+  }
+
   const limiter = getLimiter(provider, connectionId, model);
   const maxWaitMs = currentRequestQueueSettings.maxWaitMs;
   const scheduleOpts = maxWaitMs && maxWaitMs > 0 ? { expiration: maxWaitMs } : {};
+
   try {
-    return await limiter.schedule(scheduleOpts, fn);
+    if (signal) {
+      let abortListener: (() => void) | undefined;
+      const abortPromise = new Promise<never>((_, reject) => {
+        const onAbort = () => {
+          const reason = signal.reason;
+          const err =
+            reason instanceof Error
+              ? reason
+              : new Error(typeof reason === "string" ? reason : "The operation was aborted");
+          err.name = "AbortError";
+          reject(err);
+        };
+        if (signal.aborted) {
+          onAbort();
+          return;
+        }
+        abortListener = onAbort;
+        signal.addEventListener("abort", abortListener, { once: true });
+      });
+
+      try {
+        return await Promise.race([limiter.schedule(scheduleOpts, fn), abortPromise]);
+      } finally {
+        if (abortListener) {
+          signal.removeEventListener("abort", abortListener);
+        }
+      }
+    } else {
+      return await limiter.schedule(scheduleOpts, fn);
+    }
   } catch (err) {
     // Bottleneck throws when a job exceeds its expiration timeout.
     // Surface as a clear rate-limit timeout so callers can fallback.