Add PKCE support for Tetrate Agent Router Service (#4165)
Signed-off-by: John Landa <jonathanlanda@gmail.com> Co-authored-by: Michael Neale <michael.neale@gmail.com>
This commit is contained in:
John Landa
@@ -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;
|
||||
|
||||
@@ -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<oneshot::Sender<()>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TokenResponse {
|
||||
key: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct TokenRequest {
|
||||
code: String,
|
||||
code_verifier: String,
|
||||
}
|
||||
|
||||
impl PkceAuthFlow {
|
||||
pub fn new() -> Result<Self> {
|
||||
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<String> {
|
||||
let (code_tx, code_rx) = oneshot::channel::<String>();
|
||||
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<String> {
|
||||
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<String> {
|
||||
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(())
|
||||
}
|
||||
@@ -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<String>,
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
/// Run the callback server on localhost:3000
|
||||
pub async fn run_callback_server(
|
||||
code_tx: oneshot::Sender<String>,
|
||||
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<CallbackQuery>,
|
||||
state: axum::extract::State<
|
||||
std::sync::Arc<tokio::sync::Mutex<Option<oneshot::Sender<String>>>>,
|
||||
>,
|
||||
) -> 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()))
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Authentication Error - Tetrate Agent Router Service</title>
|
||||
<style>
|
||||
body {
|
||||
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
|
||||
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
|
||||
min-height: 100vh;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
margin: 0;
|
||||
padding: 20px;
|
||||
}
|
||||
.container {
|
||||
background: white;
|
||||
border-radius: 12px;
|
||||
box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1);
|
||||
padding: 40px;
|
||||
max-width: 400px;
|
||||
width: 100%;
|
||||
text-align: center;
|
||||
}
|
||||
.error-icon {
|
||||
width: 64px;
|
||||
height: 64px;
|
||||
margin: 0 auto 20px;
|
||||
background: #ef4444;
|
||||
border-radius: 50%;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
}
|
||||
.error-icon svg {
|
||||
width: 32px;
|
||||
height: 32px;
|
||||
fill: white;
|
||||
}
|
||||
h1 {
|
||||
color: #1f2937;
|
||||
font-size: 24px;
|
||||
margin: 0 0 10px;
|
||||
}
|
||||
p {
|
||||
color: #6b7280;
|
||||
font-size: 16px;
|
||||
line-height: 1.5;
|
||||
margin: 0 0 20px;
|
||||
}
|
||||
.error-message {
|
||||
background: #fef2f2;
|
||||
border: 1px solid #fee2e2;
|
||||
border-radius: 8px;
|
||||
padding: 12px;
|
||||
color: #991b1b;
|
||||
font-size: 14px;
|
||||
font-family: monospace;
|
||||
word-break: break-all;
|
||||
}
|
||||
.note {
|
||||
margin-top: 20px;
|
||||
color: #6b7280;
|
||||
font-size: 14px;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="container">
|
||||
<div class="error-icon">
|
||||
<svg viewBox="0 0 24 24">
|
||||
<path d="M6 18L18 6M6 6l12 12"/>
|
||||
</svg>
|
||||
</div>
|
||||
<h1>Authentication Failed</h1>
|
||||
<p>There was an error during the authentication process.</p>
|
||||
<div class="error-message">{{ error }}</div>
|
||||
<div class="note">
|
||||
Please close this window and try again.
|
||||
</div>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,77 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Invalid Request - Tetrate Agent Router Service</title>
|
||||
<style>
|
||||
body {
|
||||
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
|
||||
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
|
||||
min-height: 100vh;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
margin: 0;
|
||||
padding: 20px;
|
||||
}
|
||||
.container {
|
||||
background: white;
|
||||
border-radius: 12px;
|
||||
box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1);
|
||||
padding: 40px;
|
||||
max-width: 400px;
|
||||
width: 100%;
|
||||
text-align: center;
|
||||
}
|
||||
.warning-icon {
|
||||
width: 64px;
|
||||
height: 64px;
|
||||
margin: 0 auto 20px;
|
||||
background: #f59e0b;
|
||||
border-radius: 50%;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
}
|
||||
.warning-icon svg {
|
||||
width: 32px;
|
||||
height: 32px;
|
||||
fill: white;
|
||||
}
|
||||
h1 {
|
||||
color: #1f2937;
|
||||
font-size: 24px;
|
||||
margin: 0 0 10px;
|
||||
}
|
||||
p {
|
||||
color: #6b7280;
|
||||
font-size: 16px;
|
||||
line-height: 1.5;
|
||||
margin: 0 0 20px;
|
||||
}
|
||||
.note {
|
||||
background: #fffbeb;
|
||||
border: 1px solid #fef3c7;
|
||||
border-radius: 8px;
|
||||
padding: 12px;
|
||||
color: #92400e;
|
||||
font-size: 14px;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="container">
|
||||
<div class="warning-icon">
|
||||
<svg viewBox="0 0 24 24">
|
||||
<path d="M12 2L2 7v10c0 5.55 3.84 10.74 9 12 5.16-1.26 9-6.45 9-12V7l-10-5zm0 10h0m0 4h0"/>
|
||||
</svg>
|
||||
</div>
|
||||
<h1>Invalid Request</h1>
|
||||
<p>The authentication request is missing required parameters.</p>
|
||||
<div class="note">
|
||||
Please ensure you're accessing this page through the proper authentication flow.
|
||||
</div>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,76 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Authentication Successful - Tetrate Agent Router Service</title>
|
||||
<style>
|
||||
body {
|
||||
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
|
||||
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
|
||||
min-height: 100vh;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
margin: 0;
|
||||
padding: 20px;
|
||||
}
|
||||
.container {
|
||||
background: white;
|
||||
border-radius: 12px;
|
||||
box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1);
|
||||
padding: 40px;
|
||||
max-width: 400px;
|
||||
width: 100%;
|
||||
text-align: center;
|
||||
}
|
||||
.success-icon {
|
||||
width: 64px;
|
||||
height: 64px;
|
||||
margin: 0 auto 20px;
|
||||
background: #10b981;
|
||||
border-radius: 50%;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
}
|
||||
.success-icon svg {
|
||||
width: 32px;
|
||||
height: 32px;
|
||||
fill: white;
|
||||
}
|
||||
h1 {
|
||||
color: #1f2937;
|
||||
font-size: 24px;
|
||||
margin: 0 0 10px;
|
||||
}
|
||||
p {
|
||||
color: #6b7280;
|
||||
font-size: 16px;
|
||||
line-height: 1.5;
|
||||
margin: 0 0 20px;
|
||||
}
|
||||
.note {
|
||||
background: #f3f4f6;
|
||||
border-radius: 8px;
|
||||
padding: 12px;
|
||||
color: #4b5563;
|
||||
font-size: 14px;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="container">
|
||||
<div class="success-icon">
|
||||
<svg viewBox="0 0 24 24">
|
||||
<path d="M9 16.17L4.83 12l-1.42 1.41L9 19 21 7l-1.41-1.41L9 16.17z"/>
|
||||
</svg>
|
||||
</div>
|
||||
<h1>Authentication Successful!</h1>
|
||||
<p>You've successfully authenticated with Tetrate Agent Router Service.</p>
|
||||
<div class="note">
|
||||
You can now close this window and return to your terminal.
|
||||
</div>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
@@ -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::<String>("TETRATE_API_KEY").unwrap(),
|
||||
test_key
|
||||
);
|
||||
assert_eq!(
|
||||
config.get_param::<String>("GOOSE_PROVIDER").unwrap(),
|
||||
"tetrate"
|
||||
);
|
||||
assert_eq!(
|
||||
config.get_param::<String>("GOOSE_MODEL").unwrap(),
|
||||
TETRATE_DEFAULT_MODEL.to_string()
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user