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::quota;
30use super::service::run_response_safety_scan;
31use super::signature_cache::ThoughtSignatureCache;
32
33/// Shared by the streaming and buffered completion tasks so both debit quota
34/// and run the response-phase safety scan identically.
35#[derive(Debug)]
36pub struct TapFinalizeCtx {
37    pub db: DbPool,
38    pub policy: GatewayPolicySpec,
39    pub ai_request_id: AiRequestId,
40}
41
42pub fn tap(
43    upstream: BoxStream<'static, Result<CanonicalEvent, String>>,
44    inbound: Arc<dyn InboundAdapter>,
45    request_model: String,
46    audit: Arc<GatewayAudit>,
47    finalize_ctx: TapFinalizeCtx,
48) -> Body {
49    let state = Arc::new(Mutex::new(TapState::default()));
50    let tapped = TappedStream {
51        inner: upstream,
52        state: Arc::clone(&state),
53        inbound,
54        request_model,
55        audit,
56        finalize_ctx: Some(finalize_ctx),
57    };
58    Body::from_stream(tapped)
59}
60
61struct TappedStream {
62    inner: BoxStream<'static, Result<CanonicalEvent, String>>,
63    state: Arc<Mutex<TapState>>,
64    inbound: Arc<dyn InboundAdapter>,
65    request_model: String,
66    audit: Arc<GatewayAudit>,
67    finalize_ctx: Option<TapFinalizeCtx>,
68}
69
70impl Stream for TappedStream {
71    type Item = Result<Bytes, std::io::Error>;
72
73    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
74        loop {
75            match self.inner.as_mut().poll_next(cx) {
76                Poll::Pending => return Poll::Pending,
77                Poll::Ready(None) => {
78                    return self.finalize_on_eof();
79                },
80                Poll::Ready(Some(Err(e))) => {
81                    if let Ok(mut s) = self.state.lock() {
82                        s.error = Some(e.clone());
83                    }
84                    let err = std::io::Error::new(std::io::ErrorKind::BrokenPipe, e);
85                    return Poll::Ready(Some(Err(err)));
86                },
87                Poll::Ready(Some(Ok(event))) => {
88                    let terminal = matches!(
89                        event,
90                        CanonicalEvent::ContentBlockStop { .. }
91                            | CanonicalEvent::MessageStop { .. }
92                    );
93                    let snap = self.state.lock().map_or(None, |mut s| {
94                        accumulate_event(&mut s, &event);
95                        terminal.then(|| snapshot(&s))
96                    });
97                    let rendered = snap
98                        .as_ref()
99                        .and_then(|snapshot| {
100                            self.inbound.render_terminal_event(
101                                &event,
102                                snapshot,
103                                &self.request_model,
104                            )
105                        })
106                        .or_else(|| self.inbound.render_event(&event, &self.request_model));
107                    if let Some(bytes) = rendered {
108                        if let Ok(mut s) = self.state.lock() {
109                            s.final_bytes.extend_from_slice(&bytes);
110                        }
111                        return Poll::Ready(Some(Ok(bytes)));
112                    }
113                },
114            }
115        }
116    }
117}
118
119impl TappedStream {
120    fn take_summary(&mut self) -> Option<(Summary, TapFinalizeCtx)> {
121        let ctx = self.finalize_ctx.take()?;
122        self.state.lock().ok().and_then(|mut s| {
123            if s.finalized {
124                return None;
125            }
126            s.finalized = true;
127            Some((extract_summary(&mut s), ctx))
128        })
129    }
130
131    fn finalize_on_eof(&mut self) -> Poll<Option<Result<Bytes, std::io::Error>>> {
132        let Some((summary, ctx)) = self.take_summary() else {
133            return Poll::Ready(None);
134        };
135        finalize(Arc::clone(&self.audit), summary, ctx, "eof");
136        Poll::Ready(None)
137    }
138}
139
140impl Drop for TappedStream {
141    fn drop(&mut self) {
142        let Some((summary, ctx)) = self.take_summary() else {
143            return;
144        };
145        finalize(Arc::clone(&self.audit), summary, ctx, "drop");
146    }
147}
148
149#[derive(Debug, Clone, Copy, PartialEq, Eq)]
150pub enum FinalizeDecision {
151    Fail(&'static str),
152    Complete { cost_capture_miss: bool },
153}
154
155pub const fn classify(
156    error: Option<&str>,
157    saw_stop: bool,
158    has_content: bool,
159    has_usage: bool,
160) -> FinalizeDecision {
161    if error.is_some() {
162        return FinalizeDecision::Fail("upstream stream error");
163    }
164    if !saw_stop {
165        return FinalizeDecision::Fail(if has_content {
166            "stream ended without stop event"
167        } else {
168            "empty upstream stream"
169        });
170    }
171    FinalizeDecision::Complete {
172        cost_capture_miss: has_content && !has_usage,
173    }
174}
175
176fn finalize(audit: Arc<GatewayAudit>, summary: Summary, ctx: TapFinalizeCtx, origin: &'static str) {
177    if let Some(conversation) = &audit.ctx.gateway_conversation_id {
178        ThoughtSignatureCache::global().store_from_response(conversation, &summary.response);
179    }
180    tokio::spawn(async move {
181        if let Some(model) = summary.served_model.as_deref() {
182            audit.set_served_model(model).await;
183        }
184
185        let has_content = !summary.final_bytes.is_empty();
186        let has_usage = summary.usage.input_tokens > 0 || summary.usage.output_tokens > 0;
187        match classify(
188            summary.error.as_deref(),
189            summary.saw_stop,
190            has_content,
191            has_usage,
192        ) {
193            FinalizeDecision::Fail(reason) => {
194                let msg = summary.error.as_deref().unwrap_or(reason);
195                if let Err(e) = audit.fail(msg).await {
196                    tracing::warn!(origin, error = %e, "stream audit fail failed");
197                }
198            },
199            FinalizeDecision::Complete { cost_capture_miss } => {
200                if cost_capture_miss {
201                    tracing::warn!(
202                        origin,
203                        "stream completed with content but zero usage: cost capture miss"
204                    );
205                }
206                let cost_microdollars = match audit
207                    .complete(
208                        summary.usage,
209                        summary.tool_calls,
210                        &summary.response,
211                        &summary.final_bytes,
212                    )
213                    .await
214                {
215                    Ok(cost) => cost,
216                    Err(e) => {
217                        tracing::warn!(origin, error = %e, "stream audit complete failed");
218                        0
219                    },
220                };
221                quota::post_update_tokens(
222                    &ctx.db,
223                    quota::PostUpdateParams {
224                        user_id: &audit.ctx.user_id,
225                        windows: &ctx.policy.quota_windows,
226                        input_tokens: summary.usage.input_tokens,
227                        output_tokens: summary.usage.output_tokens,
228                        cost_microdollars,
229                    },
230                )
231                .await;
232                run_response_safety_scan(
233                    &ctx.db,
234                    &ctx.ai_request_id,
235                    &summary.response,
236                    &ctx.policy.safety,
237                )
238                .await;
239            },
240        }
241    });
242}
systemprompt_api/services/gateway/stream_tap/mod.rs

systemprompt_api/services/gateway/stream_tap/
mod.rs
