feat(a2a): observe how a streamed call ran, and meter calls per agent

crates/aisix-a2a/src/lib.rs

@@ -27,4 +27,4 @@ pub use bridge::{
     HttpBridge, DEFAULT_UPSTREAM_TIMEOUT,
 };
 pub use error::A2aError;
-pub use telemetry::{canonical_operation, is_streaming_operation, A2aCallFacts};
+pub use telemetry::{canonical_operation, is_stream_end, is_streaming_operation, A2aCallFacts};

crates/aisix-a2a/src/telemetry.rs

@@ -151,6 +151,45 @@ pub fn normalize_task_state(state: &str) -> &'static str {
         .unwrap_or(UNKNOWN_TASK_STATE)
 }
 
+/// Task states that mean the agent has said everything it is going to say on
+/// this stream: the task is finished, or it is waiting on the caller.
+const STREAM_ENDING_STATES: &[&str] = &[
+    "completed",
+    "canceled",
+    "failed",
+    "rejected",
+    "input-required",
+    "auth-required",
+];
+
+/// Whether a streamed event is the last one its stream will carry.
+///
+/// A caller stops reading here, which makes this — not the upstream's
+/// eventual EOF — the moment the response was fully delivered. Anything that
+/// waits for the connection to close instead reports a completed task as
+/// abandoned whenever the agent leaves the stream open afterwards.
+///
+/// 0.3 says so with `final`. 1.0's `TaskStatusUpdateEvent` has no such field
+/// at all, so there the task's own state is the only signal: a task that is
+/// neither `submitted` nor `working` has nothing further to stream until the
+/// caller acts.
+pub fn is_stream_end(event: &Value) -> bool {
+    let Some(result) = event.get("result") else {
+        return false;
+    };
+    let payload = unwrap_payload(result);
+    if payload.get("final").and_then(Value::as_bool) == Some(true) {
+        return true;
+    }
+    let Some(state) = payload
+        .get("status")
+        .and_then(|status| str_field(status, "state"))
+    else {
+        return false;
+    };
+    STREAM_ENDING_STATES.contains(&normalize_task_state(state))
+}
+
 /// What one A2A call touched, accumulated as its request and response(s) are
 /// seen.
 ///
@@ -169,7 +208,12 @@ pub struct A2aCallFacts {
     /// Last task state reported, in its 0.3 spelling. Empty when no response
     /// carried one — a state is never invented for a call that failed before
     /// the upstream answered.
-    pub task_state: String,
+    ///
+    /// `&'static str` rather than a `String`: the value is always one of the
+    /// specification's states or `unknown`, and typing it that way is what
+    /// lets it be used as a metric label without a caller having to promise
+    /// it is bounded.
+    pub task_state: &'static str,
 }
 
 impl A2aCallFacts {
@@ -223,7 +267,7 @@ impl A2aCallFacts {
         if let Some(status) = payload.get("status") {
             self.set_task_id(str_field(payload, "id"));
             if let Some(state) = str_field(status, "state") {
-                self.task_state = normalize_task_state(state).to_string();
+                self.task_state = normalize_task_state(state);
             }
         }
     }
@@ -540,6 +584,38 @@ mod tests {
         assert_eq!(facts.task_state, "");
     }
 
+    #[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
+        // status-update has no such field, so only the task's own state says
+        // so — and both spellings of that state have to work.
+        for terminal in [
+            json!({"result": {"kind": "status-update", "taskId": "t", "final": true,
+                              "status": {"state": "input-required"}}}),
+            json!({"result": {"kind": "task", "id": "t", "status": {"state": "completed"}}}),
+            json!({"result": {"statusUpdate": {"taskId": "t",
+                              "status": {"state": "TASK_STATE_FAILED"}}}}),
+            json!({"result": {"task": {"id": "t", "status": {"state": "TASK_STATE_CANCELED"}}}}),
+        ] {
+            assert!(is_stream_end(&terminal), "{terminal} ends the stream");
+        }
+        // A task still running does not, nor does an event with no state at
+        // all (an artifact chunk), nor an error envelope.
+        for ongoing in [
+            json!({"result": {"kind": "status-update", "taskId": "t", "final": false,
+                              "status": {"state": "working"}}}),
+            json!({"result": {"statusUpdate": {"taskId": "t",
+                              "status": {"state": "TASK_STATE_SUBMITTED"}}}}),
+            json!({"result": {"artifactUpdate": {"taskId": "t"}}}),
+            json!({"error": {"code": -32000, "message": "boom"}}),
+        ] {
+            assert!(
+                !is_stream_end(&ongoing),
+                "{ongoing} does not end the stream"
+            );
+        }
+    }
+
     #[test]
     fn an_empty_id_never_erases_one_already_seen() {
         let mut facts = A2aCallFacts::default();

crates/aisix-core/src/config.rs

@@ -789,6 +789,13 @@ pub struct HistogramBucketsConfig {
     pub request_ttft: Option<Vec<f64>>,
     /// `aisix_guardrail_latency_seconds`
     pub guardrail_latency: Option<Vec<f64>>,
+    /// `aisix_a2a_ttfb_seconds`
+    ///
+    /// Separate from `request_ttft` on purpose: an agent's wait for its first
+    /// event and a model's wait for its first token have the same shape but
+    /// not the same range — an A2A task may think for minutes before it says
+    /// anything. Defaults to the same edges as `request_ttft`.
+    pub a2a_ttfb: Option<Vec<f64>>,
 }
 
 /// One `client_type_rules` entry: a regex tried against the raw inbound
@@ -1307,6 +1314,7 @@ impl Config {
                 .with_list_parse_key("observability.metrics.buckets.request_e2e_latency")
                 .with_list_parse_key("observability.metrics.buckets.request_ttft")
                 .with_list_parse_key("observability.metrics.buckets.guardrail_latency")
+                .with_list_parse_key("observability.metrics.buckets.a2a_ttfb")
                 .try_parsing(true),
         );
 

crates/aisix-obs/src/lib.rs

@@ -29,9 +29,9 @@ use tracing_subscriber::{fmt, prelude::*, EnvFilter};
 
 pub use access_log::AccessLog;
 pub use metrics::{
-    client_type_from_user_agent, BudgetGauges, BudgetLabels, ClientTypeClassifier,
-    DeploymentLabels, DeploymentState, HistogramBuckets, LatencyLabels, LlmUsage, Metrics,
-    RequestLabels, RequestOutcome, UsageLabels,
+    client_type_from_user_agent, A2aCallOutcome, A2aLabels, BudgetGauges, BudgetLabels,
+    ClientTypeClassifier, DeploymentLabels, DeploymentState, HistogramBuckets, LatencyLabels,
+    LlmUsage, Metrics, RequestLabels, RequestOutcome, UsageLabels,
 };
 pub use otlp::{install_otlp_tracer, shutdown_otlp, OtlpError, OtlpHandle};
 pub use otlp_http_sink::{content_capture_cap, OtlpHttpFanOut, OtlpSink};