MCP sampling calls are not nested under agent run span

[JohnCHarrington]

Initial Checks

• I confirm that I'm using the latest version of Pydantic AI
• I confirm that I searched for my issue in https://github.com/pydantic/pydantic-ai/issues before opening this issue

Description

Run the MCP sampling example from the PydanticAI docs (https://ai.pydantic.dev/mcp/client/#mcp-sampling) with OpenTelemetry instrumentation.

I would expect the sampling LLM call to appear in the trace of my agent run, ideally nested under the tool call that made that sampling request.

Instead, the sampling LLM call appears as it's own separate trace, with no way to link it to the tool or the trace it was part of. Here the result is shown in langfuse, where agent run is the main trace, and chat acd-openai... is the sampling call:

Example Code

Simply instrument the example in the docs with opentelemetry.

Python, Pydantic AI & LLM client version

Python 3.11
Pydantic AI 1.1.0
Azure OpenAI (latest)

[mtblanton]

I was about to open a similar issue! I'm using direct model calls in a Flask application instrumented with OpenTelemetry. The chat trace should be nested in the API endpoint trace, but it's separate.

[JohnCHarrington]

I was about to open a similar issue! I'm using direct model calls in a Flask application instrumented with OpenTelemetry. The chat trace should be nested in the API endpoint trace, but it's separate.

Thanks, I have updated the title to reflect that this is a more general issue with direct model calls and not unique to the MCP Sampling (which just uses a direct model call from what I can see)

[DouweM]

@JohnCHarrington Does it work as expected if you drop async with agent: and just have the agent lazily open the MCP server connection when it needs it? That async with agent opens the connection at a higher context level than the agent.run call, and MCP sampling calls are handled at the connection/session level, so OpenTelemetry wouldn't consider the sampling call as being "under" the agent run.

@mtblanton I think that's a separate issue, as it's not about MCP sampling. The OpenTelemetry SDK is supposed to automatically set the parent trace ID for nested spans, which works as expected for FastAPI at least. In this case, it doesn't seem to realize that the chat call is made inside the POST request handler, which is odd. Do you have some example code that reproduces the issue, including how you're instrumenting it?

[mtblanton]

@mtblanton I think that's a separate issue, as it's not about MCP sampling. The OpenTelemetry SDK is supposed to automatically set the parent trace ID for nested spans, which works as expected for FastAPI at least. In this case, it doesn't seem to realize that the chat call is made inside the POST request handler, which is odd. Do you have some example code that reproduces the issue, including how you're instrumenting it?

@DouweM Ah, gotcha! I'll file a separate ticket with some example code when I get the time. We're using it synchronously with Flask which I think may be part of the problem?

[JohnCHarrington]

@DouweM exactly the same result without async with agent

[DouweM]

@JohnCHarrington OK, thanks for trying that out regardless.

With some small changes that I'll create a PR for shortly, I was able to go from this:

To this:

So now at least all the MCP spans are properly nested under agent run (as long as you don't call async with agent: yourself, because that determines the context MCP spans will be under). Ideally, MCP client handle request: sampling/createMessage would be under MCP request: tools/call use_sampling, but unfortunately I wasn't able to (in the 20 min I spent on this) find a way to pass the traceparent ID around through MCP.

Example code:

import logfire
from pydantic_ai import Agent
from pydantic_ai.mcp import MCPServerStdio
from pydantic_ai.models.instrumented import instrument_model
from pydantic_ai.models.test import TestModel

logfire.configure()
logfire.instrument_pydantic_ai()
logfire.instrument_mcp()

server = MCPServerStdio("python", ["-m", "tests.mcp_server"])
server.sampling_model = instrument_model(
    TestModel(custom_output_text="sampling model response"), True
)

agent = Agent(model=TestModel(call_tools=["use_sampling"]), toolsets=[server])
result = agent.run_sync("Hello")
print(result.output)