Add task_id to spans and fix Investigate Traces button

agentex-ui/components/task-header/investigate-traces-button.tsx

@@ -10,15 +10,15 @@ import { cn } from '@/lib/utils';
 type InvestigateTracesButtonProps = {
   className?: string;
   disabled?: boolean;
-  taskId: string;
+  traceId: string;
 };
 
 export const InvestigateTracesButton = forwardRef<
   HTMLAnchorElement,
   InvestigateTracesButtonProps
->(({ className, disabled = false, taskId, ...props }, ref) => {
+>(({ className, disabled = false, traceId, ...props }, ref) => {
   const { sgpAppURL } = useAgentexClient();
-  const sgpTracesURL = `${sgpAppURL}/beta/monitor?trace_id=${taskId}&tt-trace-id=${taskId}`;
+  const sgpTracesURL = `${sgpAppURL}/beta/monitor?trace_id=${traceId}&tt-trace-id=${traceId}`;
 
   if (!sgpAppURL) {
     return null;

agentex-ui/components/task-header/task-header.tsx

@@ -13,6 +13,7 @@ import {
   SelectValue,
 } from '@/components/ui/select';
 import { useSafeSearchParams } from '@/hooks/use-safe-search-params';
+import { useSpans } from '@/hooks/use-spans';
 
 import type { Agent } from 'agentex/resources';
 
@@ -36,6 +37,8 @@ export function TaskHeader({
 }: TaskHeaderProps) {
   const displayTaskId = taskId ? taskId.split('-')[0] : '';
   const { agentName: selectedAgentName } = useSafeSearchParams();
+  const { spans } = useSpans(taskId);
+  const traceId = spans[0]?.trace_id ?? taskId;
 
   const copyTaskId = async () => {
     if (taskId) {
@@ -109,7 +112,7 @@ export function TaskHeader({
               icon={Activity}
             />
           )}
-          {taskId && <InvestigateTracesButton taskId={taskId} />}
+          {taskId && traceId && <InvestigateTracesButton traceId={traceId} />}
         </div>
       </div>
     </motion.div>

agentex-ui/components/task-messages/task-message-reasoning-content.tsx

@@ -59,10 +59,9 @@ function TaskMessageReasoningImpl({ message }: TaskMessageReasoningProps) {
     if (message.content.type !== 'reasoning') {
       throw new Error('Message content is not a ReasoningContent');
     }
-    return [
-      ...(message.content.content ?? []),
-      ...(message.content.summary ?? []),
-    ].join('\n\n');
+    const content = message.content.content ?? [];
+    const summary = message.content.summary ?? [];
+    return (content.length > 0 ? content : summary).join('\n\n');
   }, [message.content]);
 
   const updateBlurEffects = () => {

agentex-ui/components/task-messages/task-messages.tsx

@@ -151,17 +151,34 @@ function TaskMessagesImpl({
 
   const shouldShowThinkingForLastPair = useMemo(() => {
     if (messagePairs.length === 0) return false;
+    if (rpcStatus !== 'pending' && rpcStatus !== 'success') return false;
 
     const lastPair = messagePairs[messagePairs.length - 1]!;
-    const hasNoAgentMessages = lastPair.agentMessages.length === 0;
-    const hasUserMessage = lastPair.userMessage !== null;
-
-    return (
-      hasUserMessage &&
-      hasNoAgentMessages &&
-      (rpcStatus === 'pending' || rpcStatus === 'success')
-    );
-  }, [messagePairs, rpcStatus]);
+
+    // No agent messages yet — waiting for first response
+    if (lastPair.agentMessages.length === 0) {
+      return lastPair.userMessage !== null;
+    }
+
+    const lastAgentMessage =
+      lastPair.agentMessages[lastPair.agentMessages.length - 1]!;
+    const lastType = lastAgentMessage.content.type;
+
+    // Already have text streaming or complete — not "thinking"
+    if (lastType === 'text') return false;
+
+    // Tool or reasoning still in progress — show their own indicator, not "Thinking..."
+    if (lastAgentMessage.streaming_status === 'IN_PROGRESS') return false;
+    if (
+      lastType === 'tool_request' &&
+      pendingToolCallIds.has(lastAgentMessage.content.tool_call_id)
+    )
+      return false;
+
+    // Last message is a completed tool_request, tool_response, reasoning, or data
+    // with no following text — agent is thinking about the next step
+    return true;
+  }, [messagePairs, rpcStatus, pendingToolCallIds]);
 
   // Measure container height for last-pair min-height
   useEffect(() => {

agentex-ui/hooks/use-spans.ts

@@ -8,8 +8,8 @@ import type { Span } from 'agentex/resources';
 
 export const spansKeys = {
   all: ['spans'] as const,
-  byTraceId: (traceId: string | null) =>
-    traceId ? ([...spansKeys.all, traceId] as const) : spansKeys.all,
+  byTaskId: (taskId: string | null) =>
+    taskId ? ([...spansKeys.all, 'task', taskId] as const) : spansKeys.all,
 };
 
 type UseSpansState = {
@@ -21,24 +21,37 @@ type UseSpansState = {
 /**
  * Fetches execution spans for observability and debugging of task execution.
  *
- * Spans are OpenTelemetry-style trace records that show the execution flow of an agent task.
- * The query is automatically disabled when no traceId is provided.
+ * Queries by task_id first. Falls back to trace_id=taskId for backward
+ * compatibility with spans created before the task_id column was added.
  *
- * @param traceId - string | null - The trace ID to fetch spans for, or null to disable the query
+ * @param taskId - string | null - The task ID to fetch spans for, or null to disable the query
  * @returns UseSpansState - Object containing the spans array, loading state, and any error message
  */
-export function useSpans(traceId: string | null): UseSpansState {
+export function useSpans(taskId: string | null): UseSpansState {
   const { agentexClient } = useAgentexClient();
 
   const { data, isLoading, error } = useQuery<Span[], Error>({
-    queryKey: spansKeys.byTraceId(traceId),
+    queryKey: spansKeys.byTaskId(taskId),
     queryFn: async ({ signal }) => {
-      if (!traceId) {
+      if (!taskId) {
         return [];
       }
-      return await agentexClient.spans.list({ trace_id: traceId }, { signal });
+
+      // task_id is not yet in the SDK types (SDK update pending), but the
+      // server already accepts it — cast until the SDK is regenerated.
+      const spansByTaskId = await agentexClient.spans.list(
+        { task_id: taskId } as Parameters<typeof agentexClient.spans.list>[0],
+        { signal }
+      );
+
+      if (spansByTaskId.length > 0) {
+        return spansByTaskId;
+      }
+
+      // Fallback: query by trace_id=taskId for backward compat with old spans
+      return await agentexClient.spans.list({ trace_id: taskId }, { signal });
     },
-    enabled: traceId !== null,
+    enabled: taskId !== null,
   });
 
   return {

agentex/database/migrations/alembic/versions/2026_04_14_1126_add_task_id_to_spans_57c5ed4f59ae.py

@@ -0,0 +1,52 @@
+"""add_task_id_to_spans
+
+Revision ID: 57c5ed4f59ae
+Revises: 4a9b7787ccd7
+Create Date: 2026-04-14 11:26:45.193515
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '57c5ed4f59ae'
+down_revision: Union[str, None] = '4a9b7787ccd7'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # Add nullable task_id column first (no FK yet, so backfill can run freely)
+    op.add_column('spans', sa.Column('task_id', sa.String(), nullable=True))
+
+    # Backfill task_id from trace_id where trace_id is a valid task ID.
+    # Uses a JOIN instead of a subquery for efficient matching.
+    op.execute("""
+        UPDATE spans
+        SET task_id = spans.trace_id
+        FROM tasks
+        WHERE spans.trace_id = tasks.id
+          AND spans.task_id IS NULL
+    """)
+
+    # Add FK constraint after backfill (NULL values are allowed by FK)
+    op.create_foreign_key(
+        'fk_spans_task_id_tasks',
+        'spans',
+        'tasks',
+        ['task_id'],
+        ['id'],
+        ondelete='SET NULL',
+    )
+
+    # Add index for querying spans by task_id
+    op.create_index('ix_spans_task_id', 'spans', ['task_id'])
+
+
+def downgrade() -> None:
+    op.drop_index('ix_spans_task_id', table_name='spans')
+    op.drop_constraint('fk_spans_task_id_tasks', 'spans', type_='foreignkey')
+    op.drop_column('spans', 'task_id')

agentex/src/adapters/orm.py

@@ -150,6 +150,7 @@ class SpanORM(BaseORM):
     __tablename__ = "spans"
     id = Column(String, primary_key=True, default=orm_id)  # Using UUIDs for IDs
     trace_id = Column(String, nullable=False)
+    task_id = Column(String, ForeignKey("tasks.id", ondelete="SET NULL"), nullable=True)
     parent_id = Column(String, nullable=True)
     name = Column(String, nullable=False)
     start_time = Column(DateTime(timezone=True), nullable=False)
@@ -166,6 +167,8 @@ class SpanORM(BaseORM):
         Index("ix_spans_trace_id_start_time", "trace_id", "start_time"),
         # Index for traversing span hierarchy
         Index("ix_spans_parent_id", "parent_id"),
+        # Index for filtering spans by task_id
+        Index("ix_spans_task_id", "task_id"),
     )
 
 

agentex/src/api/routes/spans.py

@@ -23,6 +23,7 @@ async def create_span(
     return await span_use_case.create(
         id=request.id,
         trace_id=request.trace_id,
+        task_id=request.task_id,
         name=request.name,
         parent_id=request.parent_id,
         start_time=request.start_time,
@@ -48,6 +49,7 @@ async def partial_update_span(
     return await span_use_case.partial_update(
         id=span_id,
         trace_id=request.trace_id,
+        task_id=request.task_id,
         name=request.name,
         parent_id=request.parent_id,
         start_time=request.start_time,
@@ -80,17 +82,19 @@ async def get_span(
 async def list_spans(
     span_use_case: DSpanUseCase,
     trace_id: str | None = None,
+    task_id: str | None = None,
     limit: int = Query(default=50, ge=1, le=1000),
     page_number: int = Query(default=1, ge=1),
     order_by: str | None = None,
     order_direction: str = "desc",
 ) -> list[Span]:
     """
-    List all spans for a given trace ID
+    List spans, optionally filtered by trace_id and/or task_id
     """
-    logger.info(f"Listing spans for trace ID: {trace_id}")
+    logger.info(f"Listing spans for trace_id={trace_id}, task_id={task_id}")
     spans = await span_use_case.list(
         trace_id=trace_id,
+        task_id=task_id,
         limit=limit,
         page_number=page_number,
         order_by=order_by,

agentex/src/api/schemas/spans.py

@@ -17,6 +17,11 @@ class CreateSpanRequest(BaseModel):
         title="The trace ID for this span",
         description="Unique identifier for the trace this span belongs to",
     )
+    task_id: str | None = Field(
+        None,
+        title="The task ID this span is associated with",
+        description="ID of the task this span belongs to",
+    )
     parent_id: str | None = Field(
         None,
         title="The parent span ID if this is a child span",
@@ -56,6 +61,11 @@ class UpdateSpanRequest(BaseModel):
         title="The trace ID for this span",
         description="Unique identifier for the trace this span belongs to",
     )
+    task_id: str | None = Field(
+        None,
+        title="The task ID this span is associated with",
+        description="ID of the task this span belongs to",
+    )
     parent_id: str | None = Field(
         None,
         title="The parent span ID if this is a child span",

agentex/src/domain/entities/spans.py

@@ -15,6 +15,10 @@ class SpanEntity(BaseModel):
         ...,
         title="The trace ID for this span",
     )
+    task_id: str | None = Field(
+        None,
+        title="The task ID this span is associated with",
+    )
     parent_id: str | None = Field(
         None,
         title="The parent span ID if this is a child span",

agentex/src/domain/use_cases/spans_use_case.py

@@ -21,6 +21,7 @@ async def create(
         name: str,
         trace_id: str,
         id: str | None = None,
+        task_id: str | None = None,
         parent_id: str | None = None,
         start_time: datetime | None = None,
         end_time: datetime | None = None,
@@ -38,6 +39,7 @@ async def create(
         span = SpanEntity(
             id=id,
             trace_id=trace_id,
+            task_id=task_id,
             parent_id=parent_id,
             name=name,
             start_time=start_time,
@@ -52,6 +54,7 @@ async def partial_update(
         self,
         id: str,
         trace_id: str | None = None,
+        task_id: str | None = None,
         name: str | None = None,
         parent_id: str | None = None,
         start_time: datetime | None = None,
@@ -70,6 +73,9 @@ async def partial_update(
         if trace_id is not None:
             span.trace_id = trace_id
 
+        if task_id is not None:
+            span.task_id = task_id
+
         if name is not None:
             span.name = name
 
@@ -108,19 +114,20 @@ async def list(
         limit: int,
         page_number: int,
         trace_id: str | None = None,
+        task_id: str | None = None,
         order_by: str | None = None,
         order_direction: str = "desc",
     ) -> list[SpanEntity]:
         """
-        List all spans for a given trace ID
+        List spans, optionally filtered by trace_id and/or task_id
         """
-        # Note: This would require custom implementation in the repository
-        # or filtering after fetching all spans
-
-        if trace_id:
-            filters = {"trace_id": trace_id}
-        else:
-            filters = None
+        filters: dict[str, str] | None = None
+        if trace_id or task_id:
+            filters = {}
+            if trace_id:
+                filters["trace_id"] = trace_id
+            if task_id:
+                filters["task_id"] = task_id
         return await self.span_repo.list(
             filters=filters,
             limit=limit,