aion-server 0.29.0

//! `GET /assistant/sessions/{id}/events` — one session's frames, live.
//!
//! # Replay and live are one stream
//!
//! The frames appended to a session's durable transcript are exactly the frames
//! broadcast to a watcher, so "read the transcript" and "watch the socket" are
//! the same stream seen at two times. A client passes back the index it last
//! saw as `?after=`, and gets everything after it — the backlog first, then the
//! live tail, with no seam it has to reconcile.
//!
//! # A hole is never silent
//!
//! Two failures would otherwise leave a client with a transcript it believes is
//! complete and is not, so neither is swallowed:
//!
//! - a REJECTION before streaming (no grant, a malformed id, a session that is
//!   not this caller's) is one terminal `{"error": …}` frame plus a close,
//!   never a dropped socket, exactly as `/events/stream` does it;
//! - a LAGGED receiver is one terminal error frame naming how many frames were
//!   missed, plus a close. A client that quietly lost frames would render a
//!   conversation with a gap in it and have no way to know.

use aion_core::{AssistantSessionFrame, AssistantSessionId};
use aion_proto::WireError;
use axum::{
    extract::{
        Path, Query, State,
        ws::{CloseFrame, Message, WebSocket, WebSocketUpgrade, close_code},
    },
    response::{IntoResponse, Response},
};
use futures::{SinkExt, StreamExt};
use serde::Deserialize;
use tokio::sync::broadcast::error::RecvError;

use super::assistant_sessions::{invalid_id, session_refusal};
use super::auth::WsCaller;
use crate::namespace::grants::{ASSISTANT_SESSIONS, require_grant};
use crate::stream::socket::send_wire_error;
use crate::{CallerIdentity, ServerError, ServerState};

/// The socket's one query parameter.
///
/// Read as TEXT and parsed in the handler rather than as a `u64` by the
/// extractor: an extractor rejection is a failed handshake with a body no
/// WebSocket client surfaces, while a parse here is the terminal error frame
/// every other rejection on this route arrives as.
#[derive(Debug, Default, Deserialize)]
#[serde(default)]
pub(crate) struct AssistantWatchQuery {
    /// The last transcript index the client already has. Absent replays the
    /// whole conversation.
    after: Option<String>,
}

/// `GET /assistant/sessions/{id}/events`.
pub(crate) async fn assistant_session_socket(
    websocket: WebSocketUpgrade,
    State(state): State<ServerState>,
    WsCaller(caller): WsCaller,
    Path(id): Path<String>,
    Query(query): Query<AssistantWatchQuery>,
) -> Response {
    websocket
        .on_upgrade(move |socket| async move {
            if let Err(error) = serve_assistant_socket(socket, state, caller, id, query).await {
                tracing::warn!(
                    error = %error,
                    "an assistant session subscription ended with an error"
                );
            }
        })
        .into_response()
}

/// Authorize the watch, replay the backlog, then forward the live tail.
async fn serve_assistant_socket(
    socket: WebSocket,
    state: ServerState,
    caller: CallerIdentity,
    id: String,
    query: AssistantWatchQuery,
) -> Result<(), ServerError> {
    let (mut socket_tx, mut socket_rx) = socket.split();
    // Every pre-stream rejection takes the same shape: one terminal error
    // frame, then close. A `?` here would drop the socket silently.
    if let Err(error) = require_grant(&caller, &ASSISTANT_SESSIONS) {
        return terminal(&mut socket_tx, error.to_wire_error()).await;
    }
    let session = match AssistantSessionId::parse(&id) {
        Ok(session) => session,
        Err(error) => return terminal(&mut socket_tx, invalid_id(&error.to_string())).await,
    };
    let after = match parse_after(query.after.as_deref()) {
        Ok(after) => after,
        Err(wire) => return terminal(&mut socket_tx, wire).await,
    };
    let watch = state
        .assistant_sessions()
        .watch(caller.subject(), session, after)
        .await;
    let (replayed, live) = match watch {
        Ok(watch) => watch,
        Err(error) => {
            // The SAME mapping the request routes use, so a session that is not
            // this caller's is the same plain not-found on both transports.
            let (_status, wire) = session_refusal(&error);
            return terminal(&mut socket_tx, wire).await;
        }
    };

    let mut last_index = after;
    for frame in replayed {
        last_index = Some(frame.index);
        if !write_frame(&mut socket_tx, &frame).await? {
            // The client is gone mid-backlog; there is no one left to tell.
            return Ok(());
        }
    }
    forward_live(&mut socket_tx, &mut socket_rx, live, last_index).await
}

/// Forward frames as they are recorded, until the client leaves, the session's
/// broadcaster closes, or this watcher falls behind.
async fn forward_live<Tx, Rx>(
    socket_tx: &mut Tx,
    socket_rx: &mut Rx,
    mut live: tokio::sync::broadcast::Receiver<AssistantSessionFrame>,
    mut last_index: Option<u64>,
) -> Result<(), ServerError>
where
    Tx: futures::Sink<Message> + Unpin,
    <Tx as futures::Sink<Message>>::Error: std::fmt::Debug,
    Rx: futures::Stream<Item = Result<Message, axum::Error>> + Unpin,
{
    loop {
        tokio::select! {
            client_message = socket_rx.next() => {
                // Read the client's half only to notice it leaving: a watcher
                // sends nothing, and a socket nobody polls would hold a
                // broadcast receiver open long after its reader went away.
                match client_message {
                    Some(Ok(Message::Close(_))) | None => return Ok(()),
                    Some(Ok(message)) => drop(message),
                    Some(Err(error)) => {
                        drop(error);
                        return Ok(());
                    }
                }
            }
            received = live.recv() => {
                match received {
                    Ok(frame) => {
                        // The backlog read and the subscription overlap by
                        // design (the subscription is taken first, so nothing
                        // can fall between them), which means a frame can
                        // arrive on both. The index is dense and assigned by
                        // the store, so the one a client already has is the one
                        // to drop.
                        if last_index.is_some_and(|seen| frame.index <= seen) {
                            continue;
                        }
                        last_index = Some(frame.index);
                        if !write_frame(socket_tx, &frame).await? {
                            return Ok(());
                        }
                    }
                    Err(RecvError::Closed) => return finish(socket_tx).await,
                    Err(RecvError::Lagged(missed)) => {
                        return terminal(socket_tx, lagged(missed, last_index)).await;
                    }
                }
            }
        }
    }
}

/// Write one frame as a JSON text message.
///
/// `Ok(false)` means the client is gone — a clean end, not a failure. An
/// unencodable frame is NOT one: it is reported as a terminal error rather than
/// skipped, because a skipped frame is the hole this route exists to prevent.
async fn write_frame<Tx>(
    socket_tx: &mut Tx,
    frame: &AssistantSessionFrame,
) -> Result<bool, ServerError>
where
    Tx: futures::Sink<Message> + Unpin,
    <Tx as futures::Sink<Message>>::Error: std::fmt::Debug,
{
    let payload = match serde_json::to_string(frame) {
        Ok(payload) => payload,
        Err(source) => {
            let wire = WireError::backend(format!(
                "the assistant session frame at index {} could not be encoded for the stream \
                 ({source}); the stream is ended rather than continued with a frame missing from \
                 it",
                frame.index
            ));
            // `terminal` never returns `Ok`, so the `false` is the type's
            // requirement rather than a reachable answer.
            return terminal(socket_tx, wire).await.map(|()| false);
        }
    };
    Ok(socket_tx.send(Message::Text(payload.into())).await.is_ok())
}

/// Parse `?after=`, refusing text that is not a transcript index.
fn parse_after(after: Option<&str>) -> Result<Option<u64>, WireError> {
    match after {
        None => Ok(None),
        Some(text) => text.parse::<u64>().map(Some).map_err(|source| {
            WireError::invalid_input(format!(
                "`after` is the transcript index a client last saw, and `{text}` is not one \
                 ({source}); omit it to replay the whole conversation"
            ))
            .with_error_type(INVALID_AFTER_TYPE)
        }),
    }
}

/// The refusal a lagged watcher gets, naming what it lost.
fn lagged(missed: u64, last_index: Option<u64>) -> WireError {
    let resume = match last_index {
        Some(index) => format!("reconnect with `?after={index}` to read what was missed"),
        None => "reconnect with no `after` to read the conversation from its start".to_owned(),
    };
    WireError::lagged(format!(
        "this assistant session produced frames faster than the socket drained them and {missed} \
         frame(s) were dropped; the stream is ended rather than continued with a gap in it — \
         {resume}"
    ))
    .with_error_type(LAGGED_TYPE)
}

/// Send the terminal error frame and close, then surface the failure typed so
/// the upgrade task logs it.
///
/// Every end this route refuses on goes through here — a pre-stream rejection
/// and a mid-stream lag alike — so a client sees one shape for both: the
/// standardized `{"error": …}` frame, then a close. It never returns `Ok`.
async fn terminal<Tx>(socket_tx: &mut Tx, wire: WireError) -> Result<(), ServerError>
where
    Tx: futures::Sink<Message> + Unpin,
    <Tx as futures::Sink<Message>>::Error: std::fmt::Debug,
{
    send_wire_error(socket_tx, &wire).await?;
    Err(ServerError::Wire { wire })
}

/// End a stream whose session will produce nothing more with a close-1000
/// frame: every SDK reads that as "stream complete" and anything else as a
/// transient drop it should reconnect against.
async fn finish<Tx>(socket_tx: &mut Tx) -> Result<(), ServerError>
where
    Tx: futures::Sink<Message> + Unpin,
    <Tx as futures::Sink<Message>>::Error: std::fmt::Debug,
{
    let close = CloseFrame {
        code: close_code::NORMAL,
        reason: SESSION_COMPLETE_REASON.into(),
    };
    let close_result = socket_tx.send(Message::Close(Some(close))).await;
    // A failed close means the client already left, which is still a clean end.
    drop(close_result);
    Ok(())
}

/// Reason carried by the graceful-end close-1000 frame.
const SESSION_COMPLETE_REASON: &str = "assistant session stream complete";

/// `error_type` for an `?after=` that is not a transcript index.
const INVALID_AFTER_TYPE: &str = "AssistantWatchCursorInvalid";
/// `error_type` for a watcher that fell behind and lost frames.
const LAGGED_TYPE: &str = "AssistantSessionLagged";