Skip to content

fix(temporal): keep pydantic-ai subagent spans nested and replay-safe - #1054

Merged
Abhijeet Prasad (AbhiPrasad) merged 3 commits into
mainfrom
abhi/temporal-subagent-spans-2610819b
Nov 3, 2026
Merged

Abhijeet Prasad (AbhiPrasad) merged 3 commits into
mainfrom
abhi/temporal-subagent-spans-2610819b

Conversation

@AbhiPrasad

@AbhiPrasad Abhijeet Prasad (AbhiPrasad) commented Nov 3, 2026 •

Copy link
Copy Markdown
Member

resolves https://linear.app/braintrustdata/issue/SDK-774

🤖 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. Three 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.
  3. With 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 pass logger=.
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.
  • Logger routing: before every write it allows, the workflow span hook re-applies the plugin's logger on the current thread, using the same __braintrust_get_logger call 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 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 the 1.19.0 and 1.20.0 sessions skip the new test file.
  • The new Pydantic AI 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. 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 the chat span sits under the model activity instead of around it. Each request is still traced once.
  • A Temporal test runs every workflow activation on a fresh thread, so the logger-routing bug reproduces every time.
  • pylint on Python 3.10 skips the new test file, matching how it already skips the pydantic_ai integration when pydantic-ai can't be installed.

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 don't get deterministic IDs yet, because those handlers run in their own context.

🤖 Generated with Claude Code

> 🤖 **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>
@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Nov 3, 2026 •

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-11-03T16:49:19.943588Z 7acf6e5 New commits
ℹ️ About Codex in GitHub

Your 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>

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment on lines +423 to +424
self._counter += 1
ids = _deterministic_span_ids(f"{self._base}:span:{self._counter}", self._legacy_ids)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

@AbhiPrasad
Abhijeet Prasad (AbhiPrasad) merged commit cf69e61 into main Nov 3, 2026
98 checks passed
@AbhiPrasad
Abhijeet Prasad (AbhiPrasad) deleted the abhi/temporal-subagent-spans-2610819b branch November 3, 2026 17:20
Abhijeet Prasad (AbhiPrasad) added a commit that referenced this pull request Nov 3, 2026
> 🤖 **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
Abhijeet Prasad (AbhiPrasad) added a commit that referenced this pull request Nov 3, 2026
> 🤖 **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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants