diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/.changelog/536.added b/instrumentation/opentelemetry-instrumentation-genai-langchain/.changelog/536.added new file mode 100644 index 000000000..1864fe243 --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/.changelog/536.added @@ -0,0 +1 @@ +Emit agent lifecycle events (paused, checkpointed, resumed) for LangGraph durable executions. diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/README.rst b/instrumentation/opentelemetry-instrumentation-genai-langchain/README.rst index 4e24a50f8..aabb5d12c 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-langchain/README.rst +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/README.rst @@ -16,6 +16,8 @@ application: * **Agent spans** for agent invocations nested inside a workflow, including the agent name, id, description, and conversation/session id when available. * **Tool spans** for tool calls made during a run. +* **Lifecycle events** for LangGraph durable executions that pause, checkpoint, + and resume, correlated with the workflow span. The spans nest to reflect the graph, so a single graph invocation produces a workflow span with the agent, tool, and model calls it triggered as children. @@ -96,6 +98,21 @@ calls nested underneath. } ) +LangGraph durable executions +---------------------------- + +When a LangGraph graph pauses on ``interrupt()``, the instrumentation emits a +``gen_ai.agent.paused`` event carrying the interrupt id LangGraph minted and the +checkpoint it paused at. Resuming the graph with ``Command(resume=...)`` emits +``gen_ai.agent.resumed`` naming that same checkpoint. If the graph was compiled +with a checkpointer, every checkpoint the saver persists emits a +``gen_ai.agent.checkpointed`` event; LangGraph writes one checkpoint per +superstep, so these are per-step records. + +No interrupt payload or graph state is recorded: the events carry only ids +LangGraph itself produced. These event and attribute names are candidate +semantic conventions and may change. + Configuration ------------- diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/README.rst b/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/README.rst new file mode 100644 index 000000000..12e60b8ae --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/README.rst @@ -0,0 +1,122 @@ +OpenTelemetry LangGraph agent lifecycle example +=============================================== + +This is an example of the agent lifecycle telemetry the instrumentation +captures from a LangGraph durable execution: a graph that pauses on +``interrupt()``, checkpoints, and resumes on a second invoke. + +`main.py `_ needs no API key and no collector. It runs the graph twice +under the instrumentor with in-memory exporters and writes every span and log +record it produced to `sample-output.json `_, which is the +captured run committed here. + +Setup +----- + +Set up a virtual environment like this: + +:: + + python3 -m venv .venv + source .venv/bin/activate + pip install -r requirements.txt + +Run +--- + +Run the example from the repository root like this: + +:: + + python instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/main.py + +You should see ``gen_ai.agent.paused``, ``gen_ai.agent.checkpointed``, and +``gen_ai.agent.resumed`` events correlated with the workflow span of the invoke +that produced them. + +Design note +----------- + +Every name below is a candidate semantic convention proposed in +open-telemetry/semantic-conventions-genai#445, and none of them is stable. +Everything stated here was measured against langgraph 1.2.9 and +langchain-core 1.5.0. + +Interception points +~~~~~~~~~~~~~~~~~~~ + +1. ``langgraph.callbacks`` lifecycle dispatch. LangGraph 1.2 calls + ``on_interrupt`` and ``on_resume`` on the handlers in a run's callback + manager, passing ``GraphInterruptEvent`` and ``GraphResumeEvent`` with the + real ``Interrupt`` objects, the checkpoint id, and the top level run id. The + released instrumentor is already in that handler list: without these two + methods LangGraph logs ``AttributeError`` on every interrupt and resume, so + adding them also removes existing log noise. +2. ``StateGraph.compile(checkpointer=...)``. The saver the compiled graph + retains is wrapped so every persisted checkpoint reports its id. + ``BaseCheckpointSaver.put`` is abstract and every saver overrides it, so the + instance is patched rather than the base class. ``aput`` may delegate to + ``put``, possibly on a worker thread, so a write is de-duplicated by holding + the identity of the ``Checkpoint`` object while the write is in flight, not + by remembering ids and not by relying on context propagation: a nested call + receiving the same object is the delegation and stays silent, and the + outermost call reports. This assumes one checkpoint object per logical + write, which LangGraph guarantees by handing the saver a freshly constructed + mapping per superstep (``_loop.py`` passes + ``copy_checkpoint(self.checkpoint)``), so two independent writes never share + an object. A saver that copies the checkpoint before delegating is reported + twice, which is the safe direction. + +The interrupt is not observable from ``on_chain_end``: LangGraph adds +``__interrupt__`` to the invoke return value after the callback fires. + +Correlation +~~~~~~~~~~~ + +``paused`` and ``resumed`` use the run id on the lifecycle event, resolved to +the nearest workflow invocation. A checkpointer call has no run id, so +``checkpointed`` maps the ``configurable.thread_id`` and +``configurable.checkpoint_ns`` in its config to a live run. Runs are tracked per +thread id, each bound with the namespace its own checkpoints use ("" for a top +level graph, otherwise the run's ``langgraph_checkpoint_ns``). A write resolves +to the run bound with that exact namespace, or to the nearest enclosing run when +none is. A nested run never displaces the run containing it, and when two +equally specific live runs are unrelated the event is dropped rather than +guessed. + +LangGraph classifies a subgraph's own graph run as a plain chain, not a +workflow, so today the child namespace resolves to the parent workflow. The +resolution is namespace aware regardless, so a nested workflow would own its own +writes. + +``resumed_from.type`` is always ``checkpoint``. That is a constant in the +instrumentation, but it is determined by LangGraph, not chosen: the resume event +supplies a checkpoint id and nothing else. + +``GenAIInvocation.emit_event`` is byte-identical to the hunk in open PR #507, +saved as ``util/opentelemetry-util-genai/tests/fixtures/pr507_emit_event.py.txt`` +and checked by a test, and is dropped when rebasing onto that PR. It is not a +competing API. + +Checkpoint volume +~~~~~~~~~~~~~~~~~ + +``put`` fires once per superstep. The captured run emits 3 ``checkpointed`` +events for the invoke that pauses, 2 for the invoke that resumes. Nested +subgraphs checkpoint independently, under their own namespace, and emit one +``resumed`` per graph level, which the flat event model does not distinguish. + +Not demonstrable +~~~~~~~~~~~~~~~~ + +``gen_ai.agent.execution.id`` is omitted. LangGraph mints no id spanning suspend +and resume: ``thread_id`` is a conversation reused across runs, and every other +id changes on the resuming invoke. The captured run shows the two invokes as +separate traces with nothing linking them, which is the gap the attribute +describes. + +``gen_ai.agent.pause.reason`` is omitted. ``interrupt(value)`` carries an opaque +application payload and nothing saying who resolves the pause, so neither +``human_input`` nor ``external_system`` is derivable. The ``pause`` member of +``resumed_from.type`` is likewise absent: LangGraph never reports a pause id at +resume time. diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/main.py b/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/main.py new file mode 100644 index 000000000..7f0efaf7d --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/main.py @@ -0,0 +1,146 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +"""LangGraph agent lifecycle telemetry example. + +Runs a small human-in-the-loop LangGraph application twice, once until it +interrupts and once to resume it, under the LangChain instrumentor, and writes +every span and log record it produced to ``sample-output.json`` beside this +file. Needs no API key and no collector. + +Usage, from this directory:: + + python main.py + +See `README.rst `_ for what the events mean and where they come +from. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, TypedDict + +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.graph import END, START, StateGraph +from langgraph.types import Command, interrupt + +from opentelemetry.instrumentation.genai.langchain import LangChainInstrumentor +from opentelemetry.sdk._logs import LoggerProvider +from opentelemetry.sdk._logs.export import ( + InMemoryLogRecordExporter, + SimpleLogRecordProcessor, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + +OUTPUT = Path(__file__).with_name("sample-output.json") + + +class ExpenseState(TypedDict, total=False): + amount: int + approval: str + status: str + + +def build_graph(): + def prepare(_state: ExpenseState) -> ExpenseState: + return {"status": "prepared"} + + def await_approval(state: ExpenseState) -> ExpenseState: + decision = interrupt( + {"question": "approve expense?", "amount": state.get("amount")} + ) + return {"approval": str(decision)} + + def submit(_state: ExpenseState) -> ExpenseState: + return {"status": "submitted"} + + builder = StateGraph(ExpenseState) + builder.add_node("prepare", prepare) + builder.add_node("await_approval", await_approval) + builder.add_node("submit", submit) + builder.add_edge(START, "prepare") + builder.add_edge("prepare", "await_approval") + builder.add_edge("await_approval", "submit") + builder.add_edge("submit", END) + return builder.compile(checkpointer=InMemorySaver()) + + +def _span_json(span: Any) -> dict[str, Any]: + return { + "name": span.name, + "trace_id": f"{span.context.trace_id:032x}", + "span_id": f"{span.context.span_id:016x}", + "parent_span_id": ( + f"{span.parent.span_id:016x}" if span.parent else None + ), + "kind": str(span.kind), + "attributes": dict(span.attributes or {}), + } + + +def _log_json(record: Any) -> dict[str, Any]: + return { + "event_name": record.event_name, + "body": record.body, + "trace_id": f"{record.trace_id:032x}" if record.trace_id else None, + "span_id": f"{record.span_id:016x}" if record.span_id else None, + "attributes": dict(record.attributes or {}), + } + + +def main() -> None: + span_exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + log_exporter = InMemoryLogRecordExporter() + logger_provider = LoggerProvider() + logger_provider.add_log_record_processor( + SimpleLogRecordProcessor(log_exporter) + ) + + instrumentor = LangChainInstrumentor() + instrumentor.instrument( + tracer_provider=tracer_provider, logger_provider=logger_provider + ) + try: + graph = build_graph() + config = {"configurable": {"thread_id": "expense-4711"}} + paused = graph.invoke({"amount": 250}, config=config) + interrupts = [ + {"id": item.id, "value": item.value} + for item in paused["__interrupt__"] + ] + resumed = graph.invoke(Command(resume="approved"), config=config) + finally: + instrumentor.uninstrument() + + document = { + "description": ( + "Telemetry captured by opentelemetry-instrumentation-genai-" + "langchain from an unmodified LangGraph human-in-the-loop " + "application: one invoke that interrupts, one that resumes." + ), + "langgraph_ground_truth": { + "interrupts_returned_by_invoke": interrupts, + "final_state": resumed, + }, + "spans": [ + _span_json(span) for span in span_exporter.get_finished_spans() + ], + "logs": [ + _log_json(item.log_record) + for item in log_exporter.get_finished_logs() + ], + } + OUTPUT.write_text(json.dumps(document, indent=2) + "\n") + print(f"wrote {OUTPUT}") + + +if __name__ == "__main__": + main() diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/requirements.txt b/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/requirements.txt new file mode 100644 index 000000000..e8429074c --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/requirements.txt @@ -0,0 +1,3 @@ +langgraph>=1.2.11 +opentelemetry-sdk>=1.43.0 +opentelemetry-instrumentation-genai-langchain~=1.0b0 \ No newline at end of file diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/sample-output.json b/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/sample-output.json new file mode 100644 index 000000000..b28c5f058 --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/examples/langgraph-lifecycle/sample-output.json @@ -0,0 +1,112 @@ +{ + "description": "Telemetry captured by opentelemetry-instrumentation-genai-langchain from an unmodified LangGraph human-in-the-loop application: one invoke that interrupts, one that resumes.", + "langgraph_ground_truth": { + "interrupts_returned_by_invoke": [ + { + "id": "87c8f19ab465c4c0e2e1bae6c078e9b7", + "value": { + "question": "approve expense?", + "amount": 250 + } + } + ], + "final_state": { + "amount": 250, + "approval": "approved", + "status": "submitted" + } + }, + "spans": [ + { + "name": "invoke_workflow LangGraph", + "trace_id": "3eac7ea19b2e4a9f8956ca033826dc83", + "span_id": "5dc78fc8da41ee72", + "parent_span_id": null, + "kind": "SpanKind.INTERNAL", + "attributes": { + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": "LangGraph", + "gen_ai.conversation.id": "expense-4711" + } + }, + { + "name": "invoke_workflow LangGraph", + "trace_id": "cf16b2377bc87086b7a47ebefc7216c2", + "span_id": "74ccb3f536f343fa", + "parent_span_id": null, + "kind": "SpanKind.INTERNAL", + "attributes": { + "gen_ai.operation.name": "invoke_workflow", + "gen_ai.workflow.name": "LangGraph", + "gen_ai.conversation.id": "expense-4711" + } + } + ], + "logs": [ + { + "event_name": "gen_ai.agent.checkpointed", + "body": "Agent execution checkpointed", + "trace_id": "3eac7ea19b2e4a9f8956ca033826dc83", + "span_id": "5dc78fc8da41ee72", + "attributes": { + "gen_ai.agent.checkpoint.id": "1f1a6616-3d25-63ee-bfff-fb017d557da4" + } + }, + { + "event_name": "gen_ai.agent.checkpointed", + "body": "Agent execution checkpointed", + "trace_id": "3eac7ea19b2e4a9f8956ca033826dc83", + "span_id": "5dc78fc8da41ee72", + "attributes": { + "gen_ai.agent.checkpoint.id": "1f1a6616-3d26-6c50-8000-200196557dff" + } + }, + { + "event_name": "gen_ai.agent.checkpointed", + "body": "Agent execution checkpointed", + "trace_id": "3eac7ea19b2e4a9f8956ca033826dc83", + "span_id": "5dc78fc8da41ee72", + "attributes": { + "gen_ai.agent.checkpoint.id": "1f1a6616-3d28-66d6-8001-7695a4dc6958" + } + }, + { + "event_name": "gen_ai.agent.paused", + "body": "Agent execution paused", + "trace_id": "3eac7ea19b2e4a9f8956ca033826dc83", + "span_id": "5dc78fc8da41ee72", + "attributes": { + "gen_ai.agent.pause.id": "87c8f19ab465c4c0e2e1bae6c078e9b7", + "gen_ai.agent.checkpoint.id": "1f1a6616-3d28-66d6-8001-7695a4dc6958" + } + }, + { + "event_name": "gen_ai.agent.resumed", + "body": "Agent execution resumed", + "trace_id": "cf16b2377bc87086b7a47ebefc7216c2", + "span_id": "74ccb3f536f343fa", + "attributes": { + "gen_ai.agent.resumed_from.type": "checkpoint", + "gen_ai.agent.resumed_from.id": "1f1a6616-3d28-66d6-8001-7695a4dc6958" + } + }, + { + "event_name": "gen_ai.agent.checkpointed", + "body": "Agent execution checkpointed", + "trace_id": "cf16b2377bc87086b7a47ebefc7216c2", + "span_id": "74ccb3f536f343fa", + "attributes": { + "gen_ai.agent.checkpoint.id": "1f1a6616-3d2d-60a8-8002-e4b41fe7ede8" + } + }, + { + "event_name": "gen_ai.agent.checkpointed", + "body": "Agent execution checkpointed", + "trace_id": "cf16b2377bc87086b7a47ebefc7216c2", + "span_id": "74ccb3f536f343fa", + "attributes": { + "gen_ai.agent.checkpoint.id": "1f1a6616-3d2e-64c0-8003-6f1956b33e67" + } + } + ] +} diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/__init__.py b/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/__init__.py index 5ddda7042..2b2c51162 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/__init__.py @@ -38,6 +38,10 @@ from opentelemetry.instrumentation.genai.langchain.callback_handler import ( OpenTelemetryLangChainCallbackHandler, ) +from opentelemetry.instrumentation.genai.langchain.lifecycle import ( + instrument_checkpointers, + uninstrument_checkpointers, +) from opentelemetry.instrumentation.genai.langchain.package import _instruments from opentelemetry.instrumentation.instrumentor import BaseInstrumentor from opentelemetry.instrumentation.utils import unwrap @@ -86,6 +90,10 @@ def _instrument(self, **kwargs: Any): ) self._instrument_agent_entry_points() + # LangGraph only: report every checkpoint the graph's saver persists. + # No-op when LangGraph is not installed. + instrument_checkpointers(otel_callback_handler) + @staticmethod def _instrument_agent_entry_points() -> None: """Recover the create_agent provenance the callback metadata does not carry.""" @@ -105,6 +113,7 @@ def _uninstrument(self, **kwargs: Any): Cleanup instrumentation (unwrap). """ unwrap("langchain_core.callbacks.base.BaseCallbackManager", "__init__") + uninstrument_checkpointers() try: import langgraph.pregel diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/callback_handler.py b/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/callback_handler.py index 054e4874c..24ecff947 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/callback_handler.py +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/callback_handler.py @@ -23,6 +23,16 @@ from opentelemetry.instrumentation.genai.langchain.invocation_manager import ( _InvocationManager, ) +from opentelemetry.instrumentation.genai.langchain.lifecycle import ( + ATTR_CHECKPOINT_ID, + ATTR_PAUSE_ID, + ATTR_RESUMED_FROM_ID, + ATTR_RESUMED_FROM_TYPE, + EVENT_AGENT_CHECKPOINTED, + EVENT_AGENT_PAUSED, + EVENT_AGENT_RESUMED, + RESUMED_FROM_TYPE_CHECKPOINT, +) from opentelemetry.instrumentation.genai.langchain.operation_mapping import ( OperationName, classify_chain_run, @@ -135,6 +145,7 @@ def on_chain_start( self._invocation_manager.add_invocation_state( run_id, parent_run_id, workflow ) + self._bind_langgraph_thread(run_id, metadata) elif operation == OperationName.INVOKE_AGENT: # agent name passed by the user suggested_agent_name = resolve_agent_name( @@ -215,6 +226,7 @@ def on_chain_end( parent_run_id: UUID | None = None, **kwargs: Any, ) -> Any: + self._invocation_manager.unbind_thread(run_id) invocation = self._invocation_manager.get_invocation(run_id=run_id) if invocation is None or not isinstance( invocation, (WorkflowInvocation, AgentInvocation) @@ -239,6 +251,7 @@ def on_chain_error( parent_run_id: UUID | None = None, **kwargs: Any, ) -> Any: + self._invocation_manager.unbind_thread(run_id) invocation = self._invocation_manager.get_invocation(run_id=run_id) if invocation is None or not isinstance( invocation, (WorkflowInvocation, AgentInvocation) @@ -744,6 +757,92 @@ def on_retriever_error( if not invocation.span.is_recording(): self._invocation_manager.delete_invocation_state(run_id=run_id) + # ------------------------------------------------------------------ + # LangGraph durable-execution lifecycle. + # + # LangGraph dispatches ``on_interrupt``/``on_resume`` to the handlers in the + # run's callback manager (see ``langgraph.callbacks``), so this handler + # receives them through the callback manager it is already injected into. + # The events are typed dataclasses in LangGraph, but they are read + # structurally here so the instrumentation keeps no LangGraph import. + # ------------------------------------------------------------------ + + def on_interrupt(self, event: Any) -> None: + """Emit one paused event per interrupt that suspended the graph.""" + workflow = self._find_nearest_workflow(getattr(event, "run_id", None)) + if workflow is None: + return + + checkpoint_id = getattr(event, "checkpoint_id", None) + for pause in getattr(event, "interrupts", ()) or (): + pause_id = getattr(pause, "id", None) + if not pause_id: + continue + attributes: dict[str, str] = {ATTR_PAUSE_ID: str(pause_id)} + # ``gen_ai.agent.pause.reason`` is deliberately absent: LangGraph's + # ``interrupt(value)`` carries an opaque application payload and + # nothing that says who or what is expected to resolve the pause. + if checkpoint_id: + attributes[ATTR_CHECKPOINT_ID] = str(checkpoint_id) + workflow.emit_event( + EVENT_AGENT_PAUSED, + attributes, + body="Agent execution paused", + ) + + def on_resume(self, event: Any) -> None: + """Emit a resumed event for a graph continuing from a checkpoint.""" + workflow = self._find_nearest_workflow(getattr(event, "run_id", None)) + checkpoint_id = getattr(event, "checkpoint_id", None) + if workflow is None or not checkpoint_id: + return + + workflow.emit_event( + EVENT_AGENT_RESUMED, + { + ATTR_RESUMED_FROM_TYPE: RESUMED_FROM_TYPE_CHECKPOINT, + ATTR_RESUMED_FROM_ID: str(checkpoint_id), + }, + body="Agent execution resumed", + ) + + def checkpoint_written( + self, thread_id: str, checkpoint_id: str, namespace: str = "" + ) -> None: + """Emit a checkpointed event for one persisted checkpoint. + + Called by the wrapped checkpointer, which has no callback run id, so + the live graph run is resolved through the LangGraph thread id and the + checkpoint namespace the write was made under. + """ + workflow = self._invocation_manager.get_thread_invocation( + thread_id, namespace + ) + if not isinstance(workflow, WorkflowInvocation): + return + + workflow.emit_event( + EVENT_AGENT_CHECKPOINTED, + {ATTR_CHECKPOINT_ID: checkpoint_id}, + body="Agent execution checkpointed", + ) + + def _bind_langgraph_thread( + self, run_id: UUID, metadata: dict[str, Any] | None + ) -> None: + if not metadata or metadata.get("ls_integration") != "langgraph": + return + thread_id = metadata.get("thread_id") + if not thread_id: + return + # A top level graph run has no ``langgraph_checkpoint_ns`` and writes + # its checkpoints under the root namespace. A nested graph run carries + # the namespace its own checkpoints are written under. + namespace = metadata.get("langgraph_checkpoint_ns") or "" + self._invocation_manager.bind_thread( + str(thread_id), run_id, str(namespace) + ) + def _find_agent_context( self, run_id: UUID | None ) -> tuple[AgentInvocation | None, set[str]]: @@ -761,3 +860,16 @@ def _find_agent_context( ancestor_agent_names.add(entity.agent_name.lower()) current = self._invocation_manager.get_parent_run_id(current) return nearest_agent, ancestor_agent_names + + def _find_nearest_workflow( + self, run_id: UUID | None + ) -> WorkflowInvocation | None: + current = run_id + visited: set[UUID] = set() + while current is not None and current not in visited: + visited.add(current) + entity = self._invocation_manager.get_invocation(current) + if isinstance(entity, WorkflowInvocation): + return entity + current = self._invocation_manager.get_parent_run_id(current) + return None diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/invocation_manager.py b/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/invocation_manager.py index b85309ce4..3d1e153cf 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/invocation_manager.py +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/invocation_manager.py @@ -1,6 +1,8 @@ # Copyright The OpenTelemetry Authors # SPDX-License-Identifier: Apache-2.0 +import logging +import threading from dataclasses import dataclass, field from uuid import UUID @@ -8,6 +10,25 @@ __all__ = ["_InvocationManager"] +_logger = logging.getLogger(__name__) + +# Upper bound on LangGraph thread ids tracked at once, and on runs tracked per +# thread id. Runs that never report an end (a crashed worker, a cancelled task) +# would otherwise accumulate forever. +_MAX_TRACKED_THREADS = 512 +_MAX_RUNS_PER_THREAD = 32 + +# LangGraph joins nested checkpoint namespaces with this separator +# (``langgraph.constants.NS_SEP``). +_NAMESPACE_SEPARATOR = "|" + + +def _namespace_contains(outer: str, inner: str) -> bool: + """Return whether ``outer`` is ``inner`` or encloses it.""" + if not outer: + return True + return inner == outer or inner.startswith(outer + _NAMESPACE_SEPARATOR) + @dataclass class _InvocationState: @@ -24,6 +45,118 @@ def __init__( # Map from run_id -> _InvocationState, to keep track of invocations and parent/child relationships # TODO: TTL cache to avoid memory leaks in long-running processes. self._invocations: dict[UUID, _InvocationState] = {} + # Map from LangGraph thread id -> the live graph runs on that thread + # id, outermost first, each paired with the checkpoint namespace its + # own checkpoints are written under ("" for a top level graph). A + # checkpointer call carries no callback run id, so the thread id and + # namespace in its config are the only handle back to a running + # invocation. This is a list rather than a single id so that a nested + # graph run never displaces the run that contains it. + self._threads: dict[str, list[tuple[str, UUID]]] = {} + # One lock guards both maps: a checkpoint lookup reads the thread stack + # and the invocation states together, and must not observe a run being + # torn down halfway through by a concurrent chain end. + self._lock = threading.RLock() + + # ------------------------------------------------------------------ + # LangGraph thread correlation + # ------------------------------------------------------------------ + + def bind_thread( + self, thread_id: str, run_id: UUID, namespace: str = "" + ) -> None: + """Record that ``run_id`` is a live graph run on ``thread_id``. + + ``namespace`` is the LangGraph checkpoint namespace the run's own + checkpoints are written under: "" for a top level graph, and the + run's ``langgraph_checkpoint_ns`` for a nested graph. + """ + with self._lock: + runs = self._threads.setdefault(thread_id, []) + # Drop runs whose invocation state is already gone, so abandoned + # runs do not pin an entry forever. + runs[:] = [ + entry for entry in runs if entry[1] in self._invocations + ] + if all(entry[1] != run_id for entry in runs): + runs.append((namespace, run_id)) + del runs[:-_MAX_RUNS_PER_THREAD] + self._prune_threads() + + def unbind_thread(self, run_id: UUID) -> None: + """Forget ``run_id``, keeping any run that contains or follows it.""" + with self._lock: + for thread_id, runs in list(self._threads.items()): + runs[:] = [entry for entry in runs if entry[1] != run_id] + if not runs: + del self._threads[thread_id] + + def get_thread_invocation( + self, thread_id: str, namespace: str = "" + ) -> GenAIInvocation | None: + """Return the invocation a checkpoint write belongs to. + + A write carries the namespace of the graph that produced it. The run + bound with that exact namespace owns it; if no run is bound there, the + nearest enclosing run does, which is the live run whose namespace is + the longest prefix of the write's namespace. If two live runs are + equally specific and unrelated, ownership is genuinely ambiguous and + nothing is returned: a miscorrelated event is worse than a missing one. + """ + with self._lock: + candidates = [ + entry + for entry in self._threads.get(thread_id, []) + if entry[1] in self._invocations + and _namespace_contains(entry[0], namespace) + ] + if not candidates: + return None + + longest = max(len(entry[0]) for entry in candidates) + matches = [ + entry for entry in candidates if len(entry[0]) == longest + ] + if len(matches) > 1: + # Equally specific runs are only unambiguous when they nest. + innermost = matches[-1][1] + ancestors = self._ancestors(innermost) + if not all(entry[1] in ancestors for entry in matches[:-1]): + _logger.debug( + "Ambiguous LangGraph thread id %s namespace %r: %d " + "unrelated live runs, dropping the checkpoint event", + thread_id, + namespace, + len(matches), + ) + return None + + state = self._invocations.get(matches[-1][1]) + return state.invocation if state else None + + def _ancestors(self, run_id: UUID) -> set[UUID]: + """Return the run ids between ``run_id`` and the root, exclusive.""" + ancestors: set[UUID] = set() + state = self._invocations.get(run_id) + while state is not None and state.parent_run_id is not None: + parent_run_id = state.parent_run_id + if parent_run_id in ancestors: + break + ancestors.add(parent_run_id) + state = self._invocations.get(parent_run_id) + return ancestors + + def _prune_threads(self) -> None: + """Drop the oldest thread ids once the tracked set grows too large.""" + overflow = len(self._threads) - _MAX_TRACKED_THREADS + if overflow <= 0: + return + for thread_id in list(self._threads)[:overflow]: + del self._threads[thread_id] + + # ------------------------------------------------------------------ + # Run state + # ------------------------------------------------------------------ def add_invocation_state( self, @@ -34,39 +167,48 @@ def add_invocation_state( invocation_state = _InvocationState(invocation=invocation) invocation_state.parent_run_id = parent_run_id - if parent_run_id is not None and parent_run_id in self._invocations: - parent_invocation_state = self._invocations[parent_run_id] - parent_invocation_state.children.append(run_id) + with self._lock: + if ( + parent_run_id is not None + and parent_run_id in self._invocations + ): + parent_invocation_state = self._invocations[parent_run_id] + parent_invocation_state.children.append(run_id) - self._invocations[run_id] = invocation_state + self._invocations[run_id] = invocation_state def get_invocation(self, run_id: UUID) -> GenAIInvocation | None: - invocation_state = self._invocations.get(run_id) - return invocation_state.invocation if invocation_state else None + with self._lock: + invocation_state = self._invocations.get(run_id) + return invocation_state.invocation if invocation_state else None def get_parent_run_id(self, run_id: UUID) -> UUID | None: - invocation_state = self._invocations.get(run_id) - return invocation_state.parent_run_id if invocation_state else None + with self._lock: + invocation_state = self._invocations.get(run_id) + return invocation_state.parent_run_id if invocation_state else None def delete_invocation_state(self, run_id: UUID) -> None: - invocation_state = self._invocations.get(run_id) - if not invocation_state: - return - - invocation_state.ended = True - - # Defer removal if any children are still live, so upward traversal - # via _find_agent_context can still walk through this node. - if any(c in self._invocations for c in invocation_state.children): - return - - self._invocations.pop(run_id, None) - - # Propagate cleanup upward: if the parent has already ended and has no - # more live children, it can now be removed too. - if invocation_state.parent_run_id: - parent_state = self._invocations.get( - invocation_state.parent_run_id - ) - if parent_state is not None and parent_state.ended: - self.delete_invocation_state(invocation_state.parent_run_id) + with self._lock: + invocation_state = self._invocations.get(run_id) + if not invocation_state: + return + + invocation_state.ended = True + + # Defer removal if any children are still live, so upward traversal + # via _find_agent_context can still walk through this node. + if any(c in self._invocations for c in invocation_state.children): + return + + self._invocations.pop(run_id, None) + + # Propagate cleanup upward: if the parent has already ended and has + # no more live children, it can now be removed too. + if invocation_state.parent_run_id: + parent_state = self._invocations.get( + invocation_state.parent_run_id + ) + if parent_state is not None and parent_state.ended: + self.delete_invocation_state( + invocation_state.parent_run_id + ) diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/lifecycle.py b/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/lifecycle.py new file mode 100644 index 000000000..1161edcd8 --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/lifecycle.py @@ -0,0 +1,339 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +"""Agent lifecycle telemetry for LangGraph durable executions. + +Every name in this module is a candidate semantic convention proposed in +open-telemetry/semantic-conventions-genai#445 (agent lifecycle events for async +and long running executions). None of them is stable, so they are kept here as +plain string constants instead of being imported from the semconv package. + +Two interception points feed these events, both of them generic over any +LangGraph application: + +* ``langgraph.callbacks`` lifecycle dispatch. LangGraph calls ``on_interrupt`` + and ``on_resume`` on the handlers registered in the run's callback manager, + passing the real ``Interrupt`` objects and the checkpoint id the graph + paused at or resumed from. +* ``StateGraph.compile(checkpointer=...)``. The checkpointer the compiled graph + retains is wrapped so that every ``put``/``aput`` (LangGraph writes one + checkpoint per superstep) reports the id it persisted. +""" + +from __future__ import annotations + +import logging +import threading +import weakref +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, Protocol +from weakref import WeakKeyDictionary + +from langchain_core.runnables import RunnableConfig +from wrapt import wrap_function_wrapper + +from opentelemetry.instrumentation.utils import unwrap + +if TYPE_CHECKING: # pragma: no cover - typing only + from opentelemetry.instrumentation.genai.langchain.callback_handler import ( + OpenTelemetryLangChainCallbackHandler, + ) + +__all__ = [ + "ATTR_CHECKPOINT_ID", + "ATTR_PAUSE_ID", + "ATTR_PAUSE_REASON", + "ATTR_RESUMED_FROM_ID", + "ATTR_RESUMED_FROM_TYPE", + "EVENT_AGENT_CHECKPOINTED", + "EVENT_AGENT_PAUSED", + "EVENT_AGENT_RESUMED", + "RESUMED_FROM_TYPE_CHECKPOINT", + "instrument_checkpointers", + "uninstrument_checkpointers", +] + +_logger = logging.getLogger(__name__) + +# Candidate event names, pending open-telemetry/semantic-conventions-genai#445. +EVENT_AGENT_PAUSED = "gen_ai.agent.paused" +EVENT_AGENT_CHECKPOINTED = "gen_ai.agent.checkpointed" +EVENT_AGENT_RESUMED = "gen_ai.agent.resumed" + +# Candidate attribute names, pending the same proposal. +ATTR_PAUSE_ID = "gen_ai.agent.pause.id" +ATTR_PAUSE_REASON = "gen_ai.agent.pause.reason" +ATTR_CHECKPOINT_ID = "gen_ai.agent.checkpoint.id" +ATTR_RESUMED_FROM_TYPE = "gen_ai.agent.resumed_from.type" +ATTR_RESUMED_FROM_ID = "gen_ai.agent.resumed_from.id" + +# LangGraph's resume event supplies a checkpoint id and nothing else, so this is +# the only member of ``gen_ai.agent.resumed_from.type`` the library determines. +RESUMED_FROM_TYPE_CHECKPOINT = "checkpoint" + +_WRAPPED_METHODS = ("put", "aput") +_MISSING = object() + +# Checkpointer instances patched by this instrumentation, mapped to the +# instance attributes they had before patching, so ``uninstrument`` restores +# rather than deletes. +_wrapped_checkpointers: WeakKeyDictionary[Any, dict[str, Any]] = ( + WeakKeyDictionary() +) +_wrapped_lock = threading.Lock() + + +class _Reporter(Protocol): + def checkpoint_written( + self, thread_id: str, checkpoint_id: str, namespace: str + ) -> None: ... + + +def _configurable_value( + config: RunnableConfig | None, key: str, default: str | None = None +) -> str | None: + """Return one ``configurable`` value from a runnable config.""" + if not config: + return default + configurable = config.get("configurable") + if not configurable: + return default + value = configurable.get(key) + if value is None: + return default + return str(value) + + +def _config_arg( + args: tuple[Any, ...], kwargs: dict[str, Any] +) -> RunnableConfig | None: + """Return the ``config`` argument of a ``put``/``aput`` call.""" + if "config" in kwargs: + return kwargs["config"] + return args[0] if args else None + + +def _checkpoint_arg(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: + """Return the ``checkpoint`` argument of a ``put``/``aput`` call.""" + if "checkpoint" in kwargs: + return kwargs["checkpoint"] + return args[1] if len(args) > 1 else None + + +class _CheckpointReporter: + """Report each persisted checkpoint exactly once. + + A saver may implement ``aput`` by delegating to ``put`` (LangGraph's own + ``InMemorySaver`` does ``return self.put(config, checkpoint, metadata, + new_versions)``), and the delegated call may run on a worker thread, so + nesting cannot be detected from the call context. Instead the identity of + the ``Checkpoint`` object being written is held while a write is in + flight: a nested call that receives the same object is the delegation and + stays silent, and the outermost call reports. + + Known limit: a saver that copies the checkpoint before delegating is not + recognised as a delegation, and the write is reported by both calls. Over + reporting is preferred to dropping a real checkpoint. + """ + + def __init__(self, reporter: _Reporter) -> None: + self._reporter = reporter + self._in_flight: set[int] = set() + self._lock = threading.Lock() + + def claim(self, checkpoint: Any) -> int | None: + """Claim a write, or return None if it is a nested delegation.""" + if checkpoint is None: + return None + key = id(checkpoint) + with self._lock: + if key in self._in_flight: + return None + self._in_flight.add(key) + return key + + def release(self, key: int | None) -> None: + if key is None: + return + with self._lock: + self._in_flight.discard(key) + + def report( + self, + config: RunnableConfig | None, + returned_config: RunnableConfig | None, + ) -> None: + thread_id = _configurable_value(config, "thread_id") + checkpoint_id = _configurable_value(returned_config, "checkpoint_id") + if not thread_id or not checkpoint_id: + return + namespace = _configurable_value(config, "checkpoint_ns", "") or "" + self._reporter.checkpoint_written(thread_id, checkpoint_id, namespace) + + +def _supports_tracking(checkpointer: Any) -> bool: + """Return whether the saver can be tracked for later restoration. + + Checked before anything is patched: a saver that cannot be put in the + registry must not be patched either, or uninstrument could never undo it. + """ + try: + hash(checkpointer) + weakref.ref(checkpointer) + except TypeError: + return False + return True + + +def _wrap_checkpointer(checkpointer: Any, reporter: _Reporter) -> None: + """Wrap one checkpointer instance's write methods. + + ``BaseCheckpointSaver.put`` is abstract, so wrapping the base class + intercepts nothing: every saver overrides it. The instance is patched + instead, which also keeps the patch scoped to savers an instrumented + application actually compiled a graph with. + """ + # Checked before the registry is touched: an unhashable or + # non-weak-referenceable saver cannot even be looked up there. + if not _supports_tracking(checkpointer): + _logger.debug( + "Checkpointer %r cannot be tracked for uninstrument, " + "skipping checkpoint events for it", + type(checkpointer).__name__, + ) + return + + with _wrapped_lock: + if checkpointer in _wrapped_checkpointers: + return + + checkpoint_reporter = _CheckpointReporter(reporter) + + def sync_put( + wrapped: Callable[..., Any], + _instance: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> Any: + key = checkpoint_reporter.claim(_checkpoint_arg(args, kwargs)) + try: + result = wrapped(*args, **kwargs) + finally: + checkpoint_reporter.release(key) + if key is not None: + checkpoint_reporter.report(_config_arg(args, kwargs), result) + return result + + async def async_put( + wrapped: Callable[..., Any], + _instance: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> Any: + key = checkpoint_reporter.claim(_checkpoint_arg(args, kwargs)) + try: + result = await wrapped(*args, **kwargs) + finally: + checkpoint_reporter.release(key) + if key is not None: + checkpoint_reporter.report(_config_arg(args, kwargs), result) + return result + + # Remember the pre-patch instance attributes so uninstrument restores + # a saver that legitimately supplied ``put`` or ``aput`` per instance. + instance_dict: dict[str, Any] = getattr(checkpointer, "__dict__", {}) + previous: dict[str, Any] = { + name: instance_dict.get(name, _MISSING) + for name in _WRAPPED_METHODS + } + + patched: list[str] = [] + for name, wrapper in (("put", sync_put), ("aput", async_put)): + if getattr(checkpointer, name, None) is None: + continue + try: + wrap_function_wrapper(checkpointer, name, wrapper) + except (AttributeError, TypeError): # pragma: no cover + continue + patched.append(name) + + if not patched: + return + _wrapped_checkpointers[checkpointer] = { + name: previous[name] for name in patched + } + + +def _restore_checkpointer(checkpointer: Any, previous: dict[str, Any]) -> None: + """Undo one instance patch, restoring any pre-existing attribute.""" + for name, original in previous.items(): + unwrap(checkpointer, name) + instance_dict: dict[str, Any] | None = getattr( + checkpointer, "__dict__", None + ) + if instance_dict is None: + continue + if original is _MISSING: + instance_dict.pop(name, None) + else: + instance_dict[name] = original + + +class _CompileWrapper: + """Wrap ``StateGraph.compile`` to reach the checkpointer it was given.""" + + def __init__(self, reporter: _Reporter) -> None: + self._reporter = reporter + + def __call__( + self, + wrapped: Callable[..., Any], + _instance: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> Any: + compiled = wrapped(*args, **kwargs) + # Prefer the saver the compiled graph actually retains, so the patch + # follows what LangGraph will call rather than what was passed in. + checkpointer = getattr(compiled, "checkpointer", None) + if checkpointer is None: + checkpointer = kwargs.get("checkpointer") + if checkpointer is None and args: + checkpointer = args[0] + # ``checkpointer`` is also allowed to be ``None`` or a bool (a subgraph + # inheriting the parent's saver); only real savers can be wrapped. + if checkpointer is not None and not isinstance(checkpointer, bool): + _wrap_checkpointer(checkpointer, self._reporter) + return compiled + + +def instrument_checkpointers( + callback_handler: OpenTelemetryLangChainCallbackHandler, +) -> bool: + """Wrap ``StateGraph.compile``. Returns False when LangGraph is absent.""" + try: + wrap_function_wrapper( + "langgraph.graph.state", + "StateGraph.compile", + _CompileWrapper(callback_handler), + ) + except (ImportError, AttributeError): + return False + return True + + +def uninstrument_checkpointers() -> None: + """Undo ``instrument_checkpointers`` and restore patched savers.""" + try: + unwrap("langgraph.graph.state.StateGraph", "compile") + except ( + ImportError, + AttributeError, + ): # pragma: no cover - langgraph absent + pass + + with _wrapped_lock: + tracked = list(_wrapped_checkpointers.items()) + _wrapped_checkpointers.clear() + for checkpointer, previous in tracked: + _restore_checkpointer(checkpointer, previous) diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_callback_handler.py b/instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_callback_handler.py index baac28150..b4e6c80a7 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_callback_handler.py +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_callback_handler.py @@ -2689,6 +2689,62 @@ def test_on_chat_model_start_captures_input_messages_when_content_enabled(): assert blob.content == _REAL_PNG_BYTES +class TestLangGraphThreadBinding: + """A LangGraph thread binding must not outlive its run.""" + + def _start_langgraph_workflow(self, handler, run_id): + handler.on_chain_start( + serialized={"name": "LangGraph"}, + inputs={}, + run_id=run_id, + metadata={"ls_integration": "langgraph", "thread_id": "bound"}, + ) + + def test_chain_error_unbinds_the_thread(self): + handler, _, _, _ = _make_handler() + run_id = _run_id() + self._start_langgraph_workflow(handler, run_id) + + # The binding exists before the failure. + assert handler._invocation_manager._threads["bound"] == [("", run_id)] + + handler.on_chain_error(ValueError("boom"), run_id=run_id) + + assert "bound" not in handler._invocation_manager._threads + assert ( + handler._invocation_manager.get_thread_invocation("bound") is None + ) + + def test_chain_end_unbinds_the_thread(self): + handler, _, _, _ = _make_handler() + run_id = _run_id() + self._start_langgraph_workflow(handler, run_id) + + assert handler._invocation_manager._threads["bound"] == [("", run_id)] + + handler.on_chain_end(outputs={}, run_id=run_id) + + assert "bound" not in handler._invocation_manager._threads + + def test_nested_graph_run_binds_its_own_checkpoint_namespace(self): + handler, _, _, _ = _make_handler() + run_id = _run_id() + handler.on_chain_start( + serialized={"name": "LangGraph"}, + inputs={}, + run_id=run_id, + metadata={ + "ls_integration": "langgraph", + "thread_id": "bound", + "langgraph_checkpoint_ns": "child:abc", + }, + ) + + assert handler._invocation_manager._threads["bound"] == [ + ("child:abc", run_id) + ] + + def test_on_chat_model_start_preserves_message_name(): run_id = _run_id() handler, telemetry, llm_inv = _make_handler_with_llm_invocation(run_id) diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_invocation_manager.py b/instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_invocation_manager.py index 7a4796205..9b9760066 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_invocation_manager.py +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_invocation_manager.py @@ -240,5 +240,183 @@ def test_none_invocation_can_be_stored_and_retrieved(invocation_manager): assert invocation_manager.get_invocation(run_id) is None +def test_thread_binding_resolves_the_running_invocation( + invocation_manager, mock_invocation +): + run_id = uuid.uuid4() + invocation_manager.add_invocation_state( + run_id=run_id, parent_run_id=None, invocation=mock_invocation + ) + + invocation_manager.bind_thread("thread-1", run_id) + + assert ( + invocation_manager.get_thread_invocation("thread-1") is mock_invocation + ) + assert invocation_manager.get_thread_invocation("thread-2") is None + + invocation_manager.unbind_thread(run_id) + + assert invocation_manager.get_thread_invocation("thread-1") is None + + def test_delete_nonexistent_run_id_does_not_raise(invocation_manager): invocation_manager.delete_invocation_state(uuid.uuid4()) # must not raise + + +# --------------------------------------------------------------------------- +# LangGraph thread correlation +# --------------------------------------------------------------------------- + + +def _bind_run(manager, thread_id, parent_run_id=None): + run_id = uuid.uuid4() + invocation = mock.Mock(spec=GenAIInvocation) + manager.add_invocation_state( + run_id=run_id, parent_run_id=parent_run_id, invocation=invocation + ) + manager.bind_thread(thread_id, run_id) + return run_id, invocation + + +def test_nested_run_never_displaces_the_run_containing_it(invocation_manager): + outer_run_id, outer_invocation = _bind_run(invocation_manager, "t") + _, inner_invocation = _bind_run( + invocation_manager, "t", parent_run_id=outer_run_id + ) + + # The innermost live run owns the checkpoint. + assert invocation_manager.get_thread_invocation("t") is inner_invocation + + # When the nested run ends, the run containing it keeps the thread id. + inner_run_id = invocation_manager._threads["t"][-1][1] + invocation_manager.unbind_thread(inner_run_id) + invocation_manager.delete_invocation_state(inner_run_id) + + assert invocation_manager._threads["t"] == [("", outer_run_id)] + assert invocation_manager.get_thread_invocation("t") is outer_invocation + + +def test_unrelated_concurrent_runs_on_one_thread_are_dropped( + invocation_manager, +): + _bind_run(invocation_manager, "t") + _bind_run(invocation_manager, "t") + + # Ownership is ambiguous, so nothing is returned rather than guessed. + assert invocation_manager.get_thread_invocation("t") is None + + +def test_dead_runs_are_pruned_from_the_thread_binding(invocation_manager): + abandoned_run_id, _ = _bind_run(invocation_manager, "t") + invocation_manager.delete_invocation_state(abandoned_run_id) + + _, live_invocation = _bind_run(invocation_manager, "t") + + assert len(invocation_manager._threads["t"]) == 1 + assert invocation_manager.get_thread_invocation("t") is live_invocation + + +def test_thread_bindings_are_bounded_for_runs_that_never_end( + invocation_manager, +): + from opentelemetry.instrumentation.genai.langchain.invocation_manager import ( + _MAX_RUNS_PER_THREAD, + _MAX_TRACKED_THREADS, + ) + + for _ in range(_MAX_RUNS_PER_THREAD + 10): + _bind_run(invocation_manager, "abandoned-thread") + assert ( + len(invocation_manager._threads["abandoned-thread"]) + == _MAX_RUNS_PER_THREAD + ) + + for index in range(_MAX_TRACKED_THREADS + 5): + _bind_run(invocation_manager, f"thread-{index}") + assert len(invocation_manager._threads) <= _MAX_TRACKED_THREADS + + +def test_thread_lookup_is_atomic_against_concurrent_run_teardown( + invocation_manager, +): + import threading + + results = [] + errors = [] + + def churn(): + try: + for _ in range(200): + run_id, invocation = _bind_run(invocation_manager, "t") + invocation_manager.unbind_thread(run_id) + invocation_manager.delete_invocation_state(run_id) + except Exception as exc: # pragma: no cover - failure path + errors.append(exc) + + def read(): + try: + for _ in range(200): + results.append(invocation_manager.get_thread_invocation("t")) + except Exception as exc: # pragma: no cover - failure path + errors.append(exc) + + threads = [threading.Thread(target=churn), threading.Thread(target=read)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert not errors + # Every answer is either a live invocation or nothing, never a torn read. + assert all( + result is None or isinstance(result, mock.Mock) for result in results + ) + + +def test_checkpoint_namespace_selects_the_run_that_owns_it(invocation_manager): + root_run_id = uuid.uuid4() + child_run_id = uuid.uuid4() + root_invocation = mock.Mock(spec=GenAIInvocation) + child_invocation = mock.Mock(spec=GenAIInvocation) + invocation_manager.add_invocation_state( + run_id=root_run_id, parent_run_id=None, invocation=root_invocation + ) + invocation_manager.add_invocation_state( + run_id=child_run_id, + parent_run_id=root_run_id, + invocation=child_invocation, + ) + invocation_manager.bind_thread("t", root_run_id, "") + invocation_manager.bind_thread("t", child_run_id, "child:abc") + + # A write made under the child namespace belongs to the child run. + assert ( + invocation_manager.get_thread_invocation("t", "child:abc") + is child_invocation + ) + # And so does a write from a graph nested inside the child. + assert ( + invocation_manager.get_thread_invocation("t", "child:abc|leaf:def") + is child_invocation + ) + # A root namespace write belongs to the root run. + assert invocation_manager.get_thread_invocation("t", "") is root_invocation + # A namespace no run is bound for falls back to the nearest enclosing run. + assert ( + invocation_manager.get_thread_invocation("t", "other:xyz") + is root_invocation + ) + + +def test_namespace_write_falls_back_when_the_child_run_has_no_binding( + invocation_manager, +): + root_run_id, root_invocation = _bind_run(invocation_manager, "t") + + # Nothing is bound for the child namespace, so the enclosing run owns it. + assert ( + invocation_manager.get_thread_invocation("t", "child:abc") + is root_invocation + ) + assert root_run_id is not None diff --git a/instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_langgraph_lifecycle_events.py b/instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_langgraph_lifecycle_events.py new file mode 100644 index 000000000..b545559cc --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-langchain/tests/test_langgraph_lifecycle_events.py @@ -0,0 +1,704 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +"""Agent lifecycle telemetry captured from unmodified LangGraph applications. + +Every id asserted here comes from LangGraph itself: interrupt ids from the +``Interrupt`` objects the graph returns, checkpoint ids from the checkpointer's +own return values. +""" + +from __future__ import annotations + +import asyncio +import threading +from typing import Any, TypedDict + +import pytest + +from opentelemetry.instrumentation.genai.langchain.lifecycle import ( + _wrap_checkpointer, +) + +InMemorySaver = pytest.importorskip( + "langgraph.checkpoint.memory" +).InMemorySaver +pytest.importorskip("langgraph.callbacks") +_graph = pytest.importorskip("langgraph.graph") +_types = pytest.importorskip("langgraph.types") +END = _graph.END +START = _graph.START +StateGraph = _graph.StateGraph +Command = _types.Command +interrupt = _types.interrupt + +PAUSED = "gen_ai.agent.paused" +CHECKPOINTED = "gen_ai.agent.checkpointed" +RESUMED = "gen_ai.agent.resumed" +LIFECYCLE_EVENTS = (PAUSED, CHECKPOINTED, RESUMED) + + +class _State(TypedDict, total=False): + value: str + approval: str + + +def _build_graph(): + def start(_state: _State) -> _State: + return {"value": "prepared"} + + def approval(_state: _State) -> _State: + return {"approval": str(interrupt({"question": "approve?"}))} + + def finish(_state: _State) -> _State: + return {"value": "submitted"} + + builder = StateGraph(_State) + builder.add_node("start", start) + builder.add_node("approval", approval) + builder.add_node("finish", finish) + builder.add_edge(START, "start") + builder.add_edge("start", "approval") + builder.add_edge("approval", "finish") + builder.add_edge("finish", END) + return builder + + +def _events(log_exporter, event_name: str) -> list[Any]: + return [ + item.log_record + for item in log_exporter.get_finished_logs() + if item.log_record.event_name == event_name + ] + + +def _workflow_spans(span_exporter) -> list[Any]: + return [ + span + for span in span_exporter.get_finished_spans() + if span.attributes + and span.attributes.get("gen_ai.operation.name") == "invoke_workflow" + ] + + +def test_interrupt_and_resume_emit_correlated_lifecycle_events( + start_instrumentation, + log_exporter, + span_exporter, +): + graph = _build_graph().compile(checkpointer=InMemorySaver()) + config = {"configurable": {"thread_id": "lifecycle-test"}} + + paused_result = graph.invoke({"value": "new"}, config=config) + interrupts = paused_result["__interrupt__"] + assert len(interrupts) == 1 + real_interrupt_id = interrupts[0].id + + # paused carries the id LangGraph minted for this interrupt. + paused = _events(log_exporter, PAUSED) + assert len(paused) == 1 + assert paused[0].attributes["gen_ai.agent.pause.id"] == real_interrupt_id + pause_checkpoint = paused[0].attributes["gen_ai.agent.checkpoint.id"] + # No pause reason: LangGraph does not report why execution paused. + assert "gen_ai.agent.pause.reason" not in paused[0].attributes + # No execution id: LangGraph has no id spanning suspend and resume. + assert "gen_ai.agent.execution.id" not in paused[0].attributes + + # LangGraph writes one checkpoint per superstep, so checkpointed is a + # per-step record: the input checkpoint plus one per completed superstep. + first_run_checkpoints = [ + record.attributes["gen_ai.agent.checkpoint.id"] + for record in _events(log_exporter, CHECKPOINTED) + ] + assert len(first_run_checkpoints) == 3 + assert len(set(first_run_checkpoints)) == 3 + # The graph pauses at the checkpoint it last persisted. + assert pause_checkpoint == first_run_checkpoints[-1] + + resumed_result = graph.invoke(Command(resume="approved"), config=config) + assert resumed_result["approval"] == "approved" + + resumed = _events(log_exporter, RESUMED) + assert len(resumed) == 1 + assert ( + resumed[0].attributes["gen_ai.agent.resumed_from.type"] == "checkpoint" + ) + # The resume continues from exactly the checkpoint the pause reported. + assert ( + resumed[0].attributes["gen_ai.agent.resumed_from.id"] + == pause_checkpoint + ) + + second_run_checkpoints = [ + record.attributes["gen_ai.agent.checkpoint.id"] + for record in _events(log_exporter, CHECKPOINTED) + ][len(first_run_checkpoints) :] + assert len(second_run_checkpoints) == 2 + + # Each event is correlated with the workflow span of its own invoke call. + workflow_spans = _workflow_spans(span_exporter) + assert len(workflow_spans) == 2 + paused_run, resumed_run = workflow_spans + for record in _events(log_exporter, PAUSED) + [ + record + for record in _events(log_exporter, CHECKPOINTED) + if record.attributes["gen_ai.agent.checkpoint.id"] + in first_run_checkpoints + ]: + assert record.trace_id == paused_run.context.trace_id + assert record.span_id == paused_run.context.span_id + for record in _events(log_exporter, RESUMED) + [ + record + for record in _events(log_exporter, CHECKPOINTED) + if record.attributes["gen_ai.agent.checkpoint.id"] + in second_run_checkpoints + ]: + assert record.trace_id == resumed_run.context.trace_id + assert record.span_id == resumed_run.context.span_id + + # The two invoke calls are separate traces, and nothing in the telemetry + # ties them together: LangGraph mints no end-to-end execution id. + assert paused_run.context.trace_id != resumed_run.context.trace_id + + +def test_graph_without_checkpointer_emits_no_durability_events( + start_instrumentation, + log_exporter, +): + graph = _build_graph().compile() + + graph.invoke({"value": "new"}, config={"configurable": {"thread_id": "x"}}) + + assert _events(log_exporter, CHECKPOINTED) == [] + assert _events(log_exporter, RESUMED) == [] + # The interrupt still happens, so paused is still reported. LangGraph + # supplies a checkpoint id for the in-memory loop checkpoint even though + # nothing persisted it. + assert len(_events(log_exporter, PAUSED)) == 1 + + +def test_plain_graph_run_emits_no_lifecycle_events( + start_instrumentation, + log_exporter, +): + builder = StateGraph(_State) + builder.add_node("only", lambda _state: {"value": "done"}) + builder.add_edge(START, "only") + builder.add_edge("only", END) + graph = builder.compile() + + assert graph.invoke({"value": "new"}) == {"value": "done"} + + for event_name in LIFECYCLE_EVENTS: + assert _events(log_exporter, event_name) == [] + + +def test_async_interrupt_and_resume_emit_lifecycle_events( + start_instrumentation, + log_exporter, +): + graph = _build_graph().compile(checkpointer=InMemorySaver()) + config = {"configurable": {"thread_id": "async-lifecycle-test"}} + + async def run() -> str: + paused_result = await graph.ainvoke({"value": "new"}, config=config) + await graph.ainvoke(Command(resume="approved"), config=config) + return paused_result["__interrupt__"][0].id + + real_interrupt_id = asyncio.run(run()) + + paused = _events(log_exporter, PAUSED) + assert len(paused) == 1 + assert paused[0].attributes["gen_ai.agent.pause.id"] == real_interrupt_id + + # aput is wrapped alongside put, so the per-superstep volume matches the + # synchronous run. + assert len(_events(log_exporter, CHECKPOINTED)) == 5 + resumed = _events(log_exporter, RESUMED) + assert len(resumed) == 1 + assert ( + resumed[0].attributes["gen_ai.agent.resumed_from.id"] + == paused[0].attributes["gen_ai.agent.checkpoint.id"] + ) + + +def test_uninstrument_restores_the_checkpointer( + start_instrumentation, + log_exporter, +): + checkpointer = InMemorySaver() + graph = _build_graph().compile(checkpointer=checkpointer) + graph.invoke({"value": "new"}, config={"configurable": {"thread_id": "u"}}) + assert _events(log_exporter, CHECKPOINTED) + + start_instrumentation.uninstrument() + + assert "put" not in checkpointer.__dict__ + assert "aput" not in checkpointer.__dict__ + + +# --------------------------------------------------------------------------- +# Correlation and checkpointer wrapping +# --------------------------------------------------------------------------- + + +class _RecordingSaver(InMemorySaver): + """Saver that records every checkpoint it actually persists.""" + + def __init__(self) -> None: + super().__init__() + self.recorded: list[tuple[str, str]] = [] + + def put(self, config, checkpoint, metadata, new_versions): + result = super().put(config, checkpoint, metadata, new_versions) + configurable = result["configurable"] + self.recorded.append( + ( + str(config.get("configurable", {}).get("checkpoint_ns", "")), + str(configurable["checkpoint_id"]), + ) + ) + return result + + +def _checkpoint_ids(log_exporter) -> list[str]: + return [ + record.attributes["gen_ai.agent.checkpoint.id"] + for record in _events(log_exporter, CHECKPOINTED) + ] + + +def test_nested_subgraph_checkpoints_correlate_to_the_owning_workflow( + start_instrumentation, + log_exporter, + span_exporter, +): + class _SubState(TypedDict, total=False): + decision: str + + sub_builder = StateGraph(_SubState) + sub_builder.add_node( + "ask", lambda _state: {"decision": str(interrupt("approve?"))} + ) + sub_builder.add_edge(START, "ask") + sub_builder.add_edge("ask", END) + + builder = StateGraph(_State) + builder.add_node("prepare", lambda _state: {"value": "prepared"}) + builder.add_node("child", sub_builder.compile()) + builder.add_edge(START, "prepare") + builder.add_edge("prepare", "child") + builder.add_edge("child", END) + + saver = _RecordingSaver() + graph = builder.compile(checkpointer=saver) + config = {"configurable": {"thread_id": "nested-test"}} + + graph.invoke({"value": "new"}, config=config) + first_run_writes = list(saver.recorded) + graph.invoke(Command(resume="approved"), config=config) + + # Both the root namespace and the subgraph namespace are exercised. + namespaces = {namespace for namespace, _ in saver.recorded} + assert "" in namespaces + assert any(namespace.startswith("child:") for namespace in namespaces) + + # Every checkpoint LangGraph persisted is reported exactly once, at every + # nesting level, and none is invented. + assert _checkpoint_ids(log_exporter) == [ + checkpoint_id for _, checkpoint_id in saver.recorded + ] + + # LangGraph classifies a subgraph's chain run as a plain chain, not a + # workflow, so no child workflow span exists and the child namespace's + # checkpoints resolve to the nearest enclosing run: the parent workflow. + # Namespace resolution itself is covered in test_invocation_manager. + workflow_spans = _workflow_spans(span_exporter) + assert len(workflow_spans) == 2 + first_ids = {checkpoint_id for _, checkpoint_id in first_run_writes} + for record in _events(log_exporter, CHECKPOINTED): + expected = ( + workflow_spans[0] + if record.attributes["gen_ai.agent.checkpoint.id"] in first_ids + else workflow_spans[1] + ) + assert record.span_id == expected.context.span_id + assert record.trace_id == expected.context.trace_id + + +def test_two_unrelated_runs_on_one_thread_id_drop_the_checkpoint( + start_instrumentation, + log_exporter, +): + from unittest import mock + from uuid import uuid4 + + from opentelemetry.instrumentation.genai.langchain.callback_handler import ( + OpenTelemetryLangChainCallbackHandler, + ) + from opentelemetry.util.genai.invocation import WorkflowInvocation + + telemetry = mock.MagicMock() + telemetry.should_capture_content.return_value = False + telemetry.workflow.side_effect = lambda **_kwargs: mock.MagicMock( + spec=WorkflowInvocation + ) + handler = OpenTelemetryLangChainCallbackHandler(telemetry) + + metadata = {"ls_integration": "langgraph", "thread_id": "shared"} + workflows = [] + for _ in range(2): + run_id = uuid4() + handler.on_chain_start( + serialized={"name": "LangGraph"}, + inputs={}, + run_id=run_id, + metadata=metadata, + ) + workflows.append(handler._invocation_manager.get_invocation(run_id)) + + handler.checkpoint_written("shared", "ckpt-1") + + for workflow in workflows: + workflow.emit_event.assert_not_called() + + +class _DelegatingSaver(InMemorySaver): + """Saver whose ``aput`` delegates to ``put``, like ``InMemorySaver``.""" + + def __init__(self) -> None: + super().__init__() + self.write_count = 0 + + def put(self, config, checkpoint, metadata, new_versions): + self.write_count += 1 + return super().put(config, checkpoint, metadata, new_versions) + + async def aput(self, config, checkpoint, metadata, new_versions): + return self.put(config, checkpoint, metadata, new_versions) + + +class _WorkerThreadSaver(_DelegatingSaver): + """Saver whose ``aput`` runs ``put`` on a worker thread. + + The worker gets no copy of the caller's context, so de-duplication cannot + rely on context propagation. + """ + + async def aput(self, config, checkpoint, metadata, new_versions): + from concurrent.futures import ThreadPoolExecutor + + loop = asyncio.get_running_loop() + with ThreadPoolExecutor(max_workers=1) as pool: + return await loop.run_in_executor( + pool, + lambda: self.put(config, checkpoint, metadata, new_versions), + ) + + +@pytest.mark.parametrize( + "saver_factory", [_DelegatingSaver, _WorkerThreadSaver] +) +def test_delegating_saver_reports_each_write_once( + start_instrumentation, + log_exporter, + saver_factory, +): + saver = saver_factory() + graph = _build_graph().compile(checkpointer=saver) + config = {"configurable": {"thread_id": "delegating-test"}} + + asyncio.run(graph.ainvoke({"value": "new"}, config=config)) + + checkpoint_ids = _checkpoint_ids(log_exporter) + assert saver.write_count == 3 + assert len(checkpoint_ids) == saver.write_count + assert len(set(checkpoint_ids)) == saver.write_count + + +def test_uninstrument_restores_instance_level_put_and_aput( + start_instrumentation, + log_exporter, +): + saver = InMemorySaver() + original_put = saver.put + original_aput = saver.aput + # A saver that legitimately supplies its write methods per instance. + saver.put = original_put + saver.aput = original_aput + + graph = _build_graph().compile(checkpointer=saver) + graph.invoke({"value": "new"}, config={"configurable": {"thread_id": "i"}}) + assert _events(log_exporter, CHECKPOINTED) + + start_instrumentation.uninstrument() + + assert saver.__dict__["put"] is original_put + assert saver.__dict__["aput"] is original_aput + + +def test_untrackable_saver_is_left_untouched( + start_instrumentation, + log_exporter, +): + class _UnhashableSaver(InMemorySaver): + # Defining __eq__ without __hash__ makes instances unhashable, so the + # saver cannot be registered for restoration. + def __eq__(self, other: object) -> bool: + return self is other + + saver = _UnhashableSaver() + graph = _build_graph().compile(checkpointer=saver) + + # Nothing was patched, so nothing is left behind to clean up. + assert "put" not in saver.__dict__ + assert "aput" not in saver.__dict__ + + graph.invoke({"value": "new"}, config={"configurable": {"thread_id": "n"}}) + + assert _events(log_exporter, CHECKPOINTED) == [] + + +def test_checkpointer_is_discovered_from_the_compiled_graph( + start_instrumentation, + log_exporter, +): + keyword_saver = InMemorySaver() + keyword_graph = _build_graph().compile(checkpointer=keyword_saver) + positional_saver = InMemorySaver() + positional_graph = _build_graph().compile(positional_saver) + + # The saver the compiled graph retains is the one that was patched. + assert keyword_graph.checkpointer is keyword_saver + assert positional_graph.checkpointer is positional_saver + assert "put" in keyword_saver.__dict__ + assert "put" in positional_saver.__dict__ + + keyword_graph.invoke( + {"value": "new"}, config={"configurable": {"thread_id": "kw"}} + ) + keyword_checkpoints = len(_checkpoint_ids(log_exporter)) + positional_graph.invoke( + {"value": "new"}, config={"configurable": {"thread_id": "pos"}} + ) + + assert keyword_checkpoints == 3 + assert len(_checkpoint_ids(log_exporter)) == 6 + + +def test_nested_resume_reports_one_event_per_graph_level( + start_instrumentation, + log_exporter, +): + class _SubState(TypedDict, total=False): + decision: str + + sub_builder = StateGraph(_SubState) + sub_builder.add_node( + "ask", lambda _state: {"decision": str(interrupt("approve?"))} + ) + sub_builder.add_edge(START, "ask") + sub_builder.add_edge("ask", END) + + builder = StateGraph(_State) + builder.add_node("prepare", lambda _state: {"value": "prepared"}) + builder.add_node("child", sub_builder.compile()) + builder.add_edge(START, "prepare") + builder.add_edge("prepare", "child") + builder.add_edge("child", END) + + saver = _RecordingSaver() + graph = builder.compile(checkpointer=saver) + config = {"configurable": {"thread_id": "nested-resume"}} + + graph.invoke({"value": "new"}, config=config) + persisted = { + namespace: checkpoint_id for namespace, checkpoint_id in saver.recorded + } + graph.invoke(Command(resume="approved"), config=config) + + resumed_ids = [ + record.attributes["gen_ai.agent.resumed_from.id"] + for record in _events(log_exporter, RESUMED) + ] + # One resumed event per graph level: the root graph and the subgraph. + assert len(resumed_ids) == 2 + assert len(set(resumed_ids)) == 2 + + # Provenance: each id is the last checkpoint that level actually persisted. + root_namespace = "" + child_namespace = next( + namespace for namespace in persisted if namespace.startswith("child:") + ) + assert set(resumed_ids) == { + persisted[root_namespace], + persisted[child_namespace], + } + assert all( + record.attributes["gen_ai.agent.resumed_from.type"] == "checkpoint" + for record in _events(log_exporter, RESUMED) + ) + + +# --------------------------------------------------------------------------- +# Checkpoint write de-duplication +# --------------------------------------------------------------------------- + + +class _StubReporter: + def __init__(self) -> None: + self.calls: list[tuple[str, str, str]] = [] + self._lock = threading.Lock() + + def checkpoint_written( + self, thread_id: str, checkpoint_id: str, namespace: str + ) -> None: + with self._lock: + self.calls.append((thread_id, checkpoint_id, namespace)) + + +class _FixedIdSaver: + """Minimal saver whose writes always persist the same checkpoint id.""" + + def __init__(self, checkpoint_id: str = "fixed-ckpt") -> None: + self._checkpoint_id = checkpoint_id + self.write_count = 0 + + def put(self, config, checkpoint, metadata, new_versions): + self.write_count += 1 + return { + "configurable": { + "thread_id": config["configurable"]["thread_id"], + "checkpoint_ns": config["configurable"].get( + "checkpoint_ns", "" + ), + "checkpoint_id": self._checkpoint_id, + } + } + + async def aput(self, config, checkpoint, metadata, new_versions): + # Delegation, passing the same checkpoint object through. + return self.put(config, checkpoint, metadata, new_versions) + + +def _fixed_id_config(thread_id: str = "dedup"): + return {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + + +def test_delegated_write_is_reported_once_by_the_outer_call(): + saver = _FixedIdSaver() + reporter = _StubReporter() + _wrap_checkpointer(saver, reporter) + + asyncio.run(saver.aput(_fixed_id_config(), {"id": "a"}, {}, {})) + + assert saver.write_count == 1 + assert reporter.calls == [("dedup", "fixed-ckpt", "")] + + +def test_separate_writes_reusing_a_checkpoint_id_are_both_reported(): + saver = _FixedIdSaver() + reporter = _StubReporter() + _wrap_checkpointer(saver, reporter) + + saver.put(_fixed_id_config(), {"id": "a"}, {}, {}) + saver.put(_fixed_id_config(), {"id": "b"}, {}, {}) + + # Two distinct writes, so two events even though the id is identical. + assert saver.write_count == 2 + assert len(reporter.calls) == 2 + + +class _SequencedSaver: + """Saver that lets a test force an exact interleaving of two writes. + + ``aput`` delegates to ``put`` with the same checkpoint object, then holds + the outer call open until a second, independent write has entered. + """ + + def __init__(self) -> None: + self.write_count = 0 + self.delegated_returned = threading.Event() + self.second_write_entered = threading.Event() + self._lock = threading.Lock() + + def put(self, config, checkpoint, metadata, new_versions): + with self._lock: + self.write_count += 1 + if checkpoint["name"] == "B": + self.second_write_entered.set() + return { + "configurable": { + "thread_id": config["configurable"]["thread_id"], + "checkpoint_ns": "", + "checkpoint_id": f"ckpt-{checkpoint['name']}", + } + } + + async def aput(self, config, checkpoint, metadata, new_versions): + result = self.put(config, checkpoint, metadata, new_versions) + self.delegated_returned.set() + # Hold the outer write open until the independent write has entered. + assert self.second_write_entered.wait(10) + return result + + +def test_forced_interleaving_reports_each_write_exactly_once(): + saver = _SequencedSaver() + reporter = _StubReporter() + _wrap_checkpointer(saver, reporter) + config = _fixed_id_config("shared") + + def write_a() -> None: + asyncio.run(saver.aput(config, {"name": "A"}, {}, {})) + + def write_b() -> None: + # Enter only after A's delegated inner call has already returned. + assert saver.delegated_returned.wait(10) + saver.put(config, {"name": "B"}, {}, {}) + + threads = [ + threading.Thread(target=write_a), + threading.Thread(target=write_b), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=15) + assert not any(thread.is_alive() for thread in threads) + + # Sequence forced above: outer A enters, delegated A enters and returns, + # outer B enters, outer A returns, outer B returns. The delegation is + # silent and neither independent write is dropped as a false duplicate. + assert saver.write_count == 2 + assert sorted(checkpoint_id for _, checkpoint_id, _ in reporter.calls) == [ + "ckpt-A", + "ckpt-B", + ] + + +def test_concurrent_writes_through_a_real_delegating_saver( + start_instrumentation, + log_exporter, +): + saver = _DelegatingSaver() + graph = _build_graph().compile(checkpointer=saver) + started = threading.Barrier(4) + + def run(index: int) -> None: + started.wait() + graph.invoke( + {"value": "new"}, + config={"configurable": {"thread_id": f"concurrent-{index}"}}, + ) + + threads = [threading.Thread(target=run, args=(i,)) for i in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + checkpoint_ids = _checkpoint_ids(log_exporter) + assert saver.write_count == 12 + assert len(checkpoint_ids) == saver.write_count + assert len(set(checkpoint_ids)) == saver.write_count diff --git a/util/opentelemetry-util-genai/.changelog/536.added b/util/opentelemetry-util-genai/.changelog/536.added new file mode 100644 index 000000000..b7785d1f1 --- /dev/null +++ b/util/opentelemetry-util-genai/.changelog/536.added @@ -0,0 +1 @@ +Add an invocation event API that preserves the active span context, mirroring the identical change in PR #507 and dropped when rebasing onto it. diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_invocation.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_invocation.py index 388a14d73..735b3039d 100644 --- a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_invocation.py +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_invocation.py @@ -5,7 +5,7 @@ import timeit from abc import abstractmethod -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from contextlib import AbstractContextManager from contextvars import Token from dataclasses import asdict @@ -157,6 +157,25 @@ def record_stream_chunk(self) -> None: self._request_stream = True self._on_stream_chunk(timeit.default_timer()) + def emit_event( + self, + event_name: str, + attributes: Mapping[str, AttributeValue], + *, + body: str | None = None, + ) -> None: + """Emit an event correlated with this active invocation.""" + if self._context_token is None: + return + self._logger.emit( + LogRecord( + event_name=event_name, + body=body, + attributes=dict(attributes), + context=self._span_context, + ) + ) + def _on_stream_chunk(self, chunk_at: float) -> None: """Record streaming timing for one output chunk as it arrives. diff --git a/util/opentelemetry-util-genai/tests/fixtures/pr507_emit_event.py.txt b/util/opentelemetry-util-genai/tests/fixtures/pr507_emit_event.py.txt new file mode 100644 index 000000000..2b9ea9e7c --- /dev/null +++ b/util/opentelemetry-util-genai/tests/fixtures/pr507_emit_event.py.txt @@ -0,0 +1,20 @@ +from collections.abc import Mapping, Sequence + + def emit_event( + self, + event_name: str, + attributes: Mapping[str, AttributeValue], + *, + body: str | None = None, + ) -> None: + """Emit an event correlated with this active invocation.""" + if self._context_token is None: + return + self._logger.emit( + LogRecord( + event_name=event_name, + body=body, + attributes=dict(attributes), + context=self._span_context, + ) + ) diff --git a/util/opentelemetry-util-genai/tests/test_workflow_invocation.py b/util/opentelemetry-util-genai/tests/test_workflow_invocation.py index e7d7a4003..486634ea1 100644 --- a/util/opentelemetry-util-genai/tests/test_workflow_invocation.py +++ b/util/opentelemetry-util-genai/tests/test_workflow_invocation.py @@ -5,6 +5,11 @@ import unittest from unittest.mock import patch +from opentelemetry.sdk._logs import LoggerProvider +from opentelemetry.sdk._logs.export import ( + InMemoryLogRecordExporter, + SimpleLogRecordProcessor, +) from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( @@ -28,11 +33,19 @@ class TestWorkflowInvocation(unittest.TestCase): def setUp(self): self.span_exporter = InMemorySpanExporter() + self.log_exporter = InMemoryLogRecordExporter() self.tracer_provider = TracerProvider() self.tracer_provider.add_span_processor( SimpleSpanProcessor(self.span_exporter) ) - self.handler = TelemetryHandler(tracer_provider=self.tracer_provider) + logger_provider = LoggerProvider() + logger_provider.add_log_record_processor( + SimpleLogRecordProcessor(self.log_exporter) + ) + self.handler = TelemetryHandler( + tracer_provider=self.tracer_provider, + logger_provider=logger_provider, + ) def test_default_values(self): invocation = self.handler.workflow(name=None) @@ -110,6 +123,32 @@ def test_full_construction(self): assert len(invocation.output_messages) == 1 assert invocation.output_messages[0].parts[0].content == "answer" + def test_emit_event_uses_invocation_context(self): + invocation = self.handler.workflow(name="durable_workflow") + invocation.emit_event( + "gen_ai.agent.checkpointed", + {"gen_ai.agent.checkpoint.id": "ckpt-1"}, + body="Agent execution checkpointed", + ) + invocation.stop() + + records = self.log_exporter.get_finished_logs() + assert len(records) == 1 + record = records[0].log_record + assert record.event_name == "gen_ai.agent.checkpointed" + assert record.body == "Agent execution checkpointed" + assert record.attributes is not None + assert record.attributes["gen_ai.agent.checkpoint.id"] == "ckpt-1" + assert record.trace_id == invocation.span.get_span_context().trace_id + assert record.span_id == invocation.span.get_span_context().span_id + + def test_emit_event_after_stop_is_ignored(self): + invocation = self.handler.workflow(name="durable_workflow") + invocation.stop() + invocation.emit_event("test.event", {}) + + assert self.log_exporter.get_finished_logs() == () + def test_with_conversation_id(self): from opentelemetry.semconv._incubating.attributes import ( gen_ai_attributes as GenAI, @@ -146,3 +185,21 @@ def test_messages_omitted_from_span_in_event_only_mode(self): attrs = span.attributes or {} assert GenAI.GEN_AI_INPUT_MESSAGES not in attrs assert GenAI.GEN_AI_OUTPUT_MESSAGES not in attrs + + +class TestEmitEventMatchesPR507(unittest.TestCase): + """`emit_event` is a copy of PR #507's hunk, not a competing API.""" + + def test_local_emit_event_matches_the_stored_pr507_source(self): + from pathlib import Path + + tests_dir = Path(__file__).resolve().parent + stored = tests_dir / "fixtures" / "pr507_emit_event.py.txt" + if not stored.is_file(): + self.skipTest("PR #507 reference source is not packaged") + + source = ( + tests_dir.parent / "src/opentelemetry/util/genai/_invocation.py" + ).read_text() + for fragment in stored.read_text().split("\n\n", 1): + assert fragment.strip("\n") in source