Repository navigation
fix(temporal): keep pydantic-ai subagent spans nested and replay-safe - #1054
Conversation
> 🤖 **This text was AI-generated.**
With `BraintrustPlugin` and Pydantic AI's Temporal durability integration, a
tool such as `delegate_task` runs in workflow code and starts a sub-agent
child workflow. Two things went wrong:
1. The child workflow was parented to the workflow span, so it showed up as
a sibling of the `delegate_task` span instead of under it. The workflow
outbound interceptor always sent the workflow span as the parent of
activities, local activities, and child workflows, ignoring the current
in-workflow span.
2. Temporal re-runs workflow code from the start whenever a workflow
replays, for example after a worker restart or when the workflow is
evicted from the worker cache. Spans opened in workflow code, such as
Pydantic AI's `agent_run`, `delegate_task`, and `chat` spans or users'
own spans, were logged again on every replay with new random IDs. The
plugin also skipped the workflow span during replay, so the copies became
separate root traces.
```
before after
temporal.workflow.Parent temporal.workflow.Parent
├─ agent_run [parent_agent] └─ agent_run [parent_agent]
│ └─ delegate_task ├─ chat test
├─ temporal.workflow.SubAgent <- sibling ├─ delegate_task
agent_run [parent_agent] <- replay copy, │ └─ temporal.workflow.SubAgent
└─ ... new root │ └─ agent_run [sub_agent]
└─ chat test
```
The fix:
- Parenting: the workflow outbound interceptor now uses the current span as
the parent when it is a real span other than the workflow span, and falls
back to the workflow span otherwise.
- Core: `logger.py` gets an internal, context-local span hook. When a hook
is set, `SpanImpl` takes default row, span, and trace IDs from it, and
`log_internal` drops writes the hook rejects. Spans only look up the hook
after one has been installed, so code that never uses Temporal pays for
one boolean check per span.
- Replay safety: each workflow run installs a hook that numbers spans per
run and derives their IDs from the run ID and that number. Workflow code
is deterministic, so a replayed span gets the ID it was first logged with.
The hook drops writes while replaying. The plugin now always creates the
workflow span, so it is the current span during replay too.
This also fixes the workflow span's end time. Evicting a workflow unwinds
its coroutine, which ended the span at the first eviction. The span now ends
when the workflow really finishes, because the later end row merges into the
same span.
`test_temporal(latest)` now also installs `pydantic-ai-slim[temporal]`,
pinned to the `pydantic-ai-integration` latest version. Pydantic AI's
temporal extra needs temporalio>=1.27 and Python 3.11+, so older temporalio
entries skip it. The new test covers both the deprecated `TemporalAgent`
wrapper and the newer `TemporalDurability` capability
(`metadata={"temporal": False}` on the delegating tool), with and without
the workflow cache. On the capability path, Pydantic AI routes the model
request to the activity before our wrapper runs, so the `chat` span sits
under the model activity instead of around it. Each request is still traced
once.
Limitations: replay safety only applies to workflows that take the existing
`braintrust-temporal-workflow-span-context-v1` patch branch. Spans opened in
signal or update handlers do not get deterministic IDs yet, because those
handlers run in their own context.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
…able
> 🤖 **This text was AI-generated.**
On Python 3.10 the `pylint` session does not install pydantic-ai, because
its latest release needs Python 3.11+. The session already drops the
`pydantic_ai` integration's own files from the lint set in that case. The
new `temporal/test_temporal_pydantic_ai.py` imports pydantic-ai from outside
that directory, so pylint still linted it and failed `static_checks (3.10)`:
```
src/braintrust/integrations/temporal/test_temporal_pydantic_ai.py:13:0:
E0401: Unable to import 'pydantic_ai.durable_exec.temporal' (import-error)
```
The session now has a small `cross_integration_files` map that ties files
outside an integration's directory to that integration. They are skipped
together whenever the integration cannot be linted:
```python
cross_integration_files = {
"pydantic_ai": {f"src/{INTEGRATION_DIR}/temporal/test_temporal_pydantic_ai.py"},
}
```
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ry thread > 🤖 **This text was AI-generated.** `BraintrustPlugin(logger=...)` sends writes to its logger through `_override_bg_logger`, which is thread-local. The plugin set it only on the thread that started the workflow. Temporal can run later activations of the same workflow on other worker threads, so writes made there skipped the plugin's logger and went to the global logger instead: ``` activation 1 (thread A, override set) -> span start row -> plugin logger activation 2 (thread B, no override) -> span end row -> global logger ``` This split spans across loggers and left rows with no start, which showed up in CI as `KeyError: 'span_attributes'` in the Pydantic AI + Temporal test. It only affects users who pass `logger=`. With the default, every thread uses the global logger anyway. The workflow span hook now re-applies the plugin's logger on the current thread before every write it allows, using the same `__braintrust_get_logger` call the workflow makes at startup. The activity path already did this per activity. A new test case runs workflow activations on a fresh thread every time, so the bug reproduces every time instead of by chance. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7acf6e5032
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "Codex (@codex) review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "Codex (@codex) address that feedback".
| self._counter += 1 | ||
| ids = _deterministic_span_ids(f"{self._base}:span:{self._counter}", self._legacy_ids) |
There was a problem hiding this comment.
Exclude query-only spans from the replay counter
When a workflow query handler opens a Braintrust span, this mutable counter advances even though Temporal queries are not recorded in workflow history and therefore are not replayed. After a cache eviction or worker restart, subsequent history-driven spans receive different IDs; they can be duplicated or, worse, reuse the earlier query span's deterministic ID and merge unrelated records. Query executions need to be excluded from this sequence or assigned IDs without affecting replayed workflow spans.
Useful? React with 👍 / 👎.
> 🤖 **This text was AI-generated.** Cache eviction could lose the original Braintrust span for a long-running Pydantic AI chat. Before the fix in PR #1054, replay created multiple workflow spans for the same chat: ```text first message -> cache eviction -> replay -> duplicate workflow spans ``` Add a two-turn chat regression for both Temporal durability APIs. Force cache eviction, stop the worker while the chat waits, and resume it on a fresh worker: ```text workflow span └─ chat_session (same IDs and start time after restart) ├─ first agent run └─ second agent run ``` Check that open spans remain unfinished during eviction and retain their input, metadata, and parent relationships when the chat completes. Share the durability API cases with the existing subagent regression. Refs: https://linear.app/braintrustdata/issue/SDK-781 Refs: #1054
> 🤖 **This text was AI-generated.** Cache eviction could lose the original Braintrust span for a long-running Pydantic AI chat. Before the fix in PR #1054, replay created multiple workflow spans for the same chat: ```text first message -> cache eviction -> replay -> duplicate workflow spans ``` Add a two-turn chat regression for both Temporal durability APIs. Force cache eviction, stop the worker while the chat waits, and resume it on a fresh worker: ```text workflow span └─ chat_session (same IDs and start time after restart) ├─ first agent run └─ second agent run ``` Check that open spans remain unfinished during eviction and retain their input, metadata, and parent relationships when the chat completes. Share the durability API cases with the existing subagent regression. Refs: https://linear.app/braintrustdata/issue/SDK-781 Refs: #1054
resolves https://linear.app/braintrustdata/issue/SDK-774
With
BraintrustPluginand Pydantic AI's Temporal durability integration, a tool such asdelegate_taskruns in workflow code and starts a sub-agent child workflow. Three things went wrong:delegate_taskspan instead of under it. The workflow outbound interceptor always sent the workflow span as the parent of activities, local activities, and child workflows, ignoring the current in-workflow span.agent_run,delegate_task, andchatspans or users' own spans, were logged again on every replay with new random IDs. The plugin also skipped the workflow span during replay, so the copies became separate root traces.BraintrustPlugin(logger=...), some workflow span writes went to the wrong logger. The plugin's logger is set through a thread-local override, and the plugin only set it on the thread that started the workflow. Temporal can run later activations of the same workflow on other worker threads, so writes made there went to the global logger instead, splitting spans across loggers. This only affects users who passlogger=.The fix:
logger.pygets an internal, context-local span hook. When a hook is set,SpanImpltakes default row, span, and trace IDs from it, andlog_internaldrops writes the hook rejects. Spans only look up the hook after one has been installed, so code that never uses Temporal pays for one boolean check per span.__braintrust_get_loggercall the workflow makes at startup. The activity path already did this per activity.This also fixes the workflow span's end time. Evicting a workflow unwinds its coroutine, which ended the span at the first eviction. The span now ends when the workflow really finishes, because the later end row merges into the same span.
Tests and CI:
test_temporal(latest)now also installspydantic-ai-slim[temporal], pinned to thepydantic-ai-integrationlatest version. Pydantic AI's temporal extra needs temporalio>=1.27 and Python 3.11+, so the 1.19.0 and 1.20.0 sessions skip the new test file.TemporalAgentwrapper and the newerTemporalDurabilitycapability (metadata={"temporal": False}on the delegating tool), with and without the workflow cache. Without the cache, every workflow task replays. On the capability path, Pydantic AI routes the model request to the activity before our wrapper runs, so thechatspan sits under the model activity instead of around it. Each request is still traced once.pylinton Python 3.10 skips the new test file, matching how it already skips thepydantic_aiintegration when pydantic-ai can't be installed.Limitations:
braintrust-temporal-workflow-span-context-v1patch branch.🤖 Generated with Claude Code