From 024e45ea751370a9d90a42b8df01ac6c9bfd1720 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 16 Sep 2026 01:43:16 +0000 Subject: [PATCH 1/4] fix(telemetry): write a streamed request's access-log line when the stream ends, not when its head does MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A streaming handler returns the moment the response head exists — often minutes before the request ends — and it wrote the request's access-log line there. So a stream the caller abandoned logged `200` beside its own `499` usage event, a delivered stream's line carried neither the token counts nor the provider response id, and one request read as two records under two statuses. Six families did this: chat, messages, responses, audio's two transcription routes, and a2a. Passthrough already wrote its line at the stream's end and is the shape the others now follow. A streaming handler no longer writes the line at all. It parks the fields only it can resolve on the request's 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 three endings produced it: delivered, abandoned mid-flight, or dropped before its first poll. Taking the parked line out of the cell is also the interlock, so however many terminal emitters race, exactly one line is written. The line and the row therefore agree on `status`, `error_class` and `error_message` by construction rather than by each family remembering to, and the streamed line now carries the counts and `provider_request_id` that only exist once the upstream has answered. Two supporting changes make that reachable: the response body is polled inside the request's attribution scope (the delivered ending runs there), and the cancel guard installs the cell around its own emit (the two abandoned endings run from `Drop`, outside every scope). Every access-log line also gains `duration_ms`, how long the request occupied the gateway. `latency_ms` keeps its meaning — what the CALLER waited for, time-to-first-token on a stream (AISIX-Cloud#1394) — so the two differ by the length of the stream and coincide on everything buffered. Passthrough's streamed line now reports its first relayed frame in `latency_ms` for the same reason, instead of the whole relay. `/v1/completions` refuses `stream: true`, so it only gains the new field. `/v1/audio/speech` relays bytes but meters at its tail, so it keeps writing its line there. On `/a2a` and the audio transcription relay, `latency_ms` equals the whole-stream figure their usage event already reports as the caller's wait; reporting a time-to-first-frame there would mean changing what `downstream_latency_ms` means, which this does not touch. Ref api7/AISIX-Cloud#1571 --- crates/aisix-obs/src/access_log.rs | 80 +++-- crates/aisix-proxy/AGENTS.md | 14 + crates/aisix-proxy/src/a2a.rs | 117 +++++-- crates/aisix-proxy/src/attribution.rs | 201 ++++++++++++ crates/aisix-proxy/src/audio.rs | 125 +++++-- crates/aisix-proxy/src/cancel.rs | 17 +- crates/aisix-proxy/src/chat.rs | 61 ++-- crates/aisix-proxy/src/completions.rs | 47 +++ crates/aisix-proxy/src/count_tokens.rs | 1 + crates/aisix-proxy/src/embeddings.rs | 1 + crates/aisix-proxy/src/images.rs | 1 + crates/aisix-proxy/src/jobs.rs | 1 + crates/aisix-proxy/src/lib.rs | 218 ++++++------- crates/aisix-proxy/src/mcp.rs | 1 + crates/aisix-proxy/src/messages.rs | 104 +++++- crates/aisix-proxy/src/passthrough_route.rs | 24 +- crates/aisix-proxy/src/realtime.rs | 1 + crates/aisix-proxy/src/reject.rs | 1 + crates/aisix-proxy/src/rerank.rs | 1 + crates/aisix-proxy/src/responses.rs | 95 ++++-- crates/aisix-proxy/src/test_log.rs | 305 ++++++++++++++++++ crates/aisix-proxy/src/usage_attr.rs | 9 + crates/aisix-proxy/src/videos.rs | 1 + .../client-cancel-usage-1571-e2e.test.ts | 125 +++++++ 24 files changed, 1299 insertions(+), 252 deletions(-) create mode 100644 crates/aisix-proxy/src/test_log.rs diff --git a/crates/aisix-obs/src/access_log.rs b/crates/aisix-obs/src/access_log.rs index 387e1b68..638c9f50 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,12 @@ 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, arrival to last byte out. + /// Equal to `latency` on everything that is not streamed. + 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 +96,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 +159,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 +233,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 +259,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 +317,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 +370,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..a317eee1 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,157 @@ 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, and the same figure the event + // reports: the first token forwarded downstream on a stream that + // delivered one. A stream that delivered nothing — the caller left + // before the first frame — has no such moment, so the line falls + // back to the whole request, which is all it waited for. + let latency = match event.downstream_latency_ms { + 0 => duration, + ms => Duration::from_millis(u64::from(ms)), + }; + let prompt = u64::from(event.prompt_tokens); + let completion = u64::from(event.completion_tokens); + let total = prompt + completion; + // 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 +385,23 @@ impl RequestAttribution { std::mem::take(&mut self.lock().cancel) } + /// 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 +648,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..f3db6f45 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -150,17 +150,34 @@ 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, - ); + 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 +331,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 +2297,7 @@ fn emit_access_log( path, status, latency, + duration: latency, provider: Some(provider), model: Some(model), upstream_model: target.upstream_model(), @@ -4200,4 +4235,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..05bd5d47 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,14 +867,12 @@ 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. + // The head phase builds its OWN line, because the handler 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. if matches!(phase, cancel::Phase::BeforeHead) { let target = attribution::AccessLogTarget::from_resolved(resolved.clone()); AccessLog { @@ -870,6 +880,9 @@ impl Drop for ClientCancelGuard { 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 +908,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 +9068,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 +9085,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 +9402,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 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..95549a7d 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -180,21 +180,36 @@ 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, - ); + if 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 +4219,7 @@ fn emit_access_log( path: "/v1/messages", status, latency, + duration: latency, provider: Some(provider), model: Some(model), upstream_model: target.upstream_model(), @@ -6994,4 +7010,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..29ca026f 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -245,24 +245,39 @@ 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, - ); + if 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 +3824,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 +6921,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..b9d50875 --- /dev/null +++ b/crates/aisix-proxy/src/test_log.rs @@ -0,0 +1,305 @@ +//! 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(); + let first = body.next().await; + assert!(first.is_some(), "premise: the stream delivered no frame"); + tokio::time::sleep(SLOW_DRAIN).await; + while body.next().await.is_some() {} + 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!(first.is_some(), "premise: the stream delivered no frame"); + 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..b6181fee 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,80 @@ 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(); + await reader.cancel().catch(() => {}); + + 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, + ); }); From 731e8b215410b27de9ba0dcf7ea4a34dd0313329 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 16 Sep 2026 01:47:24 +0000 Subject: [PATCH 2/4] test(e2e): a streamed request's access-log line carries the provider response id MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The line is written at the stream's end now, so the claim this spec was built on — that the id lands after the request's access-log line — no longer holds. Assert the line carries it, and keep the per-attempt line for what only it can say: one row per provider call of a retried or failed-over request. Ref api7/AISIX-Cloud#1571 --- .../provider-request-id-logging-e2e.test.ts | 23 ++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) 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="); }); From 2ca474ac344926408de6efc9c85afe290541049f Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 16 Sep 2026 02:04:47 +0000 Subject: [PATCH 3/4] fix(telemetry): bill the deferred line's tokens cache-inclusively, and take the caller's wait verbatim MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three review findings on the deferred access-log line, all in what the line reports rather than where it is written. `total_tokens` summed only the two visible columns. `prompt_tokens` excludes `cache_creation_tokens` and `cache_read_tokens` on the Anthropic-shaped paths, so a streamed call with prompt caching logged a total an order of magnitude under the row it was emitted beside — 89 against 102 in the test that now pins it. It uses `total_tokens_with_cache`, like every other total in the crate. `latency` treated `downstream_latency_ms == 0` as "never stamped" and fell back to the whole request. It is also a real value: the field is truncated milliseconds, so a co-located upstream whose first frame lands in under a millisecond stamps `0` — and the line would then report the length of the stream as the caller's wait, the exact AISIX-Cloud#1394 misreading. The figure is now taken verbatim off the event, so the line and the row cannot report two different waits. `duration`'s doc claimed "arrival to last byte out" everywhere. That holds wherever the line is written at completion, but `/v1/audio/speech` and `/v1/videos/{id}/content` meter at their handler tail and relay an open-ended body afterwards; theirs ends at the response head. The field says so. Also: `/v1/messages` and `/v1/responses` conjoin their own streaming flag before parking a line, the way chat already does, so reusing `usage_handled_by_stream` to mean "already emitted" on a buffered path cannot silently lose a line. The audio tails have no second predicate — the flag there means the response IS a live relay — and say so. Review (CodeRabbit): the stream-ending driver treated `Some(Err(_))` as a delivered frame, so a body that broke on its first poll would have passed as the delivered ending; and the e2e swallowed every `reader.cancel()` rejection rather than just the abort's own. Ref api7/AISIX-Cloud#1571 --- crates/aisix-obs/src/access_log.rs | 14 ++++++-- crates/aisix-proxy/src/attribution.rs | 31 +++++++++++------ crates/aisix-proxy/src/audio.rs | 7 ++++ crates/aisix-proxy/src/messages.rs | 33 ++++++++++++++++++- crates/aisix-proxy/src/responses.rs | 4 ++- crates/aisix-proxy/src/test_log.rs | 17 ++++++++-- .../client-cancel-usage-1571-e2e.test.ts | 7 +++- 7 files changed, 95 insertions(+), 18 deletions(-) diff --git a/crates/aisix-obs/src/access_log.rs b/crates/aisix-obs/src/access_log.rs index 638c9f50..2d7cb380 100644 --- a/crates/aisix-obs/src/access_log.rs +++ b/crates/aisix-obs/src/access_log.rs @@ -67,8 +67,18 @@ pub struct AccessLog<'a> { /// 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, arrival to last byte out. - /// Equal to `latency` on everything that is not streamed. + /// 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 diff --git a/crates/aisix-proxy/src/attribution.rs b/crates/aisix-proxy/src/attribution.rs index a317eee1..ef468df0 100644 --- a/crates/aisix-proxy/src/attribution.rs +++ b/crates/aisix-proxy/src/attribution.rs @@ -308,18 +308,29 @@ impl PendingAccessLog { /// 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, and the same figure the event - // reports: the first token forwarded downstream on a stream that - // delivered one. A stream that delivered nothing — the caller left - // before the first frame — has no such moment, so the line falls - // back to the whole request, which is all it waited for. - let latency = match event.downstream_latency_ms { - 0 => duration, - ms => Duration::from_millis(u64::from(ms)), - }; + // 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); - let total = prompt + completion; + // 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 diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index f3db6f45..b04535d7 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -150,6 +150,13 @@ 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(); + // 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 diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 95549a7d..0615e777 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -180,7 +180,12 @@ pub async fn messages( monitor_hits.extend(output_monitor_hits); let elapsed = started.elapsed(); let status = response.status().as_u16(); - if usage_handled_by_stream { + // 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 @@ -5636,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); @@ -5679,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 diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 29ca026f..1af4503c 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -245,7 +245,9 @@ pub async fn responses( monitor_hits.extend(success.output_monitor_hits.clone()); let elapsed = started.elapsed(); let status = success.response.status().as_u16(); - if success.usage_handled_by_stream { + // 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 diff --git a/crates/aisix-proxy/src/test_log.rs b/crates/aisix-proxy/src/test_log.rs index b9d50875..fc790ad4 100644 --- a/crates/aisix-proxy/src/test_log.rs +++ b/crates/aisix-proxy/src/test_log.rs @@ -169,10 +169,18 @@ pub(crate) async fn three_stream_endings( 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!(first.is_some(), "premise: the stream delivered no frame"); + assert!( + matches!(first, Some(Ok(_))), + "premise: the stream delivered no frame: {first:?}", + ); tokio::time::sleep(SLOW_DRAIN).await; - while body.next().await.is_some() {} + 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") }; @@ -182,7 +190,10 @@ pub(crate) async fn three_stream_endings( let response = app.clone().oneshot(request()).await.unwrap(); let mut body = response.into_body().into_data_stream(); let first = body.next().await; - assert!(first.is_some(), "premise: the stream delivered no frame"); + assert!( + matches!(first, Some(Ok(_))), + "premise: the stream delivered no frame to abandon: {first:?}", + ); drop(body); capture.only("a stream abandoned mid-flight") }; 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 b6181fee..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 @@ -403,7 +403,12 @@ describe("client cancel before the response head (AISIX-Cloud#1571)", () => { const first = await reader.read(); expect(first.done, "the stream delivered nothing to abandon").toBe(false); controller.abort(); - await reader.cancel().catch(() => {}); + // `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, From ff2da2e9bb5fc294ee22ffedae9a54cadf69766d Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 16 Sep 2026 02:09:15 +0000 Subject: [PATCH 4/4] fix(telemetry): a head-phase cancel writes the parked line, not a second one beside it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A streaming family parks its access-log line at its handler tail and chat then awaits once more — peeking the rate limiter to fill the `x-ratelimit-*` headers. A caller that hangs up in that window leaves the guard in its HEAD phase with the request's line already on the cell, so both emitters spoke: the guard built its own `499` line and `cancel::emit` then wrote the parked one under the same status and the same message. One request read as two near-identical records, and counting `499` lines doubled it. The head phase now stands down when a line is parked. The parked one is the fuller of the two — it names the model, the dispatched target and the routing counts — and it cannot fall between the two emitters: a parked line means the request authenticated on a metering surface, which is the gate `cancel::emit` applies before emitting the terminal event that carries it. Found by review; the window is reachable today on `/v1/chat/completions` with a Redis rate-limit backend, where the peek is a real round trip. Ref api7/AISIX-Cloud#1571 --- crates/aisix-proxy/src/attribution.rs | 7 ++ crates/aisix-proxy/src/lib.rs | 96 +++++++++++++++++++++++++-- 2 files changed, 96 insertions(+), 7 deletions(-) diff --git a/crates/aisix-proxy/src/attribution.rs b/crates/aisix-proxy/src/attribution.rs index ef468df0..8a90e1c5 100644 --- a/crates/aisix-proxy/src/attribution.rs +++ b/crates/aisix-proxy/src/attribution.rs @@ -396,6 +396,13 @@ 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. /// diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 05bd5d47..32066aae 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -867,13 +867,25 @@ impl Drop for ClientCancelGuard { if matches!(phase, cancel::Phase::BeforeBody) && cancel_ctx.emitted_terminal { return; } - // The head phase builds its OWN line, because the handler 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. - 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(), @@ -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