refactor: modularize codebase and add 404 unit tests (#13)

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

@@ -0,0 +1,116 @@
+use std::sync::Arc;
+
+use axum::extract::{Extension, Path};
+use axum::http::StatusCode;
+use axum::response::sse::{Event, Sse};
+use futures_util::stream;
+use mongodb::bson::doc;
+
+use compliance_core::models::pentest::*;
+
+use crate::agent::ComplianceAgent;
+
+use super::super::dto::collect_cursor_async;
+
+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.
+#[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,
+> {
+    let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
+
+    // Verify session exists
+    let _session = agent
+        .db
+        .pentest_sessions()
+        .find_one(doc! { "_id": oid })
+        .await
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
+        .ok_or(StatusCode::NOT_FOUND)?;
+
+    // Fetch recent messages for this session
+    let messages: Vec<PentestMessage> = match agent
+        .db
+        .pentest_messages()
+        .find(doc! { "session_id": &id })
+        .sort(doc! { "created_at": 1 })
+        .limit(100)
+        .await
+    {
+        Ok(cursor) => collect_cursor_async(cursor).await,
+        Err(_) => Vec::new(),
+    };
+
+    // Fetch recent attack chain nodes
+    let nodes: Vec<AttackChainNode> = match agent
+        .db
+        .attack_chain_nodes()
+        .find(doc! { "session_id": &id })
+        .sort(doc! { "started_at": 1 })
+        .limit(100)
+        .await
+    {
+        Ok(cursor) => collect_cursor_async(cursor).await,
+        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",
+            "role": msg.role,
+            "content": msg.content,
+            "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)));
+        }
+    }
+
+    for node in &nodes {
+        let event_data = serde_json::json!({
+            "type": "tool_execution",
+            "node_id": node.node_id,
+            "tool_name": node.tool_name,
+            "status": node.status,
+            "findings_produced": node.findings_produced,
+        });
+        if let Ok(data) = serde_json::to_string(&event_data) {
+            events.push(Ok(Event::default().event("tool").data(data)));
+        }
+    }
+
+    // Add session status event
+    let session = agent
+        .db
+        .pentest_sessions()
+        .find_one(doc! { "_id": oid })
+        .await
+        .ok()
+        .flatten();
+
+    if let Some(s) = session {
+        let status_data = serde_json::json!({
+            "type": "status",
+            "status": s.status,
+            "findings_count": s.findings_count,
+            "tool_invocations": s.tool_invocations,
+        });
+        if let Ok(data) = serde_json::to_string(&status_data) {
+            events.push(Ok(Event::default().event("status").data(data)));
+        }
+    }
+
+    Ok(Sse::new(stream::iter(events)))
+}