agent-runtime-infra-adapter 0.1.1

use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::sync::Arc;

use async_trait::async_trait;
use runtime_ports::{
    InteractionRequest, InteractionResponse, InteractionSession, OperationContext, PortFailure,
    PortFailureKind,
};
use runtime_types::{ExecutionId, OperationId};
use tokio::sync::{Mutex, oneshot};

struct Pending {
    request_id: String,
    sender: Option<oneshot::Sender<InteractionResponse>>,
    receiver: Option<oneshot::Receiver<InteractionResponse>>,
}
struct State {
    pending: BTreeMap<ExecutionId, Pending>,
    completed: BTreeSet<OperationId>,
    order: VecDeque<OperationId>,
}

pub struct InMemoryInteractionBroker {
    state: Mutex<State>,
    max_completed: usize,
}
impl InMemoryInteractionBroker {
    pub fn new(max_completed: usize) -> Self {
        Self {
            state: Mutex::new(State {
                pending: BTreeMap::new(),
                completed: BTreeSet::new(),
                order: VecDeque::new(),
            }),
            max_completed,
        }
    }

    pub async fn submit(
        &self,
        execution_id: &ExecutionId,
        operation_id: &OperationId,
        request_id: &str,
        response: InteractionResponse,
    ) -> Result<(), PortFailure> {
        let mut state = self.state.lock().await;
        if state.completed.contains(operation_id) {
            return Ok(());
        }
        let pending = state.pending.get_mut(execution_id).ok_or_else(|| {
            PortFailure::new(
                PortFailureKind::Conflict,
                "INTERACTION_NOT_PENDING",
                "execution has no pending interaction",
            )
        })?;
        if pending.request_id != request_id {
            return Err(PortFailure::new(
                PortFailureKind::Conflict,
                "INTERACTION_REQUEST_MISMATCH",
                "request ID does not match the pending interaction",
            ));
        }
        let sender = pending.sender.take().ok_or_else(|| {
            PortFailure::new(
                PortFailureKind::Conflict,
                "INTERACTION_ALREADY_ANSWERED",
                "pending interaction was already answered",
            )
        })?;
        sender.send(response).map_err(|_| {
            PortFailure::new(
                PortFailureKind::Conflict,
                "INTERACTION_CLOSED",
                "pending interaction is no longer active",
            )
        })?;
        if pending.receiver.is_none() {
            state.pending.remove(execution_id);
        }
        state.completed.insert(operation_id.clone());
        state.order.push_back(operation_id.clone());
        while state.order.len() > self.max_completed {
            if let Some(old) = state.order.pop_front() {
                state.completed.remove(&old);
            }
        }
        Ok(())
    }
}

pub struct PendingInteractionSession {
    broker: Arc<InMemoryInteractionBroker>,
}
impl PendingInteractionSession {
    pub fn new(broker: Arc<InMemoryInteractionBroker>) -> Self {
        Self { broker }
    }
}
#[async_trait]
impl InteractionSession for PendingInteractionSession {
    async fn prepare(
        &self,
        operation: &OperationContext,
        request: InteractionRequest,
    ) -> Result<(), PortFailure> {
        let (sender, receiver) = oneshot::channel();
        let mut state = self.broker.state.lock().await;
        if state.pending.contains_key(&operation.execution_id) {
            return Err(PortFailure::new(
                PortFailureKind::Conflict,
                "INTERACTION_ALREADY_PENDING",
                "execution already has a pending interaction",
            ));
        }
        state.pending.insert(
            operation.execution_id.clone(),
            Pending {
                request_id: request.request_id,
                sender: Some(sender),
                receiver: Some(receiver),
            },
        );
        Ok(())
    }

    async fn wait(
        &self,
        operation: &OperationContext,
        request_id: &str,
    ) -> Result<InteractionResponse, PortFailure> {
        let receiver = {
            let mut state = self.broker.state.lock().await;
            let pending = state
                .pending
                .get_mut(&operation.execution_id)
                .ok_or_else(|| {
                    PortFailure::new(
                        PortFailureKind::Conflict,
                        "INTERACTION_NOT_PREPARED",
                        "interaction was not prepared",
                    )
                })?;
            if pending.request_id != request_id {
                return Err(PortFailure::new(
                    PortFailureKind::Conflict,
                    "INTERACTION_REQUEST_MISMATCH",
                    "request ID does not match the pending interaction",
                ));
            }
            let receiver = pending.receiver.take().ok_or_else(|| {
                PortFailure::new(
                    PortFailureKind::Conflict,
                    "INTERACTION_ALREADY_WAITING",
                    "interaction already has a waiter",
                )
            })?;
            if pending.sender.is_none() {
                state.pending.remove(&operation.execution_id);
            }
            receiver
        };
        let remaining = operation.remaining()?;
        let result = tokio::select! {
            _ = operation.cancellation.cancelled() => Err(PortFailure::canceled()),
            value = tokio::time::timeout(remaining, receiver) => match value {
                Ok(Ok(response)) => Ok(response),
                Ok(Err(_)) => Err(PortFailure::new(PortFailureKind::Unavailable, "INTERACTION_CLOSED", "interaction response channel closed")),
                Err(_) => Err(PortFailure::timeout()),
            }
        };
        self.broker
            .state
            .lock()
            .await
            .pending
            .remove(&operation.execution_id);
        result
    }
}

#[cfg(test)]
mod tests {
    use std::time::{Duration, Instant};

    use runtime_types::OperationId;
    use tokio_util::sync::CancellationToken;

    use super::*;

    #[tokio::test]
    async fn input_can_commit_after_prepare_before_wait_starts() {
        let broker = Arc::new(InMemoryInteractionBroker::new(8));
        let session = PendingInteractionSession::new(broker.clone());
        let operation = OperationContext {
            id: OperationId::random(),
            execution_id: ExecutionId::random(),
            deadline: Instant::now() + Duration::from_secs(2),
            cancellation: CancellationToken::new(),
        };
        session
            .prepare(
                &operation,
                InteractionRequest {
                    request_id: "call-1".into(),
                    prompt: "approve".into(),
                },
            )
            .await
            .unwrap();
        broker
            .submit(
                &operation.execution_id,
                &OperationId::random(),
                "call-1",
                InteractionResponse { text: "yes".into() },
            )
            .await
            .unwrap();
        assert_eq!(
            session.wait(&operation, "call-1").await.unwrap().text,
            "yes"
        );
    }
}