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/export.rs

@@ -7,8 +7,12 @@ use axum::Json;
 use mongodb::bson::doc;
 use serde::Deserialize;
 
+use futures_util::StreamExt;
+
 use compliance_core::models::dast::DastFinding;
+use compliance_core::models::finding::Finding;
 use compliance_core::models::pentest::*;
+use compliance_core::models::sbom::SbomEntry;
 
 use crate::agent::ComplianceAgent;
 
@@ -103,6 +107,97 @@ pub async fn export_session_report(
         Err(_) => Vec::new(),
     };
 
+    // Fetch SAST findings, SBOM, and code context for the linked repository
+    let repo_id = session
+        .repo_id
+        .clone()
+        .or_else(|| target.as_ref().and_then(|t| t.repo_id.clone()));
+
+    let (sast_findings, sbom_entries, code_context) = if let Some(ref rid) = repo_id {
+        let sast: Vec<Finding> = match agent
+            .db
+            .findings()
+            .find(doc! {
+                "repo_id": rid,
+                "status": { "$in": ["open", "triaged"] },
+            })
+            .sort(doc! { "severity": -1 })
+            .limit(100)
+            .await
+        {
+            Ok(mut cursor) => {
+                let mut results = Vec::new();
+                while let Some(Ok(f)) = cursor.next().await {
+                    results.push(f);
+                }
+                results
+            }
+            Err(_) => Vec::new(),
+        };
+
+        let sbom: Vec<SbomEntry> = match agent
+            .db
+            .sbom_entries()
+            .find(doc! {
+                "repo_id": rid,
+                "known_vulnerabilities": { "$exists": true, "$ne": [] },
+            })
+            .limit(50)
+            .await
+        {
+            Ok(mut cursor) => {
+                let mut results = Vec::new();
+                while let Some(Ok(e)) = cursor.next().await {
+                    results.push(e);
+                }
+                results
+            }
+            Err(_) => Vec::new(),
+        };
+
+        // Build code context from graph nodes
+        let code_ctx: Vec<CodeContextHint> = match agent
+            .db
+            .graph_nodes()
+            .find(doc! { "repo_id": rid, "is_entry_point": true })
+            .limit(50)
+            .await
+        {
+            Ok(mut cursor) => {
+                let mut nodes_vec = Vec::new();
+                while let Some(Ok(n)) = cursor.next().await {
+                    let linked_vulns: Vec<String> = sast
+                        .iter()
+                        .filter(|f| f.file_path.as_deref() == Some(&n.file_path))
+                        .map(|f| {
+                            format!(
+                                "[{}] {}: {} (line {})",
+                                f.severity,
+                                f.scanner,
+                                f.title,
+                                f.line_number.unwrap_or(0)
+                            )
+                        })
+                        .collect();
+                    nodes_vec.push(CodeContextHint {
+                        endpoint_pattern: n.qualified_name.clone(),
+                        handler_function: n.name.clone(),
+                        file_path: n.file_path.clone(),
+                        code_snippet: String::new(),
+                        known_vulnerabilities: linked_vulns,
+                    });
+                }
+                nodes_vec
+            }
+            Err(_) => Vec::new(),
+        };
+
+        (sast, sbom, code_ctx)
+    } else {
+        (Vec::new(), Vec::new(), Vec::new())
+    };
+
+    let config = session.config.clone();
     let ctx = crate::pentest::report::ReportContext {
         session,
         target_name,
@@ -115,6 +210,10 @@ pub async fn export_session_report(
             body.requester_name
         },
         requester_email: body.requester_email,
+        config,
+        sast_findings,
+        sbom_entries,
+        code_context,
     };
 
     let report = crate::pentest::generate_encrypted_report(&ctx, &body.password)

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

@@ -17,10 +17,12 @@ type AgentExt = Extension<Arc<ComplianceAgent>>;
 
 #[derive(Deserialize)]
 pub struct CreateSessionRequest {
-    pub target_id: String,
+    pub target_id: Option<String>,
     #[serde(default = "default_strategy")]
     pub strategy: String,
     pub message: Option<String>,
+    /// Wizard configuration — if present, takes precedence over legacy fields
+    pub config: Option<PentestConfig>,
 }
 
 fn default_strategy() -> String {
@@ -32,83 +34,365 @@ pub struct SendMessageRequest {
     pub message: String,
 }
 
+#[derive(Deserialize)]
+pub struct LookupRepoQuery {
+    pub url: String,
+}
+
 /// POST /api/v1/pentest/sessions — Create a new pentest session and start the orchestrator
 #[tracing::instrument(skip_all)]
 pub async fn create_session(
     Extension(agent): AgentExt,
     Json(req): Json<CreateSessionRequest>,
 ) -> Result<Json<ApiResponse<PentestSession>>, (StatusCode, String)> {
-    let oid = mongodb::bson::oid::ObjectId::parse_str(&req.target_id).map_err(|_| {
-        (
-            StatusCode::BAD_REQUEST,
-            "Invalid target_id format".to_string(),
-        )
-    })?;
-
-    // Look up the target
-    let target = agent
-        .db
-        .dast_targets()
-        .find_one(doc! { "_id": oid })
-        .await
-        .map_err(|e| {
+    // Try to acquire a concurrency permit
+    let permit = agent
+        .session_semaphore
+        .clone()
+        .try_acquire_owned()
+        .map_err(|_| {
             (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                format!("Database error: {e}"),
+                StatusCode::TOO_MANY_REQUESTS,
+                "Maximum concurrent pentest sessions reached. Try again later.".to_string(),
             )
-        })?
-        .ok_or_else(|| (StatusCode::NOT_FOUND, "Target not found".to_string()))?;
+        })?;
 
-    // Parse strategy
-    let strategy = match req.strategy.as_str() {
+    if let Some(ref config) = req.config {
+        // ── Wizard path ──────────────────────────────────────────────
+        if !config.disclaimer_accepted {
+            return Err((
+                StatusCode::BAD_REQUEST,
+                "Disclaimer must be accepted".to_string(),
+            ));
+        }
+
+        // Look up or auto-create DastTarget by app_url
+        let target = match agent
+            .db
+            .dast_targets()
+            .find_one(doc! { "base_url": &config.app_url })
+            .await
+            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?
+        {
+            Some(t) => t,
+            None => {
+                use compliance_core::models::dast::{DastTarget, DastTargetType};
+                let mut t = DastTarget::new(
+                    config.app_url.clone(),
+                    config.app_url.clone(),
+                    DastTargetType::WebApp,
+                );
+                if let Some(rl) = config.rate_limit {
+                    t.rate_limit = rl;
+                }
+                t.allow_destructive = config.allow_destructive;
+                t.excluded_paths = config.scope_exclusions.clone();
+                let res = agent.db.dast_targets().insert_one(&t).await.map_err(|e| {
+                    (
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                        format!("Failed to create target: {e}"),
+                    )
+                })?;
+                t.id = res.inserted_id.as_object_id();
+                t
+            }
+        };
+
+        let target_id = target.id.map(|oid| oid.to_hex()).unwrap_or_default();
+
+        // Parse strategy from config or request
+        let strat_str = config.strategy.as_deref().unwrap_or(req.strategy.as_str());
+        let strategy = parse_strategy(strat_str);
+
+        let mut session = PentestSession::new(target_id, strategy);
+        session.config = Some(config.clone());
+        session.repo_id = target.repo_id.clone();
+
+        // Resolve repo_id from git_repo_url if provided
+        if let Some(ref git_url) = config.git_repo_url {
+            if let Ok(Some(repo)) = agent
+                .db
+                .repositories()
+                .find_one(doc! { "git_url": git_url })
+                .await
+            {
+                session.repo_id = repo.id.map(|oid| oid.to_hex());
+            }
+        }
+
+        let insert_result = agent
+            .db
+            .pentest_sessions()
+            .insert_one(&session)
+            .await
+            .map_err(|e| {
+                (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    format!("Failed to create session: {e}"),
+                )
+            })?;
+        session.id = insert_result.inserted_id.as_object_id();
+
+        let session_id_str = session.id.map(|oid| oid.to_hex()).unwrap_or_default();
+
+        // Register broadcast stream and pause control
+        let event_tx = agent.register_session_stream(&session_id_str);
+        let pause_rx = agent.register_pause_control(&session_id_str);
+
+        // Merge server-default IMAP/email settings where wizard left blanks
+        if let Some(ref mut cfg) = session.config {
+            if cfg.auth.mode == AuthMode::AutoRegister {
+                if cfg.auth.verification_email.is_none() {
+                    cfg.auth.verification_email = agent.config.pentest_verification_email.clone();
+                }
+                if cfg.auth.imap_host.is_none() {
+                    cfg.auth.imap_host = agent.config.pentest_imap_host.clone();
+                }
+                if cfg.auth.imap_port.is_none() {
+                    cfg.auth.imap_port = agent.config.pentest_imap_port;
+                }
+                if cfg.auth.imap_username.is_none() {
+                    cfg.auth.imap_username = agent.config.pentest_imap_username.clone();
+                }
+                if cfg.auth.imap_password.is_none() {
+                    cfg.auth.imap_password = agent.config.pentest_imap_password.as_ref().map(|s| {
+                        use secrecy::ExposeSecret;
+                        s.expose_secret().to_string()
+                    });
+                }
+            }
+        }
+
+        // Pre-populate test user record for auto-register sessions
+        if let Some(ref cfg) = session.config {
+            if cfg.auth.mode == AuthMode::AutoRegister {
+                let verification_email = cfg.auth.verification_email.clone();
+                // Build plus-addressed email for this session
+                let test_email = verification_email.as_deref().map(|email| {
+                    let parts: Vec<&str> = email.splitn(2, '@').collect();
+                    if parts.len() == 2 {
+                        format!("{}+{}@{}", parts[0], session_id_str, parts[1])
+                    } else {
+                        email.to_string()
+                    }
+                });
+
+                // Detect identity provider from keycloak config
+                let provider = if agent.config.keycloak_url.is_some() {
+                    Some(compliance_core::models::pentest::IdentityProvider::Keycloak)
+                } else {
+                    None
+                };
+
+                session.test_user = Some(compliance_core::models::pentest::TestUserRecord {
+                    username: None, // LLM will choose; updated after registration
+                    email: test_email,
+                    provider_user_id: None,
+                    provider,
+                    cleaned_up: false,
+                });
+            }
+        }
+
+        // Encrypt credentials before they linger in memory
+        let mut session_for_task = session.clone();
+        if let Some(ref mut cfg) = session_for_task.config {
+            cfg.auth.username = cfg
+                .auth
+                .username
+                .as_ref()
+                .map(|u| crate::pentest::crypto::encrypt(u));
+            cfg.auth.password = cfg
+                .auth
+                .password
+                .as_ref()
+                .map(|p| crate::pentest::crypto::encrypt(p));
+        }
+
+        // Persist encrypted credentials to DB
+        if session_for_task.config.is_some() {
+            if let Some(sid) = session.id {
+                let _ = agent
+                    .db
+                    .pentest_sessions()
+                    .update_one(
+                        doc! { "_id": sid },
+                        doc! { "$set": {
+                            "config.auth.username": session_for_task.config.as_ref()
+                                .and_then(|c| c.auth.username.as_deref())
+                                .map(|s| mongodb::bson::Bson::String(s.to_string()))
+                                .unwrap_or(mongodb::bson::Bson::Null),
+                            "config.auth.password": session_for_task.config.as_ref()
+                                .and_then(|c| c.auth.password.as_deref())
+                                .map(|s| mongodb::bson::Bson::String(s.to_string()))
+                                .unwrap_or(mongodb::bson::Bson::Null),
+                        }},
+                    )
+                    .await;
+            }
+        }
+
+        let initial_message = config
+            .initial_instructions
+            .clone()
+            .or(req.message.clone())
+            .unwrap_or_else(|| {
+                format!(
+                    "Begin a {} penetration test against {} ({}). \
+                     Identify vulnerabilities and provide evidence for each finding.",
+                    session.strategy, target.name, target.base_url,
+                )
+            });
+
+        let llm = agent.llm.clone();
+        let db = agent.db.clone();
+        let session_clone = session.clone();
+        let target_clone = target.clone();
+        let agent_ref = agent.clone();
+        tokio::spawn(async move {
+            let orchestrator = PentestOrchestrator::new(llm, db, event_tx, Some(pause_rx));
+            orchestrator
+                .run_session_guarded(&session_clone, &target_clone, &initial_message)
+                .await;
+            // Clean up session resources
+            agent_ref.cleanup_session(&session_id_str);
+            // Release concurrency permit
+            drop(permit);
+        });
+
+        // Redact credentials in response
+        let mut response_session = session;
+        if let Some(ref mut cfg) = response_session.config {
+            if cfg.auth.username.is_some() {
+                cfg.auth.username = Some("********".to_string());
+            }
+            if cfg.auth.password.is_some() {
+                cfg.auth.password = Some("********".to_string());
+            }
+        }
+
+        Ok(Json(ApiResponse {
+            data: response_session,
+            total: None,
+            page: None,
+        }))
+    } else {
+        // ── Legacy path ──────────────────────────────────────────────
+        let target_id = req.target_id.clone().ok_or_else(|| {
+            (
+                StatusCode::BAD_REQUEST,
+                "target_id is required for legacy creation".to_string(),
+            )
+        })?;
+
+        let oid = mongodb::bson::oid::ObjectId::parse_str(&target_id).map_err(|_| {
+            (
+                StatusCode::BAD_REQUEST,
+                "Invalid target_id format".to_string(),
+            )
+        })?;
+
+        let target = agent
+            .db
+            .dast_targets()
+            .find_one(doc! { "_id": oid })
+            .await
+            .map_err(|e| {
+                (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    format!("Database error: {e}"),
+                )
+            })?
+            .ok_or_else(|| (StatusCode::NOT_FOUND, "Target not found".to_string()))?;
+
+        let strategy = parse_strategy(&req.strategy);
+
+        let mut session = PentestSession::new(target_id, strategy);
+        session.repo_id = target.repo_id.clone();
+
+        let insert_result = agent
+            .db
+            .pentest_sessions()
+            .insert_one(&session)
+            .await
+            .map_err(|e| {
+                (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    format!("Failed to create session: {e}"),
+                )
+            })?;
+        session.id = insert_result.inserted_id.as_object_id();
+
+        let session_id_str = session.id.map(|oid| oid.to_hex()).unwrap_or_default();
+
+        // Register broadcast stream and pause control
+        let event_tx = agent.register_session_stream(&session_id_str);
+        let pause_rx = agent.register_pause_control(&session_id_str);
+
+        let initial_message = req.message.unwrap_or_else(|| {
+            format!(
+                "Begin a {} penetration test against {} ({}). \
+                 Identify vulnerabilities and provide evidence for each finding.",
+                session.strategy, target.name, target.base_url,
+            )
+        });
+
+        let llm = agent.llm.clone();
+        let db = agent.db.clone();
+        let session_clone = session.clone();
+        let target_clone = target.clone();
+        let agent_ref = agent.clone();
+        tokio::spawn(async move {
+            let orchestrator = PentestOrchestrator::new(llm, db, event_tx, Some(pause_rx));
+            orchestrator
+                .run_session_guarded(&session_clone, &target_clone, &initial_message)
+                .await;
+            agent_ref.cleanup_session(&session_id_str);
+            drop(permit);
+        });
+
+        Ok(Json(ApiResponse {
+            data: session,
+            total: None,
+            page: None,
+        }))
+    }
+}
+
+fn parse_strategy(s: &str) -> PentestStrategy {
+    match s {
         "quick" => PentestStrategy::Quick,
         "targeted" => PentestStrategy::Targeted,
         "aggressive" => PentestStrategy::Aggressive,
         "stealth" => PentestStrategy::Stealth,
         _ => PentestStrategy::Comprehensive,
+    }
+}
+
+/// GET /api/v1/pentest/lookup-repo — Look up a tracked repository by git URL
+#[tracing::instrument(skip_all)]
+pub async fn lookup_repo(
+    Extension(agent): AgentExt,
+    Query(params): Query<LookupRepoQuery>,
+) -> Result<Json<ApiResponse<serde_json::Value>>, StatusCode> {
+    let repo = agent
+        .db
+        .repositories()
+        .find_one(doc! { "git_url": &params.url })
+        .await
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+
+    let data = match repo {
+        Some(r) => serde_json::json!({
+            "name": r.name,
+            "default_branch": r.default_branch,
+            "last_scanned_commit": r.last_scanned_commit,
+        }),
+        None => serde_json::Value::Null,
     };
 
-    // Create session
-    let mut session = PentestSession::new(req.target_id.clone(), strategy);
-    session.repo_id = target.repo_id.clone();
-
-    let insert_result = agent
-        .db
-        .pentest_sessions()
-        .insert_one(&session)
-        .await
-        .map_err(|e| {
-            (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                format!("Failed to create session: {e}"),
-            )
-        })?;
-
-    // Set the generated ID back on the session so the orchestrator has it
-    session.id = insert_result.inserted_id.as_object_id();
-
-    let initial_message = req.message.unwrap_or_else(|| {
-        format!(
-            "Begin a {} penetration test against {} ({}). \
-             Identify vulnerabilities and provide evidence for each finding.",
-            session.strategy, target.name, target.base_url,
-        )
-    });
-
-    // Spawn the orchestrator on a background task
-    let llm = agent.llm.clone();
-    let db = agent.db.clone();
-    let session_clone = session.clone();
-    let target_clone = target.clone();
-    tokio::spawn(async move {
-        let orchestrator = PentestOrchestrator::new(llm, db);
-        orchestrator
-            .run_session_guarded(&session_clone, &target_clone, &initial_message)
-            .await;
-    });
-
     Ok(Json(ApiResponse {
-        data: session,
+        data,
         total: None,
         page: None,
     }))
@@ -158,7 +442,7 @@ pub async fn get_session(
 ) -> Result<Json<ApiResponse<PentestSession>>, StatusCode> {
     let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
 
-    let session = agent
+    let mut session = agent
         .db
         .pentest_sessions()
         .find_one(doc! { "_id": oid })
@@ -166,6 +450,16 @@ pub async fn get_session(
         .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
         .ok_or(StatusCode::NOT_FOUND)?;
 
+    // Redact credentials in response
+    if let Some(ref mut cfg) = session.config {
+        if cfg.auth.username.is_some() {
+            cfg.auth.username = Some("********".to_string());
+        }
+        if cfg.auth.password.is_some() {
+            cfg.auth.password = Some("********".to_string());
+        }
+    }
+
     Ok(Json(ApiResponse {
         data: session,
         total: None,
@@ -241,8 +535,20 @@ pub async fn send_message(
     let llm = agent.llm.clone();
     let db = agent.db.clone();
     let message = req.message.clone();
+
+    // Use existing broadcast sender if available, otherwise create a new one
+    let event_tx = agent
+        .subscribe_session(&session_id)
+        .and_then(|_| {
+            agent
+                .session_streams
+                .get(&session_id)
+                .map(|entry| entry.value().clone())
+        })
+        .unwrap_or_else(|| agent.register_session_stream(&session_id));
+
     tokio::spawn(async move {
-        let orchestrator = PentestOrchestrator::new(llm, db);
+        let orchestrator = PentestOrchestrator::new(llm, db, event_tx, None);
         orchestrator
             .run_session_guarded(&session, &target, &message)
             .await;
@@ -277,10 +583,10 @@ pub async fn stop_session(
         })?
         .ok_or_else(|| (StatusCode::NOT_FOUND, "Session not found".to_string()))?;
 
-    if session.status != PentestStatus::Running {
+    if session.status != PentestStatus::Running && session.status != PentestStatus::Paused {
         return Err((
             StatusCode::BAD_REQUEST,
-            format!("Session is {}, not running", session.status),
+            format!("Session is {}, not running or paused", session.status),
         ));
     }
 
@@ -303,6 +609,9 @@ pub async fn stop_session(
             )
         })?;
 
+    // Clean up session resources
+    agent.cleanup_session(&id);
+
     let updated = agent
         .db
         .pentest_sessions()
@@ -328,6 +637,92 @@ pub async fn stop_session(
     }))
 }
 
+/// POST /api/v1/pentest/sessions/:id/pause — Pause a running pentest session
+#[tracing::instrument(skip_all, fields(session_id = %id))]
+pub async fn pause_session(
+    Extension(agent): AgentExt,
+    Path(id): Path<String>,
+) -> Result<Json<ApiResponse<serde_json::Value>>, (StatusCode, String)> {
+    let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
+        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
+
+    let session = agent
+        .db
+        .pentest_sessions()
+        .find_one(doc! { "_id": oid })
+        .await
+        .map_err(|e| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Database error: {e}"),
+            )
+        })?
+        .ok_or_else(|| (StatusCode::NOT_FOUND, "Session not found".to_string()))?;
+
+    if session.status != PentestStatus::Running {
+        return Err((
+            StatusCode::BAD_REQUEST,
+            format!("Session is {}, not running", session.status),
+        ));
+    }
+
+    if !agent.pause_session(&id) {
+        return Err((
+            StatusCode::INTERNAL_SERVER_ERROR,
+            "Failed to send pause signal".to_string(),
+        ));
+    }
+
+    Ok(Json(ApiResponse {
+        data: serde_json::json!({ "status": "paused" }),
+        total: None,
+        page: None,
+    }))
+}
+
+/// POST /api/v1/pentest/sessions/:id/resume — Resume a paused pentest session
+#[tracing::instrument(skip_all, fields(session_id = %id))]
+pub async fn resume_session(
+    Extension(agent): AgentExt,
+    Path(id): Path<String>,
+) -> Result<Json<ApiResponse<serde_json::Value>>, (StatusCode, String)> {
+    let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
+        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
+
+    let session = agent
+        .db
+        .pentest_sessions()
+        .find_one(doc! { "_id": oid })
+        .await
+        .map_err(|e| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Database error: {e}"),
+            )
+        })?
+        .ok_or_else(|| (StatusCode::NOT_FOUND, "Session not found".to_string()))?;
+
+    if session.status != PentestStatus::Paused {
+        return Err((
+            StatusCode::BAD_REQUEST,
+            format!("Session is {}, not paused", session.status),
+        ));
+    }
+
+    if !agent.resume_session(&id) {
+        return Err((
+            StatusCode::INTERNAL_SERVER_ERROR,
+            "Failed to send resume signal".to_string(),
+        ));
+    }
+
+    Ok(Json(ApiResponse {
+        data: serde_json::json!({ "status": "running" }),
+        total: None,
+        page: None,
+    }))
+}
+
 /// GET /api/v1/pentest/sessions/:id/attack-chain — Get attack chain nodes for a session
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn get_attack_chain(

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"),
+    ))
 }

compliance-agent/src/api/routes.rs

@@ -100,6 +100,10 @@ pub fn build_router() -> Router {
             get(handlers::chat::embedding_status),
         )
         // Pentest API endpoints
+        .route(
+            "/api/v1/pentest/lookup-repo",
+            get(handlers::pentest::lookup_repo),
+        )
         .route(
             "/api/v1/pentest/sessions",
             get(handlers::pentest::list_sessions).post(handlers::pentest::create_session),
@@ -116,6 +120,14 @@ pub fn build_router() -> Router {
             "/api/v1/pentest/sessions/{id}/stop",
             post(handlers::pentest::stop_session),
         )
+        .route(
+            "/api/v1/pentest/sessions/{id}/pause",
+            post(handlers::pentest::pause_session),
+        )
+        .route(
+            "/api/v1/pentest/sessions/{id}/resume",
+            post(handlers::pentest::resume_session),
+        )
         .route(
             "/api/v1/pentest/sessions/{id}/stream",
             get(handlers::pentest::session_stream),