Revert "Revert "[wrangler] fix: prevent remote binding sessions expiring during long dev sessions"" (#13470)

Co-authored-by: ask-bonk[bot] <ask-bonk[bot]@users.noreply.github.com>

Author: penalosa

.changeset/fix-remote-bindings-timeout.md

@@ -0,0 +1,12 @@
+---
+"wrangler": patch
+"miniflare": patch
+---
+
+fix: prevent remote binding sessions from expiring during long-running dev sessions
+
+Preview tokens for remote bindings expire after one hour. Previously, the first request after expiry would fail before a refresh was triggered. This change proactively refreshes the token at 50 minutes so no request ever sees an expired session.
+
+The reactive recovery path is also improved: `error code: 1031` responses (returned by bindings such as Workers AI when their session times out) now correctly trigger a refresh, where previously only `Invalid Workers Preview configuration` HTML responses did.
+
+Auth credentials are now resolved lazily when a remote proxy session starts rather than at bundle-complete time. This means that if your OAuth access token has been refreshed since `wrangler dev` started, the new token is used rather than the one captured at startup.

packages/wrangler/e2e/start-worker-auth-opts.test.ts

@@ -84,7 +84,7 @@ describe("startWorker - auth options", { sequential: true }, () => {
 
 			await assertValidWorkerAiResponse(expect);
 
-			expect(validAuth).toHaveBeenCalledOnce();
+			expect(validAuth).toHaveBeenCalled();
 
 			consoleErrorMock.mockReset();
 
@@ -105,7 +105,7 @@ describe("startWorker - auth options", { sequential: true }, () => {
 
 			await assertInvalidWorkerAiResponse(expect);
 
-			expect(incorrectAuth).toHaveBeenCalledOnce();
+			expect(incorrectAuth).toHaveBeenCalled();
 		});
 
 		test("starting a worker with startWorker with invalid auth information and updating it with valid auth information", async ({
@@ -139,7 +139,7 @@ describe("startWorker - auth options", { sequential: true }, () => {
 
 			await assertInvalidWorkerAiResponse(expect);
 
-			expect(incorrectAuth).toHaveBeenCalledOnce();
+			expect(incorrectAuth).toHaveBeenCalled();
 
 			consoleErrorMock.mockReset();
 
@@ -162,7 +162,7 @@ describe("startWorker - auth options", { sequential: true }, () => {
 
 			await assertValidWorkerAiResponse(expect);
 
-			expect(validAuth).toHaveBeenCalledOnce();
+			expect(validAuth).toHaveBeenCalled();
 		});
 
 		async function assertValidWorkerAiResponse(expect: ExpectStatic) {

packages/wrangler/src/__tests__/api/startDevWorker/RemoteRuntimeController.test.ts

@@ -1,6 +1,5 @@
-import { beforeEach, describe, it, vi } from "vitest";
+import { afterEach, beforeEach, describe, it, vi } from "vitest";
 import { RemoteRuntimeController } from "../../../api/startDevWorker/RemoteRuntimeController";
-import { unwrapHook } from "../../../api/startDevWorker/utils";
 // Import the mocked functions so we can set their behavior
 import {
 	createPreviewSession,
@@ -38,10 +37,6 @@ vi.mock("../../../user/access", () => ({
 	domainUsesAccess: vi.fn(),
 }));
 
-vi.mock("../../../api/startDevWorker/utils", () => ({
-	unwrapHook: vi.fn(),
-}));
-
 function makeConfig(
 	overrides: Partial<StartDevWorkerOptions> = {}
 ): StartDevWorkerOptions {
@@ -103,12 +98,6 @@ describe("RemoteRuntimeController", () => {
 	}
 
 	beforeEach(() => {
-		// Setup mock implementations
-		vi.mocked(unwrapHook).mockResolvedValue({
-			accountId: "test-account-id",
-			apiToken: { apiToken: "test-token" },
-		});
-
 		vi.mocked(getWorkerAccountAndContext).mockResolvedValue({
 			workerAccount: {
 				accountId: "test-account-id",
@@ -159,12 +148,106 @@ describe("RemoteRuntimeController", () => {
 		vi.mocked(createWorkerPreview).mockResolvedValue({
 			value: "test-preview-token",
 			host: "test.workers.dev",
-			tailUrl: "wss://test.workers.dev/tail",
+			// No tailUrl — avoids real WebSocket connections in unit tests
 		});
 
 		vi.mocked(getAccessHeaders).mockResolvedValue({});
 	});
 
+	describe("proactive token refresh", () => {
+		afterEach(() => vi.useRealTimers());
+
+		it("should proactively refresh the token before expiry", async ({
+			expect,
+		}) => {
+			vi.useFakeTimers();
+
+			const { controller, bus } = setup();
+			const config = makeConfig();
+			const bundle = makeBundle();
+
+			controller.onBundleStart({ type: "bundleStart", config });
+			controller.onBundleComplete({ type: "bundleComplete", config, bundle });
+			await bus.waitFor("reloadComplete");
+
+			vi.mocked(createWorkerPreview).mockClear();
+			vi.mocked(createRemoteWorkerInit).mockClear();
+			vi.mocked(createWorkerPreview).mockResolvedValue({
+				value: "proactively-refreshed-token",
+				host: "test.workers.dev",
+			});
+
+			// Register the waiter before advancing so it's in place when the
+			// event fires. Use a timeout larger than the advance window so the
+			// waiter's own faked setTimeout doesn't race the refresh timer.
+			const reloadPromise = bus.waitFor(
+				"reloadComplete",
+				undefined,
+				60 * 60 * 1000
+			);
+			await vi.advanceTimersByTimeAsync(50 * 60 * 1000 + 1);
+			const reloadEvent = await reloadPromise;
+
+			expect(createWorkerPreview).toHaveBeenCalledTimes(1);
+			expect(reloadEvent).toMatchObject({
+				type: "reloadComplete",
+				proxyData: {
+					headers: {
+						"cf-workers-preview-token": "proactively-refreshed-token",
+					},
+				},
+			});
+		});
+
+		it("should cancel the proactive refresh timer on bundle start", async ({
+			expect,
+		}) => {
+			vi.useFakeTimers();
+
+			const { controller, bus } = setup();
+			const config = makeConfig();
+			const bundle = makeBundle();
+
+			controller.onBundleStart({ type: "bundleStart", config });
+			controller.onBundleComplete({ type: "bundleComplete", config, bundle });
+			await bus.waitFor("reloadComplete");
+
+			vi.mocked(createWorkerPreview).mockClear();
+
+			// A new bundleStart cancels the old timer before it fires
+			controller.onBundleStart({ type: "bundleStart", config });
+			controller.onBundleComplete({ type: "bundleComplete", config, bundle });
+			await bus.waitFor("reloadComplete");
+
+			vi.mocked(createWorkerPreview).mockClear();
+
+			// Advance to just before T2 would fire — no proactive refresh should occur
+			await vi.advanceTimersByTimeAsync(50 * 60 * 1000 - 1);
+			expect(createWorkerPreview).not.toHaveBeenCalled();
+		});
+
+		it("should cancel the proactive refresh timer on teardown", async ({
+			expect,
+		}) => {
+			vi.useFakeTimers();
+
+			const { controller, bus } = setup();
+			const config = makeConfig();
+			const bundle = makeBundle();
+
+			controller.onBundleStart({ type: "bundleStart", config });
+			controller.onBundleComplete({ type: "bundleComplete", config, bundle });
+			await bus.waitFor("reloadComplete");
+
+			vi.mocked(createWorkerPreview).mockClear();
+			await controller.teardown();
+
+			// Advance past where the timer would have fired
+			await vi.advanceTimersByTimeAsync(50 * 60 * 1000 + 1);
+			expect(createWorkerPreview).not.toHaveBeenCalled();
+		});
+	});
+
 	describe("preview token refresh", () => {
 		it("should handle missing state gracefully", async ({ expect }) => {
 			const { controller } = setup();

packages/wrangler/src/__tests__/dev/remote-bindings.test.ts

@@ -1,4 +1,5 @@
 /* eslint-disable @typescript-eslint/consistent-type-imports */
+import assert from "node:assert";
 import { seed } from "@cloudflare/workers-utils/test-helpers";
 import { fetch } from "undici";
 /* eslint-disable no-restricted-imports */
@@ -13,6 +14,7 @@ import {
 } from "vitest";
 /* eslint-enable no-restricted-imports */
 import { Binding, StartRemoteProxySessionOptions } from "../../api";
+import { unwrapHook } from "../../api/startDevWorker/utils";
 import { mockAccountId, mockApiToken } from "../helpers/mock-account-id";
 import { mockConsoleMethods } from "../helpers/mock-console";
 import {
@@ -714,16 +716,18 @@ describe("dev with remote bindings", { sequential: true, retry: 2 }, () => {
 		await vi.waitFor(() => expect(std.out).toMatch(/Ready/), {
 			timeout: 5_000,
 		});
-		expect(sessionOptions).toEqual({
-			auth: {
-				accountId: "some-account-id",
-				apiToken: {
-					apiToken: "some-api-token",
-				},
-			},
+		expect(sessionOptions).toBeDefined();
+		assert(sessionOptions);
+		const { auth, ...rest1 } = sessionOptions;
+		expect(rest1).toEqual({
 			complianceRegion: undefined,
 			workerName: "worker",
 		});
+		assert(auth);
+		expect(await unwrapHook(auth, { account_id: undefined })).toEqual({
+			accountId: "some-account-id",
+			apiToken: { apiToken: "some-api-token" },
+		});
 		await stopWrangler();
 		await wranglerStopped;
 	});
@@ -756,16 +760,18 @@ describe("dev with remote bindings", { sequential: true, retry: 2 }, () => {
 			timeout: 5_000,
 		});
 
-		expect(sessionOptions).toEqual({
-			auth: {
-				accountId: "mock-account-id",
-				apiToken: {
-					apiToken: "some-api-token",
-				},
-			},
+		expect(sessionOptions).toBeDefined();
+		assert(sessionOptions);
+		const { auth: auth2, ...rest2 } = sessionOptions;
+		expect(rest2).toEqual({
 			complianceRegion: undefined,
 			workerName: "worker",
 		});
+		assert(auth2);
+		expect(await unwrapHook(auth2, { account_id: undefined })).toEqual({
+			accountId: "mock-account-id",
+			apiToken: { apiToken: "some-api-token" },
+		});
 
 		await stopWrangler();
 

packages/wrangler/src/api/remoteBindings/index.ts

@@ -82,9 +82,9 @@ export async function maybeStartOrUpdateRemoteProxySession(
 	preExistingRemoteProxySessionData?: {
 		session: RemoteProxySession;
 		remoteBindings: Record<string, Binding>;
-		auth?: CfAccount | undefined;
+		auth?: AsyncHook<CfAccount> | undefined;
 	} | null,
-	auth?: CfAccount | undefined
+	auth?: AsyncHook<CfAccount> | undefined
 ): Promise<{
 	session: RemoteProxySession;
 	remoteBindings: Record<string, Binding>;
@@ -188,9 +188,9 @@ export async function maybeStartOrUpdateRemoteProxySession(
  * @returns the auth hook to pass to the startRemoteProxy session function if any
  */
 function getAuthHook(
-	auth: CfAccount | undefined,
+	auth: AsyncHook<CfAccount> | undefined,
 	config: Pick<Config, "account_id"> | undefined
-): AsyncHook<CfAccount, [Pick<Config, "account_id">]> | undefined {
+): AsyncHook<CfAccount> | undefined {
 	if (auth) {
 		return auth;
 	}

packages/wrangler/src/api/startDevWorker/LocalRuntimeController.ts

@@ -14,7 +14,7 @@ import * as MF from "../../dev/miniflare";
 import { logger } from "../../logger";
 import { RuntimeController } from "./BaseController";
 import { castErrorCause } from "./events";
-import { getBinaryFileContents, unwrapHook } from "./utils";
+import { getBinaryFileContents } from "./utils";
 import type { RemoteProxySession } from "../remoteBindings";
 import type {
 	BundleCompleteEvent,
@@ -209,12 +209,6 @@ export class LocalRuntimeController extends RuntimeController {
 
 				const remoteBindings = pickRemoteBindings(configBundle.bindings ?? {});
 
-				const auth =
-					Object.keys(remoteBindings).length === 0
-						? // If there are no remote bindings (this is a local only session) there's no need to get auth data
-							undefined
-						: await unwrapHook(data.config.dev.auth);
-
 				this.#remoteProxySessionData =
 					await maybeStartOrUpdateRemoteProxySession(
 						{
@@ -223,7 +217,10 @@ export class LocalRuntimeController extends RuntimeController {
 							bindings: remoteBindings,
 						},
 						this.#remoteProxySessionData ?? null,
-						auth
+						Object.keys(remoteBindings).length === 0
+							? // If there are no remote bindings (this is a local only session) there's no need to get auth data
+								undefined
+							: data.config.dev.auth
 					);
 			}