diff --git a/packages/ui/src/stores/session-actions.ts b/packages/ui/src/stores/session-actions.ts index 3fc5178..9d50394 100644 --- a/packages/ui/src/stores/session-actions.ts +++ b/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 + 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", - message: error?.message || "Request failed", + : failedProvider === "qwen-oauth" + ? "Qwen request failed" + : "Request failed", + message: normalizedErrorMessage, variant: "error", duration: 8000, }) diff --git a/packages/ui/src/stores/session-api.ts b/packages/ui/src/stores/session-api.ts index e6dd5e5..ac376bc 100644 --- a/packages/ui/src/stores/session-api.ts +++ b/packages/ui/src/stores/session-api.ts @@ -40,6 +40,20 @@ import { getUserScopedKey } from "../lib/user-storage" const log = getLogger("api") type ProviderMap = Map +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(url: string): Promise { 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 { const token = getStoredAntigravityToken() + const projectId = getStoredAntigravityProjectId() const headers: Record = { "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 { 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 { } } + 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)