Jelajahi Sumber

Merge pull request #1110 from leducmaxime/fix/foreground-fallback-per-turn-descent

fix(foreground-fallback): reset the descent when a new turn returns to the chain primary
Alvin 2 minggu lalu
induk
melakukan
73716531a9

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

@@ -1578,6 +1578,303 @@ describe('ForegroundFallbackManager session.status', () => {
 // ---------------------------------------------------------------------------
 
 describe('ForegroundFallbackManager chain exhaustion', () => {
+  test('re-walks from the second chain entry on each new user turn', async () => {
+    const { mocks } = createMockClient();
+    const mgr = new ForegroundFallbackManager(makeChains(), true, {
+      directory: '/test',
+    } as any);
+    const sessionID = 'sess-turns';
+
+    await mgr.handleEvent({
+      type: 'message.updated',
+      properties: {
+        info: {
+          sessionID,
+          agent: 'orchestrator',
+          providerID: 'anthropic',
+          modelID: 'claude-opus-4-5',
+          role: 'assistant',
+        },
+      },
+    });
+
+    const realNowFn = Date.now;
+    let fakeNow = realNowFn();
+    Date.now = () => fakeNow;
+    try {
+      fakeNow += 6_000;
+      await mgr.handleEvent({
+        type: 'session.error',
+        properties: { sessionID, error: { message: 'rate limit exceeded' } },
+      });
+      expect(mocks.promptAsync).toHaveBeenCalledWith(
+        expect.objectContaining({
+          body: expect.objectContaining({
+            model: { providerID: 'openai', modelID: 'gpt-4o' },
+          }),
+        }),
+      );
+
+      await mgr.handleEvent({
+        type: 'message.updated',
+        properties: {
+          info: {
+            sessionID,
+            agent: 'orchestrator',
+            role: 'assistant',
+            providerID: 'openai',
+            modelID: 'gpt-4o',
+            time: { created: 1, completed: 2 },
+          },
+        },
+      });
+      await mgr.handleEvent({
+        type: 'message.updated',
+        properties: {
+          info: {
+            sessionID,
+            agent: 'orchestrator',
+            role: 'assistant',
+            providerID: 'anthropic',
+            modelID: 'claude-opus-4-5',
+          },
+        },
+      });
+
+      fakeNow += 6_000;
+      await mgr.handleEvent({
+        type: 'session.error',
+        properties: { sessionID, error: { message: 'rate limit exceeded' } },
+      });
+
+      expect(mocks.promptAsync.mock.calls[1]?.[0]).toEqual(
+        expect.objectContaining({
+          body: expect.objectContaining({
+            model: { providerID: 'openai', modelID: 'gpt-4o' },
+          }),
+        }),
+      );
+    } finally {
+      Date.now = realNowFn;
+    }
+  });
+
+  test('recovers fallback after a chain-exhaustion abort when a new turn returns to the primary', async () => {
+    const { mocks } = createMockClient();
+    const mgr = new ForegroundFallbackManager(
+      { orchestrator: ['openai/gpt-b', 'openai/gpt-c'] },
+      true,
+      { directory: '/test' } as any,
+    );
+    const sessionID = 'sess-recover-after-abort';
+
+    await mgr.handleEvent({
+      type: 'message.updated',
+      properties: {
+        info: {
+          sessionID,
+          agent: 'orchestrator',
+          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,
+            error: { message: 'rate limit exceeded' },
+          },
+        });
+      };
+
+      await fail();
+      await fail();
+      await fail();
+      expect(mocks.promptAsync).toHaveBeenCalledTimes(2);
+      expect(mocks.abort).toHaveBeenCalledTimes(1);
+
+      // Deliberately omit time.completed: this is not a successful response;
+      // recovery must come from the fresh descent reset instead.
+      await mgr.handleEvent({
+        type: 'message.updated',
+        properties: {
+          info: {
+            sessionID,
+            agent: 'orchestrator',
+            providerID: 'openai',
+            modelID: 'gpt-b',
+            role: 'assistant',
+          },
+        },
+      });
+
+      await fail();
+      expect(mocks.promptAsync).toHaveBeenCalledTimes(3);
+      expect(mocks.promptAsync.mock.calls[2]?.[0]).toEqual(
+        expect.objectContaining({
+          body: expect.objectContaining({
+            model: { providerID: 'openai', modelID: 'gpt-c' },
+          }),
+        }),
+      );
+      expect(mgr.willAttemptFallback(sessionID)).toBe(true);
+    } finally {
+      Date.now = realNowFn;
+    }
+  });
+
+  test('does not fall back onto an earlier chain entry', async () => {
+    const { mocks } = createMockClient();
+    const mgr = new ForegroundFallbackManager(makeChains(), true, {
+      directory: '/test',
+    } as any);
+    const sessionID = 'sess-mid-chain';
+
+    await mgr.handleEvent({
+      type: 'message.updated',
+      properties: {
+        info: {
+          sessionID,
+          agent: 'orchestrator',
+          providerID: 'openai',
+          modelID: 'gpt-4o',
+          role: 'assistant',
+        },
+      },
+    });
+    await mgr.handleEvent({
+      type: 'session.error',
+      properties: { sessionID, error: { message: 'rate limit exceeded' } },
+    });
+
+    expect(mocks.promptAsync.mock.calls[0]?.[0]).toEqual(
+      expect.objectContaining({
+        body: expect.objectContaining({
+          model: { providerID: 'google', modelID: 'gemini-2.5-pro' },
+        }),
+      }),
+    );
+    expect(mocks.promptAsync.mock.calls[0]?.[0].body.model).not.toEqual({
+      providerID: 'anthropic',
+      modelID: 'claude-opus-4-5',
+    });
+  });
+
+  test('does not fall back onto the primary when the current model is off-chain', async () => {
+    const { mocks } = createMockClient();
+    const mgr = new ForegroundFallbackManager(makeChains(), true, {
+      directory: '/test',
+    } as any);
+    const sessionID = 'sess-off-chain';
+
+    await mgr.handleEvent({
+      type: 'message.updated',
+      properties: {
+        info: {
+          sessionID,
+          agent: 'orchestrator',
+          providerID: 'anthropic',
+          modelID: 'claude-opus-4-5',
+          role: 'assistant',
+        },
+      },
+    });
+
+    const realNowFn = Date.now;
+    let fakeNow = realNowFn();
+    Date.now = () => fakeNow;
+    try {
+      fakeNow += 6_000;
+      await mgr.handleEvent({
+        type: 'session.error',
+        properties: { sessionID, error: { message: 'rate limit exceeded' } },
+      });
+      await mgr.handleEvent({
+        type: 'message.updated',
+        properties: {
+          info: {
+            sessionID,
+            agent: 'orchestrator',
+            providerID: 'openai',
+            modelID: 'gpt-4o-mini',
+            role: 'assistant',
+          },
+        },
+      });
+
+      fakeNow += 6_000;
+      await mgr.handleEvent({
+        type: 'session.error',
+        properties: { sessionID, error: { message: 'rate limit exceeded' } },
+      });
+
+      expect(mocks.promptAsync.mock.calls[1]?.[0]).toEqual(
+        expect.objectContaining({
+          body: expect.objectContaining({
+            model: { providerID: 'google', modelID: 'gemini-2.5-pro' },
+          }),
+        }),
+      );
+      expect(mocks.promptAsync.mock.calls[1]?.[0].body.model).not.toEqual({
+        providerID: 'anthropic',
+        modelID: 'claude-opus-4-5',
+      });
+    } finally {
+      Date.now = realNowFn;
+    }
+  });
+
+  test('does not reset the descent when the current model was inferred, not observed', async () => {
+    createMockClient({ messagesData: [] });
+    const mgr = new ForegroundFallbackManager(
+      { orchestrator: ['a/1', 'b/2', 'c/3'] },
+      true,
+      { directory: '/test' } as any,
+    );
+    const sessionID = 'sess-inferred-model';
+
+    await mgr.handleEvent({
+      type: 'subagent.session.created',
+      properties: { sessionID, agentName: 'orchestrator' },
+    });
+
+    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,
+            error: { message: 'rate limit exceeded' },
+          },
+        });
+      };
+
+      await fail();
+      await fail();
+
+      expect([...(mgr as any).sessionTried.get(sessionID)]).toEqual([
+        'a/1',
+        'b/2',
+        'c/3',
+      ]);
+    } finally {
+      Date.now = realNowFn;
+    }
+  });
+
   test('does not call promptAsync when the only chain model is already the current model', async () => {
     // Scenario: chain = ['openai/gpt-b'], current model IS 'openai/gpt-b'.
     // tryFallback adds 'openai/gpt-b' to tried → chain.find() returns undefined → exhausted.
@@ -1853,6 +2150,8 @@ describe('ForegroundFallbackManager chain exhaustion', () => {
     }
   });
 
+  // Protects the tried.size > 1 invariant in execFallback: a single-model
+  // chain must not re-abort repeatedly after exhaustion.
   test('does not abort repeatedly for single-model chains after exhaustion', async () => {
     const { mocks } = createMockClient();
     const mgr = new ForegroundFallbackManager(

+ 40 - 5
src/hooks/foreground-fallback/index.ts

@@ -663,11 +663,8 @@ export class ForegroundFallbackManager {
   ): Promise<void> {
     const session = getClient(this.input).session;
     try {
-      // After the chain has been exhausted twice (reset retry failed and we
-      // aborted), do not intervene again for this session: re-entering would
-      // keep aborting in a loop. Surface errors to the user instead.
-      if (this.chainExhaustion.get(sessionID) === 2) return;
-      let currentModel = this.sessionModel.get(sessionID);
+      const observedModel = this.sessionModel.get(sessionID);
+      let currentModel = observedModel;
       const agentName = this.sessionAgent.get(sessionID);
       const chain = this.resolveChain(agentName, currentModel);
       // Callers pre-check via hasFallbackChain; keep as defensive guard only.
@@ -687,6 +684,44 @@ export class ForegroundFallbackManager {
       }
       // biome-ignore lint/style/noNonNullAssertion: We just set this above
       let tried = this.sessionTried.get(sessionID)!;
+
+      // A new user turn always re-sends the agent's configured primary:
+      // promptAsync's `model` is a per-message override, so a fallback never
+      // persists past the message it was applied to. Landing here on chain[0]
+      // with a tried set that already walked past it therefore means the
+      // previous descent has ended and its state is stale. Without this the
+      // next descent resumes one link deeper every turn (link 2, then 3, then
+      // 4...) until the chain is spent and the session aborts, instead of
+      // re-walking from link 2 each turn.
+      //
+      // This does not weaken the backward-fallback guard below: currentModel
+      // is re-added immediately after, so chain[0] still can never be picked.
+      // Only an OBSERVED chain[0] counts. execFallback infers
+      // `currentModel = chain[0]` above when no model was ever captured for
+      // this session, which is the opposite situation — resetting there would
+      // re-pick chain[1] on every error instead of descending.
+      // size > 1 means a previous descent actually selected a fallback
+      // (tried.add(nextModel) below), so there is stale state to clear. A
+      // single-entry chain never gets there and must stay terminal after its
+      // one abort rather than re-aborting on every error.
+      if (
+        observedModel !== undefined &&
+        observedModel === chain[0] &&
+        tried.size > 1
+      ) {
+        tried = new Set();
+        this.sessionTried.set(sessionID, tried);
+        // A descent that ended in a stage-2 abort is never followed by a
+        // successful assistant message, so the message.updated recovery path
+        // cannot clear chainExhaustion and fallback would stay disabled for
+        // the rest of the session. A fresh descent earns a fresh chance.
+        this.chainExhaustion.delete(sessionID);
+      }
+
+      // After the chain has been exhausted twice (reset retry failed and we
+      // aborted), do not intervene again for this session: re-entering would
+      // keep aborting in a loop. Surface errors to the user instead.
+      if (this.chainExhaustion.get(sessionID) === 2) return;
       if (currentModel) tried.add(currentModel);
       // ponytail: seed chain entries at or before the current model's index
       // to prevent backward fallback onto models the session already left.