use axum::extract::ws::{Message, WebSocket};
use futures_util::{SinkExt, StreamExt};
use portable_pty::PtySize;
use tokio::sync::broadcast::error::RecvError;
use tracing::{debug, error, info, warn};
use crate::AppState;
use crate::ws::terminal::{ClientControl, ServerControl, TerminalQuery};
pub async fn handle_pty_terminal(
ws: WebSocket,
session_id: String,
query: TerminalQuery,
state: AppState,
) {
info!("terminal WS connected (pty): session={}", session_id);
let row: Option<(String, Option<String>, Option<String>)> = sqlx::query_as(
"SELECT workspace_path, tmux_session_name, last_cwd FROM sessions WHERE id = ?",
)
.bind(&session_id)
.fetch_optional(&state.db)
.await
.ok()
.flatten();
let Some((workspace, engine_key, last_cwd)) = row else {
let (mut sender, _) = ws.split();
let msg =
serde_json::to_string(&ServerControl::Error { message: "session not found" }).unwrap();
let _ = sender.send(Message::Text(msg.into())).await;
return;
};
let key = engine_key.unwrap_or_else(|| session_id.clone());
let spawn_cwd = last_cwd.filter(|p| std::path::Path::new(p).is_dir()).unwrap_or(workspace);
let size = PtySize {
rows: query.rows.filter(|&r| r > 0 && r <= 1000).unwrap_or(24),
cols: query.cols.filter(|&c| c > 0 && c <= 1000).unwrap_or(80),
pixel_width: 0,
pixel_height: 0,
};
info!(
"terminal PTY initial size: {}x{} for session={} (pty key={})",
size.cols, size.rows, session_id, key
);
let attach = match state.engines.attach_pty(&key, &spawn_cwd, size).await {
Ok(a) => a,
Err(e) => {
error!("failed to attach pty session {}: {}", key, e);
let (mut sender, _) = ws.split();
let msg = serde_json::to_string(&ServerControl::Error {
message: "failed to start terminal",
})
.unwrap();
let _ = sender.send(Message::Text(msg.into())).await;
return;
}
};
if let Err(e) = attach.resize(size) {
warn!("initial resize failed (pty): {}", e);
}
if attach.reconnected && size.rows > 1 {
let nudged = PtySize { rows: size.rows - 1, ..size };
let _ = attach.resize(nudged);
tokio::time::sleep(std::time::Duration::from_millis(30)).await;
let _ = attach.resize(size);
}
let (mut ws_tx, mut ws_rx) = ws.split();
let attached_msg = serde_json::to_string(&ServerControl::Attached { session: &key }).unwrap();
if ws_tx.send(Message::Text(attached_msg.into())).await.is_err() {
return;
}
if !attach.replay.is_empty()
&& ws_tx.send(Message::Binary(attach.replay.clone().into())).await.is_err()
{
return;
}
let (pty_in_tx, mut pty_in_rx) = tokio::sync::mpsc::channel::<Vec<u8>>(256);
let writer_attach = attach.clone();
let resize_attach = attach.clone();
let writer_key = key.clone();
std::thread::spawn(move || {
while let Some(data) = pty_in_rx.blocking_recv() {
if let Err(e) = writer_attach.write(&data) {
warn!("PTY write failed (pty session {writer_key}): {e}");
return;
}
}
debug!("PTY writer exited (pty session {writer_key})");
});
let mut rx = attach.rx;
let mut forward_handle = tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(data) => {
if ws_tx.send(Message::Binary(data.into())).await.is_err() {
break;
}
}
Err(RecvError::Lagged(n)) => {
warn!("pty output lagged, dropped {n} frames");
continue;
}
Err(RecvError::Closed) => break,
}
}
});
let mut read_handle = tokio::spawn(async move {
while let Some(msg) = ws_rx.next().await {
match msg {
#[allow(clippy::collapsible_match)]
Ok(Message::Binary(data)) => {
if pty_in_tx.send(data.to_vec()).await.is_err() {
break;
}
}
Ok(Message::Text(text)) => {
if let Ok(ctrl) = serde_json::from_str::<ClientControl>(&text) {
match ctrl {
ClientControl::Resize { cols, rows } => {
if cols > 0 && cols <= 1000 && rows > 0 && rows <= 1000 {
let new_size =
PtySize { rows, cols, pixel_width: 0, pixel_height: 0 };
if let Err(e) = resize_attach.resize(new_size) {
warn!("PTY resize failed (pty): {}", e);
}
}
}
ClientControl::Ping => {
debug!("ping received (pty)");
}
}
}
}
Ok(Message::Close(_)) => break,
Err(_) => break,
_ => {}
}
}
});
tokio::select! {
_ = &mut forward_handle => {
debug!("PTY→WS forward ended (pty): session={key}");
}
_ = &mut read_handle => {
debug!("WS→PTY read ended (pty): session={key}");
}
}
forward_handle.abort();
read_handle.abort();
info!("terminal WS disconnected (pty): session={session_id} — 会话进程保持常驻");
}