Segment fixes (#690)

Author: nlarew

package-lock.json

@@ -125,6 +125,7 @@
     "@types/react": "^18.0.37",
     "@types/react-dom": "^18.0.11",
     "@types/react-transition-group": "^4.4.6",
+    "@types/segment-analytics": "^0.0.38",
     "@types/unist": "^3.0.2",
     "@typescript-eslint/eslint-plugin": "^5.59.0",
     "@typescript-eslint/parser": "^5.59.0",

packages/mongodb-chatbot-ui/package.json

@@ -69,7 +69,10 @@ function App() {
           serverBaseUrl={serverBaseUrl}
           shouldStream={shouldStream}
           darkMode={darkMode}
-          fetchOptions={{ credentials: "include" }}
+          fetchOptions={() => ({
+            credentials: "include",
+            headers: getSegmentIdHeaders(),
+          })}
           onOpen={() => {
             console.log("Docs Chatbot opened");
           }}
@@ -85,10 +88,10 @@ function App() {
           serverBaseUrl={serverBaseUrl}
           shouldStream={shouldStream}
           darkMode={darkMode}
-          fetchOptions={{
+          fetchOptions={() => ({
             credentials: "include",
             headers: getSegmentIdHeaders(),
-          }}
+          })}
           getClientContext={() => ({
             user: "test-user-pls-ignore",
           })}

packages/mongodb-chatbot-ui/src/App.tsx

@@ -1,9 +1,22 @@
+function stripLiteralQuotes(str?: string) {
+  return str?.replace(/^"|"$/g, "");
+}
+
 /**
  * Get Segment IDs from the browser's local storage.
  */
 export function getSegmentIds() {
-  const userId = localStorage.getItem("ajs_user_id") ?? undefined;
-  const anonymousId = localStorage.getItem("ajs_anonymous_id") ?? undefined;
+  const analytics =
+    typeof window !== "undefined" && window.analytics ? window.analytics : null;
+
+  const userId = stripLiteralQuotes(
+    analytics?.user().id() ?? localStorage.getItem("ajs_user_id") ?? undefined
+  );
+  const anonymousId = stripLiteralQuotes(
+    analytics?.user().anonymousId() ??
+      localStorage.getItem("ajs_anonymous_id") ??
+      undefined
+  );
   return {
     userId,
     anonymousId,

packages/mongodb-chatbot-ui/src/segment.ts

@@ -648,76 +648,6 @@ describe("ConversationService", () => {
     expect(commentPromise).resolves.toBeUndefined();
   });
 
-  it("appends custom fetch options", async () => {
-    const conversationService = new ConversationService({
-      serverUrl,
-      fetchOptions: {
-        headers: new Headers({
-          foo: "bar",
-        }),
-        credentials: "include",
-      },
-    });
-    let getOptions = mockFetchResponse({
-      data: {
-        _id: "650b4b260f975ef031016c8d",
-        messages: [],
-        createdAt: new Date().getTime(),
-      },
-    });
-    await conversationService.createConversation();
-    const createOptions = getOptions();
-
-    expect(createOptions.headers?.get("foo")).toBe("bar");
-    expect(createOptions.headers?.get("content-type")).toBe("application/json");
-    expect(createOptions.credentials).toBe("include");
-
-    await conversationService.addMessage({ conversationId: "", message: "" });
-    const addOptions = getOptions();
-    expect(addOptions.headers?.get("foo")).toBe("bar");
-    expect(addOptions.headers?.get("content-type")).toBe("application/json");
-    expect(addOptions.credentials).toBe("include");
-
-    await conversationService.addMessageStreaming({
-      conversationId: "",
-      message: "",
-      onResponseDelta: vi.fn(),
-      onResponseFinished: vi.fn(),
-      onReferences: vi.fn(),
-      onMetadata: vi.fn(),
-    });
-    const addStreamingOptions = getOptions();
-    expect(addStreamingOptions.headers?.get("foo")).toBe("bar");
-    expect(addStreamingOptions.headers?.get("content-type")).toBe(
-      "application/json"
-    );
-    expect(addStreamingOptions.credentials).toBe("include");
-
-    // Updating the mock fetch to return 204, which is used by the following two calls.
-    getOptions = mockFetchResponse({ data: {}, status: 204 });
-    await conversationService.rateMessage({
-      conversationId: "a",
-      messageId: "b",
-      rating: true,
-    });
-    const ratingOptions = getOptions();
-    expect(ratingOptions.headers?.get("foo")).toBe("bar");
-    expect(ratingOptions.headers?.get("content-type")).toBe("application/json");
-    expect(ratingOptions.credentials).toBe("include");
-
-    await conversationService.commentMessage({
-      conversationId: "",
-      comment: "",
-      messageId: "",
-    });
-    const commentOptions = getOptions();
-    expect(commentOptions.headers?.get("foo")).toBe("bar");
-    expect(commentOptions.headers?.get("content-type")).toBe(
-      "application/json"
-    );
-    expect(commentOptions.credentials).toBe("include");
-  });
-
   it("throws on invalid inputs", async () => {
     const conversationId = "650b4b260f975ef031016c8d";
     const messageId = "650b4be0d5a57dd66be2ccb9";
@@ -789,6 +719,172 @@ describe("ConversationService", () => {
       });
     }).rejects.toThrow(internalServerErrorMessage);
   });
+
+  describe("fetchOptions", () => {
+    it("supports static fetchOptions", async () => {
+      const conversationService = new ConversationService({
+        serverUrl,
+        fetchOptions: {
+          headers: new Headers({
+            foo: "bar",
+          }),
+          credentials: "include",
+        },
+      });
+      let getOptions = mockFetchResponse({
+        data: {
+          _id: "650b4b260f975ef031016c8d",
+          messages: [],
+          createdAt: new Date().getTime(),
+        },
+      });
+      await conversationService.createConversation();
+      const createOptions = getOptions();
+      const createOptionsHeaders = new Headers(createOptions.headers);
+
+      expect(createOptionsHeaders.get("foo")).toBe("bar");
+      expect(createOptionsHeaders.get("content-type")).toBe("application/json");
+      expect(createOptions.credentials).toBe("include");
+
+      await conversationService.addMessage({ conversationId: "", message: "" });
+      const addOptions = getOptions();
+      const addOptionsHeaders = new Headers(addOptions.headers);
+      expect(addOptionsHeaders.get("foo")).toBe("bar");
+      expect(addOptionsHeaders.get("content-type")).toBe("application/json");
+      expect(addOptions.credentials).toBe("include");
+
+      await conversationService.addMessageStreaming({
+        conversationId: "",
+        message: "",
+        onResponseDelta: vi.fn(),
+        onResponseFinished: vi.fn(),
+        onReferences: vi.fn(),
+        onMetadata: vi.fn(),
+      });
+      const addStreamingOptions = getOptions();
+      const addStreamingOptionsHeaders = new Headers(
+        addStreamingOptions.headers
+      );
+      expect(addStreamingOptionsHeaders.get("foo")).toBe("bar");
+      expect(addStreamingOptionsHeaders.get("content-type")).toBe(
+        "application/json"
+      );
+      expect(addStreamingOptions.credentials).toBe("include");
+
+      // Updating the mock fetch to return 204, which is used by the following two calls.
+      getOptions = mockFetchResponse({ data: {}, status: 204 });
+      await conversationService.rateMessage({
+        conversationId: "a",
+        messageId: "b",
+        rating: true,
+      });
+      const ratingOptions = getOptions();
+      const ratingOptionsHeaders = new Headers(ratingOptions.headers);
+      expect(ratingOptionsHeaders.get("foo")).toBe("bar");
+      expect(ratingOptionsHeaders.get("content-type")).toBe("application/json");
+      expect(ratingOptions.credentials).toBe("include");
+
+      await conversationService.commentMessage({
+        conversationId: "",
+        comment: "",
+        messageId: "",
+      });
+      const commentOptions = getOptions();
+      const commentOptionsHeaders = new Headers(commentOptions.headers);
+      expect(commentOptionsHeaders.get("foo")).toBe("bar");
+      expect(commentOptionsHeaders.get("content-type")).toBe(
+        "application/json"
+      );
+      expect(commentOptions.credentials).toBe("include");
+    });
+
+    it("generates new headers on every request with dynamic fetchOptions", async () => {
+      let callCount = 0;
+      const conversationService = new ConversationService({
+        serverUrl,
+        fetchOptions: () => ({
+          headers: new Headers({
+            foo: `dynamic-${++callCount}`,
+          }),
+          credentials: "include",
+        }),
+      });
+      let getOptions = mockFetchResponse({
+        data: {
+          _id: "650b4b260f975ef031016c8d",
+          messages: [],
+          createdAt: new Date().getTime(),
+        },
+      });
+      await conversationService.createConversation();
+      const createOptions = getOptions();
+      const createOptionsHeaders = new Headers(createOptions.headers);
+
+      expect(createOptionsHeaders.get("foo")).toBe("dynamic-1");
+      expect(createOptionsHeaders.get("content-type")).toBe("application/json");
+      expect(createOptions.credentials).toBe("include");
+
+      await conversationService.addMessage({ conversationId: "", message: "" });
+      const addOptions = getOptions();
+      const addOptionsHeaders = new Headers(addOptions.headers);
+      expect(addOptionsHeaders.get("foo")).toBe("dynamic-2");
+      expect(addOptionsHeaders.get("content-type")).toBe("application/json");
+      expect(addOptions.credentials).toBe("include");
+
+      // For addMessageStreaming, we need to mock fetchEventSource to capture headers
+      let streamingOptions: any;
+      (FetchEventSource.fetchEventSource as jest.Mock).mockImplementation(
+        async (_url, options) => {
+          streamingOptions = options;
+          return {
+            ok: true,
+            status: 200,
+            headers: new Headers({ "content-type": "text/event-stream" }),
+          };
+        }
+      );
+
+      await conversationService.addMessageStreaming({
+        conversationId: "",
+        message: "",
+        onResponseDelta: vi.fn(),
+        onResponseFinished: vi.fn(),
+        onReferences: vi.fn(),
+        onMetadata: vi.fn(),
+      });
+
+      // Check the headers from fetchEventSource
+      expect(streamingOptions.headers.foo).toBe("dynamic-3");
+      expect(streamingOptions.headers["content-type"]).toBe("application/json");
+      expect(streamingOptions.credentials).toBe("include");
+
+      // Updating the mock fetch to return 204, which is used by the following two calls.
+      getOptions = mockFetchResponse({ data: {}, status: 204 });
+      await conversationService.rateMessage({
+        conversationId: "a",
+        messageId: "b",
+        rating: true,
+      });
+      const ratingOptions = getOptions();
+      const ratingOptionsHeaders = new Headers(ratingOptions.headers);
+      expect(ratingOptionsHeaders.get("foo")).toBe("dynamic-4");
+      expect(ratingOptionsHeaders.get("content-type")).toBe("application/json");
+      expect(ratingOptions.credentials).toBe("include");
+
+      await conversationService.commentMessage({
+        conversationId: "",
+        comment: "",
+        messageId: "",
+      });
+      const commentOptions = getOptions();
+      const commentOptionsHeaders = new Headers(commentOptions.headers);
+      expect(commentOptionsHeaders.get("foo")).toBe("dynamic-5");
+      expect(commentOptionsHeaders.get("content-type")).toBe(
+        "application/json"
+      );
+      expect(commentOptions.credentials).toBe("include");
+    });
+  });
 });
 
 describe("getCustomRequestOrigin", () => {