nanocodex_oai_api/events/
stream.rs

1use std::{
2    io::Write,
3    sync::{
4        Arc, OnceLock,
5        atomic::{AtomicU64, Ordering},
6    },
7};
8
9use futures_util::Stream;
10use serde::{Deserialize, Serialize};
11use serde_json::value::{RawValue, to_raw_value};
12use tokio::sync::mpsc;
13use web_time::Instant;
14
15const PROTOCOL_VERSION: u32 = 1;
16static PROCESS_MONOTONIC_EPOCH: OnceLock<Instant> = OnceLock::new();
17
18/// Returns a process-relative monotonic timestamp for private cross-layer timing.
19#[doc(hidden)]
20#[must_use]
21pub fn monotonic_now_ns() -> u64 {
22    let elapsed = PROCESS_MONOTONIC_EPOCH
23        .get_or_init(Instant::now)
24        .elapsed()
25        .as_nanos();
26    u64::try_from(elapsed).unwrap_or(u64::MAX)
27}
28
29/// Failure while encoding, writing, or consuming the contractual event stream.
30#[derive(Debug, thiserror::Error)]
31pub enum EventError {
32    /// A typed event could not be encoded as JSON.
33    #[error("failed to encode agent event")]
34    Encode(#[source] serde_json::Error),
35
36    /// An encoded event could not be written to the supplied output.
37    #[error("failed to write agent event")]
38    Write(#[source] std::io::Error),
39
40    /// The stream closed before the accepted turn emitted a terminal event.
41    #[error("agent event stream closed before the turn emitted a terminal event")]
42    ClosedBeforeTerminal,
43}
44
45/// One ordered event emitted by an agent run.
46#[derive(Clone, Debug, Deserialize, Serialize)]
47pub struct AgentEvent {
48    /// Version of the stable event protocol.
49    pub protocol_version: u32,
50    /// Stable session/request identity shared by this event stream.
51    pub request_id: Arc<str>,
52    /// Monotonic sequence number within the stream.
53    pub seq: u64,
54    /// Stable event category.
55    #[serde(rename = "type")]
56    pub kind: AgentEventKind,
57    /// Complete typed-event payload encoded as retained raw JSON.
58    pub payload: Arc<RawValue>,
59}
60
61/// Private in-process timing carried beside an event without changing JSONL.
62#[doc(hidden)]
63#[derive(Clone, Copy, Debug)]
64pub struct AgentEventTiming {
65    /// Process-relative nanoseconds when the event was emitted.
66    pub emitted_ns: u64,
67    /// Process-relative nanoseconds when the transport observed the source event.
68    pub source_received_ns: Option<u64>,
69}
70
71/// An agent event plus private in-process delivery timing.
72#[doc(hidden)]
73#[derive(Clone, Debug)]
74pub struct TimedAgentEvent {
75    /// Contractual event visible to consumers.
76    pub event: AgentEvent,
77    /// Private in-process timing carried beside the contractual event.
78    pub timing: AgentEventTiming,
79}
80
81/// Stable event categories emitted by the agent runtime.
82#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
83pub enum AgentEventKind {
84    /// Complete provider event in original order.
85    #[serde(rename = "api.event")]
86    ApiEvent,
87    /// Incremental assistant text.
88    #[serde(rename = "assistant.delta")]
89    AssistantDelta,
90    /// Completed assistant message.
91    #[serde(rename = "assistant.message")]
92    AssistantMessage,
93    /// Incremental reasoning summary.
94    #[serde(rename = "reasoning.summary.delta")]
95    ReasoningSummaryDelta,
96    /// Accepted turn started.
97    #[serde(rename = "run.started")]
98    RunStarted,
99    /// Input was added to an active turn.
100    #[serde(rename = "run.steered")]
101    RunSteered,
102    /// Recoverable run-level error was observed.
103    #[serde(rename = "run.error")]
104    RunError,
105    /// Turn completed successfully.
106    #[serde(rename = "run.completed")]
107    RunCompleted,
108    /// Turn terminated with an error.
109    #[serde(rename = "run.failed")]
110    RunFailed,
111    /// Tool invocation started.
112    #[serde(rename = "tool.call")]
113    ToolCall,
114    /// Tool invocation completed.
115    #[serde(rename = "tool.result")]
116    ToolResult,
117    /// Optional model connection warmup started.
118    #[serde(rename = "model.warmup.started")]
119    ModelWarmupStarted,
120    /// Optional model connection warmup completed.
121    #[serde(rename = "model.warmup.completed")]
122    ModelWarmupCompleted,
123    /// Optional model connection warmup failed.
124    #[serde(rename = "model.warmup.failed")]
125    ModelWarmupFailed,
126    /// Logical model call started.
127    #[serde(rename = "model.call.started")]
128    ModelCallStarted,
129    /// Logical model call completed.
130    #[serde(rename = "model.call.completed")]
131    ModelCallCompleted,
132    /// Logical model call failed.
133    #[serde(rename = "model.call.failed")]
134    ModelCallFailed,
135    /// Model-side context compaction started.
136    #[serde(rename = "model.compaction.started")]
137    ModelCompactionStarted,
138    /// Model-side context compaction completed.
139    #[serde(rename = "model.compaction.completed")]
140    ModelCompactionCompleted,
141    /// Model-side context compaction failed.
142    #[serde(rename = "model.compaction.failed")]
143    ModelCompactionFailed,
144    /// One transport attempt started.
145    #[serde(rename = "model.attempt.started")]
146    ModelAttemptStarted,
147    /// One transport attempt failed.
148    #[serde(rename = "model.attempt.failed")]
149    ModelAttemptFailed,
150    /// The SDK scheduled another transport attempt.
151    #[serde(rename = "model.attempt.retrying")]
152    ModelAttemptRetrying,
153    /// A model transport connection started.
154    #[serde(rename = "model.connection.started")]
155    ModelConnectionStarted,
156    /// A model transport connection completed.
157    #[serde(rename = "model.connection.completed")]
158    ModelConnectionCompleted,
159    /// A model transport connection failed.
160    #[serde(rename = "model.connection.failed")]
161    ModelConnectionFailed,
162}
163
164/// The receiving half of an agent's typed event stream.
165pub struct AgentEvents {
166    request_id: Arc<str>,
167    receiver: mpsc::UnboundedReceiver<TimedAgentEvent>,
168}
169
170impl AgentEvents {
171    /// Stable session/request identifier shared by every event in this stream.
172    #[must_use]
173    pub fn request_id(&self) -> &str {
174        &self.request_id
175    }
176
177    /// Receives the next event, or `None` after all emitters are dropped.
178    pub async fn recv(&mut self) -> Option<AgentEvent> {
179        self.recv_timed().await.map(|event| event.event)
180    }
181
182    /// Receives one event with private process-relative timing metadata.
183    #[doc(hidden)]
184    pub async fn recv_timed(&mut self) -> Option<TimedAgentEvent> {
185        self.receiver.recv().await
186    }
187
188    /// Receives one immediately available event without waiting.
189    #[doc(hidden)]
190    pub fn try_recv_timed(&mut self) -> Option<TimedAgentEvent> {
191        self.receiver.try_recv().ok()
192    }
193
194    /// Writes every event as one flushed JSONL record.
195    ///
196    /// # Errors
197    ///
198    /// Returns an error when an event cannot be encoded or written.
199    pub async fn write_jsonl(mut self, mut output: impl Write) -> Result<(), EventError> {
200        while let Some(event) = self.recv().await {
201            write_jsonl_event(&mut output, &event)?;
202        }
203        Ok(())
204    }
205
206    /// Writes one turn through its terminal event and leaves the session stream
207    /// available for follow-on turns.
208    ///
209    /// # Errors
210    ///
211    /// Returns an error when an event cannot be written or the agent stops
212    /// before emitting `run.completed` or `run.failed`.
213    pub async fn write_turn_jsonl(&mut self, mut output: impl Write) -> Result<(), EventError> {
214        while let Some(event) = self.recv().await {
215            let terminal = event.kind.is_terminal();
216            write_jsonl_event(&mut output, &event)?;
217            if terminal {
218                return Ok(());
219            }
220        }
221        Err(EventError::ClosedBeforeTerminal)
222    }
223}
224
225impl Stream for AgentEvents {
226    type Item = AgentEvent;
227
228    fn poll_next(
229        mut self: std::pin::Pin<&mut Self>,
230        context: &mut std::task::Context<'_>,
231    ) -> std::task::Poll<Option<Self::Item>> {
232        self.receiver
233            .poll_recv(context)
234            .map(|event| event.map(|event| event.event))
235    }
236}
237
238impl AgentEventKind {
239    /// Returns whether this event completes a turn.
240    #[must_use]
241    pub const fn is_terminal(self) -> bool {
242        matches!(self, Self::RunCompleted | Self::RunFailed)
243    }
244}
245
246impl AgentEvent {
247    /// Returns a stable typed projection of this event.
248    ///
249    /// Raw `OpenAI` frames and lower-level transport diagnostics remain
250    /// lossless; application-facing run, assistant, reasoning, tool, model,
251    /// and context events decode into named types.
252    ///
253    /// # Errors
254    ///
255    /// Returns an error when a payload does not satisfy the contract declared
256    /// by its event kind.
257    pub fn data(&self) -> Result<crate::AgentEventData, serde_json::Error> {
258        use crate::{
259            AgentEventData, AssistantEvent, ContextEvent, ModelEvent, ReasoningEvent, RunEvent,
260            ToolEvent, TransportEvent,
261        };
262
263        Ok(match self.kind {
264            AgentEventKind::ApiEvent => AgentEventData::OpenAi(self.decode_payload()?),
265            AgentEventKind::AssistantDelta => {
266                AgentEventData::Assistant(AssistantEvent::Delta(self.decode_payload()?))
267            }
268            AgentEventKind::AssistantMessage => {
269                AgentEventData::Assistant(AssistantEvent::Message(self.decode_payload()?))
270            }
271            AgentEventKind::ReasoningSummaryDelta => {
272                AgentEventData::Reasoning(ReasoningEvent::SummaryDelta(self.decode_payload()?))
273            }
274            AgentEventKind::RunStarted => {
275                AgentEventData::Run(RunEvent::Started(self.decode_payload()?))
276            }
277            AgentEventKind::RunSteered => {
278                AgentEventData::Run(RunEvent::Steered(self.decode_payload()?))
279            }
280            AgentEventKind::RunError => {
281                AgentEventData::Run(RunEvent::Error(self.decode_payload()?))
282            }
283            AgentEventKind::RunCompleted => {
284                AgentEventData::Run(RunEvent::Completed(Box::new(self.decode_payload()?)))
285            }
286            AgentEventKind::RunFailed => {
287                AgentEventData::Run(RunEvent::Failed(Box::new(self.decode_payload()?)))
288            }
289            AgentEventKind::ToolCall => {
290                AgentEventData::Tool(ToolEvent::Call(self.decode_payload()?))
291            }
292            AgentEventKind::ToolResult => {
293                AgentEventData::Tool(ToolEvent::Result(self.decode_payload()?))
294            }
295            AgentEventKind::ModelWarmupStarted => {
296                AgentEventData::Model(ModelEvent::WarmupStarted(self.decode_payload()?))
297            }
298            AgentEventKind::ModelWarmupCompleted => {
299                AgentEventData::Model(ModelEvent::WarmupCompleted(self.decode_payload()?))
300            }
301            AgentEventKind::ModelWarmupFailed => {
302                AgentEventData::Model(ModelEvent::WarmupFailed(self.decode_payload()?))
303            }
304            AgentEventKind::ModelCallStarted => {
305                AgentEventData::Model(ModelEvent::CallStarted(self.decode_payload()?))
306            }
307            AgentEventKind::ModelCallCompleted => {
308                AgentEventData::Model(ModelEvent::CallCompleted(self.decode_payload()?))
309            }
310            AgentEventKind::ModelCallFailed => {
311                AgentEventData::Model(ModelEvent::CallFailed(self.decode_payload()?))
312            }
313            AgentEventKind::ModelCompactionStarted => {
314                AgentEventData::Context(ContextEvent::CompactionStarted(self.decode_payload()?))
315            }
316            AgentEventKind::ModelCompactionCompleted => {
317                AgentEventData::Context(ContextEvent::CompactionCompleted(self.decode_payload()?))
318            }
319            AgentEventKind::ModelCompactionFailed => {
320                AgentEventData::Context(ContextEvent::CompactionFailed(self.decode_payload()?))
321            }
322            AgentEventKind::ModelAttemptStarted
323            | AgentEventKind::ModelAttemptFailed
324            | AgentEventKind::ModelAttemptRetrying
325            | AgentEventKind::ModelConnectionStarted
326            | AgentEventKind::ModelConnectionCompleted
327            | AgentEventKind::ModelConnectionFailed => {
328                AgentEventData::Transport(TransportEvent::new(self.kind, Arc::clone(&self.payload)))
329            }
330        })
331    }
332
333    /// Decodes the event payload into a caller-selected typed shape.
334    ///
335    /// # Errors
336    ///
337    /// Returns an error when the retained payload does not match `T`.
338    pub fn decode_payload<T: serde::de::DeserializeOwned>(&self) -> Result<T, serde_json::Error> {
339        serde_json::from_str(self.payload.get())
340    }
341}
342
343fn write_jsonl_event(output: &mut impl Write, event: &AgentEvent) -> Result<(), EventError> {
344    serde_json::to_writer(&mut *output, event).map_err(EventError::Encode)?;
345    output
346        .write_all(b"\n")
347        .and_then(|()| output.flush())
348        .map_err(EventError::Write)
349}
350
351/// Internal emission handle shared by orchestration and transport crates.
352#[derive(Clone)]
353pub struct EventSink {
354    request_id: Arc<str>,
355    next_seq: Arc<AtomicU64>,
356    sender: mpsc::UnboundedSender<TimedAgentEvent>,
357    mirror: Option<mpsc::UnboundedSender<TimedAgentEvent>>,
358}
359
360impl EventSink {
361    /// Creates an emission handle and its independently consumed event stream.
362    #[must_use]
363    pub fn channel(request_id: String) -> (Self, AgentEvents) {
364        let request_id = Arc::<str>::from(request_id);
365        let (sender, receiver) = mpsc::unbounded_channel();
366        (
367            Self {
368                request_id: Arc::clone(&request_id),
369                next_seq: Arc::new(AtomicU64::new(1)),
370                sender,
371                mirror: None,
372            },
373            AgentEvents {
374                request_id,
375                receiver,
376            },
377        )
378    }
379
380    /// Returns the stable request/session identity attached to emitted events.
381    #[must_use]
382    pub fn request_id(&self) -> &str {
383        &self.request_id
384    }
385
386    /// Creates a sink that mirrors its events into one independently owned stream.
387    ///
388    /// The returned sink preserves the parent sink's request identity and
389    /// sequence counter. Dropping every clone of it closes only the mirror;
390    /// the original session stream remains available.
391    #[must_use]
392    pub fn mirrored_channel(&self) -> (Self, AgentEvents) {
393        let (mirror, receiver) = mpsc::unbounded_channel();
394        (
395            Self {
396                request_id: Arc::clone(&self.request_id),
397                next_seq: Arc::clone(&self.next_seq),
398                sender: self.sender.clone(),
399                mirror: Some(mirror),
400            },
401            AgentEvents {
402                request_id: Arc::clone(&self.request_id),
403                receiver,
404            },
405        )
406    }
407
408    /// Emits an event when a receiver is present and otherwise discards it.
409    ///
410    /// # Errors
411    ///
412    /// Returns an error when the payload cannot be converted to JSON.
413    pub fn emit<P: Serialize>(&self, kind: AgentEventKind, payload: P) -> Result<(), EventError> {
414        self.emit_with_sequence(kind, payload).map(|_| ())
415    }
416
417    /// Emits an event and returns its session-monotonic sequence number.
418    ///
419    /// This is intended for transport telemetry that must correlate the point
420    /// of emission with a downstream consumer without retaining payload data.
421    ///
422    /// # Errors
423    ///
424    /// Returns an error when the payload cannot be converted to JSON.
425    pub fn emit_with_sequence<P: Serialize>(
426        &self,
427        kind: AgentEventKind,
428        payload: P,
429    ) -> Result<u64, EventError> {
430        self.emit_with_source_sequence(kind, payload, None)
431    }
432
433    /// Emits an event correlated with the process-monotonic source receipt time.
434    pub fn emit_with_source_sequence<P: Serialize>(
435        &self,
436        kind: AgentEventKind,
437        payload: P,
438        source_received_ns: Option<u64>,
439    ) -> Result<u64, EventError> {
440        if self.sender.is_closed()
441            && self
442                .mirror
443                .as_ref()
444                .is_none_or(tokio::sync::mpsc::UnboundedSender::is_closed)
445        {
446            return Ok(self.next_seq.fetch_add(1, Ordering::Relaxed));
447        }
448        let payload = Arc::from(to_raw_value(&payload).map_err(EventError::Encode)?);
449        let seq = self.next_seq.fetch_add(1, Ordering::Relaxed);
450        let event = TimedAgentEvent {
451            event: AgentEvent {
452                protocol_version: PROTOCOL_VERSION,
453                request_id: Arc::clone(&self.request_id),
454                seq,
455                kind,
456                payload,
457            },
458            timing: AgentEventTiming {
459                emitted_ns: monotonic_now_ns(),
460                source_received_ns,
461            },
462        };
463        drop(self.sender.send(event.clone()));
464        if let Some(mirror) = &self.mirror {
465            drop(mirror.send(event));
466        }
467        Ok(seq)
468    }
469}
470
471#[cfg(test)]
472mod tests {
473    use serde::{Serialize, Serializer};
474    use serde_json::json;
475
476    use super::{AgentEventKind, EventSink};
477    use crate::{AgentEventData, AssistantEvent, ToolEvent, TransportEvent};
478
479    #[test]
480    fn events_are_ordered_and_receiver_drop_is_not_an_error() {
481        let (events, mut receiver) = EventSink::channel("request-1".to_owned());
482        assert_eq!(receiver.request_id(), "request-1");
483        events
484            .emit(AgentEventKind::RunStarted, json!({ "n": 1 }))
485            .unwrap();
486        events
487            .emit(AgentEventKind::RunCompleted, json!({ "n": 2 }))
488            .unwrap();
489        let first = receiver.receiver.try_recv().unwrap().event;
490        let second = receiver.receiver.try_recv().unwrap().event;
491        assert_eq!((first.seq, first.kind), (1, AgentEventKind::RunStarted));
492        assert_eq!((second.seq, second.kind), (2, AgentEventKind::RunCompleted));
493        assert_eq!(
494            second.decode_payload::<serde_json::Value>().unwrap()["n"],
495            2
496        );
497        drop(receiver);
498        events.emit(AgentEventKind::RunFailed, json!({})).unwrap();
499    }
500
501    #[test]
502    fn receiver_drop_skips_payload_serialization() {
503        struct MustNotSerialize;
504
505        impl Serialize for MustNotSerialize {
506            fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
507            where
508                S: Serializer,
509            {
510                panic!("closed event streams must not serialize payloads")
511            }
512        }
513
514        let (events, receiver) = EventSink::channel("request-1".to_owned());
515        drop(receiver);
516
517        assert_eq!(
518            events
519                .emit_with_sequence(AgentEventKind::ApiEvent, MustNotSerialize)
520                .unwrap(),
521            1
522        );
523    }
524
525    #[test]
526    fn timing_is_private_and_preserves_the_jsonl_contract() {
527        let (events, mut receiver) = EventSink::channel("request-1".to_owned());
528        let source_received_ns = super::monotonic_now_ns();
529        events
530            .emit_with_source_sequence(
531                AgentEventKind::AssistantDelta,
532                json!({ "text": "x" }),
533                Some(source_received_ns),
534            )
535            .unwrap();
536
537        let timed = receiver.receiver.try_recv().unwrap();
538        assert_eq!(timed.timing.source_received_ns, Some(source_received_ns));
539        assert!(timed.timing.emitted_ns >= source_received_ns);
540        let encoded = serde_json::to_value(&timed.event).unwrap();
541        assert!(encoded.get("timing").is_none());
542        assert!(encoded.get("source_received_ns").is_none());
543        assert_eq!(encoded["type"], "assistant.delta");
544    }
545
546    #[test]
547    fn timed_events_can_be_drained_without_async_receive_round_trips() {
548        let (events, mut receiver) = EventSink::channel("request-1".to_owned());
549        for n in 1..=3 {
550            events
551                .emit(AgentEventKind::AssistantDelta, json!({ "n": n }))
552                .unwrap();
553        }
554
555        let sequences = std::iter::from_fn(|| receiver.try_recv_timed())
556            .map(|event| event.event.seq)
557            .collect::<Vec<_>>();
558        assert_eq!(sequences, vec![1, 2, 3]);
559        assert!(receiver.try_recv_timed().is_none());
560    }
561
562    #[test]
563    fn mirrored_stream_preserves_session_order_and_closes_independently() {
564        let (events, mut session) = EventSink::channel("request-1".to_owned());
565        let (turn_events, mut turn) = events.mirrored_channel();
566
567        turn_events
568            .emit(AgentEventKind::RunStarted, json!({ "turn": 1 }))
569            .unwrap();
570        turn_events
571            .emit(AgentEventKind::RunCompleted, json!({ "turn": 1 }))
572            .unwrap();
573
574        let session_first = session.receiver.try_recv().unwrap().event;
575        let session_second = session.receiver.try_recv().unwrap().event;
576        let turn_first = turn.receiver.try_recv().unwrap().event;
577        let turn_second = turn.receiver.try_recv().unwrap().event;
578        assert_eq!(
579            (session_first.seq, session_second.seq),
580            (turn_first.seq, turn_second.seq)
581        );
582        assert_eq!(turn_second.kind, AgentEventKind::RunCompleted);
583
584        drop(turn_events);
585        assert!(turn.receiver.try_recv().is_err());
586        events
587            .emit(AgentEventKind::RunStarted, json!({ "turn": 2 }))
588            .unwrap();
589        assert_eq!(
590            session.receiver.try_recv().unwrap().event.seq,
591            session_second.seq + 1
592        );
593    }
594
595    #[test]
596    fn typed_projection_preserves_domain_values_and_raw_diagnostics() {
597        let (events, mut receiver) = EventSink::channel("request-1".to_owned());
598        events
599            .emit(
600                AgentEventKind::AssistantDelta,
601                json!({
602                    "model_call_index": 2,
603                    "item_id": "item-1",
604                    "phase": "final_answer",
605                    "text": "hello"
606                }),
607            )
608            .unwrap();
609        events
610            .emit(
611                AgentEventKind::ToolCall,
612                json!({
613                    "call_id": "call-1",
614                    "tool": "deployment_region",
615                    "arguments": {"service": "api"},
616                    "model_call_index": 2
617                }),
618            )
619            .unwrap();
620        events
621            .emit(
622                AgentEventKind::ModelAttemptRetrying,
623                json!({"attempt": 1, "next_attempt": 2}),
624            )
625            .unwrap();
626
627        let assistant = receiver.receiver.try_recv().unwrap().event;
628        let AgentEventData::Assistant(AssistantEvent::Delta(delta)) = assistant.data().unwrap()
629        else {
630            panic!("assistant delta should use the typed assistant projection");
631        };
632        assert_eq!(delta.text, "hello");
633        assert_eq!(delta.model_call_index, 2);
634
635        let tool = receiver.receiver.try_recv().unwrap().event;
636        let AgentEventData::Tool(ToolEvent::Call(call)) = tool.data().unwrap() else {
637            panic!("tool call should use the typed tool projection");
638        };
639        assert_eq!(call.tool, "deployment_region");
640        assert_eq!(
641            call.decode_arguments::<serde_json::Value>().unwrap()["service"],
642            "api"
643        );
644
645        let diagnostic = receiver.receiver.try_recv().unwrap().event;
646        let AgentEventData::Transport(transport) = diagnostic.data().unwrap() else {
647            panic!("retry should remain a lossless transport diagnostic");
648        };
649        assert_eq!(
650            TransportEvent::kind(&transport),
651            AgentEventKind::ModelAttemptRetrying
652        );
653        assert_eq!(
654            transport.decode_payload::<serde_json::Value>().unwrap()["next_attempt"],
655            2
656        );
657    }
658}
nanocodex_oai_api/events/stream.rs

nanocodex_oai_api/events/
stream.rs
