Add PKCE support for Tetrate Agent Router Service

crates/goose-cli/src/commands/configure.rs

@@ -72,6 +72,11 @@ pub async fn handle_configure() -> Result<(), Box<dyn Error>> {
                 "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<dyn Error>> {
                     }
                 }
             }
+            "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<dyn Error>> {
     Ok(())
 }
 
+/// Handle Tetrate Agent Router Service authentication
+pub async fn handle_tetrate_auth() -> Result<(), Box<dyn Error>> {
+    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<dyn Error>> {
     let provider_type = cliclack::select("What type of API is this?")
         .item(

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<AppState>) -> 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<Arc<AppState>>,
+) -> Result<Json<SetupResponse>, 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),
+            }))
+        }
+    }
+}

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<dyn std::error::Error>> {
+    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::<String>()
+            );
+            println!("\nYou can now use this API key with the Tetrate provider.");
+        }
+        Err(e) => {
+            eprintln!("\n❌ Authentication failed: {}", e);
+            eprintln!("Error details: {:?}", e);
+        }
+    }
+
+    Ok(())
+}

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;