@corbet-labs/ccht 0.2.0 → 0.2.3

package/source/src/native/pool.rs

@@ -0,0 +1,440 @@
+//! Keyed pool of native agent sessions sharing one client.
+//!
+//! Workspaces, projects or roles map to pool keys; each key owns exactly one
+//! session, created lazily on first use. Turns in different keys run in
+//! parallel automatically, while each key serializes its own turns. Events
+//! arrive tagged so late answers still land where they were asked.
+//!
+//! The pool owns session mechanics only: applications own prompts, permission
+//! decisions, display and storage, and read conversation snapshots back for
+//! replay or persistence.
+
+use std::collections::HashMap;
+
+use tokio::sync::mpsc;
+
+use super::{
+    AgentCommand, NativeClient, NativeError, NativeOptions, SessionHandle, SessionOptions,
+};
+use crate::{Conversation, SessionEvent, WireEvent};
+
+/// Application-chosen session key: a workspace, project or role name.
+pub type SessionKey = String;
+
+/// Streamed pool activity. Raw session events keep their full detail; the
+/// application maps them to display, storage and permission decisions.
+#[derive(Debug)]
+pub enum PoolEvent {
+    /// A session event, tagged with its key.
+    Session {
+        /// Application-chosen session key.
+        key: SessionKey,
+        /// Raw session event with full detail (boxed: the common case
+        /// dwarfs the terminal variant).
+        event: Box<SessionEvent>,
+    },
+    /// The forwarder for a removed session drained. No further events follow.
+    Ended {
+        /// Application-chosen session key.
+        key: SessionKey,
+    },
+}
+
+struct Session {
+    handle: SessionHandle,
+    conv: Conversation,
+    busy: bool,
+    seq: u64,
+}
+
+struct Inner {
+    client: Option<NativeClient>,
+    sessions: HashMap<SessionKey, Session>,
+    app_events: mpsc::UnboundedSender<PoolEvent>,
+    fwd: Option<mpsc::UnboundedSender<(SessionKey, Option<SessionEvent>)>>,
+}
+
+/// Keyed native sessions over one shared client connection.
+#[derive(Clone)]
+pub struct SessionPool {
+    inner: std::sync::Arc<tokio::sync::Mutex<Inner>>,
+}
+
+impl SessionPool {
+    /// Connect the shared client and return the pool plus its event stream.
+    /// A failed connection still returns a pool: every operation then
+    /// reports the outage instead of panicking, so applications degrade
+    /// gracefully (e.g. keep local tools working while the agent is down).
+    pub async fn connect(
+        command: AgentCommand,
+        options: NativeOptions,
+    ) -> (Self, mpsc::UnboundedReceiver<PoolEvent>) {
+        let (app_tx, app_rx) = mpsc::unbounded_channel();
+        let (fwd_tx, fwd_rx) = mpsc::unbounded_channel::<(SessionKey, Option<SessionEvent>)>();
+        let client = NativeClient::connect(command, options).await.ok();
+        let pool = Self {
+            inner: std::sync::Arc::new(tokio::sync::Mutex::new(Inner {
+                client,
+                sessions: HashMap::new(),
+                app_events: app_tx,
+                fwd: Some(fwd_tx),
+            })),
+        };
+        let router = pool.clone();
+        tokio::spawn(async move {
+            Self::route(router, fwd_rx).await;
+        });
+        (pool, app_rx)
+    }
+
+    /// Send a prompt to a key's session, creating it from `new` on first
+    /// use. Resolves when the turn ends; stream events (including the
+    /// terminal `Completed`) arrive separately, tagged with the key.
+    /// A second prompt while the key is busy fails with [`NativeError::Busy`].
+    pub async fn prompt(
+        &self,
+        key: &str,
+        new: SessionOptions,
+        text: String,
+    ) -> Result<(), NativeError> {
+        let handle = {
+            let mut inner = self.inner.lock().await;
+            Self::ensure(&mut inner, key, new).await?;
+            let session = inner.sessions.get_mut(key).expect("just ensured");
+            if session.busy {
+                return Err(NativeError::Busy);
+            }
+            session.busy = true;
+            session.handle.clone()
+        };
+        let request_id = format!("ccht-pool-{key}");
+        let result = handle.prompt(crate::Prompt::text(request_id, text)).await;
+        // Busy clears on the terminal stream event; a failed call has none,
+        // so release the slot here to avoid wedging the key.
+        if result.is_err()
+            && let Some(session) = self.inner.lock().await.sessions.get_mut(key)
+        {
+            session.busy = false;
+        }
+        result.map(|_| ())
+    }
+
+    /// Cancel the active turn of one key. Missing keys are a no-op success.
+    pub async fn cancel(&self, key: &str) -> Result<(), NativeError> {
+        let handle = {
+            let inner = self.inner.lock().await;
+            match inner.sessions.get(key) {
+                Some(session) => session.handle.clone(),
+                None => return Ok(()),
+            }
+        };
+        handle.cancel().await
+    }
+
+    /// Switch the model of an existing session, reporting the agent's
+    /// resulting configuration. Unknown keys fail without creating sessions:
+    /// model selection needs no session of its own.
+    pub async fn set_model(
+        &self,
+        key: &str,
+        model: &str,
+    ) -> Result<crate::SessionConfiguration, NativeError> {
+        let handle = {
+            let inner = self.inner.lock().await;
+            match inner.sessions.get(key) {
+                Some(session) => session.handle.clone(),
+                None => return Err(NativeError::Closed),
+            }
+        };
+        handle.set_model(model).await
+    }
+
+    /// Answer a pending permission request of one key's session.
+    pub async fn respond_permission(
+        &self,
+        key: &str,
+        request_id: &str,
+        decision: crate::acp::RequestPermissionOutcome,
+    ) -> Result<(), NativeError> {
+        let handle = {
+            let inner = self.inner.lock().await;
+            match inner.sessions.get(key) {
+                Some(session) => session.handle.clone(),
+                None => return Err(NativeError::Closed),
+            }
+        };
+        handle.respond_permission(request_id, decision).await
+    }
+
+    /// Drop one key's session and end it agent-side. Unknown keys are a no-op.
+    pub async fn close(&self, key: &str) {
+        let handle = {
+            let mut inner = self.inner.lock().await;
+            inner.sessions.remove(key).map(|s| s.handle)
+        };
+        if let Some(handle) = handle {
+            let _ = handle.close().await;
+        }
+    }
+
+    /// Whether the key holds a live session with a turn in flight.
+    pub async fn is_busy(&self, key: &str) -> bool {
+        self.inner
+            .lock()
+            .await
+            .sessions
+            .get(key)
+            .is_some_and(|s| s.busy)
+    }
+
+    /// Whether the key holds a live session at all.
+    pub async fn has_session(&self, key: &str) -> bool {
+        self.inner.lock().await.sessions.contains_key(key)
+    }
+
+    /// A snapshot of one key's conversation for replay or persistence.
+    /// Applications own storage; the pool only lends the state.
+    pub async fn conversation(&self, key: &str) -> Option<Conversation> {
+        self.inner
+            .lock()
+            .await
+            .sessions
+            .get(key)
+            .map(|s| s.conv.clone())
+    }
+
+    /// Shut every session down and close the shared client.
+    pub async fn shutdown(&self) {
+        let mut inner = self.inner.lock().await;
+        let keys: Vec<String> = inner.sessions.keys().cloned().collect();
+        for key in keys {
+            if let Some(session) = inner.sessions.remove(&key) {
+                let _ = session.handle.close().await;
+            }
+        }
+        if let Some(client) = inner.client.take() {
+            let _ = client.close().await;
+        }
+    }
+
+    async fn ensure(inner: &mut Inner, key: &str, new: SessionOptions) -> Result<(), NativeError> {
+        if inner.sessions.contains_key(key) {
+            return Ok(());
+        }
+        let Some(client) = &inner.client else {
+            return Err(NativeError::Closed);
+        };
+        let mut session = client.new_session(new).await?;
+        let handle = session.handle();
+        let mut events = session.take_events().map_err(|_| NativeError::Closed)?;
+        drop(session);
+        let key_owned = key.to_string();
+        let fwd = inner
+            .fwd
+            .clone()
+            .expect("forwarder channel installed at connect");
+        tokio::spawn(async move {
+            while let Some(ev) = events.recv().await {
+                if fwd.send((key_owned.clone(), Some(ev))).is_err() {
+                    break;
+                }
+            }
+            let _ = fwd.send((key_owned, None));
+        });
+        inner.sessions.insert(
+            key.to_string(),
+            Session {
+                handle,
+                conv: Conversation::new(format!("pool-{key}")),
+                busy: false,
+                seq: 0,
+            },
+        );
+        Ok(())
+    }
+
+    /// Single router: applies every streamed event to its key's conversation
+    /// (clearing finished turns) and re-emits it tagged for the application.
+    async fn route(
+        pool: Self,
+        mut fwd: mpsc::UnboundedReceiver<(SessionKey, Option<SessionEvent>)>,
+    ) {
+        while let Some((key, ev)) = fwd.recv().await {
+            let mut inner = pool.inner.lock().await;
+            let Some(session) = inner.sessions.get_mut(&key) else {
+                continue;
+            };
+            let Some(se) = ev else {
+                session.busy = false;
+                let _ = inner.app_events.send(PoolEvent::Ended { key });
+                continue;
+            };
+            session.seq += 1;
+            let rid = se.request_id.clone().unwrap_or_default();
+            let SessionEvent { event, .. } = &se;
+            let conv_id = session.conv.id().to_string();
+            let _ = session
+                .conv
+                .apply(WireEvent::new(conv_id, rid, session.seq, event.clone()));
+            match &event {
+                crate::Event::Completed { .. } | crate::Event::Error { .. } => {
+                    session.busy = false;
+                }
+                _ => {}
+            }
+            let _ = inner.app_events.send(PoolEvent::Session {
+                key,
+                event: Box::new(se),
+            });
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::Event;
+    use crate::acp::{ContentBlock, SessionUpdate};
+    use std::time::Duration;
+    use tokio::time::timeout;
+
+    fn options() -> NativeOptions {
+        NativeOptions {
+            operation_timeout: Duration::from_secs(5),
+            prompt_timeout: Duration::from_secs(5),
+            permission_timeout: Duration::from_secs(2),
+            shutdown_timeout: Duration::from_secs(2),
+            ..NativeOptions::default()
+        }
+    }
+
+    fn command(mode: &str) -> AgentCommand {
+        AgentCommand::new("python3").args(["-u", "-c", include_str!("fixture.py"), mode])
+    }
+
+    fn session_options() -> SessionOptions {
+        SessionOptions::new(std::env::temp_dir())
+    }
+
+    fn text(event: &SessionEvent) -> Option<&str> {
+        match &event.event {
+            Event::Update {
+                update: SessionUpdate::AgentMessageChunk(chunk),
+            } => {
+                if let ContentBlock::Text(content) = &chunk.content {
+                    Some(&content.text)
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        }
+    }
+
+    /// Collect streamed text for one key until its turn terminates.
+    /// Other keys' events are parked and replayed to the caller in order.
+    /// Collect streamed text for both keys until each turn terminates,
+    /// demultiplexing the single shared event stream.
+    async fn collect_two(rx: &mut mpsc::UnboundedReceiver<PoolEvent>) -> (String, String) {
+        let (mut a, mut b) = (String::new(), String::new());
+        let (mut done_a, mut done_b) = (false, false);
+        while !done_a || !done_b {
+            let ev = timeout(Duration::from_secs(15), rx.recv())
+                .await
+                .expect("event in time")
+                .expect("stream open");
+            if let PoolEvent::Session { key, event } = ev {
+                let terminal = matches!(event.event, Event::Completed { .. } | Event::Error { .. });
+                if let Some(t) = text(&event) {
+                    match key.as_str() {
+                        "a" => a.push_str(t),
+                        "b" => b.push_str(t),
+                        _ => {}
+                    }
+                }
+                if terminal {
+                    match key.as_str() {
+                        "a" => done_a = true,
+                        "b" => done_b = true,
+                        _ => {}
+                    }
+                }
+            }
+        }
+        (a, b)
+    }
+
+    #[tokio::test]
+    async fn parallel_keys_complete_independently() {
+        let (pool, mut rx) = SessionPool::connect(command("normal"), options()).await;
+        // Both turns run concurrently; a single shared session could not do this.
+        let p1 = pool.clone();
+        let t1 =
+            tokio::spawn(
+                async move { p1.prompt("a", session_options(), "first".to_string()).await },
+            );
+        let p2 = pool.clone();
+        let t2 = tokio::spawn(async move {
+            p2.prompt("b", session_options(), "second".to_string())
+                .await
+        });
+        let (r1, r2, texts) = tokio::join!(t1, t2, collect_two(&mut rx));
+        r1.expect("task a").expect("turn a");
+        r2.expect("task b").expect("turn b");
+        // Fixture answers every turn with the same stream.
+        assert_eq!(
+            texts,
+            ("hello world".to_string(), "hello world".to_string())
+        );
+        assert!(!pool.is_busy("a").await);
+        assert!(!pool.is_busy("b").await);
+        // Conversations accumulated per key for replay.
+        assert!(pool.conversation("a").await.is_some());
+        assert!(pool.conversation("b").await.is_some());
+        assert!(pool.conversation("ghost").await.is_none());
+        pool.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn second_prompt_on_busy_key_fails_cleanly() {
+        use crate::native::NativeError;
+
+        let (pool, _rx) = SessionPool::connect(command("hang"), options()).await;
+        // The hang fixture holds the turn open without answering, so the
+        // overlapping prompt deterministically hits the busy guard.
+        let p = pool.clone();
+        let held = tokio::spawn(async move {
+            // The hang fixture never answers, so the held turn stays open
+            // until the prompt timeout fires; the busy transition in between
+            // is what this test observes, not the outcome.
+            let _ = p.prompt("k", session_options(), "one".to_string()).await;
+        });
+        let mut busy_seen = false;
+        for _ in 0..100 {
+            if pool.is_busy("k").await {
+                busy_seen = true;
+                break;
+            }
+            tokio::time::sleep(Duration::from_millis(50)).await;
+        }
+        assert!(busy_seen, "turn never went busy");
+        let second = pool.prompt("k", session_options(), "two".to_string()).await;
+        assert!(
+            matches!(second, Err(NativeError::Busy)),
+            "expected Busy, got {second:?}"
+        );
+        let _ = held.await;
+        assert!(!pool.is_busy("k").await);
+        pool.shutdown().await;
+    }
+
+    #[tokio::test]
+    async fn missing_keys_fail_without_creating_sessions() {
+        let (pool, _rx) = SessionPool::connect(command("normal"), options()).await;
+        assert!(pool.set_model("ghost", "m").await.is_err());
+        assert!(pool.cancel("ghost").await.is_ok());
+        assert!(!pool.has_session("ghost").await);
+        pool.close("ghost").await;
+        pool.shutdown().await;
+    }
+}

package/source/src/native/session.rs

@@ -1,3 +1,9 @@
+//! Native conversation session with one ordered event stream.
+//!
+//! The worker owns prompt routing, permission validation, and timeout handling.
+//! Applications consume events through the single receiver and control the turn
+//! through [`SessionHandle`]; dropping an active prompt closes the connection.
+
 use std::collections::BTreeMap;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};

package/source/src/native/tests.rs

@@ -461,3 +461,207 @@ async fn advertised_controls_preserve_dependent_changes_and_reject_invalid_value
     ));
     client.close().await.unwrap();
 }
+
+fn codex_fixture(mode: &str) -> super::drivers::CodexDeviceDriver {
+    use std::time::Duration;
+    super::drivers::CodexDeviceDriver::new("python3")
+        .with_lead_args(vec![
+            "-u".to_owned(),
+            "-c".to_owned(),
+            include_str!("fixture.py").to_owned(),
+            mode.to_owned(),
+        ])
+        .with_call_timeout(Duration::from_secs(5))
+        .with_approval_deadline(Duration::from_secs(5))
+}
+
+fn opencode_fixture(mode: &str, key: &str) -> super::drivers::OpenCodeKeyDriver {
+    use std::time::Duration;
+    super::drivers::OpenCodeKeyDriver::new("python3", key.to_owned())
+        .unwrap()
+        .with_lead_args(vec![
+            "-u".to_owned(),
+            "-c".to_owned(),
+            include_str!("fixture.py").to_owned(),
+            mode.to_owned(),
+        ])
+        .with_operation_timeout(Duration::from_secs(5))
+        .with_readiness_timeout(Duration::from_secs(5))
+}
+
+#[tokio::test]
+async fn codex_challenge_round_trip_projects_presence() {
+    use super::drivers::{LoginDriver, LoginState};
+    let mut driver = codex_fixture("codex_ok");
+    let state = driver.start().await.unwrap();
+    let challenge = match state {
+        LoginState::ChallengeRequired(challenge) => challenge,
+        other => panic!("challenge expected, got {other:?}"),
+    };
+    assert_eq!(
+        challenge.verification_url,
+        "https://auth.openai.com/codex/device"
+    );
+    assert_eq!(challenge.user_code, "ABCD-1234");
+    challenge.validate().unwrap();
+    let state = driver.poll().await.unwrap();
+    match state {
+        LoginState::Authenticated(info) => {
+            assert_eq!(info.account.as_deref(), Some("user@example.com"));
+        }
+        other => panic!("authenticated presence expected, got {other:?}"),
+    }
+    let info = driver.account().await.unwrap();
+    assert_eq!(info.account.as_deref(), Some("user@example.com"));
+}
+
+#[tokio::test]
+async fn codex_rejects_malformed_challenge_without_leak() {
+    use super::drivers::{Challenge, DriverError, LoginDriver};
+    for mode in ["codex_bad_url", "codex_bad_code"] {
+        let mut driver = codex_fixture(mode);
+        let error = driver.start().await.unwrap_err();
+        assert!(
+            matches!(error, DriverError::InvalidOptions(_)),
+            "expected invalid options for {mode}, got {error:?}"
+        );
+        assert_eq!(error.code(), "invalid_options");
+        let rendered = format!("{error} {error:?}");
+        assert!(!rendered.contains("ABCD-1234"));
+        assert!(!rendered.contains("evil.example.com"));
+    }
+    let bad = Challenge::new(
+        "http://evil.example.com/x".to_owned(),
+        "ABCD-1234".to_owned(),
+    );
+    assert!(matches!(
+        bad.validate().unwrap_err(),
+        DriverError::InvalidOptions(_)
+    ));
+    let short = Challenge::new(
+        "https://auth.openai.com/codex/device".to_owned(),
+        "x".to_owned(),
+    );
+    assert!(matches!(
+        short.validate().unwrap_err(),
+        DriverError::InvalidOptions(_)
+    ));
+    let custom = super::drivers::CodexDeviceDriver::with_allowed_hosts(
+        "python3",
+        vec!["example.com".to_owned()],
+    )
+    .unwrap()
+    .with_lead_args(vec![
+        "-u".to_owned(),
+        "-c".to_owned(),
+        include_str!("fixture.py").to_owned(),
+        "codex_ok".to_owned(),
+    ]);
+    let mut custom = custom;
+    let error = custom.start().await.unwrap_err();
+    assert_eq!(
+        error,
+        DriverError::InvalidOptions("challenge host is not allowed")
+    );
+}
+
+#[tokio::test]
+async fn codex_poll_times_out_without_browser_approval() {
+    use super::drivers::{DriverError, LoginDriver, LoginState};
+    use std::time::Duration;
+    let mut driver = codex_fixture("codex_never");
+    driver = driver
+        .with_call_timeout(Duration::from_millis(300))
+        .with_approval_deadline(Duration::from_millis(400));
+    let state = driver.start().await.unwrap();
+    assert!(matches!(state, LoginState::ChallengeRequired(_)));
+    let error = driver.poll().await.unwrap_err();
+    assert_eq!(error, DriverError::Timeout);
+    assert_eq!(error.code(), "timeout");
+}
+
+#[tokio::test]
+async fn codex_cancel_stops_helper_and_reports_cancelled() {
+    use super::drivers::{DriverError, LoginDriver, LoginState};
+    let mut driver = codex_fixture("codex_never");
+    let state = driver.start().await.unwrap();
+    assert!(matches!(state, LoginState::ChallengeRequired(_)));
+    assert!(driver.has_helper());
+    driver.cancel().await.unwrap();
+    assert!(!driver.has_helper());
+    assert_eq!(driver.poll().await.unwrap_err(), DriverError::Cancelled);
+    assert_eq!(driver.account().await.unwrap_err(), DriverError::Cancelled);
+    assert_eq!(driver.start().await.unwrap_err(), DriverError::Cancelled);
+}
+
+#[tokio::test]
+async fn codex_declined_approval_reports_failed() {
+    use super::drivers::{LoginDriver, LoginState};
+    let mut driver = codex_fixture("codex_declined");
+    let state = driver.start().await.unwrap();
+    assert!(matches!(state, LoginState::ChallengeRequired(_)));
+    assert_eq!(driver.poll().await.unwrap(), LoginState::Failed);
+}
+
+#[tokio::test]
+async fn opencode_key_login_confirms_connected() {
+    use super::drivers::{LoginDriver, LoginState};
+    let key = "ccht-fixture-key-9f8e7d6c5b4a";
+    let mut driver = opencode_fixture("opencode_ok", key);
+    let state = driver.start().await.unwrap();
+    match state {
+        LoginState::Authenticated(info) => {
+            assert_eq!(info.account.as_deref(), Some("opencode-go"));
+        }
+        other => panic!("authenticated presence expected, got {other:?}"),
+    }
+    let state = driver.poll().await.unwrap();
+    assert!(matches!(state, LoginState::Authenticated(_)));
+    let info = driver.account().await.unwrap();
+    assert_eq!(info.account.as_deref(), Some("opencode-go"));
+    let rendered = format!("{driver:?}");
+    assert!(!rendered.contains(key));
+}
+
+#[tokio::test]
+async fn opencode_start_times_out_when_server_never_ready() {
+    use super::drivers::{DriverError, LoginDriver};
+    use std::time::Duration;
+    let mut driver = opencode_fixture("opencode_never", "ccht-fixture-key-timeout");
+    driver = driver
+        .with_operation_timeout(Duration::from_millis(300))
+        .with_readiness_timeout(Duration::from_millis(400));
+    let error = driver.start().await.unwrap_err();
+    assert_eq!(error, DriverError::Timeout);
+}
+
+#[tokio::test]
+async fn driver_debug_and_errors_omit_key_material() {
+    use super::drivers::{DriverError, LoginDriver};
+    let key = "ccht-fixture-key-secret-12345";
+    let driver = super::drivers::OpenCodeKeyDriver::new("python3", key.to_owned()).unwrap();
+    let rendered = format!("{driver:?}");
+    assert!(!rendered.contains(key));
+    assert!(!rendered.contains("OPENCODE_SERVER_PASSWORD"));
+    let codex = super::drivers::CodexDeviceDriver::new("python3");
+    let codex_rendered = format!("{codex:?}");
+    assert!(!codex_rendered.contains("ABCD-1234"));
+    for error in [
+        DriverError::InvalidOptions("bad shape"),
+        DriverError::Timeout,
+        DriverError::Closed,
+        DriverError::Spawn("helper unavailable"),
+        DriverError::Protocol(-32603),
+        DriverError::Cancelled,
+        DriverError::Unsupported("extra step"),
+    ] {
+        let text = format!("{error} {error:?} {}", error.code());
+        assert!(!text.contains(key));
+        assert!(!text.contains("ccht-fixture"));
+    }
+    let mut failing = opencode_fixture("opencode_fail", key);
+    let state = failing.start().await.unwrap();
+    assert_eq!(state, super::drivers::LoginState::Failed);
+    let rendered = format!("{failing:?}");
+    assert!(!rendered.contains(key));
+}