Make RunContext.tracer, trace_include_content and instrumentation_version follow an explicit Instrumentation capability (#8706)

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -91,6 +91,7 @@ from ..capabilities.abstract import (
 from ..capabilities.combined import bind_capabilities_tier
 from ..capabilities.hooks import EventT, Hooks, OnEventHookFunc
 from ..capabilities.instrumentation import Instrumentation as InstrumentationCap
+from ..capabilities.wrapper import WrapperCapability
 from ..models.instrumented import InstrumentationSettings, InstrumentedModel
 from ..native_tools import AbstractNativeTool
 from ..native_tools._tool_search import ToolSearchTool
@@ -1666,12 +1667,9 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         else:
             instrumentation_settings = self._resolve_instrumentation_settings()
 
-        if instrumentation_settings is not None:
-            tracer = instrumentation_settings.tracer
-            instrumentation_cap: InstrumentationCap | None = InstrumentationCap(settings=instrumentation_settings)
-        else:
-            tracer = NoOpTracer()
-            instrumentation_cap = None
+        instrumentation_cap = (
+            InstrumentationCap(settings=instrumentation_settings) if instrumentation_settings is not None else None
+        )
 
         # Allocated here rather than with the graph deps below, so the context `for_run` receives
         # shares the very mappings the run fills at setup. A capability that holds on to that
@@ -1694,11 +1692,6 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             usage_limits=usage_limits,
             prompt=user_prompt,
             messages=state.message_history,
-            tracer=tracer,
-            trace_include_content=instrumentation_settings is not None and instrumentation_settings.include_content,
-            instrumentation_version=instrumentation_settings.version
-            if instrumentation_settings
-            else DEFAULT_INSTRUMENTATION_VERSION,
             run_step=0,
             pending_messages=state.pending_messages,
             run_id=state.run_id,
@@ -1706,6 +1699,14 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             _cancellation=cancellation,
         )
 
+        # An explicit `Instrumentation` capability (agent- or call-level) replaces the one injected from
+        # `instrumentation_settings` (see `_resolve_run_capabilities`), so `for_run` hooks and metadata
+        # factories are shown the settings of the one that will actually instrument the run.
+        _set_run_context_instrumentation(
+            initial_ctx,
+            _run_instrumentation_settings([base_capability, *extra_capabilities], default=instrumentation_settings),
+        )
+
         # Resolve run metadata up front so capability and toolset `for_run` hooks
         # can see it on `RunContext.metadata`. Metadata factories receive the
         # `initial_ctx` above (no `tool_manager` / `validation_context` yet); they
@@ -1731,6 +1732,10 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             base_is_override=base_is_override,
         )
         run_capability = resolved_caps.run_capability
+        # Read back off the resolved tree, which also sees an `Instrumentation` that only a `for_run`
+        # contributed, so every later step's `RunContext` describes what the run's spans use.
+        instrumentation_settings = _run_instrumentation_settings([run_capability])
+        _set_run_context_instrumentation(initial_ctx, instrumentation_settings)
         capabilities_dict = resolved_caps.capabilities
         cap_instructions = resolved_caps.instructions
         cap_native_tools = resolved_caps.native_tools
@@ -1891,7 +1896,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             native_tools=cap_native_tools,
             tool_manager=tool_manager,
             display_banner=display_banner,
-            tracer=tracer,
+            tracer=initial_ctx.tracer,
             get_instructions=get_instructions,
             instrumentation_settings=instrumentation_settings,
             cancellation=cancellation,
@@ -3473,40 +3478,16 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
 
         # The session-level `realtime` span and per-response `chat` spans are hand-managed by
         # `RealtimeSession`; run-lifecycle hooks have no per-exchange boundary on which to build them.
-        # Drive them from the settings that
-        # will actually win: an explicit `Instrumentation` capability's (agent- or call-level) over the
+        # Drive them, and the `RunContext` (as `iter` does), from the settings that will actually win:
+        # an explicit `Instrumentation` capability's (agent- or call-level) over the
         # `instrument=`-derived ones, matching the precedence `_resolve_run_capabilities` applies to the
-        # tool spans.
-        # The winner is *selected* here, not combined: `combine` settles duplicates *within* one
-        # layer and never runs across the agent-to-run boundary, so the session reads the
-        # configuration off the instance `_resolve_run_capabilities` would keep -- the last
-        # explicit `Instrumentation` in application order (a call-level one supersedes the
-        # agent-level one, and within a layer the guarded combine resolves to the last instance's
-        # values). Taking the first match would drive the session spans from settings the
-        # effective configuration had already turned off. Id guards apply once, where the layers
-        # actually combine.
-        instrumentation_layers = [self._effective_root_capability(), *extra_capabilities]
-        explicit_instrumentations = [
-            leaf
-            for layer in instrumentation_layers
-            for leaf in leaf_capabilities(layer)
-            if isinstance(leaf, InstrumentationCap)
-        ]
-        explicit_instrumentation = explicit_instrumentations[-1] if explicit_instrumentations else None
-        session_instrumentation_settings = (
-            explicit_instrumentation.settings if explicit_instrumentation is not None else instrumentation_settings
+        # tool spans. Resolve metadata after, and before `for_run`, so capability/toolset hooks see it
+        # (same ordering as the graph run).
+        base_capability, base_is_override = self._base_run_capability()
+        _set_run_context_instrumentation(
+            run_context,
+            _run_instrumentation_settings([base_capability, *extra_capabilities], default=instrumentation_settings),
         )
-        # Mirror `iter`'s `RunContext`: expose the resolved tracer (a `NoOpTracer` when uninstrumented)
-        # and content-tracing flag, and resolve metadata before `for_run` so capability/toolset hooks see
-        # it (same ordering as the graph run).
-        run_context.tracer = (
-            session_instrumentation_settings.tracer if session_instrumentation_settings is not None else NoOpTracer()
-        )
-        run_context.trace_include_content = (
-            session_instrumentation_settings is not None and session_instrumentation_settings.include_content
-        )
-        if session_instrumentation_settings is not None:
-            run_context.instrumentation_version = session_instrumentation_settings.version
         run_context.metadata = self._get_metadata(run_context, metadata)
 
         # Resolve the capability layers and extract their contributions, exactly as `run`/`iter` do via
@@ -3514,7 +3495,6 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         # Realtime keeps its own surroundings: no `InstrumentedModel` unwrap, once-only model settings
         # (below), and the `_keep_native` drop plus the shared native ↔ local-tool swap (below). Keep
         # this in sync with the `iter` call site.
-        base_capability, base_is_override = self._base_run_capability()
         resolved_caps = await self._resolve_run_capabilities(
             run_context,
             base_capability=base_capability,
@@ -3524,6 +3504,10 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             base_is_override=base_is_override,
         )
         run_capability = resolved_caps.run_capability
+        # Read back off the resolved tree, as `iter` does, so an `Instrumentation` only a `for_run`
+        # contributed drives the session's spans and context too.
+        session_instrumentation_settings = _run_instrumentation_settings([run_capability])
+        _set_run_context_instrumentation(run_context, session_instrumentation_settings)
         # Deferred capabilities load in a session the same way they do in a graph run: the catalog
         # renders into the connect-time instructions and the loaded instructions come back as the
         # `load_capability` tool's own result, which every provider supports. What no provider
@@ -4381,6 +4365,39 @@ _AUTO_INJECT_CAPABILITY_TYPES: tuple[type[AbstractCapability[Any]], ...] = (
 """Infrastructure capabilities auto-injected when not already present."""
 
 
+def _run_instrumentation_settings(
+    capabilities: Sequence[AbstractCapability[Any]], *, default: InstrumentationSettings | None = None
+) -> InstrumentationSettings | None:
+    """The settings of the `Instrumentation` capability that instruments a run with these capabilities.
+
+    That's the last one in application order, or `default` when there is none. `_resolve_run_capabilities`
+    selects rather than combines across the agent-to-run boundary (`combine` only settles duplicates
+    *within* a layer), so a call-level one supersedes the agent-level one, and within a layer the guarded
+    combine keeps the last instance's values. Taking the first match would describe settings the effective
+    configuration had already turned off.
+
+    Before resolution, `default` is the `instrument=`-derived settings the run injects only when no explicit
+    capability is present. After resolution, the injected capability is part of the tree, so no default applies.
+    A wrapper around a single capability is itself a leaf, so it's unwrapped: that's the shape a capability
+    function's `Instrumentation` takes once resolved, and the wrapper delegates its hooks to it.
+    """
+    instrumentations: list[InstrumentationCap] = []
+    for capability in capabilities:
+        for leaf in leaf_capabilities(capability):
+            while isinstance(leaf, WrapperCapability):
+                leaf = leaf.wrapped
+            if isinstance(leaf, InstrumentationCap):
+                instrumentations.append(leaf)
+    return instrumentations[-1].settings if instrumentations else default
+
+
+def _set_run_context_instrumentation(ctx: RunContext[Any], settings: InstrumentationSettings | None) -> None:
+    """Point `ctx`'s tracer, content flag and instrumentation version at `settings`; uninstrumented if `None`."""
+    ctx.tracer = settings.tracer if settings is not None else NoOpTracer()
+    ctx.trace_include_content = settings is not None and settings.include_content
+    ctx.instrumentation_version = settings.version if settings is not None else DEFAULT_INSTRUMENTATION_VERSION
+
+
 def _registered_capability_count(capability: AbstractCapability[Any]) -> int:
     """Count the capabilities the user registered, for a `pydantic_ai._display` banner.
 

tests/test_display.py

@@ -18,7 +18,7 @@ import pydantic_ai
 import pydantic_ai._display as _display
 from pydantic_ai import Agent, ModelMessage, ModelRequest, UserPromptPart, __version__
 from pydantic_ai.agent import _registered_capability_count  # pyright: ignore[reportPrivateUsage]
-from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.capabilities import AbstractCapability, Instrumentation
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.toolsets import FunctionToolset
 
@@ -916,6 +916,15 @@ def test_instrumented_agent_run_is_silent(monkeypatch: pytest.MonkeyPatch, stder
     assert stderr.getvalue() == ''
 
 
+def test_run_instrumented_by_an_explicit_capability_is_silent(monkeypatch: pytest.MonkeyPatch, stderr: TTYStream):
+    """The banner defers to the `Instrumentation` capability that instruments the run, not just `instrument`."""
+    monkeypatch.setattr(sys, 'stderr', stderr)
+
+    Agent(TestModel()).run_sync('hello', capabilities=[Instrumentation()])
+
+    assert stderr.getvalue() == ''
+
+
 def test_an_instrumented_run_leaves_the_banner_for_another_agent(monkeypatch: pytest.MonkeyPatch, stderr: TTYStream):
     """Instrumentation is the agent's, not the process's, so it doesn't speak for the ones after it."""
     monkeypatch.setattr(sys, 'stderr', stderr)

tests/test_logfire.py

@@ -55,6 +55,8 @@ with try_import():
     # `opentelemetry-sdk` arrives with the `logfire` extra, so it is not importable in the
     # `pydantic-ai-slim` / `pydantic-evals` install groups either.
     from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+    from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
     from opentelemetry.sdk.trace.sampling import ALWAYS_OFF
 
 try:
@@ -3916,6 +3918,105 @@ def test_instrument_all_skipped_when_capability_already_present(
         Agent.instrument_all(False)
 
 
+@dataclass
+class _SeenInstrumentation:
+    tracer: Any
+    include_content: bool
+    version: int
+
+
+def _seen(ctx: RunContext[Any]) -> _SeenInstrumentation:
+    return _SeenInstrumentation(ctx.tracer, ctx.trace_include_content, ctx.instrumentation_version)
+
+
+@dataclass
+class _ObserveForRun(AbstractCapability[Any]):
+    """Records what a `for_run` hook sees of the run's instrumentation."""
+
+    seen: list[_SeenInstrumentation]
+
+    async def for_run(self, ctx: RunContext[Any]) -> _ObserveForRun:
+        self.seen.append(_seen(ctx))
+        return self
+
+
+@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
+@pytest.mark.parametrize('explicit_on', [None, 'agent', 'run'])
+def test_run_context_instrumentation_follows_the_capability_that_instruments_the_run(
+    explicit_on: Literal['agent', 'run'] | None,
+) -> None:
+    """An explicit `Instrumentation` capability replaces the `instrument_all()` one, so `RunContext` describes it.
+
+    Its `tracer`, `trace_include_content` and `instrumentation_version` are what the run's spans use, both
+    in `for_run` hooks and in the context a tool gets later in the run. Without one, `instrument_all()`'s
+    settings still apply.
+    """
+    implicit = InstrumentationSettings(include_content=True, tracer_provider=TracerProvider())
+    explicit = InstrumentationSettings(include_content=False, version=6, tracer_provider=TracerProvider())
+    expected = implicit if explicit_on is None else explicit
+
+    for_run_seen: list[_SeenInstrumentation] = []
+    agent_capabilities: list[AbstractCapability[Any]] = [_ObserveForRun(for_run_seen)]
+    run_capabilities: list[AbstractCapability[Any]] = []
+    if explicit_on == 'agent':
+        agent_capabilities.append(Instrumentation(settings=explicit))
+    elif explicit_on == 'run':
+        run_capabilities.append(Instrumentation(settings=explicit))
+
+    Agent.instrument_all(implicit)
+    try:
+        agent = Agent(TestModel(), capabilities=agent_capabilities)
+        tool_seen: list[_SeenInstrumentation] = []
+
+        @agent.tool
+        def peek(ctx: RunContext[Any]) -> str:
+            tool_seen.append(_seen(ctx))
+            return 'ok'
+
+        agent.run_sync('hi', capabilities=run_capabilities)
+    finally:
+        Agent.instrument_all(False)
+
+    assert implicit.tracer is not explicit.tracer
+    expected_seen = _SeenInstrumentation(expected.tracer, expected.include_content, expected.version)
+    assert for_run_seen == [expected_seen]
+    assert tool_seen == [expected_seen]
+
+
+@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
+def test_run_context_instrumentation_follows_a_capability_function_instrumentation() -> None:
+    """An `Instrumentation` only a capability function's `for_run` returns still instruments the run, so the
+    context the run's tools get describes it too, even though it couldn't be known before `for_run`."""
+    implicit_provider = TracerProvider()
+    implicit_exporter = InMemorySpanExporter()
+    implicit_provider.add_span_processor(SimpleSpanProcessor(implicit_exporter))
+    dynamic_provider = TracerProvider()
+    dynamic_exporter = InMemorySpanExporter()
+    dynamic_provider.add_span_processor(SimpleSpanProcessor(dynamic_exporter))
+    dynamic = InstrumentationSettings(include_content=False, version=6, tracer_provider=dynamic_provider)
+
+    Agent.instrument_all(InstrumentationSettings(include_content=True, tracer_provider=implicit_provider))
+    try:
+        agent = Agent(TestModel())
+        tool_seen: list[_SeenInstrumentation] = []
+
+        @agent.tool
+        def peek(ctx: RunContext[Any]) -> str:
+            tool_seen.append(_seen(ctx))
+            return 'ok'
+
+        agent.run_sync('hi', capabilities=[lambda ctx: Instrumentation(settings=dynamic)])
+    finally:
+        Agent.instrument_all(False)
+
+    assert tool_seen == [_SeenInstrumentation(dynamic.tracer, False, 6)]
+    # The capability function's `Instrumentation` is the one that opened the run's spans.
+    assert implicit_exporter.get_finished_spans() == ()
+    assert [span.name for span in dynamic_exporter.get_finished_spans()] == snapshot(
+        ['chat test', 'execute_tool peek', 'chat test', 'invoke_agent agent']
+    )
+
+
 @pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
 def test_agent_with_user_provided_instrumented_model(
     get_logfire_summary: Callable[[], LogfireSummary],