compliance-scanner-agent/compliance-agent/src/agent.rs

use std::sync::Arc;

use dashmap::DashMap;
use tokio::sync::{broadcast, watch, Semaphore};

use compliance_core::models::pentest::PentestEvent;
use compliance_core::AgentConfig;

use crate::database::DatabasePool;
use crate::llm::LlmClient;
use crate::pipeline::orchestrator::PipelineOrchestrator;

/// Default maximum concurrent pentest sessions.
const DEFAULT_MAX_CONCURRENT_SESSIONS: usize = 5;

#[derive(Clone)]
pub struct ComplianceAgent {
    pub config: AgentConfig,
    /// Per-tenant Mongo broker. Every code path must obtain a
    /// tenant-scoped [`crate::database::Database`] from this pool —
    /// there is no single shared database any more.
    pub db_pool: DatabasePool,
    pub llm: Arc<LlmClient>,
    pub http: reqwest::Client,
    /// Per-session broadcast senders for SSE streaming.
    pub session_streams: Arc<DashMap<String, broadcast::Sender<PentestEvent>>>,
    /// Per-session pause controls (true = paused).
    pub session_pause: Arc<DashMap<String, watch::Sender<bool>>>,
    /// Semaphore limiting concurrent pentest sessions.
    pub session_semaphore: Arc<Semaphore>,
}

impl ComplianceAgent {
    pub fn new(config: AgentConfig, db_pool: DatabasePool) -> Self {
        let llm = Arc::new(LlmClient::new(
            config.litellm_url.clone(),
            config.litellm_api_key.clone(),
            config.litellm_model.clone(),
            config.litellm_embed_model.clone(),
        ));
        let http = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(30))
            .connect_timeout(std::time::Duration::from_secs(10))
            .build()
            .unwrap_or_default();
        Self {
            config,
            db_pool,
            llm,
            http,
            session_streams: Arc::new(DashMap::new()),
            session_pause: Arc::new(DashMap::new()),
            session_semaphore: Arc::new(Semaphore::new(DEFAULT_MAX_CONCURRENT_SESSIONS)),
        }
    }

    pub async fn run_scan(
        &self,
        tenant_id: &str,
        repo_id: &str,
        trigger: compliance_core::models::ScanTrigger,
    ) -> Result<(), crate::error::AgentError> {
        let db = self.db_pool.for_tenant_id(tenant_id).await?;
        let orchestrator =
            PipelineOrchestrator::new(self.config.clone(), db, self.llm.clone(), self.http.clone());
        orchestrator.run_target(repo_id, trigger).await
    }

    /// Alias for [`Self::run_scan`] — every scan runs the unified onboarded-target
    /// pipeline. Kept as a distinct name for the `/targets/{id}/scan` endpoint's
    /// intent.
    pub async fn run_target_scan(
        &self,
        tenant_id: &str,
        target_id: &str,
        trigger: compliance_core::models::ScanTrigger,
    ) -> Result<(), crate::error::AgentError> {
        let db = self.db_pool.for_tenant_id(tenant_id).await?;
        let orchestrator =
            PipelineOrchestrator::new(self.config.clone(), db, self.llm.clone(), self.http.clone());
        orchestrator.run_target(target_id, trigger).await
    }

    /// Run a PR review: scan the diff and post review comments.
    pub async fn run_pr_review(
        &self,
        tenant_id: &str,
        repo_id: &str,
        pr_number: u64,
        base_sha: &str,
        head_sha: &str,
    ) -> Result<(), crate::error::AgentError> {
        let db = self.db_pool.for_tenant_id(tenant_id).await?;
        let oid = mongodb::bson::oid::ObjectId::parse_str(repo_id)
            .map_err(|e| crate::error::AgentError::Other(e.to_string()))?;
        let target = db
            .onboarded_targets()
            .find_one(mongodb::bson::doc! { "_id": oid })
            .await?
            .ok_or_else(|| {
                crate::error::AgentError::Other(format!("Target {repo_id} not found"))
            })?;
        let code = target.code_artifact().ok_or_else(|| {
            crate::error::AgentError::Other(format!("Target {repo_id} has no code artifact"))
        })?;
        let repo = crate::pipeline::repo_view::RepoView::from_target(&target, code);

        let orchestrator =
            PipelineOrchestrator::new(self.config.clone(), db, self.llm.clone(), self.http.clone());
        orchestrator
            .run_pr_review(&repo, repo_id, pr_number, base_sha, head_sha)
            .await
    }

    // ── Session stream management ──────────────────────────────────

    /// Register a broadcast sender for a session. Returns the sender.
    pub fn register_session_stream(&self, session_id: &str) -> broadcast::Sender<PentestEvent> {
        let (tx, _) = broadcast::channel(256);
        self.session_streams
            .insert(session_id.to_string(), tx.clone());
        tx
    }

    /// Subscribe to a session's broadcast stream.
    pub fn subscribe_session(&self, session_id: &str) -> Option<broadcast::Receiver<PentestEvent>> {
        self.session_streams
            .get(session_id)
            .map(|tx| tx.subscribe())
    }

    // ── Session pause/resume management ────────────────────────────

    /// Register a pause control for a session. Returns the watch receiver.
    pub fn register_pause_control(&self, session_id: &str) -> watch::Receiver<bool> {
        let (tx, rx) = watch::channel(false);
        self.session_pause.insert(session_id.to_string(), tx);
        rx
    }

    /// Pause a session.
    pub fn pause_session(&self, session_id: &str) -> bool {
        if let Some(tx) = self.session_pause.get(session_id) {
            tx.send(true).is_ok()
        } else {
            false
        }
    }

    /// Resume a session.
    pub fn resume_session(&self, session_id: &str) -> bool {
        if let Some(tx) = self.session_pause.get(session_id) {
            tx.send(false).is_ok()
        } else {
            false
        }
    }

    /// Clean up all per-session resources.
    pub fn cleanup_session(&self, session_id: &str) {
        self.session_streams.remove(session_id);
        self.session_pause.remove(session_id);
    }
}