oira666_pi-limits-wait 0.4.3 → 0.4.5

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "oira666_pi-limits-wait",
-  "version": "0.4.3",
+  "version": "0.4.5",
   "description": "Extension for PI that allows to automatically wait for limits to reset and retry the request",
   "type": "module",
   "license": "MIT",

package/src/index.ts

@@ -316,6 +316,25 @@ function rememberNonRetryableFailure(model: Model<Api>, errorMessage: string): v
   ensureRateLimitedModelsStatus();
 }
 
+function isAbortMessage(msg: string): boolean {
+  const lower = msg.toLowerCase();
+  return (
+    lower.includes("operation aborted") ||
+    lower.includes("request aborted") ||
+    lower.includes("aborterror") ||
+    /\baborted\b/.test(lower)
+  );
+}
+
+function isAbortErrorEvent(event: AssistantMessageEvent, signal?: AbortSignal): boolean {
+  if (signal?.aborted) return true;
+  if (event.type !== "error") return false;
+  const reason = "reason" in event ? String(event.reason) : "";
+  const stopReason = event.error.stopReason ? String(event.error.stopReason) : "";
+  const message = event.error.errorMessage ?? "";
+  return reason === "aborted" || stopReason === "aborted" || isAbortMessage(message);
+}
+
 function earliestCandidateDeadline(current: Model<Api>): number | undefined {
   const deadlines = candidateOrder(current)
     .map((entry) => activeLimit(entry.model)?.deadline ?? activeNonRetryableFailure(entry.model)?.deadline)
@@ -636,7 +655,7 @@ type StreamSimpleFn = (
   model: Model<Api>,
   context: Context,
   options?: SimpleStreamOptions,
-) => AssistantMessageEventStream;
+) => AssistantMessageEventStream | Promise<AssistantMessageEventStream>;
 
 export function __configureFallbackModelsForTests(
   models: FallbackModel[],
@@ -712,12 +731,28 @@ export function streamWithLimitsRetry(
             : builtinStreamSimpleByApi.get(attempt.model.api)
               ?? getApiProviders().find((provider) => provider.api === attempt.model.api)?.streamSimple;
           if (!attemptDelegate) throw new Error(`No stream handler registered for API ${attempt.model.api}.`);
-          const inner = attemptDelegate(attempt.model, context, attemptOptions);
+          const inner = await attemptDelegate(attempt.model, context, attemptOptions);
           for await (const event of inner) {
             if (!committed) {
               if (event.type === "error") {
                 const errMsg = event.error.errorMessage ?? "";
-                retryable = !options?.signal?.aborted ? getRetryableError(errMsg) : undefined;
+
+                if (isAbortErrorEvent(event, options?.signal)) {
+                  if (buffer.length > 0) {
+                    flush(buffer);
+                  } else {
+                    output.push({ type: "start", partial: freshMessage(attempt.model) });
+                    committed = true;
+                  }
+                  const error = freshMessage(attempt.model);
+                  error.stopReason = "aborted";
+                  error.errorMessage = errMsg || "Operation aborted";
+                  output.push({ type: "error", reason: "aborted", error });
+                  output.end();
+                  return;
+                }
+
+                retryable = getRetryableError(errMsg);
                 if (retryable) break;
 
                 if (allowFallback) {
@@ -767,26 +802,27 @@ export function streamWithLimitsRetry(
           }
         } catch (err) {
           const errMsg = err instanceof Error ? err.message : String(err);
-          retryable = !options?.signal?.aborted ? getRetryableError(errMsg) : undefined;
+          const aborted = Boolean(options?.signal?.aborted) || isAbortMessage(errMsg);
+          retryable = !aborted ? getRetryableError(errMsg) : undefined;
           if (!retryable) {
-            if (allowFallback && !options?.signal?.aborted) {
+            if (allowFallback && !aborted) {
               nonRetryableError = { message: errMsg };
             } else {
-            if (!committed) {
-              // Synthetic start+error so the stream protocol stays valid.
-              output.push({ type: "start", partial: freshMessage(attempt.model) });
-              committed = true;
-            }
-            const error = freshMessage(attempt.model);
-            error.stopReason = options?.signal?.aborted ? "aborted" : "error";
-            error.errorMessage = errMsg;
-            output.push({
-              type: "error",
-              reason: error.stopReason as "error" | "aborted",
-              error,
-            });
-            output.end();
-            return;
+              if (!committed) {
+                // Synthetic start+error so the stream protocol stays valid.
+                output.push({ type: "start", partial: freshMessage(attempt.model) });
+                committed = true;
+              }
+              const error = freshMessage(attempt.model);
+              error.stopReason = aborted ? "aborted" : "error";
+              error.errorMessage = errMsg;
+              output.push({
+                type: "error",
+                reason: error.stopReason as "error" | "aborted",
+                error,
+              });
+              output.end();
+              return;
             }
           }
         }