feat(tracing): per-step obs wrappers inside business Temporal activities

[NiteshDhanpal]

What & why

On the Temporal / async path, _begin_obs skipped the per-step obs wrapper for any Temporal activity (Option A) and only stamped the ambient RunActivity span. So every business span in an async turn collapsed onto one obs span, whereas the sync path gives each step its own (1:1). This implements the _in_temporal_activity follow-up called out in #484.

The fix

Option A is only actually required for the SDK's own dispatched start-span / end-span activities (the in_temporal_workflow() path), where start and end run as separate activities that Temporal can route to different workers — so a wrapper opened in start-span could never be closed by end-span.

Inside a business activity (an agent turn's own adk.tracing.span, e.g. process_mortgage_turn), start and end run in the same process, so a wrapper is safe there: it nests under the interceptor's ambient RunActivity span and closes in-process.

_begin_obs now discriminates on activity type via _in_tracing_dispatch_activity() (true only for start-span/end-span):

• dispatched tracing activity → tag the ambient span (Option A, unchanged),
• everything else (sync or a business activity) → open a real per-step wrapper.

The now-dead _in_temporal_activity() is removed (its TODO is exactly what this implements). The bounded _OBS_HANDLES registry backstops any mis-discrimination.

Verified on infra-staging (rocket-mock-async-agent), one async turn

	before (Option A)	after
Tempo agentex.business wrapper spans	0	52
Tempo spans carrying the reverse tag	1	52
Postgres business spans	52	52
Postgres distinct obs_span_id	1	52 (1:1)

Each step (mortgage.advisor.turn, classify_intent, retrieve_docs.kb_query, authz.check.*, tool.*, telemetry.shard.*, synthesis.draft_reply) now gets its own obs span nested under RunActivity, matching the sync path.

Tests

Updated test_temporal_obs_backend.py for the new discriminator, plus two tests asserting: inside a dispatch activity → tag ambient (no wrapper); otherwise → open a wrapper.

🤖 Generated with Claude Code

Greptile Summary

The PR narrows the Temporal observability special case to the SDK’s dispatched start-span and end-span activities, allowing business activities to create per-step wrappers.

• Replaces generic Temporal-activity detection with activity-type discrimination.
• Adds tests for wrapper creation outside dispatch activities and ambient-span tagging inside them.

Confidence Score: 4/5

The PR appears safe to merge, with one non-blocking maintainability issue in the activity-name discriminator.

The new branching behavior is consistent with the registered activity names and intended wrapper lifecycle, but duplicating those names as literals creates an avoidable drift point.

Files Needing Attention: src/agentex/lib/core/tracing/trace.py

Important Files Changed

Filename	Overview
src/agentex/lib/core/tracing/trace.py	Narrows ambient-span reuse to dispatched tracing activities; the activity-name comparison should reference the existing enum rather than duplicate its values.
tests/test_temporal_obs_backend.py	Updates existing mocks and covers both resulting _begin_obs branches without introducing a concrete behavioral defect.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[_begin_obs] --> B{Tracing dispatch activity?}
    B -->|start-span or end-span| C[Tag ambient RunActivity span]
    C --> D[Return ambient correlation]
    B -->|Business activity or sync path| E[Open per-step obs wrapper]
    E --> F[Return wrapper correlation]

Loading

Prompt To Fix All With AI

### Issue 1
src/agentex/lib/core/tracing/trace.py:127
**Duplicated tracing activity names**

This comparison duplicates the values already centralized in `TracingActivityName`, creating a drift point where renamed activity values would stop dispatch recognition and open wrappers that cannot reliably be closed across workers. Use the shared enum values here so registration and discrimination remain synchronized.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

_{Reviews (1): Last reviewed commit: "feat(tracing): per-step obs wrappers ins..." | Re-trigger Greptile}

Context used:

• Rule used - Use enum values instead of hardcoded strings when ... (source)

Learned From
scaleapi/scaleapi#126557