aion_server/stream/
transcript_stream.rs

1//! NOI-5b transcript subscription: namespace-gated durable-tail + live-splice
2//! forward loop for one `(workflow, run, activity, attempt)` agent transcript.
3//!
4//! This is the agent-observability counterpart to [`super::cluster_stream`]'s
5//! forward loop, and a NEW ARM on the existing single subscription frame of
6//! `/events/stream` (the socket stays one-subscription-per-socket; there is no
7//! multiplexing layer). A client that wants both a workflow stream and a
8//! transcript opens two `/events/stream` sockets.
9//!
10//! # Authorization: namespace-scoped (like the per-workflow event stream)
11//!
12//! A transcript belongs to the workflow the activity runs under, so it is
13//! authorized exactly like the per-workflow event subscription: the caller must
14//! hold a grant for the transcript's `namespace` AND the target `workflow_id`
15//! must be visible in it. The gate reuses the SAME
16//! [`NamespaceGuard::scope`](crate::namespace::NamespaceGuard::scope) +
17//! [`SubscriptionScope::PerWorkflow`](crate::namespace::SubscriptionScope) path
18//! the workflow stream uses, so a caller probing a foreign or nonexistent
19//! workflow receives the guard's anti-leak `not_found`, never a transcript.
20//!
21//! # Splice contract (gap-free, no duplicate)
22//!
23//! Mirrors the workflow resume path: attach the live broadcast BEFORE reading the
24//! durable `O` tail (subscribe-then-replay), so an event that races the priming
25//! read is retained by the receiver and applied after it, deduped on `store_seq`.
26//! Ephemeral token deltas (`store_seq: None`) are forwarded live and never
27//! replayed.
28
29use aion_core::{ActivityId, RunId};
30use aion_proto::{
31    PerWorkflowSubscription, ProtoRunId, ProtoWorkflowId, StreamedActivityEvent,
32    TranscriptSubscription, WireError,
33};
34use aion_store::ActivityStreamKey;
35use axum::extract::ws::{CloseFrame, Message, WebSocket, close_code};
36use futures::{SinkExt, StreamExt};
37
38use crate::activity_publisher::TranscriptStreamLagged;
39use crate::error::ServerError;
40use crate::namespace::{CallerIdentity, NamespaceOperation, SubscriptionScope, WorkflowTarget};
41use crate::state::ServerState;
42
43/// Serve a transcript subscription on an already-upgraded socket.
44///
45/// Flow: decode + namespace-gate FIRST (anti-leak: a denied caller receives one
46/// terminal wire-error frame + close, byte-identical to the workflow path); then
47/// attach the live broadcast BEFORE reading the durable `O` tail (gap-free
48/// splice); replay the durable tail; forward live events until the client closes
49/// or the subscriber lags (one typed `transcript_lagged` frame then close).
50///
51/// # Errors
52///
53/// Returns [`ServerError`] when decode/authorization fails (after the terminal
54/// frame is sent), the durable replay read fails, or the stream ends with a lag
55/// terminal frame.
56pub async fn serve_transcript_socket(
57    mut socket: WebSocket,
58    state: &ServerState,
59    caller: &CallerIdentity,
60    subscription: &TranscriptSubscription,
61) -> Result<(), ServerError> {
62    let key = match authorize_transcript(state, caller, subscription).await {
63        Ok(key) => key,
64        Err(error) => {
65            super::socket::send_wire_error(&mut socket, &error.to_wire_error()).await?;
66            return Err(error);
67        }
68    };
69
70    let publisher = state.transcript_publisher();
71    // T0: attach the live tail BEFORE reading the durable replay, so an event
72    // emitted between the replay read and the first live poll is retained by the
73    // receiver and applied after the replay (deduped on `store_seq`).
74    let mut live = publisher.subscribe(key.clone(), subscription.after_seq);
75
76    // T1 (> T0): replay the durable `O` tail from the resume cursor. `after_seq`
77    // is the highest already-applied `store_seq`; replay everything strictly
78    // after it (`None` replays the full transcript from `store_seq == 0`).
79    let from_seq = subscription
80        .after_seq
81        .map_or(0, |seq| seq.saturating_add(1));
82    let replay = publisher
83        .replay_from(&key, from_seq)
84        .await
85        .map_err(ServerError::from)?;
86    for record in replay {
87        if send_activity_frame(&mut socket, record.event)
88            .await?
89            .is_break()
90        {
91            return Ok(());
92        }
93    }
94
95    let (mut socket_tx, mut socket_rx) = socket.split();
96    loop {
97        tokio::select! {
98            client_message = socket_rx.next() => {
99                match client_message {
100                    // Close, socket error, or any inbound frame ends the read
101                    // side: the transcript channel takes no further client
102                    // frames (one-subscription-per-socket), so an inbound frame
103                    // after subscribe is a benign close.
104                    Some(Ok(Message::Close(_))) | None => {
105                        return send_normal_close(&mut socket_tx).await;
106                    }
107                    Some(Ok(_other)) => {}
108                    Some(Err(_error)) => return Ok(()),
109                }
110            }
111            item = live.next() => {
112                match item {
113                    Some(Ok(event)) => {
114                        if forward_live_frame(&mut socket_tx, event).await?.is_break() {
115                            return Ok(());
116                        }
117                    }
118                    Some(Err(TranscriptStreamLagged { skipped })) => {
119                        return deliver_transcript_terminal(&mut socket_tx, skipped).await;
120                    }
121                    None => return send_normal_close(&mut socket_tx).await,
122                }
123            }
124        }
125    }
126}
127
128/// Decode the transcript identifiers and namespace-gate the caller, returning the
129/// authorized `(workflow, run, activity, attempt)` stream key.
130///
131/// Reuses the per-workflow subscription scope so authorization is byte-identical
132/// to the workflow event stream: the caller must hold the namespace grant AND the
133/// workflow must be visible in it (anti-leak `not_found` otherwise).
134async fn authorize_transcript(
135    state: &ServerState,
136    caller: &CallerIdentity,
137    subscription: &TranscriptSubscription,
138) -> Result<ActivityStreamKey, ServerError> {
139    let workflow_id = decode_workflow_id(subscription.workflow_id.as_ref())?;
140    let run_id = decode_run_id(subscription.run_id.as_ref())?;
141    let activity_id = decode_activity_id(subscription)?;
142    gate_transcript_workflow(state, caller, &subscription.namespace, &workflow_id).await?;
143    Ok(ActivityStreamKey::new(
144        workflow_id,
145        run_id,
146        activity_id,
147        subscription.attempt,
148    ))
149}
150
151/// Per-workflow transcript gate shared by the WS subscription and the REST
152/// fetch/enumeration endpoints — byte-identical to the workflow event
153/// subscription's authorization: the caller must hold a grant for `namespace`
154/// AND `workflow_id` must be visible in it (anti-leak `not_found` otherwise).
155pub(crate) async fn gate_transcript_workflow(
156    state: &ServerState,
157    caller: &CallerIdentity,
158    namespace: &str,
159    workflow_id: &aion_core::WorkflowId,
160) -> Result<(), ServerError> {
161    let per_workflow = PerWorkflowSubscription {
162        namespace: namespace.to_owned(),
163        workflow_id: Some(ProtoWorkflowId::from(workflow_id.clone())),
164        resume_from_seq: None,
165    };
166    let target = WorkflowTarget::workflow(workflow_id);
167    let scope = SubscriptionScope::PerWorkflow(&per_workflow, target);
168    let filter = aion::EventFilter {
169        workflow_id: Some(workflow_id.clone()),
170        ..aion::EventFilter::default()
171    };
172    let operation = NamespaceOperation::subscribe(scope, &filter);
173    // Guard verdict FIRST: nothing runs for an unauthorized caller. The scoped
174    // engine handle is not needed here — transcripts read the `O` keyspace
175    // through the publisher, never the engine.
176    let scoped = state.namespace_guard().scope(caller, &operation).await?;
177    drop(scoped);
178    Ok(())
179}
180
181fn decode_workflow_id(
182    workflow_id: Option<&ProtoWorkflowId>,
183) -> Result<aion_core::WorkflowId, ServerError> {
184    workflow_id
185        .cloned()
186        .ok_or_else(|| ServerError::Wire {
187            wire: WireError::invalid_input("transcript subscription workflow_id is missing"),
188        })?
189        .try_into()
190        .map_err(|wire| ServerError::Wire { wire })
191}
192
193/// Decode the REQUIRED run axis of a transcript subscription.
194///
195/// An absent `run_id` is `invalid_input`, never a defaulted "latest run": the
196/// transcript keyspace is run-scoped, and silently resolving a generation the
197/// caller did not name is the ambiguity the axis exists to remove.
198fn decode_run_id(run_id: Option<&ProtoRunId>) -> Result<RunId, ServerError> {
199    run_id
200        .cloned()
201        .ok_or_else(|| ServerError::Wire {
202            wire: WireError::invalid_input("transcript subscription run_id is missing"),
203        })?
204        .try_into()
205        .map_err(|wire| ServerError::Wire { wire })
206}
207
208fn decode_activity_id(subscription: &TranscriptSubscription) -> Result<ActivityId, ServerError> {
209    let activity_id = subscription.activity_id.ok_or_else(|| ServerError::Wire {
210        wire: WireError::invalid_input("transcript subscription activity_id is missing"),
211    })?;
212    Ok(ActivityId::from(activity_id))
213}
214
215/// Serialize + send one transcript event on the still-unified socket (during the
216/// durable replay, before the read/write split). A send failure means the client
217/// is gone: signal a clean end.
218async fn send_activity_frame(
219    socket: &mut WebSocket,
220    event: aion_core::ActivityEvent,
221) -> Result<std::ops::ControlFlow<()>, ServerError> {
222    let frame = encode_activity_frame(&event)?;
223    if socket.send(Message::Text(frame.into())).await.is_err() {
224        return Ok(std::ops::ControlFlow::Break(()));
225    }
226    Ok(std::ops::ControlFlow::Continue(()))
227}
228
229/// Serialize + send one live transcript event on the write half of the split
230/// socket. A serialize failure sends a terminal wire-error frame and surfaces the
231/// error; a send failure is a benign client-gone end.
232async fn forward_live_frame<Tx>(
233    socket_tx: &mut Tx,
234    event: aion_core::ActivityEvent,
235) -> Result<std::ops::ControlFlow<()>, ServerError>
236where
237    Tx: futures::Sink<Message> + Unpin,
238    <Tx as futures::Sink<Message>>::Error: std::fmt::Debug,
239{
240    let frame = match encode_activity_frame(&event) {
241        Ok(frame) => frame,
242        Err(error) => {
243            super::socket::send_wire_error(socket_tx, &error.to_wire_error()).await?;
244            return Err(error);
245        }
246    };
247    if socket_tx.send(Message::Text(frame.into())).await.is_err() {
248        return Ok(std::ops::ControlFlow::Break(()));
249    }
250    Ok(std::ops::ControlFlow::Continue(()))
251}
252
253fn encode_activity_frame(event: &aion_core::ActivityEvent) -> Result<String, ServerError> {
254    let frame = StreamedActivityEvent::new(event.clone());
255    serde_json::to_string(&frame).map_err(|source| ServerError::Wire {
256        wire: WireError::backend(format!(
257            "failed to serialize transcript event frame: {source}"
258        )),
259    })
260}
261
262/// Send the typed `transcript_lagged` terminal frame + close, then surface it
263/// typed — the client re-resumes from the durable `O` tail by `store_seq`.
264async fn deliver_transcript_terminal<Tx>(
265    socket_tx: &mut Tx,
266    skipped: u64,
267) -> Result<(), ServerError>
268where
269    Tx: futures::Sink<Message> + Unpin,
270    <Tx as futures::Sink<Message>>::Error: std::fmt::Debug,
271{
272    let payload = serde_json::json!({
273        "error": { "code": "transcript_lagged", "skipped": skipped },
274    });
275    let payload = serde_json::to_string(&payload).map_err(|source| ServerError::Wire {
276        wire: WireError::backend(format!(
277            "failed to serialize transcript lag frame: {source}"
278        )),
279    })?;
280    if socket_tx.send(Message::Text(payload.into())).await.is_ok() {
281        let close = CloseFrame {
282            code: close_code::ERROR,
283            reason: "transcript_lagged".into(),
284        };
285        let close_result = socket_tx.send(Message::Close(Some(close))).await;
286        drop(close_result);
287    }
288    Err(ServerError::lagged_stream())
289}
290
291/// Finish a graceful transcript subscription end with a close-1000 frame.
292async fn send_normal_close<Tx>(socket_tx: &mut Tx) -> Result<(), ServerError>
293where
294    Tx: futures::Sink<Message> + Unpin,
295    <Tx as futures::Sink<Message>>::Error: std::fmt::Debug,
296{
297    let close = CloseFrame {
298        code: close_code::NORMAL,
299        reason: "subscription complete".into(),
300    };
301    let close_result = socket_tx.send(Message::Close(Some(close))).await;
302    drop(close_result);
303    Ok(())
304}
aion_server/stream/transcript_stream.rs

aion_server/stream/
transcript_stream.rs
