codex-rs/exec-server/src/server/session_registry.rs

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::time::Duration;

use codex_exec_server_protocol::JSONRPCErrorError;
use tokio::sync::Mutex;
use uuid::Uuid;

use crate::ExecServerRuntimePaths;
use crate::rpc::RpcNotificationSender;
use crate::rpc::invalid_request;
use crate::rpc::session_already_attached;
use crate::server::process_handler::ProcessHandler;
use crate::telemetry::ExecServerTelemetry;

#[cfg(test)]
const DETACHED_SESSION_TTL: Duration = Duration::from_millis(200);
#[cfg(not(test))]
const DETACHED_SESSION_TTL: Duration = Duration::from_secs(30);

pub(crate) struct SessionRegistry {
    sessions: Mutex<HashMap<String, Arc<SessionEntry>>>,
    telemetry: ExecServerTelemetry,
}

struct SessionEntry {
    session_id: String,
    process: ProcessHandler,
    attachment: StdMutex<AttachmentState>,
}

struct AttachmentState {
    current_connection_id: Option<ConnectionId>,
    detached_connection_id: Option<ConnectionId>,
    detached_expires_at: Option<tokio::time::Instant>,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct ConnectionId(Uuid);

impl std::fmt::Display for ConnectionId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}

#[derive(Clone)]
pub(crate) struct SessionHandle {
    registry: Arc<SessionRegistry>,
    entry: Arc<SessionEntry>,
    connection_id: ConnectionId,
}

impl SessionRegistry {
    pub(crate) fn new(telemetry: ExecServerTelemetry) -> Arc<Self> {
        Arc::new(Self {
            sessions: Mutex::new(HashMap::new()),
            telemetry,
        })
    }

    pub(crate) async fn attach(
        self: &Arc<Self>,
        resume_session_id: Option<String>,
        notifications: RpcNotificationSender,
        runtime_paths: ExecServerRuntimePaths,
    ) -> Result<SessionHandle, JSONRPCErrorError> {
        enum AttachOutcome {
            Attached(Arc<SessionEntry>),
            Expired {
                session_id: String,
                entry: Arc<SessionEntry>,
            },
        }

        let connection_id = ConnectionId(Uuid::new_v4());
        let outcome = {
            let mut sessions = self.sessions.lock().await;
            if let Some(session_id) = resume_session_id {
                let entry = sessions
                    .get(&session_id)
                    .cloned()
                    .ok_or_else(|| invalid_request(format!("unknown session id {session_id}")))?;
                if entry.is_expired(tokio::time::Instant::now()) {
                    let entry = sessions.remove(&session_id).ok_or_else(|| {
                        invalid_request(format!("unknown session id {session_id}"))
                    })?;
                    Ok(AttachOutcome::Expired { session_id, entry })
                } else if entry.has_active_connection() {
                    Err(session_already_attached(format!(
                        "session {session_id} is already attached to another connection"
                    )))
                } else {
                    entry.process.set_notification_sender(Some(notifications));
                    entry.attach(connection_id);
                    Ok(AttachOutcome::Attached(entry))
                }
            } else {
                let session_id = Uuid::new_v4().to_string();
                let entry = Arc::new(SessionEntry::new(
                    session_id.clone(),
                    ProcessHandler::new(notifications, self.telemetry.clone(), runtime_paths),
                    connection_id,
                ));
                sessions.insert(session_id, Arc::clone(&entry));
                Ok(AttachOutcome::Attached(entry))
            }
        };
        let entry = match outcome? {
            AttachOutcome::Attached(entry) => entry,
            AttachOutcome::Expired { session_id, entry } => {
                entry.process.shutdown().await;
                return Err(invalid_request(format!("unknown session id {session_id}")));
            }
        };

        Ok(SessionHandle {
            registry: Arc::clone(self),
            entry,
            connection_id,
        })
    }

    pub(crate) async fn shutdown(&self) {
        let sessions = std::mem::take(&mut *self.sessions.lock().await);
        for entry in sessions.into_values() {
            entry.process.shutdown().await;
        }
    }

    async fn expire_if_detached(&self, session_id: String, connection_id: ConnectionId) {
        tokio::time::sleep(DETACHED_SESSION_TTL).await;

        let removed = {
            let mut sessions = self.sessions.lock().await;
            let Some(entry) = sessions.get(&session_id) else {
                return;
            };
            if !entry.is_detached_connection_expired(connection_id, tokio::time::Instant::now()) {
                return;
            }
            sessions.remove(&session_id)
        };

        if let Some(entry) = removed {
            entry.process.shutdown().await;
        }
    }
}

impl Default for SessionRegistry {
    fn default() -> Self {
        Self {
            sessions: Mutex::new(HashMap::new()),
            telemetry: ExecServerTelemetry::default(),
        }
    }
}

impl SessionEntry {
    fn new(session_id: String, process: ProcessHandler, connection_id: ConnectionId) -> Self {
        Self {
            session_id,
            process,
            attachment: StdMutex::new(AttachmentState {
                current_connection_id: Some(connection_id),
                detached_connection_id: None,
                detached_expires_at: None,
            }),
        }
    }

    fn attach(&self, connection_id: ConnectionId) {
        let mut attachment = self
            .attachment
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        attachment.current_connection_id = Some(connection_id);
        attachment.detached_connection_id = None;
        attachment.detached_expires_at = None;
    }

    fn detach(&self, connection_id: ConnectionId) -> bool {
        let mut attachment = self
            .attachment
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if attachment.current_connection_id != Some(connection_id) {
            return false;
        }

        self.process.set_notification_sender(/*notifications*/ None);
        attachment.current_connection_id = None;
        attachment.detached_connection_id = Some(connection_id);
        attachment.detached_expires_at = Some(tokio::time::Instant::now() + DETACHED_SESSION_TTL);
        true
    }

    fn has_active_connection(&self) -> bool {
        self.attachment
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .current_connection_id
            .is_some()
    }

    fn is_attached_to(&self, connection_id: ConnectionId) -> bool {
        self.attachment
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .current_connection_id
            == Some(connection_id)
    }

    fn is_expired(&self, now: tokio::time::Instant) -> bool {
        self.attachment
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .detached_expires_at
            .is_some_and(|deadline| now >= deadline)
    }

    fn is_detached_connection_expired(
        &self,
        connection_id: ConnectionId,
        now: tokio::time::Instant,
    ) -> bool {
        let attachment = self
            .attachment
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        attachment.current_connection_id.is_none()
            && attachment.detached_connection_id == Some(connection_id)
            && attachment
                .detached_expires_at
                .is_some_and(|deadline| now >= deadline)
    }
}

impl SessionHandle {
    pub(crate) fn session_id(&self) -> &str {
        &self.entry.session_id
    }

    pub(crate) fn connection_id(&self) -> String {
        self.connection_id.to_string()
    }

    pub(crate) fn is_session_attached(&self) -> bool {
        self.entry.is_attached_to(self.connection_id)
    }

    pub(crate) fn process(&self) -> &ProcessHandler {
        &self.entry.process
    }

    pub(crate) async fn detach(&self) {
        if !self.entry.detach(self.connection_id) {
            return;
        }

        let registry = Arc::clone(&self.registry);
        let session_id = self.entry.session_id.clone();
        let connection_id = self.connection_id;
        tokio::spawn(async move {
            registry.expire_if_detached(session_id, connection_id).await;
        });
    }
}