Add operator Mentra memory browser

vendor/mentra/tests/responses_transport.rs

@@ -0,0 +1,254 @@
+//! Which transport a runtime's Responses requests go out on, and what happens
+//! when a provider cannot serve the one it was handed.
+
+use std::{
+    fs,
+    sync::{
+        Arc, Mutex,
+        atomic::{AtomicU64, Ordering},
+    },
+    time::{SystemTime, UNIX_EPOCH},
+};
+
+use async_trait::async_trait;
+use mentra::{
+    AgentConfig, BuiltinProvider, ContentBlock, ProviderCapabilities, Role, Runtime,
+    provider::{
+        ContentBlockDelta, ContentBlockStart, ModelInfo, Provider, ProviderDescriptor,
+        ProviderError, ProviderEvent, ProviderEventStream, ProviderId, ProviderRequestOptions,
+        Request, ResponsesTransport,
+    },
+    runtime::SqliteRuntimeStore,
+};
+use tokio::sync::mpsc;
+
+static NEXT_TEMP_ID: AtomicU64 = AtomicU64::new(1);
+
+/// Answers one short text turn and remembers the options it was handed.
+///
+/// The options are what this file is about: the transport is settled before the
+/// request leaves the runtime, so what arrives here is the only evidence of
+/// which one was chosen.
+#[derive(Clone)]
+struct RecordingProvider {
+    id: ProviderId,
+    display_name: Option<String>,
+    capabilities: ProviderCapabilities,
+    seen: Arc<Mutex<Vec<ProviderRequestOptions>>>,
+}
+
+impl RecordingProvider {
+    fn new(id: impl Into<ProviderId>, supports_websockets: bool) -> Self {
+        Self {
+            id: id.into(),
+            display_name: None,
+            capabilities: ProviderCapabilities {
+                supports_streaming: true,
+                supports_websockets,
+                ..ProviderCapabilities::default()
+            },
+            seen: Arc::new(Mutex::new(Vec::new())),
+        }
+    }
+
+    fn named(mut self, display_name: &str) -> Self {
+        self.display_name = Some(display_name.to_string());
+        self
+    }
+
+    fn transports(&self) -> Vec<ResponsesTransport> {
+        self.seen
+            .lock()
+            .expect("recorded options")
+            .iter()
+            .map(|options| options.responses.transport)
+            .collect()
+    }
+}
+
+#[async_trait]
+impl Provider for RecordingProvider {
+    fn descriptor(&self) -> ProviderDescriptor {
+        ProviderDescriptor {
+            id: self.id.clone(),
+            display_name: self.display_name.clone(),
+            description: None,
+        }
+    }
+
+    fn capabilities(&self) -> ProviderCapabilities {
+        self.capabilities
+    }
+
+    async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
+        Ok(vec![ModelInfo::new("model", self.id.clone())])
+    }
+
+    async fn stream(&self, request: Request<'_>) -> Result<ProviderEventStream, ProviderError> {
+        self.seen
+            .lock()
+            .expect("recorded options")
+            .push(request.provider_request_options.clone());
+
+        let (tx, rx) = mpsc::unbounded_channel();
+        for event in [
+            ProviderEvent::MessageStarted {
+                id: "msg-1".to_string(),
+                model: "model".to_string(),
+                role: Role::Assistant,
+            },
+            ProviderEvent::ContentBlockStarted {
+                index: 0,
+                kind: ContentBlockStart::Text,
+            },
+            ProviderEvent::ContentBlockDelta {
+                index: 0,
+                delta: ContentBlockDelta::Text("ok".to_string()),
+            },
+            ProviderEvent::ContentBlockStopped { index: 0 },
+            ProviderEvent::MessageStopped,
+        ] {
+            tx.send(Ok(event)).expect("test receiver alive");
+        }
+        Ok(rx)
+    }
+}
+
+fn temp_store() -> SqliteRuntimeStore {
+    let unique = NEXT_TEMP_ID.fetch_add(1, Ordering::Relaxed);
+    let timestamp = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .expect("system time")
+        .as_nanos();
+    let path = std::env::temp_dir().join(format!(
+        "mentra-responses-transport-{timestamp}-{unique}.sqlite"
+    ));
+    if let Some(parent) = path.parent() {
+        fs::create_dir_all(parent).expect("create temp dir");
+    }
+    SqliteRuntimeStore::new(path)
+}
+
+/// A runtime around `provider`, optionally told which transport to use.
+fn runtime_for(provider: RecordingProvider, transport: Option<ResponsesTransport>) -> Runtime {
+    let mut builder = Runtime::empty_builder()
+        .with_provider_instance(provider)
+        .with_store(temp_store());
+    if let Some(transport) = transport {
+        builder = builder.with_responses_transport(transport);
+    }
+    builder.build().expect("build runtime")
+}
+
+#[tokio::test]
+async fn a_runtime_that_chooses_nothing_still_streams_over_http_sse() {
+    let provider = RecordingProvider::new(BuiltinProvider::OpenAI, true);
+    let recorder = provider.clone();
+    let runtime = runtime_for(provider, None);
+
+    let mut agent = runtime
+        .spawn("agent", ModelInfo::new("model", BuiltinProvider::OpenAI))
+        .expect("spawn agent");
+    agent
+        .send(vec![ContentBlock::text("hello")])
+        .await
+        .expect("the turn runs");
+
+    assert_eq!(recorder.transports(), vec![ResponsesTransport::HttpSse]);
+}
+
+#[test]
+fn a_runtime_reports_the_transport_it_was_given() {
+    // Without this reader the choice is write-only: the only evidence a host's
+    // selection reached the runtime is a turn run against a provider that
+    // records what it was handed, so anything downstream can test its own field
+    // and stop at the seam — the shape of test that passes while the wiring
+    // between the two is broken.
+    let unset = runtime_for(RecordingProvider::new(BuiltinProvider::OpenAI, true), None);
+    assert_eq!(unset.responses_transport(), None);
+
+    let chosen = runtime_for(
+        RecordingProvider::new(BuiltinProvider::OpenAI, true),
+        Some(ResponsesTransport::WebSocket),
+    );
+    assert_eq!(
+        chosen.responses_transport(),
+        Some(ResponsesTransport::WebSocket)
+    );
+}
+
+#[tokio::test]
+async fn a_chosen_websocket_transport_reaches_the_request() {
+    // The gap this closes: `ResponsesRequestOptions.transport` existed and the
+    // websocket path was compiled in, but nothing in the runtime ever set the
+    // field, so every request went out over HTTP+SSE whatever the host wanted.
+    let provider = RecordingProvider::new(BuiltinProvider::OpenAI, true);
+    let recorder = provider.clone();
+    let runtime = runtime_for(provider, Some(ResponsesTransport::WebSocket));
+
+    let mut agent = runtime
+        .spawn("agent", ModelInfo::new("model", BuiltinProvider::OpenAI))
+        .expect("spawn agent");
+    agent
+        .send(vec![ContentBlock::text("hello")])
+        .await
+        .expect("the turn runs");
+
+    assert_eq!(recorder.transports(), vec![ResponsesTransport::WebSocket]);
+}
+
+#[tokio::test]
+async fn a_runtime_choice_settles_a_disagreeing_agent_config() {
+    // Two live opinions about one socket is not a state the runtime keeps: the
+    // connection-level answer is the one that holds.
+    let provider = RecordingProvider::new(BuiltinProvider::OpenAI, true);
+    let recorder = provider.clone();
+    let runtime = runtime_for(provider, Some(ResponsesTransport::HttpSse));
+
+    let mut config = AgentConfig::default();
+    config.provider_request_options.responses.transport = ResponsesTransport::WebSocket;
+    let mut agent = runtime
+        .spawn_with_config(
+            "agent",
+            ModelInfo::new("model", BuiltinProvider::OpenAI),
+            config,
+        )
+        .expect("spawn agent");
+    agent
+        .send(vec![ContentBlock::text("hello")])
+        .await
+        .expect("the turn runs");
+
+    assert_eq!(recorder.transports(), vec![ResponsesTransport::HttpSse]);
+}
+
+#[tokio::test]
+async fn a_provider_without_websockets_refuses_rather_than_pretending() {
+    // anthropic and gemini report `supports_websockets: false`. Answering over
+    // HTTP+SSE would look like success and be a transport nobody asked for.
+    let provider = RecordingProvider::new(BuiltinProvider::Anthropic, false).named("Anthropic");
+    let recorder = provider.clone();
+    let runtime = runtime_for(provider, Some(ResponsesTransport::WebSocket));
+
+    let mut agent = runtime
+        .spawn("agent", ModelInfo::new("model", BuiltinProvider::Anthropic))
+        .expect("spawn agent");
+    let error = agent
+        .send(vec![ContentBlock::text("hello")])
+        .await
+        .expect_err("a transport the provider cannot serve is refused");
+
+    let message = error.to_string();
+    assert!(
+        message.contains("Anthropic"),
+        "the refusal must name the provider: {message}"
+    );
+    assert!(
+        message.contains("websocket"),
+        "and say what it could not do: {message}"
+    );
+    assert!(
+        recorder.transports().is_empty(),
+        "the request must never have been sent"
+    );
+}