diff --git a/crates/aisix-obs/src/access_log.rs b/crates/aisix-obs/src/access_log.rs index 387e1b68..2d7cb380 100644 --- a/crates/aisix-obs/src/access_log.rs +++ b/crates/aisix-obs/src/access_log.rs @@ -5,27 +5,27 @@ //! //! # When the line is written, and what that costs //! -//! WHEN differs by path, and it decides which fields can be filled at all. -//! Four cases, and only the first is "at the end of the request": +//! Exactly one line per request, whatever the outcome — but WHEN it is +//! written differs by path, and that decides which fields can be filled at +//! all. Four cases: //! //! - **Non-streamed response** — from the handler, on its way out, with //! everything it resolved available. -//! - **Streamed response** — from the handler too, but when the SSE body is -//! handed to the server, BEFORE a single frame is polled. The upstream has -//! produced nothing yet, so the token counts and `provider_request_id` are -//! necessarily absent, and `status` is the response-OPEN status: a stream -//! that later aborts, or whose consumer walks away, still logged `200`. -//! `latency` is time-to-first-token for the same reason, NOT how long the -//! stream ran — the two differ by the whole length of the stream, which -//! for an LLM is routinely minutes. A stream's real end is only on its -//! `UsageEvent`; reading this line as the end of the request is how a -//! long-running stream gets mistaken for a connection sitting idle -//! (AISIX-Cloud#1394). +//! - **Streamed response** — NOT when the SSE head is handed to the server. +//! The handler defers the line to the request's attribution cell +//! (`attribution::defer_access_log`) and it goes out beside the request's +//! TERMINAL usage event, at the point the stream's outcome is known: +//! fully consumed, abandoned mid-stream, or dropped before its first +//! poll. It therefore reports the same `status`, `error_kind` and `error` +//! as that event — a stream whose consumer walked away reads `499` / +//! `client_disconnected` on both — and it can carry the token counts and +//! `provider_request_id`, which only exist once the upstream has answered +//! (AISIX-Cloud#1571). //! - **`/v1/realtime`** — the opposite extreme. The handler returns the //! WebSocket upgrade immediately; the line is written by `run_session` on //! a detached task once the session closes, so it carries the close status //! and the session's real token totals. -//! - **Caller hung up before the response was delivered** — written from +//! - **Caller hung up before the response head was written** — from //! `ClientCancelGuard::drop`, with no handler involved. Status is `499`, //! and the fields it can fill are the ones the request published to its //! attribution cell as it resolved: `model`, `provider`, and the @@ -33,18 +33,24 @@ //! handler-side figures — tokens, `provider_request_id`, the routing //! counts — stay `None`, because the future was dropped before it could //! produce them. Such a request also emits a `499` usage event carrying -//! the same identities (AISIX-Cloud#1571), keyed by this `request_id`. -//! Only the no-response-head case writes one: a request whose head DID go -//! out already has its handler's line, and a second one under a second -//! status would make one request read as two. That request's `499` lives -//! on its usage event alone. +//! the same identities, keyed by this `request_id`. //! //! So do not add a field whose value only exists once the upstream has //! responded and expect it on every line: it is silently empty on the -//! streamed and cancelled ones. A streamed request's completion-time figures -//! live on the per-attempt `UsageEvent` (and, for the provider response id, -//! on the `provider call completed` line `UsageSink::try_emit` writes), -//! keyed by the same `request_id`. +//! cancelled ones, where the request never got that far. +//! +//! # `latency` and `duration` answer two different questions +//! +//! - `latency_ms` is what the CALLER waited for: the first token forwarded +//! downstream on a streamed response, the complete response on a buffered +//! one. It is the same figure the request's terminal `UsageEvent` reports +//! as `downstream_latency_ms`. Deliberately not the length of the stream +//! (AISIX-Cloud#1394) — reading a minutes-long stream's wait as its +//! time-to-first-token is what makes a working stream look like a +//! connection sitting idle. +//! - `duration_ms` is how long the request occupied the gateway, arrival to +//! last byte out. On a non-streamed request the two coincide; on a +//! streamed one they differ by the whole length of the stream. use std::time::Duration; @@ -58,7 +64,22 @@ pub struct AccessLog<'a> { pub method: &'a str, pub path: &'a str, pub status: u16, + /// What the caller waited for — see the module docs. On a streamed + /// response this is time-to-first-token, NOT how long the stream ran. pub latency: Duration, + /// How long the request occupied the gateway: from arrival to the point + /// this line is written. Equal to `latency` on everything that is not + /// streamed. + /// + /// "The point this line is written" is the request's end on every + /// surface whose record is written at completion — which is all of them + /// except the two that meter at their handler tail and relay an + /// open-ended body afterwards (`/v1/audio/speech`, billed per input + /// character, and `/v1/videos/{id}/content`, metered by the + /// submission). Those two have no completion-time emitter to carry a + /// line, so theirs ends at the response head and does not span the + /// relay. + pub duration: Duration, pub provider: Option<&'a str>, /// The model name the CALLER addressed — for a routing group, the group /// itself, never the target it dispatched to. See `upstream_model` @@ -85,13 +106,14 @@ pub struct AccessLog<'a> { /// /// `None` whenever no id exists by the time this line is written: /// the request never reached an upstream (guardrail block, - /// pre-dispatch error), it was served from cache, the endpoint's + /// pre-dispatch error), it was served from cache, or the endpoint's /// provider response carries no id at all (embeddings / audio / - /// images / count_tokens), or the response is **streamed** — there the - /// id arrives in the first frame, after this line. Streamed and - /// mid-stream-failed-over calls are covered instead by the per-attempt - /// `provider call completed` line (see `UsageSink::try_emit`), which - /// shares this `request_id`. + /// images / count_tokens). A **streamed** response does carry it — + /// the id arrives in the first frame and the line is written at the + /// stream's end (AISIX-Cloud#1571) — unless the caller walked away + /// before that frame. Mid-stream-failed-over calls are covered by the + /// per-attempt `provider call completed` line (see + /// `UsageSink::try_emit`), which shares this `request_id`. pub provider_request_id: Option<&'a str>, /// Routing target that ultimately served the request (the winning /// attempt's display name). `None` for direct models / cache hits. @@ -147,6 +169,7 @@ impl AccessLog<'_> { path = self.path, status = self.status, latency_ms = self.latency.as_millis() as u64, + duration_ms = self.duration.as_millis() as u64, provider = self.provider, model = self.model, upstream_model = self.upstream_model, @@ -220,6 +243,7 @@ mod tests { path: "/v1/chat/completions", status: 200, latency: Duration::from_millis(42), + duration: Duration::from_millis(9_000), provider: Some("openai"), model: Some("my-gpt4"), upstream_model: Some("gpt-4o"), @@ -245,6 +269,12 @@ mod tests { assert!(out.contains("method=\"POST\"") || out.contains("method=POST")); assert!(out.contains("status=200")); assert!(out.contains("latency_ms=42")); + // AISIX-Cloud#1571: the two figures are separate fields because on + // a streamed line they are separate questions — what the caller + // waited for, and how long the request held the gateway. This one + // is deliberately the longer of the two, so transposing them at the + // emit site cannot pass. + assert!(out.contains("duration_ms=9000"), "{out}"); assert!(out.contains("provider=\"openai\"") || out.contains("provider=openai")); assert!(out.contains("total_tokens=3")); assert!(out.contains("request_id=\"req-abc\"") || out.contains("request_id=req-abc")); @@ -297,6 +327,7 @@ mod tests { path: "/v1/messages", status: 504, latency: Duration::from_millis(7167), + duration: Duration::from_millis(7167), provider: None, model: Some("claude-sonnet-4"), upstream_model: None, @@ -349,6 +380,7 @@ mod tests { path: "/v1/chat/completions", status: 401, latency: Duration::from_millis(1), + duration: Duration::from_millis(1), provider: None, model: None, upstream_model: None, diff --git a/crates/aisix-proxy/AGENTS.md b/crates/aisix-proxy/AGENTS.md index a9a983c5..657f60b2 100644 --- a/crates/aisix-proxy/AGENTS.md +++ b/crates/aisix-proxy/AGENTS.md @@ -200,6 +200,20 @@ nothing, and nothing errors: the caller gets a correct status while the gateway keeps no record of the request, which is indistinguishable from the request never arriving. +One exception, and it is the whole of it: **a STREAMED response's line is not +the handler's to write.** Its tail runs when the head goes out, which is not +when the request ends — so it parks the line on the attribution cell +(`attribution::PendingAccessLog`) and the line goes out from +`usage_attr::emit_usage` with the request's TERMINAL usage event, whichever of +the stream's endings produced it. That is what makes the line and the row agree +on `status`, `error_class` and `error_message` by construction; writing the line +at the tail instead reported every abandoned stream as a `200` beside its own +`499` row (AISIX-Cloud#1571). Two consequences for a new streaming family: park +the line whenever the response IS a stream (a "telemetry already emitted" flag +is NOT the same predicate — chat's buffered ensemble sets one), and make sure +the stream really does emit a terminal usage event, because that emit is now +the only thing that writes the line. + Two shapes give up early, and both must answer through `reject::reject_before_dispatch` (it renders the envelope *and* emits the telemetry, so the two can't drift apart): diff --git a/crates/aisix-proxy/src/a2a.rs b/crates/aisix-proxy/src/a2a.rs index 369f86c6..d66945cb 100644 --- a/crates/aisix-proxy/src/a2a.rs +++ b/crates/aisix-proxy/src/a2a.rs @@ -164,35 +164,53 @@ pub async fn a2a_endpoint( let elapsed = started.elapsed(); let status = response.status().as_u16(); - let target = crate::attribution::AccessLogTarget::current(); - AccessLog { - method: http_method.as_str(), - path: "/a2a", - status, - latency: elapsed, - provider: Some("a2a"), - model: None, - upstream_model: target.upstream_model(), - provider_key_id: target.provider_key_id(), - api_key_id: Some(&api_key_id), - // Counted inside `dispatch`, which hands back only a rendered - // `Response` — and for a stream, not until its drop guard fires, long - // after this line. The usage event carries them. - prompt_tokens: None, - completion_tokens: None, - total_tokens: None, - request_id: &request_id, - // Same as `/mcp`: `dispatch` returns an already-rendered `Response`, - // so no typed error reaches this point. - error_kind: None, - error: None, - provider_request_id: None, - served_by_model: None, - routing_attempt_count: None, - routing_fallback_count: None, - mcp: None, + if crate::attribution::stream_owns_access_log() { + // A streamed call ends when the agent's last event is relayed or + // the caller walks away, both of which are below this frame and + // minutes away. Park the line; `StreamUsageOnDrop` writes it beside + // the usage event that already reports that ending + // (AISIX-Cloud#1571). + crate::attribution::defer_access_log( + crate::attribution::PendingAccessLog::new( + http_method.as_str(), + "/a2a", + &request_id, + &api_key_id, + started, + ) + .with_model("a2a", ""), + ); + } else { + let target = crate::attribution::AccessLogTarget::current(); + AccessLog { + method: http_method.as_str(), + path: "/a2a", + status, + latency: elapsed, + duration: elapsed, + provider: Some("a2a"), + model: None, + upstream_model: target.upstream_model(), + provider_key_id: target.provider_key_id(), + api_key_id: Some(&api_key_id), + // Counted inside `dispatch`, which hands back only a rendered + // `Response`. The usage event carries them. + prompt_tokens: None, + completion_tokens: None, + total_tokens: None, + request_id: &request_id, + // Same as `/mcp`: `dispatch` returns an already-rendered + // `Response`, so no typed error reaches this point. + error_kind: None, + error: None, + provider_request_id: None, + served_by_model: None, + routing_attempt_count: None, + routing_fallback_count: None, + mcp: None, + } + .emit(); } - .emit(); crate::request_metrics::record( &state, "/a2a", @@ -631,6 +649,12 @@ async fn dispatch_stream( drop(guard); }); + // The endpoint tail below owns this request's access-log line, and from + // there the response is opaque — it cannot tell this stream from a + // rendered error. Say so here, so the tail parks the line for the guard + // above to write at the call's real end (AISIX-Cloud#1571). + crate::attribution::note_stream_owns_access_log(); + let mut response = axum::response::Sse::new(sse); if let Some(interval) = crate::sse_keepalive::interval() { response = response.keep_alive(axum::response::sse::KeepAlive::new().interval(interval)); @@ -2142,4 +2166,41 @@ mod tests { assert_eq!(card["version"], "2.1.0"); assert_eq!(card["skills"][0]["id"], "extract"); } + + /// A streamed `/a2a` call writes ONE access-log line, at the task's end + /// rather than when the head went out, so each of the three endings + /// reports its own outcome (AISIX-Cloud#1571). The tail that writes this + /// family's line sees only an opaque `Response`, so the streaming branch + /// tells it to park the line instead. + /// + /// `latency_ms` is deliberately NOT asserted to be a time-to-first-event + /// here: an agent's stream of task updates is the call's product rather + /// than a delivery mechanism, so `/a2a` records the WHOLE stream as what + /// the caller waited for — and the line reports the same figure its + /// usage event does. + #[tokio::test] + async fn a_streamed_call_writes_one_line_per_stream_ending() { + let agent_url = spawn_progressing_stream_agent().await; + let handle = SnapshotHandle::new(snapshot_with(&agent_url, true, serde_json::json!(["*"]))); + let hub = Arc::new(aisix_gateway::Hub::new()); + let router = build_router(ProxyState::new(handle, hub, &proxy_cfg()).without_cache()); + + let endings = crate::test_log::three_stream_endings(router, || { + HttpRequest::post("/a2a/invoice") + .header("host", "gw.example.com") + .header("content-type", "application/json") + .header("authorization", format!("Bearer {TOKEN}")) + .body(Body::from( + r#"{"jsonrpc":"2.0","id":"s","method":"message/stream"}"#, + )) + .unwrap() + }) + .await; + crate::test_log::assert_one_line_per_ending(&endings, "/a2a", "ak-1"); + assert_eq!( + endings.delivered.field("provider").as_deref(), + Some("a2a"), + "the line must keep naming the family it belongs to", + ); + } } diff --git a/crates/aisix-proxy/src/attribution.rs b/crates/aisix-proxy/src/attribution.rs index 3ab8c92a..8a90e1c5 100644 --- a/crates/aisix-proxy/src/attribution.rs +++ b/crates/aisix-proxy/src/attribution.rs @@ -38,6 +38,11 @@ //! [`crate::attempt::RoutingTelemetry`] — so no endpoint has to opt in and //! none of them can drift out. //! +//! The cell carries one more thing for the same reason: a streamed +//! response's access-log LINE ([`PendingAccessLog`]). Its handler returns +//! when the head goes out, minutes before the request ends, so the line is +//! parked here and written by whichever terminal emitter ends the request. +//! //! It is kept BESIDE [`Resolved`] rather than inside it because every //! failed request reads `Resolved` back by value for its metric labels //! ([`current`]); folding a `Vec` and a `ClientContext` into @@ -45,6 +50,7 @@ use std::future::Future; use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; use aisix_core::Model; @@ -208,10 +214,168 @@ pub(crate) struct CancelContext { pub emitted_terminal: bool, } +/// A streamed response's access-log line, parked by its handler until the +/// request has an outcome to report. +/// +/// A streaming handler returns the moment the response HEAD exists — before +/// a single frame has been polled, and often minutes before the request +/// ends. Writing the line there is what made a stream the caller abandoned +/// log `200` beside its own `499` usage event, and left every streamed line +/// without the token counts and `provider_request_id` that arrive with the +/// response (AISIX-Cloud#1571). +/// +/// So the handler parks here the fields only it can resolve, and the line +/// goes out beside the request's TERMINAL usage event — from +/// [`crate::usage_attr::emit_usage`], the one point every terminal event of +/// every family passes through. That is what makes the two agree on +/// `status`, `error_kind` and `error` by construction rather than by each +/// family remembering to. +/// +/// Taking it out of the cell is also the interlock: whichever completion +/// point gets there first emits, and the rest find nothing. A request can +/// therefore never write two lines, however many of its terminal emitters +/// race (see `crate::GuardPhase`). +pub(crate) struct PendingAccessLog { + method: String, + /// Bounded route template, never the caller's raw path (#451). + path: String, + provider: String, + /// The model name the CALLER addressed. The dispatched target is read + /// off [`Resolved`] at emit time instead, exactly as the handler's own + /// line reads it through [`AccessLogTarget`]. + model: String, + api_key_id: String, + request_id: String, + served_by_model: String, + routing_attempt_count: Option, + routing_fallback_count: Option, + /// The request clock, so the line can report how long the whole request + /// occupied the gateway — which on a stream is not what the caller + /// waited for (see `AccessLog::duration`). + started: Instant, +} + +impl PendingAccessLog { + /// What every family can say: who called, where, and when the request + /// started. `started` is the REQUEST clock, not the attempt's. + pub(crate) fn new( + method: &str, + path: &str, + request_id: &str, + api_key_id: &str, + started: Instant, + ) -> Self { + Self { + method: method.to_string(), + path: path.to_string(), + provider: String::new(), + model: String::new(), + api_key_id: api_key_id.to_string(), + request_id: request_id.to_string(), + served_by_model: String::new(), + routing_attempt_count: None, + routing_fallback_count: None, + started, + } + } + + /// The vendor and the model name the CALLER addressed. + pub(crate) fn with_model(mut self, provider: &str, model: &str) -> Self { + self.provider = provider.to_string(); + self.model = model.to_string(); + self + } + + /// The routing summary, for the families that dispatch through a group. + /// Same shape their own inline line carries: the winner's display name, + /// and the two counts, each absent when zero. + pub(crate) fn with_routing(mut self, routing: &crate::attempt::RoutingTelemetry) -> Self { + self.served_by_model = routing + .winner() + .map(|w| w.target_model.clone()) + .unwrap_or_default(); + self.routing_attempt_count = match routing.attempt_count() { + 0 => None, + n => Some(n), + }; + self.routing_fallback_count = match routing.fallback_count() { + 0 => None, + n => Some(n), + }; + self + } + + /// Write the line, taking its outcome from the terminal usage event. + fn emit(self, target: &Resolved, event: &aisix_obs::UsageEvent) { + let duration = self.started.elapsed(); + // What the CALLER waited for, taken VERBATIM off the event: the + // first token forwarded downstream on a stream that delivered one. + // Not re-derived, and no sentinel handling — the line and the row + // are one record of one request, so they must not be able to report + // two different waits, and a sub-millisecond first frame is a real + // `0` rather than a missing stamp. It reads `0` exactly where the + // event says the caller received nothing, which the `499` beside it + // explains. Deliberately NOT the length of the stream, which + // `duration` is (AISIX-Cloud#1394). + let latency = Duration::from_millis(u64::from(event.downstream_latency_ms)); + let prompt = u64::from(event.prompt_tokens); + let completion = u64::from(event.completion_tokens); + // Cache-inclusive, the way every emitter in this crate computes a + // request's total: `prompt_tokens` excludes the cache dimensions on + // the Anthropic-shaped paths, so summing the two visible columns + // would put a cached request's line an order of magnitude under the + // row it was emitted beside. + let total = crate::usage_attr::total_tokens_with_cache( + event.prompt_tokens, + event.completion_tokens, + event.cache_creation_tokens, + event.cache_read_tokens, + ); + // Keep a token-less outcome out of the token columns entirely, + // rather than logging an abandoned stream as a zero-token success — + // the same rule the rest of this line follows for `error_kind` and + // `provider_request_id`. + let counted = total > 0; + aisix_obs::AccessLog { + method: &self.method, + path: &self.path, + status: event.status_code, + latency, + duration, + provider: (!self.provider.is_empty()).then_some(self.provider.as_str()), + model: (!self.model.is_empty()).then_some(self.model.as_str()), + upstream_model: (!target.upstream_model.is_empty()) + .then_some(target.upstream_model.as_str()), + provider_key_id: (!target.provider_key_id.is_empty()) + .then_some(target.provider_key_id.as_str()), + api_key_id: (!self.api_key_id.is_empty()).then_some(self.api_key_id.as_str()), + prompt_tokens: counted.then_some(prompt), + completion_tokens: counted.then_some(completion), + total_tokens: counted.then_some(total), + request_id: &self.request_id, + provider_request_id: (!event.provider_request_id.is_empty()) + .then_some(event.provider_request_id.as_str()), + served_by_model: (!self.served_by_model.is_empty()) + .then_some(self.served_by_model.as_str()), + routing_attempt_count: self.routing_attempt_count, + routing_fallback_count: self.routing_fallback_count, + error_kind: (!event.error_class.is_empty()).then_some(event.error_class.as_str()), + error: (!event.error_message.is_empty()).then_some(event.error_message.as_str()), + mcp: None, + } + .emit(); + } +} + #[derive(Default)] struct Cell { resolved: Resolved, cancel: CancelContext, + /// See [`PendingAccessLog`]. `Some` only between a streaming handler + /// returning and the request's terminal usage event going out. + pending_log: Option, + /// See [`note_stream_owns_access_log`]. + stream_owns_log: bool, } /// The per-request cell. Attempts within a request are sequential, so the @@ -232,6 +396,30 @@ impl RequestAttribution { std::mem::take(&mut self.lock().cancel) } + /// Whether this request still has a parked line — i.e. whether a + /// terminal emitter is still owed one. Read by the cancel guard before + /// it builds a line of its own. + pub(crate) fn has_pending_access_log(&self) -> bool { + self.lock().pending_log.is_some() + } + + /// Emit the request's deferred line, if it still has one, against the + /// outcome `event` reports. Returns whether a line went out. + /// + /// Read through the cell handle rather than the task-local because the + /// cancel guard holds the handle and runs from `Drop`, outside every + /// scope. + pub(crate) fn emit_deferred_access_log(&self, event: &aisix_obs::UsageEvent) -> bool { + let mut cell = self.lock(); + let Some(pending) = cell.pending_log.take() else { + return false; + }; + let target = cell.resolved.clone(); + drop(cell); + pending.emit(&target, event); + true + } + fn lock(&self) -> std::sync::MutexGuard<'_, Cell> { self.0 .lock() @@ -478,6 +666,37 @@ pub(crate) fn note_attempt_settled(rec: &AttemptRecord) { }); } +/// Park this request's access-log line until its outcome is known. Called +/// by a streaming handler in place of writing the line itself — see +/// [`PendingAccessLog`]. +pub(crate) fn defer_access_log(pending: PendingAccessLog) { + let _ = CURRENT.try_with(|a| a.lock().pending_log = Some(pending)); +} + +/// Note that this request answered with a stream whose own terminal +/// emitter will write the access-log line. +/// +/// For the handlers that wrap their whole dispatch and log the wrapper's +/// status (`/a2a`): the streaming branch is several frames below the tail +/// that owns the line's fields, so it raises a flag there and the tail +/// parks the line instead of writing it. +pub(crate) fn note_stream_owns_access_log() { + let _ = CURRENT.try_with(|a| a.lock().stream_owns_log = true); +} + +/// Whether [`note_stream_owns_access_log`] was raised on this request. +pub(crate) fn stream_owns_access_log() -> bool { + CURRENT + .try_with(|a| a.lock().stream_owns_log) + .unwrap_or(false) +} + +/// Emit the deferred line of the request running on this task, if it has +/// one. Called from the terminal-usage-event chokepoint. +pub(crate) fn emit_deferred_access_log(event: &aisix_obs::UsageEvent) { + let _ = CURRENT.try_with(|a| a.emit_deferred_access_log(event)); +} + /// Note that a usage event has just left the emission chokepoint on this /// request's own task. See [`CancelContext::emitted_any`]. pub(crate) fn note_usage_emitted(terminal: bool) { diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index 124007e1..b04535d7 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -150,17 +150,41 @@ pub async fn transcriptions( // Actual status, not a hardcoded 200 — the #696 billed-then- // output-blocked path returns Ok(success) carrying a 422. let status = success.response.status().as_u16(); - emit_access_log( - "POST", - "/v1/audio/transcriptions", - &success.model_name, - &success.provider, - &api_key_id, - status, - elapsed, - &request_id, - None, - ); + // On this family the flag IS "the response is a live relay" — it + // is set only inside the `is_event_stream` branch and is what + // labels the metric as streaming — so there is no second + // predicate to conjoin, unlike `/v1/messages` and + // `/v1/responses`. If it ever comes to mean "already emitted" + // too, park on the relay itself instead: a parked line with no + // later emitter is a line silently lost. + if success.usage_handled_by_stream { + // A relayed transcription stream has no outcome yet — the caller may + // read it to the terminal event or walk away. Park the line and + // let the relay's own Drop emitter write it beside the usage + // event it already owns (AISIX-Cloud#1571). + crate::attribution::defer_access_log( + crate::attribution::PendingAccessLog::new( + "POST", + "/v1/audio/transcriptions", + &request_id, + &api_key_id, + started, + ) + .with_model(&success.provider, &success.model_name), + ); + } else { + emit_access_log( + "POST", + "/v1/audio/transcriptions", + &success.model_name, + &success.provider, + &api_key_id, + status, + elapsed, + &request_id, + None, + ); + } // ONE ProviderKey lookup for both terminal emits (#941). let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &success.provider_key_id); record_audio_metrics( @@ -314,17 +338,34 @@ pub async fn translations( // Actual status, not a hardcoded 200 — the #696 billed-then- // output-blocked path returns Ok(success) carrying a 422. let status = success.response.status().as_u16(); - emit_access_log( - "POST", - "/v1/audio/translations", - &success.model_name, - &success.provider, - &api_key_id, - status, - elapsed, - &request_id, - None, - ); + if success.usage_handled_by_stream { + // A relayed transcription stream has no outcome yet — the caller may + // read it to the terminal event or walk away. Park the line and + // let the relay's own Drop emitter write it beside the usage + // event it already owns (AISIX-Cloud#1571). + crate::attribution::defer_access_log( + crate::attribution::PendingAccessLog::new( + "POST", + "/v1/audio/translations", + &request_id, + &api_key_id, + started, + ) + .with_model(&success.provider, &success.model_name), + ); + } else { + emit_access_log( + "POST", + "/v1/audio/translations", + &success.model_name, + &success.provider, + &api_key_id, + status, + elapsed, + &request_id, + None, + ); + } // ONE ProviderKey lookup for both terminal emits (#941). let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &success.provider_key_id); record_audio_metrics( @@ -2263,6 +2304,7 @@ fn emit_access_log( path, status, latency, + duration: latency, provider: Some(provider), model: Some(model), upstream_model: target.upstream_model(), @@ -4200,4 +4242,50 @@ data: [DONE]\n\n"; let wire = serde_json::to_string(&ev).unwrap(); assert!(!wire.contains("9.9.9"), "{wire}"); } + + /// A streamed transcription relay writes ONE access-log line, at the + /// relay's end rather than when the head went out, so each of the three + /// endings reports its own outcome (AISIX-Cloud#1571). + /// + /// `latency_ms` is deliberately NOT asserted to be a time-to-first-frame + /// here: this relay's usage event reports the WHOLE relay as what the + /// caller waited for, and the line reports the same figure the event + /// does. Changing that would be a change to the usage event's meaning, + /// not to this line. + #[tokio::test] + async fn a_streamed_transcription_writes_one_line_per_stream_ending() { + let upstream = crate::test_log::spawn_sse_upstream(vec![ + "data: {\"type\":\"transcript.text.delta\",\"delta\":\"hello\"}\n\n".to_string(), + "data: {\"type\":\"transcript.text.delta\",\"delta\":\" world\"}\n\n".to_string(), + "data: {\"type\":\"transcript.text.done\",\"text\":\"hello world\",\ + \"usage\":{\"type\":\"tokens\",\"total_tokens\":38,\"input_tokens\":26,\ + \"output_tokens\":12}}\n\n" + .to_string(), + "data: [DONE]\n\n".to_string(), + ]) + .await; + + let snap = new_snap(&upstream); + snap.models.insert(whisper_model("my-transcribe")); + snap.apikeys.insert(apikey_entry(&["*"])); + let app = build_app(snap); + + let endings = crate::test_log::three_stream_endings(app, || { + let (ct, body) = streaming_transcription_multipart("my-transcribe"); + Request::builder() + .method("POST") + .uri("/v1/audio/transcriptions") + .header("authorization", "Bearer sk-caller") + .header("content-type", ct) + .body(body) + .unwrap() + }) + .await; + crate::test_log::assert_one_line_per_ending(&endings, "/v1/audio/transcriptions", "k-1"); + assert_eq!( + endings.delivered.num("total_tokens"), + Some(38), + "the terminal frame's counts belong on the line that reports the relay's end", + ); + } } diff --git a/crates/aisix-proxy/src/cancel.rs b/crates/aisix-proxy/src/cancel.rs index aab4cb98..20d6b081 100644 --- a/crates/aisix-proxy/src/cancel.rs +++ b/crates/aisix-proxy/src/cancel.rs @@ -44,8 +44,8 @@ use crate::client_ip::ClientContext; use crate::state::ProxyState; use crate::usage_attr::{self, ResolvedPk}; -/// `error_message` of the terminal event, and of the access-log line the -/// guard writes beside it — one sentence, the same on both, so a row in the +/// `error_message` of the terminal event, and of the access-log line that +/// goes out beside it — one sentence, the same on both, so a row in the /// usage log can be joined to the line that explains it. pub(crate) const CANCELLED_BEFORE_HEAD: &str = "client closed the request before the response head was written"; @@ -64,11 +64,9 @@ pub(crate) const CANCELLED_MID_STREAM: &str = /// It is its own phase because nothing else speaks for it. A streaming /// family's own terminal emitter lives in a `Drop` guard built INSIDE the /// stream's generator, and a generator first runs on the body's first poll -/// — so a body dropped before that emits no usage row at all -/// (AISIX-Cloud#1571). The request's access-log line is not missing: its -/// handler wrote one when it handed the stream over, saying `200`. This -/// phase therefore writes the row and nothing else — a second line under a -/// second status would make one request read as two. +/// — so a body dropped before that emits neither the usage row nor the +/// access-log line its handler parked on the cell (AISIX-Cloud#1571). This +/// phase writes both, under this message. pub(crate) const CANCELLED_BEFORE_BODY: &str = "client closed the request before the response body was streamed"; @@ -84,8 +82,9 @@ pub(crate) enum Phase { } impl Phase { - /// The `error_message` the terminal event carries, and — on the head - /// phase, the one phase that writes a line — the line's too. + /// The `error_message` the terminal event carries, and the line's too — + /// the line rides that event out of [`usage_attr::emit_usage`] on the + /// body phase, and is built beside it on the head phase. pub(crate) fn message(self) -> &'static str { match self { Phase::BeforeHead => CANCELLED_BEFORE_HEAD, diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index e4ccb2da..55fcb4ba 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -181,26 +181,46 @@ pub async fn chat_completions( &success, elapsed, ); - emit_access_log( - method, - path, - status, - elapsed, - Some(success.provider.as_str()), - Some(&model_name), - Some(&api_key_id), - success.prompt_tokens, - success.completion_tokens, - success.total_tokens, - &request_id, - // Empty on the streaming path — the id rides the first - // upstream frame, which has not arrived yet. That case is - // covered by the per-attempt `provider call completed` line - // the usage sink emits (AISIX-Cloud#1289). - Some(success.provider_request_id.as_str()), - &success.routing, - None, - ); + // `telemetry_handled_by_stream` alone is NOT "the response is a + // stream": the BUFFERED ensemble path sets it too, to mean "the + // sub-call emits already covered this request". That one has no + // later emitter to write a parked line, so the conjunction is + // what keeps its line from disappearing. + if req.is_streaming() && success.telemetry_handled_by_stream { + // A streamed response has no outcome yet: the head exists, nothing + // has been delivered, and whether the caller reads it to the end + // or walks away is minutes from being known. Park the line and + // let whichever terminal emitter ends the request write it, with + // that emitter's status, tokens and message (AISIX-Cloud#1571). + crate::attribution::defer_access_log( + crate::attribution::PendingAccessLog::new( + method, + path, + &request_id, + &api_key_id, + started, + ) + .with_model(&success.provider, &model_name) + .with_routing(&success.routing), + ); + } else { + emit_access_log( + method, + path, + status, + elapsed, + Some(success.provider.as_str()), + Some(&model_name), + Some(&api_key_id), + success.prompt_tokens, + success.completion_tokens, + success.total_tokens, + &request_id, + Some(success.provider_request_id.as_str()), + &success.routing, + None, + ); + } // Per #655: emit a zero-token event for each failed attempt // that preceded the winner (non-streaming fallover). No-op for // direct-model success, cache hits, and the single-attempt @@ -5020,6 +5040,7 @@ fn emit_access_log( path, status, latency, + duration: latency, provider, model, upstream_model: target.upstream_model(), diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index df95a3c8..6cf851e8 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -910,6 +910,7 @@ fn emit_access_log( path: "/v1/completions", status, latency, + duration: latency, provider: Some(provider), model: Some(model), upstream_model: target.upstream_model(), @@ -1901,4 +1902,50 @@ mod tests { let wire = serde_json::to_string(&ev).unwrap(); assert!(!wire.contains("9.9.9"), "{wire}"); } + + /// `/v1/completions` refuses `stream: true`, so it has no stream to + /// defer its line to and writes it where it always did — at the handler + /// tail. What AISIX-Cloud#1571 adds here is the second figure, and on a + /// buffered request the two are the same number: the caller waited for + /// the whole response, which is the whole request. + #[tokio::test] + async fn a_buffered_request_reports_one_line_whose_duration_is_its_latency() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "cmpl-abc", + "object": "text_completion", + "created": 1_700_000_000i64, + "model": "gpt-3.5-turbo-instruct", + "choices": [{"text": " is a test", "index": 0, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 4, "total_tokens": 9} + }))) + .mount(&upstream) + .await; + + let snap = new_snap(&upstream.uri()); + snap.models.insert(model_entry("instruct")); + snap.apikeys.insert(apikey_entry(&["*"])); + let app = build_app(snap); + + let capture = crate::test_log::Capture::install(); + let resp = tower::ServiceExt::oneshot( + app, + make_req(serde_json::json!({"model": "instruct", "prompt": "Say this"})), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let _ = to_bytes(resp.into_body(), 65536).await.unwrap(); + + let line = capture.only("a buffered request"); + assert_eq!(line.status(), 200); + assert_eq!(line.field("path").as_deref(), Some("/v1/completions")); + assert_eq!( + line.num("duration_ms"), + line.num("latency_ms"), + "nothing is streamed here, so the wait and the request are the same span", + ); + } } diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index 3afe1e06..130dc39a 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -901,6 +901,7 @@ fn emit_access_log( path: "/v1/messages/count_tokens", status, latency: elapsed, + duration: elapsed, provider: Some(provider), model: Some(model), upstream_model: target.upstream_model(), diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index 66bdf6bc..06838b71 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -603,6 +603,7 @@ fn emit_access_log( path: "/v1/embeddings", status, latency, + duration: latency, provider: Some(provider), model: Some(model), upstream_model: target.upstream_model(), diff --git a/crates/aisix-proxy/src/images.rs b/crates/aisix-proxy/src/images.rs index 5c6d0d91..f74e2956 100644 --- a/crates/aisix-proxy/src/images.rs +++ b/crates/aisix-proxy/src/images.rs @@ -653,6 +653,7 @@ pub(crate) fn emit_access_log( path: endpoint, status, latency, + duration: latency, provider: Some(provider), model: Some(model), upstream_model: target.upstream_model(), diff --git a/crates/aisix-proxy/src/jobs.rs b/crates/aisix-proxy/src/jobs.rs index 4a63120d..8dbd4ace 100644 --- a/crates/aisix-proxy/src/jobs.rs +++ b/crates/aisix-proxy/src/jobs.rs @@ -708,6 +708,7 @@ fn emit_access_log( path, status, latency: elapsed, + duration: elapsed, provider: target.map(|t| t.provider_label()).filter(|p| !p.is_empty()), model: target.map(|t| t.display_name()), upstream_model: log_target.upstream_model(), diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 1dcd79e6..32066aae 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -84,6 +84,8 @@ mod semantic; pub mod sse_keepalive; mod state; mod stream_timeout; +#[cfg(test)] +mod test_log; mod token_estimate; mod usage_attr; /// The `model` metric label for a request that resolved no model. Exported @@ -533,14 +535,13 @@ const CLIENT_DISCONNECTED_KIND: &str = "client_disconnected"; /// before it emits nothing anywhere. The guard rides the body precisely /// to cover that window — see `GuardPhase` and `TelemetryBody`. /// -/// All three shapes report `499` on the usage event with `error_class = +/// All three shapes report `499` with `error_class = /// "client_disconnected"`; only the message says where the caller left. -/// A streamed request's LINE keeps saying `200` in all of them, because -/// its handler writes it when the head goes out — one line per request, -/// whose latency is time-to-first-token by deliberate choice -/// (AISIX-Cloud#1394). Converging the two onto one end-of-stream line is -/// per-family work in each streaming handler, not something this layer -/// can do for them. +/// The LINE says the same, because a streaming handler does not write it +/// at all: it parks it on the request's cell +/// (`attribution::PendingAccessLog`) and whichever terminal emitter ends +/// the request writes it, with that emitter's status and message. One +/// line per request, in every ending. async fn record_request_telemetry( State(state): State, request: Request, @@ -695,8 +696,19 @@ impl axum::body::HttpBody for TelemetryBody { if let GuardPhase::Body { polled, .. } = &mut this.guard.phase { *polled = true; } + // Poll the stream inside the request's own attribution cell. The + // body is driven by the server long after the middleware returned, + // so without this the generator's end-of-stream emitter — which is + // where a streamed request's terminal usage event AND its + // access-log line go out (AISIX-Cloud#1571) — would run with no + // cell to read the parked line from, and a delivered stream would + // log nothing at all. `Drop` covers the abandoned endings; this + // covers the delivered one. + let cell = &this.guard.attribution; match this.inner.as_mut() { - Some(inner) => std::pin::Pin::new(inner).poll_frame(cx), + Some(inner) => { + attribution::sync_scope(cell, || std::pin::Pin::new(inner).poll_frame(cx)) + } None => std::task::Poll::Ready(None), } } @@ -855,21 +867,34 @@ impl Drop for ClientCancelGuard { if matches!(phase, cancel::Phase::BeforeBody) && cancel_ctx.emitted_terminal { return; } - // The LINE is the head phase's alone. A request whose head went out - // already has its handler's line, written when the stream was - // handed over; adding a second one here under a different status - // would make one request two, which is the opposite of the picture - // this change is for. Converging the streamed families on a single - // end-of-stream line is real work in each of them and reverses a - // deliberate decision about what a streamed line's latency means - // (AISIX-Cloud#1394) — it is not this change's to make. - if matches!(phase, cancel::Phase::BeforeHead) { + // The head phase builds its OWN line, because the handler normally + // never reached the tail that would have parked one. The body phase + // does not: a streamed response left its line on the cell, and that + // line rides the terminal usage event below, carrying the same + // `499` and the same message. Building a second one here would make + // one request read as two. + // + // "Normally" is why the head phase asks as well. A streaming family + // parks its line at its tail and can still be cancelled at the next + // await — chat peeks the rate limiter there, to fill the + // `x-ratelimit-*` headers — which lands here with the line already + // parked. That line is the fuller one (it names the model, the + // target and the routing counts) and `cancel::emit` below writes it + // under this same `499`, so this one stands down. It cannot fall + // between the two: a parked line means the request authenticated on + // a metering surface, which is exactly the gate `cancel::emit` + // applies before it emits the terminal event that carries the line. + if matches!(phase, cancel::Phase::BeforeHead) && !self.attribution.has_pending_access_log() + { let target = attribution::AccessLogTarget::from_resolved(resolved.clone()); AccessLog { method: self.method.as_str(), path: self.uri.path(), status: CLIENT_CLOSED_REQUEST, latency, + // Nothing was ever delivered, so what the caller waited for + // IS how long the request ran. + duration: latency, // The log line takes the RAW names: it is bounded by request // volume, not by label cardinality, so it can say exactly // which target the abandoned request was waiting on. @@ -895,17 +920,22 @@ impl Drop for ClientCancelGuard { .emit(); } // The usage events the dropped handler never got to write - // (AISIX-Cloud#1571). After the line, so the two land in the order - // an operator reads them. - cancel::emit( - &self.state, - self.endpoint, - &self.request_id, - &resolved, - cancel_ctx, - phase, - self.trace.as_ref(), - ); + // (AISIX-Cloud#1571) — and, on the body phase, the request's parked + // access-log line, which goes out of the same chokepoint as the + // terminal event so the two agree on the outcome. `Drop` runs + // outside every scope, so the cell has to be installed for the call + // or the chokepoint has nothing to take the line from. + attribution::sync_scope(&self.attribution, || { + cancel::emit( + &self.state, + self.endpoint, + &self.request_id, + &resolved, + cancel_ctx, + phase, + self.trace.as_ref(), + ) + }); // Bound the labels the same way every other emit does: the model // through the configured set, the ProviderKey name off the row its // id names — a cancelled request must not be able to mint series @@ -9050,7 +9080,7 @@ data: [DONE]\n\n", let hub = Arc::new(Hub::new()); hub.register_specialized("openai", Arc::new(openai_test_bridge())); let snap = seed_routing_group("smart", &[("m-primary", "primary", &upstream.uri())]); - let (tx, _rx) = tokio::sync::mpsc::channel(32); + let (tx, mut rx) = tokio::sync::mpsc::channel(32); let state = build_state(snap, hub).with_usage_sink(UsageSink::new(tx)); let app = build_router(state); @@ -9067,42 +9097,76 @@ data: [DONE]\n\n", req }; - // 1. Read to the end. - { - let (buf, _capture) = access_log_capture(); - let response = app.clone().oneshot(streaming_request()).await.unwrap(); - let _ = to_bytes(response.into_body(), 65536).await.unwrap(); + let endings = crate::test_log::three_stream_endings(app.clone(), streaming_request).await; + crate::test_log::assert_one_line_per_ending(&endings, "/v1/chat/completions", "key-id-1"); + crate::test_log::assert_latency_is_time_to_first_token(&endings); + + // The line and the row are ONE record of ONE request, so they agree + // on the outcome down to the sentence. They can only disagree if the + // line is written somewhere other than the terminal emit — which is + // exactly what writing it at the handler tail was. + for (what, line) in [ + ("a delivered stream", &endings.delivered), + ("a stream abandoned mid-flight", &endings.abandoned), + ("a stream dropped before its first poll", &endings.unread), + ] { + let event = next_event(&mut rx).await; assert_eq!( - access_log_lines(&buf), - 1, - "a delivered stream must write one line", + u64::from(event.status_code), + line.status(), + "{what}: the line and the usage event disagree on the status", ); - } - - // 2. Read one chunk, then walk away mid-stream. - { - let (buf, _capture) = access_log_capture(); - let response = app.clone().oneshot(streaming_request()).await.unwrap(); - let mut body = response.into_body().into_data_stream(); - let _first = futures::StreamExt::next(&mut body).await; - drop(body); assert_eq!( - access_log_lines(&buf), - 1, - "a stream abandoned mid-flight must write one line", + line.field("error").unwrap_or_default(), + event.error_message, + "{what}: the line and the usage event disagree on why", + ); + assert_eq!( + line.field("error_kind").unwrap_or_default(), + event.error_class, + "{what}: the line and the usage event disagree on the class", ); } - // 3. Never read it at all — the window this change covers. - { - let (buf, _capture) = access_log_capture(); - let response = app.clone().oneshot(streaming_request()).await.unwrap(); - drop_body_unpolled(response); + // The two abandoned endings are different phases and say so — the + // one message an operator reads to tell "left while it was + // streaming" from "never read a byte of it" apart. + assert_eq!( + endings.abandoned.field("error").as_deref(), + Some(cancel::CANCELLED_MID_STREAM), + ); + assert_eq!( + endings.unread.field("error").as_deref(), + Some(cancel::CANCELLED_BEFORE_BODY), + ); + + // The delivered line carries what only the stream's END knows: the + // upstream's response id and the token counts. At head time neither + // existed, which is why the old line had to leave them out. + assert_eq!( + endings.delivered.field("provider_request_id").as_deref(), + Some("cmpl-body"), + ); + assert!( + endings.delivered.num("total_tokens").is_some(), + "a delivered stream's line must carry the counts its event billed", + ); + // And the target it dispatched to, on every ending — a routing + // group's own name answers "which member served this" nowhere. + for (what, line) in [ + ("a delivered stream", &endings.delivered), + ("a stream abandoned mid-flight", &endings.abandoned), + ("a stream dropped before its first poll", &endings.unread), + ] { + assert_eq!(line.field("model").as_deref(), Some("smart"), "{what}"); assert_eq!( - access_log_lines(&buf), - 1, - "a stream dropped before its first poll must write one line, not a second one \ - under a different status", + line.field("upstream_model").as_deref(), + Some("gpt-4o"), + "{what}: the dispatched target is missing", + ); + assert!( + line.field("provider_key_id").is_some(), + "{what}: the ProviderKey that served is missing", ); } } @@ -9350,58 +9414,6 @@ data: [DONE]\n\n", assert_eq!(event.operation, "embeddings"); } - /// A tracing writer that appends every emitted byte into a shared buffer. - #[derive(Clone)] - struct LogBuf(Arc>>); - impl std::io::Write for LogBuf { - fn write(&mut self, buf: &[u8]) -> std::io::Result { - self.0.lock().unwrap().extend_from_slice(buf); - Ok(buf.len()) - } - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } - } - impl tracing_subscriber::fmt::MakeWriter<'_> for LogBuf { - type Writer = LogBuf; - fn make_writer(&self) -> Self::Writer { - self.clone() - } - } - - /// The access log's own `tracing` message. Counting occurrences of it is - /// how a test asks "how many lines did this request write" without - /// matching the other events the same subscriber sees. - const ACCESS_LOG_MESSAGE: &str = "proxy request completed"; - - /// Install a capturing subscriber on THIS thread and hand back the - /// buffer plus its guard, so a caller can hold it across awaits — a - /// `#[tokio::test]` runs its future on the calling thread, which is - /// where the handler's own line is written. - fn access_log_capture() -> ( - Arc>>, - tracing::subscriber::DefaultGuard, - ) { - static ONCE: std::sync::Once = std::sync::Once::new(); - ONCE.call_once(|| { - let _ = tracing::subscriber::set_global_default(tracing_subscriber::registry()); - }); - let buf = Arc::new(std::sync::Mutex::new(Vec::new())); - let subscriber = tracing_subscriber::fmt() - .with_ansi(false) - .with_writer(LogBuf(buf.clone())) - .finish(); - let guard = tracing::subscriber::set_default(subscriber); - (buf, guard) - } - - /// How many access-log lines the capture holds. - fn access_log_lines(buf: &Arc>>) -> usize { - String::from_utf8_lossy(&buf.lock().unwrap()) - .matches(ACCESS_LOG_MESSAGE) - .count() - } - /// A passthrough route names no model at all, so the row a cancelled one /// files is attributed by the ROUTE. Without that the request appears in /// the usage log as an anonymous `499` an operator cannot trace back to @@ -9500,6 +9512,76 @@ data: [DONE]\n\n", ); } + /// A cancel that lands AFTER the handler parked this request's line but + /// before it returned writes that line — not a second one beside it. + /// + /// The window is real rather than theoretical: a streaming family parks + /// its line at its tail and chat then awaits once more, peeking the rate + /// limiter to fill the `x-ratelimit-*` headers. A caller that hangs up + /// there leaves the guard in its HEAD phase with the line already on the + /// cell, and both emitters would speak — under the same `499`, with the + /// same message, so one request would read as two identical ones and a + /// count of `499` lines would double. + #[tokio::test] + async fn a_head_phase_cancel_writes_the_parked_line_instead_of_a_second_one() { + use aisix_obs::UsageSink; + + let snap = seed_snapshot("my-gpt4", &["my-gpt4"], "http://unused"); + let (tx, mut rx) = tokio::sync::mpsc::channel(4); + let state = build_state(snap, Arc::new(Hub::new())).with_usage_sink(UsageSink::new(tx)); + let cell = std::sync::Arc::new(attribution::RequestAttribution::default()); + + // What a streaming handler leaves behind on its way out: the caller + // it authenticated, and its line. + attribution::sync_scope(&cell, || { + attribution::note_client(&crate::client_ip::ClientContext::default(), "key-id-1"); + attribution::defer_access_log( + attribution::PendingAccessLog::new( + "POST", + "/v1/chat/completions", + "req-parked", + "key-id-1", + std::time::Instant::now(), + ) + .with_model("openai", "my-gpt4"), + ); + }); + + let capture = crate::test_log::Capture::install(); + drop(ClientCancelGuard { + phase: GuardPhase::Head, + state: state.clone(), + attribution: cell, + endpoint: "/v1/chat/completions", + method: axum::http::Method::POST, + uri: "/v1/chat/completions".parse().unwrap(), + request_id: "req-parked".to_string(), + trace: None, + started: std::time::Instant::now(), + }); + + let line = capture.only("a head-phase cancel with a parked line"); + assert_eq!(line.status(), u64::from(CLIENT_CLOSED_REQUEST)); + assert_eq!( + line.field("error_kind").as_deref(), + Some(CLIENT_DISCONNECTED_KIND), + ); + // The PARKED line is the one that went out — the guard's own names + // no model, because nothing resolved one into the cell here. + assert_eq!( + line.field("model").as_deref(), + Some("my-gpt4"), + "the guard wrote its own, thinner line instead of the parked one", + ); + let event = next_event(&mut rx).await; + assert_eq!(event.status_code, CLIENT_CLOSED_REQUEST); + assert_eq!( + u64::from(event.status_code), + line.status(), + "one record, one outcome", + ); + } + /// A panicking handler drops the guard mid-unwind still in its head /// phase, which looks identical to a cancel from `Drop`'s point of view. /// Recording it would invent a client disconnect that never happened and diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs index ac93827f..f3cb6cac 100644 --- a/crates/aisix-proxy/src/mcp.rs +++ b/crates/aisix-proxy/src/mcp.rs @@ -255,6 +255,7 @@ async fn serve(state: ProxyState, request: Request, scope: Option) -> Re path: endpoint, status, latency: elapsed, + duration: elapsed, provider: Some("mcp"), model: None, upstream_model: target.upstream_model(), diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 3c58db36..0615e777 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -180,21 +180,41 @@ pub async fn messages( monitor_hits.extend(output_monitor_hits); let elapsed = started.elapsed(); let status = response.status().as_u16(); - emit_access_log( - &model_name, - &provider_label, - &api_key_id, - status, - elapsed, - &request_id, - // Empty on the streaming path — the id rides the - // `message_start` frame, which has not arrived yet. That case - // is covered by the per-attempt `provider call completed` - // line the usage sink emits (AISIX-Cloud#1289). - Some(metrics.provider_request_id.as_str()), - &routing, - None, - ); + // Conjoined with the request's own streaming flag rather than + // resting on `usage_handled_by_stream` alone: a family that ever + // reuses that flag to mean "already emitted" on a BUFFERED path, + // the way chat's ensemble does, would park a line with no later + // emitter to write it — and lose it silently. + if stream_requested && usage_handled_by_stream { + // A streamed response has no outcome yet: the head exists, nothing + // has been delivered, and whether the caller reads it to the end + // or walks away is minutes from being known. Park the line and + // let whichever terminal emitter ends the request write it, with + // that emitter's status, tokens and message (AISIX-Cloud#1571). + crate::attribution::defer_access_log( + crate::attribution::PendingAccessLog::new( + "POST", + "/v1/messages", + &request_id, + &api_key_id, + started, + ) + .with_model(&provider_label, &model_name) + .with_routing(&routing), + ); + } else { + emit_access_log( + &model_name, + &provider_label, + &api_key_id, + status, + elapsed, + &request_id, + Some(metrics.provider_request_id.as_str()), + &routing, + None, + ); + } // ONE ProviderKey lookup for both the metric emit and the // winner's usage event below (#941). let pk = crate::usage_attr::ResolvedPk::resolve(&snapshot, &provider_key_id); @@ -4204,6 +4224,7 @@ fn emit_access_log( path: "/v1/messages", status, latency, + duration: latency, provider: Some(provider), model: Some(model), upstream_model: target.upstream_model(), @@ -5620,6 +5641,9 @@ event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; "max_tokens": 100, "stream": true, }); + // The access-log line goes out with the terminal usage event now + // (AISIX-Cloud#1571), so it is captured for the same request. + let capture = crate::test_log::Capture::install(); let resp = app.oneshot(make_req(body)).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); @@ -5663,6 +5687,29 @@ event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; "streaming /v1/messages telemetry must record TTFT", ); assert!(rx.try_recv().is_err(), "usage event should be emitted once"); + + // AISIX-Cloud#1571: the line is written beside that event, so its + // token columns have to be the total the row bills. On this path + // `prompt_tokens` EXCLUDES the two cache dimensions, so a line that + // summed only the two visible columns would report 89 where the row + // bills 102 — and the two are supposed to be one record. + let line = capture.only("a streamed /v1/messages call"); + assert_eq!(line.status(), 200); + assert_eq!( + line.num("total_tokens"), + Some(u64::from( + event.prompt_tokens + + event.completion_tokens + + event.cache_creation_tokens + + event.cache_read_tokens + )), + "the line and the row must agree on what the request cost", + ); + assert_eq!(line.num("total_tokens"), Some(102)); + assert_eq!( + line.field("provider_request_id").as_deref(), + Some("msg_stream_245") + ); } /// AISIX-Cloud#952: relay backends that ship NO usage on @@ -6994,4 +7041,62 @@ data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text crate::error::anthropic_kind_from_status(axum::http::StatusCode::UNPROCESSABLE_ENTITY), ); } + + /// A streamed `/v1/messages` response writes ONE access-log line, at the + /// stream's end rather than when the head went out — so each of the + /// three endings a stream has reports its own outcome + /// (AISIX-Cloud#1571). Written at the handler tail, all three said + /// `200`, including the two where the caller was already gone and the + /// request's own usage event said `499`. + #[tokio::test] + async fn a_streamed_request_writes_one_line_per_stream_ending() { + use aisix_provider_openai::OpenAiBridge; + + let upstream = MockServer::start().await; + let sse = "\ +data: {\"id\":\"cmpl-1\",\"object\":\"chat.completion.chunk\",\"created\":1715000000,\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\n\ +data: {\"id\":\"cmpl-1\",\"object\":\"chat.completion.chunk\",\"created\":1715000000,\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hel\"},\"finish_reason\":null}]}\n\n\ +data: {\"id\":\"cmpl-1\",\"object\":\"chat.completion.chunk\",\"created\":1715000000,\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"lo\"},\"finish_reason\":\"stop\"}]}\n\n\ +data: [DONE]\n\n"; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse), + ) + .mount(&upstream) + .await; + + let snap = new_snap_openai(&upstream.uri()); + snap.models.insert(openai_model("my-claude-alias")); + snap.apikeys.insert(apikey_entry(&["*"])); + + let hub = Arc::new(Hub::new()); + hub.register_specialized("anthropic", Arc::new(AnthropicBridge::new())); + hub.register_specialized("openai", Arc::new(OpenAiBridge::new())); + let handle = SnapshotHandle::new(snap); + let app = crate::build_router(crate::ProxyState::new(handle, hub, &cfg()).without_cache()); + + let endings = crate::test_log::three_stream_endings(app, || { + make_req(serde_json::json!({ + "model": "my-claude-alias", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 100, + "stream": true, + })) + }) + .await; + crate::test_log::assert_one_line_per_ending(&endings, "/v1/messages", "k-1"); + crate::test_log::assert_latency_is_time_to_first_token(&endings); + assert_eq!( + endings.delivered.field("model").as_deref(), + Some("my-claude-alias"), + ); + assert_eq!( + endings.abandoned.field("upstream_model").as_deref(), + Some("gpt-4o"), + "an abandoned stream must still name the target it was dispatched to", + ); + } } diff --git a/crates/aisix-proxy/src/passthrough_route.rs b/crates/aisix-proxy/src/passthrough_route.rs index d24f5af2..577942ce 100644 --- a/crates/aisix-proxy/src/passthrough_route.rs +++ b/crates/aisix-proxy/src/passthrough_route.rs @@ -314,6 +314,7 @@ pub async fn entry( api_key_id, status, elapsed, + elapsed, &request_id, None, Some(&error), @@ -2220,6 +2221,18 @@ impl RouteTelemetry { &self.route_name, &self.api_key_id, self.status, + // Same rule as the typed streaming endpoints, and the same + // figure this emit puts on the usage event below: a streamed + // relay reports the wait to its first relayed frame, a buffered + // one the whole response. A relay that delivered nothing waited + // the whole request for nothing, which is what `elapsed` says. + if self.streaming { + self.downstream_first_ms + .map(|ms| Duration::from_millis(u64::from(ms))) + .unwrap_or(elapsed) + } else { + elapsed + }, elapsed, &self.request_id, Some(AccessLogTokens { @@ -2381,7 +2394,13 @@ fn emit_access_log( route: &str, api_key_id: &str, status: u16, - elapsed: Duration, + // What the caller waited for: the first relayed frame on a streamed + // relay, the whole response otherwise — the same figure the usage + // event reports as `downstream_latency_ms`. + latency: Duration, + // How long the relay held the gateway, arrival to last byte out. On a + // streamed relay the two differ by the length of the stream. + duration: Duration, request_id: &str, tokens: Option, error: Option<&ProxyError>, @@ -2398,7 +2417,8 @@ fn emit_access_log( method: method.as_str(), path, status, - latency: elapsed, + latency, + duration, provider: Some(route), model: None, upstream_model: target.upstream_model(), diff --git a/crates/aisix-proxy/src/realtime.rs b/crates/aisix-proxy/src/realtime.rs index 9893867a..0e60a191 100644 --- a/crates/aisix-proxy/src/realtime.rs +++ b/crates/aisix-proxy/src/realtime.rs @@ -1152,6 +1152,7 @@ fn emit_access_log( path: "/v1/realtime", status, latency: elapsed, + duration: elapsed, provider: target.map(|(p, _)| p).filter(|p| !p.is_empty()), model: target.map(|(_, m)| m), upstream_model: log_target.upstream_model(), diff --git a/crates/aisix-proxy/src/reject.rs b/crates/aisix-proxy/src/reject.rs index 1b2408b8..0794f09d 100644 --- a/crates/aisix-proxy/src/reject.rs +++ b/crates/aisix-proxy/src/reject.rs @@ -69,6 +69,7 @@ pub(crate) fn reject_before_dispatch( path, status, latency: elapsed, + duration: elapsed, // Nothing is resolved this early: no upstream was picked, and the // body naming the model is exactly what we refused to read. provider: None, diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 3f2877c1..4b314192 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -854,6 +854,7 @@ fn emit_access_log( path: "/v1/rerank", status, latency: elapsed, + duration: elapsed, provider: Some(provider), model: Some(model), upstream_model: target.upstream_model(), diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 3b2b56bb..1af4503c 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -245,24 +245,41 @@ pub async fn responses( monitor_hits.extend(success.output_monitor_hits.clone()); let elapsed = started.elapsed(); let status = success.response.status().as_u16(); - emit_access_log( - &model_name, - &success.provider, - &api_key_id, - status, - elapsed, - &request_id, - // `None` on the streaming path — `usage` is filled by the - // stream's completion callback, long after this line. That - // case is covered by the per-attempt `provider call - // completed` line the usage sink emits (AISIX-Cloud#1289). - success - .usage - .as_ref() - .map(|u| u.provider_request_id.as_str()), - &success.routing, - None, - ); + // See the note in `messages.rs`: the flag alone is not "the + // response is a stream". + if stream_requested && success.usage_handled_by_stream { + // A streamed response has no outcome yet: the head exists, nothing + // has been delivered, and whether the caller reads it to the end + // or walks away is minutes from being known. Park the line and + // let whichever terminal emitter ends the request write it, with + // that emitter's status, tokens and message (AISIX-Cloud#1571). + crate::attribution::defer_access_log( + crate::attribution::PendingAccessLog::new( + "POST", + "/v1/responses", + &request_id, + &api_key_id, + started, + ) + .with_model(&success.provider, &model_name) + .with_routing(&success.routing), + ); + } else { + emit_access_log( + &model_name, + &success.provider, + &api_key_id, + status, + elapsed, + &request_id, + success + .usage + .as_ref() + .map(|u| u.provider_request_id.as_str()), + &success.routing, + None, + ); + } // ONE ProviderKey lookup for both the metric emit and the // winner's usage event below (#941). let pk = ResolvedPk::resolve(&snapshot, &success.provider_key_id); @@ -3809,6 +3826,7 @@ fn emit_access_log( path: "/v1/responses", status, latency: elapsed, + duration: elapsed, provider: Some(provider), model: Some(model), upstream_model: target.upstream_model(), @@ -6905,4 +6923,47 @@ data: [DONE]\n\n"; assert!(scanned.contains("REASONINGSECRET"), "got {scanned:?}"); assert!(scanned.contains("SUMMARYSECRET"), "got {scanned:?}"); } + + /// A streamed `/v1/responses` relay writes ONE access-log line, at the + /// stream's end rather than when the head went out, so each of the three + /// endings reports its own outcome (AISIX-Cloud#1571). + /// + /// The upstream is a real chunked SSE server rather than a canned body: + /// this family relays BYTES, so an upstream that answers in one chunk + /// would hand the caller the whole stream in a single frame and the + /// "walked away mid-stream" ending could not happen at all. + #[tokio::test] + async fn a_streamed_relay_writes_one_line_per_stream_ending() { + let upstream = crate::test_log::spawn_sse_upstream(vec![ + "event: response.output_text.delta\n\ + data: {\"type\":\"response.output_text.delta\",\"delta\":\"a \"}\n\n" + .to_string(), + "event: response.output_text.delta\n\ + data: {\"type\":\"response.output_text.delta\",\"delta\":\"clean answer\"}\n\n" + .to_string(), + "event: response.completed\n\ + data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_stream\",\ + \"usage\":{\"input_tokens\":5,\"output_tokens\":2,\"total_tokens\":7}}}\n\n" + .to_string(), + "data: [DONE]\n\n".to_string(), + ]) + .await; + + let snap = new_snap_openai(&upstream); + snap.models.insert(openai_model("gpt-4o-resp")); + snap.apikeys.insert(apikey_entry(&["*"])); + let app = build_app(snap); + + let endings = crate::test_log::three_stream_endings(app, || { + make_req(serde_json::json!({"model":"gpt-4o-resp","input":"hi","stream":true})) + }) + .await; + crate::test_log::assert_one_line_per_ending(&endings, "/v1/responses", "k-1"); + crate::test_log::assert_latency_is_time_to_first_token(&endings); + assert_eq!( + endings.delivered.field("provider_request_id").as_deref(), + Some("resp_stream"), + "the id rides the terminal frame, which only the end-of-stream line can see", + ); + } } diff --git a/crates/aisix-proxy/src/test_log.rs b/crates/aisix-proxy/src/test_log.rs new file mode 100644 index 00000000..fc790ad4 --- /dev/null +++ b/crates/aisix-proxy/src/test_log.rs @@ -0,0 +1,316 @@ +//! Access-log capture for the tests that pin what a request's ONE line says. +//! +//! A streamed request's line is not written where its handler returns: it is +//! parked on the request's attribution cell and written by whichever +//! terminal emitter ends the request (AISIX-Cloud#1571). That makes "how +//! many lines did this request write, and what did the one line say" a +//! question only a log capture can answer — the value is not returned +//! anywhere a test could read it. +//! +//! [`three_stream_endings`] drives one streaming request through all three +//! endings a stream has, so a family cannot pass the delivered case and +//! silently lose the other two. + +use std::sync::{Arc, Mutex}; + +use axum::body::Body; +use axum::http::Request; +use axum::Router; + +/// The access log's own `tracing` message. Counting its occurrences is how +/// a test asks "how many lines did this request write" without matching the +/// other events the same subscriber sees. +const ACCESS_LOG_MESSAGE: &str = "proxy request completed"; + +struct LogBuf(Arc>>); + +impl std::io::Write for LogBuf { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(buf); + Ok(buf.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +impl tracing_subscriber::fmt::MakeWriter<'_> for LogBuf { + type Writer = LogBuf; + fn make_writer(&self) -> Self::Writer { + LogBuf(self.0.clone()) + } +} + +/// A capturing subscriber installed on THIS thread, plus its buffer. +/// +/// Thread-local rather than global: a `#[tokio::test]` runs its future on +/// the calling thread, which is where the handler, the body polls and the +/// guard drops all write from. +pub(crate) struct Capture { + buf: Arc>>, + _guard: tracing::subscriber::DefaultGuard, +} + +impl Capture { + pub(crate) fn install() -> Self { + static ONCE: std::sync::Once = std::sync::Once::new(); + ONCE.call_once(|| { + let _ = tracing::subscriber::set_global_default(tracing_subscriber::registry()); + }); + let buf = Arc::new(Mutex::new(Vec::new())); + let subscriber = tracing_subscriber::fmt() + .with_ansi(false) + .with_writer(LogBuf(buf.clone())) + .finish(); + let _guard = tracing::subscriber::set_default(subscriber); + Self { buf, _guard } + } + + fn text(&self) -> String { + String::from_utf8_lossy(&self.buf.lock().unwrap()).into_owned() + } + + /// The request's one and only line. Panics with the captured text when + /// there is not exactly one — which is the failure this whole module + /// exists to catch. + pub(crate) fn only(&self, what: &str) -> AccessLine { + let text = self.text(); + let lines: Vec<&str> = text + .lines() + .filter(|l| l.contains(ACCESS_LOG_MESSAGE)) + .collect(); + assert_eq!( + lines.len(), + 1, + "{what}: expected exactly one access-log line, got {}:\n{text}", + lines.len(), + ); + AccessLine(lines[0].to_string()) + } +} + +/// One rendered access-log line, read field by field. +pub(crate) struct AccessLine(String); + +impl AccessLine { + /// The value of `name=`, unquoted. `None` when the field is absent — + /// which is meaningful here: this log omits `None` fields rather than + /// rendering them empty, so an operator can filter on their presence. + pub(crate) fn field(&self, name: &str) -> Option { + let needle = format!("{name}="); + let mut from = 0usize; + loop { + let idx = self.0[from..].find(&needle)? + from; + let starts_token = idx == 0 || self.0.as_bytes()[idx - 1] == b' '; + let after = &self.0[idx + needle.len()..]; + if !starts_token { + from = idx + needle.len(); + continue; + } + return Some(match after.strip_prefix('"') { + Some(quoted) => { + let mut out = String::new(); + let mut chars = quoted.chars(); + while let Some(c) = chars.next() { + match c { + '\\' => out.extend(chars.next()), + '"' => break, + _ => out.push(c), + } + } + out + } + None => after.split(' ').next().unwrap_or_default().to_string(), + }); + } + } + + pub(crate) fn num(&self, name: &str) -> Option { + self.field(name)?.parse().ok() + } + + pub(crate) fn status(&self) -> u64 { + self.num("status").expect("every line carries a status") + } +} + +/// The three ways a stream ends, each with the one line it wrote. +pub(crate) struct StreamEndings { + /// Read to the end — slowly, with a pause after the first frame, so the + /// caller's wait and the stream's length cannot be the same number by + /// accident. + pub delivered: AccessLine, + /// One frame read, then the caller walks away. + pub abandoned: AccessLine, + /// The body dropped before its first poll. + pub unread: AccessLine, +} + +/// How long the delivered ending holds the stream open after its first +/// frame. Large enough that a line reporting the whole stream in +/// `latency_ms` cannot be mistaken for one reporting the first frame. +pub(crate) const SLOW_DRAIN: std::time::Duration = std::time::Duration::from_millis(200); + +/// Drive `request` against `app` three times, once per stream ending, and +/// return the single access-log line each produced. +pub(crate) async fn three_stream_endings( + app: Router, + request: impl Fn() -> Request, +) -> StreamEndings { + use futures::StreamExt as _; + use tower::ServiceExt as _; + + let delivered = { + let capture = Capture::install(); + let response = app.clone().oneshot(request()).await.unwrap(); + assert!( + response.status().is_success(), + "premise: the stream has to open, got {}", + response.status(), + ); + let mut body = response.into_body().into_data_stream(); + // `Some(Err(_))` is a BROKEN body, not a delivered frame — accepting + // it would let a stream that failed on its first poll pass as the + // delivered ending, which is the one ending whose line says 200. + let first = body.next().await; + assert!( + matches!(first, Some(Ok(_))), + "premise: the stream delivered no frame: {first:?}", + ); + tokio::time::sleep(SLOW_DRAIN).await; + while let Some(frame) = body.next().await { + frame.expect("the delivered stream must read cleanly to its end"); + } + drop(body); + capture.only("a delivered stream") + }; + + let abandoned = { + let capture = Capture::install(); + let response = app.clone().oneshot(request()).await.unwrap(); + let mut body = response.into_body().into_data_stream(); + let first = body.next().await; + assert!( + matches!(first, Some(Ok(_))), + "premise: the stream delivered no frame to abandon: {first:?}", + ); + drop(body); + capture.only("a stream abandoned mid-flight") + }; + + let unread = { + let capture = Capture::install(); + let response = app.clone().oneshot(request()).await.unwrap(); + let (_parts, body) = response.into_parts(); + drop(body); + capture.only("a stream dropped before its first poll") + }; + + StreamEndings { + delivered, + abandoned, + unread, + } +} + +/// What every family's three endings must say, whatever it streams. +/// +/// The two abandoned endings report the SAME outcome the request's terminal +/// usage event reports — `499` / `client_disconnected` — because the line +/// now rides that event out of one chokepoint instead of being written when +/// the response head was handed over, which is what used to log an +/// abandoned stream as a `200` (AISIX-Cloud#1571). +pub(crate) fn assert_one_line_per_ending(endings: &StreamEndings, path: &str, api_key_id: &str) { + for (what, line, status) in [ + ("delivered", &endings.delivered, 200), + ( + "abandoned", + &endings.abandoned, + u64::from(crate::CLIENT_CLOSED_REQUEST), + ), + ( + "unread", + &endings.unread, + u64::from(crate::CLIENT_CLOSED_REQUEST), + ), + ] { + assert_eq!(line.status(), status, "{what}: wrong status on the line"); + assert_eq!( + line.field("path").as_deref(), + Some(path), + "{what}: wrong path", + ); + assert_eq!( + line.field("api_key_id").as_deref(), + Some(api_key_id), + "{what}: the line must still name the caller", + ); + let latency = line + .num("latency_ms") + .unwrap_or_else(|| panic!("{what}: no latency_ms")); + let duration = line + .num("duration_ms") + .unwrap_or_else(|| panic!("{what}: no duration_ms")); + assert!( + duration >= latency, + "{what}: duration_ms ({duration}) is what the request took and \ + latency_ms ({latency}) is the wait inside it — transposed", + ); + let expected_class = (status != 200).then_some(crate::CLIENT_DISCONNECTED_KIND); + assert_eq!( + line.field("error_kind").as_deref(), + expected_class, + "{what}: the line and the usage event must name the same failure class", + ); + } +} + +/// The families whose streamed line reports time-to-first-token: the +/// delivered ending held the stream open for [`SLOW_DRAIN`] after its first +/// frame, so a line reporting the whole stream cannot pass. +pub(crate) fn assert_latency_is_time_to_first_token(endings: &StreamEndings) { + let latency = endings.delivered.num("latency_ms").unwrap(); + let duration = endings.delivered.num("duration_ms").unwrap(); + assert!( + duration >= latency + SLOW_DRAIN.as_millis() as u64 / 2, + "latency_ms ({latency}) is supposed to be the wait to the FIRST token, \ + but it is within a rounding error of the whole stream ({duration})", + ); +} + +/// A local SSE upstream that emits `frames` one at a time, pausing between +/// them. +/// +/// wiremock answers with the whole body at once, which the byte-relaying +/// families (`/v1/responses`, the audio transcription relay) forward as a +/// SINGLE frame — so "read one frame, then walk away" would read the entire +/// stream and never model an abandoned one. Pausing between chunks is what +/// makes the three endings actually different. +pub(crate) async fn spawn_sse_upstream(frames: Vec) -> String { + use axum::response::IntoResponse; + use futures::StreamExt as _; + + let frames = Arc::new(frames); + let app = Router::new().fallback(axum::routing::any(move || { + let frames = frames.clone(); + async move { + let stream = futures::stream::iter(frames.as_ref().clone()).then(|frame| async move { + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + Ok::<_, std::convert::Infallible>(frame) + }); + ( + [(axum::http::header::CONTENT_TYPE, "text/event-stream")], + Body::from_stream(stream), + ) + .into_response() + } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, app.into_make_service()) + .await + .unwrap(); + }); + format!("http://{addr}") +} diff --git a/crates/aisix-proxy/src/usage_attr.rs b/crates/aisix-proxy/src/usage_attr.rs index f6371099..f986c704 100644 --- a/crates/aisix-proxy/src/usage_attr.rs +++ b/crates/aisix-proxy/src/usage_attr.rs @@ -759,6 +759,15 @@ pub(crate) fn emit_usage( event.error_class = crate::CLIENT_DISCONNECTED_KIND.to_string(); event.error_message = crate::cancel::CANCELLED_MID_STREAM.to_string(); } + // The request's access-log line, for the families that deferred it to + // their stream (AISIX-Cloud#1571). Here, and after the stamping above, + // so the line and the event cannot disagree about the outcome: one + // status, one error class, one message, whichever of a stream's three + // endings this is. A request that wrote its line inline — everything + // non-streamed — parked nothing, and this is a no-op for it. + if terminal { + crate::attribution::emit_deferred_access_log(&event); + } // Request-level guardrail blocks are recorded from the terminal event, // not from an individual timed execution. Some fail-closed paths (for // example a streamed-output buffer overflow) reject before a guardrail diff --git a/crates/aisix-proxy/src/videos.rs b/crates/aisix-proxy/src/videos.rs index a7d990e7..c9ffb477 100644 --- a/crates/aisix-proxy/src/videos.rs +++ b/crates/aisix-proxy/src/videos.rs @@ -1465,6 +1465,7 @@ impl Telemetry<'_> { path: &self.path, status, latency: elapsed, + duration: elapsed, provider: Some(provider).filter(|p| !p.is_empty()), model: Some(model_label), upstream_model: log_target.upstream_model(), diff --git a/tests/e2e/src/cases/client-cancel-usage-1571-e2e.test.ts b/tests/e2e/src/cases/client-cancel-usage-1571-e2e.test.ts index 5cc68271..40ac8ec9 100644 --- a/tests/e2e/src/cases/client-cancel-usage-1571-e2e.test.ts +++ b/tests/e2e/src/cases/client-cancel-usage-1571-e2e.test.ts @@ -42,16 +42,23 @@ const UPSTREAM_MODEL = "gpt-4o-mini"; /** A direct model on a fast upstream, for the success-path line below. */ const FAST_MODEL = "c1571-fast"; const FAST_UPSTREAM_MODEL = "gpt-4o-fast"; +/** A direct model on an upstream that trickles its stream, for the + * mid-stream abandon below — the head IS written there, which is the + * ending the handler tail used to log as a `200`. */ +const STREAM_MODEL = "c1571-stream"; +const STREAM_UPSTREAM_MODEL = "gpt-4o-stream"; describe("client cancel before the response head (AISIX-Cloud#1571)", () => { let etcdReachable = false; let slow: OpenAiUpstream | undefined; let fast: OpenAiUpstream | undefined; + let trickle: OpenAiUpstream | undefined; let sls: MockSls | undefined; let app: SpawnedApp | undefined; let targetModelId = ""; let providerKeyId = ""; let fastProviderKeyId = ""; + let streamProviderKeyId = ""; beforeAll(async () => { etcdReachable = await new EtcdClient().ping(); @@ -74,6 +81,35 @@ describe("client cancel before the response head (AISIX-Cloud#1571)", () => { usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, }, }); + // Half a second between events, so a caller can read the first one and + // abort while the rest is still coming — the mid-stream ending. + trickle = await startOpenAiUpstream({ + eventDelayMs: 500, + streamEvents: [ + JSON.stringify({ + id: "chatcmpl-c1571-stream", + object: "chat.completion.chunk", + created: 1_700_000_000, + model: STREAM_UPSTREAM_MODEL, + choices: [{ index: 0, delta: { role: "assistant", content: "one" } }], + }), + JSON.stringify({ + id: "chatcmpl-c1571-stream", + object: "chat.completion.chunk", + created: 1_700_000_000, + model: STREAM_UPSTREAM_MODEL, + choices: [{ index: 0, delta: { content: "two" } }], + }), + JSON.stringify({ + id: "chatcmpl-c1571-stream", + object: "chat.completion.chunk", + created: 1_700_000_000, + model: STREAM_UPSTREAM_MODEL, + choices: [{ index: 0, delta: { content: "three" }, finish_reason: "stop" }], + }), + "[DONE]", + ], + }); sls = await startMockSls(); app = await spawnApp({ @@ -125,6 +161,18 @@ describe("client cancel before the response head (AISIX-Cloud#1571)", () => { model_name: FAST_UPSTREAM_MODEL, provider_key_id: fastPk.id, }); + const streamPk = await seed.createProviderKey({ + display_name: "c1571-stream-pk", + secret: PROVIDER_SECRET, + api_base: `${trickle.baseUrl}/v1`, + }); + streamProviderKeyId = streamPk.id; + await seed.createModel({ + display_name: STREAM_MODEL, + provider: "openai", + model_name: STREAM_UPSTREAM_MODEL, + provider_key_id: streamPk.id, + }); await seed.createPassthroughRoute({ name: ROUTE, path_prefix: "/passthrough/c1571", @@ -146,6 +194,7 @@ describe("client cancel before the response head (AISIX-Cloud#1571)", () => { await app?.exit(); await slow?.close(); await fast?.close(); + await trickle?.close(); await sls?.close(); }); @@ -317,4 +366,85 @@ describe("client cancel before the response head (AISIX-Cloud#1571)", () => { expect(line).toContain(`upstream_model="${FAST_UPSTREAM_MODEL}"`); expect(line).toContain(`provider_key_id="${fastProviderKeyId}"`); }); + + // The other half of the same request: a caller that walks away AFTER the + // response head — the ordinary ending for a long stream. The row was + // already a 499 before this change; the LINE said 200, because the handler + // wrote it when it handed the stream over, minutes before the request + // ended. One request read as two, under two statuses, and nothing in the + // log said which one was the outcome. + test( + "a stream abandoned mid-flight writes exactly one line, and it says what the row says", + async (ctx) => { + if (!etcdReachable || !app || !trickle || !sls) { + ctx.skip(); + return; + } + const controller = new AbortController(); + const res = await fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: STREAM_MODEL, + messages: [{ role: "user", content: "start streaming" }], + stream: true, + }), + signal: controller.signal, + }); + expect(res.status, "the head must go out — this is the mid-stream ending").toBe(200); + const requestId = res.headers.get("x-aisix-request-id") ?? ""; + expect(requestId).not.toBe(""); + + // Read one event, then hang up with the rest still coming. + const reader = res.body!.getReader(); + const first = await reader.read(); + expect(first.done, "the stream delivered nothing to abandon").toBe(false); + controller.abort(); + // `cancel()` races the abort: it resolves when cancellation wins and + // rejects with the body's `AbortError` when the abort does. Any other + // rejection is a real failure and must not be swallowed. + await reader.cancel().catch((err: unknown) => { + if (!(err instanceof Error) || err.name !== "AbortError") throw err; + }); + + const row = await waitForSlsLog( + sls, + LOGSTORE, + (log) => log.get("request_id") === requestId, + `a usage row for the abandoned stream ${requestId}`, + 20_000, + ); + expect(row.get("status_code")).toBe("499"); + expect(row.get("error_class")).toBe("client_disconnected"); + expect(row.get("error_message")).toContain("while the response was streaming"); + + const lines = app + .output() + .split("\n") + .filter( + (l) => l.includes("proxy request completed") && l.includes(`request_id="${requestId}"`), + ); + expect( + lines.length, + `one request, one line — got ${lines.length} for ${requestId}:\n${lines.join("\n")}`, + ).toBe(1); + const line = lines[0]; + expect(line).toContain("status=499"); + expect(line).toContain(`error_kind="client_disconnected"`); + expect(line).toContain("while the response was streaming"); + // The target is still named — the line an operator reads to find out + // which upstream the abandoned call was costing them. + expect(line).toContain(`model="${STREAM_MODEL}"`); + expect(line).toContain(`upstream_model="${STREAM_UPSTREAM_MODEL}"`); + expect(line).toContain(`provider_key_id="${streamProviderKeyId}"`); + // And the two spans the line now separates: what the caller waited for + // (the first token) inside how long the request ran. + expect(line).toMatch(/\blatency_ms=\d+/); + expect(line).toMatch(/\bduration_ms=\d+/); + }, + 60_000, + ); }); diff --git a/tests/e2e/src/cases/provider-request-id-logging-e2e.test.ts b/tests/e2e/src/cases/provider-request-id-logging-e2e.test.ts index 506ac200..a11b5262 100644 --- a/tests/e2e/src/cases/provider-request-id-logging-e2e.test.ts +++ b/tests/e2e/src/cases/provider-request-id-logging-e2e.test.ts @@ -215,7 +215,7 @@ describe("provider_request_id reaches the access log and the plain log", () => { expect(res.requestId).not.toBe(NONSTREAM_ID); }); - test("streaming: the provider-call line carries the id the access log cannot", async (ctx) => { + test("streaming: the access-log line and the per-attempt line both carry the id", async (ctx) => { if (!etcdReachable || !app) { ctx.skip(); return; @@ -229,9 +229,22 @@ describe("provider_request_id reaches the access log and the plain log", () => { expect(res.status).toBe(200); expect(res.text).toContain("[DONE]"); - // The whole point of the per-attempt line: the id only exists once the - // first upstream frame lands, by which time the access-log line for this - // request has already been written. + // The id only exists once the first upstream frame lands — which used to + // be after this request's access-log line had been written. The line is + // written at the stream's END now (AISIX-Cloud#1571), so it carries the + // winning call's id like a buffered one does. + const access = await waitForLogLine( + app, + (l) => + l.includes("proxy request completed") && + l.includes(`request_id="${res.requestId}"`), + "the access-log line for this streamed request", + ); + expect(access).toContain(`provider_request_id="${STREAM_ID}"`); + + // The per-attempt line is still the one that identifies an INDIVIDUAL + // provider call: `request_id` + `attempt_index`, one per attempt of a + // retried or failed-over request, where the access log has one row. const line = await waitForLogLine( app, (l) => @@ -240,8 +253,6 @@ describe("provider_request_id reaches the access log and the plain log", () => { "the provider-call line for this streamed request", ); expect(line).toContain(`provider_request_id="${STREAM_ID}"`); - // `request_id` + `attempt_index` is what identifies an individual - // provider call across a retried / failed-over request. expect(line).toContain("attempt_index="); });