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

systemprompt_api/services/gateway/stream_tap/
mod.rs
