systemprompt_api/services/gateway/stream_tap/
mod.rs

1//! Streaming response tap: re-renders upstream canonical events to the inbound
2//! wire format while accumulating a full response snapshot for the audit sink.
3
4mod accumulator;
5
6use std::pin::Pin;
7use std::sync::{Arc, Mutex};
8use std::task::{Context, Poll};
9
10use axum::body::Body;
11use bytes::Bytes;
12use futures_util::stream::{BoxStream, Stream};
13
14use self::accumulator::{Summary, TapState, accumulate_event, extract_summary, snapshot};
15use super::audit::GatewayAudit;
16use super::protocol::canonical_response::CanonicalEvent;
17use super::protocol::inbound::InboundAdapter;
18
19pub fn tap(
20    upstream: BoxStream<'static, Result<CanonicalEvent, String>>,
21    inbound: Arc<dyn InboundAdapter>,
22    request_model: String,
23    audit: Arc<GatewayAudit>,
24) -> Body {
25    let state = Arc::new(Mutex::new(TapState::default()));
26    let tapped = TappedStream {
27        inner: upstream,
28        state: Arc::clone(&state),
29        inbound,
30        request_model,
31        audit,
32    };
33    Body::from_stream(tapped)
34}
35
36struct TappedStream {
37    inner: BoxStream<'static, Result<CanonicalEvent, String>>,
38    state: Arc<Mutex<TapState>>,
39    inbound: Arc<dyn InboundAdapter>,
40    request_model: String,
41    audit: Arc<GatewayAudit>,
42}
43
44impl Stream for TappedStream {
45    type Item = Result<Bytes, std::io::Error>;
46
47    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
48        loop {
49            match self.inner.as_mut().poll_next(cx) {
50                Poll::Pending => return Poll::Pending,
51                Poll::Ready(None) => {
52                    return self.finalize_on_eof();
53                },
54                Poll::Ready(Some(Err(e))) => {
55                    if let Ok(mut s) = self.state.lock() {
56                        s.error = Some(e.clone());
57                    }
58                    let err = std::io::Error::new(std::io::ErrorKind::BrokenPipe, e);
59                    return Poll::Ready(Some(Err(err)));
60                },
61                Poll::Ready(Some(Ok(event))) => {
62                    let terminal = matches!(
63                        event,
64                        CanonicalEvent::ContentBlockStop { .. }
65                            | CanonicalEvent::MessageStop { .. }
66                    );
67                    let snap = self.state.lock().map_or(None, |mut s| {
68                        accumulate_event(&mut s, &event);
69                        terminal.then(|| snapshot(&s))
70                    });
71                    let rendered = snap
72                        .as_ref()
73                        .and_then(|snapshot| {
74                            self.inbound.render_terminal_event(
75                                &event,
76                                snapshot,
77                                &self.request_model,
78                            )
79                        })
80                        .or_else(|| self.inbound.render_event(&event, &self.request_model));
81                    if let Some(bytes) = rendered {
82                        if let Ok(mut s) = self.state.lock() {
83                            s.final_bytes.extend_from_slice(&bytes);
84                        }
85                        return Poll::Ready(Some(Ok(bytes)));
86                    }
87                },
88            }
89        }
90    }
91}
92
93impl TappedStream {
94    fn take_summary(&self) -> Option<Summary> {
95        self.state.lock().ok().and_then(|mut s| {
96            if s.finalized {
97                return None;
98            }
99            s.finalized = true;
100            Some(extract_summary(&mut s))
101        })
102    }
103
104    fn finalize_on_eof(&self) -> Poll<Option<Result<Bytes, std::io::Error>>> {
105        let Some(summary) = self.take_summary() else {
106            return Poll::Ready(None);
107        };
108        finalize(Arc::clone(&self.audit), summary, "eof");
109        Poll::Ready(None)
110    }
111}
112
113impl Drop for TappedStream {
114    fn drop(&mut self) {
115        let Some(summary) = self.take_summary() else {
116            return;
117        };
118        finalize(Arc::clone(&self.audit), summary, "drop");
119    }
120}
121
122#[derive(Debug, Clone, Copy, PartialEq, Eq)]
123pub enum FinalizeDecision {
124    Fail(&'static str),
125    Complete { cost_capture_miss: bool },
126}
127
128pub const fn classify(
129    error: Option<&str>,
130    saw_stop: bool,
131    has_content: bool,
132    has_usage: bool,
133) -> FinalizeDecision {
134    if error.is_some() {
135        return FinalizeDecision::Fail("upstream stream error");
136    }
137    if !saw_stop {
138        return FinalizeDecision::Fail(if has_content {
139            "stream ended without stop event"
140        } else {
141            "empty upstream stream"
142        });
143    }
144    FinalizeDecision::Complete {
145        cost_capture_miss: has_content && !has_usage,
146    }
147}
148
149fn finalize(audit: Arc<GatewayAudit>, summary: Summary, origin: &'static str) {
150    tokio::spawn(async move {
151        if let Some(model) = summary.served_model.as_deref() {
152            audit.set_served_model(model).await;
153        }
154
155        let has_content = !summary.final_bytes.is_empty();
156        let has_usage = summary.usage.input_tokens > 0 || summary.usage.output_tokens > 0;
157        match classify(
158            summary.error.as_deref(),
159            summary.saw_stop,
160            has_content,
161            has_usage,
162        ) {
163            FinalizeDecision::Fail(reason) => {
164                let msg = summary.error.as_deref().unwrap_or(reason);
165                if let Err(e) = audit.fail(msg).await {
166                    tracing::warn!(origin, error = %e, "stream audit fail failed");
167                }
168            },
169            FinalizeDecision::Complete { cost_capture_miss } => {
170                if cost_capture_miss {
171                    tracing::warn!(
172                        origin,
173                        "stream completed with content but zero usage: cost capture miss"
174                    );
175                }
176                if let Err(e) = audit
177                    .complete(
178                        summary.usage,
179                        summary.tool_calls,
180                        &summary.response,
181                        &summary.final_bytes,
182                    )
183                    .await
184                {
185                    tracing::warn!(origin, error = %e, "stream audit complete failed");
186                }
187            },
188        }
189    });
190}
systemprompt_api/services/gateway/stream_tap/mod.rs

systemprompt_api/services/gateway/stream_tap/
mod.rs
