diff --git a/crates/goose-cli/src/commands/configure.rs b/crates/goose-cli/src/commands/configure.rs index cf809dbacc..dd155c71f2 100644 --- a/crates/goose-cli/src/commands/configure.rs +++ b/crates/goose-cli/src/commands/configure.rs @@ -72,6 +72,11 @@ pub async fn handle_configure() -> Result<(), Box> { "OpenRouter Login (Recommended)", "Sign in with OpenRouter to automatically configure models", ) + .item( + "tetrate", + "Tetrate Agent Router Service Login", + "Sign in with Tetrate Agent Router Service to automatically configure models", + ) .item( "manual", "Manual Configuration", @@ -95,6 +100,21 @@ pub async fn handle_configure() -> Result<(), Box> { } } } + "tetrate" => { + match handle_tetrate_auth().await { + Ok(_) => { + // Tetrate auth already handles everything including enabling developer extension + } + Err(e) => { + let _ = config.clear(); + println!( + "\n {} Tetrate Agent Router Service authentication failed: {} \n Please try again or use manual configuration", + style("Error").red().italic(), + e, + ); + } + } + } "manual" => { match configure_provider_dialog().await { Ok(true) => { @@ -1770,6 +1790,109 @@ pub async fn handle_openrouter_auth() -> Result<(), Box> { Ok(()) } +/// Handle Tetrate Agent Router Service authentication +pub async fn handle_tetrate_auth() -> Result<(), Box> { + use goose::config::{configure_tetrate, signup_tetrate::TetrateAuth}; + use goose::conversation::message::Message; + use goose::providers::create; + + // Use the Tetrate Agent Router Service authentication flow + let mut auth_flow = TetrateAuth::new()?; + match auth_flow.complete_flow().await { + Ok(api_key) => { + println!("\nAuthentication complete!"); + + let config = Config::global(); + + // Use the existing configure_tetrate function to set everything up + println!("\nConfiguring Tetrate Agent Router Service..."); + if let Err(e) = configure_tetrate(config, api_key) { + eprintln!("Failed to configure Tetrate Agent Router Service: {}", e); + return Err(e.into()); + } + + println!("✓ Tetrate Agent Router Service configuration complete"); + println!("✓ Models configured successfully"); + + // Test configuration - get the model that was configured + println!("\nTesting configuration..."); + let configured_model: String = config.get_param("GOOSE_MODEL")?; + let model_config = match goose::model::ModelConfig::new(&configured_model) { + Ok(config) => config, + Err(e) => { + eprintln!("⚠️ Invalid model configuration: {}", e); + eprintln!( + "Your settings have been saved. Please check your model configuration." + ); + return Ok(()); + } + }; + + match create("tetrate", model_config) { + Ok(provider) => { + // Simple test request + let test_result = provider + .complete( + "You are Goose, an AI assistant.", + &[Message::user().with_text("Say 'Configuration test successful!'")], + &[], + ) + .await; + + match test_result { + Ok(_) => { + println!("✓ Configuration test passed!"); + + // Enable the developer extension by default if not already enabled + let entries = ExtensionConfigManager::get_all()?; + let has_developer = entries + .iter() + .any(|e| e.config.name() == "developer" && e.enabled); + + if !has_developer { + match ExtensionConfigManager::set(ExtensionEntry { + enabled: true, + config: ExtensionConfig::Builtin { + name: "developer".to_string(), + display_name: Some( + goose::config::DEFAULT_DISPLAY_NAME.to_string(), + ), + timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT), + bundled: Some(true), + description: None, + available_tools: Vec::new(), + }, + }) { + Ok(_) => println!("✓ Developer extension enabled"), + Err(e) => { + eprintln!("⚠️ Failed to enable developer extension: {}", e) + } + } + } + + cliclack::outro("Tetrate Agent Router Service setup complete! You can now use Goose.")?; + } + Err(e) => { + eprintln!("⚠️ Configuration test failed: {}", e); + eprintln!("Your settings have been saved, but there may be an issue with the connection."); + } + } + } + Err(e) => { + eprintln!("⚠️ Failed to create provider for testing: {}", e); + eprintln!("Your settings have been saved. Please check your configuration."); + } + } + } + Err(e) => { + eprintln!("Authentication failed: {}", e); + return Err(e.into()); + } + } + + Ok(()) +} + fn add_provider() -> Result<(), Box> { let provider_type = cliclack::select("What type of API is this?") .item( diff --git a/crates/goose-server/src/routes/setup.rs b/crates/goose-server/src/routes/setup.rs index 017315e57d..0c716d500a 100644 --- a/crates/goose-server/src/routes/setup.rs +++ b/crates/goose-server/src/routes/setup.rs @@ -1,6 +1,7 @@ use crate::state::AppState; use axum::{extract::State, http::StatusCode, routing::post, Json, Router}; use goose::config::signup_openrouter::OpenRouterAuth; +use goose::config::signup_tetrate::{configure_tetrate, TetrateAuth}; use goose::config::{configure_openrouter, Config}; use serde::Serialize; use std::sync::Arc; @@ -14,6 +15,7 @@ pub struct SetupResponse { pub fn routes(state: Arc) -> Router { Router::new() .route("/handle_openrouter", post(start_openrouter_setup)) + .route("/handle_tetrate", post(start_tetrate_setup)) .with_state(state) } @@ -58,3 +60,45 @@ async fn start_openrouter_setup( } } } + +async fn start_tetrate_setup( + State(_state): State>, +) -> Result, StatusCode> { + tracing::info!("Starting Tetrate Agent Router Service setup flow"); + + let mut auth_flow = TetrateAuth::new().map_err(|e| { + tracing::error!("Failed to initialize auth flow: {}", e); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + tracing::info!("Auth flow initialized, starting complete_flow"); + + match auth_flow.complete_flow().await { + Ok(api_key) => { + tracing::info!("Got API key, configuring Tetrate Agent Router Service..."); + + let config = Config::global(); + + if let Err(e) = configure_tetrate(config, api_key) { + tracing::error!("Failed to configure Tetrate Agent Router Service: {}", e); + return Ok(Json(SetupResponse { + success: false, + message: format!("Failed to configure Tetrate Agent Router Service: {}", e), + })); + } + + tracing::info!("Tetrate Agent Router Service setup completed successfully"); + Ok(Json(SetupResponse { + success: true, + message: "Tetrate Agent Router Service setup completed successfully".to_string(), + })) + } + Err(e) => { + tracing::error!("Tetrate Agent Router Service setup failed: {}", e); + Ok(Json(SetupResponse { + success: false, + message: format!("Setup failed: {}", e), + })) + } + } +} diff --git a/crates/goose/examples/tetrate_auth.rs b/crates/goose/examples/tetrate_auth.rs new file mode 100644 index 0000000000..780ee269b8 --- /dev/null +++ b/crates/goose/examples/tetrate_auth.rs @@ -0,0 +1,39 @@ +// Example of Tetrate Agent Router Service PKCE authentication +// Run with: cargo run --example tetrate_auth + +use goose::config::signup_tetrate::TetrateAuth; + +#[tokio::main] +async fn main() -> Result<(), Box> { + println!("Testing Tetrate Agent Router Service PKCE flow...\n"); + + // Create new PKCE auth flow + let mut auth_flow = TetrateAuth::new()?; + + // Get the auth URL that would be opened + let auth_url = auth_flow.get_auth_url(); + println!("Auth URL: {}", auth_url); + println!("\nStarting authentication flow..."); + println!("This will:"); + println!("1. Open your browser to the auth page"); + println!("2. Start a local server on port 3000"); + println!("3. Wait for the callback\n"); + + // Complete the full flow + match auth_flow.complete_flow().await { + Ok(api_key) => { + println!("\n✅ Authentication successful!"); + println!( + "API Key received: {}...", + &api_key.chars().take(10).collect::() + ); + println!("\nYou can now use this API key with the Tetrate provider."); + } + Err(e) => { + eprintln!("\n❌ Authentication failed: {}", e); + eprintln!("Error details: {:?}", e); + } + } + + Ok(()) +} diff --git a/crates/goose/src/config/mod.rs b/crates/goose/src/config/mod.rs index dda2a92d66..f060299a6d 100644 --- a/crates/goose/src/config/mod.rs +++ b/crates/goose/src/config/mod.rs @@ -4,6 +4,7 @@ mod experiments; pub mod extensions; pub mod permission; pub mod signup_openrouter; +pub mod signup_tetrate; pub use crate::agents::ExtensionConfig; pub use base::{Config, ConfigError, APP_STRATEGY}; @@ -12,6 +13,7 @@ pub use experiments::ExperimentManager; pub use extensions::{ExtensionConfigManager, ExtensionEntry}; pub use permission::PermissionManager; pub use signup_openrouter::configure_openrouter; +pub use signup_tetrate::configure_tetrate; pub use extensions::DEFAULT_DISPLAY_NAME; pub use extensions::DEFAULT_EXTENSION; diff --git a/crates/goose/src/config/signup_tetrate/mod.rs b/crates/goose/src/config/signup_tetrate/mod.rs new file mode 100644 index 0000000000..78ca99bc3f --- /dev/null +++ b/crates/goose/src/config/signup_tetrate/mod.rs @@ -0,0 +1,174 @@ +pub mod server; + +#[cfg(test)] +mod tests; + +use anyhow::{anyhow, Result}; +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine}; +use rand::{distributions::Alphanumeric, Rng}; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::time::Duration; +use tokio::sync::oneshot; +use tokio::time::timeout; + +/// Default models for Tetrate Agent Router Service configuration +pub const TETRATE_DEFAULT_MODEL: &str = "claude-4-sonnet-20250514"; + +// Auth endpoints are on the main web domain +const TETRATE_AUTH_URL: &str = "https://router.tetrate.ai/auth"; +const TETRATE_TOKEN_URL: &str = "https://router.tetrate.ai/api/api-keys/verify"; +const CALLBACK_URL: &str = "http://localhost:3000"; +const AUTH_TIMEOUT: Duration = Duration::from_secs(180); // 3 minutes + +#[derive(Debug)] +pub struct PkceAuthFlow { + code_verifier: String, + code_challenge: String, + server_shutdown_tx: Option>, +} + +#[derive(Debug, Deserialize)] +struct TokenResponse { + key: String, +} + +#[derive(Debug, Serialize)] +struct TokenRequest { + code: String, + code_verifier: String, +} + +impl PkceAuthFlow { + pub fn new() -> Result { + let code_verifier: String = rand::thread_rng() + .sample_iter(&Alphanumeric) + .take(128) + .map(char::from) + .collect(); + + let mut hasher = Sha256::new(); + hasher.update(&code_verifier); + let hash = hasher.finalize(); + + let code_challenge = URL_SAFE_NO_PAD.encode(hash); + + Ok(Self { + code_verifier, + code_challenge, + server_shutdown_tx: None, + }) + } + + pub fn get_auth_url(&self) -> String { + format!( + "{}?callback={}&code_challenge={}", + TETRATE_AUTH_URL, + urlencoding::encode(CALLBACK_URL), + urlencoding::encode(&self.code_challenge) + ) + } + + /// Start local server and wait for callback + pub async fn start_server(&mut self) -> Result { + let (code_tx, code_rx) = oneshot::channel::(); + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + + // Store shutdown sender so we can stop the server later + self.server_shutdown_tx = Some(shutdown_tx); + + // Start the server in a background task + tokio::spawn(async move { + if let Err(e) = server::run_callback_server(code_tx, shutdown_rx).await { + eprintln!("Server error: {}", e); + } + }); + + // Wait for the authorization code with timeout + match timeout(AUTH_TIMEOUT, code_rx).await { + Ok(Ok(code)) => Ok(code), + Ok(Err(_)) => Err(anyhow!("Failed to receive authorization code")), + Err(_) => Err(anyhow!("Authentication timeout - please try again")), + } + } + + pub async fn exchange_code(&self, code: String) -> Result { + let client = Client::new(); + + let request_body = TokenRequest { + code: code.clone(), + code_verifier: self.code_verifier.clone(), + }; + + eprintln!("Exchanging code for API key..."); + eprintln!("Code: {}", code); + eprintln!("Code verifier length: {}", self.code_verifier.len()); + eprintln!("Code challenge: {}", self.code_challenge); + + let response = client + .post(TETRATE_TOKEN_URL) + .json(&request_body) + .send() + .await?; + + if !response.status().is_success() { + let status = response.status(); + let error_text = response.text().await.unwrap_or_default(); + eprintln!("Token exchange failed!"); + eprintln!("Status: {}", status); + eprintln!("Error response: {}", error_text); + return Err(anyhow!( + "Failed to exchange code: {} - {}", + status, + error_text + )); + } + + let token_response: TokenResponse = response.json().await?; + Ok(token_response.key) + } + + /// Complete flow: open browser, wait for callback, exchange code + pub async fn complete_flow(&mut self) -> Result { + let auth_url = self.get_auth_url(); + + println!("Opening browser for Tetrate Agent Router Service authentication..."); + eprintln!("Auth URL: {}", auth_url); + + if let Err(e) = webbrowser::open(&auth_url) { + eprintln!("Failed to open browser automatically: {}", e); + println!("Please open this URL manually: {}", auth_url); + } + + println!("Waiting for authentication callback..."); + let code = self.start_server().await?; + + println!("Authorization code received. Exchanging for API key..."); + eprintln!("Received code: {}", code); + + let api_key = self.exchange_code(code).await?; + + // Shutdown the server if it's still running + if let Some(tx) = self.server_shutdown_tx.take() { + let _ = tx.send(()); + } + + Ok(api_key) + } +} + +pub use self::PkceAuthFlow as TetrateAuth; + +use crate::config::Config; +use serde_json::Value; + +pub fn configure_tetrate(config: &Config, api_key: String) -> Result<()> { + config.set_secret("TETRATE_API_KEY", Value::String(api_key))?; + config.set_param("GOOSE_PROVIDER", Value::String("tetrate".to_string()))?; + config.set_param( + "GOOSE_MODEL", + Value::String(TETRATE_DEFAULT_MODEL.to_string()), + )?; + Ok(()) +} diff --git a/crates/goose/src/config/signup_tetrate/server.rs b/crates/goose/src/config/signup_tetrate/server.rs new file mode 100644 index 0000000000..e1c9b1585f --- /dev/null +++ b/crates/goose/src/config/signup_tetrate/server.rs @@ -0,0 +1,85 @@ +use anyhow::Result; +use axum::{ + extract::Query, + http::StatusCode, + response::{Html, IntoResponse}, + routing::get, + Router, +}; +use include_dir::{include_dir, Dir}; +use minijinja::{context, Environment}; +use serde::Deserialize; +use std::net::SocketAddr; +use tokio::sync::oneshot; + +static TEMPLATES_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/config/signup_tetrate/templates"); + +#[derive(Debug, Deserialize)] +struct CallbackQuery { + code: Option, + error: Option, +} + +/// Run the callback server on localhost:3000 +pub async fn run_callback_server( + code_tx: oneshot::Sender, + shutdown_rx: oneshot::Receiver<()>, +) -> Result<()> { + let app = Router::new().route("/", get(handle_callback)); + let addr = SocketAddr::from(([127, 0, 0, 1], 3000)); + let listener = tokio::net::TcpListener::bind(addr).await?; + let state = std::sync::Arc::new(tokio::sync::Mutex::new(Some(code_tx))); + + axum::serve(listener, app.with_state(state.clone()).into_make_service()) + .with_graceful_shutdown(async move { + let _ = shutdown_rx.await; + }) + .await?; + + Ok(()) +} + +async fn handle_callback( + Query(params): Query, + state: axum::extract::State< + std::sync::Arc>>>, + >, +) -> impl IntoResponse { + if let Some(error) = params.error { + let mut env = Environment::new(); + let template_content = TEMPLATES_DIR + .get_file("error.html") + .expect("error.html template not found") + .contents_utf8() + .expect("error.html is not valid UTF-8"); + + env.add_template("error", template_content).unwrap(); + let tmpl = env.get_template("error").unwrap(); + let rendered = tmpl.render(context! { error => error }).unwrap(); + + return (StatusCode::BAD_REQUEST, Html(rendered)); + } + + if let Some(code) = params.code { + let mut tx_guard = state.lock().await; + if let Some(tx) = tx_guard.take() { + let _ = tx.send(code); + } + + let success_html = TEMPLATES_DIR + .get_file("success.html") + .expect("success.html template not found") + .contents_utf8() + .expect("success.html is not valid UTF-8"); + + return (StatusCode::OK, Html(success_html.to_string())); + } + + let invalid_html = TEMPLATES_DIR + .get_file("invalid.html") + .expect("invalid.html template not found") + .contents_utf8() + .expect("invalid.html is not valid UTF-8"); + + (StatusCode::BAD_REQUEST, Html(invalid_html.to_string())) +} diff --git a/crates/goose/src/config/signup_tetrate/templates/error.html b/crates/goose/src/config/signup_tetrate/templates/error.html new file mode 100644 index 0000000000..b1875d7fcd --- /dev/null +++ b/crates/goose/src/config/signup_tetrate/templates/error.html @@ -0,0 +1,85 @@ + + + + + + Authentication Error - Tetrate Agent Router Service + + + +
+
+ + + +
+

Authentication Failed

+

There was an error during the authentication process.

+
{{ error }}
+
+ Please close this window and try again. +
+
+ + \ No newline at end of file diff --git a/crates/goose/src/config/signup_tetrate/templates/invalid.html b/crates/goose/src/config/signup_tetrate/templates/invalid.html new file mode 100644 index 0000000000..6c1fb3887f --- /dev/null +++ b/crates/goose/src/config/signup_tetrate/templates/invalid.html @@ -0,0 +1,77 @@ + + + + + + Invalid Request - Tetrate Agent Router Service + + + +
+
+ + + +
+

Invalid Request

+

The authentication request is missing required parameters.

+
+ Please ensure you're accessing this page through the proper authentication flow. +
+
+ + \ No newline at end of file diff --git a/crates/goose/src/config/signup_tetrate/templates/success.html b/crates/goose/src/config/signup_tetrate/templates/success.html new file mode 100644 index 0000000000..98fbe00bf6 --- /dev/null +++ b/crates/goose/src/config/signup_tetrate/templates/success.html @@ -0,0 +1,76 @@ + + + + + + Authentication Successful - Tetrate Agent Router Service + + + +
+
+ + + +
+

Authentication Successful!

+

You've successfully authenticated with Tetrate Agent Router Service.

+
+ You can now close this window and return to your terminal. +
+
+ + \ No newline at end of file diff --git a/crates/goose/src/config/signup_tetrate/tests.rs b/crates/goose/src/config/signup_tetrate/tests.rs new file mode 100644 index 0000000000..8335337c2f --- /dev/null +++ b/crates/goose/src/config/signup_tetrate/tests.rs @@ -0,0 +1,87 @@ +use super::*; +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine}; +use sha2::{Digest, Sha256}; + +#[test] +fn test_pkce_flow_creation() { + let flow = PkceAuthFlow::new().unwrap(); + + // Verify code_verifier is 128 characters + assert_eq!(flow.code_verifier.len(), 128); + + // Verify code_verifier is alphanumeric + assert!(flow.code_verifier.chars().all(|c| c.is_alphanumeric())); + + // Verify code_challenge is base64url encoded + assert!(!flow.code_challenge.contains('+')); + assert!(!flow.code_challenge.contains('/')); + assert!(!flow.code_challenge.contains('=')); +} + +#[test] +fn test_code_challenge_generation() { + let flow = PkceAuthFlow::new().unwrap(); + + // Manually compute the expected challenge + let mut hasher = Sha256::new(); + hasher.update(&flow.code_verifier); + let hash = hasher.finalize(); + let expected_challenge = URL_SAFE_NO_PAD.encode(hash); + + assert_eq!(flow.code_challenge, expected_challenge); +} + +#[test] +fn test_auth_url_generation() { + let flow = PkceAuthFlow::new().unwrap(); + let auth_url = flow.get_auth_url(); + + // Verify URL contains required parameters + assert!(auth_url.contains("callback=")); + assert!(auth_url.contains("code_challenge=")); + assert!(auth_url.starts_with(TETRATE_AUTH_URL)); + + // Verify callback URL is properly encoded + assert!(auth_url.contains(&*urlencoding::encode(CALLBACK_URL))); +} + +#[test] +fn test_different_verifiers_produce_different_challenges() { + let flow1 = PkceAuthFlow::new().unwrap(); + let flow2 = PkceAuthFlow::new().unwrap(); + + // Verifiers should be different (extremely high probability) + assert_ne!(flow1.code_verifier, flow2.code_verifier); + + // Challenges should also be different + assert_ne!(flow1.code_challenge, flow2.code_challenge); +} + +#[test] +fn test_configure_tetrate() { + use crate::config::Config; + use tempfile::TempDir; + + // Create a test config with temporary paths + let temp_dir = TempDir::new().unwrap(); + let config_path = temp_dir.path().join("test_config.yaml"); + let config = Config::new(&config_path, "test_service").unwrap(); + + // Configure with a test API key + let test_key = "test-api-key-123".to_string(); + configure_tetrate(&config, test_key.clone()).unwrap(); + + // Verify the configuration was set correctly + assert_eq!( + config.get_secret::("TETRATE_API_KEY").unwrap(), + test_key + ); + assert_eq!( + config.get_param::("GOOSE_PROVIDER").unwrap(), + "tetrate" + ); + assert_eq!( + config.get_param::("GOOSE_MODEL").unwrap(), + TETRATE_DEFAULT_MODEL.to_string() + ); +} diff --git a/crates/goose/src/providers/factory.rs b/crates/goose/src/providers/factory.rs index fbcab89b83..c06c8e954f 100644 --- a/crates/goose/src/providers/factory.rs +++ b/crates/goose/src/providers/factory.rs @@ -20,6 +20,7 @@ use super::{ provider_registry::ProviderRegistry, sagemaker_tgi::SageMakerTgiProvider, snowflake::SnowflakeProvider, + tetrate::TetrateProvider, venice::VeniceProvider, xai::XaiProvider, }; @@ -55,6 +56,7 @@ static REGISTRY: Lazy> = Lazy::new(|| { registry.register::(OpenRouterProvider::from_env); registry.register::(SageMakerTgiProvider::from_env); registry.register::(SnowflakeProvider::from_env); + registry.register::(TetrateProvider::from_env); registry.register::(VeniceProvider::from_env); registry.register::(XaiProvider::from_env); diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 5fe80017f7..64b29ed32e 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -29,6 +29,7 @@ mod retry; pub mod sagemaker_tgi; pub mod snowflake; pub mod testprovider; +pub mod tetrate; pub mod toolshim; pub mod usage_estimator; pub mod utils; diff --git a/crates/goose/src/providers/tetrate.rs b/crates/goose/src/providers/tetrate.rs new file mode 100644 index 0000000000..ce865f4cc5 --- /dev/null +++ b/crates/goose/src/providers/tetrate.rs @@ -0,0 +1,264 @@ +use anyhow::{Error, Result}; +use async_trait::async_trait; +use serde_json::Value; + +use super::api_client::{ApiClient, AuthMethod}; +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage}; +use super::errors::ProviderError; +use super::retry::ProviderRetry; +use super::utils::{ + emit_debug_trace, get_model, handle_response_google_compat, handle_response_openai_compat, + is_google_model, +}; +use crate::config::signup_tetrate::TETRATE_DEFAULT_MODEL; +use crate::conversation::message::Message; +use crate::impl_provider_default; +use crate::model::ModelConfig; +use crate::providers::formats::openai::{create_request, get_usage, response_to_message}; +use rmcp::model::Tool; + +// Tetrate Agent Router Service can run many models, we suggest the default +pub const TETRATE_KNOWN_MODELS: &[&str] = &[ + "claude-opus-4-1", + "claude-3-7-sonnet-latest", + "claude-3-5-sonnet-latest", + "claude-3-5-haiku-latest", + "gemini-2.5-pro", + "gemini-2.0-flash", + "gemini-2.0-flash-lite", + "gpt-5", + "gpt-5-mini", + "gpt-5-nano", + "gpt-4.1", +]; +pub const TETRATE_DOC_URL: &str = "https://router.tetrate.ai"; + +#[derive(serde::Serialize)] +pub struct TetrateProvider { + #[serde(skip)] + api_client: ApiClient, + model: ModelConfig, +} + +impl_provider_default!(TetrateProvider); + +impl TetrateProvider { + pub fn from_env(model: ModelConfig) -> Result { + let config = crate::config::Config::global(); + let api_key: String = config.get_secret("TETRATE_API_KEY")?; + // API host for LLM endpoints (/v1/chat/completions, /v1/models) + let host: String = config + .get_param("TETRATE_HOST") + .unwrap_or_else(|_| "https://api.router.tetrate.ai".to_string()); + + let auth = AuthMethod::BearerToken(api_key); + let api_client = ApiClient::new(host, auth)? + .with_header("HTTP-Referer", "https://block.github.io/goose")? + .with_header("X-Title", "Goose")?; + + Ok(Self { api_client, model }) + } + + async fn post(&self, payload: &Value) -> Result { + let response = self + .api_client + .response_post("v1/chat/completions", payload) + .await?; + + // Handle Google-compatible model responses differently + if is_google_model(payload) { + return handle_response_google_compat(response).await; + } + + // For OpenAI-compatible models, parse the response body to JSON + let response_body = handle_response_openai_compat(response) + .await + .map_err(|e| ProviderError::RequestFailed(format!("Failed to parse response: {e}")))?; + + let _debug = format!( + "Tetrate Agent Router Service request with payload: {} and response: {}", + serde_json::to_string_pretty(payload).unwrap_or_else(|_| "Invalid JSON".to_string()), + serde_json::to_string_pretty(&response_body) + .unwrap_or_else(|_| "Invalid JSON".to_string()) + ); + + // Tetrate Agent Router Service can return errors in 200 OK responses, so we have to check for errors explicitly + if let Some(error_obj) = response_body.get("error") { + // If there's an error object, extract the error message and code + let error_message = error_obj + .get("message") + .and_then(|m| m.as_str()) + .unwrap_or("Unknown Tetrate Agent Router Service error"); + + let error_code = error_obj.get("code").and_then(|c| c.as_u64()).unwrap_or(0); + + // Check for context length errors in the error message + if error_code == 400 && error_message.contains("maximum context length") { + return Err(ProviderError::ContextLengthExceeded( + error_message.to_string(), + )); + } + + // Return appropriate error based on the error code + match error_code { + 401 | 403 => return Err(ProviderError::Authentication(error_message.to_string())), + 429 => return Err(ProviderError::RateLimitExceeded(error_message.to_string())), + 500 | 503 => return Err(ProviderError::ServerError(error_message.to_string())), + _ => return Err(ProviderError::RequestFailed(error_message.to_string())), + } + } + + // No error detected, return the response body + Ok(response_body) + } +} + +fn create_request_based_on_model( + provider: &TetrateProvider, + system: &str, + messages: &[Message], + tools: &[Tool], +) -> anyhow::Result { + let payload = create_request( + &provider.model, + system, + messages, + tools, + &super::utils::ImageFormat::OpenAi, + )?; + + Ok(payload) +} + +#[async_trait] +impl Provider for TetrateProvider { + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + "tetrate", + "Tetrate Agent Router Service", + "Enterprise router for AI models", + TETRATE_DEFAULT_MODEL, + TETRATE_KNOWN_MODELS.to_vec(), + TETRATE_DOC_URL, + vec![ + ConfigKey::new("TETRATE_API_KEY", true, true, None), + ConfigKey::new( + "TETRATE_HOST", + false, + false, + Some("https://api.router.tetrate.ai"), + ), + ], + ) + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() + } + + #[tracing::instrument( + skip(self, system, messages, tools), + fields(model_config, input, output, input_tokens, output_tokens, total_tokens) + )] + async fn complete( + &self, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result<(Message, ProviderUsage), ProviderError> { + // Create the base payload + let payload = create_request_based_on_model(self, system, messages, tools)?; + + // Make request + let response = self + .with_retry(|| async { + let payload_clone = payload.clone(); + self.post(&payload_clone).await + }) + .await?; + + // Parse response + let message = response_to_message(&response)?; + let usage = response.get("usage").map(get_usage).unwrap_or_else(|| { + tracing::debug!("Failed to get usage data"); + Usage::default() + }); + let model = get_model(&response); + emit_debug_trace(&self.model, &payload, &response, &usage); + Ok((message, ProviderUsage::new(model, usage))) + } + + /// Fetch supported models from Tetrate Agent Router Service API (only models with tool support) + async fn fetch_supported_models(&self) -> Result>, ProviderError> { + // Use the existing api_client which already has authentication configured + let response = match self.api_client.response_get("v1/models").await { + Ok(response) => response, + Err(e) => { + tracing::warn!("Failed to fetch models from Tetrate Agent Router Service API: {}, falling back to manual model entry", e); + return Ok(None); + } + }; + + // Handle JSON parsing failures gracefully + let json: serde_json::Value = match response.json().await { + Ok(json) => json, + Err(e) => { + tracing::warn!("Failed to parse Tetrate Agent Router Service API response as JSON: {}, falling back to manual model entry", e); + return Ok(None); + } + }; + + // Check for error in response + if let Some(err_obj) = json.get("error") { + let msg = err_obj + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or("unknown error"); + tracing::warn!( + "Tetrate Agent Router Service API returned an error: {}", + msg + ); + return Ok(None); + } + + // The response format from /v1/models is expected to be OpenAI-compatible + // It should have a "data" field with an array of model objects + let data = json.get("data").and_then(|v| v.as_array()).ok_or_else(|| { + ProviderError::UsageError("Missing data field in JSON response".into()) + })?; + + let mut models: Vec = data + .iter() + .filter_map(|model| { + // Get the model ID + let id = model.get("id").and_then(|v| v.as_str())?; + + // Check if the model supports computer_use (which indicates tool/function support) + // The Tetrate API uses "supports_computer_use" instead of "supported_parameters" + let supports_computer_use = model + .get("supports_computer_use") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + if supports_computer_use { + Some(id.to_string()) + } else { + tracing::debug!( + "Model '{}' does not support computer_use (tool support), skipping", + id + ); + None + } + }) + .collect(); + + // If no models with tool support were found, fall back to manual entry + if models.is_empty() { + tracing::warn!("No models with tool support found in Tetrate Agent Router Service API response, falling back to manual model entry"); + return Ok(None); + } + + models.sort(); + Ok(Some(models)) + } +} diff --git a/ui/desktop/src/components/ProviderGuard.tsx b/ui/desktop/src/components/ProviderGuard.tsx index 761c000904..d505461e75 100644 --- a/ui/desktop/src/components/ProviderGuard.tsx +++ b/ui/desktop/src/components/ProviderGuard.tsx @@ -3,6 +3,7 @@ import { useNavigate } from 'react-router-dom'; import { useConfig } from './ConfigContext'; import { SetupModal } from './SetupModal'; import { startOpenRouterSetup } from '../utils/openRouterSetup'; +import { startTetrateSetup } from '../utils/tetrateSetup'; import WelcomeGooseLogo from './WelcomeGooseLogo'; import { initializeSystem } from '../utils/providerUtils'; import { toastService } from '../toasts'; @@ -31,6 +32,80 @@ export default function ProviderGuard({ children }: ProviderGuardProps) { showRetry: boolean; autoClose?: number; } | null>(null); + const [tetrateSetupState, setTetrateSetupState] = useState<{ + show: boolean; + title: string; + message: string; + showProgress: boolean; + showRetry: boolean; + autoClose?: number; + } | null>(null); + + const handleTetrateSetup = async () => { + setTetrateSetupState({ + show: true, + title: 'Setting up Tetrate Agent Router Service', + message: 'A browser window will open for authentication...', + showProgress: true, + showRetry: false, + }); + + const result = await startTetrateSetup(); + if (result.success) { + setTetrateSetupState({ + show: true, + title: 'Setup Complete!', + message: 'Tetrate Agent Router has been configured successfully. Initializing Goose...', + showProgress: true, + showRetry: false, + }); + + // After successful Tetrate setup, force reload config and initialize system + try { + // Get the latest config from disk + const config = window.electron.getConfig(); + const provider = (await read('GOOSE_PROVIDER', false)) ?? config.GOOSE_DEFAULT_PROVIDER; + const model = (await read('GOOSE_MODEL', false)) ?? config.GOOSE_DEFAULT_MODEL; + + if (provider && model) { + // Initialize the system with the new provider/model + await initializeSystem(provider as string, model as string, { + getExtensions, + addExtension, + }); + + toastService.configure({ silent: false }); + toastService.success({ + title: 'Success!', + msg: `Started goose with ${model} by Tetrate. You can change the model via the dropdown.`, + }); + + // Close the modal and mark as having provider + setTetrateSetupState(null); + setShowFirstTimeSetup(false); + setHasProvider(true); + } else { + throw new Error('Provider or model not found after Tetrate setup'); + } + } catch (error) { + console.error('Failed to initialize after Tetrate setup:', error); + toastService.configure({ silent: false }); + toastService.error({ + title: 'Initialization Failed', + msg: `Failed to initialize with Tetrate: ${error instanceof Error ? error.message : String(error)}`, + traceback: error instanceof Error ? error.stack || '' : '', + }); + } + } else { + setTetrateSetupState({ + show: true, + title: 'Tetrate setup pending', + message: result.message, + showProgress: false, + showRetry: true, + }); + } + }; const handleOpenRouterSetup = async () => { setOpenRouterSetupState({ @@ -68,14 +143,14 @@ export default function ProviderGuard({ children }: ProviderGuardProps) { toastService.configure({ silent: false }); toastService.success({ title: 'Success!', - msg: `Started goose with ${model} by OpenRouter. You can change the model via the lower right corner.`, + msg: `Started goose with ${model} by OpenRouter. You can change the model via the dropdown.`, }); // Close the modal and mark as having provider setOpenRouterSetupState(null); setShowFirstTimeSetup(false); setHasProvider(true); - + // Navigate to chat after successful setup navigate('/', { replace: true }); } else { @@ -121,7 +196,6 @@ export default function ProviderGuard({ children }: ProviderGuardProps) { } const ollamaStatus = await checkOllamaStatus(); setOllamaDetected(ollamaStatus.isRunning); - } catch (error) { // On error, assume no provider and redirect to welcome console.error('Error checking provider configuration:', error); @@ -152,7 +226,13 @@ export default function ProviderGuard({ children }: ProviderGuardProps) { }; }, [showFirstTimeSetup]); - if (isChecking && !openRouterSetupState?.show && !showFirstTimeSetup && !showOllamaSetup) { + if ( + isChecking && + !openRouterSetupState?.show && + !tetrateSetupState?.show && + !showFirstTimeSetup && + !showOllamaSetup + ) { return (
@@ -174,6 +254,20 @@ export default function ProviderGuard({ children }: ProviderGuardProps) { ); } + if (tetrateSetupState?.show) { + return ( + setTetrateSetupState(null)} + /> + ); + } + if (showOllamaSetup) { return (
@@ -210,111 +304,178 @@ export default function ProviderGuard({ children }: ProviderGuardProps) {
-

- Welcome to Goose -

+

Welcome to Goose

- Since it's your first time here, let's get your set you with a provider so we can make incredible work together. + Since it's your first time here, let's get you setup with a provider so we can + make incredible work together. Scroll down to see options.

{/* Setup options - same width container */} +
- {/* Primary OpenRouter Card with subtle shimmer - wrapped for badge positioning */} -
- {/* Recommended badge - positioned relative to wrapper */} -
- - Recommended - -
- -
- {/* Subtle shimmer effect */} -
- -
-
- -

- Automatic setup with OpenRouter -

+
+ {/* Tetrate Card */} + {/* Recommended badge - positioned relative to wrapper */} +
+ + Recommended +
-
- - - -
-
-

- Get instant access to multiple AI models including GPT-4, Claude, and more. Quick setup with just a few clicks. -

-
-
- {/* Ollama Card - outline style */} -
- {/* Detected badge - similar to recommended but green */} - {ollamaDetected && ( -
- - Detected - -
- )} - -
{ - setShowFirstTimeSetup(false); - setShowOllamaSetup(true); - }} - className="w-full p-4 sm:p-6 bg-transparent border border-background-hover rounded-xl hover:border-text-muted transition-all duration-200 cursor-pointer group" - > -
-
- -

- Ollama -

-
-
- - - +
+
+
+

+ Automatic setup with Tetrate Agent Router +

+
+
+ + + +
+
+

+ Get secure access to multiple AI models, start for free. Quick setup with just + a few clicks. +

-

- Run AI models locally on your computer. Completely free and private with no internet required. -

-
-
- {/* Other providers Card - outline style */} -
navigate('/welcome', { replace: true })} - className="w-full p-4 sm:p-6 bg-transparent border border-background-hover rounded-xl hover:border-text-muted transition-all duration-200 cursor-pointer group" - > -
-
-

- Other providers -

-
-
- - - -
-
-

- If you've already signed up for providers like Anthropic, OpenAI etc, you can enter your own keys. -

-
+ {/* Primary OpenRouter Card with subtle shimmer - wrapped for badge positioning */} +
+
+ {/* Subtle shimmer effect */} +
+
+
+ +

+ Automatic setup with OpenRouter +

+
+
+ + + +
+
+

+ Get instant access to multiple AI models including GPT-4, Claude, and more. + Quick setup with just a few clicks. +

+
+
+ + {/* Ollama Card - outline style */} +
+ {/* Detected badge - similar to recommended but green */} + {ollamaDetected && ( +
+ + Detected + +
+ )} + +
{ + setShowFirstTimeSetup(false); + setShowOllamaSetup(true); + }} + className="w-full p-4 sm:p-6 bg-transparent border border-background-hover rounded-xl hover:border-text-muted transition-all duration-200 cursor-pointer group" + > +
+
+ +

+ Ollama +

+
+
+ + + +
+
+

+ Run AI models locally on your computer. Completely free and private with no + internet required. +

+
+
+ + {/* Other providers Card - outline style */} +
navigate('/welcome', { replace: true })} + className="w-full p-4 sm:p-6 bg-transparent border border-background-hover rounded-xl hover:border-text-muted transition-all duration-200 cursor-pointer group" + > +
+
+

+ Other providers +

+
+
+ + + +
+
+

+ If you've already signed up for providers like Anthropic, OpenAI etc, you can + enter your own keys. +

+
diff --git a/ui/desktop/src/components/settings/providers/ProviderSettingsPage.tsx b/ui/desktop/src/components/settings/providers/ProviderSettingsPage.tsx index 37aa7c373a..dd2b9c8103 100644 --- a/ui/desktop/src/components/settings/providers/ProviderSettingsPage.tsx +++ b/ui/desktop/src/components/settings/providers/ProviderSettingsPage.tsx @@ -76,7 +76,7 @@ export default function ProviderSettings({ onClose, isOnboarding }: ProviderSett toastService.configure({ silent: false }); toastService.success({ title: 'Success!', - msg: `Started goose with ${model} by ${provider.metadata.display_name}. You can change the model via the lower right corner.`, + msg: `Started goose with ${model} by ${provider.metadata.display_name}. You can change the model via the dropdown.`, }); onClose(); diff --git a/ui/desktop/src/utils/tetrateSetup.ts b/ui/desktop/src/utils/tetrateSetup.ts new file mode 100644 index 0000000000..554aa67bc1 --- /dev/null +++ b/ui/desktop/src/utils/tetrateSetup.ts @@ -0,0 +1,25 @@ +export interface TetrateSetupStatus { + isRunning: boolean; + error: string | null; +} + +export async function startTetrateSetup(): Promise<{ success: boolean; message: string }> { + const baseUrl = `${window.appConfig.get('GOOSE_API_HOST')}:${window.appConfig.get('GOOSE_PORT')}`; + const response = await fetch(`${baseUrl}/handle_tetrate`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + }); + + if (!response.ok) { + console.error('Failed to start Tetrate setup:', response.statusText); + return { + success: false, + message: `Failed to start Tetrate setup ['${response.status}]`, + }; + } + + const result = await response.json(); + return result; +}