From 73bbb584ab789bcec6f8a4d9ff90e1fad2ba8022 Mon Sep 17 00:00:00 2001 From: Takis Kakalis <80459599+Takaros999@users.noreply.github.com> Date: Wed, 23 Sep 2026 19:48:19 -0700 Subject: [PATCH] fix(idkit): retry transient bridge polls --- .../core/src/__tests__/debug-report.test.ts | 57 ++++++++++ js/packages/core/src/request.ts | 58 +++++++++- .../react/src/__tests__/hooks.test.tsx | 105 ++++++++++++++++++ js/packages/react/src/hooks/common.ts | 26 +++++ js/packages/react/src/hooks/useIDKitFlow.ts | 36 +++++- .../react/src/hooks/useIDKitInviteCodeFlow.ts | 45 ++++++-- rust/core/src/bridge.rs | 32 +++++- 7 files changed, 337 insertions(+), 22 deletions(-) diff --git a/js/packages/core/src/__tests__/debug-report.test.ts b/js/packages/core/src/__tests__/debug-report.test.ts index cb4d1e16..01d804b2 100644 --- a/js/packages/core/src/__tests__/debug-report.test.ts +++ b/js/packages/core/src/__tests__/debug-report.test.ts @@ -118,4 +118,61 @@ describe("debug reports", () => { }); expect(wasmRequestMock.getDebugReport).toHaveBeenCalledTimes(1); }); + + it("retries a rejected bridge poll until a terminal status arrives", async () => { + requestMock.mockReturnValue(wasmBuilderMock); + wasmBuilderMock.preset.mockResolvedValue(wasmRequestMock); + wasmRequestMock.pollForStatus + .mockRejectedValueOnce(new Error("Failed to fetch")) + .mockResolvedValueOnce({ type: "confirmed", result: { proof: "ok" } }); + + const request = await IDKit.request({ + app_id: "app_test", + action: "test-action", + rp_context: { + rp_id: "rp_test", + nonce: "0x01", + created_at: 1, + expires_at: 2, + signature: "0x1234", + }, + allow_legacy_proofs: true, + }).preset(orbLegacy()); + + expect( + await request.pollUntilCompletion({ pollInterval: 0, timeout: 1000 }), + ).toEqual({ + success: true, + result: { proof: "ok" }, + }); + expect(wasmRequestMock.pollForStatus).toHaveBeenCalledTimes(2); + }); + + it("times out when a bridge poll never settles", async () => { + requestMock.mockReturnValue(wasmBuilderMock); + wasmBuilderMock.preset.mockResolvedValue(wasmRequestMock); + wasmRequestMock.pollForStatus.mockImplementationOnce( + () => new Promise(() => {}), + ); + + const request = await IDKit.request({ + app_id: "app_test", + action: "test-action", + rp_context: { + rp_id: "rp_test", + nonce: "0x01", + created_at: 1, + expires_at: 2, + signature: "0x1234", + }, + allow_legacy_proofs: true, + }).preset(orbLegacy()); + + expect( + await request.pollUntilCompletion({ pollInterval: 0, timeout: 20 }), + ).toEqual({ + success: false, + error: IDKitErrorCodes.Timeout, + }); + }); }); diff --git a/js/packages/core/src/request.ts b/js/packages/core/src/request.ts index e84bb8eb..87831082 100644 --- a/js/packages/core/src/request.ts +++ b/js/packages/core/src/request.ts @@ -88,6 +88,34 @@ export interface IDKitRequest { getDebugReport(): IDKitDebugReport; } +async function pollBeforeDeadline( + pollOnce: () => Promise, + deadline: number, + signal?: AbortSignal, +): Promise { + let timer: ReturnType | undefined; + let onAbort: (() => void) | undefined; + + try { + return await Promise.race([ + Promise.resolve() + .then(pollOnce) + .catch(() => null), + new Promise((resolve) => { + timer = setTimeout(resolve, Math.max(0, deadline - Date.now()), null); + if (signal) { + onAbort = () => resolve(null); + signal.addEventListener("abort", onAbort, { once: true }); + if (signal.aborted) onAbort(); + } + }), + ]); + } finally { + clearTimeout(timer); + if (onAbort && signal) signal.removeEventListener("abort", onAbort); + } +} + /** * Shared poll loop. Used by both URL-mode and invite-code-mode request impls; * the loop body is identical between the two paths because the bridge @@ -99,24 +127,37 @@ async function pollUntilCompletionLoop( ): Promise { const pollInterval = options?.pollInterval ?? 1000; const timeout = options?.timeout ?? 900_000; // 15 minutes default - const startTime = Date.now(); + const deadline = Date.now() + timeout; while (true) { if (options?.signal?.aborted) { return { success: false, error: IDKitErrorCodes.Cancelled }; } - if (Date.now() - startTime > timeout) { + if (Date.now() >= deadline) { return { success: false, error: IDKitErrorCodes.Timeout }; } - const status = await pollOnce(); + // A backgrounded browser can lose an individual network request. Keep the + // flow alive until the bridge reports a terminal status or time runs out. + const status = await pollBeforeDeadline( + pollOnce, + deadline, + options?.signal, + ); - if (status.type === "confirmed" && status.result) { + if (options?.signal?.aborted) { + return { success: false, error: IDKitErrorCodes.Cancelled }; + } + if (Date.now() >= deadline) { + return { success: false, error: IDKitErrorCodes.Timeout }; + } + + if (status?.type === "confirmed" && status.result) { return { success: true, result: status.result }; } - if (status.type === "failed") { + if (status?.type === "failed") { return { success: false, error: @@ -124,7 +165,12 @@ async function pollUntilCompletionLoop( }; } - await new Promise((resolve) => setTimeout(resolve, pollInterval)); + await new Promise((resolve) => + setTimeout( + resolve, + Math.min(pollInterval, Math.max(0, deadline - Date.now())), + ), + ); } } diff --git a/js/packages/react/src/__tests__/hooks.test.tsx b/js/packages/react/src/__tests__/hooks.test.tsx index 77dd941e..263199a3 100644 --- a/js/packages/react/src/__tests__/hooks.test.tsx +++ b/js/packages/react/src/__tests__/hooks.test.tsx @@ -3,6 +3,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { IDKitErrorCodes } from "@worldcoin/idkit-core"; import packageJson from "../../package.json"; import { toErrorCode } from "../hooks/common"; +import { useIDKitInviteCodeRequest } from "../hooks/useIDKitInviteCodeRequest"; import { useIDKitRequest } from "../hooks/useIDKitRequest"; import { useIDKitSession } from "../hooks/useIDKitSession"; @@ -498,6 +499,110 @@ describe("request/session hooks", () => { expect(result.current.errorCode).toBe(IDKitErrorCodes.ConnectionFailed); }); + it("request hook retries a failed poll and accepts a later result", async () => { + const pollOnce = vi + .fn() + .mockRejectedValueOnce(new Error("Failed to fetch")) + .mockResolvedValueOnce({ type: "confirmed", result: { proof: "ok" } }); + requestMock.mockReturnValue({ + preset: vi.fn(async () => makeRequest(pollOnce)), + }); + + const { result } = renderHook(() => + useIDKitRequest({ + app_id: "app_test", + action: "test-action", + rp_context: baseRpContext, + allow_legacy_proofs: false, + preset: { type: "OrbLegacy" }, + polling: { interval: 0, timeout: 1000 }, + }), + ); + + act(() => result.current.open()); + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(pollOnce).toHaveBeenCalledTimes(2); + expect(result.current.result).toEqual({ proof: "ok" }); + }); + + it("request hook counts flow creation against its timeout", async () => { + const pollOnce = vi.fn(); + requestMock.mockReturnValue({ + preset: vi.fn(async () => { + await new Promise((resolve) => setTimeout(resolve, 30)); + return makeRequest(pollOnce); + }), + }); + + const { result } = renderHook(() => + useIDKitRequest({ + app_id: "app_test", + action: "test-action", + rp_context: baseRpContext, + allow_legacy_proofs: false, + preset: { type: "OrbLegacy" }, + polling: { interval: 0, timeout: 10 }, + }), + ); + + act(() => result.current.open()); + await waitFor(() => expect(result.current.isError).toBe(true)); + expect(result.current.errorCode).toBe(IDKitErrorCodes.Timeout); + expect(pollOnce).not.toHaveBeenCalled(); + }); + + it("request hook times out when a poll never settles", async () => { + const pollOnce = vi.fn(() => new Promise(() => {})); + requestMock.mockReturnValue({ + preset: vi.fn(async () => makeRequest(pollOnce)), + }); + + const { result } = renderHook(() => + useIDKitRequest({ + app_id: "app_test", + action: "test-action", + rp_context: baseRpContext, + allow_legacy_proofs: false, + preset: { type: "OrbLegacy" }, + polling: { interval: 0, timeout: 20 }, + }), + ); + + act(() => result.current.open()); + await waitFor(() => expect(result.current.isError).toBe(true)); + expect(result.current.errorCode).toBe(IDKitErrorCodes.Timeout); + expect(pollOnce).toHaveBeenCalledTimes(1); + }); + + it("invite-code hook retries a failed poll and accepts a later result", async () => { + const pollOnce = vi + .fn() + .mockRejectedValueOnce(new Error("Failed to fetch")) + .mockResolvedValueOnce({ type: "confirmed", result: { proof: "ok" } }); + requestWithInviteCodeMock.mockReturnValue({ + preset: vi.fn(async () => ({ + ...makeRequest(pollOnce), + expiresAt: 123, + })), + }); + + const { result } = renderHook(() => + useIDKitInviteCodeRequest({ + app_id: "app_test", + action: "test-action", + rp_context: baseRpContext, + allow_legacy_proofs: false, + preset: { type: "OrbLegacy" }, + polling: { interval: 0, timeout: 1000 }, + }), + ); + + act(() => result.current.open()); + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(pollOnce).toHaveBeenCalledTimes(2); + expect(result.current.result).toEqual({ proof: "ok" }); + }); + it("request hook maps confirmed status without payload to unexpected_response", async () => { requestMock.mockReturnValue({ preset: vi.fn(async () => diff --git a/js/packages/react/src/hooks/common.ts b/js/packages/react/src/hooks/common.ts index 09af6c11..b2bfe755 100644 --- a/js/packages/react/src/hooks/common.ts +++ b/js/packages/react/src/hooks/common.ts @@ -53,6 +53,32 @@ export async function delay(ms: number, signal?: AbortSignal): Promise { }); } +export async function pollBeforeDeadline( + pollOnce: () => Promise, + deadline: number, + signal: AbortSignal, +): Promise { + let timer: ReturnType | undefined; + let onAbort: (() => void) | undefined; + + try { + return await Promise.race([ + Promise.resolve() + .then(pollOnce) + .catch(() => null), + new Promise((resolve) => { + timer = setTimeout(resolve, Math.max(0, deadline - Date.now()), null); + onAbort = () => resolve(null); + signal.addEventListener("abort", onAbort, { once: true }); + if (signal.aborted) onAbort(); + }), + ]); + } finally { + clearTimeout(timer); + if (onAbort) signal.removeEventListener("abort", onAbort); + } +} + const knownErrorCodes = new Set(Object.values(IDKitErrorCodes)); function asKnownErrorCode(value: unknown): IDKitErrorCodes | null { diff --git a/js/packages/react/src/hooks/useIDKitFlow.ts b/js/packages/react/src/hooks/useIDKitFlow.ts index 4a7ff97b..282efb73 100644 --- a/js/packages/react/src/hooks/useIDKitFlow.ts +++ b/js/packages/react/src/hooks/useIDKitFlow.ts @@ -11,6 +11,7 @@ import { createInitialHookState, delay, ensureNotAborted, + pollBeforeDeadline, toErrorCode, type HookState, } from "./common"; @@ -110,6 +111,8 @@ export function useIDKitFlow( }; void (async () => { + const deadline = + Date.now() + (configRef.current.polling?.timeout ?? 900_000); try { if (isDebug()) console.debug("[IDKit] Creating flow handle…"); const request = await createFlowHandleRef.current(); @@ -130,19 +133,37 @@ export function useIDKitFlow( }); const pollInterval = configRef.current.polling?.interval ?? 1000; - const timeout = configRef.current.polling?.timeout ?? 900_000; - const startedAt = Date.now(); - while (true) { ensureNotAborted(controller.signal); - if (Date.now() - startedAt > timeout) { + if (Date.now() >= deadline) { setFailed(IDKitErrorCodes.Timeout); return; } - const nextStatus = await request.pollOnce(); + const nextStatus = await pollBeforeDeadline( + () => + request.pollOnce().catch((error: unknown) => { + if (isDebug()) + console.warn("[IDKit] Poll request failed; retrying", error); + return null; + }), + deadline, + controller.signal, + ); ensureNotAborted(controller.signal); + if (Date.now() >= deadline) { + setFailed(IDKitErrorCodes.Timeout); + return; + } + + if (nextStatus === null) { + await delay( + Math.min(pollInterval, deadline - Date.now()), + controller.signal, + ); + continue; + } if (nextStatus.type === "confirmed") { const confirmedResult = nextStatus.result; @@ -174,7 +195,10 @@ export function useIDKitFlow( return { ...prev, status: nextStatus.type }; }); - await delay(pollInterval, controller.signal); + await delay( + Math.min(pollInterval, deadline - Date.now()), + controller.signal, + ); } } catch (error) { if (controller.signal.aborted) { diff --git a/js/packages/react/src/hooks/useIDKitInviteCodeFlow.ts b/js/packages/react/src/hooks/useIDKitInviteCodeFlow.ts index 19b0fa4d..dabb3be6 100644 --- a/js/packages/react/src/hooks/useIDKitInviteCodeFlow.ts +++ b/js/packages/react/src/hooks/useIDKitInviteCodeFlow.ts @@ -7,7 +7,12 @@ import { type IDKitInviteCodeRequest, } from "@worldcoin/idkit-core"; import type { FlowConfig, IDKitInviteCodeHookResult } from "../types"; -import { delay, ensureNotAborted, toErrorCode } from "./common"; +import { + delay, + ensureNotAborted, + pollBeforeDeadline, + toErrorCode, +} from "./common"; import { createInitialInviteCodeHookState, type InviteCodeHookState, @@ -109,6 +114,8 @@ export function useIDKitInviteCodeFlow( }; void (async () => { + const deadline = + Date.now() + (configRef.current.polling?.timeout ?? 900_000); try { if (isDebug()) console.debug("[IDKit] Creating invite-code flow handle…"); @@ -135,19 +142,40 @@ export function useIDKitInviteCodeFlow( }); const pollInterval = configRef.current.polling?.interval ?? 1000; - const timeout = configRef.current.polling?.timeout ?? 900_000; - const startedAt = Date.now(); - while (true) { ensureNotAborted(controller.signal); - if (Date.now() - startedAt > timeout) { + if (Date.now() >= deadline) { setFailed(IDKitErrorCodes.Timeout); return; } - const nextStatus = await request.pollOnce(); + const nextStatus = await pollBeforeDeadline( + () => + request.pollOnce().catch((error: unknown) => { + if (isDebug()) + console.warn( + "[IDKit] Invite-code poll request failed; retrying", + error, + ); + return null; + }), + deadline, + controller.signal, + ); ensureNotAborted(controller.signal); + if (Date.now() >= deadline) { + setFailed(IDKitErrorCodes.Timeout); + return; + } + + if (nextStatus === null) { + await delay( + Math.min(pollInterval, deadline - Date.now()), + controller.signal, + ); + continue; + } if (nextStatus.type === "confirmed") { const confirmedResult = nextStatus.result; @@ -182,7 +210,10 @@ export function useIDKitInviteCodeFlow( return { ...prev, status: nextStatus.type }; }); - await delay(pollInterval, controller.signal); + await delay( + Math.min(pollInterval, deadline - Date.now()), + controller.signal, + ); } } catch (error) { if (controller.signal.aborted) { diff --git a/rust/core/src/bridge.rs b/rust/core/src/bridge.rs index 36a9d479..db073d25 100644 --- a/rust/core/src/bridge.rs +++ b/rust/core/src/bridge.rs @@ -1146,10 +1146,17 @@ impl BridgeConnection { .send() .await?; - if !response.status().is_success() { + if response.status() == reqwest::StatusCode::NOT_FOUND { return Ok(Status::Failed(AppError::ConnectionFailed)); } + if !response.status().is_success() { + return Err(Error::BridgeError(format!( + "Polling bridge returned HTTP {}", + response.status() + ))); + } + let poll_response: BridgePollResponse = response.json().await?; match poll_response.status.as_str() { @@ -4446,16 +4453,17 @@ mod tests { } } - fn serve_bridge_response(body: String) -> BridgeUrl { + fn serve_bridge_http_response(status: &str, body: String) -> BridgeUrl { let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); let addr = listener.local_addr().unwrap(); + let status = status.to_string(); std::thread::spawn(move || { let (mut stream, _) = listener.accept().unwrap(); let mut request_buffer = [0; 1024]; let _ = stream.read(&mut request_buffer); write!( stream, - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", + "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", body.len(), body ) @@ -4466,12 +4474,30 @@ mod tests { BridgeUrl::new(format!("http://{addr}"), &app_id).unwrap() } + fn serve_bridge_response(body: String) -> BridgeUrl { + serve_bridge_http_response("200 OK", body) + } + fn poll_once(connection: &BridgeConnection) -> Result { tokio::runtime::Runtime::new() .unwrap() .block_on(connection.poll_for_status()) } + #[test] + fn test_poll_only_treats_missing_entry_as_terminal_http_error() { + let mut connection = sample_connection(None); + connection.bridge_url = serve_bridge_http_response("404 Not Found", String::new()); + assert_eq!( + poll_once(&connection).unwrap(), + Status::Failed(AppError::ConnectionFailed) + ); + + connection.bridge_url = + serve_bridge_http_response("503 Service Unavailable", String::new()); + assert!(matches!(poll_once(&connection), Err(Error::BridgeError(_)))); + } + #[test] fn test_connect_url_includes_return_to_when_provided() { let connection = sample_connection(Some("idkitsample://callback?step=proof".to_string()));