Skip to main content

mentra_provider/
definition.rs

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/// Builtin provider families Mentra can construct from presets.
18#[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/// Stable identifier for a registered provider implementation.
36#[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/// Human-facing metadata about a provider.
74#[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/// Capabilities advertised by a provider instance.
92#[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/// Wire protocol supported by a provider.
113#[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/// Retry configuration for provider transport calls.
134#[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/// Serializable provider definition used by runtime and adapter layers.
156#[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}