1use std::convert::Infallible;
4use std::str::FromStr;
5use std::sync::Arc;
6use std::time::Duration;
7
8use axum::extract::{Query, State};
9use axum::response::sse::{Event as SseEvent, KeepAlive, Sse};
10use chrono::Utc;
11use futures_util::stream::{Stream, StreamExt, unfold};
12use serde::Deserialize;
13use serde::de::{self, Deserializer};
14use serde_json::{json, to_string};
15use tokio::select;
16use tokio::sync::broadcast;
17use tokio::time::{Instant, Interval, MissedTickBehavior, interval_at};
18use tokio_stream::wrappers::BroadcastStream;
19use tokio_stream::wrappers::errors::BroadcastStreamRecvError;
20use tracing::warn;
21use uuid::Uuid;
22
23use crate::authorization::AuthorizeExt;
24use crate::error::ApiError;
25use crate::state::AppState;
26use ironflow_auth::extractor::{AuthMethod, Authenticated};
27use ironflow_engine::notify::Event;
28use ironflow_store::entities::ApiKeyScope;
29use ironflow_store::error::StoreError;
30use ironflow_store::store::Store;
31
32pub use ironflow_store::entities::EventKind;
33
34fn deserialize_comma_event_kinds<'de, D>(
36 deserializer: D,
37) -> Result<Option<Vec<EventKind>>, D::Error>
38where
39 D: Deserializer<'de>,
40{
41 let opt: Option<String> = Option::deserialize(deserializer)?;
42 match opt {
43 None => Ok(None),
44 Some(raw) => {
45 let kinds: Result<Vec<EventKind>, _> = raw
46 .split(',')
47 .map(|s| s.trim())
48 .filter(|s| !s.is_empty())
49 .map(EventKind::from_str)
50 .collect();
51 kinds.map(Some).map_err(de::Error::custom)
52 }
53 }
54}
55
56#[derive(Debug, Deserialize)]
71pub struct EventsQuery {
72 pub run_id: Option<Uuid>,
74 #[serde(default, deserialize_with = "deserialize_comma_event_kinds")]
76 pub types: Option<Vec<EventKind>>,
77}
78
79fn visible_to_non_admin(event: &Event) -> bool {
84 match event {
85 Event::RunCreated(_)
86 | Event::RunStatusChanged(_)
87 | Event::RunFailed(_)
88 | Event::RunBudgetExceeded(_)
89 | Event::RetryForced(_)
90 | Event::StepCompleted(_)
91 | Event::StepFailed(_)
92 | Event::ApprovalRequested(_)
93 | Event::ApprovalGranted(_)
94 | Event::ApprovalRejected(_)
95 | Event::ApprovalEscalated(_)
96 | Event::LogLine(_)
97 | Event::SignalAwaited(_)
98 | Event::SignalReceived(_) => true,
99 Event::UserSignedIn(_)
100 | Event::UserSignedUp(_)
101 | Event::UserSignedOut(_)
102 | Event::ProviderAccountUpdated(_)
103 | Event::ProviderAccountUsageUpdated(_)
104 | Event::ScheduleOccurrencesMissed(_)
105 | Event::WorkflowPaused(_)
106 | Event::WorkflowResumed(_) => false,
107 }
108}
109
110pub async fn events(
150 auth: Authenticated,
151 State(state): State<AppState>,
152 Query(query): Query<EventsQuery>,
153) -> Result<Sse<impl Stream<Item = Result<SseEvent, Infallible>>>, ApiError> {
154 auth.require(ApiKeyScope::RunsRead, false)?;
155 let connection = state.sse_limiter.try_acquire(auth.user_id)?;
156 let receiver = state.event_sender.subscribe();
157 let access = if auth.require(ApiKeyScope::Admin, true).is_ok() {
160 Access::Admin
161 } else {
162 Access::Member
163 };
164 let revalidation = Revalidation {
165 store: state.store.clone(),
166 principal: Principal::of(&auth),
167 every: REVALIDATE_INTERVAL,
168 };
169 let stream = event_stream(receiver, query.run_id, query.types, access, revalidation);
170 let stream = stream.map(move |item| {
173 let _held = &connection;
174 item
175 });
176
177 let sse = Sse::new(stream).keep_alive(KeepAlive::new().interval(REVALIDATE_INTERVAL));
178 Ok(sse)
179}
180
181const REVALIDATE_INTERVAL: Duration = Duration::from_secs(30);
183
184const LAGGED_EVENT: &str = "lagged";
186
187#[derive(Debug, Clone, Copy, PartialEq, Eq)]
189enum Access {
190 Closed,
192 Member,
194 Admin,
196}
197
198#[derive(Debug, Clone)]
200enum Principal {
201 Session { user_id: Uuid },
203 ApiKey { key_id: Uuid },
205}
206
207impl Principal {
208 fn of(auth: &Authenticated) -> Self {
210 match &auth.method {
211 AuthMethod::Jwt { .. } => Principal::Session {
212 user_id: auth.user_id,
213 },
214 AuthMethod::ApiKey { key_id, .. } => Principal::ApiKey { key_id: *key_id },
215 }
216 }
217
218 async fn access(&self, store: &dyn Store) -> Result<Access, StoreError> {
225 match self {
226 Principal::Session { user_id } => {
227 let access = match store.find_user_by_id(*user_id).await? {
228 None => Access::Closed,
229 Some(user) if user.is_admin => Access::Admin,
230 Some(_) => Access::Member,
231 };
232 Ok(access)
233 }
234 Principal::ApiKey { key_id } => {
235 let Some(key) = store.find_api_key_by_id(*key_id).await? else {
236 return Ok(Access::Closed);
237 };
238 let expired = key.expires_at.is_some_and(|at| at < Utc::now());
239 if !key.is_active
240 || expired
241 || !ApiKeyScope::has_permission(&key.scopes, &ApiKeyScope::RunsRead)
242 {
243 return Ok(Access::Closed);
244 }
245 let Some(owner) = store.find_user_by_id(key.user_id).await? else {
246 return Ok(Access::Closed);
247 };
248 let admin_scope = ApiKeyScope::has_permission(&key.scopes, &ApiKeyScope::Admin);
249 if owner.is_admin && admin_scope {
250 Ok(Access::Admin)
251 } else {
252 Ok(Access::Member)
253 }
254 }
255 }
256 }
257}
258
259struct Revalidation {
261 store: Arc<dyn Store>,
263 principal: Principal,
265 every: Duration,
267}
268
269struct StreamState {
271 receiver: BroadcastStream<Event>,
272 run_id_filter: Option<Uuid>,
273 type_filter: Option<Vec<EventKind>>,
274 access: Access,
275 revalidation: Revalidation,
276 ticker: Interval,
277}
278
279impl StreamState {
280 async fn revalidate(&mut self) -> bool {
285 let store = self.revalidation.store.as_ref();
286 match self.revalidation.principal.access(store).await {
287 Ok(Access::Closed) => false,
288 Ok(access) => {
289 self.access = access;
290 true
291 }
292 Err(e) => {
293 warn!(error = %e, "failed to revalidate SSE principal, closing the stream");
294 false
295 }
296 }
297 }
298
299 fn render(&self, result: Result<Event, BroadcastStreamRecvError>) -> Option<SseEvent> {
301 let event = match result {
302 Ok(event) => event,
303 Err(BroadcastStreamRecvError::Lagged(missed)) => {
304 let data = json!({ "missed": missed }).to_string();
305 return Some(SseEvent::default().event(LAGGED_EVENT).data(data));
306 }
307 };
308
309 if self.access != Access::Admin && !visible_to_non_admin(&event) {
310 return None;
311 }
312
313 if let Some(ref rid) = self.run_id_filter
314 && event.run_id() != Some(*rid)
315 {
316 return None;
317 }
318
319 if let Some(ref kinds) = self.type_filter {
320 let event_type = event.event_type();
321 if !kinds.iter().any(|k| k.as_str() == event_type) {
322 return None;
323 }
324 }
325
326 let data = to_string(&event).ok()?;
327 Some(SseEvent::default().event(event.event_type()).data(data))
328 }
329}
330
331fn event_stream(
340 receiver: broadcast::Receiver<Event>,
341 run_id_filter: Option<Uuid>,
342 type_filter: Option<Vec<EventKind>>,
343 access: Access,
344 revalidation: Revalidation,
345) -> impl Stream<Item = Result<SseEvent, Infallible>> {
346 let mut ticker = interval_at(Instant::now() + revalidation.every, revalidation.every);
347 ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
348 let state = StreamState {
349 receiver: BroadcastStream::new(receiver),
350 run_id_filter,
351 type_filter,
352 access,
353 revalidation,
354 ticker,
355 };
356
357 unfold(state, |mut state| async move {
358 loop {
359 select! {
362 biased;
363 _ = state.ticker.tick() => {
364 if !state.revalidate().await {
365 return None;
366 }
367 }
368 result = state.receiver.next() => {
369 if let Some(sse_event) = state.render(result?) {
370 return Some((Ok::<_, Infallible>(sse_event), state));
371 }
372 }
373 }
374 }
375 })
376}
377
378#[cfg(test)]
379mod tests {
380 use std::collections::HashMap;
381 use std::sync::Arc;
382 use std::time::Duration;
383
384 use axum::Router;
385 use axum::body::Body;
386 use axum::response::IntoResponse;
387 use axum::response::sse::Sse;
388 use axum::routing::get;
389 use chrono::Utc;
390 use http_body_util::BodyExt;
391 use ironflow_auth::extractor::{API_KEY_PREFIX, API_KEY_SUFFIX_LEN};
392 use ironflow_auth::password;
393 use ironflow_core::providers::claude::ClaudeCodeProvider;
394 use ironflow_engine::engine::Engine;
395 use ironflow_engine::notify::{
396 Event, ProviderAccountChange, ProviderAccountUpdatedEvent, RunStatusChangedEvent,
397 UserSignedInEvent, UserSignedOutEvent, UserSignedUpEvent,
398 };
399 use ironflow_store::entities::{ApiKeyScope, ApiKeyUpdate, NewApiKey, NewUser, User};
400 use ironflow_store::memory::InMemoryStore;
401 use ironflow_store::models::RunStatus;
402 use rust_decimal::Decimal;
403 use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader};
404 use tokio::net::{TcpListener, TcpStream};
405 use tokio::sync::broadcast;
406 use tokio::time::{sleep, timeout};
407 use uuid::Uuid;
408
409 use super::{
410 Access, EventKind, Principal, REVALIDATE_INTERVAL, Revalidation, event_stream, events,
411 visible_to_non_admin,
412 };
413 use crate::routes::test_helpers::create_user_auth_header;
414 use crate::sse_limiter::SseLimits;
415 use crate::state::AppState;
416
417 fn test_state() -> AppState {
418 let store = Arc::new(InMemoryStore::new());
419 let provider = Arc::new(ClaudeCodeProvider::new());
420 let engine = Arc::new(Engine::new(store.clone(), provider));
421 let jwt_config = Arc::new(ironflow_auth::jwt::JwtConfig {
422 secret: "test-secret".to_string(),
423 access_token_ttl_secs: 900,
424 refresh_token_ttl_secs: 604800,
425 cookie_domain: None,
426 cookie_secure: false,
427 });
428 let (event_sender, _) = broadcast::channel::<Event>(16);
429 AppState::new(
430 store,
431 engine,
432 jwt_config,
433 "test-worker-token".to_string(),
434 event_sender,
435 )
436 }
437
438 fn sample_run_event(run_id: Uuid) -> Event {
439 Event::RunStatusChanged(RunStatusChangedEvent {
440 run_id,
441 workflow_name: "deploy".to_string(),
442 from: RunStatus::Running,
443 to: RunStatus::Completed,
444 error: None,
445 cost_usd: Decimal::ZERO,
446 duration_ms: 1000,
447 labels: HashMap::new(),
448 at: Utc::now(),
449 })
450 }
451
452 fn sample_user_event() -> Event {
453 Event::UserSignedIn(UserSignedInEvent {
454 user_id: Uuid::now_v7(),
455 username: "alice".to_string(),
456 at: Utc::now(),
457 })
458 }
459
460 fn sample_user_signed_up_event() -> Event {
461 Event::UserSignedUp(UserSignedUpEvent {
462 user_id: Uuid::now_v7(),
463 username: "bob".to_string(),
464 at: Utc::now(),
465 })
466 }
467
468 fn sample_user_signed_out_event() -> Event {
469 Event::UserSignedOut(UserSignedOutEvent {
470 user_id: Uuid::now_v7(),
471 at: Utc::now(),
472 })
473 }
474
475 fn sample_provider_account_event() -> Event {
476 Event::ProviderAccountUpdated(ProviderAccountUpdatedEvent {
477 account_id: Uuid::now_v7(),
478 name: "acct".into(),
479 change: ProviderAccountChange::Updated,
480 at: Utc::now(),
481 })
482 }
483
484 async fn api_key_header(
490 state: &AppState,
491 owner_is_admin: bool,
492 scopes: Vec<ApiKeyScope>,
493 ) -> String {
494 let owner = stored_user(state, owner_is_admin).await;
495 let (_, header) = stored_key(state, &owner, scopes).await;
496 header
497 }
498
499 async fn start_sse_server(state: AppState) -> (String, broadcast::Sender<Event>, String) {
501 let sender = state.event_sender.clone();
502 let auth = create_user_auth_header(&state, "testuser", false).await;
503 let app = Router::new()
504 .route("/events", get(events))
505 .with_state(state);
506
507 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
508 let addr = listener.local_addr().unwrap().to_string();
509 tokio::spawn(async move {
510 axum::serve(listener, app).await.unwrap();
511 });
512 (addr, sender, auth)
513 }
514
515 async fn connect_sse(addr: &str, query: &str, auth: &str) -> BufReader<TcpStream> {
517 let stream = TcpStream::connect(addr).await.unwrap();
518 let (reader, mut writer) = stream.into_split();
519
520 use tokio::io::AsyncWriteExt;
521 writer
522 .write_all(
523 format!(
524 "GET /events{query} HTTP/1.1\r\nHost: {addr}\r\nAccept: text/event-stream\r\nAuthorization: {auth}\r\n\r\n"
525 )
526 .as_bytes(),
527 )
528 .await
529 .unwrap();
530
531 BufReader::new(reader.reunite(writer).unwrap())
532 }
533
534 async fn read_until_contains(
537 reader: &mut BufReader<TcpStream>,
538 needle: &str,
539 dur: Duration,
540 ) -> String {
541 let mut accumulated = String::new();
542 let result = timeout(dur, async {
543 loop {
544 let mut line = String::new();
545 let n = reader.read_line(&mut line).await.unwrap();
546 if n == 0 {
547 break;
548 }
549 accumulated.push_str(&line);
550 if accumulated.contains(needle) {
551 break;
552 }
553 }
554 })
555 .await;
556 if result.is_err() {
557 panic!("timeout waiting for '{needle}' in SSE stream. Data so far:\n{accumulated}");
558 }
559 accumulated
560 }
561
562 async fn wait_for_subscriber(sender: &broadcast::Sender<Event>) {
568 let ready = timeout(Duration::from_secs(5), async {
569 while sender.receiver_count() == 0 {
570 sleep(Duration::from_millis(5)).await;
571 }
572 })
573 .await;
574 assert!(ready.is_ok(), "SSE handler did not subscribe in time");
575 }
576
577 #[tokio::test]
578 async fn sse_stream_receives_events() {
579 let state = test_state();
580 let (addr, sender, auth) = start_sse_server(state).await;
581 let mut reader = connect_sse(&addr, "", &auth).await;
582
583 wait_for_subscriber(&sender).await;
584
585 let run_id = Uuid::now_v7();
586 sender.send(sample_run_event(run_id)).unwrap();
587
588 let text =
589 read_until_contains(&mut reader, &run_id.to_string(), Duration::from_secs(5)).await;
590
591 assert!(text.contains("run_status_changed"));
592 assert!(text.contains(&run_id.to_string()));
593 }
594
595 #[tokio::test]
596 async fn sse_filters_by_run_id() {
597 let state = test_state();
598 let (addr, sender, auth) = start_sse_server(state).await;
599
600 let target_run = Uuid::now_v7();
601 let other_run = Uuid::now_v7();
602
603 let mut reader = connect_sse(&addr, &format!("?run_id={target_run}"), &auth).await;
604 wait_for_subscriber(&sender).await;
605
606 sender.send(sample_run_event(other_run)).unwrap();
607 sender.send(sample_run_event(target_run)).unwrap();
608
609 let text =
610 read_until_contains(&mut reader, &target_run.to_string(), Duration::from_secs(5)).await;
611
612 assert!(text.contains(&target_run.to_string()));
613 assert!(!text.contains(&other_run.to_string()));
614 }
615
616 #[tokio::test]
617 async fn sse_filters_by_event_type() {
618 let state = test_state();
619 let admin = create_user_auth_header(&state, "admin", true).await;
621 let (addr, sender, _member) = start_sse_server(state).await;
622
623 let mut reader = connect_sse(&addr, "?types=user_signed_in", &admin).await;
624 wait_for_subscriber(&sender).await;
625
626 let run_id = Uuid::now_v7();
627 sender.send(sample_run_event(run_id)).unwrap();
628 sender.send(sample_user_event()).unwrap();
629
630 let text = read_until_contains(&mut reader, "user_signed_in", Duration::from_secs(5)).await;
631
632 assert!(text.contains("user_signed_in"));
633 assert!(!text.contains("run_status_changed"));
634 }
635
636 #[tokio::test]
637 async fn sse_returns_correct_content_type() {
638 let state = test_state();
639 let (addr, _sender, auth) = start_sse_server(state).await;
640 let mut reader = connect_sse(&addr, "", &auth).await;
641
642 let text =
643 read_until_contains(&mut reader, "text/event-stream", Duration::from_secs(5)).await;
644
645 assert!(text.contains("text/event-stream"));
646 }
647
648 #[tokio::test]
649 async fn sse_rejects_unauthenticated() {
650 let state = test_state();
651 let (addr, _sender, _auth) = start_sse_server(state).await;
652 let stream = TcpStream::connect(&addr).await.unwrap();
654 let (reader, mut writer) = stream.into_split();
655
656 use tokio::io::AsyncWriteExt;
657 writer
658 .write_all(
659 format!(
660 "GET /events HTTP/1.1\r\nHost: {addr}\r\nAccept: text/event-stream\r\n\r\n"
661 )
662 .as_bytes(),
663 )
664 .await
665 .unwrap();
666
667 let mut buf_reader = BufReader::new(reader.reunite(writer).unwrap());
668 let text = read_until_contains(&mut buf_reader, "401", Duration::from_secs(5)).await;
669
670 assert!(text.contains("401"));
671 }
672
673 async fn assert_provider_account_events_hidden(state: AppState, auth: &str) {
679 let (addr, sender, _jwt) = start_sse_server(state).await;
680 let mut reader = connect_sse(&addr, "", auth).await;
681 wait_for_subscriber(&sender).await;
682
683 let run_id = Uuid::now_v7();
684 sender.send(sample_provider_account_event()).unwrap();
685 sender.send(sample_run_event(run_id)).unwrap();
686
687 let text =
688 read_until_contains(&mut reader, &run_id.to_string(), Duration::from_secs(5)).await;
689
690 assert!(!text.contains("provider_account"));
691 }
692
693 #[tokio::test]
694 async fn sse_scope_hides_provider_account_events_from_a_runs_read_key() {
695 let state = test_state();
696 let auth = api_key_header(&state, true, vec![ApiKeyScope::RunsRead]).await;
697 assert_provider_account_events_hidden(state, &auth).await;
698 }
699
700 #[tokio::test]
701 async fn sse_scope_hides_provider_account_events_from_an_admin_scope_key_of_a_member() {
702 let state = test_state();
703 let auth = api_key_header(&state, false, vec![ApiKeyScope::Admin]).await;
704 assert_provider_account_events_hidden(state, &auth).await;
705 }
706
707 #[tokio::test]
708 async fn sse_scope_streams_provider_account_events_to_an_admin_scope_key() {
709 let state = test_state();
710 let auth = api_key_header(&state, true, vec![ApiKeyScope::Admin]).await;
711 let (addr, sender, _jwt) = start_sse_server(state).await;
712 let mut reader = connect_sse(&addr, "", &auth).await;
713 wait_for_subscriber(&sender).await;
714
715 sender.send(sample_provider_account_event()).unwrap();
716
717 let text =
718 read_until_contains(&mut reader, "provider_account", Duration::from_secs(5)).await;
719
720 assert!(text.contains("provider_account"));
721 }
722
723 #[tokio::test]
724 async fn sse_scope_rejects_a_key_without_runs_read() {
725 let state = test_state();
726 let auth = api_key_header(&state, true, vec![ApiKeyScope::StatsRead]).await;
727 let (addr, _sender, _jwt) = start_sse_server(state).await;
728 let mut reader = connect_sse(&addr, "", &auth).await;
729
730 let text = read_until_contains(&mut reader, "\r\n", Duration::from_secs(5)).await;
731
732 assert!(text.starts_with("HTTP/1.1 403"), "status line: {text}");
733 }
734
735 async fn assert_user_events_hidden(state: AppState, auth: Option<&str>) {
738 let (addr, sender, member) = start_sse_server(state).await;
739 let mut reader = connect_sse(&addr, "", auth.unwrap_or(&member)).await;
740 wait_for_subscriber(&sender).await;
741
742 let run_id = Uuid::now_v7();
743 sender.send(sample_user_event()).unwrap();
744 sender.send(sample_user_signed_up_event()).unwrap();
745 sender.send(sample_user_signed_out_event()).unwrap();
746 sender.send(sample_run_event(run_id)).unwrap();
747
748 let text =
749 read_until_contains(&mut reader, &run_id.to_string(), Duration::from_secs(5)).await;
750
751 assert!(!text.contains("user_signed_in"), "stream: {text}");
752 assert!(!text.contains("user_signed_up"), "stream: {text}");
753 assert!(!text.contains("user_signed_out"), "stream: {text}");
754 assert!(text.contains("run_status_changed"));
755 }
756
757 #[tokio::test]
758 async fn sse_hides_user_sign_in_sign_up_and_sign_out_events_from_a_member() {
759 assert_user_events_hidden(test_state(), None).await;
760 }
761
762 #[tokio::test]
763 async fn sse_hides_user_events_from_a_runs_read_key_of_an_admin() {
764 let state = test_state();
765 let auth = api_key_header(&state, true, vec![ApiKeyScope::RunsRead]).await;
766 assert_user_events_hidden(state, Some(&auth)).await;
767 }
768
769 #[tokio::test]
770 async fn sse_streams_user_events_to_an_admin() {
771 let state = test_state();
772 let admin = create_user_auth_header(&state, "admin", true).await;
773 let (addr, sender, _member) = start_sse_server(state).await;
774 let mut reader = connect_sse(&addr, "", &admin).await;
775 wait_for_subscriber(&sender).await;
776
777 sender.send(sample_user_event()).unwrap();
778
779 let text = read_until_contains(&mut reader, "user_signed_in", Duration::from_secs(5)).await;
780
781 assert!(text.contains("user_signed_in"));
782 }
783
784 #[test]
785 fn visible_to_non_admin_allows_only_run_events() {
786 assert!(visible_to_non_admin(&sample_run_event(Uuid::now_v7())));
787 assert!(!visible_to_non_admin(&sample_user_event()));
788 assert!(!visible_to_non_admin(&sample_user_signed_up_event()));
789 assert!(!visible_to_non_admin(&sample_user_signed_out_event()));
790 assert!(!visible_to_non_admin(&sample_provider_account_event()));
791 }
792
793 fn lagging_receiver() -> (broadcast::Sender<Event>, broadcast::Receiver<Event>) {
796 let (sender, receiver) = broadcast::channel::<Event>(2);
797 for _ in 0..5 {
798 sender.send(sample_run_event(Uuid::now_v7())).unwrap();
799 }
800 (sender, receiver)
801 }
802
803 async fn read_body_until(body: &mut Body, needle: &str) -> String {
805 let mut accumulated = String::new();
806 let result = timeout(Duration::from_secs(5), async {
807 while !accumulated.contains(needle) {
808 let Some(frame) = body.frame().await else {
809 break;
810 };
811 if let Ok(data) = frame.unwrap().into_data() {
812 accumulated.push_str(&String::from_utf8_lossy(&data));
813 }
814 }
815 })
816 .await;
817 assert!(
818 result.is_ok(),
819 "timeout waiting for '{needle}'. Data so far:\n{accumulated}"
820 );
821 accumulated
822 }
823
824 fn idle_revalidation() -> Revalidation {
827 Revalidation {
828 store: test_state().store,
829 principal: Principal::Session {
830 user_id: Uuid::now_v7(),
831 },
832 every: REVALIDATE_INTERVAL,
833 }
834 }
835
836 #[tokio::test]
837 async fn lagged_event_is_sent_when_the_subscriber_falls_behind() {
838 let (sender, receiver) = lagging_receiver();
839 let stream = event_stream(receiver, None, None, Access::Member, idle_revalidation());
840 let response = Sse::new(stream).into_response();
841 let mut body = response.into_body();
842
843 let text = read_body_until(&mut body, "\"missed\"").await;
844 assert!(text.contains("lagged"), "got {text}");
845 assert!(text.contains("\"missed\":3"), "got {text}");
846
847 let run_id = Uuid::now_v7();
849 sender.send(sample_run_event(run_id)).unwrap();
850 let text = read_body_until(&mut body, &run_id.to_string()).await;
851 assert!(text.contains("run_status_changed"), "got {text}");
852 }
853
854 #[tokio::test]
855 async fn lagged_event_bypasses_the_run_and_type_filters() {
856 let (_sender, receiver) = lagging_receiver();
857 let stream = event_stream(
858 receiver,
859 Some(Uuid::now_v7()),
860 Some(vec![EventKind::UserSignedIn]),
861 Access::Member,
862 idle_revalidation(),
863 );
864 let mut body = Sse::new(stream).into_response().into_body();
865
866 let text = read_body_until(&mut body, "\"missed\"").await;
867 assert!(text.contains("lagged"), "got {text}");
868 assert!(!text.contains("run_status_changed"), "got {text}");
869 }
870
871 const FAST_REVALIDATION: Duration = Duration::from_millis(50);
873
874 async fn stored_user(state: &AppState, is_admin: bool) -> User {
876 let suffix = Uuid::now_v7().simple().to_string();
877 state
878 .store
879 .create_user(NewUser {
880 email: format!("user-{suffix}@example.com"),
881 username: format!("user-{suffix}"),
882 password_hash: "argon2hash".to_string(),
883 is_admin: Some(is_admin),
884 })
885 .await
886 .unwrap()
887 }
888
889 async fn stored_key(
892 state: &AppState,
893 owner: &User,
894 scopes: Vec<ApiKeyScope>,
895 ) -> (Uuid, String) {
896 let suffix = Uuid::now_v7().simple().to_string();
897 let raw_key = format!("irfl_{}rest-of-secret-key", &suffix[24..]);
898 let prefix = &raw_key[..API_KEY_PREFIX.len() + API_KEY_SUFFIX_LEN];
899 let key = state
900 .store
901 .create_api_key(NewApiKey {
902 user_id: owner.id,
903 name: format!("key-{suffix}"),
904 key_hash: password::hash(&raw_key).unwrap(),
905 key_prefix: prefix.to_string(),
906 scopes,
907 expires_at: None,
908 rate_limit_override: None,
909 })
910 .await
911 .unwrap();
912 (key.id, format!("Bearer {raw_key}"))
913 }
914
915 fn revalidated_body(state: &AppState, access: Access, principal: Principal) -> Body {
918 let revalidation = Revalidation {
919 store: state.store.clone(),
920 principal,
921 every: FAST_REVALIDATION,
922 };
923 let receiver = state.event_sender.subscribe();
924 let stream = event_stream(receiver, None, None, access, revalidation);
925 Sse::new(stream).into_response().into_body()
926 }
927
928 async fn assert_stream_open(state: &AppState, body: &mut Body) {
930 let run_id = Uuid::now_v7();
931 state.event_sender.send(sample_run_event(run_id)).unwrap();
932 read_body_until(body, &run_id.to_string()).await;
933 }
934
935 async fn assert_stream_ends(body: &mut Body) {
937 let ended = timeout(Duration::from_secs(5), async {
938 while let Some(frame) = body.frame().await {
939 frame.unwrap();
940 }
941 })
942 .await;
943 assert!(ended.is_ok(), "the stream is still open");
944 }
945
946 async fn user_then_run_event(state: &AppState, body: &mut Body) -> String {
949 let run_id = Uuid::now_v7();
950 state.event_sender.send(sample_user_event()).unwrap();
951 state.event_sender.send(sample_run_event(run_id)).unwrap();
952 read_body_until(body, &run_id.to_string()).await
953 }
954
955 #[tokio::test]
956 async fn sse_stream_closes_when_the_api_key_is_revoked() {
957 let state = test_state();
958 let owner = stored_user(&state, true).await;
959 let (key_id, _) = stored_key(&state, &owner, vec![ApiKeyScope::RunsRead]).await;
960 let mut body = revalidated_body(&state, Access::Member, Principal::ApiKey { key_id });
961 assert_stream_open(&state, &mut body).await;
962
963 state.store.delete_api_key(key_id).await.unwrap();
964
965 assert_stream_ends(&mut body).await;
966 }
967
968 #[tokio::test]
969 async fn sse_stream_closes_when_the_api_key_is_disabled() {
970 let state = test_state();
971 let owner = stored_user(&state, true).await;
972 let (key_id, _) = stored_key(&state, &owner, vec![ApiKeyScope::RunsRead]).await;
973 let mut body = revalidated_body(&state, Access::Member, Principal::ApiKey { key_id });
974 assert_stream_open(&state, &mut body).await;
975
976 state
977 .store
978 .update_api_key(
979 key_id,
980 ApiKeyUpdate {
981 is_active: Some(false),
982 ..ApiKeyUpdate::default()
983 },
984 )
985 .await
986 .unwrap();
987
988 assert_stream_ends(&mut body).await;
989 }
990
991 #[tokio::test]
992 async fn sse_stream_closes_when_the_api_key_expires() {
993 let state = test_state();
994 let owner = stored_user(&state, true).await;
995 let (key_id, _) = stored_key(&state, &owner, vec![ApiKeyScope::RunsRead]).await;
996 let mut body = revalidated_body(&state, Access::Member, Principal::ApiKey { key_id });
997 assert_stream_open(&state, &mut body).await;
998
999 state
1000 .store
1001 .update_api_key(
1002 key_id,
1003 ApiKeyUpdate {
1004 expires_at: Some(Some(Utc::now())),
1005 ..ApiKeyUpdate::default()
1006 },
1007 )
1008 .await
1009 .unwrap();
1010
1011 assert_stream_ends(&mut body).await;
1012 }
1013
1014 #[tokio::test]
1015 async fn sse_stream_closes_when_the_api_key_loses_runs_read() {
1016 let state = test_state();
1017 let owner = stored_user(&state, true).await;
1018 let (key_id, _) = stored_key(&state, &owner, vec![ApiKeyScope::RunsRead]).await;
1019 let mut body = revalidated_body(&state, Access::Member, Principal::ApiKey { key_id });
1020 assert_stream_open(&state, &mut body).await;
1021
1022 state
1023 .store
1024 .update_api_key(
1025 key_id,
1026 ApiKeyUpdate {
1027 scopes: Some(vec![ApiKeyScope::StatsRead]),
1028 ..ApiKeyUpdate::default()
1029 },
1030 )
1031 .await
1032 .unwrap();
1033
1034 assert_stream_ends(&mut body).await;
1035 }
1036
1037 #[tokio::test]
1038 async fn sse_stream_closes_when_the_owner_is_deleted() {
1039 let state = test_state();
1040 let owner = stored_user(&state, true).await;
1041 let (key_id, _) = stored_key(&state, &owner, vec![ApiKeyScope::RunsRead]).await;
1042 let mut body = revalidated_body(&state, Access::Member, Principal::ApiKey { key_id });
1043 assert_stream_open(&state, &mut body).await;
1044
1045 state.store.delete_user(owner.id).await.unwrap();
1046
1047 assert_stream_ends(&mut body).await;
1048 }
1049
1050 #[tokio::test]
1051 async fn sse_stream_closes_when_the_session_user_is_deleted() {
1052 let state = test_state();
1053 let user = stored_user(&state, false).await;
1054 let mut body = revalidated_body(
1055 &state,
1056 Access::Member,
1057 Principal::Session { user_id: user.id },
1058 );
1059 assert_stream_open(&state, &mut body).await;
1060
1061 state.store.delete_user(user.id).await.unwrap();
1062
1063 assert_stream_ends(&mut body).await;
1064 }
1065
1066 #[tokio::test]
1067 async fn sse_stream_of_a_demoted_admin_falls_back_to_member_filtering() {
1068 let state = test_state();
1069 let admin = stored_user(&state, true).await;
1070 let mut body = revalidated_body(
1071 &state,
1072 Access::Admin,
1073 Principal::Session { user_id: admin.id },
1074 );
1075 let text = user_then_run_event(&state, &mut body).await;
1076 assert!(text.contains("user_signed_in"), "stream: {text}");
1077
1078 state.store.update_user_role(admin.id, false).await.unwrap();
1079 sleep(FAST_REVALIDATION * 3).await;
1080
1081 let text = user_then_run_event(&state, &mut body).await;
1082 assert!(!text.contains("user_signed_in"), "stream: {text}");
1083 assert!(text.contains("run_status_changed"), "stream: {text}");
1084 }
1085
1086 #[tokio::test]
1087 async fn sse_stream_of_a_demoted_admin_key_falls_back_to_member_filtering() {
1088 let state = test_state();
1089 let admin = stored_user(&state, true).await;
1090 let (key_id, _) = stored_key(
1091 &state,
1092 &admin,
1093 vec![ApiKeyScope::RunsRead, ApiKeyScope::Admin],
1094 )
1095 .await;
1096 let mut body = revalidated_body(&state, Access::Admin, Principal::ApiKey { key_id });
1097 let text = user_then_run_event(&state, &mut body).await;
1098 assert!(text.contains("user_signed_in"), "stream: {text}");
1099
1100 state.store.update_user_role(admin.id, false).await.unwrap();
1102 sleep(FAST_REVALIDATION * 3).await;
1103
1104 let text = user_then_run_event(&state, &mut body).await;
1105 assert!(!text.contains("user_signed_in"), "stream: {text}");
1106 assert!(text.contains("run_status_changed"), "stream: {text}");
1107 }
1108
1109 #[tokio::test]
1110 async fn sse_stream_of_a_demoted_admin_with_an_admin_only_key_is_closed() {
1111 let state = test_state();
1112 let admin = stored_user(&state, true).await;
1113 let (key_id, _) = stored_key(&state, &admin, vec![ApiKeyScope::Admin]).await;
1114 let mut body = revalidated_body(&state, Access::Admin, Principal::ApiKey { key_id });
1115 assert_stream_open(&state, &mut body).await;
1116
1117 state.store.update_user_role(admin.id, false).await.unwrap();
1119
1120 assert_stream_ends(&mut body).await;
1121 }
1122
1123 #[tokio::test]
1124 async fn sse_stream_keeps_flowing_while_the_principal_is_unchanged() {
1125 let state = test_state();
1126 let admin = stored_user(&state, true).await;
1127 let mut body = revalidated_body(
1128 &state,
1129 Access::Admin,
1130 Principal::Session { user_id: admin.id },
1131 );
1132
1133 for _ in 0..3 {
1134 sleep(FAST_REVALIDATION * 2).await;
1135 let text = user_then_run_event(&state, &mut body).await;
1136 assert!(text.contains("user_signed_in"), "stream: {text}");
1137 }
1138 }
1139
1140 async fn status_line(reader: &mut BufReader<TcpStream>) -> String {
1142 read_until_contains(reader, "\r\n", Duration::from_secs(5)).await
1143 }
1144
1145 async fn read_body(reader: &mut BufReader<TcpStream>) -> String {
1147 timeout(Duration::from_secs(5), async {
1148 let mut content_length = 0;
1149 loop {
1150 let mut line = String::new();
1151 reader.read_line(&mut line).await.unwrap();
1152 if line.is_empty() || line == "\r\n" {
1153 break;
1154 }
1155 if let Some(value) = line.to_ascii_lowercase().strip_prefix("content-length:") {
1156 content_length = value.trim().parse().unwrap();
1157 }
1158 }
1159 let mut body = vec![0; content_length];
1160 reader.read_exact(&mut body).await.unwrap();
1161 String::from_utf8(body).unwrap()
1162 })
1163 .await
1164 .expect("timeout reading the response body")
1165 }
1166
1167 fn limited_state(per_principal: usize) -> AppState {
1168 test_state().with_sse_limits(SseLimits {
1169 per_principal: Some(per_principal),
1170 global: None,
1171 })
1172 }
1173
1174 #[tokio::test]
1175 async fn events_returns_429_past_the_per_principal_limit() {
1176 let state = limited_state(2);
1177 let other_auth = create_user_auth_header(&state, "otheruser", false).await;
1178 let (addr, _sender, auth) = start_sse_server(state).await;
1179
1180 let mut first = connect_sse(&addr, "", &auth).await;
1181 assert!(status_line(&mut first).await.starts_with("HTTP/1.1 200"));
1182 let mut second = connect_sse(&addr, "", &auth).await;
1183 assert!(status_line(&mut second).await.starts_with("HTTP/1.1 200"));
1184
1185 let mut third = connect_sse(&addr, "", &auth).await;
1186 let status = status_line(&mut third).await;
1187 assert!(status.starts_with("HTTP/1.1 429"), "status line: {status}");
1188 let body = read_body(&mut third).await;
1189 assert!(body.contains("SSE_CONNECTION_LIMIT"), "body: {body}");
1190
1191 let mut other = connect_sse(&addr, "", &other_auth).await;
1192 let status = status_line(&mut other).await;
1193 assert!(
1194 status.starts_with("HTTP/1.1 200"),
1195 "another user keeps its own quota: {status}"
1196 );
1197 }
1198
1199 #[tokio::test]
1200 async fn events_frees_the_slot_when_the_client_disconnects() {
1201 let state = limited_state(1);
1202 let (addr, _sender, auth) = start_sse_server(state).await;
1203
1204 let mut first = connect_sse(&addr, "", &auth).await;
1205 assert!(status_line(&mut first).await.starts_with("HTTP/1.1 200"));
1206 let mut refused = connect_sse(&addr, "", &auth).await;
1207 assert!(status_line(&mut refused).await.starts_with("HTTP/1.1 429"));
1208
1209 drop(first);
1210
1211 let reopened = timeout(Duration::from_secs(5), async {
1212 loop {
1213 let mut retry = connect_sse(&addr, "", &auth).await;
1214 if status_line(&mut retry).await.starts_with("HTTP/1.1 200") {
1215 return retry;
1216 }
1217 sleep(Duration::from_millis(20)).await;
1218 }
1219 })
1220 .await;
1221 assert!(reopened.is_ok(), "the slot was not freed after disconnect");
1222 }
1223}