Skip to main content

systemprompt_api/services/gateway/audit/
mod.rs

1//! Persistence of gateway request lifecycle to the AI-request audit trail.
2//!
3//! [`GatewayAudit`] opens a record when a request arrives (see the `open`
4//! submodule), records the canonical messages and request payload, then closes
5//! it on completion with token usage, resolved cost, latency, captured tool
6//! calls, and the response payload (see the `complete` submodule) — or marks it
7//! failed. [`GatewayRequestContext`] carries the identifiers and routing
8//! metadata bound to a single request.
9
10mod complete;
11mod message_text;
12mod open;
13pub mod payload;
14
15#[cfg(feature = "test-api")]
16pub mod test_api {
17    pub use super::message_text::flatten_message_content;
18}
19
20use std::sync::{Arc, Mutex};
21use std::time::Instant;
22
23use anyhow::Result;
24use systemprompt_ai::repository::{AiRequestPayloadRepository, AiRequestRepository};
25use systemprompt_database::DbPool;
26use systemprompt_identifiers::{
27    AiRequestId, ContextId, GatewayConversationId, SessionId, TraceId, UserId,
28};
29
30#[derive(Debug, Clone)]
31pub struct GatewayRequestContext {
32    pub ai_request_id: AiRequestId,
33    pub user_id: UserId,
34    pub session_id: Option<SessionId>,
35    pub context_id: ContextId,
36    pub gateway_conversation_id: Option<GatewayConversationId>,
37    pub trace_id: Option<TraceId>,
38    pub provider: String,
39    /// The upstream model the request dispatches to (after route rewrite). The
40    /// audit `model` column is opened from this, then overwritten by
41    /// `set_served_model` with the model the provider echoes back.
42    pub model: String,
43    /// The model the client requested on the wire, before route rewrite.
44    /// Persisted to `ai_requests.requested_model` so an audit retains both.
45    pub requested_model: Option<String>,
46    pub max_tokens: Option<u32>,
47    pub is_streaming: bool,
48    pub wire_protocol: String,
49}
50
51#[expect(
52    missing_debug_implementations,
53    reason = "service type holds repository clients that intentionally do not implement Debug"
54)]
55pub struct GatewayAudit {
56    requests: Arc<AiRequestRepository>,
57    payloads: Arc<AiRequestPayloadRepository>,
58    pub ctx: GatewayRequestContext,
59    served_model: Mutex<Option<String>>,
60    started_at: Instant,
61}
62
63impl GatewayAudit {
64    pub fn new(
65        db: &DbPool,
66        ctx: GatewayRequestContext,
67    ) -> Result<Self, systemprompt_ai::error::RepositoryError> {
68        let requests = Arc::new(AiRequestRepository::new(db)?);
69        let payloads = Arc::new(AiRequestPayloadRepository::new(db)?);
70        Ok(Self {
71            requests,
72            payloads,
73            ctx,
74            served_model: Mutex::new(None),
75            started_at: Instant::now(),
76        })
77    }
78
79    pub async fn set_served_model(&self, model: &str) {
80        if model.is_empty() || model == self.ctx.model {
81            return;
82        }
83        if let Ok(mut slot) = self.served_model.lock() {
84            *slot = Some(model.to_owned());
85        }
86        if let Err(e) = self
87            .requests
88            .update_model(&self.ctx.ai_request_id, model)
89            .await
90        {
91            tracing::warn!(error = %e, "update_model failed");
92        }
93    }
94
95    pub async fn set_system_prompt_override(&self, descriptor: &str) {
96        if let Err(e) = self
97            .requests
98            .update_system_prompt_override(&self.ctx.ai_request_id, descriptor)
99            .await
100        {
101            tracing::warn!(error = %e, "update_system_prompt_override failed");
102        }
103    }
104
105    pub async fn set_route_match(&self, descriptor: &str) {
106        if let Err(e) = self
107            .requests
108            .update_route_match(&self.ctx.ai_request_id, descriptor)
109            .await
110        {
111            tracing::warn!(error = %e, "update_route_match failed");
112        }
113    }
114
115    pub async fn fail(&self, error: &str) -> Result<()> {
116        if let Err(e) = self
117            .requests
118            .update_error(&self.ctx.ai_request_id, error)
119            .await
120        {
121            tracing::warn!(error = %e, "audit fail update failed");
122        }
123        tracing::info!(
124            ai_request_id = %self.ctx.ai_request_id,
125            user_id = %self.ctx.user_id,
126            provider = %self.ctx.provider,
127            model = %self.ctx.model,
128            error,
129            "Gateway audit: request failed"
130        );
131        Ok(())
132    }
133}