agents-orchestrator/agent-orchestrator
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(agor-pty): complete PTY daemon — auth, sessions, output fanout

Browse source
This commit is contained in:
Hibryda 2026-03-20 03:10:49 +01:00
parent 4b5583430d
commit f3456bd09d
6 changed files with 1853 additions and 65 deletions
Download patch file Download diff file Expand all files Collapse all files

131
agor-pty/src/session.rs
View file

@ -1,5 +1,6 @@
use std::collections::HashMap;
use std::io::{Read, Write as IoWrite};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
@ -10,7 +11,10 @@ use crate::protocol::SessionInfo;
const OUTPUT_CHANNEL_CAP: usize = 256;
/// A live PTY session.
/// A live (or recently exited) PTY session.
///
/// All fields that cross await points are either `Send + Sync` or wrapped in
/// `Arc<Mutex<_>>` so the Session itself is `Send`.
pub struct Session {
pub id: String,
pub pid: u32,
@ -19,19 +23,20 @@ pub struct Session {
pub cols: u16,
pub rows: u16,
pub created_at: u64,
/// Used to write input into the PTY master.
/// Used to write input into the PTY master. `Box<dyn Write + Send>`.
writer: Arc<Mutex<Box<dyn IoWrite + Send>>>,
/// Broadcast channel — subscribers receive raw output chunks.
/// Broadcast channel — all subscribers receive raw output chunks.
pub tx: broadcast::Sender<Vec<u8>>,
/// Set to false when the child process exits.
pub alive: Arc<std::sync::atomic::AtomicBool>,
/// Last known exit code (populated by the reader task on process exit).
/// false once the child process exits.
pub alive: Arc<AtomicBool>,
/// Last known exit code (set by the reader task on child exit).
/// Public for callers that poll exit state after SessionClosed is received.
#[allow(dead_code)]
pub exit_code: Arc<Mutex<Option<i32>>>,
/// Keep the master alive so the PTY stays open.
_master: Box<dyn portable_pty::MasterPty + Send>,
}
impl Session {
/// Snapshot metadata for ListSessions responses.
pub fn snapshot(&self) -> SessionInfo {
SessionInfo {
id: self.id.clone(),
@ -41,40 +46,35 @@ impl Session {
cols: self.cols,
rows: self.rows,
created_at: self.created_at,
alive: self.alive.load(std::sync::atomic::Ordering::Relaxed),
alive: self.alive.load(Ordering::Relaxed),
}
}
/// Write bytes into the PTY (user keystrokes, paste, etc.).
/// Write raw bytes into the PTY master (keyboard input, paste, etc.).
pub async fn write_input(&self, data: &[u8]) -> Result<(), String> {
let mut w = self.writer.lock().await;
w.write_all(data)
.map_err(|e| format!("PTY write failed for session {}: {e}", self.id))
.map_err(|e| format!("PTY write for {}: {e}", self.id))
}
/// Send TIOCSWINSZ to resize the PTY.
pub fn resize(&mut self, cols: u16, rows: u16) -> Result<(), String> {
/// Update cached dimensions after a resize. The actual TIOCSWINSZ is issued
/// by the daemon before calling this.
pub fn note_resize(&mut self, cols: u16, rows: u16) {
self.cols = cols;
self.rows = rows;
// portable-pty exposes resize via the master handle which we've moved.
// We reach into nix directly via the stored master fd.
// portable-pty's MasterPty trait has `resize` on nightly targets; on
// stable we use nix ourselves.
log::debug!(
"session {} resize → {}x{} (handled via pty master)",
self.id, cols, rows
);
// The resize is done by the caller via `master.resize()` before this
// method; this method just updates our cached dimensions.
Ok(())
}
/// Return a new receiver subscribed to this session's broadcast output.
pub fn subscribe(&self) -> broadcast::Receiver<Vec<u8>> {
self.tx.subscribe()
}
}
/// Owns all sessions and serialises mutations.
// ---------------------------------------------------------------------------
// Session manager
// ---------------------------------------------------------------------------
/// Owns the full set of PTY sessions.
pub struct SessionManager {
sessions: HashMap<String, Session>,
default_shell: String,
@ -88,8 +88,10 @@ impl SessionManager {
}
}
/// Create and start a new PTY session. Returns the session id, pid, and a
/// receiver end of the output broadcast channel.
/// Spawn a new PTY session.
///
/// Returns `(pid, output_rx)` on success. `on_exit` is called from the
/// blocking reader task once the child process exits.
pub fn create_session(
&mut self,
id: String,
@ -98,7 +100,6 @@ impl SessionManager {
env: Option<HashMap<String, String>>,
cols: u16,
rows: u16,
// Callback invoked from the reader task when the child exits.
on_exit: impl FnOnce(String, Option<i32>) + Send + 'static,
) -> Result<(u32, broadcast::Receiver<Vec<u8>>), String> {
if self.sessions.contains_key(&id) {
@ -114,7 +115,7 @@ impl SessionManager {
pixel_width: 0,
pixel_height: 0,
})
.map_err(|e| format!("openpty failed: {e}"))?;
.map_err(|e| format!("openpty: {e}"))?;
let mut cmd = CommandBuilder::new(&shell_path);
if let Some(ref dir) = cwd {
@ -129,34 +130,40 @@ impl SessionManager {
let child = pair
.slave
.spawn_command(cmd)
.map_err(|e| format!("spawn failed: {e}"))?;
.map_err(|e| format!("spawn: {e}"))?;
let pid = child.process_id().unwrap_or(0);
let cwd_str = cwd.unwrap_or_else(|| std::env::current_dir()
.map(|p| p.to_string_lossy().into_owned())
.unwrap_or_else(|_| "/".into()));
// portable-pty requires us to take the writer from the master.
let cwd_str = cwd.unwrap_or_else(|| {
std::env::current_dir()
.map(|p| p.to_string_lossy().into_owned())
.unwrap_or_else(|_| "/".into())
});
// Take the writer before moving `pair.master` into the reader task.
let writer = pair
.master
.take_writer()
.map_err(|e| format!("take_writer failed: {e}"))?;
.map_err(|e| format!("take_writer: {e}"))?;
// Obtain a blocking reader for the reader task.
// Clone a reader; the master handle itself moves into the blocking task
// so the PTY stays open until the reader is done.
let reader = pair
.master
.try_clone_reader()
.map_err(|e| format!("clone_reader failed: {e}"))?;
.map_err(|e| format!("clone_reader: {e}"))?;
let (tx, rx) = broadcast::channel(OUTPUT_CHANNEL_CAP);
let alive = Arc::new(std::sync::atomic::AtomicBool::new(true));
let alive = Arc::new(AtomicBool::new(true));
let exit_code = Arc::new(Mutex::new(None::<i32>));
// Spawn a blocking task to drain PTY output and broadcast it.
// Spawn the blocking reader task. It takes ownership of `pair.master`
// (via `_master`) so the PTY file descriptor stays open.
let tx_clone = tx.clone();
let alive_clone = alive.clone();
let exit_code_clone = exit_code.clone();
let id_clone = id.clone();
let _master = pair.master; // keep PTY fd alive inside the task
tokio::task::spawn_blocking(move || {
read_pty_output(
reader,
@ -166,6 +173,7 @@ impl SessionManager {
id_clone,
on_exit,
child,
_master,
);
});
@ -181,7 +189,6 @@ impl SessionManager {
tx,
alive,
exit_code,
_master: pair.master,
};
log::info!("created session {id} pid={pid}");
@ -201,8 +208,8 @@ impl SessionManager {
self.sessions.values().map(|s| s.snapshot()).collect()
}
/// Close a session: the child is killed if still alive and the entry is
/// removed after a brief wait for the reader task to notice.
/// Remove a session entry. The reader task will notice the PTY is closed
/// and stop on its own.
pub fn close_session(&mut self, id: &str) -> Result<(), String> {
if self.sessions.remove(id).is_some() {
log::info!("closed session {id}");
@ -211,10 +218,6 @@ impl SessionManager {
Err(format!("session {id} not found"))
}
}
pub fn sessions(&self) -> &HashMap<String, Session> {
&self.sessions
}
}
// ---------------------------------------------------------------------------
@ -228,24 +231,27 @@ fn unix_now() -> u64 {
.unwrap_or(0)
}
/// Blocking PTY reader — lives in a `spawn_blocking` task.
/// Blocking PTY reader — lives inside `tokio::task::spawn_blocking`.
///
/// `_master` is held here so the PTY file descriptor is not closed until this
/// task finishes.
#[allow(clippy::too_many_arguments)]
fn read_pty_output(
mut reader: Box<dyn Read + Send>,
tx: broadcast::Sender<Vec<u8>>,
alive: Arc<std::sync::atomic::AtomicBool>,
alive: Arc<AtomicBool>,
exit_code_cell: Arc<Mutex<Option<i32>>>,
id: String,
on_exit: impl FnOnce(String, Option<i32>),
mut child: Box<dyn portable_pty::Child + Send>,
_master: Box<dyn portable_pty::MasterPty + Send>,
) {
let mut buf = [0u8; 4096];
loop {
match reader.read(&mut buf) {
Ok(0) => break,
Ok(n) => {
let chunk = buf[..n].to_vec();
// Non-blocking send — if all receivers are gone, ignore.
let _ = tx.send(chunk);
let _ = tx.send(buf[..n].to_vec());
}
Err(e) => {
log::debug!("session {id} reader error: {e}");
@ -254,21 +260,15 @@ fn read_pty_output(
}
}
// PTY EOF — child has exited (or master was closed).
alive.store(false, std::sync::atomic::AtomicBool::from(false).load(std::sync::atomic::Ordering::SeqCst).into());
alive.store(false, std::sync::atomic::Ordering::Relaxed);
alive.store(false, Ordering::Relaxed);
let code = child.wait().ok().and_then(|status| {
if let Some(exit) = status.exit_code() {
Some(exit as i32)
} else {
None
}
});
// `exit_code()` on portable-pty returns u32 directly (not Option).
let code = child
.wait()
.ok()
.map(|status| status.exit_code() as i32);
// Write exit code into the shared cell.
// We're in a blocking context so we use try_lock in a tight spin — the
// lock is never held for long.
// Write exit code using try_lock spin — the lock is never held for long.
loop {
if let Ok(mut guard) = exit_code_cell.try_lock() {
*guard = code;
@ -277,6 +277,7 @@ fn read_pty_output(
std::thread::sleep(std::time::Duration::from_millis(1));
}
log::info!("session {id} exited with code {:?}", code);
log::info!("session {id} exited with code {code:?}");
on_exit(id, code);
// `_master` drops here — PTY closed.
}