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
63/// Relays an upstream byte stream unchanged while accounting for what passes.
64///
65/// The client receives the provider's own SSE frames; a decoded copy feeds the
66/// same accumulator the translating tap uses, so audit, cost, quota, and the
67/// response-phase safety scan behave identically on both lanes.
68pub fn tap_raw(
69    upstream: BoxStream<'static, Result<Bytes, String>>,
70    audit: Arc<GatewayAudit>,
71    finalize_ctx: TapFinalizeCtx,
72) -> Body {
73    Body::from_stream(RawTappedStream {
74        inner: upstream,
75        state: Arc::new(Mutex::new(TapState::default())),
76        decoder: SseDecoder::default(),
77        audit,
78        finalize_ctx: Some(finalize_ctx),
79    })
80}
81
82struct RawTappedStream {
83    inner: BoxStream<'static, Result<Bytes, String>>,
84    state: Arc<Mutex<TapState>>,
85    decoder: SseDecoder,
86    audit: Arc<GatewayAudit>,
87    finalize_ctx: Option<TapFinalizeCtx>,
88}
89
90impl RawTappedStream {
91    fn take_summary(&mut self) -> Option<(Summary, TapFinalizeCtx)> {
92        let ctx = self.finalize_ctx.take()?;
93        self.state.lock().ok().and_then(|mut s| {
94            if s.finalized {
95                return None;
96            }
97            s.finalized = true;
98            Some((extract_summary(&mut s), ctx))
99        })
100    }
101}
102
103impl Stream for RawTappedStream {
104    type Item = Result<Bytes, std::io::Error>;
105
106    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
107        match self.inner.as_mut().poll_next(cx) {
108            Poll::Pending => Poll::Pending,
109            Poll::Ready(None) => {
110                if let Some((summary, ctx)) = self.take_summary() {
111                    finalize(Arc::clone(&self.audit), summary, ctx, "eof");
112                }
113                Poll::Ready(None)
114            },
115            Poll::Ready(Some(Err(e))) => {
116                if let Ok(mut s) = self.state.lock() {
117                    s.error = Some(e.clone());
118                }
119                Poll::Ready(Some(Err(std::io::Error::new(
120                    std::io::ErrorKind::BrokenPipe,
121                    e,
122                ))))
123            },
124            Poll::Ready(Some(Ok(bytes))) => {
125                let events = self.decoder.push(&bytes);
126                if let Ok(mut s) = self.state.lock() {
127                    for event in &events {
128                        accumulate_event(&mut s, event);
129                    }
130                    s.final_bytes.extend_from_slice(&bytes);
131                }
132                Poll::Ready(Some(Ok(bytes)))
133            },
134        }
135    }
136}
137
138impl Drop for RawTappedStream {
139    fn drop(&mut self) {
140        let Some((summary, ctx)) = self.take_summary() else {
141            return;
142        };
143        finalize(Arc::clone(&self.audit), summary, ctx, "drop");
144    }
145}
146
147struct TappedStream {
148    inner: BoxStream<'static, Result<CanonicalEvent, String>>,
149    state: Arc<Mutex<TapState>>,
150    inbound: Arc<dyn InboundAdapter>,
151    request_model: String,
152    audit: Arc<GatewayAudit>,
153    finalize_ctx: Option<TapFinalizeCtx>,
154}
155
156impl Stream for TappedStream {
157    type Item = Result<Bytes, std::io::Error>;
158
159    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
160        loop {
161            match self.inner.as_mut().poll_next(cx) {
162                Poll::Pending => return Poll::Pending,
163                Poll::Ready(None) => {
164                    return self.finalize_on_eof();
165                },
166                Poll::Ready(Some(Err(e))) => {
167                    if let Ok(mut s) = self.state.lock() {
168                        s.error = Some(e.clone());
169                    }
170                    let err = std::io::Error::new(std::io::ErrorKind::BrokenPipe, e);
171                    return Poll::Ready(Some(Err(err)));
172                },
173                Poll::Ready(Some(Ok(event))) => {
174                    let terminal = matches!(
175                        event,
176                        CanonicalEvent::ContentBlockStop { .. }
177                            | CanonicalEvent::MessageStop { .. }
178                    );
179                    let snap = self.state.lock().map_or(None, |mut s| {
180                        accumulate_event(&mut s, &event);
181                        terminal.then(|| snapshot(&s))
182                    });
183                    let rendered = snap
184                        .as_ref()
185                        .and_then(|snapshot| {
186                            self.inbound.render_terminal_event(
187                                &event,
188                                snapshot,
189                                &self.request_model,
190                            )
191                        })
192                        .or_else(|| self.inbound.render_event(&event, &self.request_model));
193                    if let Some(bytes) = rendered {
194                        if let Ok(mut s) = self.state.lock() {
195                            s.final_bytes.extend_from_slice(&bytes);
196                        }
197                        return Poll::Ready(Some(Ok(bytes)));
198                    }
199                },
200            }
201        }
202    }
203}
204
205impl TappedStream {
206    fn take_summary(&mut self) -> Option<(Summary, TapFinalizeCtx)> {
207        let ctx = self.finalize_ctx.take()?;
208        self.state.lock().ok().and_then(|mut s| {
209            if s.finalized {
210                return None;
211            }
212            s.finalized = true;
213            Some((extract_summary(&mut s), ctx))
214        })
215    }
216
217    fn finalize_on_eof(&mut self) -> Poll<Option<Result<Bytes, std::io::Error>>> {
218        let Some((summary, ctx)) = self.take_summary() else {
219            return Poll::Ready(None);
220        };
221        finalize(Arc::clone(&self.audit), summary, ctx, "eof");
222        Poll::Ready(None)
223    }
224}
225
226impl Drop for TappedStream {
227    fn drop(&mut self) {
228        let Some((summary, ctx)) = self.take_summary() else {
229            return;
230        };
231        finalize(Arc::clone(&self.audit), summary, ctx, "drop");
232    }
233}
234
235#[derive(Debug, Clone, Copy, PartialEq, Eq)]
236pub enum FinalizeDecision {
237    Fail(&'static str),
238    Complete { cost_capture_miss: bool },
239}
240
241pub const fn classify(
242    error: Option<&str>,
243    saw_stop: bool,
244    has_content: bool,
245    has_usage: bool,
246) -> FinalizeDecision {
247    if error.is_some() {
248        return FinalizeDecision::Fail("upstream stream error");
249    }
250    if !saw_stop {
251        return FinalizeDecision::Fail(if has_content {
252            "stream ended without stop event"
253        } else {
254            "empty upstream stream"
255        });
256    }
257    FinalizeDecision::Complete {
258        cost_capture_miss: has_content && !has_usage,
259    }
260}
261
262fn finalize(audit: Arc<GatewayAudit>, summary: Summary, ctx: TapFinalizeCtx, origin: &'static str) {
263    if let Some(conversation) = &audit.ctx.gateway_conversation_id {
264        ThoughtSignatureCache::global().store_from_response(conversation, &summary.response);
265    }
266    tokio::spawn(async move {
267        if let Some(model) = summary.served_model.as_deref() {
268            audit.set_served_model(model).await;
269        }
270
271        let has_content = !summary.final_bytes.is_empty();
272        let has_usage = summary.saw_usage_delta
273            && (summary.usage.input_tokens > 0 || summary.usage.output_tokens > 0);
274        match classify(
275            summary.error.as_deref(),
276            summary.saw_stop,
277            has_content,
278            has_usage,
279        ) {
280            FinalizeDecision::Fail(reason) => {
281                let msg = summary.error.as_deref().unwrap_or(reason);
282                if let Err(e) = audit.fail(msg).await {
283                    tracing::warn!(origin, error = %e, "stream audit fail failed");
284                }
285            },
286            FinalizeDecision::Complete { cost_capture_miss } => {
287                if cost_capture_miss {
288                    tracing::warn!(
289                        origin,
290                        "stream completed with content but zero usage: cost capture miss"
291                    );
292                }
293                let cost_microdollars = match audit
294                    .complete(
295                        summary.usage,
296                        summary.tool_calls,
297                        &summary.response,
298                        &summary.final_bytes,
299                    )
300                    .await
301                {
302                    Ok(cost) => cost,
303                    Err(e) => {
304                        tracing::warn!(origin, error = %e, "stream audit complete failed");
305                        0
306                    },
307                };
308                quota::post_update_tokens(
309                    &ctx.db,
310                    &ctx.repos.quota_buckets,
311                    quota::PostUpdateParams {
312                        user_id: &audit.ctx.user_id,
313                        windows: &ctx.policy.quota_windows,
314                        input_tokens: summary.usage.input_tokens,
315                        output_tokens: summary.usage.output_tokens,
316                        cost_microdollars,
317                    },
318                )
319                .await;
320                run_response_safety_scan(
321                    &ctx.repos.safety_findings,
322                    &ctx.ai_request_id,
323                    &summary.response,
324                    &ctx.policy.safety,
325                )
326                .await;
327            },
328        }
329    });
330}
systemprompt_api/services/gateway/stream_tap/mod.rs

systemprompt_api/services/gateway/stream_tap/
mod.rs
