@animalabs/membrane 0.5.76 → 0.5.78

package/src/membrane.ts

@@ -31,6 +31,7 @@ import {
   DEFAULT_RETRY_CONFIG,
   MembraneError,
   classifyError,
+  isOverloadedError,
   isTextContent,
   isAbortedResponse,
 } from './types/index.js';
@@ -80,7 +81,11 @@ export class Membrane {
   ) {
     this.adapter = adapter;
     this.registry = config.registry;
-    this.retryConfig = { ...DEFAULT_RETRY_CONFIG, ...config.retry };
+    this.retryConfig = {
+      ...DEFAULT_RETRY_CONFIG,
+      ...config.retry,
+      overloaded: { ...DEFAULT_RETRY_CONFIG.overloaded, ...config.retry?.overloaded },
+    };
     this.config = config;
     // Use provided formatter or default to AnthropicXmlFormatter
     this.formatter = config.formatter ?? new AnthropicXmlFormatter();
@@ -100,6 +105,10 @@ export class Membrane {
     const startTime = Date.now();
     let attempts = 0;
     let rawRequest: unknown;
+    // Counted separately from `attempts` (the transport-error budget): a
+    // refusal is a successful HTTP call with an unwanted verdict, and letting
+    // it consume error retries would couple two unrelated budgets.
+    let refusalRetriesUsed = 0;
 
     while (true) {
       attempts++;
@@ -135,6 +144,19 @@ export class Membrane {
           rawRequest
         );
 
+        // Re-issue a content-policy refusal (opt-in, default off). Safe here
+        // in a way the streaming paths are not: nothing has reached the
+        // caller yet, so the abandoned attempt leaves no trace to retract.
+        // Deliberately BEFORE afterResponse — a hook that logs or transforms
+        // should see the attempt that actually stands, not the discarded one.
+        if (
+          response.stopReason === 'refusal' &&
+          refusalRetriesUsed < Math.max(0, options.refusalRetries ?? 0)
+        ) {
+          refusalRetriesUsed++;
+          continue;
+        }
+
         // Call afterResponse hook
         if (this.config.hooks?.afterResponse) {
           return await this.config.hooks.afterResponse(response, providerResponse.raw);
@@ -146,12 +168,23 @@ export class Membrane {
         const errorInfo = classifyError(error);
         errorInfo.rawRequest = rawRequest;
 
-        // Rate limits (429) always retry up to 5 attempts regardless of config.
-        // Other retryable errors only retry when maxRetries > 0.
+        // Rate limits (429) always retry up to 5 attempts regardless of
+        // config, and overloaded (529) always retries on its own longer
+        // schedule — both are transient by definition, and the default
+        // maxRetries of 0 would otherwise turn a capacity blip into a dead
+        // turn. Other retryable errors only retry when maxRetries > 0.
+        // overloaded.maxRetries: 0 disables the dedicated policy entirely;
+        // the 529 then follows the base config like any retryable server
+        // error (exactly the pre-policy behavior), rather than being
+        // silently re-promoted to the long schedule by a positive base limit.
         const isRateLimit = errorInfo.type === 'rate_limit';
+        const isOverloaded =
+          isOverloadedError(errorInfo) && this.retryConfig.overloaded.maxRetries > 0;
         const effectiveMax = isRateLimit
           ? Math.max(this.retryConfig.maxRetries, 5)
-          : this.retryConfig.maxRetries;
+          : isOverloaded
+            ? Math.max(this.retryConfig.maxRetries, this.retryConfig.overloaded.maxRetries)
+            : this.retryConfig.maxRetries;
 
         if (errorInfo.retryable && attempts < effectiveMax) {
           // Check hook for retry decision
@@ -163,7 +196,7 @@ export class Membrane {
           }
 
           // Wait before retry (abort-aware)
-          const delay = this.calculateRetryDelay(attempts);
+          const delay = this.calculateRetryDelay(attempts, isOverloaded);
           await this.sleep(delay, options.signal);
           continue;
         }
@@ -217,11 +250,74 @@ export class Membrane {
 
     // Determine tool mode
     const toolMode = this.resolveToolMode(request);
+    const useNative = toolMode === 'native' && !!request.tools && request.tools.length > 0;
+
+    // Overloaded (529) pre-emission retry. The streaming paths have no retry
+    // loop of their own, so a capacity error used to kill the turn outright —
+    // and 529s most often arrive INSTEAD of a stream, before anything reaches
+    // the caller, where retrying is transparent. Once any callback has
+    // delivered output (tokens, blocks, usage), retrying would replay content
+    // the caller already consumed, so mid-stream errors still throw.
+    let attempts = 0;
+    const retryDelaysMs: number[] = [];
+    while (true) {
+      attempts++;
+      let emitted = false;
+      const mark = <A extends unknown[], R>(fn?: (...args: A) => R) =>
+        fn && ((...args: A): R => { emitted = true; return fn(...args); });
+      const tracked: StreamOptions = {
+        ...options,
+        onChunk: mark(options.onChunk),
+        onContentBlockUpdate: mark(options.onContentBlockUpdate),
+        onToolCalls: mark(options.onToolCalls),
+        onPreToolContent: mark(options.onPreToolContent),
+        onUsage: mark(options.onUsage),
+        onBlock: mark(options.onBlock),
+        onResponse: mark(options.onResponse),
+        // onRequest fires before the send — it is not an emission.
+      };
 
-    if (toolMode === 'native' && request.tools && request.tools.length > 0) {
-      return this.streamWithNativeTools(request, options);
-    } else {
-      return this.streamWithXmlTools(request, options);
+      try {
+        const result = useNative
+          ? await this.streamWithNativeTools(request, tracked)
+          : await this.streamWithXmlTools(request, tracked);
+        // The inner paths report attempts: 1 — they can't see this wrapper.
+        // A call that succeeded after N overloaded retries must not look like
+        // a first-attempt success in durable logs, so patch the real count
+        // (and the waits) into the response telemetry.
+        if (attempts > 1 && 'details' in result) {
+          result.details.timing.attempts = attempts;
+          result.details.timing.retryDelaysMs = retryDelaysMs;
+        }
+        return result;
+      } catch (error) {
+        const errorInfo = classifyError(error);
+        // Same semantics as complete(): maxRetries bounds total attempts,
+        // the overloaded floor applies over the base config, and
+        // overloaded.maxRetries: 0 opts out of stream retries entirely
+        // (streaming had no retry before this policy existed).
+        const overloadedEnabled = this.retryConfig.overloaded.maxRetries > 0;
+        const maxOverloaded = Math.max(
+          this.retryConfig.maxRetries,
+          this.retryConfig.overloaded.maxRetries
+        );
+        if (!emitted && overloadedEnabled && isOverloadedError(errorInfo) && attempts < maxOverloaded) {
+          // Honor the same pre-retry hook contract as complete(): hosts use
+          // onError for circuit-breaking, and its 'abort' decision must work
+          // on the streaming path too.
+          if (this.config.hooks?.onError) {
+            const decision = await this.config.hooks.onError(errorInfo, attempts);
+            if (decision === 'abort') {
+              throw error;
+            }
+          }
+          const delay = this.calculateRetryDelay(attempts, true);
+          retryDelaysMs.push(delay);
+          await this.sleep(delay, options.signal);
+          continue;
+        }
+        throw error;
+      }
     }
   }
 
@@ -330,6 +426,52 @@ export class Membrane {
     // from blocks the model itself opened during generation
     const prefillDepths = parser.getDepths();
 
+    // Resumption spin guards (issue #39). Observed live on Ash 2026-07-26:
+    // each automatic resumption re-sent ~172k input tokens, streamed ~6
+    // output tokens, and stopped on the same (dropped) stop sequence — 43
+    // rounds, ~7M input tokens, zero progress, found only because a human
+    // noticed. Two guards, both scoped to AUTOMATIC false-positive
+    // resumptions — tool rounds are real caller-governed work (maxToolDepth
+    // / the yielding API's uncapped contract) and are never counted here:
+    //   - stall guard: several CONSECUTIVE resumptions that each stream
+    //     almost nothing and stop identically end the turn ('no_progress').
+    //     One short repeated round is low progress, not proof of none — a
+    //     stop sequence inside legitimate tool-argument text can cause a
+    //     couple of short resumptions on the way to completing.
+    //   - round cap: a hard bound on resumptions per turn ('round_limit'),
+    //     the backstop for a spin that keeps technically progressing.
+    const MIN_ROUND_PROGRESS_CHARS = 16;
+    const MAX_CONSECUTIVE_STALLED_RESUMPTIONS = 3;
+    const RESUMPTION_WARN_ROUNDS = 5;
+    const maxResumptionRounds = options.maxResumptionRounds ?? 24;
+    let resumptionRounds = 0;
+    let consecutiveStalledResumptions = 0;
+    let enteredViaResumption = false;
+    let prevRoundStopSequence: string | undefined;
+    const warnLog = this.config.logger ?? console;
+
+    /** Count an automatic resumption; emits the visibility warning at the
+     *  threshold and returns false when the cap says the turn should end. */
+    const registerResumptionRound = (): boolean => {
+      resumptionRounds++;
+      if (resumptionRounds === RESUMPTION_WARN_ROUNDS) {
+        warnLog.warn(
+          `[membrane] automatic resumption at round ${resumptionRounds} ` +
+          `(${totalUsage.inputTokens} input tokens so far this turn) — ` +
+          `a spin shows up here before it shows up on the bill`
+        );
+      }
+      if (resumptionRounds > maxResumptionRounds) {
+        warnLog.warn(
+          `[membrane] automatic resumption cap (${maxResumptionRounds}) reached — ` +
+          `ending turn with stopReason 'round_limit'. ` +
+          `${totalUsage.inputTokens} input tokens spent this turn.`
+        );
+        return false;
+      }
+      return true;
+    };
+
     try {
       // Tool execution loop
       while (toolDepth <= maxToolDepth) {
@@ -338,7 +480,10 @@ export class Membrane {
         let detectedStopSequence: string | null = null;
         let truncatedAccumulated: string | null = null;
 
-        // Track where to start checking for stop sequences (skip already-processed content)
+        // Track where to start checking for stop sequences (skip already-processed content).
+        // Also the round's progress baseline: XML we pushed ourselves at the
+        // end of the previous round (tool results, closing tags) sits below
+        // this index and doesn't count as model progress.
         const checkFromIndex = parser.getAccumulated().length;
 
         // Stream from provider
@@ -465,6 +610,38 @@ export class Membrane {
         // Get accumulated text from parser
         const accumulated = parser.getAccumulated();
 
+        // Stall accounting (issue #39): only rounds ENTERED via automatic
+        // resumption can stall — tool rounds are caller-governed work and a
+        // real tool call is longer than the threshold anyway. A stall is a
+        // resumption that streamed almost nothing and stopped identically to
+        // the previous round; the turn ends only after several IN A ROW
+        // (one short repeated round is low progress, not proof of none —
+        // a stop sequence inside legitimate tool-argument text can cause a
+        // couple of short resumptions on the way to completing).
+        const streamedThisRound = accumulated.length - checkFromIndex;
+        if (
+          enteredViaResumption &&
+          lastStopReason === 'stop_sequence' &&
+          streamedThisRound < MIN_ROUND_PROGRESS_CHARS &&
+          lastStopSequence === prevRoundStopSequence
+        ) {
+          consecutiveStalledResumptions++;
+          if (consecutiveStalledResumptions >= MAX_CONSECUTIVE_STALLED_RESUMPTIONS) {
+            warnLog.warn(
+              `[membrane] ${consecutiveStalledResumptions} consecutive automatic resumptions ` +
+              `made no progress (${streamedThisRound} chars this round, stop ` +
+              `${JSON.stringify(lastStopSequence ?? null)} repeated) — ending turn with ` +
+              `stopReason 'no_progress'. ${totalUsage.inputTokens} input tokens spent this turn.`
+            );
+            lastStopReason = 'no_progress';
+            break;
+          }
+        } else {
+          consecutiveStalledResumptions = 0;
+        }
+        prevRoundStopSequence = lastStopSequence;
+        enteredViaResumption = false;
+
         // Check for tool calls (if handler provided)
         if (onToolCalls && streamResult.stopSequence === '</function_calls>') {
           // Append the closing tag (we truncated before it, or API stopped before it)
@@ -655,7 +832,9 @@ export class Membrane {
               );
             }
 
-            // Reset parser state for new streaming iteration
+            // Reset parser state for new streaming iteration. Tool rounds
+            // are the caller's work — they count against maxToolDepth only,
+            // never against the resumption guards (issue #39 review).
             parser.resetForNewIteration();
             toolDepth++;
             continue;
@@ -692,6 +871,11 @@ export class Membrane {
           if (toolDepth > maxToolDepth) {
             break;
           }
+          if (!registerResumptionRound()) {
+            lastStopReason = 'round_limit';
+            break;
+          }
+          enteredViaResumption = true;
           prefillResult.assistantPrefill = parser.getAccumulated();
           providerRequest = this.buildContinuationRequest(
             request,
@@ -1514,6 +1698,23 @@ export class Membrane {
        * separate `streamOnceWithoutHook` so the bypass is intentional.
        */
       normalizedRequest: NormalizedRequest;
+      /**
+       * Re-issue this attempt when the provider ends it with
+       * `stop_reason: 'refusal'` (see RetryingEvent). Default 0 = off, so
+       * every existing caller keeps byte-identical behaviour.
+       */
+      refusalRetries?: number;
+      /**
+       * REQUIRED to enable streaming retries. Called immediately before a
+       * re-issue so the caller can discard the abandoned attempt: reset its
+       * accumulators and tell its own consumer to drop what it emitted.
+       *
+       * Without it a retry would silently concatenate two attempts, so a
+       * caller that does not pass this simply does not get retries — an
+       * unaware consumer can never be corrupted by enabling the option
+       * somewhere upstream.
+       */
+      onRetrying?: (info: { attempt: number; maxAttempts: number; category?: string }) => void;
     }
   ) {
     // Strip `normalizedRequest` before forwarding to the adapter — it's
@@ -1521,9 +1722,21 @@ export class Membrane {
     // compatibility won't catch the excess field (checked only on object
     // literals, not on variables). Leaving it in would silently leak the
     // normalized form into every adapter's options.
-    const { normalizedRequest, ...adapterOptions } = options;
+    const { normalizedRequest, refusalRetries, onRetrying, ...adapterOptions } = options;
     const finalRequest = (await this.applyBeforeRequestHook(normalizedRequest, request)) as typeof request;
-    return await this.adapter.stream(finalRequest, callbacks, adapterOptions);
+
+    // Retries are only safe when the caller can discard the abandoned
+    // attempt, so they require BOTH a budget and an onRetrying hook.
+    const maxAttempts = onRetrying ? Math.max(0, refusalRetries ?? 0) : 0;
+    let retried = 0;
+    while (true) {
+      const result = await this.adapter.stream(finalRequest, callbacks, adapterOptions);
+      if (result.stopReason !== 'refusal' || retried >= maxAttempts) return result;
+      retried++;
+      const category = (result.raw as { response?: { stop_details?: { category?: string } } } | undefined)
+        ?.response?.stop_details?.category;
+      onRetrying!({ attempt: retried, maxAttempts, category });
+    }
   }
 
   private buildContinuationRequest(
@@ -1893,10 +2106,15 @@ export class Membrane {
     return pricing ? calculateCost(usage, pricing) : undefined;
   }
 
-  private calculateRetryDelay(attempt: number): number {
-    const { retryDelayMs, backoffMultiplier, maxRetryDelayMs } = this.retryConfig;
-    const delay = retryDelayMs * Math.pow(backoffMultiplier, attempt - 1);
-    return Math.min(delay, maxRetryDelayMs);
+  private calculateRetryDelay(attempt: number, overloaded = false): number {
+    const { retryDelayMs, backoffMultiplier, maxRetryDelayMs } = overloaded
+      ? this.retryConfig.overloaded
+      : this.retryConfig;
+    const delay = Math.min(retryDelayMs * Math.pow(backoffMultiplier, attempt - 1), maxRetryDelayMs);
+    // Equal jitter on the overloaded schedule only: a capacity storm is
+    // exactly the case where a fleet retrying in sync re-creates the
+    // stampede it's backing off from. [delay/2, delay) keeps the wait long.
+    return overloaded ? Math.floor(delay / 2 + Math.random() * (delay / 2)) : delay;
   }
 
   private attachRawRequest(error: unknown, rawRequest: unknown): Error {
@@ -2000,6 +2218,18 @@ export class Membrane {
   ): YieldingStream {
     const toolMode = this.resolveToolMode(request);
 
+    // refusalRetries is implemented on the native path only. The XML path
+    // accumulates into a streaming parser carrying prefill context and
+    // resumption depths; rolling that back mid-turn is a separate problem,
+    // and a partial implementation would corrupt the turn instead of
+    // retrying it. Fail LOUD and OFF rather than silently mis-retrying.
+    if (toolMode !== 'native' && (options.refusalRetries ?? 0) > 0) {
+      (this.config.logger ?? console).warn(
+        '[membrane] refusalRetries is ignored in XML tool mode ' +
+        '(native-only for now) — the turn will surface the refusal as before.',
+      );
+    }
+
     // Create the yielding stream with the appropriate inference runner
     const runInference = toolMode === 'native'
       ? (stream: YieldingStreamImpl) => this.runNativeToolsYielding(request, options, stream)
@@ -2033,6 +2263,29 @@ export class Membrane {
         ? Infinity
         : maxToolDepthOpt;
 
+    // Resumption spin guards (issue #39). This is the path the Ash spin ran
+    // on: tool depth here is unlimited BY DESIGN (the caller budgets its own
+    // tool work — that contract stands untouched), and the false-positive
+    // resumption path counted against that same unlimited bound — 43 rounds
+    // × ~172k input tokens of zero progress. Only AUTOMATIC resumptions are
+    // guarded: a stall guard (consecutive no-progress resumptions →
+    // 'no_progress') and a hard resumption cap ('round_limit'). Tool rounds
+    // are never counted. See streamWithXmlTools for the rationale details.
+    const MIN_ROUND_PROGRESS_CHARS = 16;
+    const MAX_CONSECUTIVE_STALLED_RESUMPTIONS = 3;
+    const RESUMPTION_WARN_ROUNDS = 5;
+    const maxResumptionRounds =
+      options.maxResumptionRounds === undefined
+        ? 24
+        : options.maxResumptionRounds === -1
+          ? Infinity
+          : options.maxResumptionRounds;
+    let resumptionRounds = 0;
+    let consecutiveStalledResumptions = 0;
+    let enteredViaResumption = false;
+    let prevRoundStopSequence: string | undefined;
+    const warnLog = this.config.logger ?? console;
+
     // Initialize parser from formatter for format-specific tracking
     const formatter = this.formatter;
     const parser = formatter.createStreamParser();
@@ -2078,6 +2331,28 @@ export class Membrane {
     // from blocks the model itself opened during generation
     const prefillDepths = parser.getDepths();
 
+    /** Count an automatic resumption; emits the visibility warning at the
+     *  threshold and returns false when the cap says the turn should end. */
+    const registerResumptionRound = (): boolean => {
+      resumptionRounds++;
+      if (resumptionRounds === RESUMPTION_WARN_ROUNDS) {
+        warnLog.warn(
+          `[membrane] automatic resumption at round ${resumptionRounds} ` +
+          `(${totalUsage.inputTokens} input tokens so far this turn) — ` +
+          `a spin shows up here before it shows up on the bill`
+        );
+      }
+      if (resumptionRounds > maxResumptionRounds) {
+        warnLog.warn(
+          `[membrane] automatic resumption cap (${maxResumptionRounds}) reached — ` +
+          `ending turn with stopReason 'round_limit'. ` +
+          `${totalUsage.inputTokens} input tokens spent this turn.`
+        );
+        return false;
+      }
+      return true;
+    };
+
     try {
       // Tool execution loop
       while (toolDepth <= maxToolDepth) {
@@ -2206,6 +2481,34 @@ export class Membrane {
           }
         }
 
+        // Stall accounting (issue #39): only rounds ENTERED via automatic
+        // resumption can stall; the turn ends only after several consecutive
+        // stalls. Tool rounds are never counted. Mirrors streamWithXmlTools —
+        // see the detailed rationale there.
+        const streamedThisRound = parser.getAccumulated().length - checkFromIndex;
+        if (
+          enteredViaResumption &&
+          lastStopReason === 'stop_sequence' &&
+          streamedThisRound < MIN_ROUND_PROGRESS_CHARS &&
+          lastStopSequence === prevRoundStopSequence
+        ) {
+          consecutiveStalledResumptions++;
+          if (consecutiveStalledResumptions >= MAX_CONSECUTIVE_STALLED_RESUMPTIONS) {
+            warnLog.warn(
+              `[membrane] ${consecutiveStalledResumptions} consecutive automatic resumptions ` +
+              `made no progress (${streamedThisRound} chars this round, stop ` +
+              `${JSON.stringify(lastStopSequence ?? null)} repeated) — ending turn with ` +
+              `stopReason 'no_progress'. ${totalUsage.inputTokens} input tokens spent this turn.`
+            );
+            lastStopReason = 'no_progress';
+            break;
+          }
+        } else {
+          consecutiveStalledResumptions = 0;
+        }
+        prevRoundStopSequence = lastStopSequence;
+        enteredViaResumption = false;
+
         // Check for tool calls
         if (streamResult.stopSequence === '</function_calls>') {
           const closeTag = '</function_calls>';
@@ -2422,6 +2725,9 @@ export class Membrane {
               );
             }
 
+            // Tool rounds are the caller's work — they count against
+            // maxToolDepth only, never against the resumption guards
+            // (issue #39 review: the uncapped tool-loop contract stands).
             parser.resetForNewIteration();
             toolDepth++;
             continue;
@@ -2453,6 +2759,11 @@ export class Membrane {
           if (toolDepth > maxToolDepth) {
             break;
           }
+          if (!registerResumptionRound()) {
+            lastStopReason = 'round_limit';
+            break;
+          }
+          enteredViaResumption = true;
           prefillResult.assistantPrefill = parser.getAccumulated();
           providerRequest = this.buildContinuationRequest(
             request,
@@ -2571,6 +2882,9 @@ export class Membrane {
         // Stream from provider
         let textAccumulated = '';
         let blockIndex = 0;
+        // Where this attempt starts inside the tool-loop-spanning buffer, so
+        // a refusal retry can roll back exactly this attempt's contribution.
+        const allTextBefore = allTextAccumulated.length;
         // Track block-type from the provider's content_block_start signal so
         // every token chunk is tagged with the membrane block it belongs to.
         // Without this, thinking_delta chunks get mislabelled as 'text' and
@@ -2649,6 +2963,25 @@ export class Membrane {
             idleTimeoutMs: options.idleTimeoutMs,
             normalizedRequest: request,
             onRequest: (req: unknown) => { rawRequest = req; },
+            refusalRetries: options.refusalRetries,
+            // Discard the refused attempt: roll the accumulators back to
+            // where this attempt began and tell the consumer to drop what it
+            // already received. `allTextAccumulated` spans the whole tool
+            // loop, so it is truncated rather than cleared.
+            onRetrying: (info) => {
+              allTextAccumulated = allTextAccumulated.slice(0, allTextBefore);
+              textAccumulated = '';
+              blockIndex = 0;
+              currentBlockType = 'text';
+              seenBlockIndices.clear();
+              stream.emit({
+                type: 'retrying',
+                attempt: info.attempt,
+                maxAttempts: info.maxAttempts,
+                reason: 'refusal',
+                ...(info.category ? { category: info.category } : {}),
+              });
+            },
           }
         );