feat: enhance model fetching and tool call handling in AI components

- Updated the model fetching logic to include persistent storage for improved performance and reliability.
- Introduced background model list updates to ensure the UI reflects the latest data.
- Enhanced the handling of parameterless tool calls in the AI provider, injecting default arguments when necessary.
- Refactored the chat adapter to support session rollbacks, improving user experience during message regeneration.

These changes significantly improve the functionality and responsiveness of the AI components.

Author: buttercannfly

packages/aipex-react/src/components/chatbot/components/input-area.tsx

@@ -2,7 +2,7 @@ import type { ChatStatus } from "ai";
 import { ClockIcon } from "lucide-react";
 import { useCallback, useEffect, useMemo, useState } from "react";
 import { useTranslation } from "../../../i18n/context";
-import { fetchModelsForSelector } from "../../../lib/models";
+import { fetchModelsForSelector, onModelListChange } from "../../../lib/models";
 import { cn } from "../../../lib/utils";
 import type { ContextItem, InputAreaProps } from "../../../types";
 import {
@@ -91,8 +91,19 @@ export function DefaultInputArea({
           setIsLoadingModels(false);
         }
       });
+
+    // Subscribe to background model list updates (e.g. server returned newer data)
+    const unsubscribe = onModelListChange((updatedModels) => {
+      if (!cancelled) {
+        setFetchedModels(
+          updatedModels.map((m) => ({ name: m.name, value: m.id })),
+        );
+      }
+    });
+
     return () => {
       cancelled = true;
+      unsubscribe();
     };
   }, []);
 

packages/aipex-react/src/hooks/use-chat.ts

@@ -285,7 +285,7 @@ export function useChat(
       return;
     }
 
-    // Remove last assistant message
+    // Remove last assistant message from UI
     const removed = adapter.removeLastAssistantMessage();
     if (!removed) return;
 
@@ -302,6 +302,9 @@ export function useChat(
     const text = textPart?.type === "text" ? textPart.text : "";
 
     if (sessionId && text) {
+      // Roll back the session so the agent doesn't see the old assistant turn
+      await agent.rollbackLastAssistantTurn(sessionId);
+
       adapter.setStatus("submitted");
       const events = agent.chat(text, { sessionId });
       await processAgentEvents(events);

packages/aipex-react/src/lib/models.ts

@@ -70,8 +70,10 @@ const FALLBACK_MODELS: ModelInfo[] = [
 ];
 
 const MODELS_API_URL = "https://www.claudechrome.com/api/models";
+const STORAGE_KEY = "cachedModelList";
+const STORAGE_TIMESTAMP_KEY = "cachedModelListTimestamp";
+const MAX_MODELS = 200;
 
-// Convert API pricing to price level
 function getPriceLevel(
   pricing: ApiModelPricing,
 ): "cheap" | "normal" | "expensive" {
@@ -81,7 +83,6 @@ function getPriceLevel(
   return "expensive";
 }
 
-// Convert API model to internal ModelInfo
 function convertApiModel(apiModel: ApiModel): ModelInfo {
   return {
     id: apiModel.id,
@@ -97,15 +98,13 @@ function convertApiModel(apiModel: ApiModel): ModelInfo {
   };
 }
 
-// Validate that the API response matches the expected schema
 function isValidApiResponse(data: unknown): data is ApiResponse {
   if (typeof data !== "object" || data === null) return false;
   const obj = data as Record<string, unknown>;
   if (typeof obj.success !== "boolean") return false;
   if (typeof obj.data !== "object" || obj.data === null) return false;
   const d = obj.data as Record<string, unknown>;
   if (!Array.isArray(d.models)) return false;
-  // Validate first model shape if present
   if (d.models.length > 0) {
     const first = d.models[0] as Record<string, unknown>;
     if (typeof first.id !== "string" || typeof first.name !== "string") {
@@ -115,53 +114,164 @@ function isValidApiResponse(data: unknown): data is ApiResponse {
   return true;
 }
 
-// Cache for models
+// --- Persistent storage helpers ---
+
+async function loadFromStorage(): Promise<ModelInfo[] | null> {
+  try {
+    if (typeof chrome !== "undefined" && chrome.storage?.local) {
+      const result = await chrome.storage.local.get([STORAGE_KEY]);
+      const models = result[STORAGE_KEY];
+      if (Array.isArray(models) && models.length > 0) {
+        return models as ModelInfo[];
+      }
+    }
+  } catch {
+    // Storage not available (e.g. in tests)
+  }
+  return null;
+}
+
+async function saveToStorage(
+  models: ModelInfo[],
+  serverTimestamp: number,
+): Promise<void> {
+  try {
+    if (typeof chrome !== "undefined" && chrome.storage?.local) {
+      await chrome.storage.local.set({
+        [STORAGE_KEY]: models,
+        [STORAGE_TIMESTAMP_KEY]: serverTimestamp,
+      });
+    }
+  } catch {
+    // Ignore storage errors
+  }
+}
+
+async function getStoredTimestamp(): Promise<number> {
+  try {
+    if (typeof chrome !== "undefined" && chrome.storage?.local) {
+      const result = await chrome.storage.local.get([STORAGE_TIMESTAMP_KEY]);
+      const ts = result[STORAGE_TIMESTAMP_KEY];
+      if (typeof ts === "number") return ts;
+    }
+  } catch {
+    // Ignore
+  }
+  return 0;
+}
+
+// --- In-memory cache (fast path) ---
+
 let cachedModels: ModelInfo[] | null = null;
-let lastFetchTime = 0;
-const CACHE_DURATION = 5 * 60 * 1000; // 5 minutes
-const MAX_MODELS = 200; // Safety cap on number of models
+let cachedServerTimestamp = 0;
+let storageLoaded = false;
 
 /**
- * Fetch models from the server API with caching and fallback.
- * Returns cached result if still valid (5 min TTL).
- * Falls back to FALLBACK_MODELS on any error.
+ * Fetch models with a two-tier cache:
+ * 1. In-memory cache (instant)
+ * 2. chrome.storage.local (survives service worker restarts)
+ *
+ * On the first call, returns storage-cached models immediately.
+ * A background fetch updates both caches when the server reports new data.
  */
 export async function fetchModels(): Promise<ModelInfo[]> {
-  // Return cached models if still valid
-  if (cachedModels && Date.now() - lastFetchTime < CACHE_DURATION) {
+  // 1. Fast path: in-memory cache
+  if (cachedModels) {
+    // Trigger background refresh (fire-and-forget)
+    void refreshFromServer();
     return cachedModels;
   }
 
+  // 2. Try loading from persistent storage
+  if (!storageLoaded) {
+    storageLoaded = true;
+    const stored = await loadFromStorage();
+    if (stored) {
+      cachedModels = stored;
+      cachedServerTimestamp = await getStoredTimestamp();
+      // Trigger background refresh
+      void refreshFromServer();
+      return cachedModels;
+    }
+  }
+
+  // 3. Nothing cached: fetch synchronously and return
+  return await fetchFromServer();
+}
+
+let refreshInFlight = false;
+
+async function refreshFromServer(): Promise<void> {
+  if (refreshInFlight) return;
+  refreshInFlight = true;
   try {
-    const response = await fetch(MODELS_API_URL);
-    console.log("response", response);
+    await fetchFromServer();
+  } finally {
+    refreshInFlight = false;
+  }
+}
 
+async function fetchFromServer(): Promise<ModelInfo[]> {
+  try {
+    const response = await fetch(MODELS_API_URL);
     if (!response.ok) {
       throw new Error(`API request failed: ${response.status}`);
     }
 
     const data: unknown = await response.json();
-    console.log("data", data);
 
     if (!isValidApiResponse(data)) {
       throw new Error("Invalid API response structure");
     }
 
     if (data.success && data.data.models.length > 0) {
-      // Apply safety cap
-      const models = data.data.models.slice(0, MAX_MODELS).map(convertApiModel);
-      cachedModels = models;
-      lastFetchTime = Date.now();
-      return cachedModels;
+      const serverTimestamp = data.data.cache?.lastUpdate ?? Date.now();
+
+      // Only update if the server data is newer
+      if (serverTimestamp > cachedServerTimestamp) {
+        const models = data.data.models
+          .slice(0, MAX_MODELS)
+          .map(convertApiModel);
+        cachedModels = models;
+        cachedServerTimestamp = serverTimestamp;
+        await saveToStorage(models, serverTimestamp);
+        // Notify listeners that models changed
+        notifyModelChange(models);
+      }
+
+      return cachedModels ?? FALLBACK_MODELS;
     }
 
     throw new Error("Empty model list from API");
-  } catch (_error) {
-    // Return fallback - do not log sensitive details
-    return FALLBACK_MODELS;
+  } catch {
+    return cachedModels ?? FALLBACK_MODELS;
   }
 }
 
+// --- Change notification for components ---
+
+type ModelChangeListener = (models: ModelInfo[]) => void;
+const modelChangeListeners = new Set<ModelChangeListener>();
+
+function notifyModelChange(models: ModelInfo[]): void {
+  for (const listener of modelChangeListeners) {
+    try {
+      listener(models);
+    } catch {
+      // Don't let listener errors break the loop
+    }
+  }
+}
+
+/**
+ * Subscribe to model list updates (triggered when server returns new data).
+ * Returns an unsubscribe function.
+ */
+export function onModelListChange(listener: ModelChangeListener): () => void {
+  modelChangeListeners.add(listener);
+  return () => modelChangeListeners.delete(listener);
+}
+
 /**
  * Fetch models and convert to the {name, value} format used by the model selector.
  */

packages/browser-ext/package.json

@@ -30,7 +30,7 @@
   "dependencies": {
     "@ai-sdk/anthropic": "^3.0.44",
     "@ai-sdk/google": "^3.0.22",
-    "@ai-sdk/openai": "^3.0.25",
+    "@ai-sdk/openai": "^3.0.41",
     "@ai-sdk/openai-compatible": "^2.0.18",
     "@aipexstudio/aipex-core": "workspace:*",
     "@aipexstudio/aipex-react": "workspace:*",