From 60e9701624ee71c6bf68974ca95d686b5630ed24 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 03:59:41 -0700 Subject: [PATCH] feat(observability): stamp langgraph Agent-Graph-view keys + host span-metadata injector on graph Langfuse exporter - node spans now carry langgraph_node (node id) and langgraph_step (superstep index) in metadata so Langfuse's Agent Graph view can lay the trace out as a graph; step spans carry langgraph_step; subgraph spans carry langgraph_step from their coordinate observation - GraphLangfuseExporter::with_span_metadata_fn installs an optional host metadata injector (Fn(&GraphObservation) -> Option) merged over built-in span metadata keys (host wins on collision); default None, non-breaking - re-export SpanMetadataFn from graph/observability, graph, and crate root Claude-Session: https://claude.ai/code/session_01VSwRwLDBXg3o6timchvKn8 --- src/graph/mod.rs | 3 +- src/graph/observability/langfuse.rs | 117 ++++++++++++++++++---- src/graph/observability/langfuse/tests.rs | 92 +++++++++++++++++ src/graph/observability/mod.rs | 2 +- src/lib.rs | 3 +- 5 files changed, 197 insertions(+), 20 deletions(-) diff --git a/src/graph/mod.rs b/src/graph/mod.rs index f369847..f0dd7f1 100644 --- a/src/graph/mod.rs +++ b/src/graph/mod.rs @@ -69,7 +69,8 @@ pub use goals::{ pub use observability::{ GraphEventJournal, GraphHealthSummary, GraphLangfuseExporter, GraphLatencyMetrics, GraphNodeHealth, GraphNodeLatency, GraphObservation, GraphStatusStore, GraphStepLatency, - InMemoryGraphEventJournal, InMemoryGraphStatusStore, JournalGraphSink, StoreGraphEventJournal, + InMemoryGraphEventJournal, InMemoryGraphStatusStore, JournalGraphSink, SpanMetadataFn, + StoreGraphEventJournal, }; pub use orchestration::{ InMemoryTaskStore, JsonlTaskStore, OrchestrationControlOutcome, OrchestrationTaskFilter, diff --git a/src/graph/observability/langfuse.rs b/src/graph/observability/langfuse.rs index 13ba672..5e96a26 100644 --- a/src/graph/observability/langfuse.rs +++ b/src/graph/observability/langfuse.rs @@ -25,8 +25,9 @@ use std::collections::HashMap; use std::collections::VecDeque; +use std::sync::Arc; -use serde_json::{Value, json}; +use serde_json::{Map, Value, json}; use crate::error::{Result, TinyAgentsError}; use crate::graph::observability::{GraphHealthSummary, GraphObservation}; @@ -39,20 +40,43 @@ use crate::harness::observability::{LangfuseClient, LangfuseTraceConfig, clean_n /// arrival order and pair with their terminals FIFO. type StartQueue = VecDeque<(usize, u64)>; +/// A host-supplied hook that contributes extra metadata to every span the +/// exporter emits. Called with the span's coordinate observation; returning +/// `None` (or an empty map) contributes nothing. Host keys are merged **over** +/// the exporter's built-in coordinate keys, so a host can also override them. +pub type SpanMetadataFn = dyn Fn(&GraphObservation) -> Option> + Send + Sync; + /// Async Langfuse exporter for durable graph observations. /// /// Wraps a shared [`LangfuseClient`] for transport (auth, endpoint /// normalization, `207 Multi-Status` handling) and adds graph-aware payload /// construction on top. -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct GraphLangfuseExporter { client: LangfuseClient, + /// Optional host metadata injector, merged into every span's metadata. + span_metadata_fn: Option>, +} + +impl std::fmt::Debug for GraphLangfuseExporter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GraphLangfuseExporter") + .field("client", &self.client) + .field( + "span_metadata_fn", + &self.span_metadata_fn.as_ref().map(|_| "Fn(..)"), + ) + .finish() + } } impl GraphLangfuseExporter { /// Wraps an existing [`LangfuseClient`] as a graph exporter. pub fn new(client: LangfuseClient) -> Self { - Self { client } + Self { + client, + span_metadata_fn: None, + } } /// Builds an exporter from the same environment variables as @@ -61,6 +85,17 @@ impl GraphLangfuseExporter { Ok(Self::new(LangfuseClient::from_env()?)) } + /// Installs a host metadata injector: `f` is called once per emitted span + /// with the span's coordinate observation, and any returned keys are merged + /// into that span's metadata (host keys win on collision). Default: none. + pub fn with_span_metadata_fn( + mut self, + f: impl Fn(&GraphObservation) -> Option> + Send + Sync + 'static, + ) -> Self { + self.span_metadata_fn = Some(Arc::new(f)); + self + } + /// Returns the underlying transport client. pub fn client(&self) -> &LangfuseClient { &self.client @@ -97,7 +132,13 @@ impl GraphLangfuseExporter { batch.push(trace_create(&trace_id, &trace_ts, &trace, first, &health)); let mut consumed = vec![false; observations.len()]; - push_span_events(&trace_id, observations, &mut consumed, &mut batch); + push_span_events( + &trace_id, + observations, + &mut consumed, + &mut batch, + self.span_metadata_fn.as_deref(), + ); push_point_events(&trace_id, observations, &consumed, &mut batch); Ok(json!({ "batch": batch })) @@ -187,6 +228,7 @@ fn push_span_events( observations: &[GraphObservation], consumed: &mut [bool], batch: &mut Vec, + injector: Option<&SpanMetadataFn>, ) { // Pending starts keyed by their span identity, each holding (index, ts). let mut step_starts: HashMap = HashMap::new(); @@ -212,6 +254,7 @@ fn push_span_events( Some(obs.ts_ms), obs, &observations[start_idx], + injector, )); } } @@ -233,6 +276,7 @@ fn push_span_events( Some(obs.ts_ms), None, obs, + injector, )); } } @@ -247,6 +291,7 @@ fn push_span_events( Some(obs.ts_ms), Some(error.as_str()), obs, + injector, )); } } @@ -269,6 +314,7 @@ fn push_span_events( start_ts, Some(obs.ts_ms), obs, + injector, )); } } @@ -280,7 +326,9 @@ fn push_span_events( for (step, mut queue) in step_starts { while let Some((start_idx, start_ts)) = queue.pop_front() { let obs = &observations[start_idx]; - batch.push(step_span(trace_id, step, start_ts, None, obs, obs)); + batch.push(step_span( + trace_id, step, start_ts, None, obs, obs, injector, + )); } } for ((node, step), mut queue) in node_starts { @@ -293,6 +341,7 @@ fn push_span_events( None, None, &observations[start_idx], + injector, )); } } @@ -305,6 +354,7 @@ fn push_span_events( start_ts, None, &observations[start_idx], + injector, )); } } @@ -342,13 +392,14 @@ fn push_point_events( "startTime": ts, "level": level, "statusMessage": status, - "metadata": span_metadata(obs), + "metadata": span_metadata(obs, None, None, None), })), })); } } /// Builds a `span-create` for a superstep, parented directly to the trace. +#[allow(clippy::too_many_arguments)] fn step_span( trace_id: &str, step: usize, @@ -356,7 +407,9 @@ fn step_span( end_ts: Option, terminal: &GraphObservation, start: &GraphObservation, + injector: Option<&SpanMetadataFn>, ) -> Value { + let metadata = span_metadata(start, None, Some(step), injector); span_event( trace_id, &format!("{trace_id}:step:{step}"), @@ -366,11 +419,15 @@ fn step_span( end_ts, None, terminal, - start, + metadata, ) } /// Builds a `span-create` for a node handler, parented to its superstep span. +/// +/// Node spans carry the Langfuse **Agent Graph view** keys `langgraph_node` +/// (the node id) and `langgraph_step` (the superstep index) in metadata, so +/// Langfuse can lay the trace out as a graph. #[allow(clippy::too_many_arguments)] fn node_span( trace_id: &str, @@ -380,7 +437,9 @@ fn node_span( end_ts: Option, error: Option<&str>, terminal: &GraphObservation, + injector: Option<&SpanMetadataFn>, ) -> Value { + let metadata = span_metadata(terminal, Some(node.as_str()), Some(step), injector); span_event( trace_id, &format!("{trace_id}:node:{}:{step}", node.as_str()), @@ -390,7 +449,7 @@ fn node_span( end_ts, error, terminal, - terminal, + metadata, ) } @@ -402,7 +461,9 @@ fn subgraph_span( start_ts: u64, end_ts: Option, terminal: &GraphObservation, + injector: Option<&SpanMetadataFn>, ) -> Value { + let metadata = span_metadata(terminal, None, Some(terminal.step), injector); span_event( trace_id, &format!("{trace_id}:subgraph:{}", namespace.join("/")), @@ -412,13 +473,13 @@ fn subgraph_span( end_ts, None, terminal, - terminal, + metadata, ) } /// Shared `span-create` builder. `error` promotes the span to `ERROR` level. -/// The batch item id comes from the terminal observation so it is unique; span -/// coordinate metadata comes from the start observation. +/// The batch item id comes from the terminal observation so it is unique; the +/// caller supplies the fully-built span `metadata`. #[allow(clippy::too_many_arguments)] fn span_event( trace_id: &str, @@ -429,7 +490,7 @@ fn span_event( end_ts: Option, error: Option<&str>, terminal: &GraphObservation, - coords: &GraphObservation, + metadata: Value, ) -> Value { let start_iso = iso_ms(start_ts); let end_iso = end_ts.map(iso_ms); @@ -450,14 +511,22 @@ fn span_event( "endTime": end_iso, "level": level, "statusMessage": status, - "metadata": span_metadata(coords), + "metadata": metadata, })), }) } -/// Extracts the correlation coordinates every span/event carries in metadata. -fn span_metadata(obs: &GraphObservation) -> Value { - json!({ +/// Extracts the correlation coordinates every span/event carries in metadata, +/// stamping the Langfuse Agent-Graph-view keys (`langgraph_node`, +/// `langgraph_step`) when the caller supplies them and merging any +/// host-injected keys last (host keys win on collision). +fn span_metadata( + obs: &GraphObservation, + langgraph_node: Option<&str>, + langgraph_step: Option, + injector: Option<&SpanMetadataFn>, +) -> Value { + let mut metadata = json!({ "run_id": obs.run_id.as_str(), "root_run_id": obs.root_run_id.as_str(), "parent_run_id": obs.parent_run_id.as_ref().map(|id| id.as_str()), @@ -467,7 +536,21 @@ fn span_metadata(obs: &GraphObservation) -> Value { "step": obs.step, "offset": obs.offset, "event": obs.event, - }) + }); + if let Value::Object(map) = &mut metadata { + if let Some(node) = langgraph_node { + map.insert("langgraph_node".to_string(), json!(node)); + } + if let Some(step) = langgraph_step { + map.insert("langgraph_step".to_string(), json!(step)); + } + if let Some(extra) = injector.and_then(|f| f(obs)) { + for (k, v) in extra { + map.insert(k, v); + } + } + } + metadata } /// Pops the FIFO-oldest pending start for `key`, cleaning up empty queues. diff --git a/src/graph/observability/langfuse/tests.rs b/src/graph/observability/langfuse/tests.rs index ced406d..790ddfb 100644 --- a/src/graph/observability/langfuse/tests.rs +++ b/src/graph/observability/langfuse/tests.rs @@ -243,6 +243,98 @@ fn health_telemetry_rides_on_trace_metadata() { assert_eq!(nodes[1]["failed"], 1); } +#[test] +fn node_spans_carry_langgraph_agent_graph_view_keys() { + let batch = exporter() + .build_ingestion_batch(LangfuseTraceConfig::default(), &sample_run()) + .unwrap(); + let events = batch["batch"].as_array().unwrap(); + + // Node spans carry both Agent-Graph-view keys: the node id and superstep. + let node_a = find(events, "span-create", "root-1:node:a:1").expect("node a span"); + assert_eq!(node_a["body"]["metadata"]["langgraph_node"], "a"); + assert_eq!(node_a["body"]["metadata"]["langgraph_step"], 1); + let node_b = find(events, "span-create", "root-1:node:b:1").expect("node b span"); + assert_eq!(node_b["body"]["metadata"]["langgraph_node"], "b"); + assert_eq!(node_b["body"]["metadata"]["langgraph_step"], 1); + + // Step spans carry the superstep index but no node id. + let step = find(events, "span-create", "root-1:step:1").expect("step span"); + assert_eq!(step["body"]["metadata"]["langgraph_step"], 1); + assert!(step["body"]["metadata"].get("langgraph_node").is_none()); + + // Existing coordinate keys are preserved alongside the new ones. + assert_eq!(node_a["body"]["metadata"]["run_id"], "run-1"); + assert_eq!(node_a["body"]["metadata"]["graph_id"], "demo-graph"); +} + +#[test] +fn subgraph_spans_carry_langgraph_step() { + let sub = NodeId::new("sub"); + let ns = vec!["sub".to_string()]; + let observations = vec![ + obs( + 0, + 2, + 1_000, + GraphEvent::SubgraphStarted { + node: sub.clone(), + namespace: ns.clone(), + }, + ), + obs( + 1, + 2, + 1_050, + GraphEvent::SubgraphCompleted { + node: sub.clone(), + namespace: ns.clone(), + }, + ), + ]; + let batch = exporter() + .build_ingestion_batch(LangfuseTraceConfig::default(), &observations) + .unwrap(); + let events = batch["batch"].as_array().unwrap(); + let span = find(events, "span-create", "root-1:subgraph:sub").expect("subgraph span"); + assert_eq!(span["body"]["metadata"]["langgraph_step"], 2); + assert!(span["body"]["metadata"].get("langgraph_node").is_none()); +} + +#[test] +fn host_span_metadata_injector_merges_and_wins_on_collision() { + let exporter = exporter().with_span_metadata_fn(|obs| { + let mut extra = serde_json::Map::new(); + extra.insert("flow_id".to_string(), json!("flow-42")); + // Host keys win over built-ins on collision. + extra.insert("graph_id".to_string(), json!("host-override")); + // Prove the injector sees the observation it decorates. + extra.insert("seen_step".to_string(), json!(obs.step)); + Some(extra) + }); + let batch = exporter + .build_ingestion_batch(LangfuseTraceConfig::default(), &sample_run()) + .unwrap(); + let events = batch["batch"].as_array().unwrap(); + + let node_a = find(events, "span-create", "root-1:node:a:1").expect("node a span"); + assert_eq!(node_a["body"]["metadata"]["flow_id"], "flow-42"); + assert_eq!(node_a["body"]["metadata"]["graph_id"], "host-override"); + assert_eq!(node_a["body"]["metadata"]["seen_step"], 1); + // Built-in Agent-Graph-view keys survive the merge. + assert_eq!(node_a["body"]["metadata"]["langgraph_node"], "a"); + + let step = find(events, "span-create", "root-1:step:1").expect("step span"); + assert_eq!(step["body"]["metadata"]["flow_id"], "flow-42"); + + // Point events (e.g. run.failed) are untouched by the injector. + let failed = events + .iter() + .find(|e| e["body"]["name"] == "run.failed") + .expect("run.failed event"); + assert!(failed["body"]["metadata"].get("flow_id").is_none()); +} + #[test] fn open_node_span_has_start_but_no_end() { let a = NodeId::new("a"); diff --git a/src/graph/observability/mod.rs b/src/graph/observability/mod.rs index 9fedcfb..0939683 100644 --- a/src/graph/observability/mod.rs +++ b/src/graph/observability/mod.rs @@ -40,7 +40,7 @@ mod langfuse; mod types; -pub use langfuse::GraphLangfuseExporter; +pub use langfuse::{GraphLangfuseExporter, SpanMetadataFn}; pub use types::*; use std::collections::{HashMap, VecDeque}; diff --git a/src/lib.rs b/src/lib.rs index 00906a9..9060a35 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -174,7 +174,8 @@ pub use graph::{ pub use graph::{ GraphEventJournal, GraphHealthSummary, GraphLangfuseExporter, GraphLatencyMetrics, GraphNodeHealth, GraphNodeLatency, GraphObservation, GraphStatusStore, GraphStepLatency, - InMemoryGraphEventJournal, InMemoryGraphStatusStore, JournalGraphSink, StoreGraphEventJournal, + InMemoryGraphEventJournal, InMemoryGraphStatusStore, JournalGraphSink, SpanMetadataFn, + StoreGraphEventJournal, }; // --- Graph: orchestration tools (ordinary harness Tool implementations) ---