Skip to content

ref(pydantic-ai): Consolidate run-scoped state into one context manager - #7195

Draft
ericapisani wants to merge 1 commit into
ep/pydantic-ai-refactor-2-extractfrom
ep/pydantic-ai-refactor-3-run-context
Draft

ref(pydantic-ai): Consolidate run-scoped state into one context manager#7195
ericapisani wants to merge 1 commit into
ep/pydantic-ai-refactor-2-extractfrom
ep/pydantic-ai-refactor-3-run-context

Conversation

@ericapisani

Copy link
Copy Markdown
Member

Description

Issues

Reminders


Stack created with GitHub Stacks CLIGive Feedback 💬

Replace the hand-rolled push/pop contextvar stack with _run_context.py: one AgentRun dataclass on a token-reset contextvar, managed by a single agent_run_scope() context manager so pairing no longer relies on manual try/finally in three places. The streaming wrapper's manual scope/span/stack bookkeeping collapses into one ExitStack, which also stops leaking the isolation scope and invoke_agent span when entering the wrapped context manager fails, and the token reset tolerates streaming runs that exit in a different asyncio task.

The request hooks deliberately keep pairing chat spans through RunContext.metadata: contextvar pairing can mispair overlapping run_stream context managers in the same task, while the metadata dict is correct per run by construction. Review-driven hardening folded in: hook bodies are wrapped in capture_internal_exceptions so span-data failures cannot abort the user's model request, invoke_agent spans now receive model data from the agent argument instead of the not-yet-pushed contextvar, the model-settings mapping has a single source of truth in _extract.py, and the usage extractor passes record_token_usage kwargs directly.
@github-actions

Copy link
Copy Markdown
Contributor

Codecov Results 📊

112337 passed | ❌ 1 failed | ⏭️ 6755 skipped | Total: 119093 | Pass Rate: 94.33% | Execution Time: 412m 58s

📊 Comparison with Base Branch

Metric Change
Total Tests 📉 -129
Passed Tests 📉 -126
Failed Tests 📈 +1
Skipped Tests 📉 -4

➕ New Tests (1)

View new tests
  • test_span_streaming_quart_auth_user_id[42-True]
    • File: tests.integrations.quart.test_quart
    • Status: ❌ Failing

❌ Failed Tests

test_span_streaming_quart_auth_user_id[42-True]

File: tests.integrations.quart.test_quart
Suite: py3.11-quart-v0.21.0
Error: KeyError: 'user.id'

Stack Trace
tests/integrations/quart/test_quart.py:1117: in test_span_streaming_quart_auth_user_id
    assert segment["attributes"]["user.id"] == user_id
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   KeyError: 'user.id'

✅ Patch coverage is 96.81%. Project has 2481 uncovered lines.
✅ Project coverage is 90.23%. Comparing base (base) to head (head).

Files with missing lines (3)
File Patch % Lines
sentry_sdk/integrations/pydantic_ai/patches/agent_run.py 100.00% ⚠️ 3 partials
sentry_sdk/integrations/pydantic_ai/init.py 85.71% ⚠️ 1 Missing and 1 partials
sentry_sdk/integrations/pydantic_ai/_run_context.py 92.59% ⚠️ 2 Missing
Coverage diff
@@            Coverage Diff             @@
##          main       #PR       +/-##
==========================================
+ Coverage    90.23%    90.23%        —%
==========================================
  Files          193       195        +2
  Lines        25361     25382       +21
  Branches      9302      9278       -24
==========================================
+ Hits         22883     22901       +18
- Misses        2478      2481        +3
- Partials      1432      1422       -10

Generated by Codecov Action

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant