cortex/crates/helexa-acp/src/session.rs

//! Per-session state for the ACP agent loop.
//!
//! Concurrency:
//!
//! - [`SessionStore`] is an `Arc<RwLock<HashMap<SessionId, …>>>`. The map
//!   itself is read-mostly: it changes only on `session/new` and never
//!   shrinks during Stage 2, so an `RwLock` keeps concurrent reads
//!   contention-free.
//! - Each session is wrapped in its own `Arc<Mutex<SessionState>>`. Holding
//!   one session's lock doesn't block requests against any other session,
//!   which matters once a client opens multiple sessions in parallel.
//!
//! All operations hold a lock only long enough to copy out (or mutate) the
//! state they need — never across an `await` that drives the upstream
//! provider stream.

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;

use agent_client_protocol::schema::{SessionId, SessionModeId};
use tokio::sync::{Mutex, RwLock};
use tokio_util::sync::CancellationToken;

use crate::provider::Message;

/// Mode id advertised as the gated default. Writes / bash prompt for
/// permission via `session/request_permission`.
pub const MODE_DEFAULT: &str = "default";

/// Mode id advertised as "auto-allow everything". Matches the
/// favorite name (`bypassPermissions`) Zed clients tend to reference.
pub const MODE_BYPASS: &str = "bypassPermissions";

/// Mode id for read-and-plan-only operation. The model may read files
/// and list directories freely, may write *only* into the per-project
/// plan directory under `$XDG_DATA_HOME/helexa-acp/plans/<project-id>/`,
/// and cannot run shell commands. Designed for "draft the
/// implementation plan, then I'll review and let you execute" flows.
pub const MODE_PLAN: &str = "plan";

/// State carried for a single ACP session.
///
/// Mutated under `Mutex<SessionState>`; never share a clone across
/// tasks expecting to see the same `cancel` token — clone the token
/// explicitly when handing it to the streaming task.
#[derive(Debug)]
pub struct SessionState {
    /// Conversation history in chronological order (user / assistant
    /// turns). The system prompt is *not* stored here — it's built
    /// fresh per request so any cwd / config changes take effect.
    pub history: Vec<Message>,
    /// Working directory the client opened the session against. Used
    /// by [`crate::prompt::build_system_prompt`] and (Stage 3) by
    /// filesystem tools.
    pub cwd: PathBuf,
    /// Currently-selected model id. Format is either a bare model id
    /// (resolved against the default endpoint) or `endpoint:model`.
    /// Mutated by `session/set_model` in Stage 4; Stage 2 sets it
    /// once at session creation and never changes it.
    pub model_id: String,
    /// Cancellation handle for the in-flight prompt, if any. A fresh
    /// token is installed at the start of every `session/prompt`
    /// request; `session/cancel` fires this one. Between prompts the
    /// token is "spent" — firing it does nothing — which is fine,
    /// `session/cancel` is a no-op when there's nothing to cancel.
    pub cancel: CancellationToken,
    /// Permission gating mode. Stage 3 advertises two ids in
    /// `NewSessionResponse.modes`: [`MODE_DEFAULT`] (writes / bash
    /// prompt the user) and [`MODE_BYPASS`] (auto-allow). Mutated by
    /// `session/set_mode`.
    pub mode_id: SessionModeId,
}

impl SessionState {
    pub fn new(cwd: PathBuf, model_id: String) -> Self {
        Self {
            history: Vec::new(),
            cwd,
            model_id,
            cancel: CancellationToken::new(),
            mode_id: SessionModeId::new(MODE_DEFAULT),
        }
    }
}

/// Concurrent map of live sessions.
///
/// Cloning is cheap (`Arc` bump). Pass clones into every handler that
/// needs session access; never hold a clone across an `.await` that
/// could outlive the request.
pub type SessionStore = Arc<RwLock<HashMap<SessionId, Arc<Mutex<SessionState>>>>>;

/// Fresh, empty session store.
pub fn new_store() -> SessionStore {
    Arc::new(RwLock::new(HashMap::new()))
}

/// Look up a session by id. Returns `None` if no such session is registered.
pub async fn get(store: &SessionStore, id: &SessionId) -> Option<Arc<Mutex<SessionState>>> {
    store.read().await.get(id).cloned()
}

/// Register a fresh session. Overwrites any prior entry with the same id
/// (which should never happen — ids are uniquely generated by the agent).
pub async fn insert(store: &SessionStore, id: SessionId, state: SessionState) {
    store.write().await.insert(id, Arc::new(Mutex::new(state)));
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::provider::{MessageContent, Role};

    fn id(s: &str) -> SessionId {
        SessionId::new(s)
    }

    #[tokio::test]
    async fn insert_then_get_round_trip() {
        let store = new_store();
        let state = SessionState::new(PathBuf::from("/tmp"), "m".into());
        insert(&store, id("s1"), state).await;
        let got = get(&store, &id("s1")).await.expect("session present");
        let locked = got.lock().await;
        assert_eq!(locked.cwd, PathBuf::from("/tmp"));
        assert_eq!(locked.model_id, "m");
        assert!(locked.history.is_empty());
    }

    #[tokio::test]
    async fn missing_session_is_none() {
        let store = new_store();
        assert!(get(&store, &id("nope")).await.is_none());
    }

    #[tokio::test]
    async fn history_is_per_session() {
        let store = new_store();
        insert(
            &store,
            id("a"),
            SessionState::new(PathBuf::from("/a"), "m".into()),
        )
        .await;
        insert(
            &store,
            id("b"),
            SessionState::new(PathBuf::from("/b"), "m".into()),
        )
        .await;

        // Appending to a's history must not affect b's.
        get(&store, &id("a"))
            .await
            .unwrap()
            .lock()
            .await
            .history
            .push(Message {
                role: Role::User,
                content: MessageContent::Text {
                    text: "hello".into(),
                },
            });

        assert_eq!(
            get(&store, &id("a"))
                .await
                .unwrap()
                .lock()
                .await
                .history
                .len(),
            1
        );
        assert_eq!(
            get(&store, &id("b"))
                .await
                .unwrap()
                .lock()
                .await
                .history
                .len(),
            0
        );
    }
}