diff --git a/README.md b/README.md index b6f3503d..d35d0e77 100644 --- a/README.md +++ b/README.md @@ -124,6 +124,7 @@ Read the [configuration docs here](/_docs/config/index.mdx). I tried crabcode specifically for these providers: - [x] **openai** (both API key and OAuth, thank you OpenAI for supporting OSS harnesses!) +- [x] **Google Antigravity / Gemini** (browser OAuth with a pasted callback URL) - [x] **xAI / Grok** (API key and SuperGrok/X Premium OAuth, thank you xAI for openly supporting OSS harnesses: based on [OpenClaw](https://x.ai/news/grok-openclaw), [OpenCode](https://x.ai/news/grok-opencode), [KiloCode](https://x.ai/news/grok-kilocode), [Hermes](https://x.ai/news/grok-hermes)) - [x] **opencode-zen** and **opencode-go** - [x] **nano-gpt** @@ -143,9 +144,32 @@ I tried crabcode specifically for these providers: > I might work harder to support these in the future. -- Gemini - It's OAuth + also very unsure. So currently no. - Claude Code Subscription - Known to explicitly not like harnesses. So never will, sorry. +### Google Antigravity sign-in + +Before starting Crabcode, provide an authorized OAuth application's client +credentials through `CRABCODE_ANTIGRAVITY_CLIENT_ID` and +`CRABCODE_ANTIGRAVITY_CLIENT_SECRET`. Both are required at runtime for browser +sign-in and token refresh. There are no bundled client credentials or build-time +fallbacks. Keep these values in your local environment or secret manager, not +in committed configuration or command-line arguments. The application's +registered redirect URI must be `http://localhost:51121/oauth-callback`. + +The project maintainer must arrange an OAuth client authorized for the required +Google scopes and Antigravity service. An arbitrary Google Cloud OAuth client +does not guarantee access. Existing refresh tokens are tied to their original +client; reconnect if you change clients. + +Run `/connect`, choose **Google Antigravity**, then **Gemini via Antigravity (browser)**. +Open the sign-in link, authorize your account, and paste the full +`http://localhost:51121/oauth-callback?...` URL into the callback input. The +localhost page may fail to load; copy the address bar URL anyway. Press Enter +to submit, or Esc to cancel. Credentials are saved and refreshed automatically. + +This is an unofficial integration. Google account eligibility, service terms, +and account restrictions apply. The API-key option remains available. + ## Performance Like any benchmark, please take this with a grain of salt. I have a cherry-picked benchmark for 1 purpose only: "Is crabcode at least as reliable as codex/opencode?". It's a useful feedback loop for my Crabcode trying to improve itself. We're honestly only chasing for at least: best parity if not better perf while having massively better TUI UX based on my personal preferences. Here is my most recent run (Jul 13, 2026): diff --git a/src/acp/login.rs b/src/acp/login.rs index 5c310485..5b71080f 100644 --- a/src/acp/login.rs +++ b/src/acp/login.rs @@ -20,9 +20,41 @@ pub async fn run(cwd: Option) -> Result<()> { if !io::stdin().is_terminal() || !io::stderr().is_terminal() { bail!("provider login requires an interactive terminal; run crabcode acp --login in a terminal"); } + + fn open_url(url: &str) -> Result<()> { + #[cfg(target_os = "macos")] + { + std::process::Command::new("open") + .arg(url) + .spawn() + .map(|_| ()) + .map_err(Into::into) + } + #[cfg(target_os = "linux")] + { + std::process::Command::new("xdg-open") + .arg(url) + .spawn() + .map(|_| ()) + .map_err(Into::into) + } + #[cfg(target_os = "windows")] + { + std::process::Command::new("cmd") + .args(["/C", "start", "", url]) + .spawn() + .map(|_| ()) + .map_err(Into::into) + } + #[cfg(not(any(target_os = "macos", target_os = "linux", target_os = "windows")))] + { + let _ = url; + Ok(()) + } + } eprintln!("Connect a provider (no chat will be started)."); eprintln!( - "Provider examples: openai, anthropic, xai, ollama; custom provider IDs are also accepted." + "Provider examples: openai, anthropic, xai, google (Google Antigravity), ollama; custom provider IDs are also accepted." ); let provider = prompt("Provider ID (empty to cancel): ")?; validate_provider_id(&provider)?; @@ -61,10 +93,27 @@ pub async fn run(cwd: Option) -> Result<()> { let method = if matches!(provider.as_str(), "openai" | "xai") { eprintln!("1: browser OAuth 2: headless/device OAuth 3: API key"); prompt("Login method (empty to cancel): ")? + } else if provider == "google" { + eprintln!("1: Google Antigravity browser OAuth 3: API key"); + prompt("Login method (empty to cancel): ")? } else { "3".into() }; let credentials = match (provider.as_str(), method.as_str()) { + ("google", "1") => { + let authorization = crate::auth::antigravity_oauth::BrowserAuthorization::new()?; + eprintln!("This unofficial third-party integration uses Antigravity. Google account eligibility, service terms, and account restrictions apply.\nOpen this URL in a browser:\n{}", authorization.authorization_url()); + let _ = open_url(authorization.authorization_url()); + eprintln!("The localhost page may be unavailable; copy its complete URL from the address bar anyway."); + let callback = + read_hidden("Complete callback URL (hidden; Esc to cancel): ", 16 * 1024)?; + oauth_config( + authorization + .complete(&callback) + .await + .map_err(|_| anyhow::anyhow!("OAuth login failed; retry provider login"))?, + ) + } ("openai", "1") => oauth_config( crate::auth::openai_oauth::authorize_browser() .await @@ -137,6 +186,7 @@ fn oauth_config(credentials: crate::auth::OAuthCredentials) -> AuthConfig { expires: credentials.expires, account_id: credentials.account_id, enterprise_url: credentials.enterprise_url, + project_id: credentials.project_id, } } @@ -161,9 +211,13 @@ impl Drop for RawMode { } fn read_api_key() -> Result { - eprint!("API key (hidden; Esc to cancel): "); + read_hidden("API key (hidden; Esc to cancel): ", 16 * 1024) +} + +fn read_hidden(label: &str, max_bytes: usize) -> Result { + eprint!("{label}"); io::stderr().flush()?; - enable_raw_mode().context("failed to hide API key input")?; + enable_raw_mode().context("failed to hide credential input")?; let guard = RawMode; let mut value = String::new(); loop { @@ -182,7 +236,12 @@ fn read_api_key() -> Result { KeyCode::Backspace => { value.pop(); } - KeyCode::Char(ch) if !key.modifiers.contains(KeyModifiers::CONTROL) => value.push(ch), + KeyCode::Char(ch) if !key.modifiers.contains(KeyModifiers::CONTROL) => { + if value.len() + ch.len_utf8() > max_bytes { + bail!("credential input is too long"); + } + value.push(ch); + } _ => {} } } @@ -190,7 +249,7 @@ fn read_api_key() -> Result { eprintln!(); let value = value.trim().to_owned(); if value.is_empty() { - bail!("API key cannot be empty"); + bail!("credential input cannot be empty"); } Ok(value) } @@ -223,6 +282,7 @@ mod tests { expires: 0, account_id: None, enterprise_url: None, + project_id: None, }); assert!(usable_credentials(&auth)); assert!(!usable_credentials(&AuthConfig::OAuth { @@ -231,6 +291,7 @@ mod tests { expires: 0, account_id: None, enterprise_url: None, + project_id: None, })); } } diff --git a/src/acp/service.rs b/src/acp/service.rs index 11a35ebb..a6bf6eb7 100644 --- a/src/acp/service.rs +++ b/src/acp/service.rs @@ -1614,6 +1614,9 @@ impl AcpService { &live_usage, )?; } + crate::llm::ChunkMessage::ProviderReplay(item) => { + crate::llm::antigravity::persist_replay(&mut assistant, item); + } crate::llm::ChunkMessage::Reasoning(text) => { assistant.append_reasoning(&text); send_text(connection, session_id, &message_id, text, true)?; diff --git a/src/agent/config.rs b/src/agent/config.rs index a8e1a1b4..e87d189f 100644 --- a/src/agent/config.rs +++ b/src/agent/config.rs @@ -15,6 +15,8 @@ pub struct OpenAIRequestOptions { /// Sticky prompt-cache routing key (Responses / chat-completions). /// Typically the crabcode session id. pub prompt_cache_key: Option, + /// Google Antigravity OAuth project routing identifier. + pub antigravity_project: Option, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -100,6 +102,7 @@ pub enum ProviderKind { OpenAI, OpenAICompatible, Anthropic, + Antigravity, } static LLM_SESSION: OnceLock>> = OnceLock::new(); diff --git a/src/agent/subagent.rs b/src/agent/subagent.rs index 1a42606f..639a36da 100644 --- a/src/agent/subagent.rs +++ b/src/agent/subagent.rs @@ -97,7 +97,11 @@ pub async fn run_subagent( use futures::StreamExt; use std::collections::HashMap; - let mut session = resolve_subagent_session(&agent, parent_session, sender.as_ref()).await?; + let mut session = tokio::select! { + biased; + _ = cancel_token.cancelled() => return Err("Subagent cancelled".to_string()), + session = resolve_subagent_session(&agent, parent_session, sender.as_ref()) => session?, + }; // Child cache key on purpose: subagents have a different system prompt and tool set, // so parent-prefix reuse would miss (and risk sticky-routing pollution). See // SessionAffinity::child_session docs. @@ -443,6 +447,40 @@ async fn start_subagent_stream( ); match session.provider_kind { + ProviderKind::Antigravity => { + let access = session + .api_key + .clone() + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| "Missing Google Antigravity OAuth access token".to_string())?; + let project = session + .openai_options + .antigravity_project + .clone() + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| "Missing Google Antigravity OAuth project".to_string())?; + let provider = crate::llm::antigravity::Antigravity::new( + session.base_url.clone(), + session.model.clone(), + access, + project, + ) + .with_stored_credentials(); + stream_with_tools_options( + provider, + messages, + tools, + max_steps, + None, + headers, + cancel_token, + StreamWithToolsOptions { + prune_tool_outputs: session.prune_tool_outputs, + }, + ) + .await + .map_err(|e| format!("Stream error: {}", e)) + } ProviderKind::OpenAICompatible => { let mut builder = OpenAICompatible::builder() .base_url(&session.base_url) @@ -1088,6 +1126,29 @@ mod tests { ); } + #[test] + fn inherited_parent_session_preserves_antigravity_routing() { + let registry = crate::agent::definition::AgentRegistry::default(); + let agent = registry.task_target("explore").unwrap(); + let mut parent = test_session(None); + parent.provider_kind = crate::agent::config::ProviderKind::Antigravity; + parent.openai_options.antigravity_project = Some("google-project".into()); + parent.base_url = crate::auth::antigravity_oauth::ANTIGRAVITY_ENDPOINT.into(); + let session = tokio_test::block_on(resolve_subagent_session(&agent, parent, None)).unwrap(); + assert_eq!( + session.provider_kind, + crate::agent::config::ProviderKind::Antigravity + ); + assert_eq!( + session.openai_options.antigravity_project.as_deref(), + Some("google-project") + ); + assert_eq!( + session.base_url, + crate::auth::antigravity_oauth::ANTIGRAVITY_ENDPOINT + ); + } + fn test_session( reasoning_effort: Option, ) -> crate::agent::config::LlmSessionConfig { diff --git a/src/app.rs b/src/app.rs index e3ffbc51..b3b67f19 100644 --- a/src/app.rs +++ b/src/app.rs @@ -263,6 +263,7 @@ enum ConnectDialogMode { ProviderSelection, OpenAIMethodSelection, XAIMethodSelection, + GoogleMethodSelection, } #[derive(Debug)] @@ -291,6 +292,7 @@ enum McpOAuthTaskMessage { enum OAuthProvider { OpenAI, XAI, + Google, } impl OAuthProvider { @@ -298,6 +300,7 @@ impl OAuthProvider { match self { Self::OpenAI => "openai", Self::XAI => "xai", + Self::Google => "google", } } @@ -305,6 +308,7 @@ impl OAuthProvider { match self { Self::OpenAI => "OpenAI", Self::XAI => "xAI", + Self::Google => "Google Antigravity", } } @@ -312,6 +316,7 @@ impl OAuthProvider { match self { Self::OpenAI => "Connected OpenAI via ChatGPT Plus/Pro OAuth", Self::XAI => "Connected xAI via Grok OAuth", + Self::Google => "Connected Gemini via Google Antigravity OAuth", } } @@ -319,6 +324,7 @@ impl OAuthProvider { match self { Self::OpenAI => "gpt-5.4", Self::XAI => "grok-build-0.1", + Self::Google => "gemini-3-flash", } } } @@ -933,6 +939,8 @@ pub struct App { pub api_key_input: crate::ui::components::api_key_input::ApiKeyInput, provider_oauth_receiver: Option>, provider_oauth_in_progress: Option, + google_browser_authorization: Option, + google_oauth_task: Option>, mcp_oauth_receiver: Option>, mcp_oauth_in_progress: Option, compaction_receiver: Option>, @@ -1193,6 +1201,8 @@ impl App { api_key_input, provider_oauth_receiver: None, provider_oauth_in_progress: None, + google_browser_authorization: None, + google_oauth_task: None, mcp_oauth_receiver: None, mcp_oauth_in_progress: None, compaction_receiver: None, @@ -4208,6 +4218,18 @@ impl App { self.overlay_focus = OverlayFocus::None; true } + ProviderOAuthFlowAction::OpenLink(url) => { + self.open_oauth_link(&url); + true + } + ProviderOAuthFlowAction::SubmitCallback(url) => { + self.submit_google_oauth_callback(url); + true + } + ProviderOAuthFlowAction::Cancel => { + self.cancel_google_oauth(); + true + } ProviderOAuthFlowAction::CopyLink(url) => { match crate::utils::clipboard::copy_text(&url) { Ok(_) => push_toast(Toast::new( @@ -5762,10 +5784,15 @@ impl App { ProviderOAuthFlowAction::Close => { self.overlay_focus = OverlayFocus::None; } + ProviderOAuthFlowAction::OpenLink(url) => self.open_oauth_link(&url), + ProviderOAuthFlowAction::SubmitCallback(url) => { + self.submit_google_oauth_callback(url); + } + ProviderOAuthFlowAction::Cancel => self.cancel_google_oauth(), ProviderOAuthFlowAction::CopyLink(url) => { match crate::utils::clipboard::copy_text(&url) { Ok(_) => push_toast(Toast::new( - "Copied OpenAI login link", + "Copied OAuth login link", ToastLevel::Info, None, )), @@ -6382,6 +6409,9 @@ impl App { (_, OverlayFocus::ApiKeyInput) => { self.api_key_input.text_area.insert_str(&text); } + (_, OverlayFocus::ProviderOAuthFlow) => { + self.provider_oauth_flow_state.insert_callback(&text); + } (_, OverlayFocus::SuggestionsPopup) => { if self.is_subagent_session_active() { clear_suggestions(&mut self.suggestions_popup_state); @@ -9548,6 +9578,129 @@ impl App { } } + fn show_google_connect_methods(&mut self) { + use crate::ui::components::dialog::DialogItem; + let items = vec![ + DialogItem { + id: "google-oauth-browser".into(), + name: "Gemini via Antigravity (browser)".into(), + group: "Google Antigravity".into(), + description: "Authorize in your browser, then paste the callback URL".into(), + tip: Some("Unofficial integration; Google account restrictions may apply".into()), + provider_id: "google".into(), + active: false, + }, + DialogItem { + id: "google-api-key".into(), + name: "Manually enter API key".into(), + group: "Google Antigravity".into(), + description: "Use a Google Gemini API key".into(), + tip: None, + provider_id: "google".into(), + active: false, + }, + ]; + self.connect_dialog_state = + crate::views::ConnectDialogState::with_items("Connect Google Antigravity", items); + self.connect_dialog_state.dialog.show(); + self.connect_dialog_mode = ConnectDialogMode::GoogleMethodSelection; + self.overlay_focus = OverlayFocus::ConnectDialog; + } + + fn open_oauth_link(&self, url: &str) { + if crate::utils::file_opener::open_url(url).is_err() { + push_toast(Toast::new( + "Could not open browser. Press ctrl+y to copy the sign-in link.", + ToastLevel::Warning, + None, + )); + } + } + + fn begin_google_oauth_browser(&mut self) { + self.begin_google_oauth_browser_with( + crate::auth::antigravity_oauth::BrowserAuthorization::new, + ); + } + + fn begin_google_oauth_browser_with( + &mut self, + authorization: impl FnOnce() -> anyhow::Result< + crate::auth::antigravity_oauth::BrowserAuthorization, + >, + ) { + if self.provider_oauth_in_progress == Some(OAuthProvider::Google) { + // Reopen the pending attempt without replacing its state/verifier. + self.provider_oauth_flow_state.visible = true; + self.overlay_focus = OverlayFocus::ProviderOAuthFlow; + return; + } + if self.provider_oauth_in_progress.is_some() || self.mcp_oauth_in_progress.is_some() { + push_toast(Toast::new( + "Another OAuth login is in progress", + ToastLevel::Info, + None, + )); + return; + } + match authorization() { + Ok(authorization) => { + self.provider_oauth_flow_state.show_browser_callback_for( + OAuthProvider::Google.label(), + authorization.authorization_url(), + ); + self.google_browser_authorization = Some(authorization); + self.provider_oauth_in_progress = Some(OAuthProvider::Google); + self.connect_dialog_mode = ConnectDialogMode::ProviderSelection; + self.connect_dialog_state = init_connect_dialog(); + self.overlay_focus = OverlayFocus::ProviderOAuthFlow; + } + Err(error) => push_toast(Toast::new(error.to_string(), ToastLevel::Error, None)), + } + } + + fn submit_google_oauth_callback(&mut self, callback_url: String) { + let Some(authorization) = self.google_browser_authorization.as_ref() else { + return; + }; + if let Err(error) = authorization.validate_callback(&callback_url) { + self.provider_oauth_flow_state + .set_callback_error(error.to_string()); + return; + } + let authorization = self + .google_browser_authorization + .take() + .expect("validated pending login"); + self.provider_oauth_flow_state.mark_submitting(); + let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); + self.provider_oauth_receiver = Some(receiver); + self.google_oauth_task = Some(tokio::spawn(async move { + let event = match authorization.complete(&callback_url).await { + Ok(credentials) => ProviderOAuthTaskMessage::Success { + provider: OAuthProvider::Google, + credentials, + }, + Err(error) => ProviderOAuthTaskMessage::Failed { + provider: OAuthProvider::Google, + error: error.to_string(), + }, + }; + let _ = sender.send(event); + })); + } + + fn cancel_google_oauth(&mut self) { + self.google_browser_authorization = None; + if let Some(task) = self.google_oauth_task.take() { + task.abort(); + } + self.provider_oauth_receiver = None; + self.provider_oauth_in_progress = None; + self.provider_oauth_flow_state.hide(); + self.overlay_focus = OverlayFocus::None; + } + fn show_openai_connect_methods(&mut self) { use crate::ui::components::dialog::DialogItem; @@ -9739,6 +9892,11 @@ impl App { return; } + if selected_item.id == "google" { + self.show_google_connect_methods(); + return; + } + self.api_key_input.show(&selected_item.id); self.overlay_focus = OverlayFocus::ApiKeyInput; } @@ -9758,6 +9916,15 @@ impl App { self.overlay_focus = OverlayFocus::None; } }, + ConnectDialogMode::GoogleMethodSelection => match selected_item.id.as_str() { + "google-oauth-browser" => self.begin_google_oauth_browser(), + "google-api-key" => { + self.api_key_input.show("google"); + self.connect_dialog_mode = ConnectDialogMode::ProviderSelection; + self.overlay_focus = OverlayFocus::ApiKeyInput; + } + _ => self.overlay_focus = OverlayFocus::None, + }, ConnectDialogMode::XAIMethodSelection => match selected_item.id.as_str() { "xai-oauth-browser" => { self.begin_provider_oauth_browser(OAuthProvider::XAI); @@ -9844,6 +10011,10 @@ impl App { } fn begin_provider_oauth_browser(&mut self, provider: OAuthProvider) { + if provider == OAuthProvider::Google { + self.begin_google_oauth_browser(); + return; + } if let Some(active_provider) = self.provider_oauth_in_progress { push_toast(Toast::new( format!("{} OAuth is already in progress", active_provider.label()), @@ -9867,6 +10038,7 @@ impl App { let result = match provider { OAuthProvider::OpenAI => crate::auth::openai_oauth::authorize_browser().await, OAuthProvider::XAI => crate::auth::xai_oauth::authorize_browser().await, + OAuthProvider::Google => unreachable!("Google uses pasted callbacks"), }; let _ = match result { @@ -9912,6 +10084,7 @@ impl App { }) .await } + OAuthProvider::Google => Err(anyhow::anyhow!("Google uses browser sign-in")), OAuthProvider::XAI => { let code_sender = sender.clone(); crate::auth::xai_oauth::authorize_headless(move |code, url| { @@ -9954,8 +10127,8 @@ impl App { provider, credentials, } => { - if let Ok(auth_dao) = crate::persistence::AuthDAO::new() { - let _ = auth_dao.set_provider( + let saved = crate::persistence::AuthDAO::new().and_then(|auth_dao| { + auth_dao.set_provider( provider.provider_id().to_string(), crate::persistence::AuthConfig::OAuth { refresh: credentials.refresh, @@ -9963,8 +10136,26 @@ impl App { expires: credentials.expires, account_id: credentials.account_id, enterprise_url: credentials.enterprise_url, + project_id: credentials.project_id, }, - ); + ) + }); + + self.google_browser_authorization = None; + self.google_oauth_task = None; + if saved.is_err() { + self.provider_oauth_in_progress = None; + self.provider_oauth_receiver = None; + self.provider_oauth_flow_state.hide(); + if self.overlay_focus == OverlayFocus::ProviderOAuthFlow { + self.overlay_focus = OverlayFocus::None; + } + push_toast(Toast::new( + "OAuth completed, but credentials could not be saved. Please reconnect.", + ToastLevel::Error, + None, + )); + continue; } let default_model = provider.default_model(); @@ -9991,6 +10182,8 @@ impl App { )); } ProviderOAuthTaskMessage::Failed { provider, error } => { + self.google_browser_authorization = None; + self.google_oauth_task = None; self.provider_oauth_in_progress = None; self.provider_oauth_receiver = None; self.provider_oauth_flow_state.hide(); @@ -10892,6 +11085,15 @@ impl App { self.mark_streaming_snapshot_pending(session_id); true } + crate::llm::ChunkMessage::ProviderReplay(item) => { + if let Some(chat) = self.chat_for_session_mut(session_id) { + if let Some(message) = chat.get_last_assistant_message_mut() { + crate::llm::antigravity::persist_replay(message, item); + } + } + self.mark_streaming_snapshot_pending(session_id); + true + } crate::llm::ChunkMessage::Reasoning(reasoning) => { self.set_session_retry_status(session_id, None); if let Some(chat) = self.chat_for_session_mut(session_id) { @@ -13553,6 +13755,157 @@ mod tests { use crate::tools::{PermissionAction, PermissionPrompt}; use serde_json::json; + #[test] + fn google_connect_keeps_api_key_and_browser_choices() { + let mut app = test_app(); + app.show_google_connect_methods(); + assert_eq!( + app.connect_dialog_mode, + ConnectDialogMode::GoogleMethodSelection + ); + assert!(app + .connect_dialog_state + .dialog + .select_item_by_id("google-api-key")); + let item = app + .connect_dialog_state + .dialog + .get_selected() + .unwrap() + .clone(); + app.handle_connect_dialog_selection(item); + assert_eq!(app.overlay_focus, OverlayFocus::ApiKeyInput); + assert_eq!( + app.connect_dialog_mode, + ConnectDialogMode::ProviderSelection + ); + + app.show_google_connect_methods(); + assert!(app + .connect_dialog_state + .dialog + .select_item_by_id("google-oauth-browser")); + let item = app + .connect_dialog_state + .dialog + .get_selected() + .unwrap() + .clone(); + assert_eq!(item.id, "google-oauth-browser"); + app.begin_google_oauth_browser_with(|| { + Ok(crate::auth::antigravity_oauth::BrowserAuthorization::for_test()) + }); + assert_eq!(app.overlay_focus, OverlayFocus::ProviderOAuthFlow); + assert_eq!(app.provider_oauth_in_progress, Some(OAuthProvider::Google)); + assert!(app + .provider_oauth_flow_state + .url + .as_deref() + .unwrap() + .starts_with("https://accounts.google.com/")); + let url = app.provider_oauth_flow_state.url.clone(); + app.begin_google_oauth_browser_with(|| { + Ok(crate::auth::antigravity_oauth::BrowserAuthorization::for_test()) + }); + assert_eq!(app.provider_oauth_flow_state.url, url); + app.cancel_google_oauth(); + } + + #[test] + fn antigravity_stream_preserves_hidden_signed_history() { + let mut app = test_app(); + let session_id = app.create_new_session(Some("Signed history".into())); + app.chat_state + .chat + .add_message(crate::session::types::Message::incomplete("answer")); + let payload = format!( + "{}{}", + crate::llm::antigravity::REPLAY_PREFIX, + json!({"parts":[{"text":"answer","thoughtSignature":"signature"}],"call_ids":[]}) + ); + app.process_streaming_chunk_for_session( + &session_id, + crate::llm::ChunkMessage::ProviderReplay(crate::aisdk::chunk::ReasoningReplayItem { + id: Some("signed-step".into()), + summary: String::new(), + encrypted_content: Some(payload), + }), + ); + let message = app.chat_state.chat.messages.last().unwrap(); + assert_eq!(message.content, "answer"); + assert_eq!(message.parts[0].part_type, "provider_replay"); + assert_eq!(message.parts[0].data["id"], "signed-step"); + } + + #[test] + fn google_callback_paste_can_be_corrected_and_cancelled_without_chat_leak() { + let mut app = test_app(); + app.input.set_text("existing draft"); + app.begin_google_oauth_browser_with(|| { + Ok(crate::auth::antigravity_oauth::BrowserAuthorization::for_test()) + }); + app.handle_paste("http://localhost:51121/oauth-callback?code=private&state=wrong".into()); + let action = handle_provider_oauth_flow_key_event( + &mut app.provider_oauth_flow_state, + KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE), + ); + let ProviderOAuthFlowAction::SubmitCallback(callback) = action else { + panic!("paste should be submitted only by Enter"); + }; + app.submit_google_oauth_callback(callback); + assert!(app.provider_oauth_flow_state.callback_error.is_some()); + assert!(!app.provider_oauth_flow_state.submitting); + assert!(app.google_browser_authorization.is_some()); + assert!(app.google_oauth_task.is_none()); + assert_eq!(app.input.get_text(), "existing draft"); + + app.handle_keys(KeyEvent::new(KeyCode::Esc, KeyModifiers::NONE)); + assert_eq!(app.overlay_focus, OverlayFocus::None); + assert!(app.google_browser_authorization.is_none()); + assert!(app.provider_oauth_in_progress.is_none()); + assert!(app.provider_oauth_flow_state.url.is_none()); + } + + #[test] + fn google_missing_client_configuration_leaves_auth_state_unchanged() { + let mut app = test_app(); + app.show_google_connect_methods(); + app.begin_google_oauth_browser_with(|| { + Err(anyhow::anyhow!( + "Set CRABCODE_ANTIGRAVITY_CLIENT_ID before connecting Google Antigravity" + )) + }); + assert_eq!(app.overlay_focus, OverlayFocus::ConnectDialog); + assert!(app.google_browser_authorization.is_none()); + assert!(app.provider_oauth_in_progress.is_none()); + assert!(!app.provider_oauth_flow_state.visible); + } + + #[tokio::test] + async fn google_callback_cancel_aborts_task_and_discards_queued_completion() { + let mut app = test_app(); + app.begin_google_oauth_browser_with(|| { + Ok(crate::auth::antigravity_oauth::BrowserAuthorization::for_test()) + }); + let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); + app.provider_oauth_receiver = Some(receiver); + sender + .send(ProviderOAuthTaskMessage::Failed { + provider: OAuthProvider::Google, + error: "queued completion".into(), + }) + .unwrap(); + let task = tokio::spawn(std::future::pending::<()>()); + let abort = task.abort_handle(); + app.google_oauth_task = Some(task); + app.cancel_google_oauth(); + tokio::task::yield_now().await; + assert!(abort.is_finished()); + assert!(app.provider_oauth_receiver.is_none()); + assert!(app.google_oauth_task.is_none()); + assert!(app.google_browser_authorization.is_none()); + } + /// Serializes the preview-flag env tests so concurrent cases cannot swap /// `CRABCODE_FORCE_UPDATE_NOTICE` / `CRABCODE_NO_UPDATE_CHECK` mid-call. fn update_preview_env_lock() -> &'static std::sync::Mutex<()> { @@ -13977,6 +14330,8 @@ mod tests { api_key_input: crate::ui::components::api_key_input::ApiKeyInput::new(), provider_oauth_receiver: None, provider_oauth_in_progress: None, + google_browser_authorization: None, + google_oauth_task: None, mcp_oauth_receiver: None, mcp_oauth_in_progress: None, compaction_receiver: None, diff --git a/src/auth/antigravity_oauth.rs b/src/auth/antigravity_oauth.rs new file mode 100644 index 00000000..ae281c59 --- /dev/null +++ b/src/auth/antigravity_oauth.rs @@ -0,0 +1,971 @@ +//! Manual-paste Google Antigravity OAuth. No callback listener is started. +use anyhow::{anyhow, bail, Result}; +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine}; +use rand::{rngs::OsRng, RngCore}; +use serde::Deserialize; +use serde_json::{json, Value}; +use sha2::{Digest, Sha256}; +use std::{collections::HashSet, time::Duration}; +use url::Url; + +use super::OAuthCredentials; + +const CLIENT_ID_ENV: &str = "CRABCODE_ANTIGRAVITY_CLIENT_ID"; +const CLIENT_SECRET_ENV: &str = "CRABCODE_ANTIGRAVITY_CLIENT_SECRET"; +pub const REDIRECT_URI: &str = "http://localhost:51121/oauth-callback"; +pub const ANTIGRAVITY_ENDPOINT: &str = "https://daily-cloudcode-pa.googleapis.com"; +pub const ANTIGRAVITY_VERSION: &str = "2.24.0"; +const TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token"; +const MAX_CALLBACK_BYTES: usize = 16 * 1024; +const MAX_RESPONSE_BYTES: usize = 1024 * 1024; +const SCOPES: &str = "https://www.googleapis.com/auth/cloud-platform https://www.googleapis.com/auth/userinfo.email https://www.googleapis.com/auth/userinfo.profile https://www.googleapis.com/auth/cclog https://www.googleapis.com/auth/experimentsandconfigs"; + +pub fn user_agent() -> &'static str { + "antigravity/hub/2.24.0 (aidev_client; os_type=darwin; arch=arm64; cl=1000335682)" +} + +struct OAuthClient { + id: String, + secret: String, +} + +impl OAuthClient { + fn from_env() -> Result { + Self::from_lookup(|name| std::env::var_os(name)) + } + + fn from_lookup(mut lookup: impl FnMut(&str) -> Option) -> Result { + let mut required = |name| -> Result { + lookup(name) + .and_then(|value| value.into_string().ok()) + .filter(|value| !value.is_empty() && !value.chars().any(char::is_whitespace) + && !value.chars().any(char::is_control)) + .ok_or_else(|| anyhow!("Set {name} to a valid OAuth client value before connecting Google Antigravity")) + }; + Ok(Self { + id: required(CLIENT_ID_ENV)?, + secret: required(CLIENT_SECRET_ENV)?, + }) + } + + fn token_form(&self) -> Vec<(&str, &str)> { + vec![ + ("client_id", self.id.as_str()), + ("client_secret", self.secret.as_str()), + ] + } +} + +pub struct BrowserAuthorization { + url: String, + verifier: String, + state: String, + client: OAuthClient, +} + +impl BrowserAuthorization { + pub fn new() -> Result { + Self::with_client(OAuthClient::from_env()?) + } + + #[cfg(test)] + pub(crate) fn for_test() -> Self { + Self::with_client(OAuthClient { + id: "test-oauth-client".into(), + secret: "test-oauth-secret".into(), + }) + .unwrap() + } + + fn with_client(client: OAuthClient) -> Result { + let mut verifier = [0u8; 32]; + let mut state = [0u8; 32]; + OsRng + .try_fill_bytes(&mut verifier) + .map_err(|_| anyhow!("OAuth random generation failed"))?; + OsRng + .try_fill_bytes(&mut state) + .map_err(|_| anyhow!("OAuth random generation failed"))?; + let verifier = URL_SAFE_NO_PAD.encode(verifier); + let state = URL_SAFE_NO_PAD.encode(state); + let challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes())); + let mut url = Url::parse("https://accounts.google.com/o/oauth2/v2/auth")?; + url.query_pairs_mut().extend_pairs([ + ("client_id", client.id.as_str()), + ("redirect_uri", REDIRECT_URI), + ("response_type", "code"), + ("scope", SCOPES), + ("code_challenge", challenge.as_str()), + ("code_challenge_method", "S256"), + ("state", state.as_str()), + ("access_type", "offline"), + ("prompt", "consent"), + ]); + Ok(Self { + url: url.into(), + verifier, + state, + client, + }) + } + + pub fn authorization_url(&self) -> &str { + &self.url + } + + pub fn validate_callback(&self, input: &str) -> Result { + if input.len() > MAX_CALLBACK_BYTES { + bail!("OAuth callback is too long"); + } + let input = input.trim(); + if input.chars().any(char::is_control) { + bail!("Invalid OAuth callback"); + } + // Compare the original spelling too: URL normalization must not admit another path/origin. + if input.split('?').next() != Some(REDIRECT_URI) { + bail!("OAuth callback must use the expected localhost address"); + } + let url = Url::parse(input).map_err(|_| anyhow!("Invalid OAuth callback URL"))?; + if url.fragment().is_some() { + bail!("OAuth callback must not contain a fragment"); + } + let mut names = HashSet::new(); + let mut code = None; + let mut state = None; + let mut failed = false; + for (name, value) in url.query_pairs() { + if !names.insert(name.to_string()) { + bail!("Duplicate OAuth callback field"); + } + match name.as_ref() { + "code" => code = Some(value.into_owned()), + "state" => state = Some(value.into_owned()), + "error" | "error_description" | "error_uri" => failed = true, + _ => {} + } + } + if state.as_deref() != Some(self.state.as_str()) { + bail!("OAuth callback state is missing or does not match this login"); + } + if failed { + bail!("Google authorization was denied or failed; start a new login"); + } + code.filter(|value| !value.trim().is_empty() && !value.chars().any(char::is_control)) + .ok_or_else(|| anyhow!("OAuth callback is missing a valid authorization code")) + } + + pub async fn complete(self, callback_url: &str) -> Result { + let code = self.validate_callback(callback_url)?; + let client = http_client()?; + let started = chrono::Utc::now().timestamp_millis(); + let mut form = self.client.token_form(); + form.extend([ + ("grant_type", "authorization_code"), + ("redirect_uri", REDIRECT_URI), + ("code", code.as_str()), + ("code_verifier", self.verifier.as_str()), + ]); + let response = client + .post(TOKEN_ENDPOINT) + .form(&form) + .send() + .await + .map_err(|_| anyhow!("Google token exchange request failed"))?; + let payload: TokenResponse = read_json(response, "Google token exchange").await?; + let mut credentials = payload.into_credentials(None, started)?; + credentials.project_id = Some(resolve_project(&credentials.access).await?); + Ok(credentials) + } +} + +fn http_client() -> Result { + reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .connect_timeout(Duration::from_secs(10)) + .timeout(Duration::from_secs(30)) + .user_agent(user_agent()) + .build() + .map_err(|_| anyhow!("Could not create Google HTTP client")) +} + +async fn read_json( + mut response: reqwest::Response, + operation: &str, +) -> Result { + if !response.status().is_success() { + bail!("{operation} failed (HTTP {})", response.status().as_u16()); + } + let mut bytes = Vec::new(); + while let Some(chunk) = response + .chunk() + .await + .map_err(|_| anyhow!("{operation} response could not be read"))? + { + if bytes.len().saturating_add(chunk.len()) > MAX_RESPONSE_BYTES { + bail!("{operation} response is too large"); + } + bytes.extend_from_slice(&chunk); + } + serde_json::from_slice(&bytes).map_err(|_| anyhow!("{operation} returned an invalid response")) +} + +#[derive(Deserialize)] +struct TokenResponse { + access_token: String, + expires_in: i64, + refresh_token: Option, +} +impl TokenResponse { + fn into_credentials( + self, + previous_refresh: Option<&str>, + started: i64, + ) -> Result { + if self.access_token.trim().is_empty() || self.expires_in <= 0 { + bail!("Google returned invalid token credentials"); + } + let refresh = self + .refresh_token + .filter(|s| !s.trim().is_empty()) + .or_else(|| previous_refresh.map(str::to_owned)) + .filter(|s| !s.trim().is_empty()) + .ok_or_else(|| { + anyhow!("Google did not return a refresh token; restart login and grant consent") + })?; + let expires = self + .expires_in + .checked_mul(1000) + .and_then(|duration| started.checked_add(duration)) + .ok_or_else(|| anyhow!("Google returned an invalid token expiry"))?; + Ok(OAuthCredentials { + access: self.access_token, + refresh, + expires, + account_id: None, + enterprise_url: None, + project_id: None, + }) + } +} + +pub async fn refresh_access_token(refresh_token: &str) -> Result { + if refresh_token.trim().is_empty() { + bail!("Google refresh token is missing"); + } + let oauth_client = OAuthClient::from_env()?; + let mut form = oauth_client.token_form(); + form.extend([ + ("grant_type", "refresh_token"), + ("refresh_token", refresh_token), + ]); + let started = chrono::Utc::now().timestamp_millis(); + let response = http_client()? + .post(TOKEN_ENDPOINT) + .form(&form) + .send() + .await + .map_err(|_| anyhow!("Google token refresh request failed"))?; + let payload: TokenResponse = read_json(response, "Google token refresh").await?; + payload.into_credentials(Some(refresh_token), started) +} + +static REQUEST_CREDENTIALS_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + +pub async fn load_request_credentials() -> Result { + let dao = crate::persistence::AuthDAO::new() + .map_err(|_| anyhow!("Could not open Google Antigravity OAuth credentials"))?; + load_request_credentials_with( + || { + let config = dao + .get_provider("google") + .map_err(|_| anyhow!("Could not read Google Antigravity OAuth credentials"))?; + let Some(crate::persistence::AuthConfig::OAuth { + refresh, + access, + expires, + account_id, + enterprise_url, + project_id, + }) = config + else { + bail!("Google Antigravity OAuth credentials are missing. Please reconnect Google Antigravity OAuth."); + }; + Ok(OAuthCredentials { + refresh, + access, + expires, + account_id, + enterprise_url, + project_id, + }) + }, + |refresh| async move { refresh_access_token(&refresh).await }, + |access| async move { resolve_project(&access).await }, + |credentials| { + dao.set_provider( + "google".to_string(), + crate::persistence::AuthConfig::OAuth { + refresh: credentials.refresh, + access: credentials.access, + expires: credentials.expires, + account_id: credentials.account_id, + enterprise_url: credentials.enterprise_url, + project_id: credentials.project_id, + }, + ) + .map_err(|_| anyhow!("Could not save Google Antigravity OAuth credentials")) + }, + || chrono::Utc::now().timestamp_millis(), + ) + .await +} + +fn valid_access(access: &str) -> bool { + !access.is_empty() + && !access.chars().any(char::is_whitespace) + && !access.chars().any(char::is_control) +} + +async fn load_request_credentials_with< + Load, + Refresh, + RefreshFuture, + Resolve, + ResolveFuture, + Persist, + Clock, +>( + load: Load, + refresh: Refresh, + resolve: Resolve, + persist: Persist, + now: Clock, +) -> Result +where + Load: FnOnce() -> Result, + Refresh: FnOnce(String) -> RefreshFuture, + RefreshFuture: std::future::Future>, + Resolve: FnOnce(String) -> ResolveFuture, + ResolveFuture: std::future::Future>, + Persist: FnOnce(OAuthCredentials) -> Result<()>, + Clock: Fn() -> i64, +{ + // Read after acquiring the lock: another generation may have rotated the refresh token. + let _guard = REQUEST_CREDENTIALS_LOCK.lock().await; + let mut credentials = load()?; + let previous = credentials.clone(); + if credentials.expires <= now().saturating_add(60_000) || !valid_access(&credentials.access) { + match refresh(credentials.refresh.clone()).await { + Ok(updated) => { + if !updated.refresh.trim().is_empty() { + credentials.refresh = updated.refresh; + } + credentials.access = updated.access; + credentials.expires = updated.expires; + } + Err(_) if credentials.expires > now() && valid_access(&credentials.access) => {} + Err(_) => bail!("Google Antigravity OAuth token refresh failed. Please reconnect Google Antigravity OAuth."), + } + } + if !valid_access(&credentials.access) || credentials.expires <= now() { + bail!("Google Antigravity OAuth access token is invalid or expired. Please reconnect Google Antigravity OAuth."); + } + if credentials + .project_id + .as_deref() + .is_none_or(|id| id.trim().is_empty()) + { + let resolved = resolve(credentials.access.clone()).await.map_err(|_| { + anyhow!("Failed to resolve Google Antigravity OAuth project. Please reconnect Google Antigravity OAuth.") + })?; + if project_id(&Value::String(resolved.clone())).is_none() { + bail!("Google Antigravity OAuth project is invalid. Please reconnect Google Antigravity OAuth."); + } + credentials.project_id = Some(resolved); + } + if credentials.expires <= now() { + bail!("Google Antigravity OAuth access token expired while resolving the project. Please retry."); + } + if credentials.access != previous.access + || credentials.refresh != previous.refresh + || credentials.expires != previous.expires + || credentials.project_id != previous.project_id + { + persist(credentials.clone())?; + } + Ok(credentials) +} + +fn project_id(value: &Value) -> Option { + value + .as_str() + .or_else(|| value.get("id").and_then(Value::as_str)) + .map(str::trim) + .filter(|s| !s.is_empty() && s.len() <= 1024 && !s.chars().any(char::is_control)) + .map(str::to_owned) +} + +fn tier_id(payload: &Value) -> &str { + let tiers = payload.get("allowedTiers").and_then(Value::as_array); + fn valid_id(tier: &Value) -> Option<&str> { + tier.get("id") + .and_then(Value::as_str) + .filter(|s| !s.trim().is_empty()) + } + tiers + .and_then(|tiers| { + tiers.iter().find(|tier| { + tier.get("isDefault").and_then(Value::as_bool) == Some(true) + && valid_id(tier).is_some() + }) + }) + .and_then(valid_id) + .or_else(|| tiers.and_then(|tiers| tiers.first()).and_then(valid_id)) + .or_else(|| payload.get("currentTier").and_then(valid_id)) + .unwrap_or("free-tier") +} + +fn onboard_body(tier: &str) -> Value { + json!({ "tier_id": tier, "metadata": { "ide_type": "ANTIGRAVITY", "ide_version": ANTIGRAVITY_VERSION, "ide_name": "antigravity" } }) +} + +pub async fn resolve_project(access: &str) -> Result { + if access.trim().is_empty() { + bail!("Google access token is missing"); + } + tokio::time::timeout( + Duration::from_secs(120), + resolve_project_at( + &http_client()?, + ANTIGRAVITY_ENDPOINT, + access, + Duration::from_secs(5), + ), + ) + .await + .map_err(|_| anyhow!("Antigravity project discovery timed out; retry login"))? +} + +async fn resolve_project_at( + client: &reqwest::Client, + endpoint: &str, + access: &str, + delay: Duration, +) -> Result { + let response = client + .post(format!("{endpoint}/v1internal:loadCodeAssist")) + .bearer_auth(access) + .json(&json!({ "metadata": { "ideType": "ANTIGRAVITY" } })) + .send() + .await + .map_err(|_| anyhow!("Antigravity project discovery request failed"))?; + let payload: Value = read_json(response, "Antigravity project discovery").await?; + if let Some(project) = project_id(&payload["cloudaicompanionProject"]) { + return Ok(project); + } + let body = onboard_body(tier_id(&payload)); + for attempt in 0..10 { + let response = client + .post(format!("{endpoint}/v1internal:onboardUser")) + .bearer_auth(access) + .header( + "User-Agent", + format!("{} google-api-nodejs-client/10.3.0", user_agent()), + ) + .header("X-Goog-Api-Client", "gl-node/22.21.1") + .json(&body) + .send() + .await + .map_err(|_| anyhow!("Antigravity onboarding request failed"))?; + let payload: Value = read_json(response, "Antigravity onboarding").await?; + if payload["done"].as_bool() == Some(true) { + return project_id(&payload["response"]["cloudaicompanionProject"]) + .ok_or_else(|| anyhow!("Antigravity onboarding completed without a managed project; check account eligibility")); + } + if attempt < 9 { + tokio::time::sleep(delay).await; + } + } + bail!("Antigravity did not provision a managed project; check account eligibility and retry login") +} + +#[cfg(test)] +mod tests { + use super::*; + + async fn fixture_server( + responses: Vec<(u16, Value)>, + ) -> (String, tokio::task::JoinHandle>) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let mut requests = Vec::new(); + for (status, body) in responses { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut bytes = Vec::new(); + loop { + let mut buffer = [0; 4096]; + let count = socket.read(&mut buffer).await.unwrap(); + assert!(count > 0); + bytes.extend_from_slice(&buffer[..count]); + if let Some(end) = bytes.windows(4).position(|s| s == b"\r\n\r\n") { + let headers = String::from_utf8_lossy(&bytes[..end]); + let length: usize = headers + .lines() + .find_map(|line| { + line.to_ascii_lowercase() + .strip_prefix("content-length: ") + .map(str::to_owned) + }) + .unwrap() + .parse() + .unwrap(); + if bytes.len() >= end + 4 + length { + break; + } + } + } + requests.push(String::from_utf8(bytes).unwrap()); + let body = body.to_string(); + let response = format!("HTTP/1.1 {status} Fixture\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()); + socket.write_all(response.as_bytes()).await.unwrap(); + } + requests + }); + (endpoint, server) + } + + fn request_credentials() -> OAuthCredentials { + OAuthCredentials { + access: "access".into(), + refresh: "refresh".into(), + expires: 120_000, + account_id: Some("account".into()), + enterprise_url: Some("enterprise".into()), + project_id: Some("managed".into()), + } + } + + #[tokio::test] + async fn request_credentials_refresh_preserves_metadata_and_saves_only_changes() { + let credentials = load_request_credentials_with( + || Ok(request_credentials()), + |_| async { panic!("fresh token must not refresh") }, + |_| async { panic!("existing project must not resolve") }, + |_| panic!("unchanged credentials must not save"), + || 0, + ) + .await + .unwrap(); + assert_eq!(credentials.access, "access"); + let credentials = load_request_credentials_with( + || Ok(request_credentials()), + |token| async move { + assert_eq!(token, "refresh"); + let mut updated = request_credentials(); + updated.access = "new-access".into(); + updated.refresh.clear(); + updated.project_id = None; + updated.expires = 240_000; + Ok(updated) + }, + |_| async { panic!("refresh must retain project") }, + |saved| { + assert_eq!(saved.access, "new-access"); + assert_eq!(saved.refresh, "refresh"); + assert_eq!(saved.account_id.as_deref(), Some("account")); + assert_eq!(saved.enterprise_url.as_deref(), Some("enterprise")); + assert_eq!(saved.project_id.as_deref(), Some("managed")); + Ok(()) + }, + || 60_000, + ) + .await + .unwrap(); + assert_eq!(credentials.expires, 240_000); + } + + #[tokio::test] + async fn request_credentials_resolve_project_and_reject_invalid_refresh() { + let credentials = load_request_credentials_with( + || { + let mut c = request_credentials(); + c.project_id = None; + Ok(c) + }, + |_| async { panic!("fresh token must not refresh") }, + |access| async move { + assert_eq!(access, "access"); + Ok("resolved".into()) + }, + |saved| { + assert_eq!(saved.project_id.as_deref(), Some("resolved")); + Ok(()) + }, + || 0, + ) + .await + .unwrap(); + assert_eq!(credentials.project_id.as_deref(), Some("resolved")); + for invalid_access in ["", "bad\ntoken"] { + let result = load_request_credentials_with( + || { + let mut c = request_credentials(); + c.access = invalid_access.into(); + Ok(c) + }, + |_| async { + let mut c = request_credentials(); + c.expires = 0; + Ok(c) + }, + |_| async { panic!("expired refresh must not resolve") }, + |_| panic!("expired refresh must not save"), + || 0, + ) + .await; + assert!(result.is_err()); + } + } + + #[tokio::test] + async fn failed_proactive_refresh_uses_only_still_valid_access() { + let c = load_request_credentials_with( + || Ok(request_credentials()), + |_| async { bail!("sensitive-upstream-error") }, + |_| async { panic!("existing project must not resolve") }, + |_| panic!("failed refresh must not save"), + || 60_000, + ) + .await + .unwrap(); + assert_eq!(c.access, "access"); + let error = load_request_credentials_with( + || Ok(request_credentials()), + |_| async { bail!("sensitive-upstream-error") }, + |_| async { panic!("expired access must not resolve") }, + |_| panic!("failed refresh must not save"), + || 120_000, + ) + .await + .unwrap_err() + .to_string(); + assert!(error.contains("reconnect Google Antigravity OAuth")); + assert!(!error.contains("sensitive-upstream-error")); + } + + #[tokio::test] + async fn cancelled_refresh_does_not_save_or_hold_lock() { + let refresh_started = std::sync::Arc::new(tokio::sync::Notify::new()); + let started = refresh_started.clone(); + let task = tokio::spawn(async move { + load_request_credentials_with( + || { + let mut c = request_credentials(); + c.expires = 0; + Ok(c) + }, + |_| async move { + started.notify_one(); + std::future::pending::>().await + }, + |_| async { panic!("cancelled refresh must not resolve") }, + |_| panic!("cancelled refresh must not save"), + || 0, + ) + .await + }); + refresh_started.notified().await; + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + drop( + tokio::time::timeout(Duration::from_secs(1), REQUEST_CREDENTIALS_LOCK.lock()) + .await + .unwrap(), + ); + } + + #[tokio::test] + async fn concurrent_generations_read_rotated_credentials_after_lock() { + let stored = std::sync::Arc::new(std::sync::Mutex::new(request_credentials())); + let refreshes = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let mut tasks = Vec::new(); + for _ in 0..2 { + let read = stored.clone(); + let save = stored.clone(); + let refreshes = refreshes.clone(); + tasks.push(tokio::spawn(async move { + load_request_credentials_with( + || Ok(read.lock().unwrap().clone()), + |_| async move { + refreshes.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + tokio::task::yield_now().await; + let mut c = request_credentials(); + c.access = "rotated".into(); + c.refresh = "rotated-refresh".into(); + c.expires = 240_000; + Ok(c) + }, + |_| async { panic!("existing project must not resolve") }, + |c| { + *save.lock().unwrap() = c; + Ok(()) + }, + || 60_000, + ) + .await + .unwrap() + .access + })); + } + for task in tasks { + assert_eq!(task.await.unwrap(), "rotated"); + } + assert_eq!(refreshes.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn project_discovery_and_bounded_onboarding() { + let (endpoint, server) = fixture_server(vec![ + ( + 200, + json!({"allowedTiers":[{"id":"fixture-tier","isDefault":true}]}), + ), + (200, json!({"done":false})), + ( + 200, + json!({"done":true,"response":{"cloudaicompanionProject":{"id":"managed"}}}), + ), + ]) + .await; + assert_eq!( + resolve_project_at( + &http_client().unwrap(), + &endpoint, + "test-token", + Duration::ZERO + ) + .await + .unwrap(), + "managed" + ); + let requests = server.await.unwrap(); + assert!(requests[0].starts_with("POST /v1internal:loadCodeAssist ")); + let body: Value = + serde_json::from_str(requests[1].split("\r\n\r\n").nth(1).unwrap()).unwrap(); + assert_eq!(body, onboard_body("fixture-tier")); + let (endpoint, server) = + fixture_server(vec![(200, json!({"cloudaicompanionProject":"direct"}))]).await; + assert_eq!( + resolve_project_at(&http_client().unwrap(), &endpoint, "test", Duration::ZERO) + .await + .unwrap(), + "direct" + ); + assert_eq!(server.await.unwrap().len(), 1); + let (endpoint, server) = + fixture_server(vec![(403, json!({"error":"sensitive-body"}))]).await; + let error = resolve_project_at(&http_client().unwrap(), &endpoint, "test", Duration::ZERO) + .await + .unwrap_err() + .to_string(); + assert!(error.contains("403")); + assert!(!error.contains("sensitive-body")); + server.await.unwrap(); + let mut responses = vec![(200, json!({}))]; + responses.extend((0..10).map(|_| (200, json!({"done":false})))); + let (endpoint, server) = fixture_server(responses).await; + assert!( + resolve_project_at(&http_client().unwrap(), &endpoint, "test", Duration::ZERO) + .await + .is_err() + ); + assert_eq!(server.await.unwrap().len(), 11); + let (endpoint, server) = + fixture_server(vec![(200, json!({})), (200, json!({"done":true}))]).await; + assert!( + resolve_project_at(&http_client().unwrap(), &endpoint, "test", Duration::ZERO) + .await + .is_err() + ); + server.await.unwrap(); + } + fn callback(session: &BrowserAuthorization, fields: &str) -> String { + format!("{REDIRECT_URI}?state={}&{fields}", session.state) + } + + #[test] + fn client_configuration_requires_both_values_without_echoing_them() { + let error = OAuthClient::from_lookup(|_| None) + .err() + .unwrap() + .to_string(); + assert!(error.contains(CLIENT_ID_ENV)); + for invalid in ["", " ", "private value", "private\nvalue"] { + let error = OAuthClient::from_lookup(|name| { + Some(if name == CLIENT_ID_ENV { + "configured-client".into() + } else { + invalid.into() + }) + }) + .err() + .unwrap() + .to_string(); + assert!(error.contains(CLIENT_SECRET_ENV)); + assert!(!error.contains("private")); + } + let client = OAuthClient::from_lookup(|name| { + Some(if name == CLIENT_ID_ENV { + "configured-client".into() + } else { + "configured-secret".into() + }) + }) + .unwrap(); + assert_eq!( + client.token_form(), + vec![ + ("client_id", "configured-client"), + ("client_secret", "configured-secret"), + ] + ); + let authorization = BrowserAuthorization::with_client(client).unwrap(); + let url = Url::parse(authorization.authorization_url()).unwrap(); + assert!(url + .query_pairs() + .any(|(key, value)| key == "client_id" && value == "configured-client")); + assert!(!authorization + .authorization_url() + .contains("configured-secret")); + assert_eq!(authorization.client.token_form()[0].1, "configured-client"); + } + #[test] + fn authorization_has_independent_random_state_and_pkce() { + let a = BrowserAuthorization::for_test(); + let b = BrowserAuthorization::for_test(); + assert_ne!(a.state, b.state); + assert_ne!(a.verifier, b.verifier); + assert_ne!(a.state, a.verifier); + assert_eq!(URL_SAFE_NO_PAD.decode(&a.state).unwrap().len(), 32); + assert_eq!(URL_SAFE_NO_PAD.decode(&a.verifier).unwrap().len(), 32); + let url = Url::parse(a.authorization_url()).unwrap(); + let pairs: std::collections::HashMap<_, _> = url.query_pairs().collect(); + assert_eq!( + pairs["code_challenge"], + URL_SAFE_NO_PAD.encode(Sha256::digest(a.verifier.as_bytes())) + ); + assert_eq!(pairs["code_challenge_method"], "S256"); + assert_eq!(pairs["redirect_uri"], REDIRECT_URI); + assert_eq!(pairs["client_id"], "test-oauth-client"); + assert!(!a.authorization_url().contains("test-oauth-secret")); + assert!(!a.authorization_url().contains(&a.verifier)); + } + #[test] + fn callback_validation() { + let a = BrowserAuthorization::for_test(); + assert_eq!( + a.validate_callback(&callback(&a, "code=a%2Fb")).unwrap(), + "a/b" + ); + for fields in [ + "code=", + "code=%20", + "code=%0A", + "code=x&code=y", + "code=x&%63ode=y", + "code=x&state=x", + "code=x&error=access_denied", + "error=denied", + "code=x#error", + "code=x&error_description=secret", + ] { + assert!( + a.validate_callback(&callback(&a, fields)).is_err(), + "{fields}" + ); + } + let good = callback(&a, "code=secret"); + for prefix in [ + "https://localhost:51121/oauth-callback", + "http://127.0.0.1:51121/oauth-callback", + "http://localhost:51122/oauth-callback", + "http://localhost:51121/other", + "http://localhost:51121/a/../oauth-callback", + "http://user@localhost:51121/oauth-callback", + "http://localhost.evil:51121/oauth-callback", + ] { + assert!(a + .validate_callback(&good.replacen(REDIRECT_URI, prefix, 1)) + .is_err()); + } + assert!(a + .validate_callback(&format!("{REDIRECT_URI}?code=secret")) + .is_err()); + assert!(a + .validate_callback(&format!("{REDIRECT_URI}?state=wrong&code=secret")) + .is_err()); + assert!(a + .validate_callback(&"x".repeat(MAX_CALLBACK_BYTES + 1)) + .is_err()); + assert!(!a + .validate_callback(&format!("{REDIRECT_URI}?code=secret")) + .unwrap_err() + .to_string() + .contains("secret")); + } + #[test] + fn token_conversion() { + let payload = || TokenResponse { + access_token: "access".into(), + expires_in: 3600, + refresh_token: None, + }; + assert!(payload().into_credentials(None, 1000).is_err()); + let credentials = payload() + .into_credentials(Some("old-refresh"), 1000) + .unwrap(); + assert_eq!(credentials.refresh, "old-refresh"); + assert_eq!(credentials.expires, 3_601_000); + assert!(credentials.project_id.is_none()); + let mut value = payload(); + value.refresh_token = Some("new-refresh".into()); + assert_eq!( + value.into_credentials(Some("old"), 0).unwrap().refresh, + "new-refresh" + ); + let mut value = payload(); + value.expires_in = i64::MAX; + assert!(value.into_credentials(Some("old"), 0).is_err()); + let mut value = payload(); + value.access_token.clear(); + assert!(value.into_credentials(Some("old"), 0).is_err()); + } + #[test] + fn project_and_onboarding_fixtures() { + assert_eq!(project_id(&json!(" managed ")).as_deref(), Some("managed")); + assert_eq!( + project_id(&json!({"id":"managed"})).as_deref(), + Some("managed") + ); + for value in [ + Value::Null, + json!(" "), + json!({"id": 1}), + json!({"name":"fake"}), + ] { + assert!(project_id(&value).is_none()); + } + assert_eq!(tier_id(&json!({})), "free-tier"); + assert_eq!(tier_id(&json!({"currentTier":{"id":"current"}})), "current"); + assert_eq!( + tier_id(&json!({"allowedTiers":[{"id":"first"},{"id":"default","isDefault":true}]})), + "default" + ); + assert_eq!( + onboard_body("free-tier"), + json!({"tier_id":"free-tier","metadata":{"ide_type":"ANTIGRAVITY","ide_version":"2.24.0","ide_name":"antigravity"}}) + ); + } +} diff --git a/src/auth/mod.rs b/src/auth/mod.rs index 20191cda..07c8ed96 100644 --- a/src/auth/mod.rs +++ b/src/auth/mod.rs @@ -1,3 +1,4 @@ +pub mod antigravity_oauth; pub mod oauth; pub mod openai_oauth; pub mod xai_oauth; diff --git a/src/auth/oauth.rs b/src/auth/oauth.rs index 9f8800ce..495936e8 100644 --- a/src/auth/oauth.rs +++ b/src/auth/oauth.rs @@ -5,4 +5,5 @@ pub struct OAuthCredentials { pub expires: i64, pub account_id: Option, pub enterprise_url: Option, + pub project_id: Option, } diff --git a/src/auth/openai_oauth.rs b/src/auth/openai_oauth.rs index 0b3a9f47..9ef0c1ff 100644 --- a/src/auth/openai_oauth.rs +++ b/src/auth/openai_oauth.rs @@ -240,6 +240,7 @@ fn credentials_from_token_response( expires, account_id, enterprise_url: None, + project_id: None, }) } diff --git a/src/auth/xai_oauth.rs b/src/auth/xai_oauth.rs index f9398e45..b0ecbcac 100644 --- a/src/auth/xai_oauth.rs +++ b/src/auth/xai_oauth.rs @@ -200,6 +200,7 @@ fn credentials_from_token_response( expires, account_id: None, enterprise_url: None, + project_id: None, }) } diff --git a/src/command/handlers.rs b/src/command/handlers.rs index 511030a6..20ec2528 100644 --- a/src/command/handlers.rs +++ b/src/command/handlers.rs @@ -131,7 +131,7 @@ pub fn handle_connect<'a>( ("anthropic", "Anthropic"), ("openai", "OpenAI"), ("xai", "xAI"), - ("google", "Google"), + ("google", "Google Antigravity"), ] { out.insert( id.to_string(), diff --git a/src/llm/antigravity.rs b/src/llm/antigravity.rs new file mode 100644 index 00000000..7399e20e --- /dev/null +++ b/src/llm/antigravity.rs @@ -0,0 +1,1321 @@ +use crate::aisdk::chunk::{ChunkType, FinishReason, ReasoningReplayItem, TokenUsage}; +use crate::aisdk::error::{Error, Result}; +use crate::aisdk::message::{ImageContent, Message}; +use crate::aisdk::provider::{Provider, ProviderStream}; +use crate::aisdk::retry::RetryError; +use crate::aisdk::tool::{Tool, ToolTransport}; +use async_trait::async_trait; +use base64::Engine; +use eventsource_stream::Eventsource; +use futures::StreamExt; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use std::collections::{HashMap, HashSet, VecDeque}; +use std::time::Duration; + +pub(crate) const REPLAY_PREFIX: &str = "antigravity:gemini-parts:v1:"; + +pub(crate) fn persist_replay( + message: &mut crate::session::types::Message, + item: ReasoningReplayItem, +) { + if !item + .encrypted_content + .as_deref() + .is_some_and(|value| value.starts_with(REPLAY_PREFIX)) + { + return; + } + // The item arrives after this step's text but before its tools execute. + let index = message + .parts + .iter() + .rposition(|part| matches!(part.part_type.as_str(), "tool_result" | "provider_replay")) + .map_or(0, |index| index + 1); + message.parts.insert( + index, + crate::session::types::MessagePart { + part_type: "provider_replay".into(), + data: json!({ + "id": item.id, + "summary": item.summary, + "encrypted_content": item.encrypted_content, + }), + }, + ); +} + +#[derive(Clone)] +pub struct Antigravity { + base_url: String, + model: String, + access: String, + project: String, + stored_credentials: bool, +} + +impl std::fmt::Debug for Antigravity { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Antigravity") + .field("model", &self.model) + .finish_non_exhaustive() + } +} + +impl Antigravity { + pub fn new(base_url: String, model: String, access: String, project: String) -> Self { + Self { + base_url, + model, + access, + project, + stored_credentials: false, + } + } + + pub fn with_stored_credentials(mut self) -> Self { + self.stored_credentials = true; + self + } + + fn request_body( + &self, + messages: &[Message], + tools: &[Tool], + headers: &HashMap, + ) -> Result { + let (contents, system) = gemini_messages(messages)?; + let mut declarations = Vec::new(); + for tool in tools { + if !matches!(tool.transport, ToolTransport::ClientFunction) { + return Err(Error::InvalidInput( + "Antigravity only supports client function tools".into(), + )); + } + let schema = serde_json::to_value(&tool.input_schema)?; + declarations.push(json!({ + "name": tool.name, + "description": tool.description, + "parameters": gemini_schema(&schema, &schema, 0)?, + })); + } + let session = [ + "x-opencode-session", + "session-id", + "x-session-id", + "x-crabcode-session", + "thread-id", + ] + .into_iter() + .find_map(|name| { + headers + .iter() + .find(|(key, value)| key.eq_ignore_ascii_case(name) && !value.trim().is_empty()) + .map(|(_, value)| value.clone()) + }) + .unwrap_or_else(cuid2::create_id); + let mut request = json!({"sessionId": session, "contents": contents}); + if !system.is_empty() { + request["systemInstruction"] = json!({"parts": system}); + } + if !declarations.is_empty() { + request["tools"] = json!([{"functionDeclarations": declarations}]); + } + Ok(json!({ + "project": self.project, + "model": self.model, + "userAgent": "antigravity", + "requestType": "agent", + "requestId": format!("agent-{}", cuid2::create_id()), + "request": request, + })) + } +} + +// The SDK's language-model interface is named Provider. Its opaque reasoning +// replay items survive stream_with_tools_options and persisted Message history. +#[async_trait] +impl Provider for Antigravity { + fn name(&self) -> &str { + "antigravity" + } + fn model_name(&self) -> &str { + &self.model + } + + async fn stream_text( + &self, + messages: &[Message], + tools: &[Tool], + headers: &HashMap, + ) -> Result { + let mut request = self.clone(); + if self.stored_credentials { + let credentials = crate::auth::antigravity_oauth::load_request_credentials() + .await + .map_err(|_| { + Error::Provider("Google Antigravity credentials are unavailable".into()) + })?; + request.access = credentials.access; + request.project = credentials.project_id.ok_or_else(|| { + Error::Provider("Google Antigravity project is unavailable".into()) + })?; + } + let body = request.request_body(messages, tools, headers)?; + let client = reqwest::Client::builder() + .connect_timeout(Duration::from_secs(30)) + .read_timeout(Duration::from_secs(90)) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|_| Error::Provider("Could not initialize Antigravity HTTP client".into()))?; + // Do not forward arbitrary provider headers (especially API keys or an + // Authorization override) to this OAuth-only endpoint. + let response = client + .post(format!( + "{}/v1internal:streamGenerateContent?alt=sse", + self.base_url.trim_end_matches('/') + )) + .bearer_auth(&request.access) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .header(reqwest::header::ACCEPT, "text/event-stream") + .header( + reqwest::header::USER_AGENT, + crate::auth::antigravity_oauth::user_agent(), + ) + .json(&body) + .send() + .await + .map_err(|_| Error::RetryableProvider(RetryError::new("Antigravity request failed")))?; + if !response.status().is_success() { + let status = response.status(); + // Error bodies may echo credentials or request contents. Keep only + // the status and retry timing, never the body or request URL. + let error = RetryError::new(format!("Antigravity HTTP error {}", status.as_u16())) + .with_status(status.as_u16()) + .with_headers(response.headers()); + return if crate::aisdk::retry::retryable(&error) { + Err(Error::RetryableProvider(error)) + } else { + Err(Error::Provider(error.message)) + }; + } + Ok(gemini_event_stream( + response.bytes_stream().eventsource(), + Duration::from_secs(90), + Duration::from_millis(250), + )) + } +} + +fn gemini_event_stream(events: S, idle: Duration, grace: Duration) -> ProviderStream +where + S: futures::Stream> + Send + 'static, + E: Send + 'static, +{ + let state = ( + Box::pin(events), + GeminiStream::default(), + VecDeque::new(), + None::, + false, + ); + Box::pin(futures::stream::unfold( + state, + move |(mut events, mut parser, mut pending, mut deadline, mut done)| async move { + loop { + if let Some(chunk) = pending.pop_front() { + return Some((chunk, (events, parser, pending, deadline, done))); + } + if done { + return None; + } + // The terminal grace is absolute: keepalives must not extend it. + let until = deadline.unwrap_or_else(|| tokio::time::Instant::now() + idle); + match tokio::time::timeout_at(until, events.next()).await { + Ok(Some(Ok(event))) if event.data.trim() == "[DONE]" => { + pending.extend(parser.finish()); + done = true; + } + Ok(Some(Ok(event))) => { + pending.extend(parser.data(&event.data)); + if parser.failed { + done = true; + } else if parser.reason.is_some() && deadline.is_none() { + deadline = Some(tokio::time::Instant::now() + grace); + } + } + Ok(None) => { + pending.extend(parser.finish()); + done = true; + } + Err(_) if deadline.is_some() => { + pending.extend(parser.finish()); + done = true; + } + result => { + parser.failed = true; + let message = if result.is_err() { + "Antigravity SSE read timed out" + } else { + "Antigravity SSE transport failed" + }; + pending.push_back(Err(Error::Stream(message.into()))); + done = true; + } + } + } + }, + )) +} + +#[derive(Serialize, Deserialize)] +struct Replay { + parts: Vec, + call_ids: Vec, +} + +fn push_parts(contents: &mut Vec, role: &str, parts: Vec) { + if parts.is_empty() { + return; + } + if let Some(last) = contents.last_mut().filter(|last| last["role"] == role) { + last["parts"] + .as_array_mut() + .expect("constructed parts array") + .extend(parts); + } else { + contents.push(json!({"role": role, "parts": parts})); + } +} + +pub(super) fn gemini_messages(messages: &[Message]) -> Result<(Vec, Vec)> { + let mut contents = Vec::new(); + let mut system = Vec::new(); + let mut replay_calls = HashSet::new(); + let mut replay_text = None::; + for message in messages { + match message { + Message::System(message) => system.push(json!({"text": message.content})), + Message::User(message) => { + replay_text = None; + replay_calls.clear(); + if !message.audios.is_empty() { + return Err(Error::InvalidInput( + "Antigravity audio input is not supported".into(), + )); + } + let mut parts = Vec::new(); + if !message.content.is_empty() { + parts.push(json!({"text": message.content})); + } + for image in &message.images { + parts.push(image_part(image)?); + } + push_parts(&mut contents, "user", parts); + } + Message::Reasoning(message) => { + if let Some(payload) = message + .encrypted_content + .as_deref() + .and_then(|value| value.strip_prefix(REPLAY_PREFIX)) + { + let replay: Replay = serde_json::from_str(payload).map_err(|_| { + Error::InvalidInput("Invalid Antigravity replay metadata".into()) + })?; + replay_text = Some( + replay + .parts + .iter() + .filter(|part| part["thought"] != true) + .filter_map(|part| part["text"].as_str()) + .collect::() + .trim() + .to_owned(), + ); + replay_calls = replay.call_ids.into_iter().collect(); + push_parts(&mut contents, "model", replay.parts); + } + } + Message::Assistant(message) => { + if let Some(remaining) = replay_text.take() { + let text = message.content.trim(); + if let Some(rest) = remaining.trim_start().strip_prefix(text) { + if !rest.trim().is_empty() { + replay_text = Some(rest.to_owned()); + } + continue; + } + } + push_parts( + &mut contents, + "model", + vec![json!({"text": message.content})], + ); + } + Message::ToolCall(message) => { + if replay_calls.remove(&message.call_id) { + continue; + } + let args: Value = serde_json::from_str(&message.arguments).map_err(|_| { + Error::InvalidInput("Invalid Antigravity tool arguments in history".into()) + })?; + if !args.is_object() { + return Err(Error::InvalidInput( + "Antigravity tool arguments must be an object".into(), + )); + } + push_parts( + &mut contents, + "model", + vec![ + json!({"functionCall": {"name": message.name, "args": args, "id": message.call_id}}), + ], + ); + } + Message::ToolOutput(message) => { + replay_text = None; + replay_calls.clear(); + let response = if message.is_error { + json!({"error": message.output}) + } else { + json!({"output": message.output}) + }; + let mut parts = vec![ + json!({"functionResponse": {"name": message.name, "id": message.call_id, "response": response}}), + ]; + for image in &message.images { + parts.push(image_part(image)?); + } + push_parts(&mut contents, "user", parts); + } + } + } + Ok((contents, system)) +} + +fn image_part(image: &ImageContent) -> Result { + let data = if let Some((prefix, data)) = image.data_url.split_once(',') { + if !prefix.starts_with("data:") || !prefix.ends_with(";base64") { + return Err(Error::InvalidInput( + "Antigravity images must contain base64 data".into(), + )); + } + data + } else { + image.data_url.as_str() + }; + if !image.media_type.starts_with("image/") + || base64::engine::general_purpose::STANDARD + .decode(data) + .is_err() + || data.is_empty() + { + return Err(Error::InvalidInput( + "Invalid Antigravity image data or media type".into(), + )); + } + Ok(json!({"inlineData": {"mimeType": image.media_type, "data": data}})) +} + +// Internal Gemini endpoints accept the Schema subset, not arbitrary JSON Schema. +// Resolve local references before stripping unsupported validation keywords; +// validation remains the tool executor's responsibility. +fn gemini_schema(schema: &Value, root: &Value, depth: usize) -> Result { + if depth > 32 { + return Err(Error::InvalidInput( + "Recursive or deeply nested Antigravity tool schema".into(), + )); + } + if schema == &Value::Bool(true) { + return Ok(json!({"type": "OBJECT"})); + } + let object = schema + .as_object() + .ok_or_else(|| Error::InvalidInput("Unsupported Antigravity tool schema".into()))?; + if let Some(reference) = object.get("$ref").and_then(Value::as_str) { + let target = reference + .strip_prefix('#') + .and_then(|pointer| root.pointer(pointer)) + .ok_or_else(|| { + Error::InvalidInput("Unresolved Antigravity tool schema reference".into()) + })?; + return gemini_schema(target, root, depth + 1); + } + let mut out = serde_json::Map::new(); + for (key, value) in object { + match key.as_str() { + "type" => { + let types: Vec<&str> = match value { + Value::String(kind) => vec![kind], + Value::Array(types) => types.iter().filter_map(Value::as_str).collect(), + _ => { + return Err(Error::InvalidInput( + "Invalid Antigravity schema type".into(), + )) + } + }; + if types.contains(&"null") { + out.insert("nullable".into(), json!(true)); + } + let kinds: Vec<_> = types.into_iter().filter(|kind| *kind != "null").collect(); + if kinds.len() > 1 { + return Err(Error::InvalidInput( + "Antigravity does not support union schema types".into(), + )); + } + if let Some(kind) = kinds.first() { + out.insert("type".into(), json!(kind.to_uppercase())); + } + } + "properties" => { + let properties = value.as_object().ok_or_else(|| { + Error::InvalidInput("Invalid Antigravity schema properties".into()) + })?; + let mut converted = serde_json::Map::new(); + for (name, property) in properties { + converted.insert(name.clone(), gemini_schema(property, root, depth + 1)?); + } + out.insert(key.clone(), Value::Object(converted)); + } + "items" => { + out.insert(key.clone(), gemini_schema(value, root, depth + 1)?); + } + "anyOf" => { + let branches = value.as_array().ok_or_else(|| { + Error::InvalidInput("Invalid Antigravity schema alternatives".into()) + })?; + out.insert( + key.clone(), + Value::Array( + branches + .iter() + .map(|branch| gemini_schema(branch, root, depth + 1)) + .collect::>>()?, + ), + ); + } + "description" | "enum" | "required" | "nullable" | "format" => { + out.insert(key.clone(), value.clone()); + } + _ => {} + } + } + Ok(Value::Object(out)) +} + +#[derive(Default)] +struct GeminiStream { + parts: Vec, + call_ids: Vec, + usage: Option, + reason: Option, + failed: bool, + finished: bool, + part_bytes: usize, +} + +impl GeminiStream { + fn data(&mut self, data: &str) -> Vec> { + if self.failed || self.finished || data.trim().is_empty() || data.trim() == "[DONE]" { + return Vec::new(); + } + if data.len() > 8 * 1024 * 1024 { + self.failed = true; + return vec![Err(Error::Stream( + "Antigravity SSE event exceeds limit".into(), + ))]; + } + let result = serde_json::from_str::(data) + .map_err(|_| Error::Stream("Malformed Antigravity SSE JSON".into())) + .and_then(|value| self.response(&value)); + match result { + Ok(chunks) => chunks.into_iter().map(Ok).collect(), + Err(error) => { + self.failed = true; + vec![Err(error)] + } + } + } + + fn response(&mut self, envelope: &Value) -> Result> { + if envelope.get("error").is_some() { + return Err(Error::Provider( + "Antigravity returned a stream error".into(), + )); + } + let response = envelope.get("response").unwrap_or(envelope); + if response.get("error").is_some() { + return Err(Error::Provider( + "Antigravity returned a stream error".into(), + )); + } + if let Some(usage) = response.get("usageMetadata") { + let token = |key| -> Result { + match usage.get(key) { + None => Ok(0), + Some(value) => value + .as_u64() + .ok_or_else(|| Error::Stream("Invalid Antigravity token usage".into())), + } + }; + let cached = token("cachedContentTokenCount")?; + self.usage = Some(TokenUsage { + input: token("promptTokenCount")?.saturating_sub(cached), + output: token("candidatesTokenCount")?.saturating_add(token("thoughtsTokenCount")?), + cache_read: cached, + cache_write: 0, + }); + } + if response + .pointer("/promptFeedback/blockReason") + .and_then(Value::as_str) + .is_some_and(|reason| reason != "BLOCK_REASON_UNSPECIFIED") + { + self.reason = Some(FinishReason::Refusal); + return Ok(Vec::new()); + } + let Some(candidates) = response.get("candidates") else { + if response.get("usageMetadata").is_some() { + return Ok(Vec::new()); + } + return Err(Error::Stream( + "Antigravity SSE response has no candidates".into(), + )); + }; + let candidates = candidates + .as_array() + .ok_or_else(|| Error::Stream("Invalid Antigravity candidates".into()))?; + if candidates.len() > 1 { + return Err(Error::Stream( + "Multiple Antigravity candidates are unsupported".into(), + )); + } + let mut chunks = Vec::new(); + for candidate in candidates { + if let Some(content) = candidate.get("content") { + let parts = content + .get("parts") + .and_then(Value::as_array) + .ok_or_else(|| Error::Stream("Invalid Antigravity content parts".into()))?; + for part in parts { + if !part.is_object() { + return Err(Error::Stream("Invalid Antigravity part".into())); + } + if let Some(signature) = part.get("thoughtSignature") { + if !signature.is_string() { + return Err(Error::Stream( + "Invalid Antigravity thought signature".into(), + )); + } + } + if let Some(text) = part.get("text") { + let text = text + .as_str() + .ok_or_else(|| Error::Stream("Invalid Antigravity text".into()))?; + if part["thought"] == true { + chunks.push(ChunkType::Reasoning(text.into())); + } else { + chunks.push(ChunkType::Text(text.into())); + } + } + if let Some(call) = part.get("functionCall") { + let name = call + .get("name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()) + .ok_or_else(|| { + Error::Stream("Antigravity function call has no name".into()) + })?; + let args = call.get("args").cloned().unwrap_or_else(|| json!({})); + if !args.is_object() { + return Err(Error::Stream( + "Invalid Antigravity function arguments".into(), + )); + } + let id = match call.get("id") { + Some(value) => value + .as_str() + .filter(|id| !id.is_empty()) + .map(str::to_owned) + .ok_or_else(|| { + Error::Stream("Invalid Antigravity function ID".into()) + })?, + None => format!("ag_{}", cuid2::create_id()), + }; + if self.call_ids.contains(&id) { + return Err(Error::Stream("Duplicate Antigravity function ID".into())); + } + chunks.push(ChunkType::ToolCall( + json!([{ + "index": self.call_ids.len(), "id": id, "type": "function", + "function": {"name": name, "arguments_done": args.to_string()} + }]) + .to_string(), + )); + self.call_ids.push(id); + } + if part.get("text").is_none() + && part.get("functionCall").is_none() + && part.get("thoughtSignature").is_none() + { + return Err(Error::Stream( + "Unsupported Antigravity response part".into(), + )); + } + self.part_bytes = self.part_bytes.saturating_add(part.to_string().len()); + if self.parts.len() >= 65_536 || self.part_bytes > 16 * 1024 * 1024 { + return Err(Error::Stream("Antigravity replay exceeds limit".into())); + } + self.parts.push(part.clone()); + } + } + if let Some(reason) = candidate.get("finishReason") { + let reason = reason + .as_str() + .ok_or_else(|| Error::Stream("Invalid Antigravity finish reason".into()))?; + self.reason = Some(match reason { + "STOP" => { + if self.call_ids.is_empty() { + FinishReason::Stop + } else { + FinishReason::ToolCalls + } + } + "MAX_TOKENS" => FinishReason::Length, + "SAFETY" | "RECITATION" | "BLOCKLIST" | "PROHIBITED_CONTENT" | "SPII" + | "IMAGE_SAFETY" => FinishReason::Refusal, + // Tool execution must never proceed after a failed generation. + _ => { + return Err(Error::Stream( + "Antigravity generation did not finish successfully".into(), + )) + } + }); + } + } + Ok(chunks) + } + + fn finish(&mut self) -> Vec> { + if self.failed || self.finished { + return Vec::new(); + } + self.finished = true; + let Some(reason) = self.reason.take() else { + return vec![Err(Error::Stream( + "Antigravity stream closed before finishReason".into(), + ))]; + }; + let mut chunks = Vec::new(); + // Store the *entire* ordered model turn, not just a signature string: + // signatures bind to their original parts, including text/thought parts. + // The SDK reorders reasoning/text/calls in history; this opaque sibling + // restores the exact wire parts and suppresses their flattened duplicates. + if self + .parts + .iter() + .any(|part| part.get("thoughtSignature").is_some()) + { + let replay = Replay { + parts: std::mem::take(&mut self.parts), + call_ids: self.call_ids.clone(), + }; + match serde_json::to_string(&replay) { + Ok(payload) => chunks.push(Ok(ChunkType::ReasoningItem(ReasoningReplayItem { + id: Some(format!("ag_replay_{}", cuid2::create_id())), + summary: String::new(), + encrypted_content: Some(format!("{REPLAY_PREFIX}{payload}")), + }))), + Err(_) => { + return vec![Err(Error::Stream( + "Could not preserve Antigravity thought signatures".into(), + ))] + } + } + } + if let Some(usage) = self.usage.take() { + chunks.push(Ok(ChunkType::Usage(usage))); + } + chunks.push(Ok(ChunkType::End { + reason: Some(reason), + })); + chunks + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn model() -> Antigravity { + Antigravity::new( + "http://localhost".into(), + "gemini-test".into(), + "secret-token".into(), + "project".into(), + ) + } + + #[test] + fn serializes_messages_images_tools_and_stable_session() { + let tool = Tool::builder().name("read").description("Read a file") + .input_schema(serde_json::from_value(json!({"type":"object","properties":{"path":{"type":"string"}},"required":["path"],"additionalProperties":false})).unwrap()) + .execute(crate::aisdk::tool::ToolExecute::new(|_| async { Ok::<_, String>("done") })).build().unwrap(); + let messages = vec![ + Message::system("Use tools"), + Message::user_with_images( + "look", + vec![ImageContent { + data_url: "data:image/png;base64,aGk=".into(), + media_type: "image/png".into(), + }], + ), + Message::assistant("reading"), + Message::tool_call("call", "read", r#"{"path":"a"}"#), + Message::tool_output("call", "read", "contents", false), + ]; + let headers = HashMap::from([ + ("X-Opencode-Session".into(), "session".into()), + ("x-api-key".into(), "never-send".into()), + ]); + let body = model().request_body(&messages, &[tool], &headers).unwrap(); + assert_eq!(body["model"], "gemini-test"); + assert_eq!(body["request"]["sessionId"], "session"); + assert_eq!( + body["request"]["systemInstruction"]["parts"][0]["text"], + "Use tools" + ); + assert_eq!( + body["request"]["contents"][0]["parts"][1]["inlineData"]["data"], + "aGk=" + ); + assert_eq!(body["request"]["contents"][1]["role"], "model"); + assert_eq!( + body["request"]["contents"][1]["parts"][1]["functionCall"]["args"]["path"], + "a" + ); + assert_eq!( + body["request"]["contents"][2]["parts"][0]["functionResponse"]["response"]["output"], + "contents" + ); + assert_eq!( + body["request"]["tools"][0]["functionDeclarations"][0]["parameters"]["type"], + "OBJECT" + ); + assert!(!body.to_string().contains("secret-token")); + assert!(!body.to_string().contains("never-send")); + assert!(!format!("{:?}", model()).contains("secret-token")); + } + + #[test] + fn signatures_restore_exact_parts_without_duplicating_calls_or_text() { + let parts = json!([ + {"text":"thinking", "thought":true,"thoughtSignature":"thought-signature"}, + {"text":"inspect", "thoughtSignature":"text-signature"}, + {"functionCall":{"name":"read","args":{"path":"a"}},"thoughtSignature":"call-signature"}, + {"functionCall":{"name":"read","args":{"path":"b"}}} + ]); + let mut state = GeminiStream::default(); + let chunks = state.response(&json!({"response":{"candidates":[{"content":{"role":"model","parts":parts},"finishReason":"STOP"}]}})).unwrap(); + assert!(matches!(&chunks[0], ChunkType::Reasoning(text) if text == "thinking")); + let calls = state.call_ids.clone(); + let terminal = state.finish(); + let replay = terminal + .into_iter() + .find_map(|chunk| match chunk.unwrap() { + ChunkType::ReasoningItem(item) => Some(item), + _ => None, + }) + .unwrap(); + let messages = vec![ + Message::user("start"), + Message::reasoning(replay.id, replay.summary, replay.encrypted_content), + Message::assistant("inspect"), + Message::tool_call(&calls[0], "read", r#"{"path":"a"}"#), + Message::tool_call(&calls[1], "read", r#"{"path":"b"}"#), + Message::tool_output(&calls[0], "read", "a", false), + Message::tool_output(&calls[1], "read", "b", false), + ]; + let (contents, _) = gemini_messages(&messages).unwrap(); + assert_eq!(contents.len(), 3); + assert_eq!(contents[1]["parts"], parts); + assert_eq!(contents[2]["parts"].as_array().unwrap().len(), 2); + } + + #[test] + fn persisted_replay_consumes_split_sanitized_assistant_text() { + for (first, second) in [("first", "second"), (" \n first ", " \tsecond\n ")] { + let parts = json!([ + {"text":first,"thoughtSignature":"signed-first"}, + {"text":"thinking","thought":true,"thoughtSignature":"signed-thought"}, + {"text":second,"thoughtSignature":"signed-second"} + ]); + let mut state = GeminiStream::default(); + state + .response( + &json!({"candidates":[{"content":{"parts":parts},"finishReason":"STOP"}]}), + ) + .unwrap(); + let item = state + .finish() + .into_iter() + .find_map(|chunk| match chunk.unwrap() { + ChunkType::ReasoningItem(item) => Some(item), + _ => None, + }) + .unwrap(); + let mut stored = crate::session::types::Message::new( + crate::session::types::MessageRole::Assistant, + "", + ); + stored.append(first); + stored.append_reasoning("thinking"); + stored.append(second); + persist_replay(&mut stored, item); + assert_eq!(stored.parts[0].part_type, "provider_replay"); + let replay = &stored.parts[0].data; + // Match the host's flattened, trimmed text messages without calling its private converter. + let mut messages = vec![Message::reasoning( + replay["id"].as_str().map(str::to_owned), + "", + replay["encrypted_content"].as_str().map(str::to_owned), + )]; + for part in &stored.parts { + if part.part_type == "text" { + messages.push(Message::assistant( + part.data["text"].as_str().unwrap().trim(), + )); + } + } + messages.push(Message::assistant("new text")); + let (contents, _) = gemini_messages(&messages).unwrap(); + let mut expected = parts.as_array().unwrap().clone(); + expected.push(json!({"text":"new text"})); + assert_eq!(contents[0]["parts"], json!(expected)); + } + } + + async fn held_open_server( + first: String, + trailing: Option, + ) -> (String, tokio::task::JoinHandle<()>) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + loop { + let mut buffer = [0; 4096]; + let count = socket.read(&mut buffer).await.unwrap(); + assert!(count > 0); + request.extend_from_slice(&buffer[..count]); + if let Some(index) = request.windows(4).position(|bytes| bytes == b"\r\n\r\n") { + let headers = String::from_utf8_lossy(&request[..index]).to_lowercase(); + let length = headers + .lines() + .find_map(|line| line.strip_prefix("content-length:")) + .map(|value| value.trim().parse::().unwrap()) + .unwrap_or(0); + if request.len() >= index + 4 + length { + break; + } + } + } + socket.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: keep-alive\r\n\r\n").await.unwrap(); + socket.write_all(first.as_bytes()).await.unwrap(); + if let Some(trailing) = trailing { + tokio::time::sleep(Duration::from_millis(10)).await; + socket.write_all(trailing.as_bytes()).await.unwrap(); + } + std::future::pending::<()>().await; + }); + (base, task) + } + + fn terminal_event() -> String { + format!( + "data: {}\n\n", + json!({"candidates":[{"content":{"parts":[{"text":"first"},{"text":"second"}]},"finishReason":"STOP"}]}) + ) + } + + #[tokio::test] + async fn held_open_terminal_preserves_chunks_and_trailing_usage() { + let (base, server) = held_open_server( + terminal_event(), + Some(format!( + "data: {}\n\n", + json!({"usageMetadata":{"candidatesTokenCount":7}}) + )), + ) + .await; + let response = reqwest::get(base).await.unwrap(); + let chunks = tokio::time::timeout( + Duration::from_secs(2), + gemini_event_stream( + response.bytes_stream().eventsource(), + Duration::from_secs(1), + Duration::from_millis(80), + ) + .collect::>(), + ) + .await + .unwrap(); + server.abort(); + assert_eq!(chunks.len(), 4); + assert!(matches!(&chunks[0], Ok(ChunkType::Text(text)) if text == "first")); + assert!(matches!(&chunks[1], Ok(ChunkType::Text(text)) if text == "second")); + assert!(matches!(&chunks[2], Ok(ChunkType::Usage(usage)) if usage.output == 7)); + assert!(matches!(&chunks[3], Ok(ChunkType::End { .. }))); + } + + #[tokio::test] + async fn done_terminates_immediately_and_failures_do_not_finish() { + for (event, success) in [ + (format!("{}data: [DONE]\n\n", terminal_event()), true), + ("data: malformed\n\n".into(), false), + ("data: [DONE]\n\n".into(), false), + ] { + let (base, server) = held_open_server(event, None).await; + let response = reqwest::get(base).await.unwrap(); + let chunks = tokio::time::timeout( + Duration::from_secs(1), + gemini_event_stream( + response.bytes_stream().eventsource(), + Duration::from_secs(30), + Duration::from_secs(30), + ) + .collect::>(), + ) + .await + .unwrap(); + server.abort(); + if success { + assert_eq!(chunks.len(), 3); + assert_eq!( + chunks + .iter() + .filter(|chunk| matches!(chunk, Ok(ChunkType::End { .. }))) + .count(), + 1 + ); + } else { + assert_eq!(chunks.len(), 1); + assert!(chunks[0].is_err()); + } + } + } + + #[tokio::test] + async fn terminal_grace_is_not_extended_by_more_events() { + let events = futures::stream::unfold(0, |index| async move { + tokio::time::sleep(Duration::from_millis(5)).await; + let data = if index == 0 { + json!({"candidates":[{"finishReason":"STOP"}]}).to_string() + } else { + json!({"usageMetadata":{"candidatesTokenCount":index}}).to_string() + }; + Some(( + Ok::<_, std::io::Error>(eventsource_stream::Event { + data, + ..Default::default() + }), + index + 1, + )) + }); + let chunks = tokio::time::timeout( + Duration::from_secs(1), + gemini_event_stream(events, Duration::from_secs(30), Duration::from_millis(40)) + .collect::>(), + ) + .await + .unwrap(); + assert_eq!(chunks.len(), 2); + assert!(matches!(&chunks[0], Ok(ChunkType::Usage(usage)) if usage.output > 0)); + assert!(matches!(&chunks[1], Ok(ChunkType::End { .. }))); + } + + #[tokio::test] + async fn eof_and_transport_errors_finalize_only_once() { + for fail in [false, true] { + let mut events = vec![Ok(eventsource_stream::Event { + data: json!({"candidates":[{"finishReason":"STOP"}]}).to_string(), + ..Default::default() + })]; + if fail { + events.push(Err(std::io::Error::other("disconnected"))); + } + let chunks = gemini_event_stream( + futures::stream::iter(events), + Duration::from_secs(1), + Duration::from_secs(1), + ) + .collect::>() + .await; + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].is_err(), fail); + } + let mut parser = GeminiStream::default(); + parser + .response(&json!({"candidates":[{"finishReason":"STOP"}]})) + .unwrap(); + assert_eq!(parser.finish().len(), 1); + assert!(parser.finish().is_empty()); + assert!(parser.data("malformed").is_empty()); + } + + #[tokio::test] + async fn stalled_reads_are_bounded_even_inside_incomplete_sse_frames() { + for request_timeout in [false, true] { + let (base, server) = held_open_server("data: {".into(), None).await; + let client = reqwest::Client::builder() + .read_timeout(if request_timeout { + Duration::from_millis(40) + } else { + Duration::from_secs(30) + }) + .build() + .unwrap(); + let response = client.get(base).send().await.unwrap(); + let chunks = tokio::time::timeout( + Duration::from_secs(2), + gemini_event_stream( + response.bytes_stream().eventsource(), + Duration::from_millis(100), + Duration::from_millis(20), + ) + .collect::>(), + ) + .await + .unwrap(); + server.abort(); + assert_eq!(chunks.len(), 1); + assert!(chunks[0].is_err()); + } + } + + #[tokio::test] + async fn sdk_cancellation_interrupts_a_stalled_provider() { + let event = format!( + "data: {}\n\n", + json!({"candidates":[{"content":{"parts":[{"text":"ready"}]}}]}) + ); + let (base, server) = held_open_server(event, None).await; + let provider = Antigravity::new(base, "test".into(), "token".into(), "project".into()); + let cancel = tokio_util::sync::CancellationToken::new(); + let mut response = crate::aisdk::response::stream_with_tools_options( + provider, + vec![Message::user("hello")], + vec![], + Some(1), + None, + HashMap::new(), + Some(cancel.clone()), + crate::aisdk::response::StreamWithToolsOptions::default(), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(2), async { + while let Some(chunk) = response.stream.next().await { + if matches!(chunk, ChunkType::Text(_)) { + break; + } + } + cancel.cancel(); + while response.stream.next().await.is_some() {} + }) + .await + .unwrap(); + server.abort(); + } + + #[test] + fn cumulative_usage_is_emitted_once_including_reasoning_and_cache() { + let mut state = GeminiStream::default(); + state.response(&json!({"response":{"usageMetadata":{"promptTokenCount":100,"cachedContentTokenCount":40,"candidatesTokenCount":3}}})).unwrap(); + state.response(&json!({"response":{"candidates":[{"finishReason":"MAX_TOKENS"}],"usageMetadata":{"promptTokenCount":100,"cachedContentTokenCount":40,"candidatesTokenCount":10,"thoughtsTokenCount":5}}})).unwrap(); + let chunks = state.finish(); + assert!(matches!( + &chunks[0], + Ok(ChunkType::Usage(TokenUsage { + input: 60, + output: 15, + cache_read: 40, + cache_write: 0 + })) + )); + assert!(matches!( + &chunks[1], + Ok(ChunkType::End { + reason: Some(FinishReason::Length) + }) + )); + } + + #[test] + fn malformed_errors_and_truncation_fail_without_echoing_payloads() { + for data in [ + "secret-token", + r#"{"error":{"message":"secret-token"}}"#, + r#"{"response":{"error":{"message":"secret-token"}}}"#, + r#"{"response":{"candidates":[{"content":{"parts":[{"functionCall":{"name":"read","args":[]}}]}}]}}"#, + ] { + let mut state = GeminiStream::default(); + let chunks = state.data(data); + assert!(chunks[0].is_err()); + assert!(!chunks[0] + .as_ref() + .unwrap_err() + .to_string() + .contains("secret-token")); + assert!(state.finish().is_empty()); + } + assert!(GeminiStream::default().finish()[0].is_err()); + let mut state = GeminiStream::default(); + state + .response(&json!({"response":{"candidates":[{"finishReason":"SAFETY"}]}})) + .unwrap(); + assert!(matches!( + &state.finish()[0], + Ok(ChunkType::End { + reason: Some(FinishReason::Refusal) + }) + )); + } + + #[tokio::test] + async fn sse_parser_handles_split_bytes_and_multiline_data() { + let bytes = futures::stream::iter(vec![ + Ok::<_, std::io::Error>(b"data: {\"response\":{\"candi".to_vec()), + Ok(b"dates\":[{\"content\":{\"parts\":[{\"text\":\"hello\"}]},\n".to_vec()), + Ok(b"data: \"finishReason\":\"STOP\"}]}}\n\n".to_vec()), + ]); + let mut events = bytes.eventsource(); + let mut state = GeminiStream::default(); + let event = events.next().await.unwrap().unwrap(); + assert!( + matches!(&state.data(&event.data)[0], Ok(ChunkType::Text(text)) if text == "hello") + ); + assert!(events.next().await.is_none()); + assert!(matches!( + &state.finish()[0], + Ok(ChunkType::End { + reason: Some(FinishReason::Stop) + }) + )); + } + #[tokio::test] + async fn http_tool_loop_replays_signatures_and_never_forwards_api_keys() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let mut requests = Vec::new(); + for step in 0..2 { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut bytes = Vec::new(); + let (header_end, length) = loop { + let mut buffer = [0u8; 4096]; + let count = socket.read(&mut buffer).await.unwrap(); + assert!(count > 0); + bytes.extend_from_slice(&buffer[..count]); + if let Some(index) = bytes.windows(4).position(|window| window == b"\r\n\r\n") { + let header = String::from_utf8_lossy(&bytes[..index]).to_lowercase(); + let length = header + .lines() + .find_map(|line| { + line.strip_prefix("content-length:") + .map(|value| value.trim().parse::().unwrap()) + }) + .unwrap(); + break (index + 4, length); + } + }; + while bytes.len() < header_end + length { + let mut buffer = [0u8; 4096]; + let count = socket.read(&mut buffer).await.unwrap(); + assert!(count > 0); + bytes.extend_from_slice(&buffer[..count]); + } + let headers = String::from_utf8_lossy(&bytes[..header_end]).to_lowercase(); + assert!(headers.starts_with("post /v1internal:streamgeneratecontent?alt=sse ")); + assert!(headers.contains("authorization: bearer secret-token")); + assert!(!headers.contains("api-key")); + assert!(!headers.contains("never-send")); + requests.push( + serde_json::from_slice::(&bytes[header_end..header_end + length]) + .unwrap(), + ); + let parts = if step == 0 { + json!([{ "functionCall": {"name":"read", "args":{}, "id":"call-1"}, "thoughtSignature":"opaque-signature" }]) + } else { + json!([{ "text":"done" }]) + }; + let event = format!( + "data: {}\n\n", + json!({"response":{"candidates":[{"content":{"role":"model","parts":parts},"finishReason":"STOP"}]}}) + ); + socket.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", event.len()).as_bytes()).await.unwrap(); + for chunk in event.as_bytes().chunks(7) { + socket.write_all(chunk).await.unwrap(); + } + } + requests + }); + let tool = Tool::builder() + .name("read") + .description("Read") + .input_schema(serde_json::from_value(json!({"type":"object","properties":{}})).unwrap()) + .execute(crate::aisdk::tool::ToolExecute::new(|_| async { + Ok::<_, String>("result") + })) + .build() + .unwrap(); + let provider = Antigravity::new( + base, + "gemini-test".into(), + "secret-token".into(), + "project".into(), + ); + let mut response = crate::aisdk::response::stream_with_tools_options( + provider, + vec![Message::user("read")], + vec![tool], + Some(2), + None, + HashMap::from([ + ("x-api-key".into(), "never-send".into()), + ("x-opencode-session".into(), "stable".into()), + ]), + None, + crate::aisdk::response::StreamWithToolsOptions::default(), + ) + .await + .unwrap(); + let mut text = String::new(); + while let Some(chunk) = response.stream.next().await { + match chunk { + ChunkType::Text(value) => text.push_str(&value), + ChunkType::Failed(error) => panic!("{error}"), + _ => {} + } + } + assert_eq!(text, "done"); + let requests = server.await.unwrap(); + assert_eq!( + requests[0]["request"]["sessionId"], + requests[1]["request"]["sessionId"] + ); + let contents = &requests[1]["request"]["contents"]; + assert_eq!( + contents[1]["parts"][0]["thoughtSignature"], + "opaque-signature" + ); + assert_eq!(contents[1]["parts"].as_array().unwrap().len(), 1); + assert_eq!( + contents[2]["parts"][0]["functionResponse"]["response"]["output"], + "result" + ); + } +} diff --git a/src/llm/client.rs b/src/llm/client.rs index 6dfe42aa..b19a8a8e 100644 --- a/src/llm/client.rs +++ b/src/llm/client.rs @@ -662,14 +662,16 @@ pub async fn stream_llm_with_cancellation( messages.len() ); let ui_model = model.clone(); - let request_config = prepare_request_config( - &provider_name, - model, - reasoning_effort, - messages_have_user_audio(&messages), - &sender, - ) - .await?; + let request_config = tokio::select! { + biased; + _ = cancel_token.cancelled() => { + let _ = sender.send(crate::llm::ChunkMessage::Cancelled); + return Ok(()); + } + config = prepare_request_config( + &provider_name, model, reasoning_effort, messages_have_user_audio(&messages), &sender, + ) => config?, + }; let mut request_config = request_config; request_config.prune_tool_outputs = compaction_config.prune(); let model_mismatch_warning = @@ -709,6 +711,7 @@ pub async fn stream_llm_with_cancellation( ProviderKind::OpenAI => crate::agent::config::ProviderKind::OpenAI, ProviderKind::OpenAICompatible => crate::agent::config::ProviderKind::OpenAICompatible, ProviderKind::Anthropic => crate::agent::config::ProviderKind::Anthropic, + ProviderKind::Antigravity => crate::agent::config::ProviderKind::Antigravity, }, base_url: request_config.base_url.clone(), reasoning_effort: request_config.reasoning_effort, @@ -971,6 +974,7 @@ pub async fn build_subagent_llm_session( ProviderKind::OpenAI => crate::agent::config::ProviderKind::OpenAI, ProviderKind::OpenAICompatible => crate::agent::config::ProviderKind::OpenAICompatible, ProviderKind::Anthropic => crate::agent::config::ProviderKind::Anthropic, + ProviderKind::Antigravity => crate::agent::config::ProviderKind::Antigravity, }, base_url: request_config.base_url, reasoning_effort: request_config.reasoning_effort, @@ -1045,14 +1049,15 @@ pub async fn summarize_for_compaction( } let (warning_sender, _warning_receiver) = tokio::sync::mpsc::unbounded_channel(); - let request_config = prepare_request_config( - &provider_name, - model, - reasoning_effort, - false, - &warning_sender, - ) - .await?; + let request_config = tokio::select! { + biased; + _ = cancel_token.cancelled() => { + return Err(anyhow::anyhow!("Compaction cancelled by user").into()); + } + config = prepare_request_config( + &provider_name, model, reasoning_effort, false, &warning_sender, + ) => config?, + }; let messages = vec![AisdkMessage::user(prompt)]; let mut response = stream_provider_request( &request_config, @@ -1394,6 +1399,15 @@ async fn prepare_request_config( &mut request_config, ); + maybe_apply_google_oauth_overrides( + provider_name, + &auth_dao, + auth_config.as_ref(), + &mut request_config, + sender, + ) + .await?; + if request_config.api_key.is_none() && !crate::model::extensions::ModelExtensions::is_runtime_provider(provider_name) { @@ -1578,6 +1592,7 @@ async fn maybe_apply_openai_oauth_overrides( expires, account_id, enterprise_url, + project_id, }) = auth_config.cloned() else { return Ok(()); @@ -1613,6 +1628,7 @@ async fn maybe_apply_openai_oauth_overrides( expires: oauth_expires, account_id: oauth_account_id.clone(), enterprise_url: oauth_enterprise_url.clone(), + project_id: project_id.clone(), }, ); } @@ -1691,6 +1707,7 @@ async fn maybe_apply_xai_oauth_overrides( expires, account_id, enterprise_url, + project_id, }) = auth_config.cloned() else { return; @@ -1720,6 +1737,7 @@ async fn maybe_apply_xai_oauth_overrides( expires: oauth_expires, account_id: oauth_account_id.clone(), enterprise_url: oauth_enterprise_url.clone(), + project_id: project_id.clone(), }, ); } @@ -1747,6 +1765,46 @@ async fn maybe_apply_xai_oauth_overrides( crate::emit_log!("Configured xAI Grok Build OAuth transport"); } +async fn maybe_apply_google_oauth_overrides( + provider_name: &str, + _auth_dao: &crate::persistence::AuthDAO, + auth_config: Option<&crate::persistence::AuthConfig>, + request_config: &mut ProviderRequestConfig, + _sender: &crate::llm::ChunkSender, +) -> Result<(), DynError> { + if !uses_google_oauth(provider_name, auth_config) { + return Ok(()); + } + let credentials = crate::auth::antigravity_oauth::load_request_credentials().await?; + apply_google_oauth_credentials(request_config, &credentials); + Ok(()) +} + +fn uses_google_oauth( + provider_name: &str, + auth_config: Option<&crate::persistence::AuthConfig>, +) -> bool { + provider_name == "google" + && matches!( + auth_config, + Some(crate::persistence::AuthConfig::OAuth { .. }) + ) +} + +fn apply_google_oauth_credentials( + request_config: &mut ProviderRequestConfig, + credentials: &crate::auth::OAuthCredentials, +) { + request_config.kind = ProviderKind::Antigravity; + request_config.base_url = crate::auth::antigravity_oauth::ANTIGRAVITY_ENDPOINT.to_string(); + request_config.api_key = Some(credentials.access.clone()); + request_config.openai_options.antigravity_project = credentials.project_id.clone(); + request_config.openai_options.additional_headers.insert( + "User-Agent".to_string(), + crate::auth::antigravity_oauth::user_agent().to_string(), + ); +} + fn send_warning(sender: &crate::llm::ChunkSender, warning: impl Into) { let _ = sender.send(crate::llm::ChunkMessage::Warning(warning.into())); } @@ -1837,13 +1895,24 @@ fn stamp_build_parent_cached_aux( ); } +fn retain_provider_replay(messages: &mut Vec, kind: ProviderKind) { + if kind != ProviderKind::Antigravity { + messages.retain(|message| { + !matches!(message, AisdkMessage::Reasoning(item) if item + .encrypted_content.as_deref().is_some_and(|value| value.starts_with( + crate::llm::antigravity::REPLAY_PREFIX))) + }); + } +} + async fn stream_provider_request( config: &ProviderRequestConfig, - messages: Vec, + mut messages: Vec, tools: Vec, max_steps: Option, cancel_token: Option, ) -> Result { + retain_provider_replay(&mut messages, config.kind); let headers = super::opencode::ensure_session_headers( &config.provider_name, &config.base_url, @@ -1852,6 +1921,40 @@ async fn stream_provider_request( &std::collections::HashMap::new(), ); match config.kind { + ProviderKind::Antigravity => { + let access = config + .api_key + .clone() + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| anyhow::anyhow!("Missing Google Antigravity OAuth access token"))?; + let project = config + .openai_options + .antigravity_project + .clone() + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| anyhow::anyhow!("Missing Google Antigravity OAuth project"))?; + let provider = super::antigravity::Antigravity::new( + config.base_url.clone(), + config.model_name.clone(), + access, + project, + ) + .with_stored_credentials(); + stream_with_tools_options( + provider, + messages, + tools, + max_steps, + None, + headers, + cancel_token, + StreamWithToolsOptions { + prune_tool_outputs: config.prune_tool_outputs, + }, + ) + .await + .map_err(|e| Box::new(e) as DynError) + } ProviderKind::OpenAICompatible => { let mut builder = OpenAICompatible::builder() .base_url(&config.base_url) @@ -2143,9 +2246,16 @@ async fn relay_stream_to_sender( crate::emit_log!("[RELAY] Reasoning chunk ({} chars)", reasoning.len()); let _ = sender.send(crate::llm::ChunkMessage::Reasoning(reasoning)); } - ChunkType::ReasoningItem(_) => { + ChunkType::ReasoningItem(item) => { let elapsed_ms = start_time.elapsed().as_millis(); stats.record_chunk("ReasoningItem", elapsed_ms); + if item + .encrypted_content + .as_deref() + .is_some_and(|value| value.starts_with(crate::llm::antigravity::REPLAY_PREFIX)) + { + let _ = sender.send(crate::llm::ChunkMessage::ProviderReplay(item)); + } } ChunkType::ToolCall(tool_call) => { let elapsed_ms = start_time.elapsed().as_millis(); @@ -2512,7 +2622,7 @@ fn convert_messages_for_model_with_audio( if msg.parts.iter().any(|part| { matches!( part.part_type.as_str(), - "text" | "reasoning" | "tool_call" | "tool_result" + "text" | "reasoning" | "tool_call" | "tool_result" | "provider_replay" ) }) { append_assistant_parts_for_model( @@ -2553,6 +2663,32 @@ fn append_assistant_parts_for_model( for part in &msg.parts { match part.part_type.as_str() { + "provider_replay" => { + let Some(obj) = part.data.as_object() else { + continue; + }; + let id = obj + .get("id") + .and_then(|value| value.as_str()) + .map(str::to_owned); + let summary = obj + .get("summary") + .and_then(|value| value.as_str()) + .unwrap_or_default() + .to_owned(); + let encrypted_content = obj + .get("encrypted_content") + .or_else(|| obj.get("encryptedContent")) + .and_then(|value| value.as_str()) + .map(str::to_owned); + if encrypted_content + .as_deref() + .is_some_and(|value| value.starts_with(crate::llm::antigravity::REPLAY_PREFIX)) + { + pending_tools.flush_complete_pairs(aisdk_messages); + aisdk_messages.push(AisdkMessage::reasoning(id, summary, encrypted_content)); + } + } "text" => { pending_tools.flush_complete_pairs(aisdk_messages); @@ -2888,6 +3024,7 @@ enum ProviderKind { OpenAI, OpenAICompatible, Anthropic, + Antigravity, } impl ProviderKind { @@ -2941,6 +3078,185 @@ mod tests { use crate::persistence::AuthConfig; use base64::Engine as _; + #[test] + fn antigravity_signed_history_survives_multiple_steps_and_next_user_turn() { + use crate::aisdk::chunk::ReasoningReplayItem; + use crate::llm::antigravity::{gemini_messages, persist_replay, REPLAY_PREFIX}; + use crate::session::types::Message; + use serde_json::json; + + let mut assistant = Message::assistant(""); + let mut signed_parts = Vec::new(); + for step in 0..2 { + let id = format!("call_{step}"); + let text = format!("step {step}"); + let parts = json!([ + {"text":"private", "thought":true,"thoughtSignature":format!("thought_{step}")}, + {"text":text,"thoughtSignature":format!("text_{step}")}, + {"functionCall":{"name":"read","args":{"path":step}},"thoughtSignature":format!("call_signature_{step}")} + ]); + assistant.append(&text); + persist_replay( + &mut assistant, + ReasoningReplayItem { + id: Some(format!("replay_{step}")), + summary: String::new(), + encrypted_content: Some(format!( + "{REPLAY_PREFIX}{}", + json!({"parts":parts,"call_ids":[id]}) + )), + }, + ); + assistant.add_tool_call_part(&id, "read", json!({"path":step})); + assistant.add_or_update_tool_result_part(json!({ + "id":id,"name":"read","args":{"path":step},"status":"ok","output_preview":format!("result {step}") + })); + signed_parts.push(parts); + } + assistant.append("done"); + // The same flattened JSON representation is stored in session snapshots. + assistant.parts = + serde_json::from_value(serde_json::to_value(&assistant.parts).unwrap()).unwrap(); + let history = vec![Message::user("start"), assistant, Message::user("continue")]; + let messages = convert_messages(&history); + let (contents, _) = gemini_messages(&messages).unwrap(); + assert_eq!(contents.len(), 7); + for (step, expected) in signed_parts.iter().enumerate() { + assert_eq!(&contents[1 + step * 2]["parts"], expected); + assert_eq!( + contents[2 + step * 2]["parts"][0]["functionResponse"]["id"], + format!("call_{step}") + ); + } + assert_eq!(contents[5]["parts"], json!([{"text":"done"}])); + assert_eq!(contents[6]["parts"], json!([{"text":"continue"}])); + + for kind in [ + ProviderKind::OpenAI, + ProviderKind::Anthropic, + ProviderKind::OpenAICompatible, + ] { + let mut switched = messages.clone(); + super::retain_provider_replay(&mut switched, kind); + assert!(!switched + .iter() + .any(|message| matches!(message, AisdkMessage::Reasoning(_)))); + assert!(!serde_json::to_string(&switched) + .unwrap() + .contains("thoughtSignature")); + assert_eq!( + switched + .iter() + .filter(|message| matches!(message, AisdkMessage::ToolCall(_))) + .count(), + 2 + ); + } + let mut retained = messages.clone(); + super::retain_provider_replay(&mut retained, ProviderKind::Antigravity); + assert_eq!(retained.len(), messages.len()); + let mut other_reasoning = vec![AisdkMessage::reasoning( + Some("rs_openai".into()), + "summary", + Some("openai-encrypted".into()), + )]; + super::retain_provider_replay(&mut other_reasoning, ProviderKind::OpenAI); + assert_eq!(other_reasoning.len(), 1); + } + + #[test] + fn google_oauth_gate_requires_google_and_oauth() { + let oauth = AuthConfig::OAuth { + refresh: "refresh".into(), + access: "access".into(), + expires: 200_000, + account_id: None, + enterprise_url: None, + project_id: Some("project".into()), + }; + let api = AuthConfig::Api { + key: "api-key".into(), + }; + for provider in [ + "google", + "Google", + "google-vertex", + "openai", + "xai", + "other", + "", + ] { + for auth in [Some(&oauth), Some(&api), Some(&AuthConfig::Local), None] { + assert_eq!( + super::uses_google_oauth(provider, auth), + provider == "google" && matches!(auth, Some(AuthConfig::OAuth { .. })), + "provider={provider}, auth={auth:?}", + ); + } + } + } + + #[test] + fn google_oauth_credentials_route_transport_and_preserve_request_settings() { + let mut config = ProviderRequestConfig::new( + ProviderKind::OpenAI, + "Google".into(), + "https://generativelanguage.googleapis.com".into(), + "gemini-3-pro".into(), + Some("old-access".into()), + None, + true, + ); + config.supports_audio_input = true; + config.gateway_caching_auto = true; + config.prune_tool_outputs = true; + config.openai_options.prompt_cache_key = Some("session-cache".into()); + config.openai_options.force_store_false = true; + config + .openai_options + .additional_headers + .insert("X-Custom".into(), "preserved".into()); + config + .openai_options + .additional_headers + .insert("User-Agent".into(), "old-agent".into()); + let original = config.clone(); + for project_id in [Some("resolved-project".to_owned()), None] { + config.openai_options.antigravity_project = Some("stale-project".into()); + let credentials = crate::auth::OAuthCredentials { + refresh: "refresh-not-for-requests".into(), + access: "new-access".into(), + expires: 300_000, + account_id: Some("account".into()), + enterprise_url: None, + project_id: project_id.clone(), + }; + + super::apply_google_oauth_credentials(&mut config, &credentials); + + assert_eq!(config.kind, ProviderKind::Antigravity); + assert_eq!( + config.base_url, + crate::auth::antigravity_oauth::ANTIGRAVITY_ENDPOINT + ); + assert_eq!(config.api_key.as_deref(), Some("new-access")); + let mut expected_options = original.openai_options.clone(); + expected_options.antigravity_project = project_id; + expected_options.additional_headers.insert( + "User-Agent".into(), + crate::auth::antigravity_oauth::user_agent().into(), + ); + assert_eq!(config.openai_options, expected_options); + assert_eq!(config.provider_name, original.provider_name); + assert_eq!(config.model_name, original.model_name); + assert_eq!(config.reasoning_effort, original.reasoning_effort); + assert_eq!(config.supports_image_input, original.supports_image_input); + assert_eq!(config.supports_audio_input, original.supports_audio_input); + assert_eq!(config.gateway_caching_auto, original.gateway_caching_auto); + assert_eq!(config.prune_tool_outputs, original.prune_tool_outputs); + } + } + #[test] fn audio_model_receives_base64_audio_attachment() { let dir = tempfile::tempdir().expect("temp dir"); diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 244c9cfd..4d99c324 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -1,3 +1,4 @@ +pub(crate) mod antigravity; pub mod client; pub(crate) mod opencode; pub mod provider; @@ -13,6 +14,8 @@ use tokio::sync::mpsc; pub enum ChunkMessage { Text(String), Reasoning(String), + /// Opaque provider-owned history that must be persisted in stream order. + ProviderReplay(crate::aisdk::chunk::ReasoningReplayItem), Retry(crate::aisdk::retry::RetryStatus), StreamRollback { text: String, diff --git a/src/main.rs b/src/main.rs index a55d66d5..9f7f87c2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -559,6 +559,7 @@ async fn run_print_mode( | crate::llm::ChunkMessage::Metrics { .. } | crate::llm::ChunkMessage::Cancelled | crate::llm::ChunkMessage::Reasoning(_) + | crate::llm::ChunkMessage::ProviderReplay(_) | crate::llm::ChunkMessage::Usage(_) | crate::llm::ChunkMessage::Retry(_) | crate::llm::ChunkMessage::StreamRollback { .. } diff --git a/src/model/discovery.rs b/src/model/discovery.rs index 5dd25233..209bf386 100644 --- a/src/model/discovery.rs +++ b/src/model/discovery.rs @@ -684,8 +684,8 @@ impl Discovery { let is_discovered = advertised .get(provider_id) .is_some_and(|ids| ids.iter().any(|id| id == model_id)); - // Only this provider/model is materialized; pricing and limit lookups - // must not clone the entire catalog on each streamed response. + // Materialize only the requested model, except Google needs its source + // metadata to resolve the OAuth catalog's wire IDs. let mut providers = HashMap::new(); if let Some(provider) = provider { providers.insert( @@ -697,11 +697,15 @@ impl Discovery { doc: provider.doc.clone(), env: provider.env.clone(), npm: provider.npm.clone(), - models: provider - .models - .get(model_id) - .map(|model| HashMap::from([(model_id.to_string(), model.clone())])) - .unwrap_or_default(), + models: if provider_id == "google" { + provider.models.clone() + } else { + provider + .models + .get(model_id) + .map(|model| HashMap::from([(model_id.to_string(), model.clone())])) + .unwrap_or_default() + }, }, ); } @@ -730,6 +734,7 @@ impl Discovery { ids.insert(provider_id.to_string(), vec![model_id.to_string()]); } self.apply_custom_catalog(&mut providers, &ids); + crate::model::extensions::antigravity::augment_catalog(&mut providers); let mut provider = providers.remove(provider_id)?; let model = provider.models.remove(model_id)?; Some((provider, model)) @@ -875,10 +880,10 @@ impl Discovery { let cached_json = fs::read_to_string(cache_path).context("Failed to read cache file")?; - let entry = Arc::new( - serde_json::from_str::(&cached_json) - .context("Failed to parse cache file")?, - ); + let mut entry = serde_json::from_str::(&cached_json) + .context("Failed to parse cache file")?; + crate::model::extensions::antigravity::rename_provider(&mut entry.data); + let entry = Arc::new(entry); if let Ok(mut cache) = memory_cache().lock() { cache.insert(cache_path.clone(), entry.clone()); @@ -922,6 +927,8 @@ impl Discovery { schema_version: CACHE_SCHEMA_VERSION, }; + let mut entry = entry; + crate::model::extensions::antigravity::rename_provider(&mut entry.data); let serialized = serde_json::to_string_pretty(&entry).context("Failed to serialize cache data")?; @@ -971,7 +978,7 @@ impl Discovery { ("opencode", "OpenCode"), ("anthropic", "Anthropic"), ("openai", "OpenAI"), - ("google", "Google"), + ("google", "Google Antigravity"), ] { providers.insert( id.to_string(), @@ -998,6 +1005,7 @@ impl Discovery { crate::model::extensions::ModelExtensions::augment_runtime_catalog(&mut providers); let advertised = self.discover_custom_model_ids().await; self.apply_custom_catalog(&mut providers, &advertised); + crate::model::extensions::antigravity::augment_catalog(&mut providers); Ok(providers) } @@ -1010,6 +1018,7 @@ impl Discovery { crate::model::extensions::ModelExtensions::augment_runtime_catalog(&mut providers); let advertised = self.discover_custom_model_ids().await; self.apply_custom_catalog(&mut providers, &advertised); + crate::model::extensions::antigravity::augment_catalog(&mut providers); Ok(providers) } @@ -1148,6 +1157,10 @@ impl Discovery { models.retain(|model| self.provider_is_enabled(&model.provider_id)); let cache_key = (self.get_cache_path().clone(), { let mut signature = self.custom_provider_dialog_signature(); + signature.push(format!( + "google-oauth:{}", + crate::model::extensions::antigravity::uses_oauth() + )); signature.extend(self.custom_provider_endpoint_signature()); if let Some(providers) = &self.custom_providers { for (id, provider) in providers { @@ -1281,6 +1294,9 @@ impl Discovery { } pub fn get_provider_name(&self, provider_id: &str) -> Option { + if provider_id == "google" { + return Some(crate::model::extensions::antigravity::PROVIDER_NAME.to_string()); + } if let Some(name) = self .custom_providers .as_ref() @@ -2590,6 +2606,43 @@ mod tests { let _ = fs::remove_file(cache_path); } + #[test] + fn google_label_is_normalized_in_disk_and_synchronous_cache_views() { + let mut discovery = Discovery::new().unwrap(); + discovery.cache_path = unique_test_cache_path("google_label"); + fs::create_dir_all(discovery.cache_path.parent().unwrap()).unwrap(); + let provider: Provider = serde_json::from_value(serde_json::json!({ + "id": "google", "name": "Google", "models": {} + })) + .unwrap(); + discovery + .save_to_cache(&HashMap::from([("google".to_string(), provider)])) + .unwrap(); + let expected = crate::model::extensions::antigravity::PROVIDER_NAME; + assert_eq!( + discovery.get_provider_name("google").as_deref(), + Some(expected) + ); + assert_eq!( + discovery.load_from_cache().unwrap().unwrap()["google"].name, + expected + ); + let disk: CacheEntry = + serde_json::from_str(&fs::read_to_string(&discovery.cache_path).unwrap()).unwrap(); + assert_eq!(disk.data["google"].name, expected); + + // Old on-disk caches receive the same label without a network refresh. + let mut old = disk; + old.data.get_mut("google").unwrap().name = "Google".to_string(); + fs::write(&discovery.cache_path, serde_json::to_string(&old).unwrap()).unwrap(); + memory_cache().lock().unwrap().remove(&discovery.cache_path); + assert_eq!( + discovery.load_cache_entry().unwrap().unwrap().data["google"].name, + expected + ); + let _ = fs::remove_file(&discovery.cache_path); + } + #[test] fn parsed_cache_is_reused_from_memory() { let mut discovery = Discovery::new().unwrap(); diff --git a/src/model/effective_catalog.rs b/src/model/effective_catalog.rs index 1160993d..9e997ea7 100644 --- a/src/model/effective_catalog.rs +++ b/src/model/effective_catalog.rs @@ -148,7 +148,14 @@ fn write_snapshot(snapshot: &Snapshot) -> Result<()> { /// `None` means this installation has not published its first snapshot yet and /// callers should use their legacy compatibility path once. pub fn models_for_dialog() -> Result>> { - Ok(load_snapshot()?.map(|snapshot| snapshot.models.into_iter().map(Model::from).collect())) + Ok(load_snapshot()?.map(|snapshot| { + let mut models = snapshot.models.into_iter().map(Model::from).collect(); + crate::model::extensions::antigravity::apply_to_dialog_models( + &mut models, + crate::model::extensions::antigravity::uses_oauth(), + ); + models + })) } /// Publishes models produced by an explicit source refresh. diff --git a/src/model/extensions/antigravity.rs b/src/model/extensions/antigravity.rs new file mode 100644 index 00000000..51736fde --- /dev/null +++ b/src/model/extensions/antigravity.rs @@ -0,0 +1,189 @@ +use std::collections::HashMap; + +use crate::model::discovery::Provider; +use crate::persistence::{AuthConfig, AuthDAO}; + +pub const PROVIDER_NAME: &str = "Google Antigravity"; +const MODELS: &[(&str, &str, &str)] = &[ + ("gemini-3-flash", "Gemini 3 Flash", "gemini-3-flash-preview"), + ( + "gemini-3-pro-low", + "Gemini 3 Pro Low", + "gemini-3-pro-preview", + ), + ( + "gemini-3-pro-high", + "Gemini 3 Pro High", + "gemini-3-pro-preview", + ), + ( + "gemini-3.1-pro-low", + "Gemini 3.1 Pro Low", + "gemini-3.1-pro-preview", + ), + ( + "gemini-3.1-pro-high", + "Gemini 3.1 Pro High", + "gemini-3.1-pro-preview", + ), +]; + +pub fn apply_to_dialog_models(models: &mut Vec, oauth: bool) { + for model in models + .iter_mut() + .filter(|model| model.provider_id == "google") + { + model.provider_name = PROVIDER_NAME.into(); + } + if !oauth { + return; + } + let mapped = MODELS + .iter() + .filter_map(|&(id, name, reference)| { + let mut model = models + .iter() + .find(|model| model.provider_id == "google" && model.id == id) + .or_else(|| { + models + .iter() + .find(|model| model.provider_id == "google" && model.id == reference) + })? + .clone(); + model.id = id.into(); + model.name = name.into(); + Some(model) + }) + .collect::>(); + models.retain(|model| model.provider_id != "google"); + models.extend(mapped); +} + +pub fn uses_oauth() -> bool { + matches!( + AuthDAO::new().and_then(|dao| dao.get_provider("google")), + Ok(Some(AuthConfig::OAuth { .. })) + ) +} + +pub fn rename_provider(providers: &mut HashMap) { + if let Some(provider) = providers.get_mut("google") { + provider.name = PROVIDER_NAME.to_string(); + } +} + +pub fn augment_catalog(providers: &mut HashMap) { + if let Some(provider) = providers.get_mut("google") { + apply_to_provider(provider, uses_oauth()); + } +} + +pub fn apply_to_provider(provider: &mut Provider, oauth: bool) { + if provider.id != "google" { + return; + } + provider.name = PROVIDER_NAME.to_string(); + if !oauth { + return; + } + + // These are backend IDs, not the plugin's antigravity-* picker aliases. + // See opencode-antigravity-auth/src/plugin/transform/model-resolver.ts. + provider.models = MODELS + .iter() + .filter_map(|&(id, name, reference)| { + let mut model = provider + .models + .get(id) + .or_else(|| provider.models.get(reference))? + .clone(); + model.id = id.to_string(); + model.name = name.to_string(); + Some((id.to_string(), model)) + }) + .collect(); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn provider() -> Provider { + serde_json::from_value(serde_json::json!({ + "id": "google", "name": "Google", + "models": { + "gemini-3-flash-preview": { + "id": "gemini-3-flash-preview", "name": "Gemini 3 Flash Preview", + "limit": {"context": 1048576, "output": 65536}, + "cost": {"input": 0.5, "output": 3.0} + }, + "gemini-3-pro-preview": {"id": "gemini-3-pro-preview", "name": "Gemini 3 Pro Preview"}, + "gemini-2.0-flash": {"id": "gemini-2.0-flash", "name": "Gemini 2.0 Flash"} + } + })).unwrap() + } + + #[test] + fn api_key_catalog_only_changes_label() { + let mut provider = provider(); + let before = serde_json::to_value(&provider.models).unwrap(); + apply_to_provider(&mut provider, false); + assert_eq!(provider.name, PROVIDER_NAME); + assert_eq!(provider.id, "google"); + assert_eq!(serde_json::to_value(&provider.models).unwrap(), before); + } + + #[test] + fn warm_dialog_rows_use_antigravity_labels_and_oauth_model_ids() { + let row = crate::model::types::Model { + id: "gemini-3-flash-preview".into(), + name: "Gemini 3 Flash Preview".into(), + family: "gemini".into(), + provider_id: "google".into(), + provider_name: "Google".into(), + attachment: true, + structured_output: true, + free: false, + local: false, + reasoning_options: vec![], + context_window: Some(1048576), + }; + let mut rows = vec![row]; + apply_to_dialog_models(&mut rows, false); + assert_eq!(rows[0].provider_name, PROVIDER_NAME); + assert_eq!(rows[0].id, "gemini-3-flash-preview"); + apply_to_dialog_models(&mut rows, true); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].id, "gemini-3-flash"); + assert_eq!(rows[0].context_window, Some(1048576)); + apply_to_dialog_models(&mut rows, true); + assert_eq!(rows.len(), 1); + } + + #[test] + fn oauth_catalog_uses_wire_ids_and_existing_metadata() { + let mut provider = provider(); + apply_to_provider(&mut provider, true); + assert_eq!(provider.name, PROVIDER_NAME); + assert_eq!(provider.models.len(), 3); + let flash = &provider.models["gemini-3-flash"]; + assert_eq!(flash.id, "gemini-3-flash"); + assert_eq!(flash.name, "Gemini 3 Flash"); + assert_eq!(flash.limit.as_ref().unwrap().context, 1048576); + assert_eq!(flash.cost.as_ref().unwrap().input, 0.5); + assert!(provider.models.contains_key("gemini-3-pro-high")); + assert!(!provider.models.contains_key("gemini-3.1-pro-high")); + assert!(!provider.models.contains_key("gemini-2.0-flash")); + let before = serde_json::to_value(&provider).unwrap(); + apply_to_provider(&mut provider, true); + assert_eq!(serde_json::to_value(provider).unwrap(), before); + } + + #[test] + fn missing_metadata_is_not_invented() { + let mut provider = provider(); + provider.models.clear(); + apply_to_provider(&mut provider, true); + assert!(provider.models.is_empty()); + } +} diff --git a/src/model/extensions/mod.rs b/src/model/extensions/mod.rs index 43313ca5..754a7f9e 100644 --- a/src/model/extensions/mod.rs +++ b/src/model/extensions/mod.rs @@ -4,6 +4,7 @@ use std::collections::HashMap; use crate::model::discovery::{Model, Provider}; +pub mod antigravity; pub mod commandcode; pub mod ollama; diff --git a/src/persistence/auth.rs b/src/persistence/auth.rs index cbba5fc3..d867dfba 100644 --- a/src/persistence/auth.rs +++ b/src/persistence/auth.rs @@ -26,6 +26,8 @@ pub enum AuthConfig { skip_serializing_if = "Option::is_none" )] enterprise_url: Option, + #[serde(rename = "projectId", default, skip_serializing_if = "Option::is_none")] + project_id: Option, }, } @@ -177,6 +179,26 @@ impl AuthDAO { mod tests { use super::*; + #[test] + fn oauth_project_id_round_trip_and_legacy_default() { + let legacy = r#"{"type":"oauth","refresh":"r","access":"a","expires":123}"#; + let auth: AuthConfig = serde_json::from_str(legacy).unwrap(); + assert!(matches!( + &auth, + AuthConfig::OAuth { + project_id: None, + .. + } + )); + assert!(serde_json::to_value(auth) + .unwrap() + .get("projectId") + .is_none()); + let value = serde_json::json!({"type":"oauth","refresh":"r","access":"a","expires":123,"projectId":"managed"}); + let auth: AuthConfig = serde_json::from_value(value.clone()).unwrap(); + assert_eq!(serde_json::to_value(auth).unwrap(), value); + } + #[test] fn parse_auth_configs_ignores_invalid_provider_entries() { let providers = parse_auth_configs( diff --git a/src/ui/components/chat.rs b/src/ui/components/chat.rs index 4041b77a..ff41b31f 100644 --- a/src/ui/components/chat.rs +++ b/src/ui/components/chat.rs @@ -2106,6 +2106,12 @@ impl Chat { self.add_message(Message::assistant(content)); } + pub fn get_last_assistant_message_mut(&mut self) -> Option<&mut Message> { + self.messages + .iter_mut() + .rfind(|message| message.role == MessageRole::Assistant) + } + fn streaming_assistant_idx(&self) -> Option { self.messages .iter() @@ -9074,11 +9080,12 @@ mod tests { #[test] fn test_edit_tool_renders_codex_style_diff_summary() { let chat = Chat::new(); + let file_path = std::env::current_dir().unwrap().join("README.md"); let content = serde_json::json!({ "name": "edit", "status": "ok", "args": { - "file_path": "/Users/carlo/Desktop/Projects/crabcode/README.md", + "file_path": file_path, "old_string": "alpha\nbeta\nomega", "new_string": "alpha\nbravo\nomega", }, diff --git a/src/views/provider_oauth_flow.rs b/src/views/provider_oauth_flow.rs index 7301403d..0191acab 100644 --- a/src/views/provider_oauth_flow.rs +++ b/src/views/provider_oauth_flow.rs @@ -1,19 +1,23 @@ use crate::theme::ThemeColors; use ratatui::crossterm::event::{ - KeyCode, KeyEvent, KeyModifiers, MouseButton, MouseEvent, MouseEventKind, + KeyCode, KeyEvent, KeyEventKind, KeyModifiers, MouseButton, MouseEvent, MouseEventKind, }; use ratatui::{ layout::{Alignment, Constraint, Direction, Layout, Position, Rect}, style::{Modifier, Style}, text::{Line, Span}, - widgets::{Clear, Paragraph}, + widgets::{Clear, Paragraph, Wrap}, Frame, }; +use tui_textarea::TextArea; use unicode_width::UnicodeWidthStr; +const MAX_CALLBACK_BYTES: usize = 16 * 1024; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ProviderOAuthFlowMode { BrowserWaiting, + BrowserCallback, HeadlessPreparing, HeadlessCode, } @@ -25,6 +29,11 @@ pub struct ProviderOAuthFlowState { pub provider_label: String, pub url: Option, pub code: Option, + callback_textarea: TextArea<'static>, + pub callback_error: Option, + pub submitting: bool, + open_link_area: Option, + submit_area: Option, dialog_area: Rect, link_area: Option, } @@ -37,6 +46,11 @@ impl ProviderOAuthFlowState { provider_label: "OpenAI".to_string(), url: None, code: None, + callback_textarea: TextArea::default(), + callback_error: None, + submitting: false, + open_link_area: None, + submit_area: None, dialog_area: Rect::default(), link_area: None, } @@ -47,6 +61,7 @@ impl ProviderOAuthFlowState { } pub fn show_browser_waiting_for(&mut self, provider_label: impl Into) { + self.hide(); self.visible = true; self.mode = ProviderOAuthFlowMode::BrowserWaiting; self.provider_label = provider_label.into(); @@ -54,11 +69,74 @@ impl ProviderOAuthFlowState { self.code = None; } + pub fn show_browser_callback_for( + &mut self, + provider_label: impl Into, + url: impl Into, + ) { + self.hide(); + self.visible = true; + self.mode = ProviderOAuthFlowMode::BrowserCallback; + self.provider_label = provider_label.into(); + self.url = Some(url.into()); + } + + /// Insert a whole paste without submitting or retaining rejected content. + pub fn insert_callback(&mut self, text: &str) { + if !self.visible || self.mode != ProviderOAuthFlowMode::BrowserCallback || self.submitting { + return; + } + if text.chars().any(char::is_control) { + self.callback_error = + Some("Paste one full callback URL without line breaks or tabs.".into()); + return; + } + if self.callback_textarea.lines()[0] + .len() + .saturating_add(text.len()) + > MAX_CALLBACK_BYTES + { + self.callback_error = Some( + "Callback URL is too long (maximum 16 KiB). Paste only the localhost URL.".into(), + ); + return; + } + self.callback_textarea.insert_str(text); + self.callback_error = None; + } + + pub fn set_callback_error(&mut self, error: String) { + self.submitting = false; + self.callback_error = Some(error); + } + + pub fn mark_submitting(&mut self) { + self.submitting = true; + self.callback_error = None; + // Replacing the editor also drops undo history and its internal yank buffer. + self.callback_textarea = TextArea::default(); + } + + fn submit_callback(&mut self) -> ProviderOAuthFlowAction { + if self.submitting { + return ProviderOAuthFlowAction::Handled; + } + let callback = self.callback_textarea.lines()[0].trim().to_string(); + if callback.is_empty() { + self.callback_error = + Some("Paste the full localhost callback URL before submitting.".into()); + return ProviderOAuthFlowAction::Handled; + } + self.mark_submitting(); + ProviderOAuthFlowAction::SubmitCallback(callback) + } + pub fn show_headless_preparing(&mut self) { self.show_headless_preparing_for("OpenAI"); } pub fn show_headless_preparing_for(&mut self, provider_label: impl Into) { + self.hide(); self.visible = true; self.mode = ProviderOAuthFlowMode::HeadlessPreparing; self.provider_label = provider_label.into(); @@ -78,6 +156,12 @@ impl ProviderOAuthFlowState { self.url = None; self.code = None; self.link_area = None; + self.open_link_area = None; + self.submit_area = None; + self.dialog_area = Rect::default(); + self.callback_textarea = TextArea::default(); + self.callback_error = None; + self.submitting = false; } pub fn is_visible(&self) -> bool { @@ -97,6 +181,9 @@ pub enum ProviderOAuthFlowAction { NotHandled, Close, CopyLink(String), + OpenLink(String), + SubmitCallback(String), + Cancel, } pub fn init_provider_oauth_flow() -> ProviderOAuthFlowState { @@ -111,9 +198,18 @@ pub fn handle_provider_oauth_flow_key_event( return ProviderOAuthFlowAction::NotHandled; } + if event.kind == KeyEventKind::Release { + return ProviderOAuthFlowAction::Handled; + } + if event.code == KeyCode::Esc { + let callback = state.mode == ProviderOAuthFlowMode::BrowserCallback; state.hide(); - return ProviderOAuthFlowAction::Close; + return if callback { + ProviderOAuthFlowAction::Cancel + } else { + ProviderOAuthFlowAction::Close + }; } if event.code == KeyCode::Char('y') && event.modifiers == KeyModifiers::CONTROL { @@ -122,6 +218,47 @@ pub fn handle_provider_oauth_flow_key_event( } } + if state.mode == ProviderOAuthFlowMode::BrowserCallback { + if event.code == KeyCode::Char('o') && event.modifiers == KeyModifiers::CONTROL { + if let Some(url) = &state.url { + return ProviderOAuthFlowAction::OpenLink(url.clone()); + } + } + if state.submitting { + return ProviderOAuthFlowAction::Handled; + } + match event.code { + KeyCode::Enter => return state.submit_callback(), + KeyCode::Char(c) + if !event + .modifiers + .intersects(KeyModifiers::CONTROL | KeyModifiers::ALT) => + { + state.insert_callback(&c.to_string()); + } + KeyCode::Backspace + | KeyCode::Delete + | KeyCode::Left + | KeyCode::Right + | KeyCode::Home + | KeyCode::End => { + state.callback_textarea.input(event); + state.callback_error = None; + } + KeyCode::Char(c) + if event.modifiers == KeyModifiers::CONTROL + && matches!( + c, + 'a' | 'e' | 'b' | 'f' | 'h' | 'd' | 'k' | 'u' | 'w' | 'z' | 'r' + ) => + { + state.callback_textarea.input(event); + state.callback_error = None; + } + _ => {} + } + } + ProviderOAuthFlowAction::Handled } @@ -140,10 +277,28 @@ pub fn handle_provider_oauth_flow_mouse_event( let point = Position::new(event.column, event.row); if !state.dialog_area.contains(point) { + // Callback auth stays open until an explicit Esc cancellation. + if state.mode == ProviderOAuthFlowMode::BrowserCallback { + return ProviderOAuthFlowAction::Handled; + } state.hide(); return ProviderOAuthFlowAction::Close; } + if state.mode == ProviderOAuthFlowMode::BrowserCallback { + if state + .open_link_area + .is_some_and(|area| area.contains(point)) + { + if let Some(url) = &state.url { + return ProviderOAuthFlowAction::OpenLink(url.clone()); + } + } + if state.submit_area.is_some_and(|area| area.contains(point)) { + return state.submit_callback(); + } + } + if let (Some(link_area), Some(url)) = (state.link_area, &state.url) { if link_area.contains(point) { return ProviderOAuthFlowAction::CopyLink(url.clone()); @@ -163,6 +318,11 @@ pub fn render_provider_oauth_flow( return; } + if state.mode == ProviderOAuthFlowMode::BrowserCallback { + render_browser_callback(frame, state, area, colors); + return; + } + const DIALOG_WIDTH: u16 = 82; const DIALOG_HEIGHT: u16 = 16; const PADDING: u16 = 3; @@ -171,8 +331,8 @@ pub fn render_provider_oauth_flow( let dialog_height = area.height.min(DIALOG_HEIGHT); state.dialog_area = Rect { - x: (area.width - dialog_width) / 2, - y: (area.height - dialog_height) / 2, + x: area.x.saturating_add((area.width - dialog_width) / 2), + y: area.y.saturating_add((area.height - dialog_height) / 2), width: dialog_width, height: dialog_height, }; @@ -184,8 +344,14 @@ pub fn render_provider_oauth_flow( ); let content_area = Rect { - x: state.dialog_area.x + PADDING, - y: state.dialog_area.y + PADDING, + x: state + .dialog_area + .x + .saturating_add(PADDING.min(state.dialog_area.width)), + y: state + .dialog_area + .y + .saturating_add(PADDING.min(state.dialog_area.height)), width: state.dialog_area.width.saturating_sub(PADDING * 2), height: state.dialog_area.height.saturating_sub(PADDING * 2), }; @@ -237,6 +403,7 @@ pub fn render_provider_oauth_flow( ); match state.mode { + ProviderOAuthFlowMode::BrowserCallback => unreachable!("rendered separately"), ProviderOAuthFlowMode::BrowserWaiting => { frame.render_widget( Paragraph::new(Line::from(vec![Span::raw( @@ -347,3 +514,358 @@ pub fn render_provider_oauth_flow( frame.render_widget(Paragraph::new(footer).alignment(Alignment::Left), chunks[8]); } + +fn render_browser_callback( + frame: &mut Frame, + state: &mut ProviderOAuthFlowState, + area: Rect, + colors: ThemeColors, +) { + let width = area.width.min(82); + let height = area.height.min(20); + state.dialog_area = Rect::new( + area.x.saturating_add((area.width - width) / 2), + area.y.saturating_add((area.height - height) / 2), + width, + height, + ); + state.link_area = None; + state.open_link_area = None; + state.submit_area = None; + frame.render_widget(Clear, state.dialog_area); + frame.render_widget( + Paragraph::new("").style(Style::default().bg(colors.dialog_background)), + state.dialog_area, + ); + let padding = if width >= 40 { 2 } else { 0 }; + let top: u16 = if height >= 20 { 1 } else { 0 }; + let content = Rect::new( + state.dialog_area.x.saturating_add(padding), + state.dialog_area.y.saturating_add(top), + width.saturating_sub(padding * 2), + height.saturating_sub(top * 2), + ); + if content.is_empty() { + return; + } + let row = |offset: u16, rows: u16| { + Rect::new( + content.x, + content.y.saturating_add(offset.min(content.height)), + content.width, + rows.min(content.height.saturating_sub(offset)), + ) + }; + let text = Style::default().fg(colors.text); + let weak = Style::default().fg(colors.text_weak); + let link = Style::default().fg(colors.primary); + frame.render_widget( + Paragraph::new(format!("{} OAuth", state.provider_label)) + .style(text.add_modifier(Modifier::BOLD)), + row(0, 1), + ); + frame.render_widget( + Paragraph::new("Esc cancel") + .style(link) + .alignment(Alignment::Right), + row(1, 1), + ); + frame.render_widget( + Paragraph::new("1. Authorize in your browser").style(text), + row(2, 1), + ); + let url_area = row(3, 1); + if !url_area.is_empty() { + state.link_area = Some(url_area); + } + frame.render_widget( + Paragraph::new(state.url.as_deref().unwrap_or_default()) + .style(link.add_modifier(Modifier::UNDERLINED)), + url_area, + ); + let open_area = row(4, 1); + if !open_area.is_empty() { + state.open_link_area = Some(open_area); + } + frame.render_widget( + Paragraph::new("[Open link] Ctrl+O Ctrl+Y copy").style(link), + open_area, + ); + frame.render_widget(Paragraph::new("2. Callback URL").style(text), row(6, 1)); + state + .callback_textarea + .set_style(text.bg(colors.background_element)); + state + .callback_textarea + .set_cursor_line_style(Style::default()); + state + .callback_textarea + .set_cursor_style(if state.submitting { + Style::default() + } else { + link.add_modifier(Modifier::REVERSED) + }); + state + .callback_textarea + .set_placeholder_text("http://localhost:.../?code=..."); + frame.render_widget(&state.callback_textarea, row(7, 1)); + let submit_area = row(8, 1); + if !submit_area.is_empty() { + state.submit_area = Some(submit_area); + } + frame.render_widget( + Paragraph::new(if state.submitting { + "Submitting..." + } else { + "[Submit] Enter" + }) + .style(if state.submitting { weak } else { link }), + submit_area, + ); + frame.render_widget( + Paragraph::new( + "Copy the full localhost URL from your browser, even if the page cannot load.", + ) + .style(weak) + .wrap(Wrap { trim: false }), + row(9, 3), + ); + if let Some(error) = &state.callback_error { + frame.render_widget( + Paragraph::new(error.as_str()) + .style(Style::default().fg(colors.error)) + .wrap(Wrap { trim: false }), + row(12, 2), + ); + } + frame.render_widget( + Paragraph::new( + "Unofficial Antigravity access may be subject to Google account restrictions.", + ) + .style(Style::default().fg(colors.warning)) + .wrap(Wrap { trim: false }), + row(14, 3), + ); +} + +#[cfg(test)] +mod tests { + use super::*; + use ratatui::{backend::TestBackend, Terminal}; + + const URL: &str = "https://example.com/authorize?state=full-state&challenge=full-challenge"; + const CALLBACK: &str = "http://localhost:1234/?code=secret&state=test"; + + fn callback_state() -> ProviderOAuthFlowState { + let mut state = ProviderOAuthFlowState::new(); + state.show_browser_callback_for("Antigravity", URL); + state + } + + fn key( + state: &mut ProviderOAuthFlowState, + code: KeyCode, + modifiers: KeyModifiers, + ) -> ProviderOAuthFlowAction { + handle_provider_oauth_flow_key_event(state, KeyEvent::new(code, modifiers)) + } + + fn render(state: &mut ProviderOAuthFlowState, width: u16, height: u16, area: Rect) -> String { + let backend = TestBackend::new(width, height); + let mut terminal = Terminal::new(backend).unwrap(); + terminal + .draw(|frame| { + render_provider_oauth_flow( + frame, + state, + area, + crate::theme::Theme::load_builtin_default().get_colors(true), + ) + }) + .unwrap(); + terminal + .backend() + .buffer() + .content + .iter() + .map(|cell| cell.symbol()) + .collect() + } + + #[test] + fn paste_is_bounded_single_line_and_never_submits() { + let mut state = callback_state(); + for invalid in ["first\nsecond", "tab\tvalue", "escape\u{1b}", "\r", "\0"] { + state.insert_callback(invalid); + assert!(state.callback_error.is_some()); + assert_eq!(state.callback_textarea.lines(), &[String::new()]); + } + state.insert_callback(&"x".repeat(MAX_CALLBACK_BYTES + 1)); + assert!(state.callback_error.is_some()); + state.insert_callback(CALLBACK); + assert_eq!(state.callback_textarea.lines(), &[CALLBACK.to_string()]); + assert!(!state.submitting); + assert!(state.callback_error.is_none()); + } + + #[test] + fn keyboard_submit_clears_sensitive_history_and_blocks_duplicates() { + let mut state = callback_state(); + assert_eq!( + key(&mut state, KeyCode::Enter, KeyModifiers::NONE), + ProviderOAuthFlowAction::Handled + ); + assert!(state.callback_error.is_some()); + state.insert_callback(CALLBACK); + key(&mut state, KeyCode::End, KeyModifiers::NONE); + key(&mut state, KeyCode::Char('x'), KeyModifiers::NONE); + key(&mut state, KeyCode::Backspace, KeyModifiers::NONE); + assert_eq!( + key(&mut state, KeyCode::Enter, KeyModifiers::NONE), + ProviderOAuthFlowAction::SubmitCallback(CALLBACK.into()) + ); + assert!(state.submitting); + key(&mut state, KeyCode::Char('z'), KeyModifiers::CONTROL); + key(&mut state, KeyCode::Char('x'), KeyModifiers::NONE); + state.insert_callback(CALLBACK); + assert_eq!(state.callback_textarea.lines(), &[String::new()]); + assert_eq!( + key(&mut state, KeyCode::Enter, KeyModifiers::NONE), + ProviderOAuthFlowAction::Handled + ); + state.set_callback_error("Authorization failed; retry login.".into()); + assert!(!state.submitting); + key(&mut state, KeyCode::Char('z'), KeyModifiers::CONTROL); + assert_eq!(state.callback_textarea.lines(), &[String::new()]); + } + + #[test] + fn release_events_do_not_submit_edit_or_cancel() { + let mut state = callback_state(); + state.insert_callback(CALLBACK); + for code in [KeyCode::Enter, KeyCode::Esc, KeyCode::Char('x')] { + let mut event = KeyEvent::new(code, KeyModifiers::NONE); + event.kind = KeyEventKind::Release; + assert_eq!( + handle_provider_oauth_flow_key_event(&mut state, event), + ProviderOAuthFlowAction::Handled + ); + } + assert!(state.visible); + assert!(!state.submitting); + assert_eq!(state.callback_textarea.lines()[0], CALLBACK); + } + + #[test] + fn copy_open_cancel_and_hide_reset() { + let mut state = callback_state(); + assert_eq!( + key(&mut state, KeyCode::Char('y'), KeyModifiers::CONTROL), + ProviderOAuthFlowAction::CopyLink(URL.into()) + ); + assert_eq!( + key(&mut state, KeyCode::Char('o'), KeyModifiers::CONTROL), + ProviderOAuthFlowAction::OpenLink(URL.into()) + ); + state.insert_callback(CALLBACK); + state.callback_error = Some("retry".into()); + assert_eq!( + key(&mut state, KeyCode::Esc, KeyModifiers::NONE), + ProviderOAuthFlowAction::Cancel + ); + assert!(state.url.is_none()); + assert!(state.callback_error.is_none()); + assert!(!state.submitting); + assert!( + state.link_area.is_none() + && state.open_link_area.is_none() + && state.submit_area.is_none() + ); + assert_eq!(state.dialog_area, Rect::default()); + state.show_browser_callback_for("Antigravity", URL); + key(&mut state, KeyCode::Char('z'), KeyModifiers::CONTROL); + assert_eq!(state.callback_textarea.lines(), &[String::new()]); + } + + #[test] + fn headless_keeps_copy_and_close_behavior() { + let mut state = ProviderOAuthFlowState::new(); + state.show_headless_preparing(); + assert_eq!(state.mode, ProviderOAuthFlowMode::HeadlessPreparing); + state.set_headless_code("device-code".into(), URL.into()); + assert_eq!( + key(&mut state, KeyCode::Char('y'), KeyModifiers::CONTROL), + ProviderOAuthFlowAction::CopyLink(URL.into()) + ); + assert_eq!( + key(&mut state, KeyCode::Esc, KeyModifiers::NONE), + ProviderOAuthFlowAction::Close + ); + assert!(state.code.is_none()); + } + + #[test] + fn responsive_render_and_clicks_preserve_full_url() { + for (width, height) in [(100, 30), (40, 18), (20, 16), (1, 1)] { + let mut state = callback_state(); + let output = render(&mut state, width, height, Rect::new(0, 0, width, height)); + if width >= 40 { + assert!(output.contains("Esc cancel")); + assert!(output.contains("2. Callback URL")); + assert!(output.contains("[Submit] Enter")); + } + if let Some(area) = state.link_area { + let event = MouseEvent { + kind: MouseEventKind::Down(MouseButton::Left), + column: area.x, + row: area.y, + modifiers: KeyModifiers::NONE, + }; + assert_eq!( + handle_provider_oauth_flow_mouse_event(&mut state, event), + ProviderOAuthFlowAction::CopyLink(URL.into()) + ); + } + if let Some(area) = state.open_link_area { + let event = MouseEvent { + kind: MouseEventKind::Down(MouseButton::Left), + column: area.x, + row: area.y, + modifiers: KeyModifiers::NONE, + }; + assert_eq!( + handle_provider_oauth_flow_mouse_event(&mut state, event), + ProviderOAuthFlowAction::OpenLink(URL.into()) + ); + } + if let Some(area) = state.submit_area { + state.insert_callback(CALLBACK); + let event = MouseEvent { + kind: MouseEventKind::Down(MouseButton::Left), + column: area.x, + row: area.y, + modifiers: KeyModifiers::NONE, + }; + assert_eq!( + handle_provider_oauth_flow_mouse_event(&mut state, event), + ProviderOAuthFlowAction::SubmitCallback(CALLBACK.into()) + ); + assert_eq!( + handle_provider_oauth_flow_mouse_event(&mut state, event), + ProviderOAuthFlowAction::Handled + ); + } + } + let mut state = callback_state(); + render(&mut state, 100, 30, Rect::new(10, 5, 40, 18)); + assert!(state.dialog_area.x >= 10 && state.dialog_area.y >= 5); + for mode in [ + ProviderOAuthFlowMode::BrowserWaiting, + ProviderOAuthFlowMode::HeadlessPreparing, + ProviderOAuthFlowMode::HeadlessCode, + ] { + state.mode = mode; + render(&mut state, 1, 1, Rect::new(0, 0, 1, 1)); + } + } +} diff --git a/tests/acp_stdio.rs b/tests/acp_stdio.rs index b16988d0..b2341aac 100644 --- a/tests/acp_stdio.rs +++ b/tests/acp_stdio.rs @@ -30,6 +30,12 @@ impl AcpProcess { .unwrap(), ) .expect("offline model fixture"); + // A connected non-runtime fixture prevents implicit Ollama discovery. + std::fs::write( + state.path().join("crabcode/auth.json"), + r#"{"opencode":{"type":"local"}}"#, + ) + .expect("offline auth fixture"); let mut child = Command::new(env!("CARGO_BIN_EXE_crabcode")) .args(["acp", "--cwd"]) .arg(workspace) @@ -357,6 +363,10 @@ Wire skill instructions. } })); let (new_session, mut notifications) = process.recv_response(2); + assert!( + new_session.get("error").is_none(), + "session setup: {new_session}" + ); let session_id = new_session["result"]["sessionId"] .as_str() .expect("session id")