diff --git a/README.md b/README.md index ca7084de..23d7bee6 100644 --- a/README.md +++ b/README.md @@ -404,6 +404,7 @@ Kit used the fewest tokens, took the least active time, and needed the least ste - [Agent Plugins](docs/user/agent-plugins.md) - [TUI and sessions](docs/user/tui-and-sessions.md) - [Provider usage](docs/user/provider-usage.md) +- [Workload identity for inference](docs/user/workload-inference.md) - [Security, limits, and troubleshooting](docs/user/security-limits-and-troubleshooting.md) - [Releasing](docs/releasing.md) ยท [Third-party notices](THIRD_PARTY_NOTICES.md) diff --git a/docs/user/getting-started-and-configuration.md b/docs/user/getting-started-and-configuration.md index 97840857..86bc8b6a 100644 --- a/docs/user/getting-started-and-configuration.md +++ b/docs/user/getting-started-and-configuration.md @@ -114,6 +114,12 @@ kit prompt --provider openrouter \ Kit uses the CLI or TOML `model`; `OPENROUTER_MODEL` does not override that selection. The adapter also accepts `OPENROUTER_BASE_URL`, `OPENROUTER_APP_NAME`, `OPENROUTER_SITE_URL`, `OPENROUTER_MAX_COMPLETION_TOKENS`, `OPENROUTER_TEMPERATURE`, and `OPENROUTER_REASONING_EFFORT`. Its model-catalog lookup for the selected model's context length is best-effort, so a catalog failure does not by itself prevent normal provider usage. +### Workload identity for inference + +For deployed workloads with projected identity tokens, use +[workload identity for inference](workload-inference.md) instead of local API-key +or browser authentication. + ### TypeSafe API key TypeSafe authentication is BYOK (bring your own key). Create a key at diff --git a/docs/user/workload-inference.md b/docs/user/workload-inference.md new file mode 100644 index 00000000..d00505f1 --- /dev/null +++ b/docs/user/workload-inference.md @@ -0,0 +1,73 @@ +# Workload identity for inference + +Use workload identity for deployed workloads that receive projected identity tokens. +For interactive local use, follow [provider authentication](getting-started-and-configuration.md#choose-a-provider-and-authenticate) instead. + +For `openrouter` or `speakeasy`, configure a provider-specific workload exchange +in `~/.kit/config.toml`. Kit reads a deployment-projected JWT file; it does not +mint the external assertion. The deployment must rotate that file before renewal. +These settings override API-key and browser credentials for that provider only. +An exchange failure stops inference: Kit never falls back to another credential. + +OpenRouter uses [workload identity federation](https://openrouter.ai/docs/guides/overview/auth/workload-identity-federation) +with an organization federation policy targeting an inference-capable key: + +```toml +provider = "openrouter" +model = "anthropic/claude-sonnet-4" + +[inference_providers.openrouter] +endpoint = "https://openrouter.ai/api/v1/chat/completions" +[inference_providers.openrouter.workload] +tokenFile = "/var/run/identity/assertion.jwt" +endpointMode = "explicit" +grant = "tokenExchange" +tokenEndpoint = "https://openrouter.ai/api/v1/oauth/token" +scope = "inference" +[inference_providers.openrouter.workload.extensions] +federation_policy_id = "" +``` + +Speakeasy uses its shared workload issuer's JWT-bearer token endpoint. Configure +issuer trust and project authorization on the server first, then use the actual +token endpoint returned for that issuer: + +```toml +provider = "speakeasy" +model = "anthropic/claude-sonnet-4" + +[inference_providers.speakeasy] +endpoint = "https://ai.speakeasy.com/chat/completions" +[inference_providers.speakeasy.workload] +tokenFile = "/var/run/identity/assertion.jwt" +endpointMode = "explicit" +grant = "jwtBearer" +tokenEndpoint = "https://ai.speakeasy.com/oauth/usi//token" +resource = "urn:speakeasy:inference" +[inference_providers.speakeasy.workload.extensions] +project_id = "" +``` + +The Speakeasy assertion's audience is the shared issuer URL or its token endpoint, +not `urn:speakeasy:inference`. Assertions with a `jti` (or `uti`) are single-use: a new +exchange needs a fresh assertion. This is ordinary workload inference, not a +replacement for managed assistant credentials or chat-bound authorization. Kit +sends bearer authentication without API-key, project, or assistant/chat headers. + +Both examples pin HTTPS token and completion endpoints explicitly; redirects are +disabled. Optional `audience`, `resource`, `scope`, and string-valued `extensions` +belong to the workload exchange, not completion headers. `endpointMode = "discovery"` +instead requires `expectedIssuer`, same-origin issuer metadata, and an HTTPS +`resource`; use explicit mode for Speakeasy's inference URN. Keep this configuration +operator-controlled: the endpoints determine where credentials are sent. + +Kit shares in-memory tokens across sessions for each configured provider, caches +only until the returned `expires_in` minus a safety margin, and rereads the source +file on renewal. An RFC 8693 response without `expires_in` is used once and not +cached. A token may outlive its source assertion when the server permits it. +Tokens and assertions are not saved in Kit's credential store. Workload requests +are not automatically replayed on authentication refusal, throttling, or server +errors; API-key retry behavior is unchanged. Select models explicitly: workload +grants do not imply model-catalog access, and Kit does not discover a context limit +through these grants. Other features with separate credentials (such as images or +evaluation) are not authenticated by this configuration. diff --git a/src/main.rs b/src/main.rs index 83f0d845..815e8189 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1382,6 +1382,8 @@ struct Config { root: Option, model: Option, provider: Option, + #[serde(default)] + inference_providers: BTreeMap, reasoning_effort: Option, a2a: Option, capture_error_spans: Option, @@ -2305,6 +2307,10 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { openrouter_api_key.as_ref().map(|(key, _)| key.clone()), )?, }; + let runtime = kit::Runtime::with_inference_providers( + runtime, + config.inference_providers.clone(), + )?; let runtime = kit::Runtime::with_compaction(runtime, config.compaction)?; let runtime = kit::Runtime::with_eval(runtime, config.experimental.eval)?; let runtime = kit::Runtime::with_ask_user(runtime, config.ask_user.unwrap_or(true))?; @@ -2406,6 +2412,10 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { openrouter_api_key.as_ref().map(|(key, _)| key.clone()), )?, }; + let runtime = kit::Runtime::with_inference_providers( + runtime, + config.inference_providers.clone(), + )?; let runtime = kit::Runtime::with_compaction(runtime, config.compaction)?; let runtime = kit::Runtime::with_eval(runtime, config.experimental.eval)?; let runtime = kit::Runtime::with_ask_user(runtime, config.ask_user.unwrap_or(true))?; @@ -2482,6 +2492,10 @@ async fn run_cli(cli: Cli) -> Result<(), Box> { reasoning_effort, openrouter_api_key.as_ref().map(|(key, _)| key.clone()), )?; + let runtime = kit::Runtime::with_inference_providers( + runtime, + config.inference_providers.clone(), + )?; let runtime = kit::Runtime::with_compaction(runtime, config.compaction)?; let runtime = kit::Runtime::with_eval(runtime, config.experimental.eval)?; let runtime = kit::Runtime::with_ask_user(runtime, config.ask_user.unwrap_or(true))?; @@ -4604,6 +4618,26 @@ future_option = true let _rebound = tokio::net::TcpListener::bind(address).await.unwrap(); } + #[test] + fn inference_provider_configuration_parses_documented_examples() { + let docs = include_str!("../docs/user/workload-inference.md"); + let mut count = 0; + for block in docs.split("```toml\n").skip(1) { + let example = block.split("```").next().unwrap(); + if !example.contains("[inference_providers.") { + continue; + } + let config: Config = toml::from_str(example).unwrap(); + assert_eq!(config.inference_providers.len(), 1); + let provider = config.provider.unwrap(); + let inference = config.inference_providers.get(provider.as_str()).unwrap(); + inference.workload.validate().unwrap(); + assert!(inference.endpoint.ends_with("/chat/completions")); + count += 1; + } + assert_eq!(count, 2); + } + #[test] fn subagent_model_policies_are_scoped_and_validated_per_harness() { let configured: Config = toml::from_str( diff --git a/src/provider/adapter.rs b/src/provider/adapter.rs index a94b4fcc..e1994350 100644 --- a/src/provider/adapter.rs +++ b/src/provider/adapter.rs @@ -207,6 +207,7 @@ pub struct SelectableAdapter { credential_storage: crate::credentials::CredentialStorage, openrouter_api_key: Option, openai_model_catalog: SubscriptionModelCatalogCache, + inference: super::inference::Instances, } impl SelectableAdapter { @@ -258,9 +259,18 @@ impl SelectableAdapter { credential_storage, openrouter_api_key, openai_model_catalog: SubscriptionModelCatalogCache::default(), + inference: Default::default(), }) } + pub(crate) fn with_inference(mut self, inference: super::inference::Instances) -> Self { + self.inference = inference; + self + } + pub(crate) fn inference(&self) -> super::inference::Instances { + self.inference.clone() + } + /// Attach persistence before publishing any clones of this session adapter. pub(crate) fn with_session_observer( mut self, @@ -305,6 +315,7 @@ impl SelectableAdapter { credential_storage: self.credential_storage.clone(), openrouter_api_key: self.openrouter_api_key.clone(), openai_model_catalog: self.openai_model_catalog.clone(), + inference: self.inference.clone(), }, revision, )) @@ -348,22 +359,47 @@ impl SelectableAdapter { } pub async fn model_catalog(&self, current: &ModelSelection) -> Vec { - model_catalog_with_openai( + let mut groups = model_catalog_with_openai( current, self.discovered_openai_models(), &self.credential_storage, self.openrouter_api_key.as_ref(), + &self.inference, ) - .await + .await; + // A workload grant does not imply permission to query a model catalog. + // Keep explicit model selection available without resolving legacy keys. + for provider in [ProviderKind::OpenRouter, ProviderKind::Speakeasy] { + if self.inference.get(provider).is_some() { + groups.retain(|group| group.provider != provider); + groups.push(ModelGroup { + provider, + models: if current.provider == provider { + vec![current.model.clone()] + } else { + Vec::new() + }, + context_windows: Default::default(), + }); + } + } + groups } /// Providers whose credentials resolve now, in catalog order. pub fn catalog_providers(&self) -> Vec { - available_catalog_providers( + let mut providers = available_catalog_providers( &self.credential_storage, self.openrouter_api_key.as_ref(), |name| std::env::var(name), - ) + ); + providers.retain(|provider| self.inference.get(*provider).is_none()); + for provider in [ProviderKind::OpenRouter, ProviderKind::Speakeasy] { + if self.inference.get(provider).is_some() { + providers.push(provider); + } + } + providers } pub fn select(&self, selection: ModelSelection) -> Result<(), String> { @@ -371,12 +407,14 @@ impl SelectableAdapter { return Err("model name is outside canonical bounds".into()); } let reasoning_effort = self.reasoning_effort()?; - KitAdapter::new_with_credentials_and_effort( + KitAdapter::new_with_inference( selection.provider, selection.model.clone(), self.credential_storage.clone(), reasoning_effort, self.openrouter_api_key.as_ref(), + self.openai_model_catalog.clone(), + &self.inference, )?; let mut current = self .selection @@ -392,12 +430,14 @@ impl SelectableAdapter { reasoning_effort: Option, ) -> Result<(), String> { let model = self.selection()?; - KitAdapter::new_with_credentials_and_effort( + KitAdapter::new_with_inference( model.provider, model.model, self.credential_storage.clone(), reasoning_effort, self.openrouter_api_key.as_ref(), + self.openai_model_catalog.clone(), + &self.inference, )?; let mut current = self .selection @@ -421,6 +461,7 @@ pub struct SelectableSession { credential_storage: crate::credentials::CredentialStorage, openrouter_api_key: Option, openai_model_catalog: SubscriptionModelCatalogCache, + inference: super::inference::Instances, config: SessionConfig, active: SessionSelection, inner: KitSession, @@ -437,13 +478,14 @@ impl ModelAdapter for SelectableAdapter { .lock() .map(|value| value.clone()) .map_err(|_| LoopError::InvalidState("session selection lock is poisoned".into()))?; - let inner = KitAdapter::new_with_credentials_effort_and_catalog( + let inner = KitAdapter::new_with_inference( active.model.provider, active.model.model.clone(), self.credential_storage.clone(), active.reasoning_effort, self.openrouter_api_key.as_ref(), self.openai_model_catalog.clone(), + &self.inference, ) .map_err(LoopError::InvalidState)? .start_session(config.clone()) @@ -453,6 +495,7 @@ impl ModelAdapter for SelectableAdapter { credential_storage: self.credential_storage.clone(), openrouter_api_key: self.openrouter_api_key.clone(), openai_model_catalog: self.openai_model_catalog.clone(), + inference: self.inference.clone(), config, active, inner, @@ -508,13 +551,14 @@ impl ModelSession for SelectableSession { .map(|value| value.clone()) .map_err(|_| LoopError::InvalidState("session selection lock is poisoned".into()))?; if selected != self.active { - let replacement = KitAdapter::new_with_credentials_effort_and_catalog( + let replacement = KitAdapter::new_with_inference( selected.model.provider, selected.model.model.clone(), self.credential_storage.clone(), selected.reasoning_effort, self.openrouter_api_key.as_ref(), self.openai_model_catalog.clone(), + &self.inference, ) .map_err(LoopError::InvalidState)? .start_session(self.config.clone()) @@ -587,6 +631,8 @@ struct SpeakeasyProvider { api_key: String, project: String, chat_id: Option, + federation: bool, + name: ProviderKind, } // Matches Gram's chat.SessionIDToChatID mapping for captured agent sessions. @@ -600,7 +646,19 @@ impl CompletionsProvider for SpeakeasyProvider { type Config = OpenRouterRequestConfig; fn provider_name(&self) -> &str { - "Speakeasy" + if self.federation { + self.name.as_str() + } else { + "Speakeasy" + } + } + + fn authentication(&self) -> Option { + if self.federation { + self.openrouter.authentication() + } else { + None + } } fn endpoint_url(&self) -> &str { @@ -615,6 +673,14 @@ impl CompletionsProvider for SpeakeasyProvider { &self, builder: agentkit_http::HttpRequestBuilder, ) -> agentkit_http::HttpRequestBuilder { + if self.federation { + let builder = self.openrouter.preprocess_request(builder); + return if self.name == ProviderKind::Speakeasy { + builder.header("X-Gram-Source", "kit") + } else { + builder + }; + } let builder = builder .header("Gram-Key", &self.api_key) .header("Gram-Project", &self.project) @@ -654,6 +720,49 @@ impl CompletionsProvider for SpeakeasyProvider { } impl KitAdapter { + #[allow(clippy::too_many_arguments)] + fn new_with_inference( + provider: ProviderKind, + model: String, + credential_storage: crate::credentials::CredentialStorage, + effort: Option, + key: Option<&OpenRouterApiKey>, + catalog: SubscriptionModelCatalogCache, + instances: &super::inference::Instances, + ) -> Result { + let Some(instance) = instances.get(provider) else { + return Self::new_with_credentials_effort_and_catalog( + provider, + model, + credential_storage, + effort, + key, + catalog, + ); + }; + let mut config = OpenRouterConfig::new(instance.authentication.clone(), model) + .with_base_url(&instance.endpoint); + apply_openrouter_reasoning_effort(&mut config, effort); + let provider = SpeakeasyProvider { + openrouter: OpenRouterProvider::from(config), + api_key: String::new(), + project: String::new(), + chat_id: None, + federation: true, + name: provider, + }; + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .connect_timeout(Duration::from_secs(10)) + .retry(reqwest::retry::never()) + .build() + .map_err(|_| "could not build inference client".to_string())?; + Ok(Self::Speakeasy(Box::new(SpeakeasyKitAdapter { + provider, + client: agentkit_http::Http::new(super::inference::Transport(client)), + }))) + } + pub fn new(provider: ProviderKind, model: String) -> Result { Self::new_with_credentials(provider, model, Default::default()) } @@ -752,6 +861,8 @@ impl KitAdapter { api_key: credentials.api_key.clone(), project: credentials.project.clone(), chat_id: None, + federation: false, + name: ProviderKind::Speakeasy, }; let client = reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none()) @@ -934,13 +1045,19 @@ impl ModelAdapter for KitAdapter { } Self::Speakeasy(adapter) => { let mut provider = adapter.provider.clone(); - provider.chat_id = Some(gram_chat_id(&config.session_id.to_string())); + let resilience = if provider.federation { + super::inference::resilience(crate::request_budget::RequestBudget::current()) + } else { + provider.chat_id = Some(gram_chat_id(&config.session_id.to_string())); + crate::request_budget::RequestBudget::current().resilience() + }; let inner = CompletionsAdapter::with_client(provider, adapter.client.clone()) - .with_resilience(crate::request_budget::RequestBudget::current().resilience()); + .with_resilience(resilience); inner.start_session(config).await.map(|inner| { KitSession::Speakeasy(SpeakeasyKitSession { inner, context_window: None, + federation: adapter.provider.federation, }) }) } @@ -951,7 +1068,7 @@ impl ModelAdapter for KitAdapter { Some(match self { Self::OpenAiSubscription(_) => "openai-subscription", Self::OpenRouter(_) => "openrouter", - Self::Speakeasy(_) => "speakeasy", + Self::Speakeasy(adapter) => adapter.provider.name.as_str(), Self::Anthropic(_) => "anthropic", }) } @@ -972,6 +1089,7 @@ pub struct OpenRouterKitSession { pub struct SpeakeasyKitSession { inner: CompletionsSession, context_window: Option, + federation: bool, } /// Project only the outbound request; the caller's canonical transcript stays typed. @@ -1255,6 +1373,7 @@ impl ModelSession for KitSession { context_window: session.context_window, media_part: None, next_media: 0, + federation: false, }) .map(KitTurn::OpenRouter), Self::Anthropic(session, context_window) => { @@ -1265,17 +1384,23 @@ impl ModelSession for KitSession { .await .map(|turn| KitTurn::Anthropic(Box::new(turn), context_window)) } - Self::Speakeasy(session) => session - .inner - .begin_turn(request, cancellation) - .await - .map(|inner| OpenRouterKitTurn { - inner, - context_window: session.context_window, - media_part: None, - next_media: 0, + Self::Speakeasy(session) => { + let future = session.inner.begin_turn(request, cancellation); + let result = if session.federation { + super::inference::request(future).await + } else { + future.await + }; + result.map(|inner| { + KitTurn::Speakeasy(OpenRouterKitTurn { + inner, + context_window: session.context_window, + media_part: None, + next_media: 0, + federation: session.federation, + }) }) - .map(KitTurn::Speakeasy), + } } } @@ -1310,6 +1435,7 @@ pub struct OpenRouterKitTurn { context_window: Option, media_part: Option, next_media: usize, + federation: bool, } #[async_trait] @@ -1336,7 +1462,13 @@ impl ModelTurn for KitTurn { Ok(event) } Self::OpenRouter(turn) | Self::Speakeasy(turn) => { - let mut event = turn.inner.next_event(cancellation).await?; + let mut event = turn.inner.next_event(cancellation).await.map_err(|error| { + if turn.federation { + super::inference::redact_response_error(error) + } else { + error + } + })?; if let Some(ModelTurnEvent::Delta(delta)) = &mut event { rewrite_openrouter_media(delta, &mut turn.media_part, &mut turn.next_media); } @@ -1526,6 +1658,7 @@ pub async fn model_catalog(current: &ModelSelection) -> Vec { std::future::ready(DiscoveredModels::default()), &Default::default(), None, + &Default::default(), ) .await } @@ -1566,9 +1699,13 @@ async fn model_catalog_with_openai( openai_catalog: impl std::future::Future, storage: &crate::credentials::CredentialStorage, openrouter_api_key: Option<&OpenRouterApiKey>, + inference: &super::inference::Instances, ) -> Vec { - let available = - available_catalog_providers(storage, openrouter_api_key, |name| std::env::var(name)); + let available: Vec<_> = + available_catalog_providers(storage, openrouter_api_key, |name| std::env::var(name)) + .into_iter() + .filter(|provider| inference.get(*provider).is_none()) + .collect(); let openrouter_available = available.contains(&ProviderKind::OpenRouter); let speakeasy_available = available.contains(&ProviderKind::Speakeasy); let openai_catalog = async { @@ -2655,6 +2792,8 @@ mod tests { api_key: format!("gram_live_{}", "ab".repeat(32)), project: "kit-test".into(), chat_id: None, + federation: false, + name: ProviderKind::Speakeasy, }; let client = reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none()) @@ -2738,6 +2877,7 @@ mod tests { credential_storage: Default::default(), openrouter_api_key: None, openai_model_catalog: Default::default(), + inference: Default::default(), config: SessionConfig::new("provider-identity-test"), active, inner, @@ -3038,6 +3178,7 @@ mod tests { }, &storage, None, + &Default::default(), ) .await; let openai = catalog @@ -3069,6 +3210,7 @@ mod tests { async { panic!("uncredentialed discovery polled") }, &storage, None, + &Default::default(), ) .await; assert!( @@ -3146,3 +3288,7 @@ mod image_tests; #[cfg(test)] #[path = "adapter_background_image_tests.rs"] mod background_image_tests; + +#[cfg(test)] +#[path = "inference_tests.rs"] +mod inference_tests; diff --git a/src/provider/adapter_image_tests.rs b/src/provider/adapter_image_tests.rs index 7dd9b1f6..8c64c6f3 100644 --- a/src/provider/adapter_image_tests.rs +++ b/src/provider/adapter_image_tests.rs @@ -212,6 +212,8 @@ async fn completions_and_openrouter_encode_parallel_mixed_replay_as_user_images( api_key: "test".into(), project: "test".into(), chat_id: None, + federation: false, + name: ProviderKind::Speakeasy, }; assert_wire(speakeasy, projected.clone()).await; assert_wire( diff --git a/src/provider/inference.rs b/src/provider/inference.rs new file mode 100644 index 00000000..afa5fd47 --- /dev/null +++ b/src/provider/inference.rs @@ -0,0 +1,179 @@ +//! Instance-scoped inference authentication without ambient credential fallback. +use agentkit_http::{Authentication, AuthenticationAttempt, AuthenticationProvider, HttpError}; +use async_trait::async_trait; +use serde::Deserialize; +use std::{collections::BTreeMap, sync::Arc}; + +#[derive(Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Config { + pub endpoint: String, + pub workload: crate::workload_auth::Config, +} +impl std::fmt::Debug for Config { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("InferenceConfig").finish_non_exhaustive() + } +} +#[derive(Clone, Default)] +pub(crate) struct Instances(Arc>); +#[derive(Clone)] +pub(crate) struct Instance { + pub endpoint: String, + pub authentication: Authentication, +} +impl Instances { + pub(crate) fn new(configs: BTreeMap) -> Result { + let mut instances = BTreeMap::new(); + for (name, config) in configs { + if !matches!(name.as_str(), "openrouter" | "speakeasy") { + return Err("workload inference requires a supported completion provider".into()); + } + let url = + url::Url::parse(&config.endpoint).map_err(|_| "invalid inference endpoint")?; + if url.scheme() != "https" + || url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.fragment().is_some() + || url.query().is_some() + { + return Err( + "inference endpoint must be HTTPS without credentials, query or fragment" + .into(), + ); + } + let session = Arc::new(crate::workload_auth::Session::new(config.workload)?); + instances.insert( + name, + Instance { + endpoint: config.endpoint, + authentication: Authentication::new(WorkloadAuthentication(session)), + }, + ); + } + Ok(Self(Arc::new(instances))) + } + pub(crate) fn get(&self, provider: super::ProviderKind) -> Option<&Instance> { + self.0.get(provider.as_str()) + } +} +struct WorkloadAuthentication(Arc); +#[async_trait] +impl AuthenticationProvider for WorkloadAuthentication { + async fn authenticate( + &self, + previous: Option<&AuthenticationAttempt>, + ) -> Result { + // A 401 never replays a potentially billed inference request. + if previous.is_some() { + return Err(HttpError::Other( + "workload inference authentication refused; request not replayed".into(), + )); + } + let token = self.0.acquisition().await.map_err(|error| { + let _ = REQUEST + .try_with(|state| state.borrow_mut().authentication_error = Some(error.clone())); + HttpError::Other(error) + })?; + let _ = REQUEST.try_with(|state| { + state.borrow_mut().acquisition = Some((self.0.clone(), token.clone())); + }); + Authentication::bearer(token.token()) + .authenticate(None) + .await + } +} + +/// Never surface remote refusal bodies or credential-bearing transport details. +pub(crate) struct Transport(pub reqwest::Client); +#[async_trait] +impl agentkit_http::HttpClient for Transport { + async fn execute( + &self, + request: agentkit_http::HttpRequest, + ) -> Result { + let response = agentkit_http::HttpClient::execute(&self.0, request) + .await + .map_err(|_| HttpError::Other("workload inference transport failed".into()))?; + if matches!(response.status().as_u16(), 401 | 403) { + let rejected = REQUEST + .try_with(|state| state.borrow().acquisition.clone()) + .ok() + .flatten(); + if let Some((session, acquisition)) = rejected { + session.invalidate_acquisition(&acquisition).await; + } + } + if !response.status().is_success() { + return Err(HttpError::Other(format!( + "workload inference refused (HTTP {})", + response.status().as_u16() + ))); + } + Ok(response) + } +} + +// Authentication and transport run in the same begin-turn future. Task-scoped +// state preserves the exact acquisition and trusted local error without a token +// lookup, cross-session mutable state, or exposing credentials in request headers. +#[derive(Default)] +struct RequestState { + acquisition: Option<( + Arc, + crate::workload_auth::Acquisition, + )>, + authentication_error: Option, +} +tokio::task_local! { + static REQUEST: std::cell::RefCell; +} + +pub(crate) async fn request( + future: impl std::future::Future>, +) -> Result { + REQUEST + .scope(Default::default(), async { + future.await.map_err(|error| { + match REQUEST.with(|state| state.borrow_mut().authentication_error.take()) { + Some(message) => agentkit_loop::LoopError::Provider(message), + None => redact_response_error(error), + } + }) + }) + .await +} + +pub(crate) fn resilience( + budget: crate::request_budget::RequestBudget, +) -> agentkit_http::ResilienceConfig { + let mut config = budget.resilience(); + config.max_retries = 0; + config +} + +// HTTP-200 error frames and malformed nonstream responses can echo credentials. +// Typed failures are already sanitized upstream; raw provider errors are not. +pub(crate) fn redact_response_error(error: agentkit_loop::LoopError) -> agentkit_loop::LoopError { + match error { + agentkit_loop::LoopError::Provider(_) => { + agentkit_loop::LoopError::Provider("workload inference response failed".into()) + } + other => other, + } +} + +#[cfg(test)] +pub(crate) mod test_support { + pub(crate) fn authentication( + config: crate::workload_auth::Config, + client: reqwest::Client, + ) -> Result { + Ok(agentkit_http::Authentication::new( + super::WorkloadAuthentication(std::sync::Arc::new( + crate::workload_auth::test_support::session(config, client)?, + )), + )) + } +} diff --git a/src/provider/inference_tests.rs b/src/provider/inference_tests.rs new file mode 100644 index 00000000..ba5abc9d --- /dev/null +++ b/src/provider/inference_tests.rs @@ -0,0 +1,748 @@ +//! Real TLS exchanges against a scripted localhost authorization server. +#![allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::disallowed_methods, + clippy::disallowed_macros +)] + +use super::*; +use crate::workload_auth::{Config, Grant}; +use agentkit_core::{Item, ItemKind, MetadataMap, SessionId, TurnId}; +use std::path::PathBuf; +use std::{collections::BTreeMap, sync::Arc}; +use tokio::io::AsyncReadExt; +use tokio::{io::AsyncWriteExt, net::TcpListener, sync::mpsc, task::JoinHandle}; +use tokio_rustls::{TlsAcceptor, rustls}; + +const CA: &[u8] = include_bytes!("../../fixtures/workload-auth/ca.pem"); +const CERT: &[u8] = include_bytes!("../../fixtures/workload-auth/server.der"); +const KEY: &[u8] = include_bytes!("../../fixtures/workload-auth/server-key.der"); +const ASSERTION: &str = "private-header.private-payload.private-signature"; + +struct Reply { + status: u16, + body: String, + content_type: &'static str, +} +impl Reply { + fn json(body: serde_json::Value) -> Self { + Self { + status: 200, + content_type: "application/json", + body: body.to_string(), + } + } + fn token(value: &str, expiry: u64) -> Self { + Self::json(serde_json::json!({ + "access_token": value, "token_type": "Bearer", "expires_in": expiry, + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + "refresh_token": "ignored-refresh-secret" + })) + } + fn status(status: u16) -> Self { + Self { + status, + content_type: "application/json", + body: format!("server echoed {ASSERTION} access-secret"), + } + } +} + +#[derive(Debug)] +struct Request { + method: String, + path: String, + headers: String, + body: String, +} +impl Request { + fn form(&self) -> BTreeMap { + url::form_urlencoded::parse(self.body.as_bytes()) + .into_owned() + .collect() + } +} + +struct Server { + origin: String, + replies: mpsc::UnboundedSender, + requests: mpsc::UnboundedReceiver, + task: JoinHandle<()>, +} +impl Drop for Server { + fn drop(&mut self) { + self.task.abort(); + } +} +impl Server { + async fn start() -> Self { + let provider = rustls::crypto::aws_lc_rs::default_provider(); + let config = rustls::ServerConfig::builder_with_provider(Arc::new(provider)) + .with_safe_default_protocol_versions() + .unwrap() + .with_no_client_auth() + .with_single_cert( + vec![rustls::pki_types::CertificateDer::from(CERT.to_vec())], + rustls::pki_types::PrivatePkcs8KeyDer::from(KEY.to_vec()).into(), + ) + .unwrap(); + let acceptor = TlsAcceptor::from(Arc::new(config)); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let origin = format!("https://{}", listener.local_addr().unwrap()); + let (replies, mut responses) = mpsc::unbounded_channel::(); + let (records, requests) = mpsc::unbounded_channel(); + let task = tokio::spawn(async move { + loop { + let (socket, _) = listener.accept().await.unwrap(); + let mut stream = acceptor.accept(socket).await.unwrap(); + let mut data = Vec::new(); + let (header_end, length) = loop { + let mut chunk = [0; 4096]; + let n = stream.read(&mut chunk).await.unwrap(); + assert!(n > 0, "request ended before headers"); + data.extend_from_slice(&chunk[..n]); + assert!(data.len() < 32768); + if let Some(end) = data.windows(4).position(|w| w == b"\r\n\r\n") { + let headers = std::str::from_utf8(&data[..end]).unwrap(); + let length = headers + .lines() + .find_map(|line| { + let (key, value) = line.split_once(':')?; + key.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().unwrap()) + }) + .unwrap_or(0); + break (end + 4, length); + } + }; + while data.len() < header_end + length { + let mut chunk = [0; 4096]; + let n = stream.read(&mut chunk).await.unwrap(); + assert!(n > 0); + data.extend_from_slice(&chunk[..n]); + } + let headers = String::from_utf8(data[..header_end].to_vec()).unwrap(); + let mut line = headers.lines().next().unwrap().split_whitespace(); + let request = Request { + method: line.next().unwrap().into(), + path: line.next().unwrap().into(), + body: String::from_utf8(data[header_end..header_end + length].to_vec()) + .unwrap(), + headers, + }; + records.send(request).unwrap(); + // An unexpected retry receives a refusal, so a broken client fails + // promptly rather than hanging while waiting for a scripted reply. + let reply = responses.try_recv().unwrap_or_else(|_| Reply::status(500)); + let response = format!( + "HTTP/1.1 {} Test\r\nContent-Type: {}\r\nLocation: /redirect-target\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + reply.status, + reply.content_type, + reply.body.len(), + reply.body + ); + stream.write_all(response.as_bytes()).await.unwrap(); + let _ = stream.shutdown().await; + } + }); + Self { + origin, + replies, + requests, + task, + } + } + fn enqueue(&self, reply: Reply) { + self.replies.send(reply).unwrap(); + } + fn discovery(&self) { + self.enqueue(Reply::json(serde_json::json!({ + "issuer": format!("{}/tenant", self.origin), + "token_endpoint": format!("{}/token", self.origin) + }))); + } + fn exchange(&self, token: &str, expiry: u64) { + self.discovery(); + self.enqueue(Reply::token(token, expiry)); + } + fn client(&self) -> reqwest::Client { + reqwest::Client::builder() + .tls_certs_only([reqwest::Certificate::from_pem(CA).unwrap()]) + .redirect(reqwest::redirect::Policy::none()) + .retry(reqwest::retry::never()) + .no_proxy() + .timeout(Duration::from_secs(5)) + .build() + .unwrap() + } + fn adapter( + &self, + path: PathBuf, + resource: &str, + name: ProviderKind, + streaming: bool, + ) -> KitAdapter { + let authentication = crate::provider::inference::test_support::authentication( + Config { + token_file: path, + expected_issuer: format!("{}/tenant", self.origin), + token_endpoint: format!("{}/token", self.origin), + resource: resource.into(), + grant: Grant::TokenExchange, + ..Default::default() + }, + self.client(), + ) + .unwrap(); + let config = OpenRouterConfig::new(authentication, "test-model") + .with_base_url(format!("{}/chat/completions", self.origin)) + .with_streaming(streaming); + KitAdapter::Speakeasy(Box::new(SpeakeasyKitAdapter { + provider: SpeakeasyProvider { + openrouter: OpenRouterProvider::from(config), + // Deliberately populate legacy values: federation must ignore them. + api_key: "legacy-key-secret".into(), + project: "legacy-project-secret".into(), + chat_id: Some("legacy-chat-secret".into()), + federation: true, + name, + }, + client: agentkit_http::Http::new(crate::provider::inference::Transport(self.client())), + })) + } + fn drain(&mut self) -> Vec { + let mut requests = Vec::new(); + while let Ok(request) = self.requests.try_recv() { + requests.push(request); + } + requests + } +} + +fn source() -> (tempfile::TempDir, PathBuf) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("assertion"); + std::fs::write(&path, format!("{ASSERTION}\n")).unwrap(); + (dir, path) +} + +fn completion(streaming: bool) -> Reply { + if streaming { + Reply { + status: 200, + content_type: "text/event-stream", + body: concat!( + "data: {\"id\":\"test\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"test-model\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"ok\"},\"finish_reason\":null}]}\n\n", + "data: {\"id\":\"test\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"test-model\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + "data: [DONE]\n\n" + ).into(), + } + } else { + Reply::json(serde_json::json!({ + "id": "test", "object": "chat.completion", "created": 1, + "model": "test-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}] + })) + } +} + +async fn run_turn(adapter: &KitAdapter) -> Result, LoopError> { + let mut session = adapter + .start_session(SessionConfig::new("federation-contract")) + .await?; + let mut turn = session + .begin_turn( + TurnRequest { + session_id: SessionId::new("federation-contract"), + turn_id: TurnId::new("turn"), + transcript: vec![Item::text(ItemKind::User, "hello")], + available_tools: Vec::new(), + cache: None, + metadata: MetadataMap::new(), + }, + None, + ) + .await?; + let mut events = Vec::new(); + while let Some(event) = turn.next_event(None).await? { + events.push(event); + } + Ok(events) +} + +fn assert_exchange(request: &Request, resource: &str) { + assert_eq!(request.method, "POST"); + assert_eq!(request.path, "/token"); + let form = request.form(); + assert_eq!( + form, + BTreeMap::from([ + ( + "grant_type".into(), + "urn:ietf:params:oauth:grant-type:token-exchange".into() + ), + ("subject_token".into(), ASSERTION.into()), + ( + "subject_token_type".into(), + "urn:ietf:params:oauth:token-type:jwt".into() + ), + ( + "requested_token_type".into(), + "urn:ietf:params:oauth:token-type:access_token".into() + ), + ("resource".into(), resource.into()), + ]) + ); + let headers = request.headers.to_ascii_lowercase(); + assert!(headers.contains("application/x-www-form-urlencoded")); + assert!(!headers.contains("authorization:")); +} + +fn assert_inference(request: &Request, token: &str, streaming: bool) { + assert_eq!(request.method, "POST"); + assert_eq!(request.path, "/chat/completions"); + let headers = request.headers.to_ascii_lowercase(); + assert_eq!( + headers + .lines() + .filter(|line| line.starts_with("user-agent:")) + .count(), + 1 + ); + assert!(headers.contains(&format!("authorization: bearer {token}\r\n"))); + for forbidden in [ + "gram-key:", + "gram-project:", + "gram-chat-id:", + "legacy-", + ASSERTION, + ] { + assert!( + !headers.contains(forbidden), + "unexpected header: {forbidden}" + ); + } + let body: Value = serde_json::from_str(&request.body).unwrap(); + assert_eq!(body["model"], "test-model"); + assert_eq!(body["messages"][0]["content"], "hello"); + assert_eq!( + body.get("stream").and_then(Value::as_bool).unwrap_or(false), + streaming + ); + assert!(!request.body.contains(ASSERTION)); + assert!(!request.body.contains(token)); +} + +#[tokio::test] +async fn federation_completion_stream_and_nonstream_use_exchanged_bearer_only() { + for name in [ProviderKind::Speakeasy, ProviderKind::OpenRouter] { + for streaming in [false, true] { + let mut server = Server::start().await; + let (_dir, path) = source(); + let resource = "https://inference.example/resource"; + let adapter = server.adapter(path, resource, name, streaming); + server.exchange("inference-secret", 3600); + server.enqueue(completion(streaming)); + let events = run_turn(&adapter).await.unwrap(); + assert!( + events + .iter() + .any(|event| matches!(event, ModelTurnEvent::Finished(_))) + ); + let requests = server.drain(); + assert_eq!(requests.len(), 3); + assert_eq!(requests[0].method, "GET"); + assert_exchange(&requests[1], resource); + assert_inference(&requests[2], "inference-secret", streaming); + assert_eq!( + requests[2] + .headers + .to_ascii_lowercase() + .contains("x-gram-source: kit"), + name == ProviderKind::Speakeasy + ); + } + } +} + +#[tokio::test] +async fn federation_refusals_are_redacted_and_never_replayed() { + for streaming in [false, true] { + for status in [401, 429, 503] { + let mut server = Server::start().await; + let (_dir, path) = source(); + let adapter = server.adapter( + path, + "https://inference.example/resource", + ProviderKind::Speakeasy, + streaming, + ); + server.exchange("access-secret", 3600); + server.enqueue(Reply::status(status)); + let error = match run_turn(&adapter).await { + Ok(_) => panic!("HTTP {status} unexpectedly succeeded"), + Err(error) => format!("{error:?}"), + }; + assert!( + error.contains("workload inference response failed"), + "{error}" + ); + for secret in [ + ASSERTION, + "access-secret", + "server echoed", + "legacy-key-secret", + ] { + assert!(!error.contains(secret), "credential leaked in refusal"); + } + let requests = server.drain(); + assert_eq!( + requests.len(), + 3, + "HTTP {status} must not replay inference or exchange" + ); + assert_inference(&requests[2], "access-secret", streaming); + } + } +} + +#[tokio::test] +async fn federation_providers_keep_independent_tokens_and_cached_grants() { + let mut server = Server::start().await; + let (_dir, path) = source(); + let first = server.adapter( + path.clone(), + "https://inference.example/first", + ProviderKind::Speakeasy, + false, + ); + let second = server.adapter( + path, + "https://inference.example/second", + ProviderKind::OpenRouter, + false, + ); + for (adapter, token, resource) in [ + (&first, "first-secret", "https://inference.example/first"), + (&second, "second-secret", "https://inference.example/second"), + ] { + server.exchange(token, 3600); + server.enqueue(completion(false)); + run_turn(adapter).await.unwrap(); + let requests = server.drain(); + assert_eq!(requests.len(), 3); + assert_exchange(&requests[1], resource); + assert_inference(&requests[2], token, false); + } + for (adapter, token) in [(&first, "first-secret"), (&second, "second-secret")] { + server.enqueue(completion(false)); + run_turn(adapter).await.unwrap(); + let requests = server.drain(); + assert_eq!( + requests.len(), + 1, + "provider must reuse only its own cached grant" + ); + assert_inference(&requests[0], token, false); + } +} + +#[tokio::test] +async fn federation_success_status_error_frames_are_redacted() { + for (streaming, body) in [ + ( + true, + "event: error\ndata: access-secret private-header.private-payload.private-signature\n\n", + ), + ( + true, + "data: {\"error\":{\"message\":\"access-secret\"}}\n\n", + ), + ( + true, + "event: error\ndata: {\"message\":\"access-secret\"}\n\n", + ), + (false, "{\"error\":{\"message\":\"access-secret\"}}"), + (false, "access-secret malformed response"), + ] { + let mut server = Server::start().await; + let (_dir, path) = source(); + let adapter = server.adapter( + path, + "https://inference.example/resource", + ProviderKind::OpenRouter, + streaming, + ); + server.exchange("access-secret", 3600); + server.enqueue(Reply { + status: 200, + content_type: if streaming { + "text/event-stream" + } else { + "application/json" + }, + body: body.into(), + }); + let error = match run_turn(&adapter).await { + Ok(_) => panic!("error frame unexpectedly succeeded"), + Err(error) => format!("{error:?}"), + }; + assert!( + error.contains("workload inference response failed"), + "{error}" + ); + assert!(!error.contains("access-secret")); + assert!(!error.contains(ASSERTION)); + assert_eq!(server.drain().len(), 3); + } +} + +#[tokio::test] +async fn federation_speakeasy_explicit_jwt_bearer_contract_and_no_auth_replay() { + let mut server = Server::start().await; + let (_dir, path) = source(); + let authentication = crate::provider::inference::test_support::authentication( + Config { + token_file: path, + token_endpoint: format!("{}/token", server.origin), + endpoint_mode: crate::workload_auth::EndpointMode::Explicit, + resource: "urn:speakeasy:inference".into(), + extensions: BTreeMap::from([( + "project_id".into(), + "00000000-0000-0000-0000-000000000001".into(), + )]), + ..Default::default() + }, + server.client(), + ) + .unwrap(); + let previous = agentkit_http::Authentication::bearer("old-secret") + .authenticate(None) + .await + .unwrap(); + let error = authentication + .authenticate(Some(&previous)) + .await + .err() + .unwrap() + .to_string(); + assert!(error.contains("not replayed")); + assert!(!error.contains("old-secret")); + assert!(server.drain().is_empty()); + let config = OpenRouterConfig::new(authentication, "test-model") + .with_base_url(format!("{}/chat/completions", server.origin)) + .with_streaming(false); + let adapter = KitAdapter::Speakeasy(Box::new(SpeakeasyKitAdapter { + provider: SpeakeasyProvider { + openrouter: OpenRouterProvider::from(config), + api_key: String::new(), + project: String::new(), + chat_id: None, + federation: true, + name: ProviderKind::Speakeasy, + }, + client: agentkit_http::Http::new(crate::provider::inference::Transport(server.client())), + })); + server.enqueue(Reply::json(serde_json::json!({ + "access_token": "speakeasy-secret", "token_type": "Bearer", + "expires_in": 900, "authorization_expires_in": 900 + }))); + server.enqueue(completion(false)); + run_turn(&adapter).await.unwrap(); + let requests = server.drain(); + assert_eq!(requests.len(), 2); + assert_eq!( + requests[0].form(), + BTreeMap::from([ + ( + "grant_type".into(), + "urn:ietf:params:oauth:grant-type:jwt-bearer".into() + ), + ("assertion".into(), ASSERTION.into()), + ("resource".into(), "urn:speakeasy:inference".into()), + ( + "project_id".into(), + "00000000-0000-0000-0000-000000000001".into() + ), + ]) + ); + assert_inference(&requests[1], "speakeasy-secret", false); +} + +fn workload_config(endpoint: &str) -> crate::provider::inference::Config { + crate::provider::inference::Config { + endpoint: endpoint.into(), + workload: Config { + token_file: PathBuf::from("/missing/projected-assertion"), + token_endpoint: "https://issuer.example/token".into(), + endpoint_mode: crate::workload_auth::EndpointMode::Explicit, + resource: "urn:speakeasy:inference".into(), + ..Default::default() + }, + } +} + +#[test] +fn federation_configuration_is_validated_redacted_and_provider_scoped() { + use crate::provider::inference::Instances; + for endpoint in [ + "http://inference.example/chat/completions", + "https://secret@inference.example/chat/completions", + "https://inference.example/chat/completions?token=secret", + "https://inference.example/chat/completions#secret", + ] { + let error = Instances::new(BTreeMap::from([( + "speakeasy".into(), + workload_config(endpoint), + )])) + .err() + .unwrap(); + assert!(!error.contains("secret")); + } + assert!( + Instances::new(BTreeMap::from([( + "anthropic".into(), + workload_config("https://inference.example") + )])) + .is_err() + ); + let config = workload_config("https://inference.example/chat/completions"); + let debug = format!("{config:?}"); + assert!(!debug.contains("inference.example")); + assert!(!debug.contains("projected-assertion")); + let instances = Instances::new(BTreeMap::from([("speakeasy".into(), config)])).unwrap(); + assert!(instances.get(ProviderKind::OpenRouter).is_none()); + assert!(instances.get(ProviderKind::Speakeasy).is_some()); + let adapter = SelectableAdapter::new(ProviderKind::Speakeasy, "test-model") + .unwrap() + .with_inference(instances); + assert!( + adapter + .catalog_providers() + .contains(&ProviderKind::Speakeasy) + ); + adapter + .select(ModelSelection::new( + ProviderKind::Speakeasy, + "another-model", + )) + .unwrap(); + adapter + .select_reasoning_effort(Some(ReasoningEffort::Low)) + .unwrap(); +} + +#[tokio::test] +async fn federation_missing_assertion_fails_without_legacy_credentials() { + let instances = crate::provider::inference::Instances::new(BTreeMap::from([( + "speakeasy".into(), + workload_config("https://inference.example/chat/completions"), + )])) + .unwrap(); + let adapter = KitAdapter::new_with_inference( + ProviderKind::Speakeasy, + "test-model".into(), + Default::default(), + None, + None, + Default::default(), + &instances, + ) + .unwrap(); + let error = match run_turn(&adapter).await { + Ok(_) => panic!("missing assertion unexpectedly succeeded"), + Err(error) => error.to_string(), + }; + assert!( + error.contains("could not open workload token file"), + "{error}" + ); + assert!(!error.contains("projected-assertion")); + assert!(!error.contains("login")); +} + +#[test] +fn federation_preserves_configured_deadlines_without_retries() { + for seconds in [1, 600] { + let budget = crate::request_budget::RequestBudget::try_from(seconds).unwrap(); + let mut expected = budget.resilience(); + expected.max_retries = 0; + assert_eq!(crate::provider::inference::resilience(budget), expected); + } +} + +#[tokio::test] +async fn federation_refusal_invalidates_only_for_a_new_turn() { + for status in [401, 403] { + let mut server = Server::start().await; + let (_dir, path) = source(); + let adapter = server.adapter( + path.clone(), + "https://inference.example/resource", + ProviderKind::Speakeasy, + false, + ); + server.exchange("refused-secret", 3600); + server.enqueue(Reply::status(status)); + assert!(run_turn(&adapter).await.is_err()); + let requests = server.drain(); + assert_eq!(requests.len(), 3, "refused completion must not replay"); + assert_inference(&requests[2], "refused-secret", false); + std::fs::write(path, "rotated.payload.signature").unwrap(); + server.exchange("renewed-secret", 3600); + server.enqueue(completion(false)); + run_turn(&adapter).await.unwrap(); + let requests = server.drain(); + assert_eq!( + requests.len(), + 3, + "new turn must exchange the rotated source" + ); + assert_eq!( + requests[1].form()["subject_token"], + "rotated.payload.signature" + ); + assert_inference(&requests[2], "renewed-secret", false); + } +} + +#[tokio::test] +async fn federation_and_legacy_catalogs_keep_the_same_provider_order() { + let dir = tempfile::tempdir().unwrap(); + let storage = crate::credentials::CredentialStorage::Filesystem(dir.path().join("credentials")); + storage + .entry("speakeasy", "default") + .save( + &serde_json::to_vec(&serde_json::json!({ + "api_key": format!("gram_live_{}", "ab".repeat(32)), "project": "test-project", + "email": null, "organization_id": null + })) + .unwrap(), + ) + .unwrap(); + let instances = crate::provider::inference::Instances::new(BTreeMap::from([( + "openrouter".into(), + workload_config("https://inference.example/chat/completions"), + )])) + .unwrap(); + let adapter = SelectableAdapter::new_with_credentials_effort_and_openrouter_key( + ProviderKind::OpenRouter, + "test-model", + storage, + None, + Some(OpenRouterApiKey::non_empty("legacy-test-key").unwrap()), + ) + .unwrap() + .with_inference(instances); + let current = ModelSelection::new(ProviderKind::OpenRouter, "test-model"); + let groups = adapter.model_catalog(¤t).await; + assert_eq!( + groups + .iter() + .map(|group| group.provider) + .collect::>(), + adapter.catalog_providers() + ); +} diff --git a/src/provider/mod.rs b/src/provider/mod.rs index 522f1fd6..3b9e14c7 100644 --- a/src/provider/mod.rs +++ b/src/provider/mod.rs @@ -1,6 +1,7 @@ mod adapter; pub(crate) mod anthropic_auth; mod anthropic_catalog; +pub mod inference; #[doc(hidden)] pub use anthropic_auth::{AnthropicAuthCommand, execute_anthropic_auth}; diff --git a/src/runtime.rs b/src/runtime.rs index aab2d5e1..094b0f7c 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -748,6 +748,17 @@ impl Runtime { Ok(Arc::new(runtime)) } + /// Configures workload inference before publishing the runtime. + pub fn with_inference_providers( + mut runtime: Arc, + configs: std::collections::BTreeMap, + ) -> Result, String> { + let instances = crate::provider::inference::Instances::new(configs)?; + let inner = Arc::get_mut(&mut runtime).ok_or("runtime already shared")?; + inner.adapter = inner.adapter.clone().with_inference(instances); + Ok(runtime) + } + /// Configures compaction before publishing the runtime to any sessions. pub fn with_compaction( runtime: Arc, @@ -865,7 +876,8 @@ impl Runtime { } pub(crate) fn supports_terminal_authentication(&self, provider: ProviderKind) -> bool { - self.credential_storage.is_persistent() + self.adapter.inference().get(provider).is_none() + && self.credential_storage.is_persistent() && (provider != ProviderKind::OpenRouter || (self.openrouter_api_key.is_none() && !self.ambient_openrouter_api_key)) } @@ -1520,6 +1532,7 @@ impl Runtime { .unwrap_or(self.adapter.reasoning_effort()?), self.openrouter_api_key.clone(), )? + .with_inference(self.adapter.inference()) .with_session_observer(opened.observer.clone(), !request.resume)?; let pending_creation = (!request.resume).then(|| opened.observer.clone()); let skills = self.fresh_skills(); @@ -1939,7 +1952,9 @@ impl Runtime { self.openrouter_api_key.clone(), ) .and_then(|adapter| { - adapter.with_session_observer(opened.observer.clone(), is_fork || !request.resume) + adapter + .with_inference(self.adapter.inference()) + .with_session_observer(opened.observer.clone(), is_fork || !request.resume) }) .map_err(AcpRuntimeError::Loop)?; let current_skills = self.current_skills().await.map_err(AcpRuntimeError::Loop)?;