From dc0fc56fc42125cad4891a6537a7bad7e40df920 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 13 Feb 2026 11:16:47 +0000 Subject: [PATCH 1/4] Initial plan From efc7f0b80bd9044449aad0e74ab8eecb8e49851d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 13 Feb 2026 11:25:36 +0000 Subject: [PATCH 2/4] Implement session restoration from Docker containers - Add metadata labels to containers (session_id, owner_id, shell, etc.) - Restore sessions on server startup from existing containers - Attempt session restoration on WebSocket reconnection - Update container creation to include comprehensive labels - Fix pre-existing containers accessibility after upgrade Co-authored-by: Yashb404 <139128977+Yashb404@users.noreply.github.com> --- server/src/main.rs | 5 + server/src/services/websocket.rs | 194 +++++++++++++++++++++++++++++-- 2 files changed, 190 insertions(+), 9 deletions(-) diff --git a/server/src/main.rs b/server/src/main.rs index a596739..7e953b8 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -16,6 +16,7 @@ mod handlers { } use services::docker::start_background_reaper; +use services::websocket::restore_sessions_from_containers; #[tokio::main] async fn main() -> Result<(), Box> { @@ -24,6 +25,10 @@ async fn main() -> Result<(), Box> { // Setup database and Docker let state = config::setup_database_and_docker().await?; + // Restore sessions from existing Docker containers (critical for reconnecting to pre-existing containers) + println!("Restoring sessions from existing containers..."); + restore_sessions_from_containers(&state).await; + // Spawn background reaper let docker_reaper = state.docker.clone(); let sessions_reaper = state.sessions.clone(); diff --git a/server/src/services/websocket.rs b/server/src/services/websocket.rs index f75f45a..aa4e616 100644 --- a/server/src/services/websocket.rs +++ b/server/src/services/websocket.rs @@ -3,7 +3,7 @@ use axum::{ response::Response, http::StatusCode, }; -use bollard::container::{CreateContainerOptions, Config}; +use bollard::container::{CreateContainerOptions, Config, ListContainersOptions}; use bollard::models::{HostConfig, Mount, MountTypeEnum, MountTmpfsOptions}; use bollard::image::CreateImageOptions; use bollard::exec::{CreateExecOptions, StartExecResults}; @@ -16,6 +16,99 @@ use tower_sessions::Session; use crate::state::{AppState, SessionContext}; use crate::models::{User, AnalyticsEventType, log_analytics_event}; +/// Helper function to create container labels with session metadata for restoration +fn create_container_labels( + session_id: &str, + owner_id: Option, + project_owner_id: Option, + project_slug: Option<&str>, + shell: &str, + container_type: &str, // "builder" or "viewer" +) -> HashMap { + let mut labels = HashMap::from([ + ("managed_by".to_string(), "TryCli Studio".to_string()), + ("session_id".to_string(), session_id.to_string()), + ("shell".to_string(), shell.to_string()), + ("container_type".to_string(), container_type.to_string()), + ]); + + if let Some(id) = owner_id { + labels.insert("owner_id".to_string(), id.to_string()); + } + + if let Some(id) = project_owner_id { + labels.insert("project_owner_id".to_string(), id.to_string()); + } + + if let Some(slug) = project_slug { + labels.insert("project_slug".to_string(), slug.to_string()); + } + + labels +} + +/// Restore sessions from existing Docker containers on server startup +/// This allows pre-existing containers to be reconnected after server restart +pub async fn restore_sessions_from_containers(state: &AppState) { + let filters = HashMap::from([ + ("label".to_string(), vec!["managed_by=TryCli Studio".to_string()]) + ]); + + let opts = ListContainersOptions { + all: false, // Only running containers + filters, + ..Default::default() + }; + + match state.docker.list_containers(Some(opts)).await { + Ok(containers) => { + let mut restored = 0; + for container in containers { + if let Some(labels) = container.labels { + // Extract session metadata from labels + let session_id = labels.get("session_id").map(|s| s.clone()); + let shell = labels.get("shell").map(|s| s.clone()).unwrap_or_else(|| "/bin/bash".to_string()); + let owner_id = labels.get("owner_id").and_then(|s| s.parse::().ok()); + let project_owner_id = labels.get("project_owner_id").and_then(|s| s.parse::().ok()); + let project_slug = labels.get("project_slug").map(|s| s.clone()); + + if let (Some(session_id), Some(names)) = (session_id, container.names) { + let container_name = names.first() + .map(|n| n.trim_start_matches('/').to_string()) + .unwrap_or_default(); + + if !container_name.is_empty() { + // Restore session to in-memory map + let mut map = state.lock_sessions(); + + // Only restore if not already present (shouldn't happen, but be defensive) + if !map.contains_key(&session_id) { + map.insert(session_id.clone(), SessionContext { + container_name: container_name.clone(), + shell, + pending_image_tag: None, + owner_id, + project_owner_id, + is_publishing: false, + project_slug, + created_at: std::time::Instant::now(), // Reset timer + is_ws_connected: false, // Will be set to true on reconnection + }); + restored += 1; + println!("Restored session {} with container {}", session_id, container_name); + } + } + } + } + } + println!("Session restoration complete: {} sessions restored", restored); + } + Err(e) => { + eprintln!("Failed to restore sessions from containers: {}", e); + } + } +} + pub async fn ws_handler( ws: WebSocketUpgrade, Path(session_id): Path, @@ -25,6 +118,17 @@ pub async fn ws_handler( let user: Option = session.get("user").await.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let user_id = user.map(|u| u.id); + // Check if session exists in memory, if not, try to restore from Docker + let session_exists = { + let map = state.lock_sessions(); + map.contains_key(&session_id) + }; + + if !session_exists { + // Try to find and restore this specific session from Docker containers + restore_specific_session(&state, &session_id).await; + } + { let map = state.lock_sessions(); @@ -47,6 +151,57 @@ pub async fn ws_handler( Ok(ws.on_upgrade(move |socket| handle_socket(socket, state, session_id, user_id))) } +/// Attempt to restore a specific session from Docker containers +/// This is called when a client tries to connect to a session that isn't in memory +async fn restore_specific_session(state: &AppState, session_id: &str) { + let filters = HashMap::from([ + ("label".to_string(), vec![ + "managed_by=TryCli Studio".to_string(), + format!("session_id={}", session_id) + ]) + ]); + + let opts = ListContainersOptions { + all: false, // Only running containers + filters, + ..Default::default() + }; + + if let Ok(containers) = state.docker.list_containers(Some(opts)).await { + for container in containers { + if let Some(labels) = container.labels { + let shell = labels.get("shell").map(|s| s.clone()).unwrap_or_else(|| "/bin/bash".to_string()); + let owner_id = labels.get("owner_id").and_then(|s| s.parse::().ok()); + let project_owner_id = labels.get("project_owner_id").and_then(|s| s.parse::().ok()); + let project_slug = labels.get("project_slug").map(|s| s.clone()); + + if let Some(names) = container.names { + let container_name = names.first() + .map(|n| n.trim_start_matches('/').to_string()) + .unwrap_or_default(); + + if !container_name.is_empty() { + let mut map = state.lock_sessions(); + map.insert(session_id.to_string(), SessionContext { + container_name: container_name.clone(), + shell, + pending_image_tag: None, + owner_id, + project_owner_id, + is_publishing: false, + project_slug, + created_at: std::time::Instant::now(), + is_ws_connected: false, + }); + println!("Restored session {} from container {}", session_id, container_name); + return; + } + } + } + } + } +} + async fn handle_socket(mut socket: WebSocket, state: AppState, session_id: String, user_id: Option) { // Track if this is a first-time connection for view counting @@ -103,7 +258,13 @@ async fn handle_socket(mut socket: WebSocket, state: AppState, session_id: Strin if let Some(ctx) = map.get(&session_id) { // If we have an image tag but no container name, it's a viewer waiting to start if ctx.container_name.is_empty() && ctx.pending_image_tag.is_some() { - Some((ctx.pending_image_tag.clone().unwrap(), ctx.shell.clone())) + Some(( + ctx.pending_image_tag.clone().unwrap(), + ctx.shell.clone(), + ctx.owner_id, + ctx.project_owner_id, + ctx.project_slug.clone(), + )) } else { None } @@ -112,15 +273,22 @@ async fn handle_socket(mut socket: WebSocket, state: AppState, session_id: Strin } }; - if let Some((image_tag, shell)) = pending_spawn { + if let Some((image_tag, shell, owner_id, project_owner_id, project_slug)) = pending_spawn { // Perform the spawn that used to be in get_project let container_name = format!("trycli-studio-viewer-{}", Uuid::new_v4()); + let labels = create_container_labels( + &session_id, + owner_id, + project_owner_id, + project_slug.as_deref(), + &shell, + "viewer", + ); + let config = Config { image: Some(image_tag), - labels: Some(HashMap::from([ - ("managed_by".to_string(), "TryCli Studio".to_string()) - ])), + labels: Some(labels), tty: Some(true), user: Some("root".to_string()), // FIX: Run sleep infinity as PID 1. This uses almost 0 CPU/RAM. @@ -332,6 +500,16 @@ async fn run_setup_wizard(mut socket: WebSocket, state: AppState, session_id: St } let container_name = format!("trycli-studio-session-{}", Uuid::new_v4()); + + let labels = create_container_labels( + &session_id, + _user_id, + None, // Builder sessions don't have project context yet + None, + final_shell, + "builder", + ); + let config = Config { image: Some(image.to_string()), tty: Some(true), @@ -343,9 +521,7 @@ async fn run_setup_wizard(mut socket: WebSocket, state: AppState, session_id: St "LC_ALL=C.UTF-8".to_string(), "TERM=xterm-256color".to_string() ]), - labels: Some(HashMap::from([ - ("managed_by".to_string(), "TryCli Studio".to_string()) - ])), + labels: Some(labels), host_config: Some(HostConfig { runtime: Some("runsc".to_string()), memory: Some(512 * 1024 * 1024), From 773b3be19bfb7e99b7cd30955ccf0c317f442ab5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 13 Feb 2026 11:26:06 +0000 Subject: [PATCH 3/4] Security: sanitize Docker error messages sent to clients - Replace detailed Docker/bollard errors with generic messages - Keep detailed error logging server-side for debugging - Prevent leaking internal operational details to clients Co-authored-by: Yashb404 <139128977+Yashb404@users.noreply.github.com> --- server/src/services/websocket.rs | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/server/src/services/websocket.rs b/server/src/services/websocket.rs index aa4e616..2622b67 100644 --- a/server/src/services/websocket.rs +++ b/server/src/services/websocket.rs @@ -358,12 +358,10 @@ async fn handle_socket(mut socket: WebSocket, state: AppState, session_id: Strin if let Err(e) = state.docker.start_container::(&container_name, None).await { - // Surface detailed error to client and logs - let msg = format!( - "\r\n\x1b[31m[!] Failed to start viewer container: {}\x1b[0m\r\n", - e - ); + // Log detailed error server-side only eprintln!("Viewer start error for session {}: {}", session_id, e); + // Send generic error message to client + let msg = "\r\n\x1b[31m[!] Failed to start viewer container. Please try again later.\x1b[0m\r\n"; let _ = socket.send(Message::Text(msg.into())).await; return; } @@ -376,11 +374,10 @@ async fn handle_socket(mut socket: WebSocket, state: AppState, session_id: Strin } } Err(e) => { - let msg = format!( - "\r\n\x1b[31m[!] Failed to create viewer container: {}\x1b[0m\r\n", - e - ); + // Log detailed error server-side only eprintln!("Viewer create error for session {}: {}", session_id, e); + // Send generic error message to client + let msg = "\r\n\x1b[31m[!] Failed to create viewer container. Please try again later.\x1b[0m\r\n"; let _ = socket.send(Message::Text(msg.into())).await; return; } From 2d7d9a303b6a50f62e9af24330742f81c015ca2b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 13 Feb 2026 11:27:57 +0000 Subject: [PATCH 4/4] Address code review feedback - Fix timer reset issue: preserve container age on restoration - Replace println!/eprintln! with tracing::info!/error! - Use container creation timestamp for accurate session age tracking Co-authored-by: Yashb404 <139128977+Yashb404@users.noreply.github.com> --- server/src/main.rs | 2 +- server/src/services/websocket.rs | 48 ++++++++++++++++++++++++++------ 2 files changed, 41 insertions(+), 9 deletions(-) diff --git a/server/src/main.rs b/server/src/main.rs index 7e953b8..07439a5 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -26,7 +26,7 @@ async fn main() -> Result<(), Box> { let state = config::setup_database_and_docker().await?; // Restore sessions from existing Docker containers (critical for reconnecting to pre-existing containers) - println!("Restoring sessions from existing containers..."); + tracing::info!("Restoring sessions from existing containers..."); restore_sessions_from_containers(&state).await; // Spawn background reaper diff --git a/server/src/services/websocket.rs b/server/src/services/websocket.rs index 2622b67..760077e 100644 --- a/server/src/services/websocket.rs +++ b/server/src/services/websocket.rs @@ -83,6 +83,23 @@ pub async fn restore_sessions_from_containers(state: &AppState) { // Only restore if not already present (shouldn't happen, but be defensive) if !map.contains_key(&session_id) { + // Calculate elapsed time from container creation + let created_at = if let Some(created_ts) = container.created { + // Container age in seconds + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() as i64; + let age_secs = now - created_ts; + + // Set created_at to approximate original time + std::time::Instant::now() + .checked_sub(std::time::Duration::from_secs(age_secs.max(0) as u64)) + .unwrap_or_else(|| std::time::Instant::now()) + } else { + std::time::Instant::now() + }; + map.insert(session_id.clone(), SessionContext { container_name: container_name.clone(), shell, @@ -91,20 +108,20 @@ pub async fn restore_sessions_from_containers(state: &AppState) { project_owner_id, is_publishing: false, project_slug, - created_at: std::time::Instant::now(), // Reset timer + created_at, is_ws_connected: false, // Will be set to true on reconnection }); restored += 1; - println!("Restored session {} with container {}", session_id, container_name); + tracing::info!("Restored session {} with container {}", session_id, container_name); } } } } } - println!("Session restoration complete: {} sessions restored", restored); + tracing::info!("Session restoration complete: {} sessions restored", restored); } Err(e) => { - eprintln!("Failed to restore sessions from containers: {}", e); + tracing::error!("Failed to restore sessions from containers: {}", e); } } } @@ -181,6 +198,21 @@ async fn restore_specific_session(state: &AppState, session_id: &str) { .unwrap_or_default(); if !container_name.is_empty() { + // Calculate elapsed time from container creation + let created_at = if let Some(created_ts) = container.created { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() as i64; + let age_secs = now - created_ts; + + std::time::Instant::now() + .checked_sub(std::time::Duration::from_secs(age_secs.max(0) as u64)) + .unwrap_or_else(|| std::time::Instant::now()) + } else { + std::time::Instant::now() + }; + let mut map = state.lock_sessions(); map.insert(session_id.to_string(), SessionContext { container_name: container_name.clone(), @@ -190,10 +222,10 @@ async fn restore_specific_session(state: &AppState, session_id: &str) { project_owner_id, is_publishing: false, project_slug, - created_at: std::time::Instant::now(), + created_at, is_ws_connected: false, }); - println!("Restored session {} from container {}", session_id, container_name); + tracing::info!("Restored session {} from container {}", session_id, container_name); return; } } @@ -359,7 +391,7 @@ async fn handle_socket(mut socket: WebSocket, state: AppState, session_id: Strin state.docker.start_container::(&container_name, None).await { // Log detailed error server-side only - eprintln!("Viewer start error for session {}: {}", session_id, e); + tracing::error!("Viewer start error for session {}: {}", session_id, e); // Send generic error message to client let msg = "\r\n\x1b[31m[!] Failed to start viewer container. Please try again later.\x1b[0m\r\n"; let _ = socket.send(Message::Text(msg.into())).await; @@ -375,7 +407,7 @@ async fn handle_socket(mut socket: WebSocket, state: AppState, session_id: Strin } Err(e) => { // Log detailed error server-side only - eprintln!("Viewer create error for session {}: {}", session_id, e); + tracing::error!("Viewer create error for session {}: {}", session_id, e); // Send generic error message to client let msg = "\r\n\x1b[31m[!] Failed to create viewer container. Please try again later.\x1b[0m\r\n"; let _ = socket.send(Message::Text(msg.into())).await;