Kaynağa Gözat

Merge pull request #672 from mhenke/fix/delegate-task-429-backoff

Alvin 1 ay önce
ebeveyn
işleme
4e46c2f856

+ 10 - 0
src/config/schema.ts

@@ -180,6 +180,16 @@ export const FailoverConfigSchema = z
     enabled: z.boolean().default(true),
     timeoutMs: z.number().min(0).default(15000),
     retryDelayMs: z.number().min(0).default(500),
+    maxRetries: z
+      .number()
+      .int()
+      .min(0)
+      .default(3)
+      .describe(
+        'Number of consecutive 429/rate-limit responses tolerated on the ' +
+          'same model before aborting (or swapping to the next fallback ' +
+          'model when a chain is configured).',
+      ),
     retry_on_empty: z
       .boolean()
       .default(true)

+ 72 - 9
src/hooks/foreground-fallback/index.test.ts

@@ -187,7 +187,8 @@ describe('ForegroundFallbackManager session.error', () => {
       },
     });
 
-    expect(mocks.abort).toHaveBeenCalledTimes(1);
+    // promptAsync is called directly (no abort needed when it succeeds)
+    expect(mocks.abort).toHaveBeenCalledTimes(0);
     expect(mocks.promptAsync).toHaveBeenCalledTimes(1);
 
     const call = mocks.promptAsync.mock.calls[0] as [
@@ -268,6 +269,7 @@ describe('ForegroundFallbackManager session.error', () => {
       },
     });
 
+    expect(mocks.abort).not.toHaveBeenCalled();
     expect(mocks.promptAsync).not.toHaveBeenCalled();
   });
 
@@ -287,10 +289,13 @@ describe('ForegroundFallbackManager session.error', () => {
     expect(mocks.promptAsync).not.toHaveBeenCalled();
   });
 
-  test('continues fallback when abort rejects', async () => {
+  test('falls back to abort+retry when promptAsync fails on busy session', async () => {
     const { client, mocks } = createMockClient({
+      promptAsyncImpl: async () => {
+        throw new Error('session busy');
+      },
       abortImpl: async () => {
-        throw new Error('abort failed');
+        // abort succeeds on first call
       },
     });
     const mgr = new ForegroundFallbackManager(client, makeChains(), true);
@@ -298,13 +303,14 @@ describe('ForegroundFallbackManager session.error', () => {
     await mgr.handleEvent({
       type: 'session.error',
       properties: {
-        sessionID: 'sess-abort-rejects',
+        sessionID: 'sess-busy',
         error: { message: 'Rate limit exceeded' },
       },
     });
 
+    // First promptAsync attempt failed → abort called, then promptAsync retried
     expect(mocks.abort).toHaveBeenCalledTimes(1);
-    expect(mocks.promptAsync).toHaveBeenCalledTimes(1);
+    expect(mocks.promptAsync).toHaveBeenCalledTimes(2);
   });
 });
 
@@ -377,9 +383,8 @@ describe('ForegroundFallbackManager message.updated', () => {
 describe('ForegroundFallbackManager session.status', () => {
   test('triggers fallback on retry status with rate limit message', async () => {
     const { client, mocks } = createMockClient();
-    const mgr = new ForegroundFallbackManager(client, makeChains(), true);
+    const mgr = new ForegroundFallbackManager(client, makeChains(), true, 1);
 
-    // Pre-seed model
     await mgr.handleEvent({
       type: 'message.updated',
       properties: {
@@ -404,7 +409,7 @@ describe('ForegroundFallbackManager session.status', () => {
 
   test('triggers fallback on retry status with insufficient balance message', async () => {
     const { client, mocks } = createMockClient();
-    const mgr = new ForegroundFallbackManager(client, makeChains(), true);
+    const mgr = new ForegroundFallbackManager(client, makeChains(), true, 1);
 
     await mgr.handleEvent({
       type: 'message.updated',
@@ -442,6 +447,62 @@ describe('ForegroundFallbackManager session.status', () => {
 
     expect(mocks.promptAsync).not.toHaveBeenCalled();
   });
+
+  test('tracks retries and only intervenes after maxRetries', async () => {
+    const { client, mocks } = createMockClient();
+    const mgr = new ForegroundFallbackManager(client, makeChains(), true, 3);
+
+    // Pre-seed model
+    await mgr.handleEvent({
+      type: 'message.updated',
+      properties: {
+        info: {
+          sessionID: 'sess-retry',
+          providerID: 'anthropic',
+          modelID: 'claude-opus-4-5',
+        },
+      },
+    });
+
+    // First two retries should be absorbed (maxRetries - 1 = 2)
+    await mgr.handleEvent({
+      type: 'session.status',
+      properties: {
+        sessionID: 'sess-retry',
+        status: {
+          type: 'retry',
+          attempt: 1,
+          message: 'rate limit, retrying...',
+        },
+      },
+    });
+    await mgr.handleEvent({
+      type: 'session.status',
+      properties: {
+        sessionID: 'sess-retry',
+        status: {
+          type: 'retry',
+          attempt: 2,
+          message: 'rate limit, retrying...',
+        },
+      },
+    });
+    expect(mocks.promptAsync).not.toHaveBeenCalled();
+
+    // Third retry exhausts the budget → tryFallback intervenes
+    await mgr.handleEvent({
+      type: 'session.status',
+      properties: {
+        sessionID: 'sess-retry',
+        status: {
+          type: 'retry',
+          attempt: 3,
+          message: 'rate limit, retrying...',
+        },
+      },
+    });
+    expect(mocks.promptAsync).toHaveBeenCalledTimes(1);
+  });
 });
 
 // ---------------------------------------------------------------------------
@@ -480,7 +541,7 @@ describe('ForegroundFallbackManager chain exhaustion', () => {
     expect(mocks.promptAsync).not.toHaveBeenCalled();
   });
 
-  test('does not call promptAsync when all chain models have been tried', async () => {
+  test('aborts when all chain models have been tried', async () => {
     // Scenario: chain = ['anthropic/claude-a', 'openai/gpt-b'].
     // Current model is 'openai/gpt-b' (the last fallback already in use).
     // tried will contain: 'openai/gpt-b' (current) → chain.find() → 'anthropic/claude-a'
@@ -513,6 +574,7 @@ describe('ForegroundFallbackManager chain exhaustion', () => {
 
     // Session B (fresh session, different ID): only model-y is in chain and it IS
     // the current model → tried gets model-y → chain.find() = undefined → exhausted
+    // → abort called to stop the freeze
     const { client: client2, mocks: mocks2 } = createMockClient();
     const mgr2 = new ForegroundFallbackManager(
       client2,
@@ -531,6 +593,7 @@ describe('ForegroundFallbackManager chain exhaustion', () => {
         },
       },
     });
+    expect(mocks2.abort).toHaveBeenCalledTimes(1);
     expect(mocks2.promptAsync).not.toHaveBeenCalled();
   });
 });

+ 79 - 23
src/hooks/foreground-fallback/index.ts

@@ -99,6 +99,9 @@ export class ForegroundFallbackManager {
    *  when the model has changed, allowing the cascade to continue when a
    *  new fallback model also fails within the dedup window. */
   private readonly lastTriggerModel = new Map<string, string>();
+  /** sessionID → consecutive 429 count for the current model.
+   *  Reset on model swap or session deletion. */
+  private readonly sessionRetries = new Map<string, number>();
 
   /** Exposed for task-session-manager: prevents idle reconciliation
    *  while a fallback abort/re-prompt is in flight for this session. */
@@ -115,6 +118,8 @@ export class ForegroundFallbackManager {
      */
     private readonly chains: Record<string, string[]>,
     private readonly enabled: boolean,
+    /** Consecutive 429s tolerated on the same model before swap/abort. */
+    private readonly maxRetries: number = 3,
   ) {}
 
   /**
@@ -150,7 +155,12 @@ export class ForegroundFallbackManager {
         }
         // Rate-limit on an individual message
         if (info.error && isRateLimitError(info.error)) {
-          await this.tryFallback(sessionID);
+          if (this.shouldIntervene(sessionID)) {
+            await this.tryFallback(sessionID);
+          }
+        } else {
+          // Successful response: clear retry count so recovery is not forgotten.
+          this.sessionRetries.delete(sessionID);
         }
         break;
       }
@@ -159,7 +169,12 @@ export class ForegroundFallbackManager {
         const props = event.properties as
           | { sessionID?: string; error?: unknown }
           | undefined;
-        if (props?.sessionID && props.error && isRateLimitError(props.error)) {
+        if (
+          props?.sessionID &&
+          props.error &&
+          isRateLimitError(props.error) &&
+          this.shouldIntervene(props.sessionID)
+        ) {
           await this.tryFallback(props.sessionID);
         }
         break;
@@ -169,15 +184,11 @@ export class ForegroundFallbackManager {
         const props = event.properties as
           | {
               sessionID?: string;
-              status?: { type?: string; message?: string };
+              status?: { type?: string; message?: string; attempt?: number };
             }
           | undefined;
         if (!props?.sessionID || !props.status?.message) break;
         const msg = props.status.message.toLowerCase();
-        // Check for rate-limit signals in the status message regardless of
-        // status type. OpenCode proxies may emit monthly/weekly/5-hour usage
-        // limit errors with type 'error' instead of 'retry' on fresh sessions
-        // where no retry is attempted - the retry-type guard would miss them.
         if (
           msg.includes('rate limit') ||
           msg.includes('usage limit') ||
@@ -189,7 +200,14 @@ export class ForegroundFallbackManager {
           msg.includes('high concurrency') ||
           msg.includes('reduce concurrency')
         ) {
-          await this.tryFallback(props.sessionID);
+          // session.status retry path always counts toward the budget
+          // — even the first retry is absorbed before intervening.
+          if (this.checkRetryBudget(props.sessionID)) {
+            await this.tryFallback(props.sessionID);
+          }
+        } else {
+          // Non-rate-limit status: clear retry count (recovery).
+          this.sessionRetries.delete(props.sessionID);
         }
         break;
       }
@@ -223,12 +241,45 @@ export class ForegroundFallbackManager {
           this.inProgress.delete(id);
           this.lastTrigger.delete(id);
           this.lastTriggerModel.delete(id);
+          this.sessionRetries.delete(id);
         }
         break;
       }
     }
   }
 
+  // ---------------------------------------------------------------------------
+  // Retry budget
+  // ---------------------------------------------------------------------------
+
+  /** Increment retry counter and return true when the budget is exhausted.
+   *  Used by the session.status retry path — each retry counts toward the
+   *  budget and only triggers fallback after maxRetries - 1 absorptions.
+   *  Non-retry paths (session.error / message.updated) use shouldIntervene(),
+   *  which bypasses the counter on first occurrence. */
+  private checkRetryBudget(sessionID: string): boolean {
+    const tried = this.sessionRetries.get(sessionID) ?? 0;
+    if (tried < this.maxRetries - 1) {
+      this.sessionRetries.set(sessionID, tried + 1);
+      log('[foreground-fallback] rate-limit retry', {
+        sessionID,
+        attempt: tried + 1,
+        remaining: this.maxRetries - tried - 1,
+      });
+      return false;
+    }
+    this.sessionRetries.delete(sessionID);
+    return true;
+  }
+
+  /** For non-retry paths (session.error, message.updated): intervene immediately
+   *  unless the session is already in a retry window (has prior retries). */
+  private shouldIntervene(sessionID: string): boolean {
+    const tried = this.sessionRetries.get(sessionID) ?? 0;
+    if (tried === 0) return true;
+    return this.checkRetryBudget(sessionID);
+  }
+
   // ---------------------------------------------------------------------------
   // Core fallback logic
   // ---------------------------------------------------------------------------
@@ -305,15 +356,18 @@ export class ForegroundFallbackManager {
           this.sessionTried.set(sessionID, tried);
           nextModel = stickyFallback;
         } else {
-          log('[foreground-fallback] fallback chain exhausted', {
+          log('[foreground-fallback] fallback chain exhausted, aborting', {
             sessionID,
             agentName,
             tried: [...tried],
           });
+          await abortSessionWithTimeout(this.client, sessionID);
           return;
         }
       }
       tried.add(nextModel);
+      // Reset retry count on model switch — the new model starts fresh.
+      this.sessionRetries.delete(sessionID);
 
       const ref = parseModelReference(nextModel);
       if (!ref) {
@@ -357,25 +411,27 @@ export class ForegroundFallbackManager {
         return;
       }
 
-      // Abort the currently rate-limited prompt so the session becomes idle.
+      // Try queuing the fallback prompt without aborting first. If OpenCode
+      // accepts it (204), the fallback model replaces the retry loop
+      // transparently — no dialog, no session error shown to the user.
+      // If promptAsync throws (e.g. session busy), fall back to abort+retry.
       try {
-        await abortSessionWithTimeout(this.client, sessionID);
-      } catch (error) {
-        // Session may already be idle or abort may be slow; keep fallback best-effort.
-        log('[foreground-fallback] abort did not complete cleanly', {
+        await sessionClient.promptAsync({
+          path: { id: sessionID },
+          body: { parts: lastUser.parts, model: ref },
+        });
+      } catch (_promptErr) {
+        log('[foreground-fallback] promptAsync on busy session, aborting', {
           sessionID,
-          error: error instanceof Error ? error.message : String(error),
+        });
+        await abortSessionWithTimeout(this.client, sessionID);
+        await new Promise((r) => setTimeout(r, REPROMPT_DELAY_MS));
+        await sessionClient.promptAsync({
+          path: { id: sessionID },
+          body: { parts: lastUser.parts, model: ref },
         });
       }
 
-      // Give the server a moment to finalise the abort before re-prompting.
-      await new Promise((r) => setTimeout(r, REPROMPT_DELAY_MS));
-
-      await sessionClient.promptAsync({
-        path: { id: sessionID },
-        body: { parts: lastUser.parts, model: ref },
-      });
-
       this.sessionModel.set(sessionID, nextModel);
       log('[foreground-fallback] switched to fallback model', {
         sessionID,

+ 5 - 3
src/index.ts

@@ -297,12 +297,14 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
     // Initialize JSON parse error recovery hook
     jsonErrorRecoveryHook = createJsonErrorRecoveryHook(ctx);
 
-    // Initialize foreground fallback manager for runtime model switching
+    // Initialize foreground fallback manager for runtime model switching.
+    // Enabled by default even without fallback chains — the manager can still
+    // abort rate-limited sessions after maxRetries to prevent infinite freezes.
     foregroundFallback = new ForegroundFallbackManager(
       ctx.client,
       runtimeChains,
-      config.fallback?.enabled !== false &&
-        Object.keys(runtimeChains).length > 0,
+      config.fallback?.enabled !== false,
+      config.fallback?.maxRetries ?? 3,
     );
 
     deepworkCommandHook = createDeepworkCommandHook();