feat: auto-instrument LangGraph nodes and routing decisions

CHANGELOG.md

@@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- **`chronicle.instrument(graph)`**: one call auto-instruments every node *and*
+  every `add_conditional_edges` routing function on a LangGraph `StateGraph`,
+  before or after `.compile()`. Routing decisions are now recorded as
+  `kind="router"` boundaries and replay deterministically (which branch was
+  taken), not just each node's input/output. (#22)
+
 ## [0.3.0] - 2026-07-24
 
 ### Added

chronicle/__init__.py

@@ -24,7 +24,7 @@
 from chronicle.redaction import apply_redactors, default_redactors, redact_secrets
 from chronicle.replay.plan import BoundaryMode, ReplayPlan
 from chronicle.session import ChronicleSession, SessionMode, get_session, reset_session
-from chronicle.wrap import instrument_langgraph, wrap
+from chronicle.wrap import instrument, instrument_langgraph, wrap
 
 __version__ = "0.3.0"
 
@@ -51,6 +51,7 @@
     "boundary",
     "default_redactors",
     "get_session",
+    "instrument",
     "instrument_langgraph",
     "open_store",
     "record",

chronicle/session.py

@@ -271,18 +271,46 @@ def envelope_to_return_value(envelope: Envelope, kind: str) -> Any:
         state["completion"] = envelope.action_result.completion
         state["finish_reason"] = envelope.action_result.finish_reason
         return state
+    if kind == "router":
+        # A router's return value is a plain node-name (or list of names), not a
+        # dict, so it lives inside raw_response under a fixed key rather than
+        # being raw_response itself — the generic dict-passthrough below would
+        # otherwise hand back {"decision": ...} instead of the decision itself.
+        raw = envelope.action_result.raw_response
+        if raw is not None and "decision" in raw:
+            return raw["decision"]
+        return envelope.action_result.completion
     raw = envelope.action_result.raw_response
     if raw is not None:
         return raw
     return envelope.action_result.completion
 
 
+def _router_decision(result: Any) -> Any:
+    """Coerce a routing function's return value into a JSON-safe, faithfully
+    replayable shape. LangGraph routing functions return a node name or a list
+    of node names (``Hashable | list[Hashable]``); node names are strings in
+    practice (``END`` included), so this covers the real range without needing
+    general-purpose JSON coercion."""
+    if isinstance(result, str):
+        return result
+    if isinstance(result, (list, tuple)):
+        return [str(v) for v in result]
+    return str(result)
+
+
 def result_to_action_result(result: Any, kind: str) -> ActionResult:
     if kind == "tool" and isinstance(result, dict):
         return ActionResult(
             completion=result.get("status", str(result)),
             raw_response=result,
         )
+    if kind == "router":
+        decision = _router_decision(result)
+        return ActionResult(
+            completion=decision if isinstance(decision, str) else str(decision),
+            raw_response={"decision": decision},
+        )
     if kind == "llm" and isinstance(result, dict):
         tool_calls = [
             ToolCall(

chronicle/wrap.py

@@ -12,6 +12,8 @@
 
 from __future__ import annotations
 
+import asyncio
+import contextvars
 import functools
 import inspect
 from collections.abc import Callable, Mapping
@@ -32,10 +34,105 @@ def instrument_langgraph(nodes: Mapping[str, Callable], *, kind: str = "custom")
     Each node keeps its behavior (transparent) and gains record / stub-replay /
     cut-point. Async nodes are supported. Use ``kind="llm"`` for nodes that call a
     model so their envelopes capture model metadata.
+
+    For a graph you already build with ``StateGraph`` (rather than a bare dict of
+    node functions), prefer ``chronicle.instrument(graph)``: it wraps every node in
+    place *and* records routing decisions from ``add_conditional_edges``.
     """
     return {name: boundary(name, kind=kind)(fn) for name, fn in nodes.items()}
 
 
+def instrument(graph: Any, *, kind: str = "custom") -> Any:
+    """Auto-instrument every node and every conditional-edge routing decision on
+    a LangGraph graph, in one call, before or after ``.compile()``.
+
+        graph = StateGraph(State)
+        graph.add_node("agent", agent_node)
+        graph.add_conditional_edges("agent", route_fn, {"tools": "tools", END: END})
+        app = chronicle.instrument(graph).compile()
+
+    A ``CompiledStateGraph`` keeps a live reference to the ``StateGraph`` builder
+    it was compiled from (``compiled.builder``), and each node's underlying
+    callable is mutated in place rather than replaced wholesale, so
+    ``chronicle.instrument(app)`` also works *after* ``.compile()`` — the
+    already-compiled graph's execution is rerouted too.
+
+    Every node becomes a ``@boundary(name, kind=kind)`` crossing: same
+    transparent record / stub-replay / cut-point contract as everywhere else in
+    Chronicle, sync or async. Every ``add_conditional_edges`` routing function
+    becomes a ``kind="router"`` boundary, so *which branch was taken* is
+    captured and replays deterministically too, not just each node's
+    input/output. Static (unconditional) edges need no boundary — there is no
+    decision to record.
+
+    Idempotent: instrumenting the same graph twice wraps each node/router once.
+    Requires ``langgraph`` (``pip install chronicle[langgraph]``); this function
+    only touches the public ``StateGraph`` builder attributes (``nodes``,
+    ``branches``), never langgraph's compiled Pregel internals.
+    """
+    builder = getattr(graph, "builder", graph)
+    for node_id, spec in getattr(builder, "nodes", {}).items():
+        runnable = getattr(spec, "runnable", None)
+        if runnable is not None:
+            _instrument_runnable(runnable, node_id, kind)
+    for source, branches in getattr(builder, "branches", {}).items():
+        for name, branch_spec in branches.items():
+            path = getattr(branch_spec, "path", None)
+            if path is not None:
+                _instrument_runnable(path, f"{source}:{name}", "router")
+    return graph
+
+
+def _instrument_runnable(runnable: Any, boundary_id: str, kind: str) -> None:
+    """Wrap a langgraph ``RunnableCallable``'s underlying function(s) as a
+    Chronicle boundary, in place, so both ``invoke`` and ``ainvoke`` record.
+
+    A sync-only node gets an auto-generated ``afunc`` from langgraph — a
+    ``functools.partial`` that runs ``func`` in a thread executor. Rewiring
+    ``func`` alone would leave that partial's closure pointing at the
+    *original*, unwrapped function, silently skipping the boundary whenever the
+    graph is run with ``ainvoke``/``astream``. So the executor shim is rebuilt
+    around the wrapped function instead of touching ``func`` and ``afunc``
+    independently.
+    """
+    if getattr(runnable, "_chronicle_instrumented", False):
+        return
+    func = getattr(runnable, "func", None)
+    afunc = getattr(runnable, "afunc", None)
+    if func is not None:
+        wrapped_sync = boundary(boundary_id, kind=kind)(func)
+        runnable.func = wrapped_sync
+        if isinstance(afunc, functools.partial):
+            runnable.afunc = _executor_shim(wrapped_sync)
+        elif afunc is not None:
+            runnable.afunc = boundary(boundary_id, kind=kind)(afunc)
+    elif afunc is not None:
+        runnable.afunc = boundary(boundary_id, kind=kind)(afunc)
+    runnable._chronicle_instrumented = True
+
+
+def _executor_shim(sync_fn: Callable) -> Callable[..., Any]:
+    """Stdlib rebuild of langgraph's 'run this sync function in a thread
+    executor' shim, around an already-instrumented function — so async graph
+    execution keeps the same off-loop behavior without depending on
+    langgraph's private executor helper.
+
+    The thread pool executor does not inherit the calling thread's
+    ``ContextVar`` state (that is how ``ChronicleSession`` is scoped), so the
+    context is copied explicitly and run inside the worker thread — the same
+    fix langgraph's own ``run_in_executor`` applies, for the same reason.
+    """
+
+    async def _call(*args: Any, **kwargs: Any) -> Any:
+        loop = asyncio.get_event_loop()
+        ctx = contextvars.copy_context()
+        return await loop.run_in_executor(
+            None, ctx.run, functools.partial(sync_fn, *args, **kwargs)
+        )
+
+    return _call
+
+
 def wrap(client: Any, *, boundary_id: str = "llm") -> Any:
     """Record every model call an OpenAI- or Anthropic-style client makes.
 

examples/langgraph_demo/routing_demo.py

@@ -0,0 +1,83 @@
+"""chronicle.instrument(graph): one call records every node AND every
+add_conditional_edges routing decision on a compiled LangGraph graph.
+
+Run:
+    pip install -e ".[dev]"
+    python examples/langgraph_demo/routing_demo.py
+"""
+
+from __future__ import annotations
+
+from typing import TypedDict
+
+import chronicle
+
+try:
+    from langgraph.graph import END, StateGraph
+except ImportError:
+    print("Install langgraph: pip install chronicle[langgraph]")
+    raise
+
+
+class AgentState(TypedDict):
+    query: str
+    needs_search: bool
+    completion: str
+
+
+def agent_node(state: AgentState) -> dict:
+    needs_search = "reset" in state["query"].lower()
+    return {"needs_search": needs_search}
+
+
+def search_node(state: AgentState) -> dict:
+    return {"completion": "You can reset your API key from Settings > API Keys."}
+
+
+def answer_node(state: AgentState) -> dict:
+    return {"completion": "I can help — could you say more about what you need?"}
+
+
+def route_after_agent(state: AgentState) -> str:
+    return "search" if state["needs_search"] else "answer"
+
+
+def build_app():
+    graph = StateGraph(AgentState)
+    graph.add_node("agent", agent_node)
+    graph.add_node("search", search_node)
+    graph.add_node("answer", answer_node)
+    graph.set_entry_point("agent")
+    graph.add_conditional_edges(
+        "agent", route_after_agent, {"search": "search", "answer": "answer"}
+    )
+    graph.add_edge("search", END)
+    graph.add_edge("answer", END)
+
+    # One call: every node above, plus route_after_agent's decision, is now a
+    # Chronicle boundary — recorded live, and stub-replayable from a fixture.
+    return chronicle.instrument(graph).compile()
+
+
+def main() -> None:
+    app = build_app()
+
+    trace_dir = "fixtures/traces/langgraph-routing-demo"
+    with chronicle.record("langgraph-routing-demo", export=trace_dir) as session:
+        result = app.invoke({"query": "How do I reset my API key?", "needs_search": False, "completion": ""})
+
+    print(f"Completion: {result['completion']}")
+    print(f"Recorded {len(session._recorded_envelopes)} envelope(s) to {trace_dir}/")
+    for e in session._recorded_envelopes:
+        print(f"  node={e.node_id}  kind={e.boundary_kind}  completion={e.action_result.completion!r}")
+
+    # Replay: no node function and no routing function runs — every crossing,
+    # including which branch was taken, comes back from the fixture.
+    with chronicle.replay_trace(trace_dir) as session:
+        replayed = app.invoke({"query": "How do I reset my API key?", "needs_search": False, "completion": ""})
+    assert replayed["completion"] == result["completion"]
+    print("Replay reproduced the same completion with no node/router code executed.")
+
+
+if __name__ == "__main__":
+    main()