Browse Source

fix(foreground-fallback): require completed recovery

Alvin Unreal 1 week ago
parent
commit
6ba74c4317
2 changed files with 76 additions and 19 deletions
  1. 64 0
      src/hooks/foreground-fallback/index.test.ts
  2. 12 19
      src/hooks/foreground-fallback/index.ts

+ 64 - 0
src/hooks/foreground-fallback/index.test.ts

@@ -1429,6 +1429,7 @@ describe('ForegroundFallbackManager chain exhaustion', () => {
             providerID: 'google',
             modelID: 'gemini-2.5-pro',
             role: 'assistant',
+            time: { created: 1, completed: 2 },
           },
         },
       });
@@ -1442,6 +1443,69 @@ describe('ForegroundFallbackManager chain exhaustion', () => {
     }
   });
 
+  test('does not recover from an incomplete assistant message', async () => {
+    const { mocks } = createMockClient();
+    const mgr = new ForegroundFallbackManager(
+      { orchestrator: ['openai/gpt-b', 'openai/gpt-c'] },
+      true,
+      { directory: '/test' } as any,
+    );
+
+    await mgr.handleEvent({
+      type: 'message.updated',
+      properties: {
+        info: {
+          sessionID: 'sess-incomplete-recovery',
+          providerID: 'openai',
+          modelID: 'gpt-b',
+          role: 'assistant',
+        },
+      },
+    });
+
+    const realNowFn = Date.now;
+    let fakeNow = realNowFn();
+    Date.now = () => fakeNow;
+    try {
+      const fail = async () => {
+        fakeNow += 6_000;
+        await mgr.handleEvent({
+          type: 'session.error',
+          properties: {
+            sessionID: 'sess-incomplete-recovery',
+            error: { message: 'Rate limit exceeded' },
+          },
+        });
+      };
+
+      // Reach stage 1: gpt-b → gpt-c, then the sticky gpt-c retry.
+      await fail();
+      await fail();
+      expect(mocks.promptAsync).toHaveBeenCalledTimes(2);
+
+      // A streaming assistant update is not proof of recovery.
+      await mgr.handleEvent({
+        type: 'message.updated',
+        properties: {
+          info: {
+            sessionID: 'sess-incomplete-recovery',
+            providerID: 'openai',
+            modelID: 'gpt-c',
+            role: 'assistant',
+            time: { created: 1 },
+          },
+        },
+      });
+
+      // Stage 1 remains terminal on the next exhaustion: abort, no third prompt.
+      await fail();
+      expect(mocks.promptAsync).toHaveBeenCalledTimes(2);
+      expect(mocks.abort).toHaveBeenCalledTimes(1);
+    } finally {
+      Date.now = realNowFn;
+    }
+  });
+
   test('does not abort repeatedly for single-model chains after exhaustion', async () => {
     const { mocks } = createMockClient();
     const mgr = new ForegroundFallbackManager(

+ 12 - 19
src/hooks/foreground-fallback/index.ts

@@ -332,13 +332,21 @@ export class ForegroundFallbackManager {
             `${info.providerID}/${info.modelID}`,
           );
         }
+        const messageTime = info.time;
+        const isCompletedSuccessfulAssistant =
+          info.role === 'assistant' &&
+          !info.error &&
+          typeof messageTime === 'object' &&
+          messageTime !== null &&
+          'completed' in messageTime &&
+          typeof messageTime.completed === 'number';
         // Failover-worthy error on an individual message
         if (info.error && isFailoverError(info.error)) {
           if (this.shouldTriggerFallback(sessionID)) {
             await this.tryFallback(sessionID);
           }
-        } else {
-          // Successful response: clear retry count so recovery is not forgotten.
+        } else if (isCompletedSuccessfulAssistant) {
+          // Only a completed, successful assistant response proves recovery.
           this.sessionRetries.delete(sessionID);
           this.chainExhaustion.delete(sessionID);
         }
@@ -410,18 +418,13 @@ export class ForegroundFallbackManager {
           break;
         }
 
-        if (this.isRecoveredStatus(props.status?.type)) {
-          // Recovered/terminal status: clear retry count.
-          this.sessionRetries.delete(sessionID);
-          this.chainExhaustion.delete(sessionID);
-        }
         // Note: do NOT clear sessionRetries here on non-rate-limit statuses.
         // Abort events triggered by our own fallback carry non-rate-limit
         // messages and would reset the counter, creating an infinite loop:
         // abort → fallback → set retries to 1 → abort event clears retries
         // → next retry sees tried=0 → abort+fallback again → repeat.
-        // Retries are only cleared on successful response (message.updated
-        // without error) or session deletion.
+        // Retries are only cleared on a completed successful assistant
+        // response or session deletion.
         break;
       }
 
@@ -482,16 +485,6 @@ export class ForegroundFallbackManager {
     return this.consumeRetryBudget(sessionID);
   }
 
-  private isRecoveredStatus(statusType: string | undefined): boolean {
-    return (
-      statusType === 'idle' ||
-      statusType === 'complete' ||
-      statusType === 'completed' ||
-      statusType === 'success' ||
-      statusType === 'terminal'
-    );
-  }
-
   // ---------------------------------------------------------------------------
   // Core fallback logic
   // ---------------------------------------------------------------------------