From 95325f4c5e55d9dd3364a11f2369b0da0de5ebe6 Mon Sep 17 00:00:00 2001 From: raj-subhankar Date: Sun, 15 Feb 2026 11:28:18 +0530 Subject: [PATCH] Use dynamic port Signed-off-by: raj-subhankar --- crates/goose/examples/tetrate_auth.rs | 9 +-- crates/goose/src/config/signup_tetrate/mod.rs | 56 +++++++++---------- .../goose/src/config/signup_tetrate/server.rs | 6 +- .../goose/src/config/signup_tetrate/tests.rs | 7 ++- 4 files changed, 35 insertions(+), 43 deletions(-) diff --git a/crates/goose/examples/tetrate_auth.rs b/crates/goose/examples/tetrate_auth.rs index 780ee269b8ec..6f55ddbed588 100644 --- a/crates/goose/examples/tetrate_auth.rs +++ b/crates/goose/examples/tetrate_auth.rs @@ -10,13 +10,10 @@ async fn main() -> Result<(), Box> { // 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!("Starting authentication flow..."); println!("This will:"); - println!("1. Open your browser to the auth page"); - println!("2. Start a local server on port 3000"); + println!("1. Start a local server on a dynamic port"); + println!("2. Open your browser to the auth page"); println!("3. Wait for the callback\n"); // Complete the full flow diff --git a/crates/goose/src/config/signup_tetrate/mod.rs b/crates/goose/src/config/signup_tetrate/mod.rs index fc68442447a0..58508ca30925 100644 --- a/crates/goose/src/config/signup_tetrate/mod.rs +++ b/crates/goose/src/config/signup_tetrate/mod.rs @@ -19,7 +19,7 @@ pub const TETRATE_DEFAULT_MODEL: &str = "claude-haiku-4-5"; // 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 CALLBACK_BASE: &str = "http://localhost"; const AUTH_TIMEOUT: Duration = Duration::from_secs(180); // 3 minutes #[derive(Debug)] @@ -61,38 +61,16 @@ impl PkceAuthFlow { }) } - pub fn get_auth_url(&self) -> String { + pub fn get_auth_url(&self, port: u16) -> String { + let callback_url = format!("{}:{}", CALLBACK_BASE, port); format!( "{}?callback={}&code_challenge={}", TETRATE_AUTH_URL, - urlencoding::encode(CALLBACK_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(); @@ -131,9 +109,22 @@ impl PkceAuthFlow { Ok(token_response.key) } - /// Complete flow: open browser, wait for callback, exchange code + /// Complete flow: start server, open browser, wait for callback, exchange code pub async fn complete_flow(&mut self) -> Result { - let auth_url = self.get_auth_url(); + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)).await?; + let port = listener.local_addr()?.port(); + + let (code_tx, code_rx) = oneshot::channel::(); + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + self.server_shutdown_tx = Some(shutdown_tx); + + tokio::spawn(async move { + if let Err(e) = server::run_callback_server(listener, code_tx, shutdown_rx).await { + eprintln!("Server error: {}", e); + } + }); + + let auth_url = self.get_auth_url(port); println!("Opening browser for Tetrate Agent Router Service authentication..."); eprintln!("Auth URL: {}", auth_url); @@ -143,8 +134,13 @@ impl PkceAuthFlow { println!("Please open this URL manually: {}", auth_url); } - println!("Waiting for authentication callback..."); - let code = self.start_server().await?; + println!("Waiting for authentication callback on port {}...", port); + + let code = 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")), + }?; println!("Authorization code received. Exchanging for API key..."); eprintln!("Received code: {}", code); diff --git a/crates/goose/src/config/signup_tetrate/server.rs b/crates/goose/src/config/signup_tetrate/server.rs index e1c9b1585f93..bf72c92e2eec 100644 --- a/crates/goose/src/config/signup_tetrate/server.rs +++ b/crates/goose/src/config/signup_tetrate/server.rs @@ -9,7 +9,6 @@ use axum::{ 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"); @@ -20,14 +19,13 @@ struct CallbackQuery { error: Option, } -/// Run the callback server on localhost:3000 +/// Run the callback server using the provided listener. pub async fn run_callback_server( + listener: tokio::net::TcpListener, 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()) diff --git a/crates/goose/src/config/signup_tetrate/tests.rs b/crates/goose/src/config/signup_tetrate/tests.rs index 75124e7f9969..975df4ec65b6 100644 --- a/crates/goose/src/config/signup_tetrate/tests.rs +++ b/crates/goose/src/config/signup_tetrate/tests.rs @@ -34,15 +34,16 @@ fn test_code_challenge_generation() { #[test] fn test_auth_url_generation() { let flow = PkceAuthFlow::new().unwrap(); - let auth_url = flow.get_auth_url(); + let auth_url = flow.get_auth_url(12345); // 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))); + // Verify callback URL contains the dynamic port + let expected_callback = format!("{}:{}", CALLBACK_BASE, 12345); + assert!(auth_url.contains(&*urlencoding::encode(&expected_callback))); } #[test]