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