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}
121
122impl Display for WireApi {
123 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124 let value = match self {
125 Self::Responses => "responses",
126 Self::AnthropicMessages => "anthropic_messages",
127 Self::GeminiGenerateContent => "gemini_generate_content",
128 };
129 f.write_str(value)
130 }
131}
132
133#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
135pub struct RetryPolicy {
136 pub max_attempts: u64,
137 pub base_delay: Duration,
138 pub retry_429: bool,
139 pub retry_5xx: bool,
140 pub retry_transport: bool,
141}
142
143impl Default for RetryPolicy {
144 fn default() -> Self {
145 Self {
146 max_attempts: 5,
147 base_delay: Duration::from_millis(200),
148 retry_429: false,
149 retry_5xx: true,
150 retry_transport: true,
151 }
152 }
153}
154
155#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
157pub struct ProviderDefinition {
158 pub descriptor: ProviderDescriptor,
159 #[serde(default)]
160 pub wire_api: WireApi,
161 #[serde(default)]
162 pub auth_scheme: crate::AuthScheme,
163 #[serde(default)]
164 pub capabilities: ProviderCapabilities,
165 pub base_url: Option<String>,
166 #[serde(default)]
167 pub query_params: Option<HashMap<String, String>>,
168 #[serde(default)]
169 pub headers: Option<HashMap<String, String>>,
170 #[serde(default)]
171 pub retry: RetryPolicy,
172 #[serde(default = "default_stream_idle_timeout")]
173 pub stream_idle_timeout: Duration,
174 #[serde(default = "default_websocket_connect_timeout")]
175 pub websocket_connect_timeout: Duration,
176}
177
178fn default_stream_idle_timeout() -> Duration {
179 Duration::from_millis(300_000)
180}
181
182fn default_websocket_connect_timeout() -> Duration {
183 Duration::from_millis(15_000)
184}
185
186impl ProviderDefinition {
187 pub fn new(id: impl Into<ProviderId>) -> Self {
188 Self {
189 descriptor: ProviderDescriptor::new(id),
190 wire_api: WireApi::default(),
191 auth_scheme: crate::AuthScheme::default(),
192 capabilities: ProviderCapabilities {
193 supports_model_listing: true,
194 supports_streaming: true,
195 supports_websockets: false,
196 supports_tool_calls: true,
197 supports_images: true,
198 supports_history_compaction: false,
199 supports_memory_summarization: false,
200 supports_deferred_tools: false,
201 supports_hosted_tool_search: false,
202 supports_hosted_web_search: false,
203 supports_image_generation: false,
204 supports_reasoning_effort: false,
205 reports_reasoning_tokens: false,
206 reports_thoughts_tokens: false,
207 supports_structured_tool_results: false,
208 supports_embeddings: false,
209 },
210 base_url: None,
211 query_params: None,
212 headers: None,
213 retry: RetryPolicy::default(),
214 stream_idle_timeout: default_stream_idle_timeout(),
215 websocket_connect_timeout: default_websocket_connect_timeout(),
216 }
217 }
218
219 pub fn descriptor(&self) -> ProviderDescriptor {
220 self.descriptor.clone()
221 }
222
223 pub fn provider_id(&self) -> &ProviderId {
224 &self.descriptor.id
225 }
226
227 pub fn url_for_path(&self, path: &str) -> String {
228 let base = self
229 .base_url
230 .as_deref()
231 .unwrap_or_default()
232 .trim_end_matches('/');
233 let path = path.trim_start_matches('/');
234 let mut url = if path.is_empty() {
235 base.to_string()
236 } else {
237 format!("{base}/{path}")
238 };
239
240 if let Some(params) = self
241 .query_params
242 .as_ref()
243 .filter(|params| !params.is_empty())
244 {
245 let qs = params
246 .iter()
247 .map(|(key, value)| format!("{key}={value}"))
248 .collect::<Vec<_>>()
249 .join("&");
250 url.push('?');
251 url.push_str(&qs);
252 }
253
254 url
255 }
256
257 pub fn build_headers(
258 &self,
259 credentials: &crate::ProviderCredentials,
260 ) -> Result<HeaderMap, crate::ProviderError> {
261 let mut headers = HeaderMap::new();
262
263 if let Some(configured_headers) = &self.headers {
264 for (name, value) in configured_headers {
265 insert_header(&mut headers, name, value)?;
266 }
267 }
268
269 for (name, value) in &credentials.headers {
270 insert_header(&mut headers, name, value)?;
271 }
272
273 match &self.auth_scheme {
274 crate::AuthScheme::None | crate::AuthScheme::QueryParam { .. } => {}
275 crate::AuthScheme::BearerToken => {
276 let token = required_auth_value(credentials)?;
277 let auth_value =
278 HeaderValue::from_str(&format!("Bearer {token}")).map_err(|error| {
279 crate::ProviderError::InvalidRequest(format!(
280 "invalid bearer token header: {error}"
281 ))
282 })?;
283 headers.insert(header::AUTHORIZATION, auth_value);
284 }
285 crate::AuthScheme::Header { name } => {
286 let token = required_auth_value(credentials)?;
287 insert_header(&mut headers, name, token)?;
288 }
289 }
290
291 Ok(headers)
292 }
293
294 pub fn build_headers_for_session(
295 &self,
296 credentials: &crate::ProviderCredentials,
297 session: Option<&SessionRequestOptions>,
298 fallback_turn_state: Option<&str>,
299 ) -> Result<HeaderMap, crate::ProviderError> {
300 let mut headers = self.build_headers(credentials)?;
301
302 if let Some(value) = session
303 .and_then(|session| session.sticky_turn_state.as_deref())
304 .or(fallback_turn_state)
305 .and_then(|turn_state| HeaderValue::from_str(turn_state).ok())
306 {
307 headers.insert("x-mentra-turn-state", value.clone());
308 headers.insert("x-codex-turn-state", value);
309 }
310 if let Some(value) = session
311 .and_then(|session| session.turn_metadata.as_deref())
312 .and_then(|value| HeaderValue::from_str(value).ok())
313 {
314 headers.insert("x-mentra-turn-metadata", value.clone());
315 headers.insert("x-codex-turn-metadata", value);
316 }
317 if let Some(value) = session
318 .and_then(|session| session.session_affinity.as_deref())
319 .and_then(|value| HeaderValue::from_str(value).ok())
320 {
321 headers.insert("x-mentra-session-affinity", value);
322 }
323 if let Some(prefer_connection_reuse) =
324 session.and_then(|session| session.prefer_connection_reuse)
325 {
326 headers.insert(
327 "x-mentra-connection-reuse",
328 HeaderValue::from_static(if prefer_connection_reuse {
329 "prefer-reuse"
330 } else {
331 "prefer-fresh"
332 }),
333 );
334 }
335 if let Some(value) = session
336 .and_then(|session| session.subagent.as_deref())
337 .and_then(|value| HeaderValue::from_str(value).ok())
338 {
339 headers.insert("x-openai-subagent", value);
340 }
341 if let Some(extra_headers) = session.map(|session| &session.extra_headers) {
342 for (name, value) in extra_headers {
343 if let (Ok(name), Ok(value)) = (
344 name.parse::<http::HeaderName>(),
345 HeaderValue::from_str(value),
346 ) {
347 headers.insert(name, value);
348 }
349 }
350 }
351
352 Ok(headers)
353 }
354
355 pub fn request_url_with_auth_for_path(
356 &self,
357 path: &str,
358 credentials: &crate::ProviderCredentials,
359 ) -> Result<Url, crate::ProviderError> {
360 let mut url = Url::parse(&self.url_for_path(path))
361 .map_err(|error| crate::ProviderError::InvalidRequest(error.to_string()))?;
362
363 if let crate::AuthScheme::QueryParam { name } = &self.auth_scheme {
364 let token = required_auth_value(credentials)?;
365 url.query_pairs_mut().append_pair(name, token);
366 }
367
368 Ok(url)
369 }
370
371 pub fn websocket_url_for_path(&self, path: &str) -> Result<Url, url::ParseError> {
372 let mut url = Url::parse(&self.url_for_path(path))?;
373
374 let scheme = match url.scheme() {
375 "http" => "ws",
376 "https" => "wss",
377 "ws" | "wss" => return Ok(url),
378 _ => return Ok(url),
379 };
380 let _ = url.set_scheme(scheme);
381 Ok(url)
382 }
383
384 pub fn websocket_url_with_auth_for_path(
385 &self,
386 path: &str,
387 credentials: &crate::ProviderCredentials,
388 ) -> Result<Url, crate::ProviderError> {
389 let mut url = self
390 .websocket_url_for_path(path)
391 .map_err(|error| crate::ProviderError::InvalidRequest(error.to_string()))?;
392
393 if let crate::AuthScheme::QueryParam { name } = &self.auth_scheme {
394 let token = required_auth_value(credentials)?;
395 url.query_pairs_mut().append_pair(name, token);
396 }
397
398 Ok(url)
399 }
400}
401
402fn insert_header(
403 headers: &mut HeaderMap,
404 name: &str,
405 value: &str,
406) -> Result<(), crate::ProviderError> {
407 let header_name = HeaderName::from_bytes(name.as_bytes()).map_err(|error| {
408 crate::ProviderError::InvalidRequest(format!(
409 "invalid provider header name {name:?}: {error}"
410 ))
411 })?;
412 let header_value = HeaderValue::from_str(value).map_err(|error| {
413 crate::ProviderError::InvalidRequest(format!(
414 "invalid provider header value for {name:?}: {error}"
415 ))
416 })?;
417 headers.insert(header_name, header_value);
418 Ok(())
419}
420
421fn required_auth_value(
422 credentials: &crate::ProviderCredentials,
423) -> Result<&str, crate::ProviderError> {
424 credentials.bearer_token.as_deref().ok_or_else(|| {
425 crate::ProviderError::InvalidRequest("missing provider auth credential".to_string())
426 })
427}
428
429#[cfg(test)]
430mod tests {
431 use super::*;
432
433 #[test]
434 fn build_headers_applies_bearer_auth_and_static_headers() {
435 let mut definition = ProviderDefinition::new("test");
436 definition.auth_scheme = crate::AuthScheme::BearerToken;
437 definition.headers = Some(HashMap::from([(
438 "x-provider-header".to_string(),
439 "static".to_string(),
440 )]));
441
442 let headers = definition
443 .build_headers(&crate::ProviderCredentials {
444 bearer_token: Some("secret".to_string()),
445 account_id: None,
446 headers: HashMap::from([("x-runtime-header".to_string(), "dynamic".to_string())]),
447 })
448 .expect("headers should build");
449
450 assert_eq!(headers["x-provider-header"], "static");
451 assert_eq!(headers["x-runtime-header"], "dynamic");
452 assert_eq!(headers[header::AUTHORIZATION], "Bearer secret");
453 }
454
455 #[test]
456 fn request_url_with_auth_appends_query_param_auth() {
457 let mut definition = ProviderDefinition::new("test");
458 definition.base_url = Some("https://example.com/v1".to_string());
459 definition.query_params = Some(HashMap::from([(
460 "api-version".to_string(),
461 "2026".to_string(),
462 )]));
463 definition.auth_scheme = crate::AuthScheme::QueryParam {
464 name: "api-key".to_string(),
465 };
466
467 let url = definition
468 .request_url_with_auth_for_path(
469 "responses",
470 &crate::ProviderCredentials {
471 bearer_token: Some("secret".to_string()),
472 account_id: None,
473 headers: HashMap::new(),
474 },
475 )
476 .expect("url should build");
477
478 assert_eq!(
479 url.as_str(),
480 "https://example.com/v1/responses?api-version=2026&api-key=secret"
481 );
482 }
483}