feat(adk): support task_id on span creation

[declan-scale]

Summary

• Thread an optional task_id through ADK span creation (TracingModule.start_span / span context manager → TracingService → StartSpanParams activity → Trace/AsyncTrace.start_span), so spans can be associated with the task that produced them
• end_span intentionally unchanged — it already forwards the existing Span object which carries task_id from creation, matching how other fields (name, parent_id, input) are handled
• Tests at each layer (core trace, service, activity, ADK module) verify task_id is set on the Span and plumbed through correctly, and that it's preserved through end_span

Test plan

• rye run pytest tests/lib/adk/test_tracing_module.py tests/lib/adk/test_tracing_activities.py tests/lib/adk/test_tracing_service.py tests/lib/core/tracing/test_trace_task_id.py — 17 passed
• rye run pytest tests/lib/core/tracing/ tests/lib/adk/ — no regressions (86 passed, 2 pre-existing skips)
• rye run ruff format + rye run ruff check on touched files — clean

🤖 Generated with Claude Code

Greptile Summary

This PR threads an optional task_id parameter through the full ADK span creation stack — from TracingModule.start_span / span context manager down through StartSpanParams, TracingActivities, TracingService, and into both Trace and AsyncTrace.start_span — so spans can be associated with the task that produced them. end_span is intentionally left unchanged since the Span object already carries task_id from creation. Coverage is thorough with 17 new tests across all four layers.

Confidence Score: 5/5

Safe to merge — additive, backward-compatible change with comprehensive test coverage at every layer.

All changed files introduce only an optional parameter with a None default, preserving backward compatibility throughout. The Span model already had task_id. Tests cover all four layers (core trace, service, activity, module). The one P2 finding (AgentexLangGraphTracingHandler not yet forwarding task_id) is a gap in coverage, not a regression.

No files require special attention. Optionally consider updating _langgraph_tracing.py in a follow-up to propagate task_id through LLM/tool child spans.

Important Files Changed

Filename	Overview
src/agentex/lib/adk/_modules/tracing.py	Adds task_id param to both start_span and span context manager, threading it through to StartSpanParams and direct TracingService calls. Clean, backward-compatible addition.
src/agentex/lib/core/services/adk/tracing.py	Adds task_id parameter and forwards it to trace.start_span. Minimal, correct change.
src/agentex/lib/core/temporal/activities/adk/tracing_activities.py	Adds `task_id: str
src/agentex/lib/core/tracing/trace.py	Adds task_id to Trace.start_span, AsyncTrace.start_span, and both span context managers; stored on the Span object at creation.
tests/lib/adk/test_tracing_module.py	New tests cover task_id forwarding at the module level for start_span, end_span, and the context manager, including the no-op case.
tests/lib/adk/test_tracing_activities.py	Activity-level tests verify task_id is plumbed through StartSpanParams and preserved across end_span.
tests/lib/adk/test_tracing_service.py	Service-level tests confirm task_id is forwarded to trace.start_span and that end_span correctly passes through the span object.
tests/lib/core/tracing/test_trace_task_id.py	Core trace tests verify both Trace and AsyncTrace set task_id on the Span and that end_span preserves it.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant TracingModule
    participant StartSpanParams
    participant TracingActivities
    participant TracingService
    participant AsyncTrace
    participant Span

    Caller->>TracingModule: start_span(trace_id, name, task_id)
    alt in Temporal workflow
        TracingModule->>StartSpanParams: StartSpanParams(task_id=task_id)
        TracingModule->>TracingActivities: execute_activity(START_SPAN, params)
        TracingActivities->>TracingService: start_span(..., task_id=params.task_id)
    else direct call
        TracingModule->>TracingService: start_span(..., task_id=task_id)
    end
    TracingService->>AsyncTrace: trace.start_span(..., task_id=task_id)
    AsyncTrace->>Span: Span(task_id=task_id, ...)
    Span-->>Caller: span (with task_id set)

Loading

Prompt To Fix All With AI

This is a comment left during a code review.
Path: src/agentex/lib/adk/_modules/_langgraph_tracing.py
Line: 36-50

Comment:
**`task_id` not forwarded in LangGraph tracing handler**

`AgentexLangGraphTracingHandler` creates child spans for LLM calls and tool executions via `TracingModule.start_span`, but it never accepts or forwards a `task_id`. Spans produced by this handler will always have `task_id=None`, even when the caller has a task context.

If the intent is for *all* spans in a trace to carry the task association, consider accepting `task_id` in `__init__` alongside `parent_span_id` and threading it through `on_chat_model_start` / `on_tool_start`.

```python
def __init__(
    self,
    trace_id: str,
    parent_span_id: str | None = None,
    task_id: str | None = None,
    tracing: TracingModule | None = None,
) -> None:
    ...
    self._task_id = task_id
```

How can I resolve this? If you propose a fix, please make it concise.

_{Reviews (1): Last reviewed commit: "Add task_id to span creation" | Re-trigger Greptile}