Skip to main content

bhtune_server/routes/
stream.rs

1//! `GET /api/runs/{id}/stream` pushes per-tick engine state to the browser over Server-Sent
2//! Events as a run progresses, so the live trend chart does not have to re-fetch the whole
3//! (ever-growing) `samples` array once a second.
4//!
5//! SSE, not WebSocket: the flow is strictly server -> client, and SSE gives every browser
6//! automatic reconnection on a dropped connection, survives ordinary HTTP proxies, and is
7//! trivially inspectable with `curl -N`.
8//!
9//! Implemented as an internal poll of `tune_samples`/`tune_runs`, **not** a broadcast channel
10//! threaded through `bhtune-runtime`'s already-tested tick loop -- deliberately, so this
11//! endpoint adds zero risk to the shared CLI/server tune-execution code path
12//! (`run_polling_loop` keeps its existing, already-proven signature and test suite untouched).
13//! At the expected data volumes (about 9,000 samples for a two-hour run), polling the
14//! database every `POLL_INTERVAL` is negligible cost.
15
16use std::convert::Infallible;
17use std::time::Duration;
18use std::time::Instant;
19
20use async_stream::stream;
21use axum::Router;
22use axum::extract::{Path, State};
23use axum::response::IntoResponse;
24use axum::response::sse::{Event, KeepAlive, Sse};
25use axum::routing::get;
26use bhtune_db::models::{TuneOutcome, TuneRunRow, TuneSampleRow};
27use serde::Serialize;
28use utoipa::ToSchema;
29
30use crate::error::ApiError;
31use crate::routes::history::{InitialReadingsResponse, SampleResponse};
32use crate::state::AppState;
33
34/// How often the stream re-polls the database for newly recorded samples or a changed run
35/// outcome. Well under the CLI's own default 800ms tick interval, so a connected client never
36/// perceives added latency beyond this poll interval itself; see the module doc for why
37/// polling the database is an acceptable (not merely expedient) choice at this project's data
38/// volumes.
39const POLL_INTERVAL: Duration = Duration::from_millis(300);
40
41/// The final event emitted on every `GET /api/runs/{id}/stream` connection, named `done`,
42/// immediately before the stream closes. Deliberately carries only the outcome rather than
43/// duplicating the full run detail: the frontend already has `GET /api/runs/{id}` (via
44/// `useRun`) for the config/results/write-back-audit shape, so `done` just signals "stop
45/// listening and go refetch that" rather than growing this endpoint a second response shape.
46#[derive(Debug, Serialize, ToSchema)]
47pub(crate) struct RunStreamDone {
48    outcome: TuneOutcome,
49}
50
51/// The stream's item type is `Result<Event, Infallible>` -- see the module doc for why a
52/// real transport-level error is never produced (a DB failure mid-stream is represented as a
53/// named `error` SSE event instead, followed by ending the stream). A tiny helper rather
54/// than a turbofish (`Ok::<Event, Infallible>(event)`) at every `yield` site, since the
55/// `stream!` macro can't infer `Infallible` from context on its own (every `yield` only ever
56/// produces `Ok`, so there's no `Err(...)` arm anywhere for type inference to anchor on).
57fn ok_event(event: Event) -> Result<Event, Infallible> {
58    Ok(event)
59}
60
61/// Stream per-tick engine state for one run over Server-Sent Events.
62///
63/// `GET /api/runs/{id}/stream` -- 404 if no run has that id. Emits one `initial` event (JSON
64/// body: [`InitialReadingsResponse`]) as soon as the driver's initial snapshot is persisted,
65/// then a `sample` event (JSON body: [`SampleResponse`], the same shape
66/// `GET /api/runs/{id}`'s `samples` array already uses) for every tick recorded so far and
67/// every new tick recorded while connected, followed by exactly one final `done` event (JSON
68/// body: [`RunStreamDone`]) once the run reaches a terminal outcome -- after which the
69/// connection closes. Safe to open at any point in a run's lifecycle, including after it has
70/// already finished: the initial snapshot (when available) and every sample are replayed once,
71/// immediately followed by `done`.
72#[utoipa::path(
73    get,
74    path = "/api/runs/{id}/stream",
75    tag = "runs",
76    params(
77        ("id" = i64, Path, description = "Run id"),
78    ),
79    responses(
80        (
81            status = 200,
82            description = "A `text/event-stream` with an optional `initial` event \
83                (data: InitialReadingsResponse), `sample` events (data: SampleResponse), \
84                and one final `done` event (data: RunStreamDone).",
85            content_type = "text/event-stream",
86            body = SampleResponse,
87        ),
88        (status = 404, description = "No run with that id.", body = crate::error::ErrorBody),
89    ),
90)]
91pub(crate) async fn stream_run(
92    State(state): State<AppState>,
93    Path(run_id): Path<i64>,
94) -> Result<impl IntoResponse, ApiError> {
95    stream_run_with_permit(state, run_id, None, None).await
96}
97
98pub(crate) async fn stream_run_with_permit(
99    state: AppState,
100    run_id: i64,
101    permit: Option<Vec<tokio::sync::OwnedSemaphorePermit>>,
102    timeout: Option<Duration>,
103) -> Result<impl IntoResponse, ApiError> {
104    if TuneRunRow::get(&state.pool, run_id).await?.is_none() {
105        return Err(ApiError::NotFound(format!("no run with id {run_id}")));
106    }
107
108    let pool = state.pool.clone();
109    let events = stream! {
110        let _permit = permit;
111        let started = Instant::now();
112        // `-1`: nothing sent yet, so the first poll below fetches every tick recorded so
113        // far -- see `TuneSampleRow::list_for_run_since`'s own doc comment for why `-1` is a
114        // safe "everything" sentinel (`tick` is always `>= 0`).
115        let mut last_tick: i64 = -1;
116        let mut sent_initial = false;
117        loop {
118            if timeout.is_some_and(|limit| started.elapsed() >= limit) {
119                yield ok_event(Event::default().event("error").data("SSE demo time limit exceeded"));
120                break;
121            }
122            let run_outcome = match TuneRunRow::get(&pool, run_id).await {
123                Ok(Some(run)) => {
124                    if !sent_initial && let Some(initial) = run.initial_readings.as_ref() {
125                        let response = InitialReadingsResponse::from(initial.clone());
126                        match Event::default().event("initial").json_data(response) {
127                            Ok(event) => {
128                                sent_initial = true;
129                                yield ok_event(event);
130                            }
131                            Err(err) => tracing::error!(
132                                run_id,
133                                error = %err,
134                                "failed to encode initial readings as an SSE event"
135                            ),
136                        }
137                    }
138                    Some(run.outcome)
139                }
140                // Still running (or -- defensively -- the run vanished from under us mid-
141                // stream, which cannot happen in practice since nothing ever deletes a
142                // `tune_runs` row): continue to the sample query.
143                Ok(_) => None,
144                Err(err) => {
145                    tracing::error!(
146                        run_id,
147                        error = %err,
148                        "failed to poll tune_runs for the run stream; ending the stream"
149                    );
150                    yield ok_event(Event::default().event("error").data(err.to_string()));
151                    break;
152                }
153            };
154
155            match TuneSampleRow::list_for_run_since(&pool, run_id, last_tick).await {
156                Ok(samples) => {
157                    for sample in &samples {
158                        last_tick = sample.tick_index;
159                        let response = SampleResponse::from(sample);
160                        match Event::default().event("sample").json_data(response) {
161                            Ok(event) => yield ok_event(event),
162                            Err(err) => tracing::error!(
163                                run_id,
164                                error = %err,
165                                "failed to encode a tune sample as an SSE event"
166                            ),
167                        }
168                    }
169                }
170                Err(err) => {
171                    tracing::error!(
172                        run_id,
173                        error = %err,
174                        "failed to poll tune_samples for the run stream; ending the stream"
175                    );
176                    yield ok_event(Event::default().event("error").data(err.to_string()));
177                    break;
178                }
179            }
180
181            if let Some(outcome) = run_outcome
182                && outcome != TuneOutcome::Running
183            {
184                let done = RunStreamDone { outcome };
185                if let Ok(event) = Event::default().event("done").json_data(done) {
186                    yield ok_event(event);
187                }
188                break;
189            }
190
191            tokio::time::sleep(POLL_INTERVAL).await;
192        }
193    };
194
195    Ok(Sse::new(events).keep_alive(KeepAlive::default()))
196}
197
198pub fn router() -> Router<AppState> {
199    Router::new().route("/api/runs/{id}/stream", get(stream_run))
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205    use axum::body::{Body, to_bytes};
206    use axum::http::{Request, StatusCode};
207    use bhtune_core::{ControllerType, LoopConfig, LoopTags, MrftState, ProcessType, Tick};
208    use bhtune_db::models::{SampleQuality, TuneDriver, TuneRunInitialReadings};
209    use chrono::Utc;
210    use tower::ServiceExt;
211
212    async fn seed_running_run(state: &AppState, name: &str) -> i64 {
213        let template_row =
214            bhtune_db::models::DcsTemplateRow::get_by_name(&state.pool, "Yokogawa CentumVP")
215                .await
216                .unwrap()
217                .unwrap();
218        let template = template_row.template;
219        let config = LoopConfig {
220            process_type: ProcessType::Flow,
221            controller_type: ControllerType::Pi,
222            relay_amp_percent: 5.0,
223            num_cycles_skip: 1,
224            num_cycles_count: 3,
225            noise_protection_secs: 0,
226            mrft_delay_secs: 0,
227        };
228        let tags = LoopTags::derive_from_pv_tag(&format!("{name}.PV"), &template);
229        let run = TuneRunRow::start(
230            &state.pool,
231            None,
232            name,
233            TuneDriver::Simulator,
234            config,
235            template_row.origin,
236            &template,
237            &tags,
238            Utc::now(),
239        )
240        .await
241        .unwrap();
242        run.id
243    }
244
245    async fn insert_sample(state: &AppState, run_id: i64, tick: i64) {
246        let sample = Tick {
247            time: Utc::now(),
248            pv: 50.0 + tick as f32,
249        };
250        let mrft_state = MrftState {
251            hysteresis: 1.0,
252            mv_value_current: 55.0,
253            mv_sign_next_step: 1,
254            counter_all_switches: tick as u32,
255            cycles_completed: 0,
256            cycles_remaining: 2,
257        };
258        TuneSampleRow::insert(
259            &state.pool,
260            run_id,
261            tick,
262            sample,
263            mrft_state,
264            SampleQuality::Good,
265        )
266        .await
267        .unwrap();
268    }
269
270    async fn record_initial_readings(state: &AppState, run_id: i64) {
271        TuneRunRow::record_initial_readings(
272            &state.pool,
273            run_id,
274            TuneRunInitialReadings {
275                pv_ini: 48.0,
276                mv_ini: 42.0,
277                mv_range_low: 0.0,
278                mv_range_high: 100.0,
279                pv_range_high: 100.0,
280                pv_range_low: 0.0,
281                controller_direction: bhtune_core::ControllerDirection::Reverse,
282                mode_raw: Some("AUTO".to_string()),
283                mode_attribute_raw: None,
284                setpoint_ini: Some(50.0),
285            },
286        )
287        .await
288        .unwrap();
289    }
290
291    /// Parses a `text/event-stream` body into `(event, data)` pairs, in order. A minimal,
292    /// deliberately non-general parser matching exactly the shape this module's own handler
293    /// produces (one `event:`/`data:` line pair per event, `\n\n`-terminated) -- good enough
294    /// for asserting on our own output without pulling in a full SSE-client crate just for
295    /// tests.
296    fn parse_sse(body: &str) -> Vec<(String, String)> {
297        body.split("\n\n")
298            .filter(|chunk| !chunk.trim().is_empty())
299            .map(|chunk| {
300                let mut event = String::new();
301                let mut data = String::new();
302                for line in chunk.lines() {
303                    if let Some(rest) = line.strip_prefix("event:") {
304                        event = rest.trim().to_string();
305                    } else if let Some(rest) = line.strip_prefix("data:") {
306                        data = rest.trim().to_string();
307                    }
308                }
309                (event, data)
310            })
311            .collect()
312    }
313
314    #[test]
315    fn sse_parser_ignores_unrecognized_lines_between_event_and_data() {
316        assert_eq!(
317            parse_sse("event: ping\nid: 42\ndata: payload\n\n"),
318            vec![("ping".to_string(), "payload".to_string())]
319        );
320    }
321
322    #[tokio::test]
323    async fn streaming_an_unknown_run_returns_404() {
324        let state = crate::test_support::in_memory_state().await;
325        let app = crate::build_router(state);
326        let response = app
327            .oneshot(
328                Request::get("/api/runs/999/stream")
329                    .body(Body::empty())
330                    .unwrap(),
331            )
332            .await
333            .unwrap();
334        assert_eq!(response.status(), StatusCode::NOT_FOUND);
335    }
336
337    #[tokio::test]
338    async fn streaming_after_its_deadline_emits_an_error_event_and_closes() {
339        let state = crate::test_support::in_memory_state().await;
340        let run_id = seed_running_run(&state, "LIC-STREAM-TIMEOUT").await;
341        let response = stream_run_with_permit(state, run_id, None, Some(Duration::ZERO))
342            .await
343            .unwrap()
344            .into_response();
345
346        let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
347        let events = parse_sse(core::str::from_utf8(&bytes).unwrap());
348
349        assert_eq!(
350            events,
351            vec![(
352                "error".to_string(),
353                "SSE demo time limit exceeded".to_string()
354            )]
355        );
356    }
357
358    #[tokio::test]
359    async fn streaming_a_completed_run_replays_every_sample_then_emits_done() {
360        let state = crate::test_support::in_memory_state().await;
361        let run_id = seed_running_run(&state, "LIC-STREAM-1").await;
362        insert_sample(&state, run_id, 0).await;
363        insert_sample(&state, run_id, 1).await;
364        TuneRunRow::complete(&state.pool, run_id, Utc::now())
365            .await
366            .unwrap();
367
368        let app = crate::build_router(state);
369        let response = app
370            .oneshot(
371                Request::get(format!("/api/runs/{run_id}/stream"))
372                    .body(Body::empty())
373                    .unwrap(),
374            )
375            .await
376            .unwrap();
377        assert_eq!(response.status(), StatusCode::OK);
378        assert_eq!(
379            response
380                .headers()
381                .get(axum::http::header::CONTENT_TYPE)
382                .unwrap(),
383            "text/event-stream"
384        );
385
386        let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
387        let body = String::from_utf8(bytes.to_vec()).unwrap();
388        let events = parse_sse(&body);
389
390        let sample_events: Vec<_> = events.iter().filter(|(e, _)| e == "sample").collect();
391        assert_eq!(
392            sample_events.len(),
393            2,
394            "both pre-recorded samples must be replayed"
395        );
396        let first: serde_json::Value = serde_json::from_str(&sample_events[0].1).unwrap();
397        assert_eq!(first["tick_index"], 0);
398        let second: serde_json::Value = serde_json::from_str(&sample_events[1].1).unwrap();
399        assert_eq!(second["tick_index"], 1);
400
401        let (last_event, last_data) = events.last().expect("stream must emit at least `done`");
402        assert_eq!(last_event, "done");
403        let done: serde_json::Value = serde_json::from_str(last_data).unwrap();
404        assert_eq!(done["outcome"], "completed");
405    }
406
407    #[tokio::test]
408    async fn streaming_emits_initial_readings_before_replayed_samples() {
409        let state = crate::test_support::in_memory_state().await;
410        let run_id = seed_running_run(&state, "LIC-STREAM-INITIAL").await;
411        record_initial_readings(&state, run_id).await;
412        insert_sample(&state, run_id, 0).await;
413        TuneRunRow::complete(&state.pool, run_id, Utc::now())
414            .await
415            .unwrap();
416
417        let app = crate::build_router(state);
418        let response = app
419            .oneshot(
420                Request::get(format!("/api/runs/{run_id}/stream"))
421                    .body(Body::empty())
422                    .unwrap(),
423            )
424            .await
425            .unwrap();
426        let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
427        let events = parse_sse(core::str::from_utf8(&bytes).unwrap());
428
429        assert_eq!(events[0].0, "initial");
430        let initial: serde_json::Value = serde_json::from_str(&events[0].1).unwrap();
431        assert_eq!(initial["pv_ini"], 48.0);
432        assert_eq!(initial["mv_ini"], 42.0);
433        assert_eq!(events[1].0, "sample");
434        assert_eq!(events[2].0, "done");
435    }
436
437    #[tokio::test]
438    async fn streaming_a_run_with_no_samples_still_terminates_with_done() {
439        let state = crate::test_support::in_memory_state().await;
440        let run_id = seed_running_run(&state, "LIC-STREAM-2").await;
441        TuneRunRow::abort(&state.pool, run_id, Utc::now())
442            .await
443            .unwrap();
444
445        let app = crate::build_router(state);
446        let response = app
447            .oneshot(
448                Request::get(format!("/api/runs/{run_id}/stream"))
449                    .body(Body::empty())
450                    .unwrap(),
451            )
452            .await
453            .unwrap();
454        let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
455        let events = parse_sse(core::str::from_utf8(&bytes).unwrap());
456        assert_eq!(
457            events.len(),
458            1,
459            "no samples were recorded, so only `done` is emitted"
460        );
461        assert_eq!(events[0].0, "done");
462        let done: serde_json::Value = serde_json::from_str(&events[0].1).unwrap();
463        assert_eq!(done["outcome"], "aborted");
464    }
465
466    #[tokio::test]
467    async fn streaming_a_still_running_run_waits_for_it_to_finish_before_closing() {
468        let state = crate::test_support::in_memory_state().await;
469        let run_id = seed_running_run(&state, "LIC-STREAM-3").await;
470        insert_sample(&state, run_id, 0).await;
471        // Left `running` -- the stream must poll past at least one `POLL_INTERVAL` sleep
472        // before observing a terminal outcome. A background task flips it to `completed`
473        // shortly after the stream has almost certainly already seen it as `running` at
474        // least once (the handler's own pre-flight `TuneRunRow::get` existence check runs
475        // synchronously before this task is even spawned, so there is no race on "does the
476        // run exist").
477        let pool = state.pool.clone();
478        tokio::spawn(async move {
479            tokio::time::sleep(Duration::from_millis(400)).await;
480            TuneRunRow::complete(&pool, run_id, Utc::now())
481                .await
482                .unwrap();
483        });
484
485        let app = crate::build_router(state);
486        let response = tokio::time::timeout(
487            Duration::from_secs(5),
488            app.oneshot(
489                Request::get(format!("/api/runs/{run_id}/stream"))
490                    .body(Body::empty())
491                    .unwrap(),
492            ),
493        )
494        .await
495        .expect("the stream must eventually close on its own")
496        .unwrap();
497
498        let bytes = tokio::time::timeout(
499            Duration::from_secs(5),
500            to_bytes(response.into_body(), usize::MAX),
501        )
502        .await
503        .expect("reading the full (finite) SSE body must not hang")
504        .unwrap();
505        let events = parse_sse(core::str::from_utf8(&bytes).unwrap());
506
507        assert!(
508            events
509                .iter()
510                .any(|(e, d)| e == "sample" && d.contains("\"tick_index\":0")),
511            "the pre-recorded sample must have been replayed: {events:?}"
512        );
513        let (last_event, last_data) = events.last().unwrap();
514        assert_eq!(last_event, "done");
515        let done: serde_json::Value = serde_json::from_str(last_data).unwrap();
516        assert_eq!(done["outcome"], "completed");
517    }
518}