chore: minor fixes in session-actions and session-api

packages/ui/src/stores/session-actions.ts

@@ -19,7 +19,14 @@ import {
   clearCompactionSuggestion,
   type CompactionResult,
 } from "./session-compaction"
-import { createSession, loadMessages, getStoredAntigravityToken, isAntigravityTokenValid } from "./session-api"
+import {
+  ANTIGRAVITY_MODEL_IDS,
+  createSession,
+  getStoredAntigravityProjectId,
+  getStoredAntigravityToken,
+  isAntigravityTokenValid,
+  loadMessages,
+} from "./session-api"
 import { showToastNotification } from "../lib/notifications"
 import { QwenOAuthManager } from "../lib/integrations/qwen-oauth"
 import { getUserScopedKey } from "../lib/user-storage"
@@ -496,6 +503,7 @@ async function readSseStream(
     if (idleTimer) clearTimeout(idleTimer)
     idleTimer = setTimeout(() => {
       timedOut = true
+      shouldStop = true
       reader.cancel().catch(() => { })
     }, idleTimeoutMs)
   }
@@ -505,9 +513,15 @@ async function readSseStream(
     let chunkCount = 0
     let lastYieldTime = performance.now()
     while (!shouldStop) {
-      const { done, value } = await reader.read()
+      let readResult: ReadableStreamReadResult<Uint8Array>
+      try {
+        readResult = await reader.read()
+      } catch (error) {
+        if (timedOut) break
+        throw error
+      }
+      const { done, value } = readResult
       if (done) break
-      resetIdleTimer()
       buffer += decoder.decode(value, { stream: true })
       const lines = buffer.split("\n")
       buffer = lines.pop() || ""
@@ -517,6 +531,7 @@ async function readSseStream(
         if (!trimmed.startsWith("data:")) continue
         const data = trimmed.slice(5).trim()
         if (!data) continue
+        resetIdleTimer()
         if (data === "[DONE]") {
           shouldStop = true
           break
@@ -1175,6 +1190,10 @@ async function streamAntigravityChat(
     "Content-Type": "application/json",
     Authorization: `Bearer ${token.access_token}`,
   }
+  const projectId = getStoredAntigravityProjectId()
+  if (projectId) {
+    headers["X-Antigravity-Project"] = projectId
+  }
 
   const response = await fetch("/api/antigravity/chat", {
     method: "POST",
@@ -1491,6 +1510,10 @@ async function sendMessage(
   })
 
   const providerId = effectiveModel.providerId
+  const useAntigravity =
+    providerId === "antigravity" ||
+    (providerId === "google" && ANTIGRAVITY_MODEL_IDS.has(effectiveModel.modelId))
+  const routingProviderId = useAntigravity ? "antigravity" : providerId
   const tPre1 = performance.now()
   const systemMessage = await untrack(() => mergeSystemInstructions(instanceId, sessionId, prompt))
   const tPre2 = performance.now()
@@ -1498,7 +1521,7 @@ async function sendMessage(
     addDebugLog(`Merge System Instructions: ${Math.round(tPre2 - tPre1)}ms`, "warn")
   }
 
-  if (providerId === "ollama-cloud" || providerId === "qwen-oauth" || providerId === "opencode-zen" || providerId === "zai" || providerId === "antigravity") {
+  if (providerId === "ollama-cloud" || providerId === "qwen-oauth" || providerId === "opencode-zen" || providerId === "zai" || useAntigravity) {
     const store = messageStoreBus.getOrCreate(instanceId)
     const now = Date.now()
     const assistantMessageId = createId("msg")
@@ -1530,7 +1553,7 @@ async function sendMessage(
       store.setMessageInfo(assistantMessageId, {
         id: assistantMessageId,
         role: "assistant",
-        providerID: effectiveModel.providerId,
+        providerID: routingProviderId,
         modelID: effectiveModel.modelId,
         time: { created: now, completed: 0 },
       } as any)
@@ -1582,11 +1605,11 @@ async function sendMessage(
           assistantMessageId,
           assistantPartId,
         )
-      } else if (providerId === "antigravity") {
+      } else if (useAntigravity) {
         assistantText = await streamAntigravityChat(
           instanceId,
           sessionId,
-          providerId,
+          routingProviderId,
           effectiveModel.modelId,
           externalMessages,
           messageId,
@@ -1695,26 +1718,33 @@ async function sendMessage(
         updatedAt: Date.now(),
         isEphemeral: false,
       })
+      const rawErrorMessage = error?.message || "Request failed"
+      const normalizedErrorMessage = /aborted|abort/i.test(rawErrorMessage)
+        ? "Request timed out. The provider may be unavailable."
+        : rawErrorMessage
       store.setMessageInfo(assistantMessageId, {
         id: assistantMessageId,
         role: "assistant",
-        providerID: effectiveModel.providerId,
+        providerID: routingProviderId,
         modelID: effectiveModel.modelId,
         time: { created: now, completed: Date.now() },
-        error: { name: "UnknownError", message: error?.message || "Request failed" },
+        error: { name: "UnknownError", message: normalizedErrorMessage },
       } as any)
+      const failedProvider = useAntigravity ? "antigravity" : providerId
       showToastNotification({
         title:
-          providerId === "ollama-cloud"
+          failedProvider === "ollama-cloud"
             ? "Ollama request failed"
-            : providerId === "zai"
+            : failedProvider === "zai"
               ? "Z.AI request failed"
-              : providerId === "opencode-zen"
+              : failedProvider === "opencode-zen"
                 ? "OpenCode Zen request failed"
-                : providerId === "antigravity"
+                : failedProvider === "antigravity"
                   ? "Antigravity request failed"
-                  : "Qwen request failed",
+                  : failedProvider === "qwen-oauth"
-        message: error?.message || "Request failed",
+                    ? "Qwen request failed"
+                    : "Request failed",
+        message: normalizedErrorMessage,
         variant: "error",
         duration: 8000,
       })

packages/ui/src/stores/session-api.ts

@@ -40,6 +40,20 @@ import { getUserScopedKey } from "../lib/user-storage"
 const log = getLogger("api")
 
 type ProviderMap = Map<string, Provider>
+export const ANTIGRAVITY_MODEL_IDS = new Set([
+  "gemini-3-pro-low",
+  "gemini-3-pro-high",
+  "gemini-3-flash",
+  "claude-sonnet-4-5",
+  "claude-sonnet-4-5-thinking-low",
+  "claude-sonnet-4-5-thinking-medium",
+  "claude-sonnet-4-5-thinking-high",
+  "claude-opus-4-5-thinking-low",
+  "claude-opus-4-5-thinking-medium",
+  "claude-opus-4-5-thinking-high",
+  "gpt-oss-120b-medium",
+])
+const ANTIGRAVITY_PROJECT_KEY = "antigravity_project_id"
 
 async function fetchJson<T>(url: string): Promise<T | null> {
   try {
@@ -271,12 +285,26 @@ export function isAntigravityTokenValid(token: { expires_in: number; created_at:
   return Date.now() < expiresAt
 }
 
+export function getStoredAntigravityProjectId(): string | undefined {
+  if (typeof window === "undefined") return undefined
+  try {
+    const value = window.localStorage.getItem(getUserScopedKey(ANTIGRAVITY_PROJECT_KEY))
+    return value && value.trim().length > 0 ? value.trim() : undefined
+  } catch {
+    return undefined
+  }
+}
+
 async function fetchAntigravityProvider(): Promise<Provider | null> {
   const token = getStoredAntigravityToken()
+  const projectId = getStoredAntigravityProjectId()
   const headers: Record<string, string> = { "Content-Type": "application/json" }
   if (token?.access_token) {
     headers["Authorization"] = `Bearer ${token.access_token}`
   }
+  if (projectId) {
+    headers["X-Antigravity-Project"] = projectId
+  }
 
   try {
     const response = await fetch("/api/antigravity/models", { headers })
@@ -295,7 +323,7 @@ async function fetchAntigravityProvider(): Promise<Provider | null> {
 
     return {
       id: "antigravity",
-      name: "Antigravity (Google OAuth)",
+      name: "Antigravity",
       models: models.map((model) => ({
         id: model.id,
         name: model.name,
@@ -994,14 +1022,37 @@ async function fetchProviders(instanceId: string): Promise<void> {
       }
     }
 
+    const normalizedProviders = providerList
+      .map((provider) => {
+        if (provider.id !== "google") return provider
+        const filteredModels = provider.models.filter((model: Model) => ANTIGRAVITY_MODEL_IDS.has(model.id))
+        if (filteredModels.length === 0) return null
+        const defaultModelId = filteredModels.some((model: Model) => model.id === provider.defaultModelId)
+          ? provider.defaultModelId
+          : filteredModels[0]?.id
+        return {
+          ...provider,
+          name: "Antigravity",
+          defaultModelId,
+          models: filteredModels,
+        }
+      })
+      .filter(Boolean) as typeof providerList
+
     // Filter out Z.AI providers from SDK to use our custom routing with full message history
-    const filteredBaseProviders = providerList.filter((provider) =>
+    const filteredBaseProviders = normalizedProviders.filter((provider) =>
       !provider.id.toLowerCase().includes("zai") &&
       !provider.id.toLowerCase().includes("z.ai") &&
       !provider.id.toLowerCase().includes("glm")
     )
 
-    const extraProviders = await fetchExtraProviders()
+    let extraProviders = await fetchExtraProviders()
+    if (!isNative) {
+      const hasSdkAntigravity = normalizedProviders.some((provider) => provider.id === "google")
+      if (hasSdkAntigravity) {
+        extraProviders = extraProviders.filter((provider) => provider.id !== "antigravity")
+      }
+    }
     const baseProviders = removeDuplicateProviders(filteredBaseProviders, extraProviders)
     const mergedProviders = mergeProviders(baseProviders, extraProviders)