retry logic

hugodutka · hugodutka · commit 4058858ab5b8 · 2025-03-24T17:45:44.000Z
diff --git a/site/src/components/GitDeviceAuth/GitDeviceAuth.tsx b/site/src/components/GitDeviceAuth/GitDeviceAuth.tsx
@@ -5,6 +5,7 @@ import CircularProgress from "@mui/material/CircularProgress";
 import Link from "@mui/material/Link";
 import type { ApiErrorResponse } from "api/errors";
 import type { ExternalAuthDevice } from "api/typesGenerated";
+import { isAxiosError } from "axios";
 import { Alert, AlertDetail } from "components/Alert/Alert";
 import { CopyButton } from "components/CopyButton/CopyButton";
 import type { FC } from "react";
@@ -14,6 +15,59 @@ interface GitDeviceAuthProps {
 	deviceExchangeError?: ApiErrorResponse;
 }
 
+const DeviceExchangeError = {
+	AuthorizationPending: "authorization_pending",
+	SlowDown: "slow_down",
+	ExpiredToken: "expired_token",
+	AccessDenied: "access_denied",
+} as const;
+
+export const isExchangeErrorRetryable = (_: number, error: unknown) => {
+	if (!isAxiosError(error)) {
+		return false;
+	}
+	const detail = error.response?.data?.detail;
+	return (
+		detail === DeviceExchangeError.AuthorizationPending ||
+		detail === DeviceExchangeError.SlowDown
+	);
+};
+
+/**
+ * The OAuth2 specification (https://datatracker.ietf.org/doc/html/rfc8628)
+ * describes how the client should handle retries. This function returns a
+ * closure that implements the retry logic described in the specification.
+ * The closure should be memoized because it stores state.
+ */
+export const newRetryDelay = (initialInterval: number | undefined) => {
+	// "If no value is provided, clients MUST use 5 as the default."
+	// https://datatracker.ietf.org/doc/html/rfc8628#section-3.2
+	let interval = initialInterval ?? 5;
+	let lastFailureCountHandled = 0;
+	return (failureCount: number, error: unknown) => {
+		const isSlowDown =
+			isAxiosError(error) &&
+			error.response?.data.detail === DeviceExchangeError.SlowDown;
+		// We check the failure count to ensure we increase the interval
+		// at most once per failure.
+		if (isSlowDown && lastFailureCountHandled < failureCount) {
+			lastFailureCountHandled = failureCount;
+			// https://datatracker.ietf.org/doc/html/rfc8628#section-3.5
+			// "the interval MUST be increased by 5 seconds for this and all subsequent requests"
+			interval += 5;
+		}
+		let extraDelay = 0;
+		if (isSlowDown) {
+			// I found GitHub is very strict about their rate limits, and they'll block
+			// even if the request is 500ms earlier than they expect. This may happen due to
+			// e.g. network latency, so it's best to cool down for longer if GitHub just
+			// rejected our request.
+			extraDelay = 5;
+		}
+		return (interval + extraDelay) * 1000;
+	};
+};
+
 export const GitDeviceAuth: FC<GitDeviceAuthProps> = ({
 	externalAuthDevice,
 	deviceExchangeError,
@@ -27,16 +81,26 @@ export const GitDeviceAuth: FC<GitDeviceAuthProps> = ({
 	if (deviceExchangeError) {
 		// See https://datatracker.ietf.org/doc/html/rfc8628#section-3.5
 		switch (deviceExchangeError.detail) {
-			case "authorization_pending":
+			case DeviceExchangeError.AuthorizationPending:
+				break;
+			case DeviceExchangeError.SlowDown:
+				status = (
+					<div>
+						{status}
+						<Alert severity="warning">
+							Rate limit reached. Waiting a few seconds before retrying...
+						</Alert>
+					</div>
+				);
 				break;
-			case "expired_token":
+			case DeviceExchangeError.ExpiredToken:
 				status = (
 					<Alert severity="error">
 						The one-time code has expired. Refresh to get a new one!
 					</Alert>
 				);
 				break;
-			case "access_denied":
+			case DeviceExchangeError.AccessDenied:
 				status = (
 					<Alert severity="error">Access to the Git provider was denied.</Alert>
 				);
diff --git a/site/src/pages/ExternalAuthPage/ExternalAuthPage.tsx b/site/src/pages/ExternalAuthPage/ExternalAuthPage.tsx
@@ -6,10 +6,15 @@ import {
 	externalAuthProvider,
 } from "api/queries/externalAuth";
 import { isAxiosError } from "axios";
+import {
+	isExchangeErrorRetryable,
+	newRetryDelay,
+} from "components/GitDeviceAuth/GitDeviceAuth";
 import { SignInLayout } from "components/SignInLayout/SignInLayout";
 import { Welcome } from "components/Welcome/Welcome";
 import { useAuthenticated } from "contexts/auth/RequireAuth";
 import type { FC } from "react";
+import { useMemo } from "react";
 import { useQuery, useQueryClient } from "react-query";
 import { useParams, useSearchParams } from "react-router-dom";
 import ExternalAuthPageView from "./ExternalAuthPageView";
@@ -32,17 +37,22 @@ const ExternalAuthPage: FC = () => {
 			Boolean(externalAuthProviderQuery.data?.device),
 		refetchOnMount: false,
 	});
+	const retryDelay = useMemo(
+		() => newRetryDelay(externalAuthDeviceQuery.data?.interval),
+		[externalAuthDeviceQuery.data],
+	);
 	const exchangeExternalAuthDeviceQuery = useQuery({
 		...exchangeExternalAuthDevice(
 			provider,
 			externalAuthDeviceQuery.data?.device_code ?? "",
 			queryClient,
 		),
 		enabled: Boolean(externalAuthDeviceQuery.data),
-		retry: true,
-		retryDelay: (externalAuthDeviceQuery.data?.interval || 5) * 1000,
-		refetchOnWindowFocus: (query) =>
-			query.state.status === "success" ? false : "always",
+		retry: isExchangeErrorRetryable,
+		retryDelay,
+		// We don't want to refetch the query outside of the standard retry
+		// logic, because the device auth flow is very strict about rate limits.
+		refetchOnWindowFocus: false,
 	});
 
 	if (externalAuthProviderQuery.isLoading || !externalAuthProviderQuery.data) {
diff --git a/site/src/pages/LoginOAuthDevicePage/LoginOAuthDevicePage.tsx b/site/src/pages/LoginOAuthDevicePage/LoginOAuthDevicePage.tsx
@@ -4,21 +4,18 @@ import {
 	getGitHubDeviceFlowCallback,
 } from "api/queries/oauth2";
 import { isAxiosError } from "axios";
+import {
+	isExchangeErrorRetryable,
+	newRetryDelay,
+} from "components/GitDeviceAuth/GitDeviceAuth";
 import { SignInLayout } from "components/SignInLayout/SignInLayout";
 import { Welcome } from "components/Welcome/Welcome";
-import { useEffect } from "react";
+import { useEffect, useMemo } from "react";
 import type { FC } from "react";
 import { useQuery } from "react-query";
 import { useSearchParams } from "react-router-dom";
 import LoginOAuthDevicePageView from "./LoginOAuthDevicePageView";
 
-const isErrorRetryable = (error: unknown) => {
-	if (!isAxiosError(error)) {
-		return false;
-	}
-	return error.response?.data?.detail === "authorization_pending";
-};
-
 // The page is hardcoded to only use GitHub,
 // as that's the only OAuth2 login provider in our backend
 // that currently supports the device flow.
@@ -38,19 +35,23 @@ const LoginOAuthDevicePage: FC = () => {
 		...getGitHubDevice(),
 		refetchOnMount: false,
 	});
+
+	const retryDelay = useMemo(
+		() => newRetryDelay(externalAuthDeviceQuery.data?.interval),
+		[externalAuthDeviceQuery.data],
+	);
+
 	const exchangeExternalAuthDeviceQuery = useQuery({
 		...getGitHubDeviceFlowCallback(
 			externalAuthDeviceQuery.data?.device_code ?? "",
 			state,
 		),
 		enabled: Boolean(externalAuthDeviceQuery.data),
-		retry: (_, error) => isErrorRetryable(error),
-		retryDelay: (externalAuthDeviceQuery.data?.interval || 5) * 1000,
-		refetchOnWindowFocus: (query) =>
-			query.state.status === "success" ||
-			(query.state.error != null && !isErrorRetryable(query.state.error))
-				? false
-				: "always",
+		retry: isExchangeErrorRetryable,
+		retryDelay,
+		// We don't want to refetch the query outside of the standard retry
+		// logic, because the device auth flow is very strict about rate limits.
+		refetchOnWindowFocus: false,
 	});
 
 	useEffect(() => {
