feat: pentest onboarding — streaming, browser automation, reports, user cleanup (#16)

Complete pentest feature overhaul: SSE streaming, session-persistent browser tool (CDP), AES-256 credential encryption, auto-screenshots in reports, code-level remediation correlation, SAST triage chunking, context window optimization, test user cleanup (Keycloak/Auth0/Okta), wizard dropdowns, attack chain improvements, architecture docs with Mermaid diagrams.

Co-authored-by: Sharang Parnerkar <parnerkarsharang@gmail.com>
Reviewed-on: #16

compliance-agent/src/api/handlers/pentest_handlers/stream.rs

@@ -1,10 +1,14 @@
+use std::convert::Infallible;
 use std::sync::Arc;
+use std::time::Duration;
 
 use axum::extract::{Extension, Path};
 use axum::http::StatusCode;
-use axum::response::sse::{Event, Sse};
+use axum::response::sse::{Event, KeepAlive, Sse};
 use futures_util::stream;
 use mongodb::bson::doc;
+use tokio_stream::wrappers::BroadcastStream;
+use tokio_stream::StreamExt;
 
 use compliance_core::models::pentest::*;
 
@@ -16,16 +20,13 @@ type AgentExt = Extension<Arc<ComplianceAgent>>;
 
 /// GET /api/v1/pentest/sessions/:id/stream — SSE endpoint for real-time events
 ///
-/// Returns recent messages as SSE events (polling approach).
-/// True real-time streaming with broadcast channels will be added in a future iteration.
+/// Replays stored messages/nodes as initial burst, then subscribes to the
+/// broadcast channel for live updates. Sends keepalive comments every 15s.
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn session_stream(
     Extension(agent): AgentExt,
     Path(id): Path<String>,
-) -> Result<
-    Sse<impl futures_util::Stream<Item = Result<Event, std::convert::Infallible>>>,
-    StatusCode,
-> {
+) -> Result<Sse<impl futures_util::Stream<Item = Result<Event, Infallible>>>, StatusCode> {
     let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
 
     // Verify session exists
@@ -37,6 +38,10 @@ pub async fn session_stream(
         .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
         .ok_or(StatusCode::NOT_FOUND)?;
 
+    // ── Initial burst: replay stored data ──────────────────────────
+
+    let mut initial_events: Vec<Result<Event, Infallible>> = Vec::new();
+
     // Fetch recent messages for this session
     let messages: Vec<PentestMessage> = match agent
         .db
@@ -63,9 +68,6 @@ pub async fn session_stream(
         Err(_) => Vec::new(),
     };
 
-    // Build SSE events from stored data
-    let mut events: Vec<Result<Event, std::convert::Infallible>> = Vec::new();
-
     for msg in &messages {
         let event_data = serde_json::json!({
             "type": "message",
@@ -74,7 +76,7 @@ pub async fn session_stream(
             "created_at": msg.created_at.to_rfc3339(),
         });
         if let Ok(data) = serde_json::to_string(&event_data) {
-            events.push(Ok(Event::default().event("message").data(data)));
+            initial_events.push(Ok(Event::default().event("message").data(data)));
         }
     }
 
@@ -87,11 +89,11 @@ pub async fn session_stream(
             "findings_produced": node.findings_produced,
         });
         if let Ok(data) = serde_json::to_string(&event_data) {
-            events.push(Ok(Event::default().event("tool").data(data)));
+            initial_events.push(Ok(Event::default().event("tool").data(data)));
         }
     }
 
-    // Add session status event
+    // Add current session status event
     let session = agent
         .db
         .pentest_sessions()
@@ -108,9 +110,49 @@ pub async fn session_stream(
             "tool_invocations": s.tool_invocations,
         });
         if let Ok(data) = serde_json::to_string(&status_data) {
-            events.push(Ok(Event::default().event("status").data(data)));
+            initial_events.push(Ok(Event::default().event("status").data(data)));
         }
     }
 
-    Ok(Sse::new(stream::iter(events)))
+    // ── Live stream: subscribe to broadcast ────────────────────────
+
+    let live_stream = if let Some(rx) = agent.subscribe_session(&id) {
+        let broadcast = BroadcastStream::new(rx).filter_map(|result| match result {
+            Ok(event) => {
+                if let Ok(data) = serde_json::to_string(&event) {
+                    let event_type = match &event {
+                        PentestEvent::ToolStart { .. } => "tool_start",
+                        PentestEvent::ToolComplete { .. } => "tool_complete",
+                        PentestEvent::Finding { .. } => "finding",
+                        PentestEvent::Message { .. } => "message",
+                        PentestEvent::Complete { .. } => "complete",
+                        PentestEvent::Error { .. } => "error",
+                        PentestEvent::Thinking { .. } => "thinking",
+                        PentestEvent::Paused => "paused",
+                        PentestEvent::Resumed => "resumed",
+                    };
+                    Some(Ok(Event::default().event(event_type).data(data)))
+                } else {
+                    None
+                }
+            }
+            Err(_) => None,
+        });
+        // Box to unify types
+        Box::pin(broadcast)
+            as std::pin::Pin<Box<dyn futures_util::Stream<Item = Result<Event, Infallible>> + Send>>
+    } else {
+        // No active broadcast — return empty stream
+        Box::pin(stream::empty())
+            as std::pin::Pin<Box<dyn futures_util::Stream<Item = Result<Event, Infallible>> + Send>>
+    };
+
+    // Chain initial burst + live stream
+    let combined = stream::iter(initial_events).chain(live_stream);
+
+    Ok(Sse::new(combined).keep_alive(
+        KeepAlive::new()
+            .interval(Duration::from_secs(15))
+            .text("keepalive"),
+    ))
 }