From 3bddc1f3211379eaa7024df52b7ddba77adc466d Mon Sep 17 00:00:00 2001 From: DudeFromMars Date: Wed, 29 Jul 2026 12:57:33 +0000 Subject: [PATCH] feat: auto-instrument LangGraph nodes and routing decisions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit chronicle.instrument(graph) wraps every node and every add_conditional_edges routing function on a StateGraph in one call, before or after .compile() — a CompiledStateGraph keeps a live reference to its builder, and node callables are mutated in place, so instrumenting late still reroutes already-compiled execution. Routing functions are recorded as kind="router" boundaries, so which branch was taken is captured and replays deterministically, not just each node's input/output. A sync-only node's auto-generated async executor shim closes over the original, unwrapped function, so it's rebuilt around the instrumented one instead — otherwise ainvoke/astream would silently skip recording. That shim also now copies the caller's contextvars into the executor thread (matching langgraph's own run_in_executor), since Chronicle scopes its session via ContextVar and threads don't inherit it. Closes #22 Co-Authored-By: Claude Sonnet 5 Signed-off-by: DudeFromMars --- CHANGELOG.md | 7 ++ chronicle/__init__.py | 3 +- chronicle/session.py | 28 ++++++ chronicle/wrap.py | 97 +++++++++++++++++++ examples/langgraph_demo/routing_demo.py | 83 ++++++++++++++++ tests/test_langgraph_instrument.py | 123 ++++++++++++++++++++++++ 6 files changed, 340 insertions(+), 1 deletion(-) create mode 100644 examples/langgraph_demo/routing_demo.py create mode 100644 tests/test_langgraph_instrument.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 21332aa..bbecd91 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/chronicle/__init__.py b/chronicle/__init__.py index f44b2fd..1c0c858 100644 --- a/chronicle/__init__.py +++ b/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", diff --git a/chronicle/session.py b/chronicle/session.py index 3210a6b..04ea468 100644 --- a/chronicle/session.py +++ b/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( diff --git a/chronicle/wrap.py b/chronicle/wrap.py index 2bd0fda..60b8879 100644 --- a/chronicle/wrap.py +++ b/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. diff --git a/examples/langgraph_demo/routing_demo.py b/examples/langgraph_demo/routing_demo.py new file mode 100644 index 0000000..856deb8 --- /dev/null +++ b/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() diff --git a/tests/test_langgraph_instrument.py b/tests/test_langgraph_instrument.py new file mode 100644 index 0000000..5bc244f --- /dev/null +++ b/tests/test_langgraph_instrument.py @@ -0,0 +1,123 @@ +"""chronicle.instrument(graph) auto-instruments every LangGraph node and every +add_conditional_edges routing function in one call — sync, async, whether +called before or after .compile(), and idempotently. See issue #22.""" + +from __future__ import annotations + +import asyncio +from typing import TypedDict + +import pytest +from langgraph.graph import END, StateGraph + +import chronicle + + +class State(TypedDict): + x: int + path: list[str] + + +def _build_graph() -> StateGraph: + def step_a(state: State) -> dict: + return {"x": state["x"] + 1, "path": [*state["path"], "a"]} + + def step_b(state: State) -> dict: + return {"x": state["x"] + 10, "path": [*state["path"], "b"]} + + def route(state: State) -> str: + return "b" if state["x"] > 0 else END + + graph = StateGraph(State) + graph.add_node("a", step_a) + graph.add_node("b", step_b) + graph.set_entry_point("a") + graph.add_conditional_edges("a", route, {"b": "b", END: END}) + graph.add_edge("b", END) + return graph + + +@pytest.mark.layer1 +def test_instrument_records_nodes_and_router_decision(): + session = chronicle.reset_session() + session.begin_trace("t-instrument") + + app = chronicle.instrument(_build_graph()).compile() + result = app.invoke({"x": 1, "path": []}) + + assert result["path"] == ["a", "b"] + node_ids = {e.node_id for e in session._recorded_envelopes} + assert node_ids == {"a", "a:route", "b"} + + router_env = next(e for e in session._recorded_envelopes if e.node_id == "a:route") + assert router_env.boundary_kind == "router" + assert router_env.action_result.completion == "b" + + +@pytest.mark.layer1 +def test_instrument_after_compile_still_reroutes_execution(): + """Instrumenting late (after .compile()) must still record — the compiled + graph keeps a live reference to the builder this mutates in place.""" + session = chronicle.reset_session() + session.begin_trace("t-late") + + compiled = _build_graph().compile() + chronicle.instrument(compiled) + compiled.invoke({"x": 1, "path": []}) + + assert len(session._recorded_envelopes) == 3 + + +@pytest.mark.layer1 +def test_instrument_records_through_async_invocation(): + """A sync node's auto-generated executor shim must be rebuilt around the + wrapped function, not left pointing at the original — otherwise ainvoke + silently skips recording.""" + session = chronicle.reset_session() + session.begin_trace("t-async") + + app = chronicle.instrument(_build_graph()).compile() + result = asyncio.run(app.ainvoke({"x": 1, "path": []})) + + assert result["path"] == ["a", "b"] + assert len(session._recorded_envelopes) == 3 + + +@pytest.mark.layer1 +def test_instrument_is_idempotent(): + graph = _build_graph() + chronicle.instrument(graph) + chronicle.instrument(graph) # calling twice must not double-wrap + app = graph.compile() + + session = chronicle.reset_session() + session.begin_trace("t-idempotent") + app.invoke({"x": 1, "path": []}) + + assert len(session._recorded_envelopes) == 3 + + +@pytest.mark.layer1 +def test_replay_stubs_the_router_without_calling_the_route_function(tmp_path): + """Layer 1 replay must reproduce which branch was taken from the fixture, + never by actually running the routing function again.""" + trace = tmp_path / "trace" + with chronicle.record("incident", export=str(trace)): + app = chronicle.instrument(_build_graph()).compile() + app.invoke({"x": 1, "path": []}) + + calls: list[State] = [] + + def spy_route(state: State) -> str: + calls.append(state) + return "b" + + graph = _build_graph() + graph.branches["a"]["route"].path.func = spy_route + + with chronicle.replay_trace(str(trace)): + app = chronicle.instrument(graph).compile() + result = app.invoke({"x": 1, "path": []}) + + assert calls == [] # the router was stubbed, never executed live + assert result["path"] == ["a", "b"]