aion-server 0.30.0

//! One live harness process, and the recorder every frame goes through.
//!
//! # Commit before broadcast, always
//!
//! [`Recorder::record`] appends the frame to the durable transcript FIRST and
//! broadcasts it SECOND. Nothing a client saw is therefore absent from the
//! record, and the socket's `?after=` replay is a read of the same rows the
//! socket streamed. The order costs a durable commit per frame — at
//! token-chunk rate that is a commit per chunk — and that cost is deliberate
//! and measured rather than traded away for a coalescing window nobody chose.
//!
//! An append that FAILS is not broadcast. The frame the client would have seen
//! does not exist, so sending it would make the live stream and the record
//! disagree in exactly the direction that cannot be repaired later. The failure
//! is a typed error to the caller and, where a turn is open, that turn fails.

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use aion_core::{
    AssistantSessionEvent, AssistantSessionFrame, AssistantSessionId, ContentType, Payload,
};
use aion_integration_acp::AcpSession;
use aion_store::AssistantSessionStore;
use chrono::Utc;
use tokio::process::{ChildStdin, ChildStdout};
use tokio::sync::{Mutex, broadcast};

use super::error::AssistantSessionError;

/// The concrete session type an assistant harness produces: a spawned agent's
/// piped stdio.
pub(crate) type HarnessSession = AcpSession<ChildStdout, ChildStdin>;

/// Appends a frame and then publishes it, in that order, for one session.
///
/// Held by the registry AND by each turn's driver task, so both write through
/// one path: a second recorder would be a second chance to broadcast something
/// that was never stored.
#[derive(Clone)]
pub(crate) struct Recorder {
    session_id: AssistantSessionId,
    store: Arc<dyn AssistantSessionStore>,
    events: broadcast::Sender<AssistantSessionFrame>,
}

impl Recorder {
    /// Build a recorder for one session.
    pub(crate) fn new(
        session_id: AssistantSessionId,
        store: Arc<dyn AssistantSessionStore>,
        events: broadcast::Sender<AssistantSessionFrame>,
    ) -> Self {
        Self {
            session_id,
            store,
            events,
        }
    }

    /// The session this recorder writes for.
    pub(crate) const fn session_id(&self) -> AssistantSessionId {
        self.session_id
    }

    /// Append one event to the durable transcript, then publish it live.
    ///
    /// Returns the index the store assigned, which is what a client passes back
    /// as `?after=`.
    ///
    /// # Errors
    ///
    /// [`AssistantSessionError::Store`] when the append fails. NOTHING is
    /// broadcast in that case: a frame a client saw and the store never took
    /// would make the live stream and the record disagree, and no later read
    /// could repair it.
    pub(crate) async fn record(
        &self,
        event: AssistantSessionEvent,
    ) -> Result<u64, AssistantSessionError> {
        let bytes = serde_json::to_vec(&event).map_err(|error| {
            AssistantSessionError::Internal(format!(
                "an assistant session frame is not encodable: {error}"
            ))
        })?;
        let index = self
            .store
            .append_assistant_transcript_event(
                &self.session_id,
                Utc::now(),
                Payload::new(ContentType::Json, bytes),
            )
            .await?;
        // Between the append and the broadcast, and never after it: the listing
        // cache is written from the frame that is already durable, so a reader
        // that saw the frame live cannot find a record that has not caught up.
        self.cache_commands(&event).await?;
        self.cache_config_options(&event).await?;
        // A send with no subscribers is not a failure: nobody is watching, and
        // the frame is already durable. The receiver count is the only thing
        // `send` reports, and it is deliberately discarded.
        let _ = self.events.send(AssistantSessionFrame { index, event });
        Ok(index)
    }

    /// Keep the record's command cache in step with the transcript.
    ///
    /// The commands a listing shows come from the record because an
    /// advertisement can arrive at the first turn of a conversation that runs
    /// for hours, and a listing must not read every session's whole transcript
    /// to find it. The TRANSCRIPT is still the authority: this writes what the
    /// frame just committed said, in the same append path, so nothing else in
    /// the server can put a different list there.
    ///
    /// A record that vanished between the append and here is not invented: the
    /// frame is durable and the projection can rebuild the list, so the absence
    /// is logged and the append stands.
    async fn cache_commands(
        &self,
        event: &AssistantSessionEvent,
    ) -> Result<(), AssistantSessionError> {
        let AssistantSessionEvent::AvailableCommands { commands } = event else {
            return Ok(());
        };
        let Some(mut record) = self.store.get_assistant_session(&self.session_id).await? else {
            tracing::warn!(
                session = %self.session_id,
                "an assistant session advertised commands after its record disappeared; the \
                 transcript holds them and the listing cache does not"
            );
            return Ok(());
        };
        record.commands.clone_from(commands);
        record.updated_at = Utc::now();
        self.store.put_assistant_session(record).await?;
        Ok(())
    }

    /// Keep the record's configuration-option cache in step with the
    /// transcript — the same contract as [`Self::cache_commands`], for the
    /// same listing-cost reason, with the same authority (the transcript).
    async fn cache_config_options(
        &self,
        event: &AssistantSessionEvent,
    ) -> Result<(), AssistantSessionError> {
        let AssistantSessionEvent::ConfigOptions { options } = event else {
            return Ok(());
        };
        let Some(mut record) = self.store.get_assistant_session(&self.session_id).await? else {
            tracing::warn!(
                session = %self.session_id,
                "an assistant session advertised configuration options after its record \
                 disappeared; the transcript holds them and the listing cache does not"
            );
            return Ok(());
        };
        record.config_options.clone_from(options);
        record.updated_at = Utc::now();
        self.store.put_assistant_session(record).await?;
        Ok(())
    }

    /// A subscription to the live frames from this moment on.
    pub(crate) fn subscribe(&self) -> broadcast::Receiver<AssistantSessionFrame> {
        self.events.subscribe()
    }
}

/// A harness process this server is holding open for one session.
///
/// The ACP session is behind a mutex because ending it CONSUMES it: closing an
/// agent means dropping the last reference to its connection, which cannot be
/// expressed against a shared borrow. Every other operation takes `&self` and so
/// only holds the lock long enough to write a frame.
pub(crate) struct LiveSession {
    session_id: AssistantSessionId,
    session: Mutex<Option<HarnessSession>>,
    recorder: Recorder,
    /// Whether a turn is open. The adapter refuses a second concurrent prompt
    /// itself; this is what lets the HTTP surface answer `409` BEFORE spawning
    /// any work, with a message about sessions rather than about ACP.
    busy: AtomicBool,
    /// The harness's own session handle — what `session/load` resumes.
    acp_session_ref: String,
    /// Whether this agent advertised `loadSession` at `initialize`.
    load_session: bool,

    /// The raw `configOptions` advertisement the opening response carried, if
    /// any — taken exactly once by the open path to record.
    initial_config_options: std::sync::Mutex<Option<serde_json::Value>>,
}

impl LiveSession {
    /// Adopt a started harness session.
    pub(crate) fn new(
        session_id: AssistantSessionId,
        session: HarnessSession,
        recorder: Recorder,
    ) -> Self {
        let acp_session_ref = session.session_id().to_string();
        let load_session = session.supports_load_session();
        let initial_config_options = session.initial_config_options().cloned();
        Self {
            session_id,
            session: Mutex::new(Some(session)),
            recorder,
            busy: AtomicBool::new(false),
            acp_session_ref,
            load_session,
            initial_config_options: std::sync::Mutex::new(initial_config_options),
        }
    }

    /// The opening advertisement, exactly once: the first call takes it, every
    /// later call finds `None`. Taken rather than borrowed so the open path
    /// that records it cannot be followed by a second recorder of the same
    /// frame.
    pub(crate) fn take_initial_config_options(&self) -> Option<serde_json::Value> {
        match self.initial_config_options.lock() {
            Ok(mut held) => held.take(),
            Err(poisoned) => poisoned.into_inner().take(),
        }
    }

    /// The session this process serves.
    pub(crate) const fn session_id(&self) -> AssistantSessionId {
        self.session_id
    }

    /// The harness's own session handle.
    pub(crate) fn acp_session_ref(&self) -> &str {
        &self.acp_session_ref
    }

    /// Whether this agent advertised `loadSession`.
    pub(crate) const fn supports_load_session(&self) -> bool {
        self.load_session
    }

    /// The recorder every frame for this session goes through.
    pub(crate) const fn recorder(&self) -> &Recorder {
        &self.recorder
    }

    /// Claim the session's single turn slot, or refuse because one is open.
    ///
    /// A compare-and-exchange rather than a check-then-set: two turns arriving
    /// at once must not both see an idle session. The claim is released by
    /// [`Self::release_turn`] — on every failure path of the request that made
    /// it, and by the turn driver when the turn ends.
    ///
    /// # Errors
    ///
    /// [`AssistantSessionError::Busy`] when a turn is already open.
    pub(crate) fn claim_turn(&self) -> Result<(), AssistantSessionError> {
        if self
            .busy
            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
            .is_err()
        {
            return Err(AssistantSessionError::Busy {
                session_id: self.session_id,
            });
        }
        Ok(())
    }

    /// Release the turn slot.
    ///
    /// Idempotent: releasing an unclaimed slot is a no-op, which is what lets
    /// the request path release on a failure and the driver release at the end
    /// without either needing to know whether the other already did.
    pub(crate) fn release_turn(&self) {
        self.busy.store(false, Ordering::SeqCst);
    }

    /// Borrow the harness session for one operation.
    ///
    /// `None` once the session has been closed: a caller that finds it gone
    /// reports a stale target rather than waiting on a process that is not
    /// there.
    pub(crate) async fn with_session<T>(
        &self,
        operation: impl AsyncFnOnce(&HarnessSession) -> T,
    ) -> Option<T> {
        let guard = self.session.lock().await;
        let session = guard.as_ref()?;
        Some(operation(session).await)
    }

    /// Whether the harness's reader is still running.
    ///
    /// `false` once the agent's stdout has ended — the honest answer to "can
    /// this session still take a turn", read from the adapter rather than
    /// assumed from the absence of an error.
    pub(crate) async fn is_alive(&self) -> bool {
        self.session
            .lock()
            .await
            .as_ref()
            .is_some_and(AcpSession::is_live)
    }

    /// Shut the agent down, giving it the configured grace to exit before it is
    /// killed.
    ///
    /// Idempotent: a second close finds the session already taken and does
    /// nothing, which is what lets a delete race a crash without either path
    /// having to know about the other.
    pub(crate) async fn close(&self) {
        let taken = self.session.lock().await.take();
        if let Some(session) = taken {
            session.close().await;
        }
    }
}