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/// Serializable provider definition used by runtime and adapter layers.
138#[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    /// How long a stream may go without producing anything before it is
153    /// treated as failed.
154    ///
155    /// This bounds the gap between chunks, not the length of a turn: a
156    /// streamed response can legitimately take minutes, while silence means
157    /// the other end stopped talking. It applies to the SSE transports and to
158    /// the Responses websocket alike. A stream that trips it fails with a
159    /// transport error, which the runtime retries like any other.
160    #[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}