use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64};
use crate::domain::errors::{AgentError, AgentResult, ErrorCode};
use crate::domain::policy::CommandPolicy;
use crate::domain::resize::TerminalSize;
use crate::infrastructure::pty_session::PtySession;
#[derive(Clone)]
pub(crate) struct SessionOutput {
pub session_id: String,
pub session: Arc<PtySession>,
pub next_seq: Arc<AtomicU64>,
pub active: Arc<AtomicBool>,
pub exit_queued: Arc<AtomicBool>,
}
pub(crate) struct SessionSpawn {
pub session_id: String,
pub command: Vec<String>,
pub cwd: Option<String>,
}
struct SessionEntry {
session: Arc<PtySession>,
next_seq: Arc<AtomicU64>,
active: Arc<AtomicBool>,
exit_queued: Arc<AtomicBool>,
}
pub struct SessionManager {
sessions: HashMap<String, SessionEntry>,
policy: CommandPolicy,
default_command: Vec<String>,
}
#[allow(dead_code)]
impl SessionManager {
pub fn new(policy: CommandPolicy, default_command: Vec<String>) -> Self {
Self {
sessions: HashMap::new(),
policy,
default_command,
}
}
pub fn policy_snapshot(&self) -> CommandPolicy {
self.policy.clone()
}
pub async fn create_session(
&mut self,
session_id: String,
command: Option<Vec<String>>,
cwd: Option<String>,
size: TerminalSize,
) -> AgentResult<u32> {
let spawn = self.prepare_session(session_id, command, cwd)?;
let session =
PtySession::spawn(spawn.session_id.clone(), spawn.command, spawn.cwd, size).await?;
let pid = session.pid();
self.insert_session(spawn.session_id, session)?;
Ok(pid)
}
pub(crate) fn prepare_session(
&self,
session_id: String,
command: Option<Vec<String>>,
cwd: Option<String>,
) -> AgentResult<SessionSpawn> {
if self.sessions.contains_key(&session_id) {
return Err(session_exists());
}
let command = command.unwrap_or_else(|| self.default_command.clone());
let validated = self.policy.validate(&command, cwd.as_deref())?;
Ok(SessionSpawn {
session_id,
command: validated.command,
cwd: validated.cwd,
})
}
pub(crate) fn insert_session(
&mut self,
session_id: String,
session: PtySession,
) -> AgentResult<SessionOutput> {
if self.sessions.contains_key(&session_id) {
return Err(session_exists());
}
let session = Arc::new(session);
let next_seq = Arc::new(AtomicU64::new(0));
let active = Arc::new(AtomicBool::new(true));
let exit_queued = Arc::new(AtomicBool::new(false));
self.sessions.insert(
session_id.clone(),
SessionEntry {
session: session.clone(),
next_seq: next_seq.clone(),
active: active.clone(),
exit_queued: exit_queued.clone(),
},
);
Ok(SessionOutput {
session_id,
session,
next_seq,
active,
exit_queued,
})
}
pub(crate) fn session_output(&self, session_id: &str) -> AgentResult<SessionOutput> {
self.sessions
.get(session_id)
.map(|entry| SessionOutput {
session_id: session_id.to_string(),
session: entry.session.clone(),
next_seq: entry.next_seq.clone(),
active: entry.active.clone(),
exit_queued: entry.exit_queued.clone(),
})
.ok_or_else(session_missing)
}
pub(crate) fn output_sessions(&self) -> Vec<SessionOutput> {
self.sessions
.iter()
.map(|(session_id, entry)| SessionOutput {
session_id: session_id.clone(),
session: entry.session.clone(),
next_seq: entry.next_seq.clone(),
active: entry.active.clone(),
exit_queued: entry.exit_queued.clone(),
})
.collect()
}
pub(crate) fn remove_session(&mut self, session_id: &str) -> AgentResult<SessionOutput> {
self.sessions
.remove(session_id)
.map(|entry| SessionOutput {
session_id: session_id.to_string(),
session: entry.session,
next_seq: entry.next_seq,
active: entry.active,
exit_queued: entry.exit_queued,
})
.ok_or_else(session_missing)
}
pub(crate) fn restore_session(&mut self, output: SessionOutput) {
self.sessions.insert(
output.session_id,
SessionEntry {
session: output.session,
next_seq: output.next_seq,
active: output.active,
exit_queued: output.exit_queued,
},
);
}
pub async fn write_input(&self, session_id: &str, bytes: &[u8]) -> AgentResult<()> {
self.session_output(session_id)?
.session
.write_input(bytes)
.await
}
pub async fn resize_session(&self, session_id: &str, size: TerminalSize) -> AgentResult<()> {
self.session_output(session_id)?.session.resize(size).await
}
pub async fn kill_session(&mut self, session_id: &str) -> AgentResult<()> {
let output = self.remove_session(session_id)?;
if let Err(err) = output.session.kill().await {
self.restore_session(output);
return Err(err);
}
Ok(())
}
pub async fn shutdown(&mut self) {
for (session_id, entry) in std::mem::take(&mut self.sessions) {
if entry.session.kill().await.is_err() {
self.sessions.insert(session_id, entry);
}
}
}
}
fn session_missing() -> AgentError {
AgentError::new(ErrorCode::SessionNotFound, "session not found")
}
fn session_exists() -> AgentError {
AgentError::new(ErrorCode::SessionAlreadyExists, "session already exists")
}