diff --git a/crates/aisix-a2a/src/lib.rs b/crates/aisix-a2a/src/lib.rs index 63623561..eee6326b 100644 --- a/crates/aisix-a2a/src/lib.rs +++ b/crates/aisix-a2a/src/lib.rs @@ -27,4 +27,7 @@ pub use bridge::{ HttpBridge, DEFAULT_UPSTREAM_TIMEOUT, }; pub use error::A2aError; -pub use telemetry::{canonical_operation, is_stream_end, is_streaming_operation, A2aCallFacts}; +pub use telemetry::{ + canonical_operation, is_stream_end, is_streaming_operation, request_text, A2aCallFacts, + ResultText, +}; diff --git a/crates/aisix-a2a/src/telemetry.rs b/crates/aisix-a2a/src/telemetry.rs index 6688b94e..9af1ab0a 100644 --- a/crates/aisix-a2a/src/telemetry.rs +++ b/crates/aisix-a2a/src/telemetry.rs @@ -285,6 +285,142 @@ impl A2aCallFacts { } } +/// The text a caller sent an agent: every text part of the request's message, +/// newline-joined. +/// +/// A `Part` carries its text under `text` in both wire versions (0.3 tags the +/// part with `kind`, 1.0 uses a protobuf `oneof` whose set field has the same +/// name), so one reader serves both. File and data parts contribute nothing — +/// their bytes are not language, and a base64 blob would wreck both the token +/// estimate and any captured content it lands in. +/// +/// `push` bounds the buffer; a request body has no size limit by default, and +/// the result is retained for the lifetime of a streamed call. +pub fn request_text(request: &Value, push: impl Fn(&mut String, &str)) -> String { + let mut buf = String::new(); + if let Some(message) = request.pointer("/params/message") { + collect_part_text(message, &mut buf, &push); + } + buf +} + +/// The text an agent produced on one call, kept as two segments because the +/// protocol updates them by two different rules. +/// +/// Artifacts are the answer, delivered in chunks that may continue one +/// another. Statements — a Message, or the message an agent attaches to a +/// status update — are what it says ABOUT the task: "still working", "report +/// generated". Keeping them apart is not tidiness. A single buffer forces one +/// rule on both, and either choice is wrong: appending multiplies an answer +/// that the agent resends as the task progresses, while replacing lets a +/// one-line progress note wipe a report that took a thousand chunks to build. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub struct ResultText { + /// The artifact currently being streamed. `append` is scoped to one + /// artifact id by the specification ("appended to a previously sent + /// artifact with the same ID"), so a chunk for a DIFFERENT artifact opens + /// a new segment rather than discarding what came before. + artifact_id: String, + /// The answer so far. + artifacts: String, + /// The agent's latest word about the task. Only ever the latest: a + /// progress note supersedes the one before it and restates nothing. + statement: String, +} + +impl ResultText { + /// Read one response envelope — or one streamed event — for the words it + /// carries. + /// + /// `push` bounds each segment; an A2A task may stream for hours, and past + /// the bound the text becomes a prefix rather than the buffer growing + /// without limit. + pub fn observe(&mut self, response: &Value, push: impl Fn(&mut String, &str)) { + let Some(result) = response.get("result") else { + return; + }; + let payload = unwrap_payload(result); + + // An artifact update: one chunk of the answer, appended or standalone + // per the event's own flag, scoped to the artifact it names. + if let Some(artifact) = payload.get("artifact") { + let mut fresh = String::new(); + collect_part_text(artifact, &mut fresh, &push); + if fresh.is_empty() { + return; + } + let id = str_field(artifact, "artifactId").unwrap_or_default(); + if id == self.artifact_id { + // The same artifact again. Without `append` it is a + // replacement of itself, not a continuation. + if payload.get("append").and_then(Value::as_bool) != Some(true) { + self.artifacts.clear(); + } + } else { + id.clone_into(&mut self.artifact_id); + if !self.artifacts.is_empty() { + push(&mut self.artifacts, "\n"); + } + } + push(&mut self.artifacts, &fresh); + return; + } + + // A Task snapshot carries the complete artifact set, so it replaces + // that segment outright — it restates the chunks rather than adding + // to them. + if let Some(artifacts) = payload.get("artifacts").and_then(Value::as_array) { + let mut fresh = String::new(); + for artifact in artifacts { + collect_part_text(artifact, &mut fresh, &push); + } + if !fresh.is_empty() { + self.artifacts.clear(); + self.artifact_id.clear(); + push(&mut self.artifacts, &fresh); + } + } + + // ...and whatever the agent says about the task replaces only the + // previous such statement. + let mut fresh = String::new(); + collect_part_text(payload, &mut fresh, &push); + if let Some(message) = payload.pointer("/status/message") { + collect_part_text(message, &mut fresh, &push); + } + if !fresh.is_empty() { + self.statement.clear(); + push(&mut self.statement, &fresh); + } + } + + /// Everything the agent produced, for counting or capture. Empty when it + /// produced nothing. + pub fn joined(&self) -> String { + match (self.artifacts.is_empty(), self.statement.is_empty()) { + (true, _) => self.statement.clone(), + (_, true) => self.artifacts.clone(), + _ => format!("{}\n{}", self.artifacts, self.statement), + } + } +} + +/// Append every text part of a parts-bearing object (a message or an +/// artifact), through the caller's bounding `push`. +fn collect_part_text(owner: &Value, buf: &mut String, push: &impl Fn(&mut String, &str)) { + let Some(parts) = owner.get("parts").and_then(Value::as_array) else { + return; + }; + for part in parts { + if let Some(text) = str_field(part, "text") { + if !buf.is_empty() { + push(buf, "\n"); + } + push(buf, text); + } + } +} + /// The 1.0 payload `oneof` field names, in `SendMessageResponse` / /// `StreamResponse` order. const V1_PAYLOAD_KEYS: [&str; 4] = ["task", "message", "statusUpdate", "artifactUpdate"]; @@ -584,6 +720,182 @@ mod tests { assert_eq!(facts.task_state, ""); } + /// The bounded push the proxy passes in; unbounded here so what is under + /// test is the append/replace rule, not the cap. + fn push(buf: &mut String, s: &str) { + buf.push_str(s); + } + + fn observed(events: &[Value]) -> String { + let mut text = ResultText::default(); + for event in events { + text.observe(event, push); + } + text.joined() + } + + fn artifact_chunk(id: &str, text: &str, append: Option) -> Value { + let mut event = json!({"result": {"kind": "artifact-update", "taskId": "t", + "artifact": {"artifactId": id, "parts": [{"text": text}]}}}); + if let Some(append) = append { + event["result"]["append"] = json!(append); + } + event + } + + fn status(state: &str, note: Option<&str>) -> Value { + let mut st = json!({"state": state}); + if let Some(note) = note { + st["message"] = json!({"role": "agent", "parts": [{"text": note}]}); + } + json!({"result": {"kind": "status-update", "taskId": "t", "status": st}}) + } + + #[test] + fn text_is_read_from_parts_in_both_versions() { + // 0.3 tags each part with `kind`, 1.0 sets a protobuf oneof — both + // put the words under `text`, and neither file nor data parts carry + // language worth counting. + let v03 = request_text( + &json!({"params": {"message": {"role": "user", "parts": [ + {"kind": "text", "text": "invoice 42"}, + {"kind": "file", "file": {"bytes": "AAAA", "mimeType": "application/pdf"}}, + {"kind": "text", "text": "please summarise"} + ]}}}), + push, + ); + assert_eq!(v03, "invoice 42\nplease summarise"); + + let v10 = request_text( + &json!({"params": {"message": {"role": "user", "parts": [ + {"text": "invoice 42"}, + {"raw": "AAAA", "mediaType": "application/pdf"}, + {"text": "please summarise"} + ]}}}), + push, + ); + assert_eq!(v10, "invoice 42\nplease summarise"); + + assert_eq!(request_text(&json!({"params": {"id": "t-1"}}), push), ""); + assert_eq!( + request_text( + &json!({"params": {"message": {"parts": [{"data": {"a": 1}}]}}}), + push + ), + "" + ); + } + + #[test] + fn a_progress_note_between_chunks_does_not_erase_the_answer() { + // The reference pattern: artifact chunks with `update_status(working, + // message=…)` interleaved, finished by `complete(message=…)`. A single + // buffer replaced by every statement keeps only the closing note and + // loses the whole report. + let text = observed(&[ + artifact_chunk("report", "the first half", None), + status("working", Some("still working")), + artifact_chunk("report", " and the second", Some(true)), + status("completed", Some("Report generated.")), + ]); + assert!( + text.contains("the first half and the second"), + "the answer survives the notes around it: {text}" + ); + assert!(text.contains("Report generated.")); + assert!( + !text.contains("still working"), + "only the LAST statement is kept: {text}" + ); + } + + #[test] + fn append_is_scoped_to_one_artifact() { + // "appended to a previously sent artifact with the same ID" — a chunk + // for a different artifact opens a new segment instead of discarding + // the one before it. + let text = observed(&[ + artifact_chunk("summary", "the summary", None), + artifact_chunk("table", "the table", None), + ]); + assert!(text.contains("the summary"), "{text}"); + assert!(text.contains("the table"), "{text}"); + } + + #[test] + fn a_resent_artifact_replaces_itself() { + // Without `append` a chunk for the same artifact is a replacement, so + // an agent that resends its answer is not counted twice. + let text = observed(&[ + artifact_chunk("report", "the answer", None), + artifact_chunk("report", "the answer", None), + ]); + assert_eq!(text, "the answer"); + } + + #[test] + fn an_appending_chunk_continues_the_previous_one() { + let text = observed(&[ + artifact_chunk("r", "Hello", Some(false)), + artifact_chunk("r", ", world", Some(true)), + artifact_chunk("r", "!", Some(true)), + ]); + assert_eq!(text, "Hello, world!"); + } + + #[test] + fn a_task_snapshot_restates_its_artifacts_rather_than_adding_to_them() { + // A terminal Task carries the complete set the chunks already + // delivered; adding it would report the answer twice. + let text = observed(&[ + artifact_chunk("r", "the answer", None), + json!({"result": {"kind": "task", "id": "t", "status": {"state": "completed"}, + "artifacts": [{"artifactId": "r", "parts": [{"text": "the answer"}]}]}}), + ]); + assert_eq!(text, "the answer"); + } + + #[test] + fn a_1_0_wrapped_result_yields_its_text_too() { + // The payload wrapper has to be seen through here as well, or the + // default wire version contributes no text at all. + let mut message = ResultText::default(); + message.observe( + &json!({"result": {"message": {"role": "agent", "parts": [{"text": "done"}]}}}), + push, + ); + assert_eq!(message.joined(), "done"); + + let mut streamed = ResultText::default(); + streamed.observe( + &json!({"result": {"statusUpdate": {"taskId": "t", "status": { + "state": "TASK_STATE_COMPLETED", + "message": {"parts": [{"text": "finished"}]} + }}}}), + push, + ); + assert_eq!(streamed.joined(), "finished"); + + let mut artifact = ResultText::default(); + artifact.observe( + &json!({"result": {"artifactUpdate": {"taskId": "t", + "artifact": {"artifactId": "r", "parts": [{"text": "chunk"}]}}}}), + push, + ); + assert_eq!(artifact.joined(), "chunk"); + } + + #[test] + fn an_event_with_no_text_leaves_what_came_before() { + // A bare progress ping must not wipe the answer already collected. + let text = observed(&[ + artifact_chunk("r", "answer", None), + status("working", None), + json!({"error": {"code": -1, "message": "x"}}), + ]); + assert_eq!(text, "answer"); + } + #[test] fn the_terminal_event_is_recognised_in_both_versions() { // The moment a caller stops reading. 0.3 marks it with `final`; 1.0's diff --git a/crates/aisix-proxy/src/a2a.rs b/crates/aisix-proxy/src/a2a.rs index b0194c43..5d7f7e5a 100644 --- a/crates/aisix-proxy/src/a2a.rs +++ b/crates/aisix-proxy/src/a2a.rs @@ -31,10 +31,10 @@ use std::time::{Duration, Instant}; use aisix_a2a::{ - canonical_operation, is_stream_end, is_streaming_operation, upstream_from_a2a_agent, A2aBridge, - A2aCallFacts, A2aError, HttpBridge, + canonical_operation, is_stream_end, is_streaming_operation, request_text, + upstream_from_a2a_agent, A2aBridge, A2aCallFacts, A2aError, HttpBridge, ResultText, }; -use aisix_obs::{AccessLog, UsageEvent}; +use aisix_obs::{content_capture_cap, AccessLog, CapturedContent, UsageEvent}; use axum::body::to_bytes; use axum::extract::{Request, State}; use axum::http::{header, HeaderMap, StatusCode}; @@ -69,6 +69,24 @@ struct A2aCall { /// How the stream behaved, for a streaming call. Left at its default for /// a unary one, which observes none of the stream series. stream: A2aStreamProgress, + /// What was said, for token metering and opt-in content capture. + text: A2aCallText, +} + +/// The words exchanged on one A2A call. +/// +/// An agent reports no token usage — there is no `usage` block in the +/// protocol — so the only way a call can be metered at all is to count what +/// passed through. Both buffers are bounded by +/// [`crate::token_estimate::push_capped`]: a request body has no size limit by +/// default and an A2A task may stream for hours, and both are retained for the +/// life of the call. Past the bound the text becomes a prefix and the estimate +/// a lower bound, rather than the buffer growing without limit. +struct A2aCallText { + /// The caller's message text. + request: String, + /// The agent's, under the protocol's own per-artifact append/replace rule. + response: ResultText, } /// What a streamed A2A call did on the wire, accumulated as events pass. @@ -90,6 +108,18 @@ struct A2aStreamProgress { reached_end: bool, } +/// Whether an operation's words are the agent GENERATING something, and so +/// worth metering and capturing. +/// +/// `message/send` and `message/stream` are the two that make an agent produce. +/// Everything else reads back a task the gateway has already accounted for: +/// `tasks/get` returns the whole task on every poll and `tasks/resubscribe` +/// replays its stream from the start, so counting those would report one +/// answer as many times as a client cared to look at it. +fn meters_content(operation: &str) -> bool { + matches!(operation, "message/send" | "message/stream") +} + /// Serve a JSON-RPC request to `/a2a/:agent`. Authentication (`401`), per-agent /// ACL (`403`), and rate-limit + budget (`429` / budget error) gate the call /// before the request is forwarded to the upstream agent; a usage event is @@ -124,6 +154,9 @@ pub async fn a2a_endpoint( provider: Some("a2a"), model: None, 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, @@ -207,19 +240,38 @@ async fn dispatch( .unwrap_or_default() .to_string(); let rpc_id = value.get("id").cloned(); + let operation = canonical_operation(&method); let mut call = A2aCall { - operation: canonical_operation(&method), + operation, method, protocol_version: upstream.protocol_version.as_wire_str(), facts: A2aCallFacts::default(), stream: A2aStreamProgress::default(), + text: A2aCallText { + // Only the operations that make an agent GENERATE are metered. + // A read (`tasks/get`, `tasks/resubscribe`) hands back the same + // answer on every poll, and counting those would let a client + // polling a ten-minute task report its answer six hundred times — + // swamping exactly the per-agent figures this exists to produce. + request: if meters_content(operation) { + request_text(&value, crate::token_estimate::push_capped) + } else { + String::new() + }, + response: ResultText::default(), + }, }; // Read before the upstream is contacted, so a call that never lands still // records which task the caller was asking about. call.facts.observe_request(&value); // Reuse the LLM path's rate-limit + budget gate. The reservation is held - // for the call (an A2A call carries no token cost yet). On 429 / + // for the call and released without committing tokens: the counts this + // endpoint reports are the gateway's own reading of the words, not an + // agent's billed usage, so they are REPORTED but never CHARGED. Token + // windows and token budgets therefore do not move on A2A traffic, which + // is deliberate — inferring a spend limit from an estimate would throttle + // callers on a number no provider ever confirmed. On 429 / // budget-exceeded this returns before the upstream is contacted. let reservation = match crate::quota::enforce(state, &auth, None).await { Ok(reservation) => reservation, @@ -262,6 +314,11 @@ async fn dispatch( match result { Ok(response_value) => { call.facts.observe_result(&response_value); + if meters_content(call.operation) { + call.text + .response + .observe(&response_value, crate::token_estimate::push_capped); + } emit_a2a_usage( state, &auth, @@ -401,7 +458,11 @@ async fn dispatch_stream( }; let agent_label = agent.to_string(); - let sse = async_stream::stream! { + // Re-attach the request span: the body is polled after the request-id + // middleware has returned, so a mid-stream failure would otherwise be + // logged without its `request_id` and could not be joined to the rest of + // the request (AISIX-Cloud#1060). + let sse = crate::request_id::in_request_span(async_stream::stream! { let mut events = events; while let Some(event) = events.next().await { match event { @@ -416,6 +477,13 @@ async fn dispatch_stream( // inflate the event counter and put a near-zero // observation into the time-to-first-event histogram, // dragging the percentile of real streams down with it. + if meters_content(guard.call.operation) { + guard + .call + .text + .response + .observe(&value, crate::token_estimate::push_capped); + } if value.get("result").is_some() { guard.call.stream.event_count += 1; // The agent's own time to first byte. Stamped on the @@ -456,7 +524,7 @@ async fn dispatch_stream( // finished. Either way the caller was handed everything there was. guard.call.stream.reached_end = true; drop(guard); - }; + }); let mut response = axum::response::Sse::new(sse); if let Some(interval) = crate::sse_keepalive::interval() { @@ -601,8 +669,16 @@ fn a2a_error_envelope(id: Option, message: &str) -> serde_jso } /// Emit a usage event for a single A2A call into the same sink as LLM usage. -/// A2A calls carry no token cost yet, so token/cost fields stay zero; the event -/// records who called which agent with which method, the outcome, and latency. +/// +/// The event records who called which agent with which operation, which task +/// it touched, the outcome, and latency. Token counts are the gateway's own +/// reading of the words that passed through — an agent reports none of its +/// own — and are flagged `usage_estimated`; `cost_usd` stays zero, since what +/// an agent charges is not something the gateway can know. +/// +/// This is the chokepoint every A2A path emits through, so the metric +/// families ride here too: a path that accounts for a call cannot skip +/// metering it. fn emit_a2a_usage( state: &ProxyState, auth: &AuthenticatedKey, @@ -612,6 +688,11 @@ fn emit_a2a_usage( status_code: u16, latency: Duration, ) { + // No model resolves on this endpoint, so the estimator falls back to its + // default encoding — the same thing it does for any non-OpenAI model. + let response_text = call.text.response.joined(); + let prompt_tokens = crate::token_estimate::count_text("", &call.text.request); + let completion_tokens = crate::token_estimate::count_text("", &response_text); let event = UsageEvent { request_id: request_id.to_string(), occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), @@ -630,6 +711,17 @@ fn emit_a2a_usage( a2a_context_id: call.facts.context_id.clone(), a2a_task_state: call.facts.task_state.to_string(), a2a_stream_event_count: call.stream.event_count, + // An A2A agent reports no usage of its own, so these are the + // gateway's own count of the words that passed through — flagged as + // estimated, which is exactly what `usage_estimated` is for. Cost + // stays zero: what an agent charges is not something the gateway can + // know, and inventing a number would be worse than reporting none. + prompt_tokens, + completion_tokens, + // Set on every A2A row, not only the counted ones: a zero here is + // the gateway's own count too, so a consumer filtering for + // provider-billed exactness must not pick these up as exact. + usage_estimated: true, upstream_ttft_ms: call .stream .ttfb @@ -671,9 +763,19 @@ fn emit_a2a_usage( state.usage_sink.try_emit("a2a", event.clone()); let snap = state.snapshot.load(); let exporters = snap.observability_exporters.entries(); - state - .otlp_fan_out - .fan_out(&event, None, exporters.iter().map(|e| &e.value)); + // Opt-in content capture, on the same terms as every other endpoint: only + // an exporter configured for full content sees the words, and they never + // travel to the control plane — the usage event above carries counts + // only. The captured text is the message parts, not the JSON-RPC + // envelopes, so it reads as a prompt and a completion rather than as + // protocol scaffolding. + let captured = content_capture_cap(exporters.iter().map(|e| &e.value)) + .map(|cap| CapturedContent::new(&call.text.request, &response_text, cap as usize)); + state.otlp_fan_out.fan_out( + &event, + captured.as_ref(), + exporters.iter().map(|e| &e.value), + ); } #[cfg(test)] @@ -1194,6 +1296,114 @@ mod tests { assert_eq!(event.a2a_protocol_version, "1.0"); } + #[tokio::test] + async fn a_call_is_metered_from_the_words_that_passed_through() { + // An A2A agent reports no usage of its own, so a call that is not + // counted here is not counted anywhere: every agent's spend looks + // identical and zero. + let event = usage_event_for( + &spawn_task_agent().await, + serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "message/send", + "params": {"message": {"role": "user", "parts": [{"kind": "text", + "text": "summarise invoice 42"}]}} + }), + ) + .await; + + assert!(event.prompt_tokens > 0, "the caller's words are counted"); + assert!( + event.usage_estimated, + "the gateway counted these, not the agent — the flag says so" + ); + // What an agent charges is not something the gateway can know. + assert_eq!(event.cost_usd, 0.0); + } + + /// An agent that answers with words rather than a bare task record. + async fn spawn_talking_agent() -> String { + let app = axum::Router::new().route( + "/a2a", + axum::routing::post(|body: axum::Json| async move { + axum::Json(serde_json::json!({ + "jsonrpc": "2.0", + "id": body.0["id"], + "result": { + "kind": "message", + "messageId": "m-1", + "role": "agent", + "parts": [{"kind": "text", "text": "The invoice totals four hundred."}], + } + })) + }), + ); + 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}/a2a") + } + + #[tokio::test] + async fn the_agents_own_words_are_counted_too() { + // Prompt-side only would report every agent as producing nothing. + let event = usage_event_for( + &spawn_talking_agent().await, + serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "message/send", + "params": {"message": {"role": "user", "parts": [{"text": "how much?"}]}} + }), + ) + .await; + + assert!(event.prompt_tokens > 0); + assert!(event.completion_tokens > 0, "the agent's reply is counted"); + assert!(event.usage_estimated); + } + + #[tokio::test] + async fn a_call_with_nothing_to_count_reports_no_tokens() { + // A task lookup carries no words at all. Reporting a token count for + // it would be inventing one. + let event = usage_event_for( + &spawn_task_agent().await, + serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "tasks/get", "params": {"id": "t-1"} + }), + ) + .await; + + assert_eq!(event.prompt_tokens, 0); + assert_eq!(event.completion_tokens, 0); + // The flag rides every A2A row, zero included: the zero is the + // gateway's own count too, and a consumer filtering for + // provider-billed exactness must not mistake it for one. + assert!(event.usage_estimated); + } + + #[tokio::test] + async fn reading_a_task_back_does_not_re_meter_its_answer() { + // `tasks/get` returns the whole task on every poll. Counting it would + // let a client polling a long task report one answer as many times as + // it cared to look, swamping the per-agent figures this exists for. + let event = usage_event_for( + &spawn_talking_agent().await, + serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "tasks/get", "params": {"id": "t-1"} + }), + ) + .await; + + assert_eq!(event.a2a_operation, "tasks/get"); + assert_eq!( + event.completion_tokens, 0, + "a read re-states an answer already counted" + ); + } + #[tokio::test] async fn a_1_0_caller_aggregates_with_its_0_3_twin() { // `SendMessage` and `message/send` are one operation. A gateway may diff --git a/tests/e2e/src/cases/a2a-protocol-telemetry-e2e.test.ts b/tests/e2e/src/cases/a2a-protocol-telemetry-e2e.test.ts index bf3a710a..d993bad5 100644 --- a/tests/e2e/src/cases/a2a-protocol-telemetry-e2e.test.ts +++ b/tests/e2e/src/cases/a2a-protocol-telemetry-e2e.test.ts @@ -10,6 +10,7 @@ import { type SpawnedApp, } from "../harness/index.js"; import { startMockOtlp, type MockOtlp } from "../harness/otlp-mock.js"; +import { STREAM_ANSWER } from "../harness/upstream-a2a.js"; // E2E for AISIX-Cloud#1215: protocol-level observability for the A2A gateway. // @@ -38,6 +39,7 @@ describe("a2a protocol telemetry e2e (AISIX-Cloud#1215)", () => { let upstream10: A2aUpstream | undefined; let upstream03: A2aUpstream | undefined; let otlp: MockOtlp | undefined; + let otlpFull: MockOtlp | undefined; let etcdReachable = false; const call = async (agent: string, body: unknown) => { @@ -77,8 +79,17 @@ describe("a2a protocol telemetry e2e (AISIX-Cloud#1215)", () => { if (!etcdReachable) return; otlp = await startMockOtlp(); - upstream10 = await startA2aUpstream({ cardMount: "origin", wireShape: "1.0" }); - upstream03 = await startA2aUpstream({ cardMount: "origin", wireShape: "0.3" }); + otlpFull = await startMockOtlp(); + upstream10 = await startA2aUpstream({ + cardMount: "origin", + wireShape: "1.0", + streamAnswer: true, + }); + upstream03 = await startA2aUpstream({ + cardMount: "origin", + wireShape: "0.3", + streamAnswer: true, + }); app = await spawnApp(); const seed = new SeedClient(etcd, app.etcdPrefix); @@ -88,6 +99,16 @@ describe("a2a protocol telemetry e2e (AISIX-Cloud#1215)", () => { kind: "otlp_http", endpoint: otlp.url, }); + // A second exporter that opts into content. Both receive every event, so + // the pair is what proves capture is opt-in rather than merely present: + // the words must reach this one and no other. + await seed.createObservabilityExporter({ + name: "a2a-telemetry-otlp-full", + enabled: true, + kind: "otlp_http", + endpoint: otlpFull.url, + content_mode: "full", + }); // One agent per wire version, each against the stub that speaks it. Both // run the same operations, so what has to come out the same (the canonical // operation) and what has to differ (the announced version) are both @@ -127,10 +148,12 @@ describe("a2a protocol telemetry e2e (AISIX-Cloud#1215)", () => { await upstream10?.close(); await upstream03?.close(); await otlp?.close(); + await otlpFull?.close(); }); test("a completed call records the task, the context and how it ended", async (ctx) => { - if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp) return ctx.skip(); + if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp || !otlpFull) + return ctx.skip(); const contextId = `ctx-${randomUUID()}`; expect( @@ -156,7 +179,8 @@ describe("a2a protocol telemetry e2e (AISIX-Cloud#1215)", () => { }); test("both wire vocabularies aggregate under one operation", async (ctx) => { - if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp) return ctx.skip(); + if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp || !otlpFull) + return ctx.skip(); // The same operation, spelled the way each agent's version spells it. const v10Context = `ctx-${randomUUID()}`; @@ -200,7 +224,8 @@ describe("a2a protocol telemetry e2e (AISIX-Cloud#1215)", () => { }); test("a streamed task is recorded with the state its stream ended on", async (ctx) => { - if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp) return ctx.skip(); + if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp || !otlpFull) + return ctx.skip(); const contextId = `ctx-${randomUUID()}`; expect( @@ -222,7 +247,8 @@ describe("a2a protocol telemetry e2e (AISIX-Cloud#1215)", () => { }); test("a streamed call records how it ran, not just how it ended", async (ctx) => { - if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp) return ctx.skip(); + if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp || !otlpFull) + return ctx.skip(); // The stub paces its three events apart, so a wait for the first one is // separable from the total: a stream that is slow to start and one that is @@ -238,7 +264,7 @@ describe("a2a protocol telemetry e2e (AISIX-Cloud#1215)", () => { ).toBe(200); const span = await awaitSpan((s) => s.attributes["gen_ai.conversation.id"] === contextId); - expect(span.attributes["aisix.a2a.stream_event_count"]).toBe(3); + expect(span.attributes["aisix.a2a.stream_event_count"]).toBe(6); // The agent's own time to first event. The stub pauses before each of its // three events, so a figure that actually stopped at the first one is a // fraction of the call; one that quietly measured the whole stream would @@ -258,7 +284,8 @@ describe("a2a protocol telemetry e2e (AISIX-Cloud#1215)", () => { }); test("the a2a metric family slices by agent and operation", async (ctx) => { - if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp) return ctx.skip(); + if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp || !otlpFull) + return ctx.skip(); // `aisix_proxy_requests_total` already counts these calls, but only by // route — it cannot answer "is the invoices agent's stream failing?". @@ -302,8 +329,83 @@ describe("a2a protocol telemetry e2e (AISIX-Cloud#1215)", () => { expect(text).not.toMatch(/aisix_a2a_[a-z_]*\{[^}]*context_id=/); }); + test("the words are metered always and captured only on request", async (ctx) => { + if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp || !otlpFull) + return ctx.skip(); + + const contextId = `ctx-${randomUUID()}`; + const prompt = "summarise invoice forty two for the finance team"; + expect( + await call("invoices", { + jsonrpc: "2.0", + id: 10, + method: "message/send", + params: { + message: { role: "user", contextId, parts: [{ kind: "text", text: prompt }] }, + }, + }), + ).toBe(200); + + // An agent reports no usage of its own, so without the gateway counting + // them every agent's spend would read as identical and zero. + const metered = await awaitSpan((s) => s.attributes["gen_ai.conversation.id"] === contextId); + expect(metered.attributes["gen_ai.usage.input_tokens"]).toBeGreaterThan(0); + + // The words themselves reach only the exporter that asked for them. + const deadline = Date.now() + EXPORT_TIMEOUT_MS; + let captured: (typeof otlpFull.spans)[number] | undefined; + while (Date.now() < deadline && !captured) { + captured = otlpFull.spans.find((s) => s.attributes["gen_ai.conversation.id"] === contextId); + if (!captured) await new Promise((resolve) => setTimeout(resolve, 250)); + } + expect(captured, "the full-content exporter received the call").toBeDefined(); + expect(captured!.attributes["gen_ai.prompt"]).toContain(prompt); + expect(captured!.attributes["gen_ai.completion"]).toContain("The invoice is settled."); + // ...and never the default one, whatever else it carries. + expect(metered.attributes["gen_ai.prompt"]).toBeUndefined(); + expect(metered.attributes["gen_ai.completion"]).toBeUndefined(); + }); + + test("a streamed answer survives the progress notes around it", async (ctx) => { + if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp || !otlpFull) + return ctx.skip(); + + // The agent delivers its answer as two continued artifact chunks with a + // progress note between them and a closing statement after — the shape a + // reference agent produces. An accumulator that treats every statement as + // a replacement keeps only "Report generated." and loses the report. + const contextId = `ctx-${randomUUID()}`; + expect( + await call("invoices", { + jsonrpc: "2.0", + id: 11, + method: "message/stream", + params: { message: { role: "user", contextId, parts: [{ kind: "text", text: "report?" }] } }, + }), + ).toBe(200); + + const deadline = Date.now() + EXPORT_TIMEOUT_MS; + let captured: (typeof otlpFull.spans)[number] | undefined; + while (Date.now() < deadline && !captured) { + captured = otlpFull.spans.find((s) => s.attributes["gen_ai.conversation.id"] === contextId); + if (!captured) await new Promise((resolve) => setTimeout(resolve, 250)); + } + expect(captured, "the full-content exporter received the stream").toBeDefined(); + + const completion = String(captured!.attributes["gen_ai.completion"]); + expect(completion).toContain(STREAM_ANSWER.join("")); + expect(completion).toContain("Report generated."); + expect(completion).not.toContain("halfway"); + // The answer appears once, not once per chunk that restated it. + expect(completion.split(STREAM_ANSWER[0]).length - 1).toBe(1); + // And it is counted, so a long answer is not metered as a short one. + const metered = await awaitSpan((s) => s.attributes["gen_ai.conversation.id"] === contextId); + expect(metered.attributes["gen_ai.usage.output_tokens"]).toBeGreaterThan(5); + }); + test("an unrecognised method cannot become an unbounded label", async (ctx) => { - if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp) return ctx.skip(); + if (!etcdReachable || !app || !upstream10 || !upstream03 || !otlp || !otlpFull) + return ctx.skip(); // The method is caller-chosen. The raw value stays available for // forensics, but the aggregating field must collapse to `unknown`. diff --git a/tests/e2e/src/harness/upstream-a2a.ts b/tests/e2e/src/harness/upstream-a2a.ts index 81ccfb38..d6e35cc7 100644 --- a/tests/e2e/src/harness/upstream-a2a.ts +++ b/tests/e2e/src/harness/upstream-a2a.ts @@ -54,6 +54,13 @@ export interface A2aUpstreamOptions { token?: string; /** Wire shape of the results this agent returns. Defaults to `0.3`. */ wireShape?: A2aWireShape; + /** + * Stream an actual answer: [`STREAM_ANSWER`] as two continued artifact + * chunks with a progress note between them and a closing statement — the + * shape a reference agent produces. Off by default, because it changes the + * event sequence the streaming-relay suite counts. + */ + streamAnswer?: boolean; } const PATH_PREFIX = "/v3/agents/serve/tenant-42"; @@ -68,6 +75,16 @@ const PATH_PREFIX = "/v3/agents/serve/tenant-42"; */ const STREAM_GAP_MS = 120; +/** + * The two halves of the streamed answer under `streamAnswer`, delivered as one + * artifact continued across two chunks. Exported so a test can assert the whole + * of it survived the progress note and the terminal statement around it. + */ +export const STREAM_ANSWER = [ + "The invoice totals four hundred and twenty euros", + ", payable within thirty days of receipt.", +] as const; + /** JSON-RPC methods this stub answers with an SSE stream, in both spellings. */ const STREAMING_METHODS = new Set([ "message/stream", @@ -103,6 +120,7 @@ export async function startA2aUpstream( cardPath, token: options.token, wireShape: options.wireShape ?? "0.3", + streamAnswer: options.streamAnswer ?? false, }).catch((err: unknown) => { // `handle` rejects on a malformed body, and on a write to an already // closed socket. Unhandled, that terminates the test process; worse, the @@ -138,6 +156,7 @@ async function handle( cardPath: string; token?: string; wireShape: A2aWireShape; + streamAnswer: boolean; }, ): Promise { const path = new URL(req.url ?? "/", "http://127.0.0.1").pathname; @@ -180,13 +199,16 @@ async function handle( * flat under `result` with a `kind` tag on 0.3, inside the response's * payload wrapper on 1.0. */ + const V1_PAYLOAD_KEY = { + task: "task", + "status-update": "statusUpdate", + "artifact-update": "artifactUpdate", + } as const; const payload = ( - kind: "task" | "status-update", + kind: keyof typeof V1_PAYLOAD_KEY, obj: Record, ): Record => - ctx.wireShape === "1.0" - ? { [kind === "task" ? "task" : "statusUpdate"]: obj } - : { kind, ...obj }; + ctx.wireShape === "1.0" ? { [V1_PAYLOAD_KEY[kind]]: obj } : { kind, ...obj }; if (ctx.token !== undefined && header("authorization") !== `Bearer ${ctx.token}`) { send(401, { error: "unauthorized" }); @@ -250,13 +272,59 @@ async function handle( res.write( envelope(payload("status-update", { taskId: "task-e2e-stream", contextId, seq: 2 })), ); + if (ctx.streamAnswer) { + // The answer in two continued chunks with a progress note between them: + // an accumulator that lets a note replace the answer keeps only the note. + res.write( + envelope( + payload("artifact-update", { + taskId: "task-e2e-stream", + contextId, + artifact: { + artifactId: "report", + parts: [{ kind: "text", text: STREAM_ANSWER[0] }], + }, + }), + ), + ); + res.write( + envelope( + payload("status-update", { + taskId: "task-e2e-stream", + contextId, + status: { + state: "working", + message: { role: "agent", parts: [{ kind: "text", text: "halfway" }] }, + }, + }), + ), + ); + res.write( + envelope( + payload("artifact-update", { + taskId: "task-e2e-stream", + contextId, + append: true, + artifact: { + artifactId: "report", + parts: [{ kind: "text", text: STREAM_ANSWER[1] }], + }, + }), + ), + ); + } await new Promise((resolve) => setTimeout(resolve, STREAM_GAP_MS)); res.write( envelope( payload("task", { id: "task-e2e-stream", contextId, - status: { state: "completed" }, + status: ctx.streamAnswer + ? { + state: "completed", + message: { role: "agent", parts: [{ kind: "text", text: "Report generated." }] }, + } + : { state: "completed" }, seq: 3, }), ), @@ -272,7 +340,10 @@ async function handle( result: payload("task", { id: "task-e2e-1", contextId, - status: { state: "completed" }, + status: { + state: "completed", + message: { role: "agent", parts: [{ kind: "text", text: "The invoice is settled." }] }, + }, // Echoed so the gateway's forwarding can be asserted from the caller // side as well as from `requests`. sawVersion: header("a2a-version"),