klieo-a2a 2.2.0

Durable A2A v1.0 protocol layer atop klieo-bus traits.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

//! Cluster-0.23 tracecontext propagation across the A2A bus.
//!
//! One dispatcher publishes a `TaskEvent` inside a known root span
//! while a parallel `Pubsub::subscribe` consumer pulls the raw `Msg`
//! off the bus. The test asserts:
//!
//! 1. The publisher injected the W3C `traceparent` header into the
//!    bus `Headers` bag (the [`klieo_core::inject_traceparent`] seam
//!    fires from inside [`TaskEventSink::send`]).
//! 2. The extracted [`opentelemetry::Context`] carries a valid span
//!    context (so a subscriber's call to `Span::set_parent` would
//!    stitch under the publisher's span).
//! 3. The exported publisher span shares its `trace_id` with the
//!    test root span — confirming the `tracing-opentelemetry`
//!    bridge propagates the OTEL context through `#[instrument]`.

#![cfg(feature = "test-fixtures")]

use klieo_a2a::server::{A2aDispatcher, TaskEvent};
use klieo_a2a::types::TaskStatus;
use klieo_a2a::EchoHandler;
use klieo_auth_common::AllowAnonymous;
use klieo_auth_common::Authenticator;
use klieo_bus_memory::MemoryBus;
use klieo_core::{DurableName, Pubsub};
use opentelemetry::global;
use opentelemetry::trace::TraceContextExt as _;
use opentelemetry_sdk::testing::trace::InMemorySpanExporter;
use opentelemetry_sdk::trace::TracerProvider;
use std::sync::Arc;
use std::time::Duration;
use tokio_stream::StreamExt as _;
use tracing_subscriber::layer::SubscriberExt as _;
use tracing_subscriber::Registry;

#[tokio::test]
async fn publisher_span_parents_subscriber_span_across_replicas() {
    let exporter = InMemorySpanExporter::default();
    let provider = TracerProvider::builder()
        .with_simple_exporter(exporter.clone())
        .build();
    let tracer = {
        use opentelemetry::trace::TracerProvider as _;
        provider.tracer("klieo-test")
    };
    global::set_tracer_provider(provider.clone());

    let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer);
    let subscriber = Registry::default().with(otel_layer);
    let _guard = tracing::subscriber::set_default(subscriber);

    let bus = MemoryBus::new();
    let pubsub: Arc<dyn Pubsub> = bus.pubsub.clone();
    let handler: Arc<dyn klieo_a2a::handler::A2aHandler> = Arc::new(EchoHandler::default());
    let auth: Arc<dyn Authenticator> = Arc::new(AllowAnonymous);

    let dispatcher = A2aDispatcher::new(handler, auth, pubsub.clone());
    let sink = dispatcher.event_sink();

    // Subscribe before publish so we deterministically receive the message.
    let durable = DurableName::new("trace-stitch-test");
    let mut stream = pubsub
        .subscribe("klieo.a2a.task.t-trace", durable)
        .await
        .unwrap();

    // Publish inside a root span — `TaskEventSink::send` runs under
    // `#[instrument]` and inherits this as parent context.
    let test_root_trace_id = {
        let root = tracing::info_span!("test_root");
        let _enter = root.enter();
        let event = TaskEvent::new("t-trace", TaskStatus::Working, None, false).with_event_id(1);
        sink.send(event).await.unwrap();
        // Capture the trace id while we're still inside the span so
        // the assertion has a stable reference.
        use tracing_opentelemetry::OpenTelemetrySpanExt as _;
        root.context().span().span_context().trace_id()
    };

    let msg = tokio::time::timeout(Duration::from_millis(500), stream.next())
        .await
        .expect("subscribe timeout")
        .expect("stream ended")
        .expect("bus error");

    assert!(
        msg.headers.contains_key("traceparent"),
        "publisher must inject W3C traceparent into bus headers; got headers: {:?}",
        msg.headers.keys().collect::<Vec<_>>(),
    );

    let extracted_cx = klieo_core::extract_traceparent(&msg.headers);
    let span_ctx = extracted_cx.span().span_context().clone();
    assert!(
        span_ctx.is_valid(),
        "extracted span context must be valid (trace_id={:?}, span_id={:?})",
        span_ctx.trace_id(),
        span_ctx.span_id(),
    );

    msg.ack.ack().await.unwrap();

    // Flush exported spans and verify the publisher span shares the
    // root's trace id (the cross-replica stitch invariant).
    provider.force_flush();
    let spans = exporter.get_finished_spans().unwrap();
    let publisher_span = spans
        .iter()
        .find(|s| s.name == "send")
        .or_else(|| spans.iter().find(|s| s.name.contains("send")))
        .unwrap_or_else(|| {
            panic!(
                "expected publisher span named `send`; got: {:?}",
                spans.iter().map(|s| &s.name).collect::<Vec<_>>(),
            )
        });
    assert_eq!(
        publisher_span.span_context.trace_id(),
        test_root_trace_id,
        "publisher span must inherit test_root's trace_id for cross-replica stitch",
    );
}