diff --git a/Cargo.toml b/Cargo.toml index 92ab29d..de34378 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,6 +42,7 @@ chrono = { version = "0.4", features = ["serde"] } directories = "5" bytes = "1" tokio-util = "0.7" +futures-util = "0.3" async-trait = "0.1" dotenvy = "0.15" diff --git a/frontend/src/App.svelte b/frontend/src/App.svelte index 3fac963..9992adf 100644 --- a/frontend/src/App.svelte +++ b/frontend/src/App.svelte @@ -10,8 +10,8 @@ import SettingsPanel from './lib/components/SettingsPanel.svelte'; // Configuration constants - const POLL_INTERVAL_MS = 500; // Progress polling interval during job execution const STORAGE_KEY_PERSON = 'immich-timelapse-selected-person'; + const WS_RECONNECT_DELAY_MS = 1000; // WebSocket reconnection delay // Load persisted state from localStorage function loadPersistedPersonId() { @@ -44,7 +44,8 @@ let initialSelectedPersonId = $state(loadPersistedPersonId()); let jobStatus = $state('idle'); let progress = $state({ completed: 0, total: 0, message: '' }); - let pollInterval = $state(null); + let ws = $state(null); + let wsReconnectTimeout = $state(null); // View state management for gallery let currentView = $state('main'); // 'main' | 'gallery' @@ -130,66 +131,82 @@ function handleJobUpdate(data) { jobStatus = data.status; progress = data; - - // Start polling if job just started - if (data.status === 'running' || data.status === 'compiling_video') { - startPolling(); - } } async function checkAndPollProgress() { - try { - const res = await fetch('/api/progress'); - const data = await res.json(); - - // Only restore status if a job is actively running, or if we're not in idle state. - // This prevents restoring 'completed' status after we've manually dismissed it. - const isActiveJob = data.status === 'running' || data.status === 'compiling_video' || data.status === 'cancelling'; - if (isActiveJob || jobStatus !== 'idle') { - jobStatus = data.status; - progress = data; - } - - // If a job is running, start polling - if (isActiveJob) { - startPolling(); - } - } catch (e) { - console.error('Failed to check progress:', e); - } + // Connect WebSocket if not already connected + connectWebSocket(); } - async function pollProgress() { - try { - const res = await fetch('/api/progress'); - const data = await res.json(); + function connectWebSocket() { + if (ws && (ws.readyState === WebSocket.CONNECTING || ws.readyState === WebSocket.OPEN)) { + return; // Already connected or connecting + } - const previousStatus = jobStatus; - jobStatus = data.status; - progress = data; + const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; + const wsUrl = `${protocol}//${window.location.host}/api/ws`; - // Stop polling when job completes - if (data.status === 'completed' || data.status === 'cancelled' || data.status === 'error' || data.status === 'idle') { - stopPolling(); - // Refresh output folders when job finishes (was running before) - if (previousStatus === 'running' || previousStatus === 'compiling_video' || previousStatus === 'cancelling') { - outputFolderRefreshKey++; + ws = new WebSocket(wsUrl); + + ws.onopen = () => { + console.log('WebSocket connected'); + }; + + ws.onmessage = (event) => { + try { + const data = JSON.parse(event.data); + const previousStatus = jobStatus; + + // Only restore status if a job is actively running, or if we're not in idle state + const isActiveJob = data.status === 'running' || data.status === 'compiling_video' || data.status === 'cancelling'; + if (isActiveJob || jobStatus !== 'idle') { + jobStatus = data.status; + progress = data; } + + // Refresh output folders when job finishes (was running before) + if (data.status === 'completed' || data.status === 'cancelled' || data.status === 'error') { + if (previousStatus === 'running' || previousStatus === 'compiling_video' || previousStatus === 'cancelling') { + outputFolderRefreshKey++; + } + } + } catch (e) { + console.error('Failed to parse WebSocket message:', e); } - } catch (e) { - console.error('Poll failed:', e); + }; + + ws.onclose = () => { + console.log('WebSocket disconnected'); + ws = null; + // Reconnect if connection was established (connectionOk is true) + if (connectionOk) { + scheduleReconnect(); + } + }; + + ws.onerror = (error) => { + console.error('WebSocket error:', error); + }; + } + + function scheduleReconnect() { + if (wsReconnectTimeout) return; // Already scheduled + wsReconnectTimeout = setTimeout(() => { + wsReconnectTimeout = null; + if (connectionOk) { + connectWebSocket(); + } + }, WS_RECONNECT_DELAY_MS); + } + + function disconnectWebSocket() { + if (wsReconnectTimeout) { + clearTimeout(wsReconnectTimeout); + wsReconnectTimeout = null; } - } - - function startPolling() { - if (pollInterval) return; // Already polling - pollInterval = setInterval(pollProgress, POLL_INTERVAL_MS); - } - - function stopPolling() { - if (pollInterval) { - clearInterval(pollInterval); - pollInterval = null; + if (ws) { + ws.close(); + ws = null; } } @@ -198,7 +215,7 @@ }); onDestroy(() => { - stopPolling(); + disconnectWebSocket(); }); diff --git a/frontend/src/lib/components/ProcessingControls.svelte b/frontend/src/lib/components/ProcessingControls.svelte index 7963c34..e3b5806 100644 --- a/frontend/src/lib/components/ProcessingControls.svelte +++ b/frontend/src/lib/components/ProcessingControls.svelte @@ -71,7 +71,7 @@ const data = await res.json(); if (res.ok && data.success) { - // Notify parent to start polling + // Notify parent of job start onupdate?.({ status: 'running', completed: 0, diff --git a/src/web/handlers/mod.rs b/src/web/handlers/mod.rs index 47730b2..39ac9ab 100644 --- a/src/web/handlers/mod.rs +++ b/src/web/handlers/mod.rs @@ -12,6 +12,7 @@ mod health; mod output; mod people; mod processing; +mod ws; use crate::web::state::AppState; use axum::{ @@ -34,6 +35,7 @@ use output::{ }; use people::{get_people, get_person_asset_count, get_person_thumbnail}; use processing::{cancel_processing, get_progress, start_processing}; +use ws::ws_handler; // Re-export types that may be needed by other modules pub use config::ConfigResponse; @@ -72,6 +74,7 @@ pub fn create_router(state: AppState) -> Router { get(get_person_asset_count), ) .route("/api/progress", get(get_progress)) + .route("/api/ws", get(ws_handler)) .route("/api/start", post(start_processing)) .route("/api/cancel", post(cancel_processing)) .route("/api/output", get(list_output_folders)) diff --git a/src/web/handlers/ws.rs b/src/web/handlers/ws.rs new file mode 100644 index 0000000..552e61a --- /dev/null +++ b/src/web/handlers/ws.rs @@ -0,0 +1,100 @@ +//! WebSocket handler for real-time progress updates. + +use crate::web::handlers::ProgressResponse; +use crate::web::state::{AppState, JobStatus}; +use axum::{ + extract::{ + ws::{Message, WebSocket}, + State, WebSocketUpgrade, + }, + response::IntoResponse, +}; +use futures_util::{stream::StreamExt, SinkExt}; + +/// WebSocket upgrade handler. +pub async fn ws_handler( + ws: WebSocketUpgrade, + State(state): State, +) -> impl IntoResponse { + ws.on_upgrade(|socket| handle_socket(socket, state)) +} + +/// Handle an individual WebSocket connection. +async fn handle_socket(socket: WebSocket, state: AppState) { + let (mut sender, mut receiver) = socket.split(); + + // Subscribe to progress updates + let mut progress_rx = state.progress_tx.subscribe(); + + // Send current progress immediately on connect + { + let progress = state.progress.read().await; + let response = progress_to_response(&progress); + if let Ok(json) = serde_json::to_string(&response) { + if sender.send(Message::Text(json.into())).await.is_err() { + return; + } + } + } + + // Spawn a task to handle incoming messages (ping/pong, close) + let mut recv_task = tokio::spawn(async move { + while let Some(msg) = receiver.next().await { + match msg { + Ok(Message::Close(_)) => break, + Err(_) => break, + _ => {} // Ignore other messages + } + } + }); + + // Send progress updates to the client + let mut send_task = tokio::spawn(async move { + while let Ok(progress) = progress_rx.recv().await { + let response = progress_to_response(&progress); + match serde_json::to_string(&response) { + Ok(json) => { + if sender.send(Message::Text(json.into())).await.is_err() { + break; + } + } + Err(_) => break, + } + } + }); + + // Wait for either task to complete + tokio::select! { + _ = &mut recv_task => { + send_task.abort(); + } + _ = &mut send_task => { + recv_task.abort(); + } + } +} + +/// Convert Progress to ProgressResponse for JSON serialization. +fn progress_to_response(progress: &crate::web::state::Progress) -> ProgressResponse { + use crate::web::handlers::SkipStatsResponse; + + let status_str = match &progress.status { + JobStatus::Idle => "idle", + JobStatus::Running => "running", + JobStatus::Cancelling => "cancelling", + JobStatus::CompilingVideo => "compiling_video", + JobStatus::Completed => "completed", + JobStatus::Cancelled => "cancelled", + JobStatus::Error(_) => "error", + }; + + ProgressResponse { + status: status_str.to_string(), + completed: progress.completed, + total: progress.total, + message: progress.message.clone(), + skip_stats: SkipStatsResponse::from(&progress.skip_stats), + person_id: progress.person_id.clone(), + person_name: progress.person_name.clone(), + } +}