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    /// The wire the OpenAI-compatible ecosystem implements, as distinct from
121    /// OpenAI's own `v1/responses`.
122    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/// Retry configuration for provider transport calls.
138#[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/// Serializable provider definition used by runtime and adapter layers.
160#[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    /// How long a stream may go without producing anything before it is
177    /// treated as failed.
178    ///
179    /// This bounds the gap between chunks, not the length of a turn: a
180    /// streamed response can legitimately take minutes, while silence means
181    /// the other end stopped talking. It applies to the SSE transports and to
182    /// the Responses websocket alike. A stream that trips it fails with a
183    /// transport error, which the runtime retries like any other.
184    #[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}