gate4agent 0.4.5

//! Thread-safe registry of in-flight host → agent requests.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use tokio::sync::oneshot;
use serde_json::Value;

use super::message::{RpcError, RpcId};

/// Result type for a pending host → agent request.
pub type RpcResult = Result<Value, RpcError>;

/// Thread-safe registry of in-flight requests sent by the host to the agent.
///
/// When the host sends a request with ID N, it registers a oneshot sender here.
/// When the agent's response for ID N arrives, the reader loop looks it up and
/// delivers the result, waking the waiting `rpc_call` future.
///
/// `PendingRequests` is `Clone` (wraps `Arc`). One clone lives in the reader
/// loop; one stays in `RpcSession` for the `rpc_call` method.
#[derive(Clone, Default)]
pub struct PendingRequests {
    inner: Arc<Mutex<HashMap<RpcId, oneshot::Sender<RpcResult>>>>,
}

impl PendingRequests {
    /// Create a new empty registry.
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a new pending request. Returns the receiver to await.
    pub fn register(&self, id: RpcId) -> oneshot::Receiver<RpcResult> {
        let (tx, rx) = oneshot::channel();
        if let Ok(mut guard) = self.inner.lock() {
            guard.insert(id, tx);
        }
        rx
    }

    /// Deliver a response (called from reader loop).
    ///
    /// Returns `Ok(())` if a waiter was found and notified.
    /// Returns `Err(result)` if no waiter was registered for `id` (stale or
    /// unsolicited response).
    pub fn resolve(&self, id: RpcId, result: RpcResult) -> Result<(), RpcResult> {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(_) => return Err(result),
        };
        match guard.remove(&id) {
            Some(tx) => {
                let _ = tx.send(result);
                Ok(())
            }
            None => Err(result),
        }
    }

    /// Remove a request whose serialized write failed before the agent could
    /// ever produce a response.
    pub fn remove(&self, id: &RpcId) -> bool {
        self.inner
            .lock()
            .map(|mut guard| guard.remove(id).is_some())
            .unwrap_or(false)
    }

    /// Cancel all pending requests with the given error.
    ///
    /// Called on session shutdown to wake all waiting `rpc_call` futures.
    /// Takes a full `RpcError` (rather than synthesizing one internally) so
    /// callers that classified WHY the session closed -- e.g. the reader
    /// loop recognizing an authentication-required stderr signature -- can
    /// carry that classification through `RpcError::code`/`data` instead of
    /// it being flattened to a fixed internal-error code before it reaches
    /// the waiter.
    pub fn cancel_all(&self, error: RpcError) {
        let mut guard = match self.inner.lock() {
            Ok(g) => g,
            Err(_) => return,
        };
        for (_id, tx) in guard.drain() {
            let _ = tx.send(Err(error.clone()));
        }
    }

    /// Number of currently pending requests.
    pub fn len(&self) -> usize {
        self.inner.lock().map(|g| g.len()).unwrap_or(0)
    }

    /// Returns true if there are no pending requests.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

// ---------------------------------------------------------------------------
// Unit tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[tokio::test]
    async fn register_and_resolve_success() {
        let pending = PendingRequests::new();
        let rx = pending.register(RpcId::Number(1));
        pending.resolve(RpcId::Number(1), Ok(json!({"ok": true}))).unwrap();
        let result = rx.await.unwrap();
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn resolve_unknown_id_returns_err() {
        let pending = PendingRequests::new();
        let result = pending.resolve(RpcId::Number(999), Ok(json!(null)));
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn cancel_all_wakes_waiters() {
        let pending = PendingRequests::new();
        let rx1 = pending.register(RpcId::Number(1));
        let rx2 = pending.register(RpcId::Number(2));
        pending.cancel_all(RpcError::internal("session closed"));
        let r1 = rx1.await.unwrap();
        let r2 = rx2.await.unwrap();
        assert!(r1.is_err());
        assert!(r2.is_err());
        assert_eq!(r1.unwrap_err().code, RpcError::INTERNAL_ERROR);
    }

    #[tokio::test]
    async fn two_concurrent_waiters_resolve_independently() {
        let pending = PendingRequests::new();
        let rx1 = pending.register(RpcId::Number(1));
        let rx2 = pending.register(RpcId::Number(2));

        pending
            .resolve(RpcId::Number(2), Ok(json!({"which": "two"})))
            .unwrap();
        pending
            .resolve(RpcId::Number(1), Ok(json!({"which": "one"})))
            .unwrap();

        let r1 = rx1.await.unwrap().unwrap();
        let r2 = rx2.await.unwrap().unwrap();
        assert_eq!(r1["which"], "one");
        assert_eq!(r2["which"], "two");
    }

    #[test]
    fn failed_write_can_remove_pending_request() {
        let pending = PendingRequests::new();
        let id = RpcId::Number(7);
        let _receiver = pending.register(id.clone());
        assert!(pending.remove(&id));
        assert!(pending.is_empty());
    }
}