`AGUIAdapter.streaming_response` doesn't preserve context vars

[hudyweas]

Description

Hello, I'm trying to add additional information to the pydantic_ai agent, but it's missing using this configuration

@router.post("/agui", name="gigenius_canvas_agui")
@inject
async def gigenius_canvas_agui(
    request: Request,
    history_manager: HistoryManager = Depends(AsyncProvide[Container.history_manager]),  # ty:ignore[invalid-type-form]
) -> Response:
    accept = request.headers.get("accept", SSE_CONTENT_TYPE)

    try:
        run_input = AGUIAdapter.build_run_input(await request.body())
    except ValidationError as e:
        return Response(
            content=json.dumps(e.json()),
            media_type="application/json",
            status_code=HTTPStatus.UNPROCESSABLE_ENTITY,
        )

    with set_baggage(project_id=run_input.thread_id):
        run_input.tools = []

        messages = history_manager.get_ag_ui_history(UUID(run_input.thread_id))

        run_input.messages = messages + [run_input.messages[-1]]

        history_manager.add_message(
            id=run_input.messages[-1].id,
            thread_id=UUID(run_input.thread_id),
            role="user",
            message=str(run_input.messages[-1].content),
        )

        adapter = AGUIAdapter(agent=gigenius_agent, run_input=run_input, accept=accept)

        original_stream = adapter.run_stream(deps=StateDeps(GigeniusDependencies()))

        async def stream_and_capture():
            accumulated_text = []

            async for event in original_stream:
                if hasattr(event, "type") and hasattr(event, "delta"):
                    if event.type == "TEXT_MESSAGE_CONTENT":
                        accumulated_text.append(event.delta)

                yield event

            full_response = "".join(accumulated_text)

            if full_response:
                history_manager.add_message(
                    id=uuid.uuid4(),
                    thread_id=UUID(run_input.thread_id),
                    role="assistant",
                    message=full_response,
                )

        return adapter.streaming_response(stream_and_capture())

Python, Logfire & OS Versions, related packages (not required)

logfire="4.22.0"
platform="macOS-26.2-arm64-arm-64bit-Mach-O"
python="3.13.2 (main, Mar 17 2025, 21:26:38) [Clang 20.1.0 ]"
[related_packages]
requests="2.32.5"
pydantic="2.12.5"
fastapi="0.128.2"
openai="2.17.0"
protobuf="6.33.5"
rich="14.3.2"
executing="2.2.1"
opentelemetry-api="1.39.1"
opentelemetry-exporter-otlp="1.39.1"
opentelemetry-exporter-otlp-proto-common="1.39.1"
opentelemetry-exporter-otlp-proto-grpc="1.39.1"
opentelemetry-exporter-otlp-proto-http="1.39.1"
opentelemetry-instrumentation="0.60b1"
opentelemetry-instrumentation-asgi="0.60b1"
opentelemetry-instrumentation-dbapi="0.60b1"
opentelemetry-instrumentation-fastapi="0.60b1"
opentelemetry-instrumentation-httpx="0.60b1"
opentelemetry-instrumentation-psycopg2="0.60b1"
opentelemetry-instrumentation-system-metrics="0.60b1"
opentelemetry-proto="1.39.1"
opentelemetry-sdk="1.39.1"
opentelemetry-semantic-conventions="0.60b1"
opentelemetry-util-http="0.60b1"

[alexmojaki]

How's ty?

I think that the additional columns UI is still buggy. Start by checking if this attribute just exists when you look at a span directly in the details panel or the raw data.

[alexmojaki]

Also please reduce this to a minimal reproducible example

[hudyweas]

Ty is pretty good now. Not perfect, but the speed compared to mypy makes up for it.

There is no info about project_id in the raw data.

import json
from http import HTTPStatus

import logfire
from fastapi import FastAPI, Request, Response
from logfire import set_baggage
from pydantic import ValidationError
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.azure import AzureProvider
from pydantic_ai.ui import SSE_CONTENT_TYPE
from pydantic_ai.ui.ag_ui import AGUIAdapter

app = FastAPI()

agent = Agent(
    OpenAIChatModel(
        "gpt-5.1",
        provider=AzureProvider(),
    ),
    instructions="Be nice",
)

logfire.configure(service_name="issue_repro")
logfire.instrument_pydantic_ai()


@app.post("/chat")
async def chat(
    request: Request,
) -> Response:
    accept = request.headers.get("accept", SSE_CONTENT_TYPE)

    try:
        run_input = AGUIAdapter.build_run_input(await request.body())
    except ValidationError as e:
        return Response(
            content=json.dumps(e.json()),
            media_type="application/json",
            status_code=HTTPStatus.UNPROCESSABLE_ENTITY,
        )

    with set_baggage(project_id=run_input.thread_id):
        adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
        return adapter.streaming_response(adapter.run_stream())

[alexmojaki]

Is it possible to reproduce this without FastAPI?

[hudyweas]

Well, yeah. Same issue.

import logfire
from dotenv import load_dotenv
from logfire import set_baggage
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.azure import AzureProvider

load_dotenv()

agent = Agent(
    OpenAIChatModel(
        "gpt-5.1",
        provider=AzureProvider(),
    ),
    instructions="Be nice",
)

logfire.configure(service_name="issue_repro")
logfire.instrument_pydantic_ai()


with set_baggage(project_id="abc"):
    agent.run_sync("Hello")

[alexmojaki]

I reduced your code further, especially to get rid of azure, although I don't think it matters. project_id is definitely in the attributes:

import logfire
from logfire import set_baggage
from pydantic_ai import Agent

agent = Agent("gpt-5.1")

logfire.configure()
logfire.instrument_pydantic_ai()

with set_baggage(project_id="abc"):
    agent.run_sync("Hello")

[hudyweas]

You are right, it was my mistake. I missed project_id in the bare example. I found the problem tho. It is related to the FastAPI or AGUI adapter.

import logfire
from ag_ui.core import RunAgentInput
from dotenv import load_dotenv
from fastapi import FastAPI
from logfire import set_baggage
from pydantic_ai import Agent
from pydantic_ai.ui.ag_ui import AGUIAdapter

load_dotenv()

agent = Agent("gpt-5.1")

logfire.configure(service_name="issue_repro")
logfire.instrument_pydantic_ai()

app = FastAPI()

@app.post("/run_agent")
async def run_agent():
    with set_baggage(project_id="test"):
        adapter = AGUIAdapter(
            agent=agent,
            run_input=RunAgentInput(
                thread_id="test",
                run_id="test",
                state={},
                messages=[{"role": "user", "content": "hi", "id": "test"}],
                tools=[],
                context=[],
                forwardedProps={},
            ),
        )
        async for _ in adapter.run_stream():
            pass
        # return adapter.streaming_response(adapter.run_stream())

When run_stream is run like that I have project_id

        async for _ in adapter.run_stream():
            pass

When run in streaming response, project_id is lost:

 return adapter.streaming_response(adapter.run_stream())

[alexmojaki]

@hudyweas how about this?

async def run_agent():
    adapter = AGUIAdapter(...)
    async def stream():
        with set_baggage(project_id="test"):
            async for event in adapter.run_stream():
                yield event
    return adapter.streaming_response(stream())

@Kludex FastAPI expert, is there something to be improved here? Should streaming_response or StreamingResponse preserve context vars?

[hudyweas]

@alexmojaki, it works now. Thanks for this workaround. Do you have any explanation for what happened inside that context was lost?

[Kludex]

Probably some issue with the fact that it's not async, but it should preserve in async.

[DouweM]

• This smells related to How to properly span async generators/streams that are consumed outside the span context? #4173

@hudyweas Have a look at https://ai.pydantic.dev/agent/#run-metadata which is handled directly by Pydantic AI instead of using the baggage context var

[alexmojaki]

No they're not related. The span ending when it ends in the other issue is normal and not related to context vars.