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

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

@@ -7,11 +7,13 @@ use mongodb::bson::doc;
 
 use compliance_core::models::chat::{ChatRequest, ChatResponse, SourceReference};
 use compliance_core::models::embedding::EmbeddingBuildRun;
+use compliance_core::tenant_ctx::TenantCtx;
 use compliance_graph::graph::embedding_store::EmbeddingStore;
 
 use crate::agent::ComplianceAgent;
 use crate::rag::pipeline::RagPipeline;
 
+use super::dto::tenant_db;
 use super::ApiResponse;
 
 type AgentExt = Extension<Arc<ComplianceAgent>>;
@@ -20,10 +22,12 @@ type AgentExt = Extension<Arc<ComplianceAgent>>;
 #[tracing::instrument(skip_all, fields(repo_id = %repo_id))]
 pub async fn chat(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     Path(repo_id): Path<String>,
     Json(req): Json<ChatRequest>,
 ) -> Result<Json<ApiResponse<ChatResponse>>, StatusCode> {
-    let pipeline = RagPipeline::new(agent.llm.clone(), agent.db.inner());
+    let db = tenant_db(&agent, &tenant).await?;
+    let pipeline = RagPipeline::new(agent.llm.clone(), db.inner());
 
     // Step 1: Embed the user's message
     let query_vectors = agent
@@ -133,12 +137,15 @@ pub async fn chat(
 #[tracing::instrument(skip_all, fields(repo_id = %repo_id))]
 pub async fn build_embeddings(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     Path(repo_id): Path<String>,
 ) -> Result<Json<serde_json::Value>, StatusCode> {
+    // Resolve the tenant DB up front so we can move it into the spawn;
+    // the JWT/dev context isn't available inside detached tasks.
+    let db = tenant_db(&agent, &tenant).await?;
     let agent_clone = (*agent).clone();
     tokio::spawn(async move {
-        let repo = match agent_clone
-            .db
+        let repo = match db
             .repositories()
             .find_one(doc! { "_id": mongodb::bson::oid::ObjectId::parse_str(&repo_id).ok() })
             .await
@@ -151,8 +158,7 @@ pub async fn build_embeddings(
         };
 
         // Get latest graph build
-        let build = match agent_clone
-            .db
+        let build = match db
             .graph_builds()
             .find_one(doc! { "repo_id": &repo_id })
             .sort(doc! { "started_at": -1 })
@@ -171,26 +177,22 @@ pub async fn build_embeddings(
             .unwrap_or_else(|| "unknown".to_string());
 
         // Get nodes
-        let nodes: Vec<compliance_core::models::graph::CodeNode> = match agent_clone
-            .db
-            .graph_nodes()
-            .find(doc! { "repo_id": &repo_id })
-            .await
-        {
-            Ok(cursor) => {
-                use futures_util::StreamExt;
-                let mut items = Vec::new();
-                let mut cursor = cursor;
-                while let Some(Ok(item)) = cursor.next().await {
-                    items.push(item);
+        let nodes: Vec<compliance_core::models::graph::CodeNode> =
+            match db.graph_nodes().find(doc! { "repo_id": &repo_id }).await {
+                Ok(cursor) => {
+                    use futures_util::StreamExt;
+                    let mut items = Vec::new();
+                    let mut cursor = cursor;
+                    while let Some(Ok(item)) = cursor.next().await {
+                        items.push(item);
+                    }
+                    items
                 }
-                items
-            }
-            Err(e) => {
-                tracing::error!("[{repo_id}] Failed to fetch nodes: {e}");
-                return;
-            }
-        };
+                Err(e) => {
+                    tracing::error!("[{repo_id}] Failed to fetch nodes: {e}");
+                    return;
+                }
+            };
 
         let creds = crate::pipeline::git::RepoCredentials {
             ssh_key_path: Some(agent_clone.config.ssh_key_path.clone()),
@@ -207,7 +209,7 @@ pub async fn build_embeddings(
             }
         };
 
-        let pipeline = RagPipeline::new(agent_clone.llm.clone(), agent_clone.db.inner());
+        let pipeline = RagPipeline::new(agent_clone.llm.clone(), db.inner());
         match pipeline
             .build_embeddings(&repo_id, &repo_path, &graph_build_id, &nodes)
             .await
@@ -234,9 +236,11 @@ pub async fn build_embeddings(
 #[tracing::instrument(skip_all, fields(repo_id = %repo_id))]
 pub async fn embedding_status(
     Extension(agent): AgentExt,
+    tenant: TenantCtx,
     Path(repo_id): Path<String>,
 ) -> Result<Json<ApiResponse<Option<EmbeddingBuildRun>>>, StatusCode> {
-    let store = EmbeddingStore::new(agent.db.inner());
+    let db = tenant_db(&agent, &tenant).await?;
+    let store = EmbeddingStore::new(db.inner());
     let build = store.get_latest_build(&repo_id).await.map_err(|e| {
         tracing::error!("Failed to get embedding status: {e}");
         StatusCode::INTERNAL_SERVER_ERROR