From 18d037e4110123886f4848986fb740f3aa8e3bce Mon Sep 17 00:00:00 2001
From: Fredrik Rombach Ekelund <fredrik.rombach.ekelund@automattic.com>
Date: Tue, 1 Oct 2024 15:50:19 +0200
Subject: [PATCH] Refetch defensive mode endpoint after it expires

---
 client/data/hosting/use-cache.ts | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/client/data/hosting/use-cache.ts b/client/data/hosting/use-cache.ts
index cd6c2e6e66652..1f36960decc4e 100644
--- a/client/data/hosting/use-cache.ts
+++ b/client/data/hosting/use-cache.ts
@@ -7,7 +7,7 @@ import {
 } from '@tanstack/react-query';
 import { useI18n } from '@wordpress/react-i18n';
 import { useTranslate } from 'i18n-calypso';
-import { useCallback } from 'react';
+import { useCallback, useEffect } from 'react';
 import wp from 'calypso/lib/wp';
 import { useDispatch } from 'calypso/state';
 import { successNotice, errorNotice, plainNotice } from 'calypso/state/notices/actions';
@@ -184,7 +184,7 @@ type EdgeCacheDefensiveModeQueryData = {
 };
 
 export function useEdgeCacheDefensiveModeQuery( siteId: number | null ) {
-	return useQuery< EdgeCacheDefensiveModeQueryData >( {
+	const query = useQuery< EdgeCacheDefensiveModeQueryData >( {
 		queryKey: getEdgeCacheDefensiveModeQueryKey( siteId ),
 		queryFn: () =>
 			wp.req.get( {
@@ -193,6 +193,27 @@ export function useEdgeCacheDefensiveModeQuery( siteId: number | null ) {
 			} ),
 		enabled: siteId !== null,
 	} );
+
+	const { data, refetch } = query;
+	const enabled = data?.enabled ?? false;
+	const enabledUntil = data?.enabled_until ?? 0;
+
+	// Refetch the defensive mode API endpoint ten seconds after it's set to disable
+	useEffect( () => {
+		if ( ! enabled ) {
+			return;
+		}
+
+		const msDiff = enabledUntil * 1000 - Date.now();
+		const delay = Math.max( 0, msDiff ) + 10000;
+		const timeoutId = setTimeout( () => refetch(), delay );
+
+		return () => {
+			clearTimeout( timeoutId );
+		};
+	}, [ enabled, enabledUntil, refetch ] );
+
+	return query;
 }
 
 type EdgeCacheDefensiveModeMutationVariables = { active: true; ttl: number } | { active: false };