feat: AI-driven automated penetration testing (#12)

compliance-agent/src/api/handlers/mod.rs

@@ -1,6 +1,7 @@
 pub mod chat;
 pub mod dast;
 pub mod graph;
+pub mod pentest;
 
 use std::sync::Arc;
 
@@ -1108,7 +1109,7 @@ pub async fn list_scan_runs(
     }))
 }
 
-async fn collect_cursor_async<T: serde::de::DeserializeOwned + Unpin + Send>(
+pub(crate) async fn collect_cursor_async<T: serde::de::DeserializeOwned + Unpin + Send>(
     mut cursor: mongodb::Cursor<T>,
 ) -> Vec<T> {
     use futures_util::StreamExt;

compliance-agent/src/api/handlers/pentest.rs

@@ -0,0 +1,718 @@
+use std::sync::Arc;
+
+use axum::extract::{Extension, Path, Query};
+use axum::http::StatusCode;
+use axum::response::sse::{Event, Sse};
+use axum::response::IntoResponse;
+use axum::Json;
+use futures_util::stream;
+use mongodb::bson::doc;
+use serde::Deserialize;
+
+use compliance_core::models::dast::DastFinding;
+use compliance_core::models::pentest::*;
+
+use crate::agent::ComplianceAgent;
+use crate::pentest::PentestOrchestrator;
+
+use super::{collect_cursor_async, ApiResponse, PaginationParams};
+
+type AgentExt = Extension<Arc<ComplianceAgent>>;
+
+#[derive(Deserialize)]
+pub struct CreateSessionRequest {
+    pub target_id: String,
+    #[serde(default = "default_strategy")]
+    pub strategy: String,
+    pub message: Option<String>,
+}
+
+fn default_strategy() -> String {
+    "comprehensive".to_string()
+}
+
+#[derive(Deserialize)]
+pub struct SendMessageRequest {
+    pub message: 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| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Database error: {e}"),
+            )
+        })?
+        .ok_or_else(|| (StatusCode::NOT_FOUND, "Target not found".to_string()))?;
+
+    // Parse strategy
+    let strategy = match req.strategy.as_str() {
+        "quick" => PentestStrategy::Quick,
+        "targeted" => PentestStrategy::Targeted,
+        "aggressive" => PentestStrategy::Aggressive,
+        "stealth" => PentestStrategy::Stealth,
+        _ => PentestStrategy::Comprehensive,
+    };
+
+    // 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,
+        total: None,
+        page: None,
+    }))
+}
+
+/// GET /api/v1/pentest/sessions — List pentest sessions
+#[tracing::instrument(skip_all)]
+pub async fn list_sessions(
+    Extension(agent): AgentExt,
+    Query(params): Query<PaginationParams>,
+) -> Result<Json<ApiResponse<Vec<PentestSession>>>, StatusCode> {
+    let db = &agent.db;
+    let skip = (params.page.saturating_sub(1)) * params.limit as u64;
+    let total = db
+        .pentest_sessions()
+        .count_documents(doc! {})
+        .await
+        .unwrap_or(0);
+
+    let sessions = match db
+        .pentest_sessions()
+        .find(doc! {})
+        .sort(doc! { "started_at": -1 })
+        .skip(skip)
+        .limit(params.limit)
+        .await
+    {
+        Ok(cursor) => collect_cursor_async(cursor).await,
+        Err(e) => {
+            tracing::warn!("Failed to fetch pentest sessions: {e}");
+            Vec::new()
+        }
+    };
+
+    Ok(Json(ApiResponse {
+        data: sessions,
+        total: Some(total),
+        page: Some(params.page),
+    }))
+}
+
+/// GET /api/v1/pentest/sessions/:id — Get a single pentest session
+#[tracing::instrument(skip_all, fields(session_id = %id))]
+pub async fn get_session(
+    Extension(agent): AgentExt,
+    Path(id): Path<String>,
+) -> Result<Json<ApiResponse<PentestSession>>, StatusCode> {
+    let oid =
+        mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
+
+    let session = agent
+        .db
+        .pentest_sessions()
+        .find_one(doc! { "_id": oid })
+        .await
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
+        .ok_or(StatusCode::NOT_FOUND)?;
+
+    Ok(Json(ApiResponse {
+        data: session,
+        total: None,
+        page: None,
+    }))
+}
+
+/// POST /api/v1/pentest/sessions/:id/chat — Send a user message and trigger next orchestrator iteration
+#[tracing::instrument(skip_all, fields(session_id = %id))]
+pub async fn send_message(
+    Extension(agent): AgentExt,
+    Path(id): Path<String>,
+    Json(req): Json<SendMessageRequest>,
+) -> Result<Json<ApiResponse<PentestMessage>>, (StatusCode, String)> {
+    let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
+        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
+
+    // Verify session exists and is running
+    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 && session.status != PentestStatus::Paused {
+        return Err((
+            StatusCode::BAD_REQUEST,
+            format!("Session is {}, cannot send messages", session.status),
+        ));
+    }
+
+    // Look up the target
+    let target_oid =
+        mongodb::bson::oid::ObjectId::parse_str(&session.target_id).map_err(|_| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Invalid target_id in session".to_string(),
+            )
+        })?;
+
+    let target = agent
+        .db
+        .dast_targets()
+        .find_one(doc! { "_id": target_oid })
+        .await
+        .map_err(|e| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Database error: {e}"),
+            )
+        })?
+        .ok_or_else(|| {
+            (
+                StatusCode::NOT_FOUND,
+                "Target for session not found".to_string(),
+            )
+        })?;
+
+    // Store user message
+    let session_id = id.clone();
+    let user_msg = PentestMessage::user(session_id.clone(), req.message.clone());
+    let _ = agent.db.pentest_messages().insert_one(&user_msg).await;
+
+    let response_msg = user_msg.clone();
+
+    // Spawn orchestrator to continue the session
+    let llm = agent.llm.clone();
+    let db = agent.db.clone();
+    let message = req.message.clone();
+    tokio::spawn(async move {
+        let orchestrator = PentestOrchestrator::new(llm, db);
+        orchestrator
+            .run_session_guarded(&session, &target, &message)
+            .await;
+    });
+
+    Ok(Json(ApiResponse {
+        data: response_msg,
+        total: None,
+        page: None,
+    }))
+}
+
+/// 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)))
+}
+
+/// POST /api/v1/pentest/sessions/:id/stop — Stop a running pentest session
+#[tracing::instrument(skip_all, fields(session_id = %id))]
+pub async fn stop_session(
+    Extension(agent): AgentExt,
+    Path(id): Path<String>,
+) -> Result<Json<ApiResponse<PentestSession>>, (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),
+        ));
+    }
+
+    agent
+        .db
+        .pentest_sessions()
+        .update_one(
+            doc! { "_id": oid },
+            doc! { "$set": {
+                "status": "failed",
+                "completed_at": mongodb::bson::DateTime::now(),
+                "error_message": "Stopped by user",
+            }},
+        )
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Database error: {e}")))?;
+
+    let updated = 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 after update".to_string()))?;
+
+    Ok(Json(ApiResponse {
+        data: updated,
+        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(
+    Extension(agent): AgentExt,
+    Path(id): Path<String>,
+) -> Result<Json<ApiResponse<Vec<AttackChainNode>>>, StatusCode> {
+    // Verify the session ID is valid
+    let _oid =
+        mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
+
+    let nodes = match agent
+        .db
+        .attack_chain_nodes()
+        .find(doc! { "session_id": &id })
+        .sort(doc! { "started_at": 1 })
+        .await
+    {
+        Ok(cursor) => collect_cursor_async(cursor).await,
+        Err(e) => {
+            tracing::warn!("Failed to fetch attack chain nodes: {e}");
+            Vec::new()
+        }
+    };
+
+    let total = nodes.len() as u64;
+    Ok(Json(ApiResponse {
+        data: nodes,
+        total: Some(total),
+        page: None,
+    }))
+}
+
+/// GET /api/v1/pentest/sessions/:id/messages — Get messages for a session
+#[tracing::instrument(skip_all, fields(session_id = %id))]
+pub async fn get_messages(
+    Extension(agent): AgentExt,
+    Path(id): Path<String>,
+    Query(params): Query<PaginationParams>,
+) -> Result<Json<ApiResponse<Vec<PentestMessage>>>, StatusCode> {
+    let _oid =
+        mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
+
+    let skip = (params.page.saturating_sub(1)) * params.limit as u64;
+    let total = agent
+        .db
+        .pentest_messages()
+        .count_documents(doc! { "session_id": &id })
+        .await
+        .unwrap_or(0);
+
+    let messages = match agent
+        .db
+        .pentest_messages()
+        .find(doc! { "session_id": &id })
+        .sort(doc! { "created_at": 1 })
+        .skip(skip)
+        .limit(params.limit)
+        .await
+    {
+        Ok(cursor) => collect_cursor_async(cursor).await,
+        Err(e) => {
+            tracing::warn!("Failed to fetch pentest messages: {e}");
+            Vec::new()
+        }
+    };
+
+    Ok(Json(ApiResponse {
+        data: messages,
+        total: Some(total),
+        page: Some(params.page),
+    }))
+}
+
+/// GET /api/v1/pentest/stats — Aggregated pentest statistics
+#[tracing::instrument(skip_all)]
+pub async fn pentest_stats(
+    Extension(agent): AgentExt,
+) -> Result<Json<ApiResponse<PentestStats>>, StatusCode> {
+    let db = &agent.db;
+
+    let running_sessions = db
+        .pentest_sessions()
+        .count_documents(doc! { "status": "running" })
+        .await
+        .unwrap_or(0) as u32;
+
+    // Count DAST findings from pentest sessions
+    let total_vulnerabilities = db
+        .dast_findings()
+        .count_documents(doc! { "session_id": { "$exists": true, "$ne": null } })
+        .await
+        .unwrap_or(0) as u32;
+
+    // Aggregate tool invocations from all sessions
+    let sessions: Vec<PentestSession> = match db.pentest_sessions().find(doc! {}).await {
+        Ok(cursor) => collect_cursor_async(cursor).await,
+        Err(_) => Vec::new(),
+    };
+
+    let total_tool_invocations: u32 = sessions.iter().map(|s| s.tool_invocations).sum();
+    let total_successes: u32 = sessions.iter().map(|s| s.tool_successes).sum();
+    let tool_success_rate = if total_tool_invocations == 0 {
+        100.0
+    } else {
+        (total_successes as f64 / total_tool_invocations as f64) * 100.0
+    };
+
+    // Severity distribution from pentest-related DAST findings
+    let critical = db
+        .dast_findings()
+        .count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "critical" })
+        .await
+        .unwrap_or(0) as u32;
+    let high = db
+        .dast_findings()
+        .count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "high" })
+        .await
+        .unwrap_or(0) as u32;
+    let medium = db
+        .dast_findings()
+        .count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "medium" })
+        .await
+        .unwrap_or(0) as u32;
+    let low = db
+        .dast_findings()
+        .count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "low" })
+        .await
+        .unwrap_or(0) as u32;
+    let info = db
+        .dast_findings()
+        .count_documents(doc! { "session_id": { "$exists": true, "$ne": null }, "severity": "info" })
+        .await
+        .unwrap_or(0) as u32;
+
+    Ok(Json(ApiResponse {
+        data: PentestStats {
+            running_sessions,
+            total_vulnerabilities,
+            total_tool_invocations,
+            tool_success_rate,
+            severity_distribution: SeverityDistribution {
+                critical,
+                high,
+                medium,
+                low,
+                info,
+            },
+        },
+        total: None,
+        page: None,
+    }))
+}
+
+/// GET /api/v1/pentest/sessions/:id/findings — Get DAST findings for a pentest session
+#[tracing::instrument(skip_all, fields(session_id = %id))]
+pub async fn get_session_findings(
+    Extension(agent): AgentExt,
+    Path(id): Path<String>,
+    Query(params): Query<PaginationParams>,
+) -> Result<Json<ApiResponse<Vec<DastFinding>>>, StatusCode> {
+    let _oid =
+        mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
+
+    let skip = (params.page.saturating_sub(1)) * params.limit as u64;
+    let total = agent
+        .db
+        .dast_findings()
+        .count_documents(doc! { "session_id": &id })
+        .await
+        .unwrap_or(0);
+
+    let findings = match agent
+        .db
+        .dast_findings()
+        .find(doc! { "session_id": &id })
+        .sort(doc! { "created_at": -1 })
+        .skip(skip)
+        .limit(params.limit)
+        .await
+    {
+        Ok(cursor) => collect_cursor_async(cursor).await,
+        Err(e) => {
+            tracing::warn!("Failed to fetch pentest session findings: {e}");
+            Vec::new()
+        }
+    };
+
+    Ok(Json(ApiResponse {
+        data: findings,
+        total: Some(total),
+        page: Some(params.page),
+    }))
+}
+
+#[derive(Deserialize)]
+pub struct ExportBody {
+    pub password: String,
+    /// Requester display name (from auth)
+    #[serde(default)]
+    pub requester_name: String,
+    /// Requester email (from auth)
+    #[serde(default)]
+    pub requester_email: String,
+}
+
+/// POST /api/v1/pentest/sessions/:id/export — Export an encrypted pentest report archive
+#[tracing::instrument(skip_all, fields(session_id = %id))]
+pub async fn export_session_report(
+    Extension(agent): AgentExt,
+    Path(id): Path<String>,
+    Json(body): Json<ExportBody>,
+) -> Result<axum::response::Response, (StatusCode, String)> {
+    let oid = mongodb::bson::oid::ObjectId::parse_str(&id)
+        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID".to_string()))?;
+
+    if body.password.len() < 8 {
+        return Err((
+            StatusCode::BAD_REQUEST,
+            "Password must be at least 8 characters".to_string(),
+        ));
+    }
+
+    // Fetch session
+    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()))?;
+
+    // Resolve target name
+    let target = if let Ok(tid) = mongodb::bson::oid::ObjectId::parse_str(&session.target_id) {
+        agent
+            .db
+            .dast_targets()
+            .find_one(doc! { "_id": tid })
+            .await
+            .ok()
+            .flatten()
+    } else {
+        None
+    };
+    let target_name = target
+        .as_ref()
+        .map(|t| t.name.clone())
+        .unwrap_or_else(|| "Unknown Target".to_string());
+    let target_url = target
+        .as_ref()
+        .map(|t| t.base_url.clone())
+        .unwrap_or_default();
+
+    // Fetch attack chain nodes
+    let nodes: Vec<AttackChainNode> = match agent
+        .db
+        .attack_chain_nodes()
+        .find(doc! { "session_id": &id })
+        .sort(doc! { "started_at": 1 })
+        .await
+    {
+        Ok(cursor) => collect_cursor_async(cursor).await,
+        Err(_) => Vec::new(),
+    };
+
+    // Fetch DAST findings for this session
+    let findings: Vec<DastFinding> = match agent
+        .db
+        .dast_findings()
+        .find(doc! { "session_id": &id })
+        .sort(doc! { "severity": -1, "created_at": -1 })
+        .await
+    {
+        Ok(cursor) => collect_cursor_async(cursor).await,
+        Err(_) => Vec::new(),
+    };
+
+    let ctx = crate::pentest::report::ReportContext {
+        session,
+        target_name,
+        target_url,
+        findings,
+        attack_chain: nodes,
+        requester_name: if body.requester_name.is_empty() {
+            "Unknown".to_string()
+        } else {
+            body.requester_name
+        },
+        requester_email: body.requester_email,
+    };
+
+    let report = crate::pentest::generate_encrypted_report(&ctx, &body.password)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
+
+    let response = serde_json::json!({
+        "archive_base64": base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &report.archive),
+        "sha256": report.sha256,
+        "filename": format!("pentest-report-{id}.zip"),
+    });
+
+    Ok(Json(response).into_response())
+}

compliance-agent/src/api/routes.rs

@@ -99,6 +99,44 @@ pub fn build_router() -> Router {
             "/api/v1/chat/{repo_id}/status",
             get(handlers::chat::embedding_status),
         )
+        // Pentest API endpoints
+        .route(
+            "/api/v1/pentest/sessions",
+            get(handlers::pentest::list_sessions).post(handlers::pentest::create_session),
+        )
+        .route(
+            "/api/v1/pentest/sessions/{id}",
+            get(handlers::pentest::get_session),
+        )
+        .route(
+            "/api/v1/pentest/sessions/{id}/chat",
+            post(handlers::pentest::send_message),
+        )
+        .route(
+            "/api/v1/pentest/sessions/{id}/stop",
+            post(handlers::pentest::stop_session),
+        )
+        .route(
+            "/api/v1/pentest/sessions/{id}/stream",
+            get(handlers::pentest::session_stream),
+        )
+        .route(
+            "/api/v1/pentest/sessions/{id}/attack-chain",
+            get(handlers::pentest::get_attack_chain),
+        )
+        .route(
+            "/api/v1/pentest/sessions/{id}/messages",
+            get(handlers::pentest::get_messages),
+        )
+        .route(
+            "/api/v1/pentest/sessions/{id}/findings",
+            get(handlers::pentest::get_session_findings),
+        )
+        .route(
+            "/api/v1/pentest/sessions/{id}/export",
+            post(handlers::pentest::export_session_report),
+        )
+        .route("/api/v1/pentest/stats", get(handlers::pentest::pentest_stats))
         // Webhook endpoints (proxied through dashboard)
         .route(
             "/webhook/github/{repo_id}",

compliance-agent/src/database.rs

@@ -166,6 +166,38 @@ impl Database {
             )
             .await?;
 
+        // pentest_sessions: compound (target_id, started_at DESC)
+        self.pentest_sessions()
+            .create_index(
+                IndexModel::builder()
+                    .keys(doc! { "target_id": 1, "started_at": -1 })
+                    .build(),
+            )
+            .await?;
+
+        // pentest_sessions: status index
+        self.pentest_sessions()
+            .create_index(IndexModel::builder().keys(doc! { "status": 1 }).build())
+            .await?;
+
+        // attack_chain_nodes: compound (session_id, node_id)
+        self.attack_chain_nodes()
+            .create_index(
+                IndexModel::builder()
+                    .keys(doc! { "session_id": 1, "node_id": 1 })
+                    .build(),
+            )
+            .await?;
+
+        // pentest_messages: compound (session_id, created_at)
+        self.pentest_messages()
+            .create_index(
+                IndexModel::builder()
+                    .keys(doc! { "session_id": 1, "created_at": 1 })
+                    .build(),
+            )
+            .await?;
+
         tracing::info!("Database indexes ensured");
         Ok(())
     }
@@ -235,6 +267,19 @@ impl Database {
         self.inner.collection("embedding_builds")
     }
 
+    // Pentest collections
+    pub fn pentest_sessions(&self) -> Collection<PentestSession> {
+        self.inner.collection("pentest_sessions")
+    }
+
+    pub fn attack_chain_nodes(&self) -> Collection<AttackChainNode> {
+        self.inner.collection("attack_chain_nodes")
+    }
+
+    pub fn pentest_messages(&self) -> Collection<PentestMessage> {
+        self.inner.collection("pentest_messages")
+    }
+
     #[allow(dead_code)]
     pub fn raw_collection(&self, name: &str) -> Collection<mongodb::bson::Document> {
         self.inner.collection(name)

compliance-agent/src/llm/client.rs

@@ -12,10 +12,16 @@ pub struct LlmClient {
     http: reqwest::Client,
 }
 
-#[derive(Serialize)]
-struct ChatMessage {
-    role: String,
-    content: String,
+// ── Request types ──────────────────────────────────────────────
+
+#[derive(Serialize, Clone, Debug)]
+pub struct ChatMessage {
+    pub role: String,
+    pub content: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_calls: Option<Vec<ToolCallRequest>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_call_id: Option<String>,
 }
 
 #[derive(Serialize)]
@@ -26,8 +32,25 @@ struct ChatCompletionRequest {
     temperature: Option<f64>,
     #[serde(skip_serializing_if = "Option::is_none")]
     max_tokens: Option<u32>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    tools: Option<Vec<ToolDefinitionPayload>>,
 }
 
+#[derive(Serialize)]
+struct ToolDefinitionPayload {
+    r#type: String,
+    function: ToolFunctionPayload,
+}
+
+#[derive(Serialize)]
+struct ToolFunctionPayload {
+    name: String,
+    description: String,
+    parameters: serde_json::Value,
+}
+
+// ── Response types ─────────────────────────────────────────────
+
 #[derive(Deserialize)]
 struct ChatCompletionResponse {
     choices: Vec<ChatChoice>,
@@ -40,29 +63,85 @@ struct ChatChoice {
 
 #[derive(Deserialize)]
 struct ChatResponseMessage {
-    content: String,
+    #[serde(default)]
+    content: Option<String>,
+    #[serde(default)]
+    tool_calls: Option<Vec<ToolCallResponse>>,
 }
 
-/// Request body for the embeddings API
+#[derive(Deserialize)]
+struct ToolCallResponse {
+    id: String,
+    function: ToolCallFunction,
+}
+
+#[derive(Deserialize)]
+struct ToolCallFunction {
+    name: String,
+    arguments: String,
+}
+
+// ── Public types for tool calling ──────────────────────────────
+
+/// Definition of a tool that the LLM can invoke
+#[derive(Debug, Clone, Serialize)]
+pub struct ToolDefinition {
+    pub name: String,
+    pub description: String,
+    pub parameters: serde_json::Value,
+}
+
+/// A tool call request from the LLM
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LlmToolCall {
+    pub id: String,
+    pub name: String,
+    pub arguments: serde_json::Value,
+}
+
+/// A tool call in the request message format (for sending back tool_calls in assistant messages)
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ToolCallRequest {
+    pub id: String,
+    pub r#type: String,
+    pub function: ToolCallRequestFunction,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ToolCallRequestFunction {
+    pub name: String,
+    pub arguments: String,
+}
+
+/// Response from the LLM — either content or tool calls
+#[derive(Debug, Clone)]
+pub enum LlmResponse {
+    Content(String),
+    /// Tool calls with optional reasoning text from the LLM
+    ToolCalls { calls: Vec<LlmToolCall>, reasoning: String },
+}
+
+// ── Embedding types ────────────────────────────────────────────
+
 #[derive(Serialize)]
 struct EmbeddingRequest {
     model: String,
     input: Vec<String>,
 }
 
-/// Response from the embeddings API
 #[derive(Deserialize)]
 struct EmbeddingResponse {
     data: Vec<EmbeddingData>,
 }
 
-/// A single embedding result
 #[derive(Deserialize)]
 struct EmbeddingData {
     embedding: Vec<f64>,
     index: usize,
 }
 
+// ── Implementation ─────────────────────────────────────────────
+
 impl LlmClient {
     pub fn new(
         base_url: String,
@@ -83,98 +162,142 @@ impl LlmClient {
         &self.embed_model
     }
 
+    fn chat_url(&self) -> String {
+        format!(
+            "{}/v1/chat/completions",
+            self.base_url.trim_end_matches('/')
+        )
+    }
+
+    fn auth_header(&self) -> Option<String> {
+        let key = self.api_key.expose_secret();
+        if key.is_empty() {
+            None
+        } else {
+            Some(format!("Bearer {key}"))
+        }
+    }
+
+    /// Simple chat: system + user prompt → text response
     pub async fn chat(
         &self,
         system_prompt: &str,
         user_prompt: &str,
         temperature: Option<f64>,
     ) -> Result<String, AgentError> {
-        let url = format!(
-            "{}/v1/chat/completions",
-            self.base_url.trim_end_matches('/')
-        );
+        let messages = vec![
+            ChatMessage {
+                role: "system".to_string(),
+                content: Some(system_prompt.to_string()),
+                tool_calls: None,
+                tool_call_id: None,
+            },
+            ChatMessage {
+                role: "user".to_string(),
+                content: Some(user_prompt.to_string()),
+                tool_calls: None,
+                tool_call_id: None,
+            },
+        ];
 
         let request_body = ChatCompletionRequest {
             model: self.model.clone(),
-            messages: vec![
-                ChatMessage {
-                    role: "system".to_string(),
-                    content: system_prompt.to_string(),
-                },
-                ChatMessage {
-                    role: "user".to_string(),
-                    content: user_prompt.to_string(),
-                },
-            ],
+            messages,
             temperature,
             max_tokens: Some(4096),
+            tools: None,
         };
 
-        let mut req = self
-            .http
-            .post(&url)
-            .header("content-type", "application/json")
-            .json(&request_body);
-
-        let key = self.api_key.expose_secret();
-        if !key.is_empty() {
-            req = req.header("Authorization", format!("Bearer {key}"));
-        }
-
-        let resp = req
-            .send()
-            .await
-            .map_err(|e| AgentError::Other(format!("LiteLLM request failed: {e}")))?;
-
-        if !resp.status().is_success() {
-            let status = resp.status();
-            let body = resp.text().await.unwrap_or_default();
-            return Err(AgentError::Other(format!(
-                "LiteLLM returned {status}: {body}"
-            )));
-        }
-
-        let body: ChatCompletionResponse = resp
-            .json()
-            .await
-            .map_err(|e| AgentError::Other(format!("Failed to parse LiteLLM response: {e}")))?;
-
-        body.choices
-            .first()
-            .map(|c| c.message.content.clone())
-            .ok_or_else(|| AgentError::Other("Empty response from LiteLLM".to_string()))
+        self.send_chat_request(&request_body).await.map(|resp| {
+            match resp {
+                LlmResponse::Content(c) => c,
+                LlmResponse::ToolCalls { .. } => String::new(), // shouldn't happen without tools
+            }
+        })
     }
 
+    /// Chat with a list of (role, content) messages → text response
     #[allow(dead_code)]
     pub async fn chat_with_messages(
         &self,
         messages: Vec<(String, String)>,
         temperature: Option<f64>,
     ) -> Result<String, AgentError> {
-        let url = format!(
-            "{}/v1/chat/completions",
-            self.base_url.trim_end_matches('/')
-        );
+        let messages = messages
+            .into_iter()
+            .map(|(role, content)| ChatMessage {
+                role,
+                content: Some(content),
+                tool_calls: None,
+                tool_call_id: None,
+            })
+            .collect();
 
         let request_body = ChatCompletionRequest {
             model: self.model.clone(),
-            messages: messages
-                .into_iter()
-                .map(|(role, content)| ChatMessage { role, content })
-                .collect(),
+            messages,
             temperature,
             max_tokens: Some(4096),
+            tools: None,
         };
 
+        self.send_chat_request(&request_body).await.map(|resp| {
+            match resp {
+                LlmResponse::Content(c) => c,
+                LlmResponse::ToolCalls { .. } => String::new(),
+            }
+        })
+    }
+
+    /// Chat with tool definitions — returns either content or tool calls.
+    /// Use this for the AI pentest orchestrator loop.
+    pub async fn chat_with_tools(
+        &self,
+        messages: Vec<ChatMessage>,
+        tools: &[ToolDefinition],
+        temperature: Option<f64>,
+        max_tokens: Option<u32>,
+    ) -> Result<LlmResponse, AgentError> {
+        let tool_payloads: Vec<ToolDefinitionPayload> = tools
+            .iter()
+            .map(|t| ToolDefinitionPayload {
+                r#type: "function".to_string(),
+                function: ToolFunctionPayload {
+                    name: t.name.clone(),
+                    description: t.description.clone(),
+                    parameters: t.parameters.clone(),
+                },
+            })
+            .collect();
+
+        let request_body = ChatCompletionRequest {
+            model: self.model.clone(),
+            messages,
+            temperature,
+            max_tokens: Some(max_tokens.unwrap_or(8192)),
+            tools: if tool_payloads.is_empty() {
+                None
+            } else {
+                Some(tool_payloads)
+            },
+        };
+
+        self.send_chat_request(&request_body).await
+    }
+
+    /// Internal method to send a chat completion request and parse the response
+    async fn send_chat_request(
+        &self,
+        request_body: &ChatCompletionRequest,
+    ) -> Result<LlmResponse, AgentError> {
         let mut req = self
             .http
-            .post(&url)
+            .post(&self.chat_url())
             .header("content-type", "application/json")
-            .json(&request_body);
+            .json(request_body);
 
-        let key = self.api_key.expose_secret();
-        if !key.is_empty() {
-            req = req.header("Authorization", format!("Bearer {key}"));
+        if let Some(auth) = self.auth_header() {
+            req = req.header("Authorization", auth);
         }
 
         let resp = req
@@ -195,10 +318,39 @@ impl LlmClient {
             .await
             .map_err(|e| AgentError::Other(format!("Failed to parse LiteLLM response: {e}")))?;
 
-        body.choices
+        let choice = body
+            .choices
             .first()
-            .map(|c| c.message.content.clone())
-            .ok_or_else(|| AgentError::Other("Empty response from LiteLLM".to_string()))
+            .ok_or_else(|| AgentError::Other("Empty response from LiteLLM".to_string()))?;
+
+        // Check for tool calls first
+        if let Some(tool_calls) = &choice.message.tool_calls {
+            if !tool_calls.is_empty() {
+                let calls: Vec<LlmToolCall> = tool_calls
+                    .iter()
+                    .map(|tc| {
+                        let arguments = serde_json::from_str(&tc.function.arguments)
+                            .unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
+                        LlmToolCall {
+                            id: tc.id.clone(),
+                            name: tc.function.name.clone(),
+                            arguments,
+                        }
+                    })
+                    .collect();
+                // Capture any reasoning text the LLM included alongside tool calls
+                let reasoning = choice.message.content.clone().unwrap_or_default();
+                return Ok(LlmResponse::ToolCalls { calls, reasoning });
+            }
+        }
+
+        // Otherwise return content
+        let content = choice
+            .message
+            .content
+            .clone()
+            .unwrap_or_default();
+        Ok(LlmResponse::Content(content))
     }
 
     /// Generate embeddings for a batch of texts
@@ -216,9 +368,8 @@ impl LlmClient {
             .header("content-type", "application/json")
             .json(&request_body);
 
-        let key = self.api_key.expose_secret();
-        if !key.is_empty() {
-            req = req.header("Authorization", format!("Bearer {key}"));
+        if let Some(auth) = self.auth_header() {
+            req = req.header("Authorization", auth);
         }
 
         let resp = req
@@ -239,7 +390,6 @@ impl LlmClient {
             .await
             .map_err(|e| AgentError::Other(format!("Failed to parse embedding response: {e}")))?;
 
-        // Sort by index to maintain input order
         let mut data = body.data;
         data.sort_by_key(|d| d.index);
 

compliance-agent/src/main.rs

@@ -4,6 +4,7 @@ mod config;
 mod database;
 mod error;
 mod llm;
+mod pentest;
 mod pipeline;
 mod rag;
 mod scheduler;

compliance-agent/src/pentest/mod.rs

@@ -0,0 +1,5 @@
+pub mod orchestrator;
+pub mod report;
+
+pub use orchestrator::PentestOrchestrator;
+pub use report::generate_encrypted_report;

compliance-agent/src/pentest/orchestrator.rs

@@ -0,0 +1,761 @@
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures_util::StreamExt;
+use mongodb::bson::doc;
+use tokio::sync::broadcast;
+
+use compliance_core::models::dast::DastTarget;
+use compliance_core::models::finding::{Finding, FindingStatus, Severity};
+use compliance_core::models::pentest::*;
+use compliance_core::models::sbom::SbomEntry;
+use compliance_core::traits::pentest_tool::PentestToolContext;
+use compliance_dast::ToolRegistry;
+
+use crate::database::Database;
+use crate::llm::client::{
+    ChatMessage, LlmResponse, ToolCallRequest, ToolCallRequestFunction, ToolDefinition,
+};
+use crate::llm::LlmClient;
+
+/// Maximum duration for a single pentest session before timeout
+const SESSION_TIMEOUT: Duration = Duration::from_secs(30 * 60); // 30 minutes
+
+pub struct PentestOrchestrator {
+    tool_registry: ToolRegistry,
+    llm: Arc<LlmClient>,
+    db: Database,
+    event_tx: broadcast::Sender<PentestEvent>,
+}
+
+impl PentestOrchestrator {
+    pub fn new(llm: Arc<LlmClient>, db: Database) -> Self {
+        let (event_tx, _) = broadcast::channel(256);
+        Self {
+            tool_registry: ToolRegistry::new(),
+            llm,
+            db,
+            event_tx,
+        }
+    }
+
+    pub fn subscribe(&self) -> broadcast::Receiver<PentestEvent> {
+        self.event_tx.subscribe()
+    }
+
+    pub fn event_sender(&self) -> broadcast::Sender<PentestEvent> {
+        self.event_tx.clone()
+    }
+
+    /// Run a pentest session with timeout and automatic failure marking on errors.
+    pub async fn run_session_guarded(
+        &self,
+        session: &PentestSession,
+        target: &DastTarget,
+        initial_message: &str,
+    ) {
+        let session_id = session.id;
+
+        match tokio::time::timeout(
+            SESSION_TIMEOUT,
+            self.run_session(session, target, initial_message),
+        )
+        .await
+        {
+            Ok(Ok(())) => {
+                tracing::info!(?session_id, "Pentest session completed successfully");
+            }
+            Ok(Err(e)) => {
+                tracing::error!(?session_id, error = %e, "Pentest session failed");
+                self.mark_session_failed(session_id, &format!("Error: {e}"))
+                    .await;
+                let _ = self.event_tx.send(PentestEvent::Error {
+                    message: format!("Session failed: {e}"),
+                });
+            }
+            Err(_) => {
+                tracing::warn!(?session_id, "Pentest session timed out after 30 minutes");
+                self.mark_session_failed(session_id, "Session timed out after 30 minutes")
+                    .await;
+                let _ = self.event_tx.send(PentestEvent::Error {
+                    message: "Session timed out after 30 minutes".to_string(),
+                });
+            }
+        }
+    }
+
+    async fn mark_session_failed(
+        &self,
+        session_id: Option<mongodb::bson::oid::ObjectId>,
+        reason: &str,
+    ) {
+        if let Some(sid) = session_id {
+            let _ = self
+                .db
+                .pentest_sessions()
+                .update_one(
+                    doc! { "_id": sid },
+                    doc! { "$set": {
+                        "status": "failed",
+                        "completed_at": mongodb::bson::DateTime::now(),
+                        "error_message": reason,
+                    }},
+                )
+                .await;
+        }
+    }
+
+    async fn run_session(
+        &self,
+        session: &PentestSession,
+        target: &DastTarget,
+        initial_message: &str,
+    ) -> Result<(), crate::error::AgentError> {
+        let session_id = session
+            .id
+            .map(|oid| oid.to_hex())
+            .unwrap_or_default();
+
+        // Gather code-awareness context from linked repo
+        let (sast_findings, sbom_entries, code_context) =
+            self.gather_repo_context(target).await;
+
+        // Build system prompt with code context
+        let system_prompt = self
+            .build_system_prompt(session, target, &sast_findings, &sbom_entries, &code_context)
+            .await;
+
+        // Build tool definitions for LLM
+        let tool_defs: Vec<ToolDefinition> = self
+            .tool_registry
+            .all_definitions()
+            .into_iter()
+            .map(|td| ToolDefinition {
+                name: td.name,
+                description: td.description,
+                parameters: td.input_schema,
+            })
+            .collect();
+
+        // Initialize messages
+        let mut messages = vec![
+            ChatMessage {
+                role: "system".to_string(),
+                content: Some(system_prompt),
+                tool_calls: None,
+                tool_call_id: None,
+            },
+            ChatMessage {
+                role: "user".to_string(),
+                content: Some(initial_message.to_string()),
+                tool_calls: None,
+                tool_call_id: None,
+            },
+        ];
+
+        // Store user message
+        let user_msg = PentestMessage::user(session_id.clone(), initial_message.to_string());
+        let _ = self.db.pentest_messages().insert_one(&user_msg).await;
+
+        // Build tool context with real data
+        let tool_context = PentestToolContext {
+            target: target.clone(),
+            session_id: session_id.clone(),
+            sast_findings,
+            sbom_entries,
+            code_context,
+            rate_limit: target.rate_limit,
+            allow_destructive: target.allow_destructive,
+        };
+
+        let max_iterations = 50;
+        let mut total_findings = 0u32;
+        let mut total_tool_calls = 0u32;
+        let mut total_successes = 0u32;
+        let mut prev_node_ids: Vec<String> = Vec::new();
+
+        for _iteration in 0..max_iterations {
+            let response = self
+                .llm
+                .chat_with_tools(messages.clone(), &tool_defs, Some(0.2), Some(8192))
+                .await?;
+
+            match response {
+                LlmResponse::Content(content) => {
+                    let msg =
+                        PentestMessage::assistant(session_id.clone(), content.clone());
+                    let _ = self.db.pentest_messages().insert_one(&msg).await;
+                    let _ = self.event_tx.send(PentestEvent::Message {
+                        content: content.clone(),
+                    });
+
+                    messages.push(ChatMessage {
+                        role: "assistant".to_string(),
+                        content: Some(content.clone()),
+                        tool_calls: None,
+                        tool_call_id: None,
+                    });
+
+                    let done_indicators = [
+                        "pentest complete",
+                        "testing complete",
+                        "scan complete",
+                        "analysis complete",
+                        "finished",
+                        "that concludes",
+                    ];
+                    let content_lower = content.to_lowercase();
+                    if done_indicators
+                        .iter()
+                        .any(|ind| content_lower.contains(ind))
+                    {
+                        break;
+                    }
+                    break;
+                }
+                LlmResponse::ToolCalls { calls: tool_calls, reasoning } => {
+                    let tc_requests: Vec<ToolCallRequest> = tool_calls
+                        .iter()
+                        .map(|tc| ToolCallRequest {
+                            id: tc.id.clone(),
+                            r#type: "function".to_string(),
+                            function: ToolCallRequestFunction {
+                                name: tc.name.clone(),
+                                arguments: serde_json::to_string(&tc.arguments)
+                                    .unwrap_or_default(),
+                            },
+                        })
+                        .collect();
+
+                    messages.push(ChatMessage {
+                        role: "assistant".to_string(),
+                        content: if reasoning.is_empty() { None } else { Some(reasoning.clone()) },
+                        tool_calls: Some(tc_requests),
+                        tool_call_id: None,
+                    });
+
+                    let mut current_batch_node_ids: Vec<String> = Vec::new();
+
+                    for tc in &tool_calls {
+                        total_tool_calls += 1;
+                        let node_id = uuid::Uuid::new_v4().to_string();
+
+                        let mut node = AttackChainNode::new(
+                            session_id.clone(),
+                            node_id.clone(),
+                            tc.name.clone(),
+                            tc.arguments.clone(),
+                            reasoning.clone(),
+                        );
+                        // Link to previous iteration's nodes
+                        node.parent_node_ids = prev_node_ids.clone();
+                        node.status = AttackNodeStatus::Running;
+                        node.started_at = Some(chrono::Utc::now());
+                        let _ = self.db.attack_chain_nodes().insert_one(&node).await;
+                        current_batch_node_ids.push(node_id.clone());
+
+                        let _ = self.event_tx.send(PentestEvent::ToolStart {
+                            node_id: node_id.clone(),
+                            tool_name: tc.name.clone(),
+                            input: tc.arguments.clone(),
+                        });
+
+                        let result = if let Some(tool) = self.tool_registry.get(&tc.name) {
+                            match tool.execute(tc.arguments.clone(), &tool_context).await {
+                                Ok(result) => {
+                                    total_successes += 1;
+                                    let findings_count = result.findings.len() as u32;
+                                    total_findings += findings_count;
+
+                                    let mut finding_ids: Vec<String> = Vec::new();
+                                    for mut finding in result.findings {
+                                        finding.scan_run_id = session_id.clone();
+                                        finding.session_id = Some(session_id.clone());
+                                        let insert_result =
+                                            self.db.dast_findings().insert_one(&finding).await;
+                                        if let Ok(res) = &insert_result {
+                                            finding_ids.push(res.inserted_id.as_object_id().map(|oid| oid.to_hex()).unwrap_or_default());
+                                        }
+                                        let _ =
+                                            self.event_tx.send(PentestEvent::Finding {
+                                                finding_id: finding
+                                                    .id
+                                                    .map(|oid| oid.to_hex())
+                                                    .unwrap_or_default(),
+                                                title: finding.title.clone(),
+                                                severity: finding.severity.to_string(),
+                                            });
+                                    }
+
+                                    // Compute risk score based on findings severity
+                                    let risk_score: Option<u8> = if findings_count > 0 {
+                                        Some(std::cmp::min(
+                                            100,
+                                            (findings_count as u8).saturating_mul(15).saturating_add(20),
+                                        ))
+                                    } else {
+                                        None
+                                    };
+
+                                    let _ = self.event_tx.send(PentestEvent::ToolComplete {
+                                        node_id: node_id.clone(),
+                                        summary: result.summary.clone(),
+                                        findings_count,
+                                    });
+
+                                    let finding_ids_bson: Vec<mongodb::bson::Bson> = finding_ids
+                                        .iter()
+                                        .map(|id| mongodb::bson::Bson::String(id.clone()))
+                                        .collect();
+
+                                    let mut update_doc = doc! {
+                                        "status": "completed",
+                                        "tool_output": mongodb::bson::to_bson(&result.data)
+                                            .unwrap_or(mongodb::bson::Bson::Null),
+                                        "completed_at": mongodb::bson::DateTime::now(),
+                                        "findings_produced": finding_ids_bson,
+                                    };
+                                    if let Some(rs) = risk_score {
+                                        update_doc.insert("risk_score", rs as i32);
+                                    }
+
+                                    let _ = self
+                                        .db
+                                        .attack_chain_nodes()
+                                        .update_one(
+                                            doc! {
+                                                "session_id": &session_id,
+                                                "node_id": &node_id,
+                                            },
+                                            doc! { "$set": update_doc },
+                                        )
+                                        .await;
+
+                                    serde_json::json!({
+                                        "summary": result.summary,
+                                        "findings_count": findings_count,
+                                        "data": result.data,
+                                    })
+                                    .to_string()
+                                }
+                                Err(e) => {
+                                    let _ = self
+                                        .db
+                                        .attack_chain_nodes()
+                                        .update_one(
+                                            doc! {
+                                                "session_id": &session_id,
+                                                "node_id": &node_id,
+                                            },
+                                            doc! { "$set": {
+                                                "status": "failed",
+                                                "completed_at": mongodb::bson::DateTime::now(),
+                                            }},
+                                        )
+                                        .await;
+                                    format!("Tool execution failed: {e}")
+                                }
+                            }
+                        } else {
+                            format!("Unknown tool: {}", tc.name)
+                        };
+
+                        messages.push(ChatMessage {
+                            role: "tool".to_string(),
+                            content: Some(result),
+                            tool_calls: None,
+                            tool_call_id: Some(tc.id.clone()),
+                        });
+                    }
+
+                    // Advance parent links so next iteration's nodes connect to this batch
+                    prev_node_ids = current_batch_node_ids;
+
+                    if let Some(sid) = session.id {
+                        let _ = self
+                            .db
+                            .pentest_sessions()
+                            .update_one(
+                                doc! { "_id": sid },
+                                doc! { "$set": {
+                                    "tool_invocations": total_tool_calls as i64,
+                                    "tool_successes": total_successes as i64,
+                                    "findings_count": total_findings as i64,
+                                }},
+                            )
+                            .await;
+                    }
+                }
+            }
+        }
+
+        if let Some(sid) = session.id {
+            let _ = self
+                .db
+                .pentest_sessions()
+                .update_one(
+                    doc! { "_id": sid },
+                    doc! { "$set": {
+                        "status": "completed",
+                        "completed_at": mongodb::bson::DateTime::now(),
+                        "tool_invocations": total_tool_calls as i64,
+                        "tool_successes": total_successes as i64,
+                        "findings_count": total_findings as i64,
+                    }},
+                )
+                .await;
+        }
+
+        let _ = self.event_tx.send(PentestEvent::Complete {
+            summary: format!(
+                "Pentest complete. {} findings from {} tool invocations.",
+                total_findings, total_tool_calls
+            ),
+        });
+
+        Ok(())
+    }
+
+    // ── Code-Awareness: Gather context from linked repo ─────────
+
+    /// Fetch SAST findings, SBOM entries (with CVEs), and code graph entry points
+    /// for the repo linked to this DAST target.
+    async fn gather_repo_context(
+        &self,
+        target: &DastTarget,
+    ) -> (Vec<Finding>, Vec<SbomEntry>, Vec<CodeContextHint>) {
+        let Some(repo_id) = &target.repo_id else {
+            return (Vec::new(), Vec::new(), Vec::new());
+        };
+
+        let sast_findings = self.fetch_sast_findings(repo_id).await;
+        let sbom_entries = self.fetch_vulnerable_sbom(repo_id).await;
+        let code_context = self.fetch_code_context(repo_id, &sast_findings).await;
+
+        tracing::info!(
+            repo_id,
+            sast_findings = sast_findings.len(),
+            vulnerable_deps = sbom_entries.len(),
+            code_hints = code_context.len(),
+            "Gathered code-awareness context for pentest"
+        );
+
+        (sast_findings, sbom_entries, code_context)
+    }
+
+    /// Fetch open/triaged SAST findings for the repo (not false positives or resolved)
+    async fn fetch_sast_findings(&self, repo_id: &str) -> Vec<Finding> {
+        let cursor = self
+            .db
+            .findings()
+            .find(doc! {
+                "repo_id": repo_id,
+                "status": { "$in": ["open", "triaged"] },
+            })
+            .sort(doc! { "severity": -1 })
+            .limit(100)
+            .await;
+
+        match cursor {
+            Ok(mut c) => {
+                let mut results = Vec::new();
+                while let Some(Ok(f)) = c.next().await {
+                    results.push(f);
+                }
+                results
+            }
+            Err(e) => {
+                tracing::warn!("Failed to fetch SAST findings for pentest: {e}");
+                Vec::new()
+            }
+        }
+    }
+
+    /// Fetch SBOM entries that have known vulnerabilities
+    async fn fetch_vulnerable_sbom(&self, repo_id: &str) -> Vec<SbomEntry> {
+        let cursor = self
+            .db
+            .sbom_entries()
+            .find(doc! {
+                "repo_id": repo_id,
+                "known_vulnerabilities": { "$exists": true, "$ne": [] },
+            })
+            .limit(50)
+            .await;
+
+        match cursor {
+            Ok(mut c) => {
+                let mut results = Vec::new();
+                while let Some(Ok(e)) = c.next().await {
+                    results.push(e);
+                }
+                results
+            }
+            Err(e) => {
+                tracing::warn!("Failed to fetch vulnerable SBOM entries: {e}");
+                Vec::new()
+            }
+        }
+    }
+
+    /// Build CodeContextHint objects from the code knowledge graph.
+    /// Maps entry points to their source files and links SAST findings.
+    async fn fetch_code_context(
+        &self,
+        repo_id: &str,
+        sast_findings: &[Finding],
+    ) -> Vec<CodeContextHint> {
+        // Get entry point nodes from the code graph
+        let cursor = self
+            .db
+            .graph_nodes()
+            .find(doc! {
+                "repo_id": repo_id,
+                "is_entry_point": true,
+            })
+            .limit(50)
+            .await;
+
+        let nodes = match cursor {
+            Ok(mut c) => {
+                let mut results = Vec::new();
+                while let Some(Ok(n)) = c.next().await {
+                    results.push(n);
+                }
+                results
+            }
+            Err(_) => return Vec::new(),
+        };
+
+        // Build hints by matching graph nodes to SAST findings by file path
+        nodes
+            .into_iter()
+            .map(|node| {
+                // Find SAST findings in the same file
+                let linked_vulns: Vec<String> = sast_findings
+                    .iter()
+                    .filter(|f| {
+                        f.file_path.as_deref() == Some(&node.file_path)
+                    })
+                    .map(|f| {
+                        format!(
+                            "[{}] {}: {} (line {})",
+                            f.severity,
+                            f.scanner,
+                            f.title,
+                            f.line_number.unwrap_or(0)
+                        )
+                    })
+                    .collect();
+
+                CodeContextHint {
+                    endpoint_pattern: node.qualified_name.clone(),
+                    handler_function: node.name.clone(),
+                    file_path: node.file_path.clone(),
+                    code_snippet: String::new(), // Could fetch from embeddings
+                    known_vulnerabilities: linked_vulns,
+                }
+            })
+            .collect()
+    }
+
+    // ── System Prompt Builder ───────────────────────────────────
+
+    async fn build_system_prompt(
+        &self,
+        session: &PentestSession,
+        target: &DastTarget,
+        sast_findings: &[Finding],
+        sbom_entries: &[SbomEntry],
+        code_context: &[CodeContextHint],
+    ) -> String {
+        let tool_names = self.tool_registry.list_names().join(", ");
+        let strategy_guidance = match session.strategy {
+            PentestStrategy::Quick => {
+                "Focus on the most common and impactful vulnerabilities. Run a quick recon, then target the highest-risk areas."
+            }
+            PentestStrategy::Comprehensive => {
+                "Perform a thorough assessment covering all vulnerability types. Start with recon, then systematically test each attack surface."
+            }
+            PentestStrategy::Targeted => {
+                "Focus specifically on areas highlighted by SAST findings and known CVEs. Prioritize exploiting known weaknesses."
+            }
+            PentestStrategy::Aggressive => {
+                "Use all available tools aggressively. Test with maximum payloads and attempt full exploitation."
+            }
+            PentestStrategy::Stealth => {
+                "Minimize noise. Use fewer requests, avoid aggressive payloads. Focus on passive analysis and targeted probes."
+            }
+        };
+
+        // Build SAST findings section
+        let sast_section = if sast_findings.is_empty() {
+            String::from("No SAST findings available for this target.")
+        } else {
+            let critical = sast_findings
+                .iter()
+                .filter(|f| f.severity == Severity::Critical)
+                .count();
+            let high = sast_findings
+                .iter()
+                .filter(|f| f.severity == Severity::High)
+                .count();
+
+            let mut section = format!(
+                "{} open findings ({} critical, {} high):\n",
+                sast_findings.len(),
+                critical,
+                high
+            );
+
+            // List the most important findings (critical/high first, up to 20)
+            for f in sast_findings.iter().take(20) {
+                let file_info = f
+                    .file_path
+                    .as_ref()
+                    .map(|p| {
+                        format!(
+                            " in {}:{}",
+                            p,
+                            f.line_number.unwrap_or(0)
+                        )
+                    })
+                    .unwrap_or_default();
+                let status_note = match f.status {
+                    FindingStatus::Triaged => " [TRIAGED]",
+                    _ => "",
+                };
+                section.push_str(&format!(
+                    "- [{sev}] {title}{file}{status}\n",
+                    sev = f.severity,
+                    title = f.title,
+                    file = file_info,
+                    status = status_note,
+                ));
+                if let Some(cwe) = &f.cwe {
+                    section.push_str(&format!("  CWE: {cwe}\n"));
+                }
+            }
+            if sast_findings.len() > 20 {
+                section.push_str(&format!(
+                    "... and {} more findings\n",
+                    sast_findings.len() - 20
+                ));
+            }
+            section
+        };
+
+        // Build SBOM/CVE section
+        let sbom_section = if sbom_entries.is_empty() {
+            String::from("No vulnerable dependencies identified.")
+        } else {
+            let mut section = format!(
+                "{} dependencies with known vulnerabilities:\n",
+                sbom_entries.len()
+            );
+            for entry in sbom_entries.iter().take(15) {
+                let cve_ids: Vec<&str> = entry
+                    .known_vulnerabilities
+                    .iter()
+                    .map(|v| v.id.as_str())
+                    .collect();
+                section.push_str(&format!(
+                    "- {} {} ({}): {}\n",
+                    entry.name,
+                    entry.version,
+                    entry.package_manager,
+                    cve_ids.join(", ")
+                ));
+            }
+            if sbom_entries.len() > 15 {
+                section.push_str(&format!(
+                    "... and {} more vulnerable dependencies\n",
+                    sbom_entries.len() - 15
+                ));
+            }
+            section
+        };
+
+        // Build code context section
+        let code_section = if code_context.is_empty() {
+            String::from("No code knowledge graph available for this target.")
+        } else {
+            let with_vulns = code_context
+                .iter()
+                .filter(|c| !c.known_vulnerabilities.is_empty())
+                .count();
+
+            let mut section = format!(
+                "{} entry points identified ({} with linked SAST findings):\n",
+                code_context.len(),
+                with_vulns
+            );
+
+            for hint in code_context.iter().take(20) {
+                section.push_str(&format!(
+                    "- {} ({})\n",
+                    hint.endpoint_pattern, hint.file_path
+                ));
+                for vuln in &hint.known_vulnerabilities {
+                    section.push_str(&format!("  SAST: {vuln}\n"));
+                }
+            }
+            section
+        };
+
+        format!(
+            r#"You are an expert penetration tester conducting an authorized security assessment.
+
+## Target
+- **Name**: {target_name}
+- **URL**: {base_url}
+- **Type**: {target_type}
+- **Rate Limit**: {rate_limit} req/s
+- **Destructive Tests Allowed**: {allow_destructive}
+- **Linked Repository**: {repo_linked}
+
+## Strategy
+{strategy_guidance}
+
+## SAST Findings (Static Analysis)
+{sast_section}
+
+## Vulnerable Dependencies (SBOM)
+{sbom_section}
+
+## Code Entry Points (Knowledge Graph)
+{code_section}
+
+## Available Tools
+{tool_names}
+
+## Instructions
+1. Start by running reconnaissance (recon tool) to fingerprint the target and discover technologies.
+2. Run the OpenAPI parser to discover API endpoints from specs.
+3. Check infrastructure: DNS, DMARC, TLS, security headers, cookies, CSP, CORS.
+4. Based on SAST findings, prioritize testing endpoints where vulnerabilities were found in code.
+5. For each vulnerability type found in SAST, use the corresponding DAST tool to verify exploitability.
+6. If vulnerable dependencies are listed, try to trigger known CVE conditions against the running application.
+7. Test rate limiting on critical endpoints (login, API).
+8. Check for console.log leakage in frontend JavaScript.
+9. Analyze tool results and chain findings — if one vulnerability enables others, explore the chain.
+10. When testing is complete, provide a structured summary with severity and remediation.
+11. Always explain your reasoning before invoking each tool.
+12. When done, say "Testing complete" followed by a final summary.
+
+## Important
+- This is an authorized penetration test. All testing is permitted within the target scope.
+- Respect the rate limit of {rate_limit} requests per second.
+- Only use destructive tests if explicitly allowed ({allow_destructive}).
+- Use SAST findings to guide your testing — they tell you WHERE in the code vulnerabilities exist.
+- Use SBOM data to understand what technologies and versions the target runs.
+"#,
+            target_name = target.name,
+            base_url = target.base_url,
+            target_type = target.target_type,
+            rate_limit = target.rate_limit,
+            allow_destructive = target.allow_destructive,
+            repo_linked = target.repo_id.as_deref().unwrap_or("None"),
+        )
+    }
+}