fix(realtime): refresh tokens in background without dropping the WebSocket connection

libs/client/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@fal-ai/client",
   "description": "The fal.ai client for JavaScript and TypeScript",
-  "version": "1.9.3",
+  "version": "1.9.4",
   "license": "MIT",
   "repository": {
     "type": "git",

libs/client/src/realtime.ts

@@ -163,6 +163,7 @@
     ),
     active: state(
       transition("send", "active", reduce(sendMessage)),
+      transition("authenticated", "active", reduce(setToken)),
       transition("unauthorized", "idle", reduce(expireToken)),
       transition("connectionClosed", "idle", reduce(closeConnection)),
       transition("close", "idle", reduce(closeConnection)),
@@ -539,6 +540,7 @@
 
       let previousState: string | undefined;
       let latestEnqueuedMessage: any;
+      let tokenRefreshTimer: ReturnType<typeof setTimeout> | undefined;
 
       // Although the state machine is cached so we don't open multiple connections,
       // we still need to update the callbacks so we can call the correct references
@@ -586,28 +588,52 @@
                   return getTemporaryAuthToken(app, config);
                 };
 
+            const effectiveExpiration = tokenProvider
+              ? tokenExpirationSeconds
+              : TOKEN_EXPIRATION_SECONDS;
+
+            const scheduleTokenRefresh =
+              effectiveExpiration !== undefined
+                ? () => {
+                    clearTimeout(tokenRefreshTimer);
+                    const refreshMs = Math.round(
+                      effectiveExpiration * 0.9 * 1000,
+                    );
+                    tokenRefreshTimer = setTimeout(() => {
+                      if (machine.current !== "active") {
+                        return;
+                      }
+                      fetchToken()
+                        .then((newToken) => {
+                          if (machine.current !== "active") {
+                            return;
+                          }
+                          queueMicrotask(() => {
+                            send({ type: "authenticated", token: newToken });
+                          });
+                          scheduleTokenRefresh();
+                        })
+                        .catch(() => {
+                          if (machine.current !== "active") {
+                            return;
+                          }
+                          const retryMs = Math.round(
+                            effectiveExpiration * 0.05 * 1000,
+                          );
+                          tokenRefreshTimer = setTimeout(() => {
+                            scheduleTokenRefresh();
+                          }, retryMs);
+                        });
+                    }, refreshMs);
+                  }
+                : noop;
+
             fetchToken()
               .then((token) => {
-                // Use queueMicrotask to ensure the state machine processes
-                // this on a clean call stack, not nested inside onChange.
-                // robot3's interpret can lose track of state changes when
-                // send() is called synchronously inside an onChange handler.
                 queueMicrotask(() => {
                   send({ type: "authenticated", token });
                 });
-                // Only schedule token refresh if we know the expiration time.
-                // For custom tokenProvider without tokenExpirationSeconds, skip auto-refresh.
-                const effectiveExpiration = tokenProvider
-                  ? tokenExpirationSeconds
-                  : TOKEN_EXPIRATION_SECONDS;
-                if (effectiveExpiration !== undefined) {
-                  const tokenRefreshInterval = Math.round(
-                    effectiveExpiration * 0.9 * 1000,
-                  );
-                  setTimeout(() => {
-                    send({ type: "expireToken" });
-                  }, tokenRefreshInterval);
-                }
+                scheduleTokenRefresh();
               })
               .catch((error) => {
                 queueMicrotask(() => {
@@ -648,7 +674,7 @@
               }
               send({ type: "connectionClosed", code: event.code });
             };
             ws.onerror = (event) => {
               // TODO specify error protocol for identified errors
               const { onError = noop } = getCallbacks();
               onError(new ApiError({ message: "Unknown error", status: 500 }));
@@ -669,6 +695,10 @@
               });
             };
           }
+          if (previousState === "active" && machine.current !== "active") {
+            clearTimeout(tokenRefreshTimer);
+            tokenRefreshTimer = undefined;
+          }
           previousState = machine.current;
         },
       );