Skip to main content

systemprompt_api/services/gateway/stream_tap/
finalize.rs

1//! Terminal accounting for a tapped stream: classifies how the stream ended
2//! and spawns the audit, quota, and response-scan completion work.
3//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7use std::sync::Arc;
8
9use systemprompt_identifiers::TraceId;
10use systemprompt_logging::LogActor;
11
12use crate::routes::gateway::{TerminalOutcome, log_gateway_terminal};
13
14use super::super::audit::GatewayAudit;
15use super::super::quota;
16use super::super::service::run_response_safety_scan;
17use super::super::signature_cache::ThoughtSignatureCache;
18use super::TapFinalizeCtx;
19use super::accumulator::Summary;
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum FinalizeDecision {
23    Fail(FailCause),
24    Complete { cost_capture_miss: bool },
25}
26
27/// Why a tapped stream ended without a usable response.
28///
29/// Both land as `failed` in the audit row, so they are separated here:
30/// `Upstream` is the provider failing mid-stream, `Truncated` is the stream
31/// stopping without a terminal event — typically the client hanging up.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum FailCause {
34    Upstream,
35    Truncated { has_content: bool },
36}
37
38impl FailCause {
39    pub const fn reason(self) -> &'static str {
40        match self {
41            Self::Upstream => "upstream stream error",
42            Self::Truncated { has_content: true } => "stream ended without stop event",
43            Self::Truncated { has_content: false } => "empty upstream stream",
44        }
45    }
46
47    // Why: SSE errors arrive after the upstream HTTP status has already been sent.
48    const fn status(self) -> u16 {
49        match self {
50            Self::Upstream => 502,
51            Self::Truncated { .. } => 499,
52        }
53    }
54}
55
56pub const fn classify(
57    error: Option<&str>,
58    saw_stop: bool,
59    has_content: bool,
60    has_usage: bool,
61) -> FinalizeDecision {
62    if error.is_some() {
63        return FinalizeDecision::Fail(FailCause::Upstream);
64    }
65    if !saw_stop {
66        return FinalizeDecision::Fail(FailCause::Truncated { has_content });
67    }
68    FinalizeDecision::Complete {
69        cost_capture_miss: has_content && !has_usage,
70    }
71}
72
73pub(super) fn finalize(
74    audit: Arc<GatewayAudit>,
75    summary: Summary,
76    ctx: TapFinalizeCtx,
77    origin: &'static str,
78) {
79    tokio::spawn(async move {
80        capture_signatures(&ctx, &audit, &summary).await;
81        if let Some(model) = summary.served_model.as_deref() {
82            audit.set_served_model(model).await;
83        }
84
85        let has_content = !summary.final_bytes.is_empty();
86        let has_usage = summary.saw_usage_delta
87            && (summary.usage.input_tokens > 0 || summary.usage.output_tokens > 0);
88        match classify(
89            summary.error.as_deref(),
90            summary.saw_stop,
91            has_content,
92            has_usage,
93        ) {
94            FinalizeDecision::Fail(cause) => {
95                let msg = summary.error.as_deref().unwrap_or_else(|| cause.reason());
96                if let Err(e) = audit.fail(msg).await {
97                    tracing::warn!(origin, error = %e, "stream audit fail failed");
98                }
99                log_terminal(&audit, cause.status(), Some(msg));
100            },
101            FinalizeDecision::Complete { cost_capture_miss } => {
102                if cost_capture_miss {
103                    tracing::warn!(
104                        origin,
105                        "stream completed with content but zero usage: cost capture miss"
106                    );
107                }
108                let cost_microdollars = match audit
109                    .complete(
110                        summary.usage,
111                        summary.tool_calls,
112                        &summary.response,
113                        &summary.final_bytes,
114                    )
115                    .await
116                {
117                    Ok(cost) => cost,
118                    Err(e) => {
119                        tracing::warn!(origin, error = %e, "stream audit complete failed");
120                        0
121                    },
122                };
123                quota::post_update_tokens(
124                    &ctx.db,
125                    &ctx.repos.quota_buckets,
126                    quota::PostUpdateParams {
127                        user_id: &audit.ctx.user_id,
128                        windows: &ctx.policy.quota_windows,
129                        input_tokens: summary.usage.input_tokens,
130                        output_tokens: summary.usage.output_tokens,
131                        cost_microdollars,
132                    },
133                )
134                .await;
135                run_response_safety_scan(
136                    &ctx.repos.safety_findings,
137                    &ctx.ai_request_id,
138                    &summary.response,
139                    &ctx.policy.safety,
140                )
141                .await;
142                log_terminal(&audit, 200, None);
143            },
144        }
145    });
146}
147
148fn log_terminal(audit: &GatewayAudit, status: u16, error: Option<&str>) {
149    let Some(access) = audit.ctx.access_log.as_ref() else {
150        return;
151    };
152    log_gateway_terminal(TerminalOutcome {
153        access,
154        status,
155        actor: terminal_actor(audit),
156        error,
157    });
158}
159
160fn terminal_actor(audit: &GatewayAudit) -> Option<LogActor> {
161    if let (Some(session), Some(trace)) =
162        (audit.ctx.session_id.as_ref(), audit.ctx.trace_id.as_ref())
163    {
164        return Some(LogActor::new(
165            audit.ctx.user_id.clone(),
166            session.clone(),
167            trace.clone(),
168        ));
169    }
170    LogActor::platform(TraceId::system()).ok()
171}
172
173async fn capture_signatures(ctx: &TapFinalizeCtx, audit: &GatewayAudit, summary: &Summary) {
174    match &audit.ctx.gateway_conversation_id {
175        Some(conversation) => {
176            ctx.repos
177                .thought_signatures
178                .store_from_response(conversation, &summary.response)
179                .await;
180        },
181        None => {
182            ThoughtSignatureCache::note_uncacheable_response(
183                &summary.response,
184                "no_conversation_id",
185            );
186        },
187    }
188}