fix: preserve pi assistant fallback context

altaywtf · altaywtf · commit 36f648dcdc23 · 2026-05-25T13:31:54.000+03:00
diff --git a/src/agents/pi-embedded-runner/run.cross-provider-fallback-error-context.test.ts b/src/agents/pi-embedded-runner/run.cross-provider-fallback-error-context.test.ts
@@ -174,6 +174,53 @@ describe("runEmbeddedPiAgent cross-provider fallback error handling", () => {
     });
   });
 
+  it("keeps PI-stamped session assistant errors for the current candidate after compaction", async () => {
+    const getLastFormattedAssistant = captureFormattedAssistant();
+    const sameCandidateErrorMessage = "429 current PI-stamped candidate rate limit";
+    mockedIsFailoverAssistantError.mockImplementation((...args: unknown[]) => {
+      const assistant = args[0];
+      return isCurrentAttemptAssistant(assistant) && assistant.provider === "pi";
+    });
+    mockedIsRateLimitAssistantError.mockImplementation((...args: unknown[]) => {
+      const assistant = args[0];
+      return isCurrentAttemptAssistant(assistant) && assistant.provider === "pi";
+    });
+    mockedRunEmbeddedAttempt.mockResolvedValueOnce(
+      makeAttemptResult({
+        assistantTexts: [],
+        lastAssistant: makeAssistantMessageFixture({
+          stopReason: "error",
+          errorMessage: sameCandidateErrorMessage,
+          provider: "pi",
+          model: "pi",
+          content: [],
+        }),
+        currentAttemptAssistant: undefined,
+      }),
+    );
+
+    const promise = runEmbeddedPiAgent({
+      ...overflowBaseRunParams,
+      runId: "run-compaction-pi-stamped-fallback-error-context",
+      config: makeCrossProviderFallbackConfig(),
+    });
+
+    await expect(promise).rejects.toBeInstanceOf(MockedFailoverError);
+    await expect(promise).rejects.toThrow(sameCandidateErrorMessage);
+    expect(mockedIsRateLimitAssistantError).toHaveBeenCalledTimes(1);
+    const rateLimitCalls = mockedIsRateLimitAssistantError.mock.calls as unknown[][];
+    expect(rateLimitCalls.at(-1)?.[0]).toMatchObject({
+      provider: "pi",
+      model: "pi",
+      errorMessage: sameCandidateErrorMessage,
+    });
+    expect(getLastFormattedAssistant()).toMatchObject({
+      provider: "pi",
+      model: "pi",
+      errorMessage: sameCandidateErrorMessage,
+    });
+  });
+
   it("does not reuse a prior provider session assistant when the current candidate times out", async () => {
     const getLastFormattedAssistant = captureFormattedAssistant();
     mockedRunEmbeddedAttempt.mockResolvedValueOnce(
diff --git a/src/agents/pi-embedded-runner/run/helpers.ts b/src/agents/pi-embedded-runner/run/helpers.ts
@@ -105,7 +105,14 @@ export function isAssistantForModelRef(
   assistant: { provider?: string; model?: string } | undefined,
   ref: { provider: string; model: string },
 ): boolean {
-  return assistant?.provider?.trim() === ref.provider && assistant?.model?.trim() === ref.model;
+  if (!assistant) {
+    return false;
+  }
+  const resolved = resolveReportedModelRef({
+    ...ref,
+    assistant,
+  });
+  return resolved.provider === ref.provider && resolved.model === ref.model;
 }
 
 function isEmbeddedHarnessProvider(provider: string): boolean {
