langgraph_tracing/
context.rs

1use crate::event_bus::EventBus;
2use crate::observer::TracingGraphObserver;
3use crate::store::InMemoryTracingStore;
4use crate::types::TraceStatus;
5use langgraph_checkpoint::config::RunnableConfig;
6use serde_json::Value as JsonValue;
7use std::sync::Arc;
8
9/// High-level tracing context for LangGraph applications.
10///
11/// Bundles store, event bus, and observer into a single convenient API.
12/// Use with `#[derive(Traceable)]` on your state struct for automatic setup.
13pub struct TracingContext {
14    store: Arc<InMemoryTracingStore>,
15    event_bus: EventBus,
16    observer: TracingGraphObserver,
17}
18
19impl TracingContext {
20    /// Access the underlying observer for low-level tracing operations.
21    pub fn observer(&mut self) -> &mut TracingGraphObserver {
22        &mut self.observer
23    }
24
25    /// Access the underlying store.
26    pub fn store(&self) -> &Arc<InMemoryTracingStore> {
27        &self.store
28    }
29
30    /// Access the underlying event bus.
31    pub fn event_bus(&self) -> &EventBus {
32        &self.event_bus
33    }
34
35    /// Create a new tracing context with in-memory store.
36    pub fn new() -> Self {
37        let store = Arc::new(InMemoryTracingStore::new());
38        let event_bus = EventBus::new();
39        let observer = TracingGraphObserver::new(store.clone(), event_bus.clone());
40        Self {
41            store,
42            event_bus,
43            observer,
44        }
45    }
46
47    /// Start the tracing web server in a background task.
48    pub fn start_server(&self, addr: &str) {
49        let store = self.store.clone();
50        let event_bus = self.event_bus.clone();
51        let addr = addr.to_string();
52        tokio::spawn(async move {
53            crate::server::start(
54                &addr,
55                store,
56                event_bus,
57                Some("crates/langgraph-tracing/frontend/dist"),
58            )
59            .await
60            .unwrap();
61        });
62    }
63
64    /// Run a graph with automatic tracing.
65    ///
66    /// The closure receives a `RunnableConfig` (with `trace_id` injected into
67    /// `configurable`) and should execute the graph, returning the output.
68    ///
69    /// `on_graph_start` and `on_graph_end` are called automatically.
70    pub async fn run_with_tracing<F, Fut>(
71        &mut self,
72        name: &str,
73        input: JsonValue,
74        base_config: RunnableConfig,
75        f: F,
76    ) -> JsonValue
77    where
78        F: FnOnce(RunnableConfig) -> Fut,
79        Fut: std::future::Future<Output = JsonValue>,
80    {
81        let trace_id = self.observer.on_graph_start(name, input);
82
83        // Merge trace_id into the base config's configurable
84        let mut config = base_config;
85        let configurable = config
86            .entry("configurable".to_string())
87            .or_insert_with(|| serde_json::json!({}));
88        if let Some(obj) = configurable.as_object_mut() {
89            obj.insert("trace_id".to_string(), serde_json::json!(trace_id));
90        }
91
92        let output = f(config).await;
93
94        self.observer
95            .on_graph_end(&trace_id, output.clone(), TraceStatus::Success);
96        output
97    }
98}
99
100impl Default for TracingContext {
101    fn default() -> Self {
102        Self::new()
103    }
104}
langgraph_tracing/context.rs

langgraph_tracing/
context.rs
