Update axum, tracing integrations and tests for groups and named payloads (#62)

Integration layer and test updates:
- Axum HTTP observer uses groups for request/response hierarchy
- Tracing layer uses groups for span-based observation grouping
- All tests updated for multi-payload and group APIs

Author: Dig-Doug

crates/observation-tools-client/src/axum/request_observer.rs

@@ -1,6 +1,7 @@
 use crate::context;
-use crate::execution::ExecutionHandle;
+use crate::group::GroupBuilder;
 use crate::observation::ObservationBuilder;
+use crate::observation_handle::ObservationPayloadHandle;
 use axum::body::Body;
 use axum::extract::Request;
 use axum::response::Response;
@@ -28,30 +29,21 @@ use tower::Layer;
 use tower::Service;
 
 /// State shared between the streaming body and the observation emitter.
-/// This is used to collect data as it streams and emit the observation when
-/// complete. The observation is emitted in the Drop implementation when the
-/// body is finished.
+/// This is used to collect data as it streams and add the body payload when
+/// complete. The body payload is added in the Drop implementation via the
+/// ObservationPayloadHandle.
 struct StreamingObserverState {
   buffer: BytesMut,
   content_type: String,
-  log_level: LogLevel,
-  status_code: u16,
-  execution: ExecutionHandle,
+  payload_handle: ObservationPayloadHandle,
 }
 
 impl StreamingObserverState {
-  fn new(
-    content_type: String,
-    log_level: LogLevel,
-    status_code: u16,
-    execution: ExecutionHandle,
-  ) -> Self {
+  fn new(content_type: String, payload_handle: ObservationPayloadHandle) -> Self {
     Self {
       buffer: BytesMut::new(),
       content_type,
-      log_level,
-      status_code,
-      execution,
+      payload_handle,
     }
   }
 
@@ -64,7 +56,7 @@ impl Drop for StreamingObserverState {
   fn drop(&mut self) {
     let bytes = self.buffer.clone().freeze();
     tracing::debug!(
-      "StreamingObserverBody: emitting observation with {} bytes on drop",
+      "StreamingObserverBody: adding body payload with {} bytes on drop",
       bytes.len()
     );
     let payload = Payload {
@@ -73,19 +65,14 @@ impl Drop for StreamingObserverState {
       size: bytes.len(),
     };
 
-    ObservationBuilder::new("http/response/body")
-      .label("http/response")
-      .label("http/response/body")
-      .metadata("status", &self.status_code.to_string())
-      .log_level(self.log_level)
-      .execution(&self.execution)
-      .payload(payload);
+    self.payload_handle.raw_payload("body", payload);
   }
 }
 
 pin_project! {
   /// A body wrapper that streams data through while collecting it for observation.
-  /// When the stream completes, it emits the observation with the collected body.
+  /// When the stream completes, it adds the body as a named payload to the
+  /// existing response observation.
   pub struct StreamingObserverBody {
     #[pin]
     inner: Body,
@@ -94,20 +81,12 @@ pin_project! {
 }
 
 impl StreamingObserverBody {
-  fn new(
-    inner: Body,
-    content_type: String,
-    log_level: LogLevel,
-    status_code: u16,
-    execution: ExecutionHandle,
-  ) -> Self {
+  fn new(inner: Body, content_type: String, payload_handle: ObservationPayloadHandle) -> Self {
     Self {
       inner,
       state: Arc::new(Mutex::new(StreamingObserverState::new(
         content_type,
-        log_level,
-        status_code,
-        execution,
+        payload_handle,
       ))),
     }
   }
@@ -247,27 +226,33 @@ where
       };
 
       let (parts, body) = req.into_parts();
-      ObservationBuilder::new("http/request/headers")
-        .label("http/request")
-        .label("http/request/headers")
+
+      // Create a single group for the HTTP exchange
+      let http_group = GroupBuilder::new("http_request")
         .metadata("method", parts.method.to_string())
         .metadata("uri", parts.uri.to_string())
-        .serde(&json!(filter_headers(
-          &parts.headers,
-          &config.excluded_headers
-        )));
+        .build_with_execution(&execution)
+        .into_handle();
 
+      // Collect request body
       let request_body_bytes = body
         .collect()
         .await
         .map(|collected| collected.to_bytes())
         .unwrap_or_else(|_| Bytes::new());
-      ObservationBuilder::new("http/request/body")
-        .label("http/request")
-        .label("http/request/body")
+
+      // Single request observation with named payloads: "headers" + "body"
+      let headers_json = json!(filter_headers(&parts.headers, &config.excluded_headers));
+      let headers_payload = Payload::json(
+        serde_json::to_string(&headers_json).unwrap_or_default(),
+      );
+      ObservationBuilder::new("http/request")
+        .group(&http_group)
         .metadata("method", parts.method.to_string())
         .metadata("uri", parts.uri.to_string())
-        .payload(bytes_to_payload(&request_body_bytes, &parts.headers));
+        .execution(&execution)
+        .named_raw_payload("headers", headers_payload)
+        .raw_payload("body", bytes_to_payload(&request_body_bytes, &parts.headers));
 
       let response = inner
         .call(Request::from_parts(parts, Body::from(request_body_bytes)))
@@ -280,31 +265,30 @@ where
         500..=599 => LogLevel::Error,
         _ => LogLevel::Info,
       };
-      ObservationBuilder::new("http/response/headers")
-        .label("http/response")
-        .label("http/response/headers")
+
+      // Single response observation with named payload "headers" sent immediately,
+      // "body" added later when streaming completes via the payload handle
+      let resp_headers_json = json!(filter_headers(&parts.headers, &config.excluded_headers));
+      let resp_headers_payload = Payload::json(
+        serde_json::to_string(&resp_headers_json).unwrap_or_default(),
+      );
+      let payload_handle = ObservationBuilder::new("http/response")
+        .group(&http_group)
         .metadata("status", &parts.status.as_u16().to_string())
         .log_level(log_level)
-        .serde(&json!(filter_headers(
-          &parts.headers,
-          &config.excluded_headers
-        )));
+        .execution(&execution)
+        .named_raw_payload("headers", resp_headers_payload);
 
       // Wrap the response body in a streaming observer that captures data as it flows
-      // through and emits the observation when the stream completes
+      // through and adds the body payload when the stream completes
       let content_type = parts
         .headers
         .get(CONTENT_TYPE)
         .and_then(|v| v.to_str().ok())
         .unwrap_or("application/octet-stream")
         .to_string();
-      let streaming_body = StreamingObserverBody::new(
-        body,
-        content_type,
-        log_level,
-        parts.status.as_u16(),
-        execution,
-      );
+      let streaming_body =
+        StreamingObserverBody::new(body, content_type, payload_handle);
 
       Ok(Response::from_parts(parts, Body::new(streaming_body)))
     })

crates/observation-tools-client/src/tracing/layer.rs

@@ -1,6 +1,8 @@
 use super::span_data::SpanData;
 use crate::context;
+use crate::group::GroupHandle;
 use crate::observation::ObservationBuilder;
+use observation_tools_shared::GroupId;
 use observation_tools_shared::LogLevel;
 use observation_tools_shared::ObservationType;
 use observation_tools_shared::Payload;
@@ -70,9 +72,9 @@ where
 
   fn on_close(&self, id: Id, ctx: Context<'_, S>) {
     // Skip if no execution context
-    if context::get_current_execution().is_none() {
+    let Some(execution) = context::get_current_execution() else {
       return;
-    }
+    };
 
     let Some(span) = ctx.span(&id) else {
       return;
@@ -89,17 +91,20 @@ where
     let duration_ms = duration.as_secs_f64() * 1000.0;
 
     // Get the span's own ID and parent span ID
-    let span_id = id.into_u64().to_string();
+    let span_id = id.into_u64();
     let parent_span_id = span
       .parent()
       .map(|parent| parent.id().into_u64().to_string());
 
+    // Create a group handle from the span ID
+    let group_handle = GroupHandle::from_id(GroupId::from(span_id.to_string()), &execution);
+
     // Build and send observation
     let builder = ObservationBuilder::new(&data.name)
       .observation_type(ObservationType::Span)
       .log_level(tracing_level_to_log_level(data.level))
-      .label(format!("tracing/spans/{}", data.target))
-      .metadata("span_id", &span_id)
+      .group(&group_handle)
+      .metadata("span_id", &span_id.to_string())
       .metadata("duration_ms", format!("{:.3}", duration_ms))
       .metadata("target", &data.target);
 
@@ -127,9 +132,9 @@ where
       }
     }
 
-    if context::get_current_execution().is_none() {
+    let Some(execution) = context::get_current_execution() else {
       return;
-    }
+    };
 
     let mut visitor = FieldVisitor::new();
     event.record(&mut visitor);
@@ -153,13 +158,19 @@ where
       .unwrap_or_default();
 
     // Get current span from tracing's span stack for parent attribution
-    let parent_span_id = ctx.current_span().id().map(|id| id.into_u64().to_string());
+    let current_span_id = ctx.current_span().id().map(|id| id.into_u64());
+    let parent_span_id = current_span_id.map(|id| id.to_string());
 
     // Build and send observation
-    let builder = ObservationBuilder::new(metadata.name())
+    let mut builder = ObservationBuilder::new(metadata.name())
       .observation_type(ObservationType::LogEntry)
-      .log_level(tracing_level_to_log_level(*metadata.level()))
-      .label(format!("tracing/events/{}", metadata.target()));
+      .log_level(tracing_level_to_log_level(*metadata.level()));
+
+    // Add group reference from current span
+    if let Some(span_id) = current_span_id {
+      let group_handle = GroupHandle::from_id(GroupId::from(span_id.to_string()), &execution);
+      builder = builder.group(&group_handle);
+    }
 
     let builder = if let (Some(file), Some(line)) = (metadata.file(), metadata.line()) {
       builder.source(file, line)