fix(dashboard): attach Keycloak token on agent API calls (#90)

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

@@ -13,10 +13,11 @@ use compliance_core::models::dast::DastFinding;
 use compliance_core::models::finding::Finding;
 use compliance_core::models::pentest::*;
 use compliance_core::models::sbom::SbomEntry;
+use compliance_core::tenant_ctx::TenantCtx;
 
 use crate::agent::ComplianceAgent;
 
-use super::super::dto::collect_cursor_async;
+use super::super::dto::{collect_cursor_async, tenant_db};
 
 type AgentExt = Extension<Arc<ComplianceAgent>>;
 
@@ -35,11 +36,15 @@ pub struct ExportBody {
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn export_session_report(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     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()))?;
+    let db = tenant_db(&agent, &tenant)
+        .await
+        .map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
 
     if body.password.len() < 8 {
         return Err((
@@ -49,8 +54,7 @@ pub async fn export_session_report(
     }
 
     // Fetch session
-    let session = agent
-        .db
+    let session = db
         .pentest_sessions()
         .find_one(doc! { "_id": oid })
         .await
@@ -64,9 +68,7 @@ pub async fn export_session_report(
 
     // Resolve target name
     let target = if let Ok(tid) = mongodb::bson::oid::ObjectId::parse_str(&session.target_id) {
-        agent
-            .db
-            .dast_targets()
+        db.dast_targets()
             .find_one(doc! { "_id": tid })
             .await
             .ok()
@@ -84,8 +86,7 @@ pub async fn export_session_report(
         .unwrap_or_default();
 
     // Fetch attack chain nodes
-    let nodes: Vec<AttackChainNode> = match agent
-        .db
+    let nodes: Vec<AttackChainNode> = match db
         .attack_chain_nodes()
         .find(doc! { "session_id": &id })
         .sort(doc! { "started_at": 1 })
@@ -96,8 +97,7 @@ pub async fn export_session_report(
     };
 
     // Fetch DAST findings for this session, then deduplicate
-    let raw_findings: Vec<DastFinding> = match agent
-        .db
+    let raw_findings: Vec<DastFinding> = match db
         .dast_findings()
         .find(doc! { "session_id": &id })
         .sort(doc! { "severity": -1, "created_at": -1 })
@@ -122,8 +122,7 @@ pub async fn export_session_report(
         .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
+        let sast: Vec<Finding> = match db
             .findings()
             .find(doc! {
                 "repo_id": rid,
@@ -143,8 +142,7 @@ pub async fn export_session_report(
             Err(_) => Vec::new(),
         };
 
-        let sbom: Vec<SbomEntry> = match agent
-            .db
+        let sbom: Vec<SbomEntry> = match db
             .sbom_entries()
             .find(doc! {
                 "repo_id": rid,
@@ -164,8 +162,7 @@ pub async fn export_session_report(
         };
 
         // Build code context from graph nodes
-        let code_ctx: Vec<CodeContextHint> = match agent
-            .db
+        let code_ctx: Vec<CodeContextHint> = match db
             .graph_nodes()
             .find(doc! { "repo_id": rid, "is_entry_point": true })
             .limit(50)

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

@@ -7,11 +7,12 @@ use mongodb::bson::doc;
 use serde::Deserialize;
 
 use compliance_core::models::pentest::*;
+use compliance_core::tenant_ctx::TenantCtx;
 
 use crate::agent::ComplianceAgent;
 use crate::pentest::PentestOrchestrator;
 
-use super::super::dto::{collect_cursor_async, ApiResponse, PaginationParams};
+use super::super::dto::{collect_cursor_async, tenant_db, ApiResponse, PaginationParams};
 
 type AgentExt = Extension<Arc<ComplianceAgent>>;
 
@@ -43,6 +44,7 @@ pub struct LookupRepoQuery {
 #[tracing::instrument(skip_all)]
 pub async fn create_session(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     Json(req): Json<CreateSessionRequest>,
 ) -> Result<Json<ApiResponse<PentestSession>>, (StatusCode, String)> {
     // Try to acquire a concurrency permit
@@ -57,6 +59,10 @@ pub async fn create_session(
             )
         })?;
 
+    let db = tenant_db(&agent, &tenant)
+        .await
+        .map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
+
     if let Some(ref config) = req.config {
         // ── Wizard path ──────────────────────────────────────────────
         if !config.disclaimer_accepted {
@@ -67,8 +73,7 @@ pub async fn create_session(
         }
 
         // Look up or auto-create DastTarget by app_url
-        let target = match agent
-            .db
+        let target = match db
             .dast_targets()
             .find_one(doc! { "base_url": &config.app_url })
             .await
@@ -87,7 +92,7 @@ pub async fn create_session(
                 }
                 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| {
+                let res = db.dast_targets().insert_one(&t).await.map_err(|e| {
                     (
                         StatusCode::INTERNAL_SERVER_ERROR,
                         format!("Failed to create target: {e}"),
@@ -110,8 +115,7 @@ pub async fn create_session(
 
         // 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
+            if let Ok(Some(repo)) = db
                 .repositories()
                 .find_one(doc! { "git_url": git_url })
                 .await
@@ -120,8 +124,7 @@ pub async fn create_session(
             }
         }
 
-        let insert_result = agent
-            .db
+        let insert_result = db
             .pentest_sessions()
             .insert_one(&session)
             .await
@@ -212,8 +215,7 @@ pub async fn create_session(
         // Persist encrypted credentials to DB
         if session_for_task.config.is_some() {
             if let Some(sid) = session.id {
-                let _ = agent
-                    .db
+                let _ = db
                     .pentest_sessions()
                     .update_one(
                         doc! { "_id": sid },
@@ -245,12 +247,13 @@ pub async fn create_session(
             });
 
         let llm = agent.llm.clone();
-        let db = agent.db.clone();
+        let db_for_orchestrator = 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));
+            let orchestrator =
+                PentestOrchestrator::new(llm, db_for_orchestrator, event_tx, Some(pause_rx));
             orchestrator
                 .run_session_guarded(&session_clone, &target_clone, &initial_message)
                 .await;
@@ -292,8 +295,7 @@ pub async fn create_session(
             )
         })?;
 
-        let target = agent
-            .db
+        let target = db
             .dast_targets()
             .find_one(doc! { "_id": oid })
             .await
@@ -310,8 +312,7 @@ pub async fn create_session(
         let mut session = PentestSession::new(target_id, strategy);
         session.repo_id = target.repo_id.clone();
 
-        let insert_result = agent
-            .db
+        let insert_result = db
             .pentest_sessions()
             .insert_one(&session)
             .await
@@ -338,12 +339,13 @@ pub async fn create_session(
         });
 
         let llm = agent.llm.clone();
-        let db = agent.db.clone();
+        let db_for_orchestrator = 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));
+            let orchestrator =
+                PentestOrchestrator::new(llm, db_for_orchestrator, event_tx, Some(pause_rx));
             orchestrator
                 .run_session_guarded(&session_clone, &target_clone, &initial_message)
                 .await;
@@ -373,10 +375,11 @@ fn parse_strategy(s: &str) -> PentestStrategy {
 #[tracing::instrument(skip_all)]
 pub async fn lookup_repo(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     Query(params): Query<LookupRepoQuery>,
 ) -> Result<Json<ApiResponse<serde_json::Value>>, StatusCode> {
-    let repo = agent
-        .db
+    let db = tenant_db(&agent, &tenant).await?;
+    let repo = db
         .repositories()
         .find_one(doc! { "git_url": &params.url })
         .await
@@ -402,9 +405,11 @@ pub async fn lookup_repo(
 #[tracing::instrument(skip_all)]
 pub async fn list_sessions(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     Query(params): Query<PaginationParams>,
 ) -> Result<Json<ApiResponse<Vec<PentestSession>>>, StatusCode> {
-    let db = &agent.db;
+    let db = tenant_db(&agent, &tenant).await?;
+    let db = &db;
     let skip = (params.page.saturating_sub(1)) * params.limit as u64;
     let total = db
         .pentest_sessions()
@@ -438,12 +443,13 @@ pub async fn list_sessions(
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn get_session(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     Path(id): Path<String>,
 ) -> Result<Json<ApiResponse<PentestSession>>, StatusCode> {
     let oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
+    let db = tenant_db(&agent, &tenant).await?;
 
-    let mut session = agent
-        .db
+    let mut session = db
         .pentest_sessions()
         .find_one(doc! { "_id": oid })
         .await
@@ -471,15 +477,18 @@ pub async fn get_session(
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn send_message(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     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()))?;
+    let db = tenant_db(&agent, &tenant)
+        .await
+        .map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
 
     // Verify session exists and is running
-    let session = agent
-        .db
+    let session = db
         .pentest_sessions()
         .find_one(doc! { "_id": oid })
         .await
@@ -506,8 +515,7 @@ pub async fn send_message(
         )
     })?;
 
-    let target = agent
-        .db
+    let target = db
         .dast_targets()
         .find_one(doc! { "_id": target_oid })
         .await
@@ -527,13 +535,13 @@ pub async fn send_message(
     // 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 _ = 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 db_for_orchestrator = db.clone();
     let message = req.message.clone();
 
     // Use existing broadcast sender if available, otherwise create a new one
@@ -548,7 +556,7 @@ pub async fn send_message(
         .unwrap_or_else(|| agent.register_session_stream(&session_id));
 
     tokio::spawn(async move {
-        let orchestrator = PentestOrchestrator::new(llm, db, event_tx, None);
+        let orchestrator = PentestOrchestrator::new(llm, db_for_orchestrator, event_tx, None);
         orchestrator
             .run_session_guarded(&session, &target, &message)
             .await;
@@ -565,13 +573,16 @@ pub async fn send_message(
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn stop_session(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     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 db = tenant_db(&agent, &tenant)
+        .await
+        .map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
 
-    let session = agent
-        .db
+    let session = db
         .pentest_sessions()
         .find_one(doc! { "_id": oid })
         .await
@@ -590,9 +601,7 @@ pub async fn stop_session(
         ));
     }
 
-    agent
-        .db
-        .pentest_sessions()
+    db.pentest_sessions()
         .update_one(
             doc! { "_id": oid },
             doc! { "$set": {
@@ -612,8 +621,7 @@ pub async fn stop_session(
     // Clean up session resources
     agent.cleanup_session(&id);
 
-    let updated = agent
-        .db
+    let updated = db
         .pentest_sessions()
         .find_one(doc! { "_id": oid })
         .await
@@ -641,13 +649,16 @@ pub async fn stop_session(
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn pause_session(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     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 db = tenant_db(&agent, &tenant)
+        .await
+        .map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
 
-    let session = agent
-        .db
+    let session = db
         .pentest_sessions()
         .find_one(doc! { "_id": oid })
         .await
@@ -684,13 +695,16 @@ pub async fn pause_session(
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn resume_session(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     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 db = tenant_db(&agent, &tenant)
+        .await
+        .map_err(|s| (s, "failed to acquire tenant database".to_string()))?;
 
-    let session = agent
-        .db
+    let session = db
         .pentest_sessions()
         .find_one(doc! { "_id": oid })
         .await
@@ -727,12 +741,13 @@ pub async fn resume_session(
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn get_attack_chain(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     Path(id): Path<String>,
 ) -> Result<Json<ApiResponse<Vec<AttackChainNode>>>, StatusCode> {
     let _oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
+    let db = tenant_db(&agent, &tenant).await?;
 
-    let nodes = match agent
-        .db
+    let nodes = match db
         .attack_chain_nodes()
         .find(doc! { "session_id": &id })
         .sort(doc! { "started_at": 1 })
@@ -757,21 +772,21 @@ pub async fn get_attack_chain(
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn get_messages(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     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 db = tenant_db(&agent, &tenant).await?;
 
     let skip = (params.page.saturating_sub(1)) * params.limit as u64;
-    let total = agent
-        .db
+    let total = db
         .pentest_messages()
         .count_documents(doc! { "session_id": &id })
         .await
         .unwrap_or(0);
 
-    let messages = match agent
-        .db
+    let messages = match db
         .pentest_messages()
         .find(doc! { "session_id": &id })
         .sort(doc! { "created_at": 1 })
@@ -797,21 +812,21 @@ pub async fn get_messages(
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn get_session_findings(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     Path(id): Path<String>,
     Query(params): Query<PaginationParams>,
 ) -> Result<Json<ApiResponse<Vec<compliance_core::models::dast::DastFinding>>>, StatusCode> {
     let _oid = mongodb::bson::oid::ObjectId::parse_str(&id).map_err(|_| StatusCode::BAD_REQUEST)?;
+    let db = tenant_db(&agent, &tenant).await?;
 
     let skip = (params.page.saturating_sub(1)) * params.limit as u64;
-    let total = agent
-        .db
+    let total = db
         .dast_findings()
         .count_documents(doc! { "session_id": &id })
         .await
         .unwrap_or(0);
 
-    let findings = match agent
-        .db
+    let findings = match db
         .dast_findings()
         .find(doc! { "session_id": &id })
         .sort(doc! { "created_at": -1 })

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

@@ -6,10 +6,11 @@ use axum::Json;
 use mongodb::bson::doc;
 
 use compliance_core::models::pentest::*;
+use compliance_core::tenant_ctx::TenantCtx;
 
 use crate::agent::ComplianceAgent;
 
-use super::super::dto::{collect_cursor_async, ApiResponse};
+use super::super::dto::{collect_cursor_async, tenant_db, ApiResponse};
 
 type AgentExt = Extension<Arc<ComplianceAgent>>;
 
@@ -17,8 +18,10 @@ type AgentExt = Extension<Arc<ComplianceAgent>>;
 #[tracing::instrument(skip_all)]
 pub async fn pentest_stats(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
 ) -> Result<Json<ApiResponse<PentestStats>>, StatusCode> {
-    let db = &agent.db;
+    let db = tenant_db(&agent, &tenant).await?;
+    let db = &db;
 
     let running_sessions = db
         .pentest_sessions()

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

@@ -11,10 +11,11 @@ use tokio_stream::wrappers::BroadcastStream;
 use tokio_stream::StreamExt;
 
 use compliance_core::models::pentest::*;
+use compliance_core::tenant_ctx::TenantCtx;
 
 use crate::agent::ComplianceAgent;
 
-use super::super::dto::collect_cursor_async;
+use super::super::dto::{collect_cursor_async, tenant_db};
 
 type AgentExt = Extension<Arc<ComplianceAgent>>;
 
@@ -25,13 +26,14 @@ type AgentExt = Extension<Arc<ComplianceAgent>>;
 #[tracing::instrument(skip_all, fields(session_id = %id))]
 pub async fn session_stream(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     Path(id): Path<String>,
 ) -> 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)?;
+    let db = tenant_db(&agent, &tenant).await?;
 
     // Verify session exists
-    let _session = agent
-        .db
+    let _session = db
         .pentest_sessions()
         .find_one(doc! { "_id": oid })
         .await
@@ -43,8 +45,7 @@ pub async fn session_stream(
     let mut initial_events: Vec<Result<Event, Infallible>> = Vec::new();
 
     // Fetch recent messages for this session
-    let messages: Vec<PentestMessage> = match agent
-        .db
+    let messages: Vec<PentestMessage> = match db
         .pentest_messages()
         .find(doc! { "session_id": &id })
         .sort(doc! { "created_at": 1 })
@@ -56,8 +57,7 @@ pub async fn session_stream(
     };
 
     // Fetch recent attack chain nodes
-    let nodes: Vec<AttackChainNode> = match agent
-        .db
+    let nodes: Vec<AttackChainNode> = match db
         .attack_chain_nodes()
         .find(doc! { "session_id": &id })
         .sort(doc! { "started_at": 1 })
@@ -94,8 +94,7 @@ pub async fn session_stream(
     }
 
     // Add current session status event
-    let session = agent
-        .db
+    let session = db
         .pentest_sessions()
         .find_one(doc! { "_id": oid })
         .await