@astra-code/astra-ai 0.1.6 → 0.1.7

package/src/lib/backendClient.ts

@@ -7,7 +7,7 @@ import {
 } from "./config.js";
 import type {AgentEvent, AuthSession, ChatMessage} from "../types/events.js";
 import type {WorkspaceFile} from "./workspaceScanner.js";
-import {loadSession} from "./sessionStore.js";
+import {loadSession, saveSession} from "./sessionStore.js";
 
 type JsonRecord = Record<string, unknown>;
 
@@ -22,14 +22,55 @@ export type SessionSummary = {
 export class BackendClient {
   private readonly baseUrl: string;
   private authToken: string | null;
+  private refreshToken: string | null;
+  private onTokenRefreshed: ((session: AuthSession) => void) | null = null;
 
   public constructor(baseUrl = getBackendUrl()) {
     this.baseUrl = baseUrl;
-    this.authToken = loadSession()?.access_token?.trim() || null;
+    const stored = loadSession();
+    this.authToken = stored?.access_token?.trim() || null;
+    this.refreshToken = stored?.refresh_token?.trim() || null;
   }
 
   public setAuthSession(session: AuthSession | null): void {
     this.authToken = session?.access_token?.trim() || null;
+    this.refreshToken = session?.refresh_token?.trim() || null;
+  }
+
+  public setOnTokenRefreshed(cb: (session: AuthSession) => void): void {
+    this.onTokenRefreshed = cb;
+  }
+
+  private async tryRefreshToken(): Promise<boolean> {
+    if (!this.refreshToken) return false;
+    try {
+      const res = await fetch(`${this.baseUrl}/api/auth/refresh`, {
+        method: "POST",
+        headers: {"content-type": "application/json"},
+        body: JSON.stringify({refresh_token: this.refreshToken})
+      });
+      if (!res.ok) return false;
+      const data = (await res.json()) as Record<string, unknown>;
+      const newAccess = typeof data.access_token === "string" ? data.access_token.trim() : null;
+      if (!newAccess) return false;
+      const newRefresh = typeof data.refresh_token === "string" ? data.refresh_token.trim() : this.refreshToken;
+      this.authToken = newAccess;
+      this.refreshToken = newRefresh;
+      // Merge with stored session and persist
+      const stored = loadSession();
+      if (stored) {
+        const refreshed: AuthSession = {
+          ...stored,
+          access_token: newAccess,
+          refresh_token: newRefresh ?? undefined
+        };
+        saveSession(refreshed);
+        this.onTokenRefreshed?.(refreshed);
+      }
+      return true;
+    } catch {
+      return false;
+    }
   }
 
   public async get(path: string, params?: Record<string, string>): Promise<JsonRecord> {
@@ -209,6 +250,7 @@ export class BackendClient {
     workspaceTree?: string[];
     workspaceFiles?: WorkspaceFile[];
     model?: string;
+    signal?: AbortSignal;
   }): AsyncGenerator<AgentEvent, void, void> {
     const model = payload.model ?? getDefaultModel();
     const body = {
@@ -230,20 +272,33 @@ export class BackendClient {
       user_id: payload.user.user_id,
       org_id: payload.user.org_id
     };
-    return this.streamSse("/api/agent/chat/stream", body);
+    return this.streamSse("/api/agent/chat/stream", body, payload.signal);
   }
 
   private async request(url: string, method: string, payload?: JsonRecord): Promise<JsonRecord> {
-    const headers: Record<string, string> = {"content-type": "application/json"};
-    if (this.authToken) {
-      headers.authorization = `Bearer ${this.authToken}`;
-    }
-    const response = await fetch(url, {
+    const makeHeaders = (): Record<string, string> => {
+      const h: Record<string, string> = {"content-type": "application/json"};
+      if (this.authToken) h.authorization = `Bearer ${this.authToken}`;
+      return h;
+    };
+
+    let response = await fetch(url, {
       method,
-      headers,
+      headers: makeHeaders(),
       body: payload ? JSON.stringify(payload) : null
     });
 
+    if (response.status === 401) {
+      const refreshed = await this.tryRefreshToken();
+      if (refreshed) {
+        response = await fetch(url, {
+          method,
+          headers: makeHeaders(),
+          body: payload ? JSON.stringify(payload) : null
+        });
+      }
+    }
+
     if (!response.ok) {
       const detail = (await response.text()).trim();
       throw new Error(`Backend error ${response.status}: ${detail || response.statusText}`);
@@ -256,17 +311,33 @@ export class BackendClient {
     return JSON.parse(text) as JsonRecord;
   }
 
-  private async *streamSse(path: string, payload: JsonRecord): AsyncGenerator<AgentEvent> {
-    const headers: Record<string, string> = {"content-type": "application/json"};
-    if (this.authToken) {
-      headers.authorization = `Bearer ${this.authToken}`;
-    }
-    const response = await fetch(`${this.baseUrl}${path}`, {
+  private async *streamSse(path: string, payload: JsonRecord, signal?: AbortSignal): AsyncGenerator<AgentEvent> {
+    const makeHeaders = (): Record<string, string> => {
+      const h: Record<string, string> = {"content-type": "application/json"};
+      if (this.authToken) h.authorization = `Bearer ${this.authToken}`;
+      return h;
+    };
+    const sig = signal ?? null;
+
+    let response = await fetch(`${this.baseUrl}${path}`, {
       method: "POST",
-      headers,
-      body: JSON.stringify(payload)
+      headers: makeHeaders(),
+      body: JSON.stringify(payload),
+      signal: sig
     });
 
+    if (response.status === 401) {
+      const refreshed = await this.tryRefreshToken();
+      if (refreshed) {
+        response = await fetch(`${this.baseUrl}${path}`, {
+          method: "POST",
+          headers: makeHeaders(),
+          body: JSON.stringify(payload),
+          signal: sig
+        });
+      }
+    }
+
     if (!response.ok) {
       const detail = (await response.text()).trim();
       throw new Error(`Backend error ${response.status}: ${detail || response.statusText}`);

package/src/types/events.ts

@@ -5,6 +5,7 @@ export type AuthSession = {
   role?: string;
   display_name?: string;
   access_token?: string;
+  refresh_token?: string;
   [key: string]: unknown;
 };