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//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7mod accumulator;
8
9#[cfg(feature = "test-api")]
10pub mod test_api {
11    pub use super::accumulator::{Summary, TapState, accumulate_event, extract_summary, snapshot};
12}
13
14use std::pin::Pin;
15use std::sync::{Arc, Mutex};
16use std::task::{Context, Poll};
17
18use axum::body::Body;
19use bytes::Bytes;
20use futures_util::stream::{BoxStream, Stream};
21use systemprompt_database::DbPool;
22use systemprompt_identifiers::AiRequestId;
23
24use self::accumulator::{Summary, TapState, accumulate_event, extract_summary, snapshot};
25use super::audit::GatewayAudit;
26use super::policy::GatewayPolicySpec;
27use super::protocol::canonical_response::CanonicalEvent;
28use super::protocol::inbound::InboundAdapter;
29use super::protocol::outbound::anthropic::streaming::SseDecoder;
30use super::quota;
31use super::service::run_response_safety_scan;
32use super::signature_cache::ThoughtSignatureCache;
33
34/// Shared by the streaming and buffered completion tasks so both debit quota
35/// and run the response-phase safety scan identically.
36#[derive(Debug)]
37pub struct TapFinalizeCtx {
38    pub db: DbPool,
39    pub repos: crate::services::gateway::GatewayRepositories,
40    pub policy: GatewayPolicySpec,
41    pub ai_request_id: AiRequestId,
42}
43
44pub fn tap(
45    upstream: BoxStream<'static, Result<CanonicalEvent, String>>,
46    inbound: Arc<dyn InboundAdapter>,
47    request_model: String,
48    audit: Arc<GatewayAudit>,
49    finalize_ctx: TapFinalizeCtx,
50) -> Body {
51    let state = Arc::new(Mutex::new(TapState::default()));
52    let tapped = TappedStream {
53        inner: upstream,
54        state: Arc::clone(&state),
55        inbound,
56        request_model,
57        audit,
58        finalize_ctx: Some(finalize_ctx),
59    };
60    Body::from_stream(tapped)
61}
62
63pub fn tap_raw(
64    upstream: BoxStream<'static, Result<Bytes, String>>,
65    audit: Arc<GatewayAudit>,
66    finalize_ctx: TapFinalizeCtx,
67) -> Body {
68    Body::from_stream(RawTappedStream {
69        inner: upstream,
70        state: Arc::new(Mutex::new(TapState::default())),
71        decoder: SseDecoder::default(),
72        audit,
73        finalize_ctx: Some(finalize_ctx),
74    })
75}
76
77struct RawTappedStream {
78    inner: BoxStream<'static, Result<Bytes, String>>,
79    state: Arc<Mutex<TapState>>,
80    decoder: SseDecoder,
81    audit: Arc<GatewayAudit>,
82    finalize_ctx: Option<TapFinalizeCtx>,
83}
84
85impl RawTappedStream {
86    fn take_summary(&mut self) -> Option<(Summary, TapFinalizeCtx)> {
87        let ctx = self.finalize_ctx.take()?;
88        self.state.lock().ok().and_then(|mut s| {
89            if s.finalized {
90                return None;
91            }
92            s.finalized = true;
93            Some((extract_summary(&mut s), ctx))
94        })
95    }
96}
97
98impl Stream for RawTappedStream {
99    type Item = Result<Bytes, std::io::Error>;
100
101    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
102        match self.inner.as_mut().poll_next(cx) {
103            Poll::Pending => Poll::Pending,
104            Poll::Ready(None) => {
105                if let Some((summary, ctx)) = self.take_summary() {
106                    finalize(Arc::clone(&self.audit), summary, ctx, "eof");
107                }
108                Poll::Ready(None)
109            },
110            Poll::Ready(Some(Err(e))) => {
111                if let Ok(mut s) = self.state.lock() {
112                    s.error = Some(e.clone());
113                }
114                Poll::Ready(Some(Err(std::io::Error::new(
115                    std::io::ErrorKind::BrokenPipe,
116                    e,
117                ))))
118            },
119            Poll::Ready(Some(Ok(bytes))) => {
120                let events = self.decoder.push(&bytes);
121                if let Ok(mut s) = self.state.lock() {
122                    for event in &events {
123                        accumulate_event(&mut s, event);
124                    }
125                    s.final_bytes.extend_from_slice(&bytes);
126                }
127                Poll::Ready(Some(Ok(bytes)))
128            },
129        }
130    }
131}
132
133impl Drop for RawTappedStream {
134    fn drop(&mut self) {
135        let Some((summary, ctx)) = self.take_summary() else {
136            return;
137        };
138        finalize(Arc::clone(&self.audit), summary, ctx, "drop");
139    }
140}
141
142struct TappedStream {
143    inner: BoxStream<'static, Result<CanonicalEvent, String>>,
144    state: Arc<Mutex<TapState>>,
145    inbound: Arc<dyn InboundAdapter>,
146    request_model: String,
147    audit: Arc<GatewayAudit>,
148    finalize_ctx: Option<TapFinalizeCtx>,
149}
150
151impl Stream for TappedStream {
152    type Item = Result<Bytes, std::io::Error>;
153
154    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
155        loop {
156            match self.inner.as_mut().poll_next(cx) {
157                Poll::Pending => return Poll::Pending,
158                Poll::Ready(None) => {
159                    return self.finalize_on_eof();
160                },
161                Poll::Ready(Some(Err(e))) => {
162                    if let Ok(mut s) = self.state.lock() {
163                        s.error = Some(e.clone());
164                    }
165                    let err = std::io::Error::new(std::io::ErrorKind::BrokenPipe, e);
166                    return Poll::Ready(Some(Err(err)));
167                },
168                Poll::Ready(Some(Ok(event))) => {
169                    let terminal = matches!(
170                        event,
171                        CanonicalEvent::ContentBlockStop { .. }
172                            | CanonicalEvent::MessageStop { .. }
173                    );
174                    let snap = self.state.lock().map_or(None, |mut s| {
175                        accumulate_event(&mut s, &event);
176                        terminal.then(|| snapshot(&s))
177                    });
178                    let rendered = snap
179                        .as_ref()
180                        .and_then(|snapshot| {
181                            self.inbound.render_terminal_event(
182                                &event,
183                                snapshot,
184                                &self.request_model,
185                            )
186                        })
187                        .or_else(|| self.inbound.render_event(&event, &self.request_model));
188                    if let Some(bytes) = rendered {
189                        if let Ok(mut s) = self.state.lock() {
190                            s.final_bytes.extend_from_slice(&bytes);
191                        }
192                        return Poll::Ready(Some(Ok(bytes)));
193                    }
194                },
195            }
196        }
197    }
198}
199
200impl TappedStream {
201    fn take_summary(&mut self) -> Option<(Summary, TapFinalizeCtx)> {
202        let ctx = self.finalize_ctx.take()?;
203        self.state.lock().ok().and_then(|mut s| {
204            if s.finalized {
205                return None;
206            }
207            s.finalized = true;
208            Some((extract_summary(&mut s), ctx))
209        })
210    }
211
212    fn finalize_on_eof(&mut self) -> Poll<Option<Result<Bytes, std::io::Error>>> {
213        let Some((summary, ctx)) = self.take_summary() else {
214            return Poll::Ready(None);
215        };
216        finalize(Arc::clone(&self.audit), summary, ctx, "eof");
217        Poll::Ready(None)
218    }
219}
220
221impl Drop for TappedStream {
222    fn drop(&mut self) {
223        let Some((summary, ctx)) = self.take_summary() else {
224            return;
225        };
226        finalize(Arc::clone(&self.audit), summary, ctx, "drop");
227    }
228}
229
230#[derive(Debug, Clone, Copy, PartialEq, Eq)]
231pub enum FinalizeDecision {
232    Fail(&'static str),
233    Complete { cost_capture_miss: bool },
234}
235
236pub const fn classify(
237    error: Option<&str>,
238    saw_stop: bool,
239    has_content: bool,
240    has_usage: bool,
241) -> FinalizeDecision {
242    if error.is_some() {
243        return FinalizeDecision::Fail("upstream stream error");
244    }
245    if !saw_stop {
246        return FinalizeDecision::Fail(if has_content {
247            "stream ended without stop event"
248        } else {
249            "empty upstream stream"
250        });
251    }
252    FinalizeDecision::Complete {
253        cost_capture_miss: has_content && !has_usage,
254    }
255}
256
257fn finalize(audit: Arc<GatewayAudit>, summary: Summary, ctx: TapFinalizeCtx, origin: &'static str) {
258    if let Some(conversation) = &audit.ctx.gateway_conversation_id {
259        ThoughtSignatureCache::global().store_from_response(conversation, &summary.response);
260    }
261    tokio::spawn(async move {
262        if let Some(model) = summary.served_model.as_deref() {
263            audit.set_served_model(model).await;
264        }
265
266        let has_content = !summary.final_bytes.is_empty();
267        let has_usage = summary.saw_usage_delta
268            && (summary.usage.input_tokens > 0 || summary.usage.output_tokens > 0);
269        match classify(
270            summary.error.as_deref(),
271            summary.saw_stop,
272            has_content,
273            has_usage,
274        ) {
275            FinalizeDecision::Fail(reason) => {
276                let msg = summary.error.as_deref().unwrap_or(reason);
277                if let Err(e) = audit.fail(msg).await {
278                    tracing::warn!(origin, error = %e, "stream audit fail failed");
279                }
280            },
281            FinalizeDecision::Complete { cost_capture_miss } => {
282                if cost_capture_miss {
283                    tracing::warn!(
284                        origin,
285                        "stream completed with content but zero usage: cost capture miss"
286                    );
287                }
288                let cost_microdollars = match audit
289                    .complete(
290                        summary.usage,
291                        summary.tool_calls,
292                        &summary.response,
293                        &summary.final_bytes,
294                    )
295                    .await
296                {
297                    Ok(cost) => cost,
298                    Err(e) => {
299                        tracing::warn!(origin, error = %e, "stream audit complete failed");
300                        0
301                    },
302                };
303                quota::post_update_tokens(
304                    &ctx.db,
305                    &ctx.repos.quota_buckets,
306                    quota::PostUpdateParams {
307                        user_id: &audit.ctx.user_id,
308                        windows: &ctx.policy.quota_windows,
309                        input_tokens: summary.usage.input_tokens,
310                        output_tokens: summary.usage.output_tokens,
311                        cost_microdollars,
312                    },
313                )
314                .await;
315                run_response_safety_scan(
316                    &ctx.repos.safety_findings,
317                    &ctx.ai_request_id,
318                    &summary.response,
319                    &ctx.policy.safety,
320                )
321                .await;
322            },
323        }
324    });
325}
systemprompt_api/services/gateway/stream_tap/mod.rs

systemprompt_api/services/gateway/stream_tap/
mod.rs
