PicoBot/src/agent/gate.rs

use std::sync::{Arc, Weak};

use dashmap::DashMap;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use tokio_util::sync::CancellationToken;

/// Step permits are held while a provider request or tool invocation is in
/// flight and released as soon as the step finishes.  Session permits are
/// released before global permits (reverse acquisition order).  The fields
/// are never read; ownership alone keeps the quotas reserved (RAII).
#[derive(Debug)]
pub struct StepPermit {
    #[allow(dead_code)]
    session: Option<OwnedSemaphorePermit>,
    #[allow(dead_code)]
    global: OwnedSemaphorePermit,
}

/// Run quota permits cover a whole Agent Run from admission until its
/// terminal commit.  Foreground delegation does not take run permits; the
/// quota gates background admission so a waiting parent can never deadlock
/// nested foreground children.
#[derive(Debug)]
pub struct RunPermit {
    #[allow(dead_code)]
    session: Option<OwnedSemaphorePermit>,
    #[allow(dead_code)]
    global: OwnedSemaphorePermit,
}

#[derive(Debug, thiserror::Error)]
pub enum GateError {
    #[error("execution gate wait cancelled")]
    Cancelled,
}

#[derive(Default)]
struct KeyedSemaphores {
    permits: usize,
    map: DashMap<String, Weak<Semaphore>>,
}

impl KeyedSemaphores {
    fn new(permits: usize) -> Self {
        Self {
            permits,
            map: DashMap::new(),
        }
    }

    fn semaphore(self: &Arc<Self>, key: &str) -> Arc<Semaphore> {
        if let Some(existing) = self.map.get(key)
            && let Some(semaphore) = existing.upgrade()
        {
            return semaphore;
        }
        let semaphore = Arc::new(Semaphore::new(self.permits));
        self.map.insert(key.to_string(), Arc::downgrade(&semaphore));
        // Drop registry entries whose owners are gone so long-lived gateways
        // do not accumulate one semaphore per historical session.
        self.map.retain(|_, weak| weak.strong_count() > 0);
        semaphore
    }
}

/// Concurrency gates shared by one runtime generation.  Run quota and
/// provider/tool step permits have independent lifecycles; acquisition order
/// is always global -> session and release order is reversed.
pub struct ExecutionGate {
    max_concurrent_runs: usize,
    run_global: Arc<Semaphore>,
    run_session: Arc<KeyedSemaphores>,
    provider_global: Arc<Semaphore>,
    provider_session: Arc<KeyedSemaphores>,
    tool_global: Arc<Semaphore>,
    tool_session: Arc<KeyedSemaphores>,
}

impl std::fmt::Debug for ExecutionGate {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ExecutionGate")
            .field("run_available", &self.run_global.available_permits())
            .field(
                "provider_available",
                &self.provider_global.available_permits(),
            )
            .field("tool_available", &self.tool_global.available_permits())
            .finish()
    }
}

impl ExecutionGate {
    pub fn new(config: &crate::config::AgentOrchestrationConfig) -> Arc<Self> {
        Arc::new(Self {
            max_concurrent_runs: config.max_concurrent_runs,
            run_global: Arc::new(Semaphore::new(config.max_concurrent_runs)),
            run_session: Arc::new(KeyedSemaphores::new(config.max_concurrent_runs_per_session)),
            provider_global: Arc::new(Semaphore::new(config.max_concurrent_provider_steps)),
            provider_session: Arc::new(KeyedSemaphores::new(
                config.max_concurrent_provider_steps_per_session,
            )),
            tool_global: Arc::new(Semaphore::new(config.max_concurrent_tool_steps)),
            tool_session: Arc::new(KeyedSemaphores::new(
                config.max_concurrent_tool_steps_per_session,
            )),
        })
    }

    /// Unlimited gate used when orchestration is disabled; root Turns still
    /// route through it so the code path stays uniform.
    pub fn unbounded() -> Arc<Self> {
        Arc::new(Self {
            max_concurrent_runs: usize::MAX,
            run_global: Arc::new(Semaphore::new(Semaphore::MAX_PERMITS)),
            run_session: Arc::new(KeyedSemaphores::new(Semaphore::MAX_PERMITS)),
            provider_global: Arc::new(Semaphore::new(Semaphore::MAX_PERMITS)),
            provider_session: Arc::new(KeyedSemaphores::new(Semaphore::MAX_PERMITS)),
            tool_global: Arc::new(Semaphore::new(Semaphore::MAX_PERMITS)),
            tool_session: Arc::new(KeyedSemaphores::new(Semaphore::MAX_PERMITS)),
        })
    }

    /// Global run quota ceiling; background batches larger than this are
    /// rejected up front instead of queueing indefinitely.
    pub fn max_concurrent_runs(&self) -> usize {
        self.max_concurrent_runs
    }

    pub async fn acquire_run(
        self: &Arc<Self>,
        session_id: &str,
        cancellation: &CancellationToken,
    ) -> Result<RunPermit, GateError> {
        let global = acquire_owned(self.run_global.clone(), cancellation).await?;
        let session = acquire_keyed(&self.run_session, session_id, cancellation).await?;
        Ok(RunPermit {
            session: Some(session),
            global,
        })
    }

    pub async fn acquire_provider(
        self: &Arc<Self>,
        session_id: Option<&str>,
        cancellation: &CancellationToken,
    ) -> Result<StepPermit, GateError> {
        let global = acquire_owned(self.provider_global.clone(), cancellation).await?;
        let session = match session_id {
            Some(session_id) => {
                Some(acquire_keyed(&self.provider_session, session_id, cancellation).await?)
            }
            None => None,
        };
        Ok(StepPermit { session, global })
    }

    pub async fn acquire_tool(
        self: &Arc<Self>,
        session_id: Option<&str>,
        cancellation: &CancellationToken,
    ) -> Result<StepPermit, GateError> {
        let global = acquire_owned(self.tool_global.clone(), cancellation).await?;
        let session = match session_id {
            Some(session_id) => {
                Some(acquire_keyed(&self.tool_session, session_id, cancellation).await?)
            }
            None => None,
        };
        Ok(StepPermit { session, global })
    }

    pub fn available_provider_permits(&self) -> usize {
        self.provider_global.available_permits()
    }

    pub fn available_tool_permits(&self) -> usize {
        self.tool_global.available_permits()
    }
}

async fn acquire_owned(
    semaphore: Arc<Semaphore>,
    cancellation: &CancellationToken,
) -> Result<OwnedSemaphorePermit, GateError> {
    tokio::select! {
        biased;
        _ = cancellation.cancelled() => Err(GateError::Cancelled),
        result = semaphore.acquire_owned() => {
            result.map_err(|_| GateError::Cancelled)
        }
    }
}

async fn acquire_keyed(
    keyed: &Arc<KeyedSemaphores>,
    key: &str,
    cancellation: &CancellationToken,
) -> Result<OwnedSemaphorePermit, GateError> {
    acquire_owned(keyed.semaphore(key), cancellation).await
}

/// Test helper: wait until a predicate holds or fail after the timeout.
#[cfg(test)]
async fn eventually<F: Fn() -> bool>(predicate: F) {
    let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
    while std::time::Instant::now() < deadline {
        if predicate() {
            return;
        }
        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
    }
    panic!("condition did not hold within timeout");
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    use crate::config::AgentOrchestrationConfig;

    /// Global caps stay larger than session caps so a session-level waiter
    /// can never exhaust the global permits while blocked (acquisition order
    /// is global -> session by design).
    fn gate(provider_session: usize, tool_session: usize) -> Arc<ExecutionGate> {
        let config = AgentOrchestrationConfig {
            max_concurrent_runs: 8,
            max_concurrent_runs_per_session: 1,
            max_concurrent_provider_steps: 8,
            max_concurrent_provider_steps_per_session: provider_session,
            max_concurrent_tool_steps: 8,
            max_concurrent_tool_steps_per_session: tool_session,
            ..Default::default()
        };
        ExecutionGate::new(&config)
    }

    #[tokio::test]
    async fn provider_permits_are_released_after_drop() {
        let gate = gate(1, 1);
        let token = CancellationToken::new();
        let permit = gate.acquire_provider(None, &token).await.unwrap();
        assert_eq!(gate.available_provider_permits(), 7);
        drop(permit);
        assert_eq!(gate.available_provider_permits(), 8);
    }

    #[tokio::test]
    async fn session_quota_blocks_second_session_step_until_release() {
        let gate = gate(1, 8);
        let token = CancellationToken::new();
        let first = gate
            .acquire_provider(Some("cli:a:d1"), &token)
            .await
            .unwrap();
        // Global capacity remains, but the same session is capped at one.
        let blocked = tokio::spawn({
            let gate = gate.clone();
            let token = token.clone();
            async move { gate.acquire_provider(Some("cli:a:d1"), &token).await }
        });
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert!(!blocked.is_finished());
        // A different session is unaffected.
        let other = gate
            .acquire_provider(Some("cli:b:d2"), &token)
            .await
            .unwrap();
        drop(other);
        drop(first);
        let second = blocked.await.unwrap().unwrap();
        drop(second);
        assert_eq!(gate.available_provider_permits(), 8);
    }

    #[tokio::test]
    async fn cancellation_aborts_permit_wait() {
        let gate = gate(8, 1);
        let token = CancellationToken::new();
        let held = gate.acquire_tool(Some("cli:a:d1"), &token).await.unwrap();

        let waiter_token = CancellationToken::new();
        let waiter = tokio::spawn({
            let gate = gate.clone();
            let waiter_token = waiter_token.clone();
            async move { gate.acquire_tool(Some("cli:a:d1"), &waiter_token).await }
        });
        tokio::time::sleep(Duration::from_millis(50)).await;
        waiter_token.cancel();
        let error = waiter.await.unwrap().unwrap_err();
        assert!(matches!(error, GateError::Cancelled));
        drop(held);
        assert_eq!(gate.available_tool_permits(), 8);
    }

    #[tokio::test]
    async fn run_quota_is_scoped_per_session() {
        let gate = gate(8, 8);
        let token = CancellationToken::new();
        let held = gate.acquire_run("cli:a:d1", &token).await.unwrap();

        let blocked = tokio::spawn({
            let gate = gate.clone();
            let token = token.clone();
            async move { gate.acquire_run("cli:a:d1", &token).await }
        });
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert!(!blocked.is_finished());

        let other = gate.acquire_run("cli:b:d2", &token).await.unwrap();
        drop(other);
        drop(held);
        blocked.await.unwrap().unwrap();
    }

    #[tokio::test]
    async fn keyed_semaphores_are_reclaimed_when_idle() {
        let keyed = Arc::new(KeyedSemaphores::new(1));
        let semaphore = keyed.semaphore("cli:a:d1");
        assert_eq!(keyed.map.len(), 1);
        drop(semaphore);
        keyed.semaphore("cli:b:d2");
        assert_eq!(keyed.map.len(), 1, "idle session entry must be reclaimed");
    }

    #[tokio::test]
    async fn waiting_parent_holds_no_step_permits() {
        // With a session cap of 1, a parent that merely waits for its child
        // must leave the session's single provider permit free; otherwise
        // nested foreground delegation would deadlock.
        let gate = gate(1, 1);
        let token = CancellationToken::new();
        let parent_step = gate
            .acquire_provider(Some("cli:a:d1"), &token)
            .await
            .unwrap();
        drop(parent_step);
        // Parent is now in the waiting_children state: no permits held.
        eventually(|| gate.available_provider_permits() == 8).await;
        let child = gate
            .acquire_provider(Some("cli:a:d1"), &token)
            .await
            .unwrap();
        drop(child);
        assert_eq!(gate.available_provider_permits(), 8);
    }
}