feat(ui): auth token storage (#802)

Author: dazy-ds

packages/ragbits-chat/CHANGELOG.md

@@ -1,6 +1,7 @@
 # CHANGELOG
 
 ## Unreleased
+- Add auth token storage and automatic logout on 401 (#802)
 - Improve user settings storage when history is disabled (#799)
 - Remove redundant test for `/api/config` endpoint (#795)
 - Fix bug causing infinite initialization screen (#793)

typescript/@ragbits/api-client-react/tsconfig.json

@@ -3,7 +3,7 @@
         "target": "ES2020",
         "module": "ESNext",
         "lib": ["DOM", "DOM.Iterable", "ESNext"],
-        "moduleResolution": "node",
+        "moduleResolution": "bundler",
         "jsx": "react-jsx",
         "strict": true,
         "declaration": true,

typescript/@ragbits/api-client/README.md

@@ -25,6 +25,12 @@ import { RagbitsClient } from '@ragbits/api-client'
 // Initialize the client
 const client = new RagbitsClient({
     baseUrl: 'http://127.0.0.1:8000', // Optional, defaults to http://127.0.0.1:8000
+    auth: {
+        getToken: () => 'token',
+        onUnauthorized: () => {
+            console.warn('⚠️ Unauthorized')
+        },
+    }, // Optional, used to provide auth details
 })
 
 // Get API configuration
@@ -80,6 +86,7 @@ new RagbitsClient(config?: ClientConfig)
 **Parameters:**
 
 - `config.baseUrl` (optional): Base URL for the API. Defaults to 'http://127.0.0.1:8000'
+- `config.auth` (optional): An object containing authentication details. Provide `getToken` to automatically attach `Authorization: Bearer <token>` to every request. `onUnauthorized` is called if the library encounters a 401 status code.
 
 **Throws:** `Error` if the base URL is invalid
 
@@ -172,8 +179,8 @@ const cleanup = client.makeStreamRequest(
     {
         message: 'Tell me about AI',
         history: [
-            { role: 'user', content: 'Hello', id: 'msg-1' },
-            { role: 'assistant', content: 'Hi there!', id: 'msg-2' },
+            { role: 'user', content: 'Hello' },
+            { role: 'assistant', content: 'Hi there!' },
         ],
         context: { user_id: 'user-123' },
     },

typescript/@ragbits/api-client/src/index.ts

@@ -15,6 +15,7 @@ import type {
  */
 export class RagbitsClient {
     private readonly baseUrl: string
+    private readonly auth: ClientConfig['auth']
     private chunkQueue: Map<
         string,
         {
@@ -29,6 +30,7 @@ export class RagbitsClient {
      */
     constructor(config: ClientConfig = {}) {
         this.baseUrl = config.baseUrl ?? ''
+        this.auth = config.auth
 
         if (this.baseUrl.endsWith('/')) {
             this.baseUrl = this.baseUrl.slice(0, -1)
@@ -71,13 +73,25 @@ export class RagbitsClient {
         url: string,
         options: RequestInit = {}
     ): Promise<Response> {
-        const defaultOptions: RequestInit = {
-            headers: {
-                'Content-Type': 'application/json',
-            },
+        const defaultHeaders: Record<string, string> = {
+            'Content-Type': 'application/json',
+        }
+
+        const headers = {
+            ...defaultHeaders,
+            ...this.normalizeHeaders(options.headers),
+        }
+
+        if (this.auth?.getToken) {
+            headers['Authorization'] = `Bearer ${this.auth.getToken()}`
+        }
+
+        const response = await fetch(url, { ...options, headers })
+
+        if (response.status === 401) {
+            this.auth?.onUnauthorized?.()
         }
 
-        const response = await fetch(url, { ...defaultOptions, ...options })
         if (!response.ok) {
             throw new Error(`HTTP error! status: ${response.status}`)
         }
@@ -206,20 +220,34 @@ export class RagbitsClient {
 
         const startStream = async (): Promise<void> => {
             try {
+                const defaultHeaders: Record<string, string> = {
+                    'Content-Type': 'application/json',
+                    Accept: 'text/event-stream',
+                }
+
+                const headers = {
+                    ...defaultHeaders,
+                    ...customHeaders,
+                }
+
+                if (this.auth?.getToken) {
+                    headers['Authorization'] = `Bearer ${this.auth.getToken()}`
+                }
+
                 const response = await fetch(
                     this._buildApiUrl(endpoint.toString()),
                     {
                         method: 'POST',
-                        headers: {
-                            'Content-Type': 'application/json',
-                            Accept: 'text/event-stream',
-                            ...customHeaders,
-                        },
+                        headers,
                         body: JSON.stringify(data),
                         signal,
                     }
                 )
 
+                if (response.status === 401) {
+                    this.auth?.onUnauthorized?.()
+                }
+
                 if (!response.ok) {
                     throw new Error(`HTTP error! status: ${response.status}`)
                 }
@@ -254,6 +282,17 @@ export class RagbitsClient {
         }
     }
 
+    private normalizeHeaders(init?: HeadersInit): Record<string, string> {
+        if (!init) return {}
+        if (init instanceof Headers) {
+            return Object.fromEntries(init.entries())
+        }
+        if (Array.isArray(init)) {
+            return Object.fromEntries(init)
+        }
+        return init
+    }
+
     private async handleChunkedContent<T>(
         data: T,
         callbacks: StreamCallbacks<T>

typescript/@ragbits/api-client/src/types.ts

@@ -18,6 +18,10 @@ export interface GenericResponse {
  */
 export interface ClientConfig {
     baseUrl?: string
+    auth?: {
+        getToken?: () => string
+        onUnauthorized?: () => Promise<void> | void
+    }
 }
 
 /**

typescript/@ragbits/api-client/tsconfig.json

@@ -3,7 +3,7 @@
         "target": "ES2020",
         "module": "ESNext",
         "lib": ["DOM", "DOM.Iterable", "ESNext"],
-        "moduleResolution": "node",
+        "moduleResolution": "bundler",
         "strict": true,
         "declaration": true,
         "esModuleInterop": true,

typescript/ui/__tests__/integration/intergation.test.tsx

@@ -35,6 +35,7 @@ import { createHistoryStore } from "../../src/core/stores/HistoryStore/historySt
 import { createStore } from "zustand";
 import { useHistoryStore } from "../../src/core/stores/HistoryStore/useHistoryStore";
 import { HistoryStore } from "../../src/types/history";
+import { API_URL } from "../../src/config";
 
 vi.mock("../../src/core/stores/HistoryStore/useHistoryStore", () => {
   return {
@@ -51,6 +52,7 @@ vi.mock("idb-keyval", () => ({
 }));
 
 const historyStore = createStore(createHistoryStore);
+const ragbitsClient = new RagbitsClient({ baseUrl: API_URL });
 historyStore.getState()._internal._setHasHydrated(true);
 
 (useHistoryStore as Mock).mockImplementation(
@@ -72,7 +74,9 @@ describe("Integration tests", () => {
           "makeStreamRequest",
         );
         await act(() => {
-          historyStore.getState().actions.sendMessage("Test message");
+          historyStore
+            .getState()
+            .actions.sendMessage("Test message", ragbitsClient);
         });
 
         expect(makeStreamRequestSpy).toHaveBeenCalledWith(
@@ -84,7 +88,6 @@ describe("Integration tests", () => {
           },
           expect.anything(), // We don't care about callbacks
           expect.anything(), // We don't care about AbortSignal
-          expect.anything(), // We don't care about headers
         );
 
         await waitFor(
@@ -106,7 +109,9 @@ describe("Integration tests", () => {
           "makeStreamRequest",
         );
         await act(() => {
-          historyStore.getState().actions.sendMessage("Test message 2");
+          historyStore
+            .getState()
+            .actions.sendMessage("Test message 2", ragbitsClient);
         });
 
         expect(makeStreamRequestSpy).toHaveBeenCalledWith(
@@ -128,7 +133,6 @@ describe("Integration tests", () => {
           },
           expect.anything(), // We don't care about callbacks
           expect.anything(), // We don't care about AbortSignal
-          expect.anything(), // We don't care about headers
         );
 
         await waitFor(
@@ -155,7 +159,7 @@ describe("Integration tests", () => {
             <ConfigContextProvider>
               <PromptInput
                 isLoading={false}
-                submit={sendMessage}
+                submit={(text) => sendMessage(text, ragbitsClient)}
                 stopAnswering={stopAnswering}
                 followupMessages={getCurrentConversation().followupMessages}
               />
@@ -218,7 +222,6 @@ describe("Integration tests", () => {
           },
           expect.anything(), // We don't care about callbacks
           expect.anything(), // We don't care about AbortSignal
-          expect.anything(), // We don't care about headers
         );
         await waitFor(
           () => {
@@ -256,7 +259,9 @@ describe("Integration tests", () => {
       );
 
       await act(() => {
-        historyStore.getState().actions.sendMessage("Test message");
+        historyStore
+          .getState()
+          .actions.sendMessage("Test message", ragbitsClient);
       });
 
       await waitFor(

typescript/ui/__tests__/unit/plugin/AuthPlugin/AuthGuard.test.tsx

@@ -50,9 +50,31 @@ describe("AuthGuard", () => {
     mockPathname = "/"; // default
   });
 
+  it("renders Initialization screen when not hydrated", () => {
+    mockPathname = "/login";
+    useStoreMock.mockImplementation((_, selector) =>
+      selector({
+        hasHydrated: false,
+      }),
+    );
+
+    render(
+      <AuthGuard>
+        <div data-testid="child">Login page child</div>
+      </AuthGuard>,
+    );
+
+    expect(screen.getByText("Initializing...")).toBeInTheDocument();
+  });
+
   it("renders children if route is /login regardless of auth state", () => {
     mockPathname = "/login";
-    useStoreMock.mockReturnValue(false);
+    useStoreMock.mockImplementation((_, selector) =>
+      selector({
+        hasHydrated: true,
+        isAuthenticated: false,
+      }),
+    );
 
     render(
       <AuthGuard>
@@ -67,7 +89,15 @@ describe("AuthGuard", () => {
 
   it("wraps children and renders AuthWatcher if authenticated", () => {
     mockPathname = "/dashboard";
-    useStoreMock.mockReturnValue(true);
+    useStoreMock.mockImplementation((_, selector) =>
+      selector({
+        hasHydrated: true,
+        isAuthenticated: true,
+        token: {
+          access_token: "token",
+        },
+      }),
+    );
 
     render(
       <AuthGuard>
@@ -83,7 +113,12 @@ describe("AuthGuard", () => {
 
   it("renders Navigate to /login if not authenticated", () => {
     mockPathname = "/dashboard";
-    useStoreMock.mockReturnValue(false);
+    useStoreMock.mockImplementation((_, selector) =>
+      selector({
+        hasHydrated: true,
+        isAuthenticated: false,
+      }),
+    );
 
     render(
       <AuthGuard>