diff --git a/crates/goose-cli/src/commands/configure.rs b/crates/goose-cli/src/commands/configure.rs index 7c9700690a5c..6cf9ec2dda66 100644 --- a/crates/goose-cli/src/commands/configure.rs +++ b/crates/goose-cli/src/commands/configure.rs @@ -238,6 +238,11 @@ async fn handle_first_time_setup(config: &Config) -> anyhow::Result<()> { cliclack::intro(style(" goose-configure ").on_cyan().black())?; let setup_method = cliclack::select("How would you like to set up your provider?") + .item( + "aimlapi", + "AI/ML API Login (Recommended)", + "Sign in with AI/ML API to automatically configure models", + ) .item( "openrouter", "OpenRouter Login (Recommended)", @@ -249,10 +254,6 @@ async fn handle_first_time_setup(config: &Config) -> anyhow::Result<()> { "Sign in with Tetrate Agent Router Service to automatically configure models", ) .item( - "aimlapi", - "AI/ML API Login", - "Sign in with AI/ML API to automatically configure models", - ) .item( "manual", "Manual Configuration", "Choose a provider and enter credentials manually", @@ -1971,13 +1972,17 @@ pub async fn handle_aimlapi_auth() -> anyhow::Result<()> { let mut auth_flow = AimlapiAuth::new()?; let api_key = auth_flow.complete_flow().await?; - println!(" -Sign-in complete!"); + println!( + " +Sign-in complete!" + ); let config = Config::global(); - println!(" -Configuring AI/ML API..."); + println!( + " +Configuring AI/ML API..." + ); configure_aimlapi(config, api_key)?; println!("AI/ML API configuration complete"); diff --git a/crates/goose/src/config/signup_aimlapi/mod.rs b/crates/goose/src/config/signup_aimlapi/mod.rs index 63335e6a8643..c84e44487726 100644 --- a/crates/goose/src/config/signup_aimlapi/mod.rs +++ b/crates/goose/src/config/signup_aimlapi/mod.rs @@ -1,5 +1,8 @@ pub mod server; +#[cfg(test)] +mod tests; + use anyhow::{anyhow, Result}; use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine}; use rand::{distr::Alphanumeric, RngExt}; @@ -21,11 +24,17 @@ const AIMLAPI_DEFAULT_MODEL: &str = "anthropic/claude-sonnet-5"; /// non-production environment for testing; the default is production and no /// user ever needs to set it. const AIMLAPI_APP_URL_ENV: &str = "AIMLAPI_APP_URL"; -const AIMLAPI_APP_URL_DEFAULT: &str = "https://app.aimlapi.com"; +pub(crate) const AIMLAPI_APP_URL_DEFAULT: &str = "https://app.aimlapi.com"; /// Where the browser consent screen lives (the page the user actually sees). +/// +/// This is sent as `verificationBaseUrl`, and the server hands it straight back +/// with `/agent/authorize` appended. The web app is served under an `/app/` +/// base path, so the base has to carry it: dropping the `/app` yields +/// `https://aimlapi.com/agent/authorize`, which is a 404 and strands the user +/// on the very first step of the flow. const AIMLAPI_WEB_URL_ENV: &str = "AIMLAPI_WEB_URL"; -const AIMLAPI_WEB_URL_DEFAULT: &str = "https://aimlapi.com"; +pub(crate) const AIMLAPI_WEB_URL_DEFAULT: &str = "https://aimlapi.com/app"; /// Partner attribution. AI/ML API expects a registered partner id on the /// authorization request; it identifies goose as the integration that brought @@ -35,7 +44,7 @@ const AIMLAPI_WEB_URL_DEFAULT: &str = "https://aimlapi.com"; /// configuration. The environment variable exists to point a build at another /// AI/ML API environment during testing, alongside the two URLs above. const AIMLAPI_PARTNER_ID_ENV: &str = "AIMLAPI_PARTNER_ID"; -const AIMLAPI_PARTNER_ID_DEFAULT: &str = "part_R2KG8QMDBjtWAubVMgG0GF9L"; +pub(crate) const AIMLAPI_PARTNER_ID_DEFAULT: &str = "part_R2KG8QMDBjtWAubVMgG0GF9L"; /// Loopback port the consent screen redirects back to. Chosen from the /// ephemeral range and fixed, because it has to be registered with the diff --git a/crates/goose/src/config/signup_aimlapi/server.rs b/crates/goose/src/config/signup_aimlapi/server.rs index ed46fe04e275..559f8492be58 100644 --- a/crates/goose/src/config/signup_aimlapi/server.rs +++ b/crates/goose/src/config/signup_aimlapi/server.rs @@ -12,8 +12,7 @@ use serde::Deserialize; use std::net::SocketAddr; use tokio::sync::oneshot; -static TEMPLATES_DIR: Dir = - include_dir!("$CARGO_MANIFEST_DIR/src/config/signup_aimlapi/templates"); +static TEMPLATES_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/config/signup_aimlapi/templates"); #[derive(Debug, Deserialize)] struct CallbackQuery { @@ -108,6 +107,7 @@ mod tests { let response = handle_callback( Query(CallbackQuery { code: None, + state: None, error: Some(error.to_string()), }), axum::extract::State(state), diff --git a/crates/goose/src/config/signup_aimlapi/tests.rs b/crates/goose/src/config/signup_aimlapi/tests.rs new file mode 100644 index 000000000000..9974677e5c67 --- /dev/null +++ b/crates/goose/src/config/signup_aimlapi/tests.rs @@ -0,0 +1,74 @@ +use crate::config::signup_aimlapi::{ + PkceAuthFlow, AIMLAPI_APP_URL_DEFAULT, AIMLAPI_PARTNER_ID_DEFAULT, AIMLAPI_WEB_URL_DEFAULT, +}; +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine}; +use sha2::{Digest, Sha256}; + +#[test] +fn challenge_is_the_s256_hash_of_the_verifier() { + let flow = PkceAuthFlow::new().expect("Failed to create PKCE flow"); + + let mut hasher = Sha256::new(); + hasher.update(flow.code_verifier.as_bytes()); + let expected = URL_SAFE_NO_PAD.encode(hasher.finalize()); + + assert_eq!(flow.code_challenge, expected); +} + +#[test] +fn challenge_and_state_are_url_safe_and_unpadded() { + let flow = PkceAuthFlow::new().expect("Failed to create PKCE flow"); + + for value in [&flow.code_challenge, &flow.state] { + assert!(!value.contains('='), "{value} is padded"); + assert!(!value.contains('+'), "{value} is not url-safe"); + assert!(!value.contains('/'), "{value} is not url-safe"); + } +} + +#[test] +fn each_flow_gets_its_own_verifier_state_and_challenge() { + let a = PkceAuthFlow::new().expect("Failed to create PKCE flow 1"); + let b = PkceAuthFlow::new().expect("Failed to create PKCE flow 2"); + + assert_ne!(a.code_verifier, b.code_verifier); + assert_ne!(a.code_challenge, b.code_challenge); + assert_ne!(a.state, b.state); +} + +#[test] +fn consent_base_keeps_the_app_path() { + // The server appends "/agent/authorize" to whatever base it is handed. The + // web app is served under an "/app/" base path, so dropping it produces + // https://aimlapi.com/agent/authorize — a 404 that strands the user on the + // first step of the flow. This has to survive future tidying of the URL. + assert!( + AIMLAPI_WEB_URL_DEFAULT.ends_with("/app"), + "consent base must carry the /app path, got {AIMLAPI_WEB_URL_DEFAULT}" + ); +} + +#[test] +fn api_and_consent_hosts_are_distinct() { + // The registration/exchange calls go to the API host; only the browser is + // sent to the consent host. Collapsing the two would send API calls to the + // marketing site. + assert_ne!(AIMLAPI_APP_URL_DEFAULT, AIMLAPI_WEB_URL_DEFAULT); + assert!(AIMLAPI_APP_URL_DEFAULT.starts_with("https://")); + assert!(AIMLAPI_WEB_URL_DEFAULT.starts_with("https://")); +} + +#[test] +fn partner_id_matches_the_gateway_pattern() { + // The gateway only attributes ids shaped part_; anything else is + // treated as untagged usage and earns nothing. + let id = AIMLAPI_PARTNER_ID_DEFAULT; + + assert!(id.starts_with("part_"), "{id} lacks the part_ prefix"); + let rest = &id["part_".len()..]; + assert!(!rest.is_empty(), "{id} has an empty body"); + assert!( + rest.chars().all(|c| c.is_ascii_alphanumeric()), + "{id} has non-alphanumeric characters after the prefix" + ); +} diff --git a/crates/goose/src/providers/init.rs b/crates/goose/src/providers/init.rs index 9c6ce653029e..a879eb33c387 100644 --- a/crates/goose/src/providers/init.rs +++ b/crates/goose/src/providers/init.rs @@ -8,7 +8,6 @@ use super::local_inference::LocalInferenceProvider; #[cfg(feature = "aws-providers")] use super::sagemaker_tgi::SageMakerTgiProvider; use super::{ - aimlapi::AimlapiProvider, amp_acp::AmpAcpProvider, avian::AvianProvider, azure::AzureProvider, @@ -70,7 +69,6 @@ async fn init_registry() -> RwLock { true, Some(registrations::anthropic_inventory()), ); - registry.register::(false); registry.register::(false); registry.register::(false); registry.register_with_inventory::( diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 305315864a36..9987b8da297f 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -1,5 +1,4 @@ mod acp_tooling; -pub mod aimlapi; pub mod amp_acp; pub mod anthropic { pub use goose_providers::anthropic::*; diff --git a/documentation/docs/getting-started/providers.md b/documentation/docs/getting-started/providers.md index 168090dc3d12..957aca8444c7 100644 --- a/documentation/docs/getting-started/providers.md +++ b/documentation/docs/getting-started/providers.md @@ -22,7 +22,7 @@ goose is compatible with a wide range of LLM providers, allowing you to choose a | Provider | Description | Parameters | |-----------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| [AI/ML API](https://aimlapi.com/) | One API key for 300+ chat, image, video, audio, and embedding models from many providers, OpenAI-compatible. | `AIMLAPI_API_KEY`, `AIMLAPI_HOST` (optional) | +| [AI/ML API](https://aimlapi.com/) | One API key for 300+ chat, image, video, audio, and embedding models from many providers, OpenAI-compatible. | `AIMLAPI_API_KEY` | | [Amazon Bedrock](https://aws.amazon.com/bedrock/) | Offers a variety of foundation models, including Claude, Jurassic-2, and others. **AWS environment variables must be set in advance, not configured through `goose configure`** | Credential auth: `AWS_PROFILE`, or `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_REGION`

Bearer token auth: `AWS_BEARER_TOKEN_BEDROCK` and `AWS_REGION`, `AWS_DEFAULT_REGION`, or `AWS_PROFILE` | | [Amazon SageMaker TGI](https://docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints.html) | Run Text Generation Inference models through Amazon SageMaker endpoints. **AWS credentials must be configured in advance.** | `SAGEMAKER_ENDPOINT_NAME`, `AWS_REGION` (optional), `AWS_PROFILE` (optional) | | [Anthropic](https://www.anthropic.com/) | Offers Claude, an advanced AI model for natural language tasks. | `ANTHROPIC_API_KEY`, `ANTHROPIC_HOST` (optional) |