1use http::HeaderMap;
2use http::HeaderName;
3use http::HeaderValue;
4use http::header;
5use serde::Deserialize;
6use serde::Serialize;
7use std::borrow::Cow;
8use std::collections::HashMap;
9use std::fmt::Display;
10use std::time::Duration;
11use strum::Display as StrumDisplay;
12use strum::IntoStaticStr;
13use url::Url;
14
15use crate::request::SessionRequestOptions;
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, StrumDisplay, IntoStaticStr)]
19#[strum(serialize_all = "lowercase")]
20pub enum BuiltinProvider {
21 Anthropic,
22 Gemini,
23 OpenAI,
24 OpenRouter,
25 Ollama,
26 LmStudio,
27}
28
29impl From<BuiltinProvider> for ProviderId {
30 fn from(value: BuiltinProvider) -> Self {
31 Self(Cow::Borrowed(value.into()))
32 }
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord)]
37pub struct ProviderId(Cow<'static, str>);
38
39impl ProviderId {
40 pub fn new(id: impl Into<String>) -> Self {
41 Self(Cow::Owned(id.into()))
42 }
43
44 pub fn as_str(&self) -> &str {
45 self.0.as_ref()
46 }
47}
48
49impl Display for ProviderId {
50 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51 f.write_str(self.as_str())
52 }
53}
54
55impl From<&str> for ProviderId {
56 fn from(value: &str) -> Self {
57 Self::new(value)
58 }
59}
60
61impl From<String> for ProviderId {
62 fn from(value: String) -> Self {
63 Self(Cow::Owned(value))
64 }
65}
66
67impl From<&String> for ProviderId {
68 fn from(value: &String) -> Self {
69 Self::new(value.as_str())
70 }
71}
72
73#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
75pub struct ProviderDescriptor {
76 pub id: ProviderId,
77 pub display_name: Option<String>,
78 pub description: Option<String>,
79}
80
81impl ProviderDescriptor {
82 pub fn new(id: impl Into<ProviderId>) -> Self {
83 Self {
84 id: id.into(),
85 display_name: None,
86 description: None,
87 }
88 }
89}
90
91#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
93pub struct ProviderCapabilities {
94 pub supports_model_listing: bool,
95 pub supports_streaming: bool,
96 pub supports_websockets: bool,
97 pub supports_tool_calls: bool,
98 pub supports_images: bool,
99 pub supports_history_compaction: bool,
100 pub supports_memory_summarization: bool,
101 pub supports_deferred_tools: bool,
102 pub supports_hosted_tool_search: bool,
103 pub supports_hosted_web_search: bool,
104 pub supports_image_generation: bool,
105 pub supports_reasoning_effort: bool,
106 pub reports_reasoning_tokens: bool,
107 pub reports_thoughts_tokens: bool,
108 pub supports_structured_tool_results: bool,
109 pub supports_embeddings: bool,
110}
111
112#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
114#[serde(rename_all = "lowercase")]
115pub enum WireApi {
116 #[default]
117 Responses,
118 AnthropicMessages,
119 GeminiGenerateContent,
120 OpenAiChatCompletions,
123}
124
125impl Display for WireApi {
126 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
127 let value = match self {
128 Self::Responses => "responses",
129 Self::AnthropicMessages => "anthropic_messages",
130 Self::GeminiGenerateContent => "gemini_generate_content",
131 Self::OpenAiChatCompletions => "openai_chat_completions",
132 };
133 f.write_str(value)
134 }
135}
136
137#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
139pub struct RetryPolicy {
140 pub max_attempts: u64,
141 pub base_delay: Duration,
142 pub retry_429: bool,
143 pub retry_5xx: bool,
144 pub retry_transport: bool,
145}
146
147impl Default for RetryPolicy {
148 fn default() -> Self {
149 Self {
150 max_attempts: 5,
151 base_delay: Duration::from_millis(200),
152 retry_429: false,
153 retry_5xx: true,
154 retry_transport: true,
155 }
156 }
157}
158
159#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
161pub struct ProviderDefinition {
162 pub descriptor: ProviderDescriptor,
163 #[serde(default)]
164 pub wire_api: WireApi,
165 #[serde(default)]
166 pub auth_scheme: crate::AuthScheme,
167 #[serde(default)]
168 pub capabilities: ProviderCapabilities,
169 pub base_url: Option<String>,
170 #[serde(default)]
171 pub query_params: Option<HashMap<String, String>>,
172 #[serde(default)]
173 pub headers: Option<HashMap<String, String>>,
174 #[serde(default)]
175 pub retry: RetryPolicy,
176 #[serde(default = "default_stream_idle_timeout")]
185 pub stream_idle_timeout: Duration,
186 #[serde(default = "default_websocket_connect_timeout")]
187 pub websocket_connect_timeout: Duration,
188}
189
190fn default_stream_idle_timeout() -> Duration {
191 Duration::from_millis(300_000)
192}
193
194fn default_websocket_connect_timeout() -> Duration {
195 Duration::from_millis(15_000)
196}
197
198impl ProviderDefinition {
199 pub fn new(id: impl Into<ProviderId>) -> Self {
200 Self {
201 descriptor: ProviderDescriptor::new(id),
202 wire_api: WireApi::default(),
203 auth_scheme: crate::AuthScheme::default(),
204 capabilities: ProviderCapabilities {
205 supports_model_listing: true,
206 supports_streaming: true,
207 supports_websockets: false,
208 supports_tool_calls: true,
209 supports_images: true,
210 supports_history_compaction: false,
211 supports_memory_summarization: false,
212 supports_deferred_tools: false,
213 supports_hosted_tool_search: false,
214 supports_hosted_web_search: false,
215 supports_image_generation: false,
216 supports_reasoning_effort: false,
217 reports_reasoning_tokens: false,
218 reports_thoughts_tokens: false,
219 supports_structured_tool_results: false,
220 supports_embeddings: false,
221 },
222 base_url: None,
223 query_params: None,
224 headers: None,
225 retry: RetryPolicy::default(),
226 stream_idle_timeout: default_stream_idle_timeout(),
227 websocket_connect_timeout: default_websocket_connect_timeout(),
228 }
229 }
230
231 pub fn descriptor(&self) -> ProviderDescriptor {
232 self.descriptor.clone()
233 }
234
235 pub fn provider_id(&self) -> &ProviderId {
236 &self.descriptor.id
237 }
238
239 pub fn url_for_path(&self, path: &str) -> String {
240 let base = self
241 .base_url
242 .as_deref()
243 .unwrap_or_default()
244 .trim_end_matches('/');
245 let path = path.trim_start_matches('/');
246 let mut url = if path.is_empty() {
247 base.to_string()
248 } else {
249 format!("{base}/{path}")
250 };
251
252 if let Some(params) = self
253 .query_params
254 .as_ref()
255 .filter(|params| !params.is_empty())
256 {
257 let qs = params
258 .iter()
259 .map(|(key, value)| format!("{key}={value}"))
260 .collect::<Vec<_>>()
261 .join("&");
262 url.push('?');
263 url.push_str(&qs);
264 }
265
266 url
267 }
268
269 pub fn build_headers(
270 &self,
271 credentials: &crate::ProviderCredentials,
272 ) -> Result<HeaderMap, crate::ProviderError> {
273 let mut headers = HeaderMap::new();
274
275 if let Some(configured_headers) = &self.headers {
276 for (name, value) in configured_headers {
277 insert_header(&mut headers, name, value)?;
278 }
279 }
280
281 for (name, value) in &credentials.headers {
282 insert_header(&mut headers, name, value)?;
283 }
284
285 match &self.auth_scheme {
286 crate::AuthScheme::None | crate::AuthScheme::QueryParam { .. } => {}
287 crate::AuthScheme::BearerToken => {
288 let token = required_auth_value(credentials)?;
289 let auth_value =
290 HeaderValue::from_str(&format!("Bearer {token}")).map_err(|error| {
291 crate::ProviderError::InvalidRequest(format!(
292 "invalid bearer token header: {error}"
293 ))
294 })?;
295 headers.insert(header::AUTHORIZATION, auth_value);
296 }
297 crate::AuthScheme::Header { name } => {
298 let token = required_auth_value(credentials)?;
299 insert_header(&mut headers, name, token)?;
300 }
301 }
302
303 Ok(headers)
304 }
305
306 pub fn build_headers_for_session(
307 &self,
308 credentials: &crate::ProviderCredentials,
309 session: Option<&SessionRequestOptions>,
310 fallback_turn_state: Option<&str>,
311 ) -> Result<HeaderMap, crate::ProviderError> {
312 let mut headers = self.build_headers(credentials)?;
313
314 if let Some(value) = session
315 .and_then(|session| session.sticky_turn_state.as_deref())
316 .or(fallback_turn_state)
317 .and_then(|turn_state| HeaderValue::from_str(turn_state).ok())
318 {
319 headers.insert("x-mentra-turn-state", value.clone());
320 headers.insert("x-codex-turn-state", value);
321 }
322 if let Some(value) = session
323 .and_then(|session| session.turn_metadata.as_deref())
324 .and_then(|value| HeaderValue::from_str(value).ok())
325 {
326 headers.insert("x-mentra-turn-metadata", value.clone());
327 headers.insert("x-codex-turn-metadata", value);
328 }
329 if let Some(value) = session
330 .and_then(|session| session.session_affinity.as_deref())
331 .and_then(|value| HeaderValue::from_str(value).ok())
332 {
333 headers.insert("x-mentra-session-affinity", value);
334 }
335 if let Some(prefer_connection_reuse) =
336 session.and_then(|session| session.prefer_connection_reuse)
337 {
338 headers.insert(
339 "x-mentra-connection-reuse",
340 HeaderValue::from_static(if prefer_connection_reuse {
341 "prefer-reuse"
342 } else {
343 "prefer-fresh"
344 }),
345 );
346 }
347 if let Some(value) = session
348 .and_then(|session| session.subagent.as_deref())
349 .and_then(|value| HeaderValue::from_str(value).ok())
350 {
351 headers.insert("x-openai-subagent", value);
352 }
353 if let Some(extra_headers) = session.map(|session| &session.extra_headers) {
354 for (name, value) in extra_headers {
355 if let (Ok(name), Ok(value)) = (
356 name.parse::<http::HeaderName>(),
357 HeaderValue::from_str(value),
358 ) {
359 headers.insert(name, value);
360 }
361 }
362 }
363
364 Ok(headers)
365 }
366
367 pub fn request_url_with_auth_for_path(
368 &self,
369 path: &str,
370 credentials: &crate::ProviderCredentials,
371 ) -> Result<Url, crate::ProviderError> {
372 let mut url = Url::parse(&self.url_for_path(path))
373 .map_err(|error| crate::ProviderError::InvalidRequest(error.to_string()))?;
374
375 if let crate::AuthScheme::QueryParam { name } = &self.auth_scheme {
376 let token = required_auth_value(credentials)?;
377 url.query_pairs_mut().append_pair(name, token);
378 }
379
380 Ok(url)
381 }
382
383 pub fn websocket_url_for_path(&self, path: &str) -> Result<Url, url::ParseError> {
384 let mut url = Url::parse(&self.url_for_path(path))?;
385
386 let scheme = match url.scheme() {
387 "http" => "ws",
388 "https" => "wss",
389 "ws" | "wss" => return Ok(url),
390 _ => return Ok(url),
391 };
392 let _ = url.set_scheme(scheme);
393 Ok(url)
394 }
395
396 pub fn websocket_url_with_auth_for_path(
397 &self,
398 path: &str,
399 credentials: &crate::ProviderCredentials,
400 ) -> Result<Url, crate::ProviderError> {
401 let mut url = self
402 .websocket_url_for_path(path)
403 .map_err(|error| crate::ProviderError::InvalidRequest(error.to_string()))?;
404
405 if let crate::AuthScheme::QueryParam { name } = &self.auth_scheme {
406 let token = required_auth_value(credentials)?;
407 url.query_pairs_mut().append_pair(name, token);
408 }
409
410 Ok(url)
411 }
412}
413
414fn insert_header(
415 headers: &mut HeaderMap,
416 name: &str,
417 value: &str,
418) -> Result<(), crate::ProviderError> {
419 let header_name = HeaderName::from_bytes(name.as_bytes()).map_err(|error| {
420 crate::ProviderError::InvalidRequest(format!(
421 "invalid provider header name {name:?}: {error}"
422 ))
423 })?;
424 let header_value = HeaderValue::from_str(value).map_err(|error| {
425 crate::ProviderError::InvalidRequest(format!(
426 "invalid provider header value for {name:?}: {error}"
427 ))
428 })?;
429 headers.insert(header_name, header_value);
430 Ok(())
431}
432
433fn required_auth_value(
434 credentials: &crate::ProviderCredentials,
435) -> Result<&str, crate::ProviderError> {
436 credentials.bearer_token.as_deref().ok_or_else(|| {
437 crate::ProviderError::InvalidRequest("missing provider auth credential".to_string())
438 })
439}
440
441#[cfg(test)]
442mod tests {
443 use super::*;
444
445 #[test]
446 fn build_headers_applies_bearer_auth_and_static_headers() {
447 let mut definition = ProviderDefinition::new("test");
448 definition.auth_scheme = crate::AuthScheme::BearerToken;
449 definition.headers = Some(HashMap::from([(
450 "x-provider-header".to_string(),
451 "static".to_string(),
452 )]));
453
454 let headers = definition
455 .build_headers(&crate::ProviderCredentials {
456 bearer_token: Some("secret".to_string()),
457 account_id: None,
458 headers: HashMap::from([("x-runtime-header".to_string(), "dynamic".to_string())]),
459 })
460 .expect("headers should build");
461
462 assert_eq!(headers["x-provider-header"], "static");
463 assert_eq!(headers["x-runtime-header"], "dynamic");
464 assert_eq!(headers[header::AUTHORIZATION], "Bearer secret");
465 }
466
467 #[test]
468 fn request_url_with_auth_appends_query_param_auth() {
469 let mut definition = ProviderDefinition::new("test");
470 definition.base_url = Some("https://example.com/v1".to_string());
471 definition.query_params = Some(HashMap::from([(
472 "api-version".to_string(),
473 "2026".to_string(),
474 )]));
475 definition.auth_scheme = crate::AuthScheme::QueryParam {
476 name: "api-key".to_string(),
477 };
478
479 let url = definition
480 .request_url_with_auth_for_path(
481 "responses",
482 &crate::ProviderCredentials {
483 bearer_token: Some("secret".to_string()),
484 account_id: None,
485 headers: HashMap::new(),
486 },
487 )
488 .expect("url should build");
489
490 assert_eq!(
491 url.as_str(),
492 "https://example.com/v1/responses?api-version=2026&api-key=secret"
493 );
494 }
495}