Add task_id to spans and fix Investigate Traces button

[declan-scale]

Summary

• Add task_id column to spans: New nullable task_id foreign key on the spans table with a migration that backfills existing spans where trace_id matches a valid task ID. This lets the UI and API query spans by task directly instead of relying on the convention that trace_id equals the task ID.
• Fix Investigate Traces button: The button now uses the actual trace_id from spans (fetched via the new task_id query) rather than assuming task_id == trace_id. This makes the SGP Monitor link work correctly for agents whose trace IDs differ from their task IDs.
• Fix thinking indicator logic: The "Thinking..." indicator now correctly appears after tool calls complete (not just before the first agent response), and reasoning content no longer duplicates by concatenating both content and summary — it shows content when available, falling back to summary.

Backend changes

• New Alembic migration adds task_id column to spans table with FK to tasks, backfills from trace_id, and adds index
• SpanORM, SpanEntity, create/update schemas, and use case all support task_id
• /spans list endpoint accepts task_id query parameter
• Integration tests cover task_id on create, update, and null default

Frontend changes

• useSpans hook queries by task_id first, falls back to trace_id for backward compat
• InvestigateTracesButton accepts traceId instead of taskId
• TaskHeader fetches spans and extracts the real trace_id for the button
• Thinking indicator and reasoning dedup logic fixed

Follow-up required

The agentex-python SDK needs to be updated after this lands to support the new task_id field on span creation/listing. Until the SDK is updated, agents will continue to work (the field is nullable), but won't be able to explicitly set task_id when creating spans.

Test plan

• Verify spans can be created with and without task_id
• Verify /spans?task_id=<id> returns the correct spans
• Verify the Investigate Traces button links to the correct SGP Monitor trace
• Verify thinking indicator appears at the right times (after tool calls, before text response)
• Verify reasoning content doesn't show duplicated text
• Run make test for backend integration tests

🤖 Generated with Claude Code

Greptile Summary

This PR adds a task_id foreign key to the spans table with a backfill migration, threads the new field through the full backend stack (ORM → entity → schema → use case → route), and updates the frontend to query spans by task_id first (falling back to trace_id) so the Investigate Traces button resolves the real trace_id instead of assuming it equals task_id. Two frontend fixes land alongside: the reasoning content no longer concatenates content and summary, and the "Thinking…" indicator now correctly fires after completed tool calls.

Confidence Score: 4/5

Safe to merge after addressing the loading-state race on the Investigate Traces button (flagged in prior review thread) and adding a workaround comment for the SDK type cast.

The FK constraint issue in tests has been fixed. The only remaining P1 (button renders with the wrong trace_id URL during the initial spans fetch — flagged in prior thread) is still unresolved; all other findings are P2.

agentex-ui/components/task-header/task-header.tsx (loading-state race), agentex-ui/hooks/use-spans.ts (SDK type cast comment)

Important Files Changed

Filename	Overview
agentex-ui/components/task-header/investigate-traces-button.tsx	Simple prop rename from taskId to traceId; no logic changes.
agentex-ui/components/task-header/task-header.tsx	Adds useSpans to extract real trace_id; traceId falls back to taskId while spans are loading (pre-existing issue flagged in prior review thread).
agentex-ui/components/task-messages/task-messages.tsx	Thinking indicator logic correctly extended to show after completed tool calls; well-commented and dependency array is correct.
agentex-ui/components/task-messages/task-message-reasoning-content.tsx	Reasoning dedup fix: shows content when available, falls back to summary — correctly prevents double-display.
agentex-ui/hooks/use-spans.ts	task_id-first query with trace_id fallback works correctly; type cast for untyped SDK parameter should be documented.
agentex/database/migrations/alembic/versions/2026_04_14_1126_add_task_id_to_spans_57c5ed4f59ae.py	Column add + FK + index are correctly ordered; backfill comment claims batching the SQL does not implement (noted in prior thread).
agentex/src/adapters/orm.py	task_id FK and index added correctly to SpanORM; consistent with migration.
agentex/src/domain/use_cases/spans_use_case.py	task_id threading through create/partial_update/list is clean; filter dict build logic is correct.
agentex/tests/integration/api/spans/test_spans_api.py	FK violation from prior thread fixed via proper test_tasks fixture; new task_id filtering tests are thorough.
agentex/tests/unit/repositories/test_span_repository.py	Unit test creates a real TaskORM row to satisfy FK; all TaskORM columns outside id are nullable so the minimal insert is valid.

Sequence Diagram

sequenceDiagram
    participant TH as TaskHeader
    participant Hook as useSpans
    participant API as /spans endpoint
    participant DB as Database

    TH->>Hook: useSpans(taskId)
    Hook->>API: GET /spans?task_id={taskId}
    API->>DB: SELECT * FROM spans WHERE task_id = ?
    DB-->>API: spans[]
    API-->>Hook: spans[]

    alt spans found by task_id (new path)
        Hook-->>TH: spans with task_id
        TH->>TH: traceId = spans[0].trace_id
    else no spans found — backward compat fallback
        Hook->>API: GET /spans?trace_id={taskId}
        API->>DB: SELECT * FROM spans WHERE trace_id = ?
        DB-->>API: legacy spans[]
        API-->>Hook: legacy spans[]
        Hook-->>TH: legacy spans
        TH->>TH: traceId = spans[0].trace_id ?? taskId
    end

    TH->>TH: InvestigateTracesButton traceId={traceId}

Loading

Prompt To Fix All With AI

This is a comment left during a code review.
Path: agentex-ui/hooks/use-spans.ts
Line: 41-44

Comment:
**Type cast should document the SDK limitation**

The `as Parameters<...>[0]` cast bypasses TypeScript's type safety because `task_id` isn't in the current SDK types yet. Per the project's workaround-comment policy, this should include a note explaining why the cast is necessary and what resolves it, ideally with a linked ticket.

```suggestion
      // TODO(<TICKET>): Remove cast once agentex-python SDK is updated to include task_id
      // in the spans.list() parameters type. See PR description for details.
      const spansByTaskId = await agentexClient.spans.list(
        { task_id: taskId } as Parameters<typeof agentexClient.spans.list>[0],
        { signal }
      );
```

**Rule Used:** Add comments to explain complex or non-obvious log... ([source](https://app.greptile.com/review/custom-context?memory=928586f9-9432-435e-a385-026fa49318a2))

**Learned From**
[scaleapi/scaleapi#126958](https://github.com/scaleapi/scaleapi/pull/126958)

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

_{Reviews (2): Last reviewed commit: "Add more tests" | Re-trigger Greptile}

[greptile-apps]

Comment claims batching that the SQL does not implement

The comment says the backfill "batches updates to avoid long-held locks on large tables," but the UPDATE … FROM tasks WHERE … is a single unbounded statement. It acquires row-level locks on every matching span for the entire transaction duration — exactly the behaviour the comment says it avoids. On a large spans table this can stall concurrent writers for minutes and trigger statement timeouts.

Either remove the misleading claim, or implement actual batching, e.g.:

op.execute("""
    DO $$
    DECLARE
        updated_count INT;
    BEGIN
        LOOP
            UPDATE spans
            SET task_id = spans.trace_id
            FROM (
                SELECT spans.id
                FROM spans
                JOIN tasks ON spans.trace_id = tasks.id
                WHERE spans.task_id IS NULL
                LIMIT 1000
                FOR UPDATE SKIP LOCKED
            ) batch
            WHERE spans.id = batch.id;
            GET DIAGNOSTICS updated_count = ROW_COUNT;
            EXIT WHEN updated_count = 0;
        END LOOP;
    END
    $$;
""")

Prompt To Fix With AI

This is a comment left during a code review.
Path: agentex/database/migrations/alembic/versions/2026_04_14_1126_add_task_id_to_spans_57c5ed4f59ae.py
Line: 25-34

Comment:
**Comment claims batching that the SQL does not implement**

The comment says the backfill "batches updates to avoid long-held locks on large tables," but the `UPDATE … FROM tasks WHERE …` is a single unbounded statement. It acquires row-level locks on every matching span for the entire transaction duration — exactly the behaviour the comment says it avoids. On a large `spans` table this can stall concurrent writers for minutes and trigger statement timeouts.

Either remove the misleading claim, or implement actual batching, e.g.:

```python
op.execute("""
    DO $$
    DECLARE
        updated_count INT;
    BEGIN
        LOOP
            UPDATE spans
            SET task_id = spans.trace_id
            FROM (
                SELECT spans.id
                FROM spans
                JOIN tasks ON spans.trace_id = tasks.id
                WHERE spans.task_id IS NULL
                LIMIT 1000
                FOR UPDATE SKIP LOCKED
            ) batch
            WHERE spans.id = batch.id;
            GET DIAGNOSTICS updated_count = ROW_COUNT;
            EXIT WHEN updated_count = 0;
        END LOOP;
    END
    $$;
""")
```

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