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

systemprompt_api/services/gateway/stream_tap/
finalize.rs
