aion-server 0.31.0

//! Recording the agent's structured edits to a session's shared document.
//!
//! The `assistant_document_edit` tool ends here: one method that validates a
//! batch of replace-exactly-once operations against the shared document as the
//! transcript now holds it, and — only when the WHOLE batch applies — records
//! one [`AssistantSessionEvent::DocumentEdit`] through the session's recorder.
//! The recorder is the single append path, so the same act persists the batch
//! and streams it to the console, where the editor applies it and the operator
//! keeps or reverts it.
//!
//! # Validate-then-append is one critical section
//!
//! Two things make the revision number trustworthy, and both live in this file.
//! First, a per-session async lock is held across the projection read, the
//! validation, and the append — an agent that issues two edit calls in
//! parallel gets them serialized here, so two batches can never both be
//! assigned the same revision. Second, the revision is computed FROM the
//! projection (`document_revision + 1`), never from a counter beside it, so it
//! cannot drift from the transcript it numbers.
//!
//! # The one race that remains, and why it is honest
//!
//! A `ContextShared` append (a turn, or an explicit context push) is not under
//! this lock, so a batch can be validated against a document the very next
//! record rebases. The projection stays deterministic — its fold applies a
//! batch through the same [`apply_document_edits`] and leaves the text
//! unchanged when it no longer matches — and the console's twin applier makes
//! the same call visibly: a batch that no longer matches the operator's own
//! buffer is refused there, the operator's bytes win, and the operator is told
//! the agent's edits did not land. Serializing context shares behind this lock
//! would put an agent's tool call in front of the operator's own screen, which
//! is the wrong party to make wait.

use std::sync::Arc;

use aion_core::{
    AssistantDocumentEditError, AssistantDocumentEditOp, AssistantSessionEvent, AssistantSessionId,
    apply_document_edits,
};
use tokio::sync::Mutex;

use super::error::AssistantSessionError;
use super::registry::AssistantSessions;

/// What a recorded batch reports back to the tool that submitted it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct DocumentEditReceipt {
    /// The shared document's revision after this batch.
    pub(crate) revision: u64,
    /// The document the batch applied to, by the path the console shared.
    pub(crate) path: String,
    /// How many operations the batch carried.
    pub(crate) applied: usize,
}

/// Why a batch was not recorded. Nothing was appended in any of these cases.
#[derive(Debug)]
pub(crate) enum DocumentEditRefusal {
    /// The session's shared context names no document — the operator is not in
    /// an editor, or has shared nothing yet.
    NoDocument,
    /// The batch does not apply to the document as it stands.
    Edits(AssistantDocumentEditError),
    /// The transcript could not be read or written.
    Session(AssistantSessionError),
}

impl AssistantSessions {
    /// Validate `edits` against the shared document and record them as one
    /// [`AssistantSessionEvent::DocumentEdit`].
    ///
    /// # Errors
    ///
    /// [`DocumentEditRefusal`] naming what stopped the batch; on error nothing
    /// was appended and the document is untouched.
    pub(crate) async fn record_document_edit(
        &self,
        session_id: AssistantSessionId,
        edits: Vec<AssistantDocumentEditOp>,
    ) -> Result<DocumentEditReceipt, DocumentEditRefusal> {
        let lock = self.document_edit_lock(session_id);
        let guard = lock.lock().await;
        let projection = self
            .projection(session_id)
            .await
            .map_err(DocumentEditRefusal::Session)?;
        let Some(document) = projection
            .latest_context
            .as_ref()
            .and_then(|context| context.document.as_ref())
        else {
            return Err(DocumentEditRefusal::NoDocument);
        };
        apply_document_edits(&document.text, &edits).map_err(DocumentEditRefusal::Edits)?;
        let revision = projection.document_revision.saturating_add(1);
        let receipt = DocumentEditReceipt {
            revision,
            path: document.path.clone(),
            applied: edits.len(),
        };
        self.recorder(session_id)
            .record(AssistantSessionEvent::DocumentEdit { edits, revision })
            .await
            .map_err(DocumentEditRefusal::Session)?;
        drop(guard);
        Ok(receipt)
    }

    /// The per-session lock that serializes validate-then-append, minted on
    /// first use exactly as the live-frame channel is.
    fn document_edit_lock(&self, session_id: AssistantSessionId) -> Arc<Mutex<()>> {
        self.document_edit_locks()
            .entry(session_id)
            .or_insert_with(|| Arc::new(Mutex::new(())))
            .value()
            .clone()
    }
}