fix(core): wire token refresh into run shape streams and bound re-minting

Pass the refresh resolver to the runShapeStream factory so subscribeToRun,
subscribeToRunsWithTag and subscribeToBatch recover from an expired token.
Only re-arm the refresh after a connection delivers a record, keep the auth
error terminal when the refresher throws, and don't report a 401 the refresh
recovered from.

Author: claude

packages/core/src/v3/apiClient/index.ts

@@ -299,6 +299,17 @@ export class ApiClient {
     };
   }
 
+  /** As {@link ApiClient.#resolveStreamHeaders}, for the leaner realtime header set. */
+  #resolveRealtimeHeaders(): (() => Promise<Record<string, string>>) | undefined {
+    const refreshAccessToken = this.refreshAccessToken;
+    if (!refreshAccessToken) return undefined;
+
+    return async () => {
+      const accessToken = await refreshAccessTokenOnce(refreshAccessToken);
+      return { ...this.#getRealtimeHeaders(), Authorization: `Bearer ${accessToken}` };
+    };
+  }
+
   async getRunResult(
     runId: string,
     requestOptions?: ZodFetchOptions
@@ -1685,6 +1696,7 @@ export class ApiClient {
         closeOnComplete:
           typeof options?.closeOnComplete === "boolean" ? options.closeOnComplete : true,
         headers: this.#getRealtimeHeaders(),
+        resolveHeaders: this.#resolveRealtimeHeaders(),
         client: this,
         signal: options?.signal,
         onFetchError: options?.onFetchError,
@@ -1708,6 +1720,7 @@ export class ApiClient {
       {
         closeOnComplete: false,
         headers: this.#getRealtimeHeaders(),
+        resolveHeaders: this.#resolveRealtimeHeaders(),
         client: this,
         signal: options?.signal,
         onFetchError: options?.onFetchError,
@@ -1734,6 +1747,7 @@ export class ApiClient {
       {
         closeOnComplete: false,
         headers: this.#getRealtimeHeaders(),
+        resolveHeaders: this.#resolveRealtimeHeaders(),
         client: this,
         signal: options?.signal,
         onFetchError: options?.onFetchError,

packages/core/src/v3/apiClient/refreshAccessToken.test.ts

@@ -0,0 +1,63 @@
+import { describe, expect, it } from "vitest";
+import { refreshAccessTokenOnce } from "./refreshAccessToken.js";
+
+describe("refreshAccessTokenOnce", () => {
+  it("shares one in-flight mint between concurrent callers", async () => {
+    let calls = 0;
+    let release: (token: string) => void = () => {};
+    const refresh = () => {
+      calls++;
+      return new Promise<string>((resolve) => {
+        release = resolve;
+      });
+    };
+
+    const results = Promise.all([
+      refreshAccessTokenOnce(refresh),
+      refreshAccessTokenOnce(refresh),
+      refreshAccessTokenOnce(refresh),
+    ]);
+    release("fresh");
+
+    expect(await results).toEqual(["fresh", "fresh", "fresh"]);
+    expect(calls).toBe(1);
+  });
+
+  it("does not share a mint between different refreshers", async () => {
+    const a = async () => "a";
+    const b = async () => "b";
+
+    expect(await Promise.all([refreshAccessTokenOnce(a), refreshAccessTokenOnce(b)])).toEqual([
+      "a",
+      "b",
+    ]);
+  });
+
+  it("mints again once the previous call has settled", async () => {
+    let calls = 0;
+    const refresh = async () => `token-${++calls}`;
+
+    expect(await refreshAccessTokenOnce(refresh)).toBe("token-1");
+    expect(await refreshAccessTokenOnce(refresh)).toBe("token-2");
+  });
+
+  it("rejects every concurrent caller and does not poison later calls", async () => {
+    let calls = 0;
+    const refresh = async () => {
+      calls++;
+      if (calls === 1) throw new Error("mint failed");
+      return "recovered";
+    };
+
+    const first = refreshAccessTokenOnce(refresh);
+    const second = refreshAccessTokenOnce(refresh);
+
+    await expect(first).rejects.toThrow("mint failed");
+    await expect(second).rejects.toThrow("mint failed");
+    expect(calls).toBe(1);
+
+    // The failed mint was evicted, so the next call tries again.
+    expect(await refreshAccessTokenOnce(refresh)).toBe("recovered");
+    expect(calls).toBe(2);
+  });
+});

packages/core/src/v3/apiClient/runStream.test.ts

@@ -43,6 +43,33 @@ describe("SSEStreamSubscription retry behavior", () => {
     });
   }
 
+  /** An accepted connection that dies before delivering a single record. */
+  function makeDroppedResponse() {
+    const body = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.error(new Error("connection dropped"));
+      },
+    });
+    return new Response(body, {
+      status: 200,
+      headers: { "Content-Type": "text/event-stream", "X-Stream-Version": "v1" },
+    });
+  }
+
+  /** One delivered record, then the connection dies. */
+  function makeChunkThenDropResponse() {
+    const body = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(new TextEncoder().encode(`id: 1\ndata: {"hello":1}\n\n`));
+        setTimeout(() => controller.error(new Error("connection dropped")), 20);
+      },
+    });
+    return new Response(body, {
+      status: 200,
+      headers: { "Content-Type": "text/event-stream", "X-Stream-Version": "v1" },
+    });
+  }
+
   // Drain a ReadableStream<SSEStreamPart> until it closes or errors.
   // Returns received chunks plus terminal state.
   async function drain(stream: ReadableStream<{ id: string; chunk: unknown }>) {
@@ -496,6 +523,127 @@ describe("SSEStreamSubscription retry behavior", () => {
     expect(result.error).toBeDefined();
   });
 
+  it("retries a 403 once with the headers from resolveHeaders", async () => {
+    const seenTokens: Array<string | null> = [];
+    globalThis.fetch = vi.fn().mockImplementation(async (_url: string, init: RequestInit) => {
+      const token = new Headers(init.headers).get("Authorization");
+      seenTokens.push(token);
+      if (token !== "Bearer fresh") return new Response("forbidden", { status: 403 });
+      return makeSSEResponse();
+    });
+
+    const sub = new SSEStreamSubscription("http://example.test/sse", {
+      headers: { Authorization: "Bearer expired" },
+      retryDelayMs: 1,
+      maxRetryDelayMs: 5,
+      resolveHeaders: async () => ({ Authorization: "Bearer fresh" }),
+    });
+
+    const result = await sub.subscribe().then(drain);
+    expect(seenTokens).toEqual(["Bearer expired", "Bearer fresh"]);
+    expect(result.error).toBeUndefined();
+    expect(result.chunks).toHaveLength(1);
+  });
+
+  it("does not report a 401 that the refresh recovered from", async () => {
+    let attempts = 0;
+    globalThis.fetch = vi.fn().mockImplementation(async () => {
+      attempts++;
+      if (attempts === 1) return new Response("unauthorized", { status: 401 });
+      return makeSSEResponse();
+    });
+
+    const errors: Error[] = [];
+    const sub = new SSEStreamSubscription("http://example.test/sse", {
+      headers: { Authorization: "Bearer expired" },
+      retryDelayMs: 1,
+      maxRetryDelayMs: 5,
+      onError: (e) => errors.push(e),
+      resolveHeaders: async () => ({ Authorization: "Bearer fresh" }),
+    });
+
+    const result = await sub.subscribe().then(drain);
+    expect(errors).toHaveLength(0);
+    expect(result.error).toBeUndefined();
+  });
+
+  it("terminates on a 401 when the refresher itself throws", async () => {
+    let attempts = 0;
+    globalThis.fetch = vi.fn().mockImplementation(async () => {
+      attempts++;
+      return new Response("unauthorized", { status: 401 });
+    });
+
+    const errors: Error[] = [];
+    const sub = new SSEStreamSubscription("http://example.test/sse", {
+      headers: { Authorization: "Bearer expired" },
+      retryDelayMs: 1,
+      maxRetryDelayMs: 5,
+      onError: (e) => errors.push(e),
+      resolveHeaders: async () => {
+        throw new Error("mint failed");
+      },
+    });
+
+    const result = await sub.subscribe().then(drain);
+    expect(attempts).toBe(1);
+    expect(errors).toHaveLength(1);
+    expect(result.error).toBeDefined();
+  });
+
+  it("does not re-mint for a connection that is accepted but delivers nothing", async () => {
+    let attempts = 0;
+    globalThis.fetch = vi.fn().mockImplementation(async () => {
+      attempts++;
+      // Accept the second attempt, then drop it without sending a record.
+      if (attempts === 2) return makeDroppedResponse();
+      return new Response("unauthorized", { status: 401 });
+    });
+
+    let refreshes = 0;
+    const sub = new SSEStreamSubscription("http://example.test/sse", {
+      headers: { Authorization: "Bearer expired" },
+      retryDelayMs: 1,
+      maxRetryDelayMs: 5,
+      resolveHeaders: async () => {
+        refreshes++;
+        return { Authorization: `Bearer fresh-${refreshes}` };
+      },
+    });
+
+    const result = await sub.subscribe().then(drain);
+    expect(refreshes).toBe(1);
+    expect(attempts).toBe(3);
+    expect(result.error).toBeDefined();
+  });
+
+  it("allows another refresh once a connection has delivered a record", async () => {
+    let attempts = 0;
+    globalThis.fetch = vi.fn().mockImplementation(async () => {
+      attempts++;
+      if (attempts === 2) return makeChunkThenDropResponse();
+      if (attempts === 4) return makeSSEResponse();
+      return new Response("unauthorized", { status: 401 });
+    });
+
+    let refreshes = 0;
+    const sub = new SSEStreamSubscription("http://example.test/sse", {
+      headers: { Authorization: "Bearer expired" },
+      retryDelayMs: 1,
+      maxRetryDelayMs: 5,
+      resolveHeaders: async () => {
+        refreshes++;
+        return { Authorization: `Bearer fresh-${refreshes}` };
+      },
+    });
+
+    const result = await sub.subscribe().then(drain);
+    expect(refreshes).toBe(2);
+    expect(attempts).toBe(4);
+    expect(result.error).toBeUndefined();
+    expect(result.chunks).toHaveLength(2);
+  });
+
   it("retries on 503 (caller-tunable nonRetryableStatuses)", async () => {
     let attempts = 0;
     globalThis.fetch = vi.fn().mockImplementation(async () => {

packages/core/src/v3/apiClient/runStream.ts

@@ -82,6 +82,7 @@ export type RunStreamCallback<TRunTypes extends AnyRunTypes> = (
 
 export type RunShapeStreamOptions = {
   headers?: Record<string, string>;
+  resolveHeaders?: () => Promise<Record<string, string>>;
   fetchClient?: typeof fetch;
   closeOnComplete?: boolean;
   signal?: AbortSignal;
@@ -114,6 +115,7 @@ export function runShapeStream<TRunTypes extends AnyRunTypes>(
     getEnvVar("TRIGGER_STREAM_URL", getEnvVar("TRIGGER_API_URL")) ?? "https://api.trigger.dev",
     {
       headers: options?.headers,
+      resolveHeaders: options?.resolveHeaders,
       signal: abortController.signal,
     }
   );
@@ -420,12 +422,16 @@ export class SSEStreamSubscription implements StreamSubscription {
           "Could not subscribe to stream",
           Object.fromEntries(response.headers)
         );
-        this.options.onError?.(error);
         if (this.nonRetryableStatuses.has(response.status)) {
+          this.options.onError?.(error);
           controller.error(error);
           return;
         }
-        await this.refreshHeadersForAuthError(response.status);
+        // Only surface the error if we can't recover from it — a 401 that a
+        // token refresh fixes shouldn't reach the consumer.
+        if (!(await this.refreshHeadersForAuthError(response.status))) {
+          this.options.onError?.(error);
+        }
         throw error;
       }
 
@@ -438,7 +444,6 @@ export class SSEStreamSubscription implements StreamSubscription {
       const streamVersion = response.headers.get("X-Stream-Version") ?? "v1";
       this.sessionSettled = response.headers.get("X-Session-Settled") === "true";
       this.retryCount = 0; // reset on success
-      this.authRefreshed = false;
       armStall();
 
       // Dedup window for record ids. Bounded with FIFO eviction so a
@@ -554,6 +559,10 @@ export class SSEStreamSubscription implements StreamSubscription {
           }
 
           armStall(); // any chunk (including server keepalives) resets the silence timer
+          // The connection is genuinely live, so a later expiry may refresh
+          // again. Gated on a delivered chunk: a server that accepts then
+          // immediately drops must not win an unbounded mint loop.
+          this.authRefreshed = false;
           controller.enqueue(value);
         }
       } catch (error) {
@@ -590,16 +599,25 @@ export class SSEStreamSubscription implements StreamSubscription {
 
   /**
    * Re-resolve the headers after a 401/403 so the retry carries a fresh token.
-   * At most once per connection: if the refreshed token is rejected too, the
-   * auth error stays terminal.
+   * At most once per live connection: if the refreshed token is rejected too,
+   * the auth error stays terminal. Returns true when a retry should follow.
    */
-  private async refreshHeadersForAuthError(status: number): Promise<void> {
-    if (status !== 401 && status !== 403) return;
-    if (!this.options.resolveHeaders || this.authRefreshed) return;
+  private async refreshHeadersForAuthError(status: number): Promise<boolean> {
+    if (status !== 401 && status !== 403) return false;
+    if (!this.options.resolveHeaders || this.authRefreshed) return false;
 
     this.authRefreshed = true;
-    this.currentHeaders = await this.options.resolveHeaders();
+
+    try {
+      this.currentHeaders = await this.options.resolveHeaders();
+    } catch {
+      // A refresher that can't mint leaves us with the rejected token —
+      // keep the auth error terminal rather than retrying with it.
+      return false;
+    }
+
     this.retryAfterAuthRefresh = true;
+    return true;
   }
 
   private async retryConnection(