fix process manager follow stream regression

Author: wkoszek

src/cli/mod.rs

@@ -4,6 +4,13 @@ pub mod serve;
 use crate::cli::client::{new_client, new_streaming_client, DEFAULT_ADDR};
 use anyhow::{Context, Result};
 
+fn is_active_session(session: &serde_json::Value) -> bool {
+    matches!(
+        session["state"].as_str(),
+        Some("starting" | "running" | "stopping")
+    )
+}
+
 async fn read_json_response(resp: reqwest::Response) -> Result<serde_json::Value> {
     let status = resp.status();
     let body = resp.json::<serde_json::Value>().await?;
@@ -38,9 +45,10 @@ async fn resolve_id(id: Option<String>) -> Result<String> {
     let arr = sessions["sessions"]
         .as_array()
         .ok_or_else(|| anyhow::anyhow!("unexpected response from daemon"))?;
-    match arr.len() {
+    let active: Vec<&serde_json::Value> = arr.iter().filter(|s| is_active_session(s)).collect();
+    match active.len() {
         0 => anyhow::bail!("no active sessions"),
-        1 => arr[0]["id"]
+        1 => active[0]["id"]
             .as_str()
             .map(|s| s.to_string())
             .ok_or_else(|| anyhow::anyhow!("session has no id")),
@@ -131,6 +139,7 @@ pub async fn cmd_stop_all() -> Result<()> {
         .as_array()
         .map(|arr| {
             arr.iter()
+                .filter(|s| is_active_session(s))
                 .filter_map(|s| s["id"].as_str().map(|s| s.to_string()))
                 .collect()
         })

src/daemon/handlers.rs

@@ -10,7 +10,6 @@ use futures::StreamExt;
 use serde::Deserialize;
 use std::collections::HashMap;
 use std::path::{Path as FsPath, PathBuf};
-use tokio_stream::wrappers::BroadcastStream;
 use uuid::Uuid;
 
 use crate::daemon::buffer::{LogEntry, Stream};
@@ -327,44 +326,57 @@ pub async fn get_logs(
         .ok_or_else(|| AppError::NotFound("session not found".into()))?;
 
     // Streaming follow mode
-    // Both stream_name and format are owned Strings — safe to move into the closure
-    let stream_name_clone = stream_name.clone();
     let format_clone = format.clone();
-    let stream_mode = stream_name.clone();
-    let stream = BroadcastStream::new(rx).flat_map(move |msg| {
-        let line = match msg {
-            Ok(entry) => {
-                if entry.seq < snapshot_next_seq {
-                    return futures::stream::empty().left_stream();
-                }
-                // Filter: for blended, show all; for stdout/stderr, show only that stream; for unknown, show nothing
-                let should_include = match stream_name_clone.as_str() {
-                    "blended" => true,
-                    "stdout" => entry.stream == Stream::Stdout,
-                    "stderr" => entry.stream == Stream::Stderr,
-                    _ => false, // unknown stream name → no entries
-                };
-                if !should_include {
-                    return futures::stream::empty().left_stream();
-                }
-                format_entry(&entry, &stream_mode, &format_clone)
-            }
-            Err(tokio_stream::wrappers::errors::BroadcastStreamRecvError::Lagged(n)) => {
-                if format_clone == "text" {
-                    format!("[korun: dropped {n} lines]\n")
-                } else {
-                    let e = serde_json::json!({
-                        "seq": 0u64,
-                        "ts": Utc::now().to_rfc3339(),
-                        "stream": "system",
-                        "line": format!("dropped {n} lines"),
-                    });
-                    format!("{e}\n")
+    let mgr_for_follow = mgr.clone();
+    let stream = futures::stream::unfold(
+        FollowState {
+            rx,
+            mgr: mgr_for_follow,
+            session_id: id,
+            stream_name: stream_name.clone(),
+            format: format_clone,
+            cutoff_seq: snapshot_next_seq,
+            last_seen_seq: snapshot_next_seq,
+        },
+        |mut state| async move {
+            loop {
+                tokio::select! {
+                    msg = state.rx.recv() => {
+                        match msg {
+                            Ok(entry) => {
+                                if entry.seq < state.cutoff_seq {
+                                    continue;
+                                }
+                                if !should_include_stream(&state.stream_name, entry.stream) {
+                                    continue;
+                                }
+                                state.last_seen_seq = state.last_seen_seq.max(entry.seq.saturating_add(1));
+                                let line = format_entry(&entry, &state.stream_name, &state.format);
+                                return Some((Ok::<_, std::convert::Infallible>(line), state));
+                            }
+                            Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
+                                let line = format_lag_notice(n, &state.format);
+                                return Some((Ok::<_, std::convert::Infallible>(line), state));
+                            }
+                            Err(tokio::sync::broadcast::error::RecvError::Closed) => return None,
+                        }
+                    }
+                    _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
+                        let (is_terminal, next_seq) = match state
+                            .mgr
+                            .with(&state.session_id, |s| (matches!(s.state, SessionState::Exited | SessionState::Failed), s.next_seq()))
+                        {
+                            Some(values) => values,
+                            None => return None,
+                        };
+                        if is_terminal && next_seq <= state.last_seen_seq {
+                            return None;
+                        }
+                    }
                 }
             }
-        };
-        futures::stream::once(async move { Ok::<_, std::convert::Infallible>(line) }).right_stream()
-    });
+        },
+    );
 
     // First send buffered entries, then stream new ones.
     // Use into_iter() + move to avoid capturing &format (which would prevent 'static bound).
@@ -454,6 +466,39 @@ fn format_entry(entry: &LogEntry, requested_stream: &str, format: &str) -> Strin
     }
 }
 
+fn should_include_stream(stream_name: &str, stream: Stream) -> bool {
+    match stream_name {
+        "blended" => true,
+        "stdout" => stream == Stream::Stdout,
+        "stderr" => stream == Stream::Stderr,
+        _ => false,
+    }
+}
+
+fn format_lag_notice(n: u64, format: &str) -> String {
+    if format == "text" {
+        format!("[korun: dropped {n} lines]\n")
+    } else {
+        let e = serde_json::json!({
+            "seq": 0u64,
+            "ts": Utc::now().to_rfc3339(),
+            "stream": "system",
+            "line": format!("dropped {n} lines"),
+        });
+        format!("{e}\n")
+    }
+}
+
+struct FollowState {
+    rx: tokio::sync::broadcast::Receiver<LogEntry>,
+    mgr: AppState,
+    session_id: Uuid,
+    stream_name: String,
+    format: String,
+    cutoff_seq: u64,
+    last_seen_seq: u64,
+}
+
 fn validate_create_session_request(
     command: &[String],
     cwd: &str,

tests/http_test.rs

@@ -5,6 +5,7 @@ use korun::daemon::session_mgr::SessionManager;
 use serde_json::Value;
 use tempfile::tempdir;
 use tower::util::ServiceExt;
+use tokio::time::{timeout, Duration};
 
 #[tokio::test]
 async fn healthz_returns_200() {
@@ -156,3 +157,54 @@ async fn create_session_retains_watcher_handle() {
     let id: uuid::Uuid = json["id"].as_str().unwrap().parse().unwrap();
     assert!(mgr.with(&id, |s| s.watcher.is_some()).unwrap());
 }
+
+#[tokio::test]
+async fn follow_logs_closes_after_process_exits() {
+    let mgr = SessionManager::new();
+    let app = build_router(mgr);
+
+    let body = serde_json::json!({
+        "command": ["sh", "-c", "echo hi"]
+    });
+
+    let response = app
+        .clone()
+        .oneshot(
+            Request::builder()
+                .method("POST")
+                .uri("/v1/sessions")
+                .header("content-type", "application/json")
+                .body(Body::from(serde_json::to_vec(&body).unwrap()))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(response.status(), StatusCode::CREATED);
+    let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
+        .await
+        .unwrap();
+    let json: Value = serde_json::from_slice(&bytes).unwrap();
+    let id = json["id"].as_str().unwrap();
+
+    let response = app
+        .oneshot(
+            Request::builder()
+                .uri(format!("/v1/sessions/{id}/logs?follow=1&format=text"))
+                .body(Body::empty())
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(response.status(), StatusCode::OK);
+    let bytes = timeout(
+        Duration::from_secs(2),
+        axum::body::to_bytes(response.into_body(), usize::MAX),
+    )
+    .await
+    .expect("follow stream should close after process exit")
+    .unwrap();
+    let text = String::from_utf8(bytes.to_vec()).unwrap();
+    assert!(text.contains("[stdout] hi"));
+}