fix(browser-sdk) set isLoading on refetch due to context update (#547)

Author: roncohen

packages/browser-sdk/src/client.ts

@@ -676,24 +676,20 @@ export class ReflagClient {
    * @param user
    */
   async updateUser(user: { [key: string]: string | number | undefined }) {
-    const userIdChanged = user.id && user.id !== this.context.user?.id;
     const newUserContext = {
       ...this.context.user,
       ...user,
       id: user.id ?? this.context.user?.id,
     };
 
-    // Nothing has changed, skipping update
-    if (deepEqual(this.context.user, newUserContext)) return;
-    this.context.user = newUserContext;
-    void this.user();
-
-    // Update the feedback user if the user ID has changed
-    if (userIdChanged) {
-      void this.updateAutoFeedbackUser(String(user.id));
-    }
-
-    await this.flagsClient.setContext(this.context);
+    return this.applyContext(
+      {
+        user: newUserContext,
+        company: this.context.company,
+        other: this.context.other,
+      },
+      { warnOnMissingIds: false },
+    );
   }
 
   /**
@@ -710,12 +706,14 @@ export class ReflagClient {
       id: company.id ?? this.context.company?.id,
     };
 
-    // Nothing has changed, skipping update
-    if (deepEqual(this.context.company, newCompanyContext)) return;
-    this.context.company = newCompanyContext;
-    void this.company();
-
-    await this.flagsClient.setContext(this.context);
+    return this.applyContext(
+      {
+        user: this.context.user,
+        company: newCompanyContext,
+        other: this.context.other,
+      },
+      { warnOnMissingIds: false },
+    );
   }
 
   /**
@@ -733,11 +731,14 @@ export class ReflagClient {
       ...otherContext,
     };
 
-    // Nothing has changed, skipping update
-    if (deepEqual(this.context.other, newOtherContext)) return;
-    this.context.other = newOtherContext;
-
-    await this.flagsClient.setContext(this.context);
+    return this.applyContext(
+      {
+        user: this.context.user,
+        company: this.context.company,
+        other: newOtherContext,
+      },
+      { warnOnMissingIds: false },
+    );
   }
 
   /**
@@ -746,9 +747,15 @@ export class ReflagClient {
    *
    * @param context The context to update.
    */
-  async setContext({ otherContext, ...context }: ReflagDeprecatedContext) {
-    const userIdChanged =
-      context.user?.id && context.user.id !== this.context.user?.id;
+  async setContext(context: ReflagDeprecatedContext) {
+    return this.applyContext(context, { warnOnMissingIds: true });
+  }
+
+  private async applyContext(
+    { otherContext, ...context }: ReflagDeprecatedContext,
+    { warnOnMissingIds }: { warnOnMissingIds: boolean },
+  ) {
+    const previousContext = this.context;
 
     // Create a new context object making sure to clone the user and company objects
     const newContext = {
@@ -757,32 +764,49 @@ export class ReflagClient {
       other: { ...otherContext, ...context.other },
     };
 
-    if (!context.user?.id) {
+    if (warnOnMissingIds && !context.user?.id) {
       this.logger.warn("No user Id provided in context, user will be ignored");
     }
-    if (!context.company?.id) {
+    if (warnOnMissingIds && !context.company?.id) {
       this.logger.warn(
         "No company Id provided in context, company will be ignored",
       );
     }
 
     // Nothing has changed, skipping update
-    if (deepEqual(this.context, newContext)) return;
+    if (deepEqual(previousContext, newContext)) return;
+
+    const userChanged = !deepEqual(previousContext.user, newContext.user);
+    const companyChanged = !deepEqual(
+      previousContext.company,
+      newContext.company,
+    );
+    const userIdChanged =
+      !!newContext.user?.id && newContext.user.id !== previousContext.user?.id;
+
     this.context = newContext;
 
-    if (context.company) {
+    if (companyChanged) {
       void this.company();
     }
 
-    if (context.user) {
+    if (userChanged) {
       void this.user();
       // Update the automatic feedback user if the user ID has changed
       if (userIdChanged) {
-        void this.updateAutoFeedbackUser(String(context.user.id));
+        void this.updateAutoFeedbackUser(String(newContext.user!.id));
       }
     }
 
-    await this.flagsClient.setContext(this.context);
+    const shouldTrackLoading = this.state === "initialized";
+    if (shouldTrackLoading) {
+      this.setState("initializing");
+    }
+
+    const didApply = await this.flagsClient.setContext(this.context);
+    if (didApply && this.state === "initializing") {
+      this.setState("initialized");
+    }
   }
 
   /**

packages/browser-sdk/src/flag/flags.ts

@@ -209,6 +209,7 @@ export class FlagsClient {
   private flagOverrides: FlagOverrides = {};
   private flags: RawFlags = {};
   private fallbackFlags: FallbackFlags = {};
+  private contextFetchVersion = 0;
   private storage: StorageAdapter;
   private refreshEvents: number[] = [];
   private enqueueBulkEvent?: (event: BulkEvent) => Promise<void>;
@@ -297,7 +298,15 @@ export class FlagsClient {
 
   async setContext(context: ReflagContext) {
     this.context = context;
-    this.setFetchedFlags((await this.maybeFetchFlags()) || {});
+    const requestVersion = ++this.contextFetchVersion;
+    const fetchedFlags = (await this.maybeFetchFlags()) || {};
+
+    if (requestVersion !== this.contextFetchVersion) {
+      return false;
+    }
+
+    this.setFetchedFlags(fetchedFlags);
+    return true;
   }
 
   updateFlags(triggerEvent = true) {

packages/browser-sdk/test/client.test.ts

@@ -57,6 +57,20 @@ describe("ReflagClient", () => {
       );
       expect(flagClientSetContext).toHaveBeenCalledWith(client["context"]);
     });
+
+    it("does not warn about a missing company when updating a user-only context", async () => {
+      client = new ReflagClient({
+        publishableKey: "test-key-user-only",
+        user: { id: "user1" },
+      });
+      const warnSpy = vi.spyOn(client.logger, "warn");
+
+      await client.updateUser({ name: "Updated User" });
+
+      expect(warnSpy).not.toHaveBeenCalledWith(
+        "No company Id provided in context, company will be ignored",
+      );
+    });
   });
 
   describe("updateCompany", () => {
@@ -194,6 +208,101 @@ describe("ReflagClient", () => {
       expect(checkHook).not.toHaveBeenCalled();
       expect(flagsUpdated).not.toHaveBeenCalled();
     });
+
+    it("sets state to initializing while refetching flags after initialization", async () => {
+      await client.initialize();
+
+      let resolveFetch: (() => void) | undefined;
+      const setContextPromise = new Promise<boolean>((resolve) => {
+        resolveFetch = () => resolve(true);
+      });
+      const setContext = vi
+        .spyOn(FlagsClient.prototype, "setContext")
+        .mockImplementation(async () => {
+          return setContextPromise;
+        });
+
+      const stateUpdated = vi.fn();
+      client.on("stateUpdated", stateUpdated);
+
+      const updatePromise = client.updateOtherContext({ workspaceId: "ws-1" });
+
+      expect(client.getState()).toBe("initializing");
+      expect(stateUpdated).toHaveBeenCalledWith("initializing");
+      expect(setContext).toHaveBeenCalledWith(client["context"]);
+
+      resolveFetch?.();
+      await updatePromise;
+
+      expect(client.getState()).toBe("initialized");
+      expect(stateUpdated).toHaveBeenLastCalledWith("initialized");
+    });
+
+    it("keeps loading tied to the latest context update", async () => {
+      await client.initialize();
+
+      let resolveFirstFetch: (() => void) | undefined;
+      let resolveSecondFetch: (() => void) | undefined;
+      const firstFetch = new Promise<boolean>((resolve) => {
+        resolveFirstFetch = () => resolve(false);
+      });
+      const secondFetch = new Promise<boolean>((resolve) => {
+        resolveSecondFetch = () => resolve(true);
+      });
+
+      vi.spyOn(FlagsClient.prototype, "setContext")
+        .mockImplementationOnce(async () => firstFetch)
+        .mockImplementationOnce(async () => secondFetch);
+
+      const firstUpdate = client.updateOtherContext({ workspaceId: "ws-1" });
+      const secondUpdate = client.updateOtherContext({ workspaceId: "ws-2" });
+
+      expect(client.getState()).toBe("initializing");
+
+      resolveSecondFetch?.();
+      await secondUpdate;
+
+      expect(client.getState()).toBe("initialized");
+
+      resolveFirstFetch?.();
+      await firstUpdate;
+
+      expect(client.getState()).toBe("initialized");
+    });
+  });
+
+  describe("setContext warnings", () => {
+    it("does not warn about missing ids when updating anonymous other context", async () => {
+      client = new ReflagClient({
+        publishableKey: "test-key-anon",
+      });
+      const warnSpy = vi.spyOn(client.logger, "warn");
+
+      await client.updateOtherContext({ workspaceId: "ws-1" });
+
+      expect(warnSpy).not.toHaveBeenCalledWith(
+        "No user Id provided in context, user will be ignored",
+      );
+      expect(warnSpy).not.toHaveBeenCalledWith(
+        "No company Id provided in context, company will be ignored",
+      );
+    });
+
+    it("still warns when setContext replaces context without user or company ids", async () => {
+      client = new ReflagClient({
+        publishableKey: "test-key-set-context",
+      });
+      const warnSpy = vi.spyOn(client.logger, "warn");
+
+      await client.setContext({ other: { workspaceId: "ws-1" } });
+
+      expect(warnSpy).toHaveBeenCalledWith(
+        "No user Id provided in context, user will be ignored",
+      );
+      expect(warnSpy).toHaveBeenCalledWith(
+        "No company Id provided in context, company will be ignored",
+      );
+    });
   });
 
   describe("stop", () => {

packages/browser-sdk/test/flags.test.ts

@@ -141,6 +141,75 @@ describe("FlagsClient", () => {
     expect(timeoutMs).toEqual(5000);
   });
 
+  test("only applies flags from the latest context update", async () => {
+    const { newFlagsClient } = flagsClientFactory();
+    const flagsClient = newFlagsClient();
+
+    let resolveFirstFetch:
+      | ((flags: Record<string, RawFlag>) => void)
+      | undefined;
+    let resolveSecondFetch:
+      | ((flags: Record<string, RawFlag>) => void)
+      | undefined;
+    const firstFetch = new Promise<Record<string, RawFlag>>((resolve) => {
+      resolveFirstFetch = resolve;
+    });
+    const secondFetch = new Promise<Record<string, RawFlag>>((resolve) => {
+      resolveSecondFetch = resolve;
+    });
+
+    vi.spyOn(flagsClient as any, "maybeFetchFlags")
+      .mockImplementationOnce(async () => firstFetch)
+      .mockImplementationOnce(async () => secondFetch);
+
+    const firstUpdate = flagsClient.setContext({
+      user: { id: "user-1" },
+      company: { id: "company-1" },
+      other: { workspaceId: "workspace-1" },
+    });
+    const secondUpdate = flagsClient.setContext({
+      user: { id: "user-2" },
+      company: { id: "company-2" },
+      other: { workspaceId: "workspace-2" },
+    });
+
+    resolveSecondFetch?.({
+      latestFlag: {
+        key: "latestFlag",
+        isEnabled: true,
+        targetingVersion: 2,
+      },
+    });
+
+    await expect(secondUpdate).resolves.toBe(true);
+    expect(flagsClient.getFlags()).toEqual({
+      latestFlag: {
+        key: "latestFlag",
+        isEnabled: true,
+        targetingVersion: 2,
+        isEnabledOverride: null,
+      },
+    });
+
+    resolveFirstFetch?.({
+      staleFlag: {
+        key: "staleFlag",
+        isEnabled: false,
+        targetingVersion: 1,
+      },
+    });
+
+    await expect(firstUpdate).resolves.toBe(false);
+    expect(flagsClient.getFlags()).toEqual({
+      latestFlag: {
+        key: "latestFlag",
+        isEnabled: true,
+        targetingVersion: 2,
+        isEnabledOverride: null,
+      },
+    });
+  });
+
   test("return fallback flags on failure (string list)", async () => {
     const { newFlagsClient, httpClient } = flagsClientFactory();