fix bugs in token refresh, add tokenStore example

Author: azrosen92

gusto_embedded/README.md

@@ -97,11 +97,38 @@ For supported JavaScript runtimes, please consult [RUNTIMES.md](RUNTIMES.md).
 ## SDK Example Usage
 
 ### Example
+Automatic token refresh using a persistent data store.
 
 ```typescript
 import { GustoEmbedded } from "@gusto/embedded-api";
 import { CompanyAuthenticatedClient } = "@gusto/embedded-api/CompanyAuthenticatedClient";
 
+class PersistentTokenStore implements TokenStore {
+  constructor() {}
+
+  async get() {
+    const { token, expires, refreshToken } = retrieveToken();
+
+    return {
+      token,
+      expires,
+      refreshToken,
+    };
+  }
+
+  async set({
+    token,
+    expires,
+    refreshToken,
+  }: {
+    token: string;
+    expires: number;
+    refreshToken: string;
+  }) {
+    saveToken(token, refreshToken, expires);
+  }
+}
+
 const client = new GustoEmbedded();
 const clientId = process.env["GUSTOEMBEDDED_CLIENT_ID"]
 const clientSecret = process.env["GUSTOEMBEDDED_CLIENT_SECRET"];
@@ -132,6 +159,8 @@ async function run() {
 
   const { accessToken, refreshToken, companyUuid, expiresIn } = response.object;
 
+  const tokenStore = PersistentTokenStore();
+
   const companyAuthClient = CompanyAuthenticatedClient({
     clientId,
     clientSecret,
@@ -140,6 +169,7 @@ async function run() {
     expiresIn,
     options: {
       server: "demo",
+      tokenStore,
     },
   });
 

gusto_embedded/src/CompanyAuthenticatedClient.ts

@@ -1,5 +1,11 @@
-import { TokenRefreshOptions, withTokenRefresh } from "./companyAuth.js";
-import { GustoEmbeddedCore } from "./core.js";
+import { HTTPClient } from "./lib/http.js";
+import {
+  InMemoryTokenStore,
+  refreshAndSaveAuthToken,
+  TokenRefreshOptions,
+  withTokenRefresh,
+} from "./companyAuth.js";
+import { GustoEmbedded } from "./sdk/sdk.js";
 import { SDKOptions, ServerDemo, ServerList } from "./lib/config.js";
 
 type ClientArguments = {
@@ -20,16 +26,33 @@ export function CompanyAuthenticatedClient({
   options,
 }: ClientArguments) {
   const authUrl = constructAuthUrl(options);
+  const tokenStore = new InMemoryTokenStore();
 
-  return new GustoEmbeddedCore({
+  const httpClientWithTokenRefresh = new HTTPClient();
+
+  httpClientWithTokenRefresh.addHook("response", async (res) => {
+    if (res.status === 401) {
+      console.log("Unauthorized, attempting to refresh token");
+
+      await refreshAndSaveAuthToken(
+        authUrl,
+        { clientId, clientSecret, refreshToken },
+        tokenStore
+      );
+    }
+  });
+
+  return new GustoEmbedded({
     ...options,
+    httpClient: httpClientWithTokenRefresh,
     companyAccessAuth: withTokenRefresh(
       clientId,
       clientSecret,
       accessToken,
       refreshToken,
       expiresIn,
       {
+        tokenStore,
         ...options,
         url: authUrl,
       }

gusto_embedded/src/companyAuth.ts

@@ -49,43 +49,65 @@ export function withTokenRefresh(
       return session.token;
     }
 
-    try {
-      const response = await fetch(url, {
-        method: "POST",
-        headers: {
-          "content-type": "application/x-www-form-urlencoded",
-          // Include the SDK's user agent in the request so requests can be
-          // tracked using observability infrastructure.
-          "user-agent": SDK_METADATA.userAgent,
-        },
-        body: new URLSearchParams({
-          client_id: clientId,
-          client_secret: clientSecret,
-          grant_type: "refresh_token",
-          refresh_token: refreshToken,
-        }),
-      });
-
-      if (!response.ok) {
-        throw new Error("Unexpected status code: " + response.status);
-      }
-
-      const json = await response.json();
-      const data = tokenResponseSchema.parse(json);
-
-      await tokenStore.set({
-        token: data.access_token,
-        expires: Date.now() + data.expires_in * 1000 - tolerance,
-        refreshToken: data.refresh_token,
-      });
-
-      return data.access_token;
-    } catch (error) {
-      throw new Error("Failed to obtain OAuth token: " + error);
-    }
+    return await refreshAndSaveAuthToken(
+      url,
+      {
+        clientId,
+        clientSecret,
+        refreshToken: session?.refreshToken ?? refreshToken,
+      },
+      tokenStore
+    );
   };
 }
 
+export async function refreshAndSaveAuthToken(
+  authUrl: string,
+  refreshCredentials: {
+    clientId: string;
+    clientSecret: string;
+    refreshToken: string;
+  },
+  tokenStore: TokenStore
+): Promise<string> {
+  const { clientId, clientSecret, refreshToken } = refreshCredentials;
+
+  try {
+    const response = await fetch(authUrl, {
+      method: "POST",
+      headers: {
+        "content-type": "application/x-www-form-urlencoded",
+        // Include the SDK's user agent in the request so requests can be
+        // tracked using observability infrastructure.
+        "user-agent": SDK_METADATA.userAgent,
+      },
+      body: new URLSearchParams({
+        client_id: clientId,
+        client_secret: clientSecret,
+        grant_type: "refresh_token",
+        refresh_token: refreshToken,
+      }),
+    });
+
+    if (!response.ok) {
+      throw new Error("Unexpected status code: " + response.status);
+    }
+
+    const json = await response.json();
+    const data = tokenResponseSchema.parse(json);
+
+    await tokenStore.set({
+      token: data.access_token,
+      expires: Date.now() + data.expires_in * 1000 - tolerance,
+      refreshToken: data.refresh_token,
+    });
+
+    return data.access_token;
+  } catch (error) {
+    throw new Error("Failed to obtain OAuth token: " + error);
+  }
+}
+
 /**
  * A TokenStore is used to save and retrieve OAuth tokens for use across SDK
  * method calls. This interface can be implemented to store tokens in memory,

gusto_embedded/src/hooks/clientcredentials.ts

@@ -164,6 +164,12 @@ export class ClientCredentialsHook
     if (typeof source === "function") {
       security = await source();
     }
+
+    // The client was passed a raw access token, no need to fetch one.
+    if (typeof security === "string") {
+      return null;
+    }
+
     const out = parse(
       security,
       (val) =>