From ff44e59edfcc77d4d458958361d418f15a5848e4 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Fri, 11 Sep 2026 09:23:57 -0400 Subject: [PATCH 1/4] refactor(agentex): remove the legacy Postgres spans API, table and UI reader Agent spans live in Scale GenAI Platform through the SDK's SGP tracing processor, so the Postgres-backed /spans routes, ORM model, repository, use case, schemas and tests go, and a migration drops the spans table. The UI traces sidebar was the last reader. It now fetches the task's trace from the platform through a scoped same-origin BFF route, the same way feedback already reaches the platform. The regenerated OpenAPI spec drops the spans surface, so the next SDK generation removes the client resource. Co-Authored-By: Claude Fable 5.1 --- CLAUDE.md | 5 +- agentex-ui/README.md | 4 +- .../api/traces/[traceId]/spans/route.test.ts | 83 ++++ .../app/api/traces/[traceId]/spans/route.ts | 47 ++ .../traces-sidebar/traces-sidebar.tsx | 11 +- agentex-ui/example.env.development | 2 +- agentex-ui/hooks/use-spans.test.tsx | 114 +++++ agentex-ui/hooks/use-spans.ts | 84 +++- ...2026_09_11_1239_drop_spans_78384970fed5.py | 63 +++ .../docs/runbooks/spans-task-id-backfill.md | 232 ---------- agentex/openapi.yaml | 388 ---------------- agentex/src/adapters/orm.py | 26 -- agentex/src/api/app.py | 2 - agentex/src/api/routes/spans.py | 103 ----- agentex/src/api/schemas/spans.py | 108 ----- agentex/src/domain/entities/spans.py | 49 -- .../domain/repositories/span_repository.py | 78 ---- .../src/domain/use_cases/spans_use_case.py | 140 ------ agentex/tests/fixtures/repositories.py | 22 - .../integration/api/spans/test_spans_api.py | 428 ------------------ .../fixtures/integration_client.py | 11 - .../unit/repositories/test_span_repository.py | 393 ---------------- 22 files changed, 381 insertions(+), 2012 deletions(-) create mode 100644 agentex-ui/app/api/traces/[traceId]/spans/route.test.ts create mode 100644 agentex-ui/app/api/traces/[traceId]/spans/route.ts create mode 100644 agentex-ui/hooks/use-spans.test.tsx create mode 100644 agentex/database/migrations/alembic/versions/2026_09_11_1239_drop_spans_78384970fed5.py delete mode 100644 agentex/docs/runbooks/spans-task-id-backfill.md delete mode 100644 agentex/src/api/routes/spans.py delete mode 100644 agentex/src/api/schemas/spans.py delete mode 100644 agentex/src/domain/entities/spans.py delete mode 100644 agentex/src/domain/repositories/span_repository.py delete mode 100644 agentex/src/domain/use_cases/spans_use_case.py delete mode 100644 agentex/tests/integration/api/spans/test_spans_api.py delete mode 100644 agentex/tests/unit/repositories/test_span_repository.py diff --git a/CLAUDE.md b/CLAUDE.md index 3d1522d2..fc378909 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -208,7 +208,7 @@ The backend (`agentex/src/`) follows a clean architecture with strict layer sepa ``` src/ ├── api/ # FastAPI routes, middleware, request/response schemas -│ ├── routes/ # API endpoints (agents, tasks, messages, spans, etc.) +│ ├── routes/ # API endpoints (agents, tasks, messages, states, etc.) │ ├── schemas/ # Pydantic request/response models │ ├── authentication_middleware.py │ └── app.py # FastAPI application setup @@ -285,7 +285,6 @@ Tests are organized by type and use different strategies: - **Agents**: Autonomous entities that execute tasks, managed via ACP protocol - **Tasks**: Work units with lifecycle states (pending → running → completed/failed). Identified by a UUID `id`; the human-readable `name` is **optional** (nullable) and, when set, globally unique. `task/create` is get-or-create keyed on `name`, so reusing an existing name returns that task with its prior history instead of creating a new one — omit `name` (or make it unique) whenever each call should produce a fresh task. - **Messages**: Communication between system and agents (stored in MongoDB) -- **Spans**: Execution traces for observability (OpenTelemetry-style) - **Events**: Domain events for async communication - **States**: Key-value state storage for agents - **Deployment History**: Track agent deployment versions and changes @@ -347,7 +346,7 @@ For any migration that adds a backfilled column with an FK and an index on a lar | Step | What | Why | |---|---|---| | **M1 (Alembic)** | `ADD COLUMN` (nullable) + `ADD CONSTRAINT ... NOT VALID` + `CREATE INDEX CONCURRENTLY` (in `autocommit_block()`) | Schema-only, all metadata-cheap or non-blocking. Each operation is idempotent (`IF NOT EXISTS` / `pg_constraint` guard) so the migration is safe to re-run on environments that already ran a previous (broken) version. | -| **Out-of-band runbook** | Chunked backfill script with `lock_timeout`, small batches, `COMMIT` between batches, `pg_sleep` between batches | Operator-driven; runs during a low-traffic window, can be cancelled cleanly, doesn't block pod startup. Pattern: `agentex/docs/runbooks/spans-task-id-backfill.md`. | +| **Out-of-band runbook** | Chunked backfill script with `lock_timeout`, small batches, `COMMIT` between batches, `pg_sleep` between batches | Operator-driven; runs during a low-traffic window, can be cancelled cleanly, doesn't block pod startup. | | **M2 (Alembic)** | `ALTER TABLE ... VALIDATE CONSTRAINT` (only if a fully validated FK state is actually needed) | Runs after the backfill so the scan finds no violations. `ShareUpdateExclusiveLock` is non-blocking against reads/writes but still scans the table — usually optional. | The application should also tolerate the partially-backfilled state at read time (e.g. ORing the new column against the legacy column where they overlap) so deployment of M1 is decoupled from the backfill's completion. diff --git a/agentex-ui/README.md b/agentex-ui/README.md index 9eb53860..90e9c0db 100644 --- a/agentex-ui/README.md +++ b/agentex-ui/README.md @@ -33,7 +33,7 @@ A modern web interface for building, testing, and monitoring intelligent agents. ### Observability -- **Execution Traces** - View OpenTelemetry-style spans for task execution +- **Execution Traces** - View a task's spans from Scale GenAI Platform (needs `SGP_API_URL`) - **Span Visualization** - Hierarchical view of execution flow - **Performance Metrics** - Timing and duration information for each execution step - **Error Tracking** - Detailed error information when tasks fail @@ -178,7 +178,7 @@ For Docker-related commands, see the Docker section in `build.ps1 help`. - `hooks/use-tasks.ts` - Task list with infinite scroll pagination - `hooks/use-task-messages.ts` - Message fetching and sending with message streaming for sync agents - `hooks/use-task-subscription.ts` - Real-time task updates via WebSocket for async agents -- `hooks/use-spans.ts` - Execution trace data +- `hooks/use-spans.ts` - Execution trace data (via `/api/traces`) **Components:** diff --git a/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts b/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts new file mode 100644 index 00000000..40a1d9da --- /dev/null +++ b/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts @@ -0,0 +1,83 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; + +import { GET } from './route'; + +const bff = vi.hoisted(() => ({ + baseURL: 'https://sgp.example/api' as string | undefined, + applyBffCredentials: vi.fn(async (_req: Request, headers: Headers) => { + headers.set('authorization', 'Bearer server-side'); + }), +})); + +vi.mock('@/app/api/_lib/bff', () => ({ + get SGP_BASE_URL() { + return bff.baseURL; + }, + applyBffCredentials: bff.applyBffCredentials, +})); + +function call(traceId: string) { + return GET(new Request(`http://ui.local/api/traces/${traceId}/spans`), { + params: Promise.resolve({ traceId }), + }); +} + +describe('GET /api/traces/[traceId]/spans', () => { + afterEach(() => { + vi.unstubAllGlobals(); + bff.baseURL = 'https://sgp.example/api'; + }); + + it('searches the platform for the trace with server-attached credentials', async () => { + const page = { items: [{ id: 's1', trace_id: 't1' }], has_more: false }; + const fetchMock = vi.fn().mockResolvedValue( + new Response(JSON.stringify(page), { + status: 200, + headers: { 'content-type': 'application/json' }, + }) + ); + vi.stubGlobal('fetch', fetchMock); + + const res = await call('t1'); + + expect(res.status).toBe(200); + expect(await res.json()).toEqual(page); + expect(bff.applyBffCredentials).toHaveBeenCalledTimes(1); + const [url, init] = fetchMock.mock.calls[0]!; + expect(url).toBe( + 'https://sgp.example/api/v5/spans/search?limit=100&sort_by=start_timestamp&sort_order=asc' + ); + expect(init.method).toBe('POST'); + expect(JSON.parse(init.body)).toEqual({ trace_ids: ['t1'] }); + expect(new Headers(init.headers).get('authorization')).toBe( + 'Bearer server-side' + ); + }); + + it('passes the upstream status through', async () => { + vi.stubGlobal( + 'fetch', + vi + .fn() + .mockResolvedValue( + new Response(JSON.stringify({ detail: 'forbidden' }), { status: 403 }) + ) + ); + + const res = await call('t1'); + + expect(res.status).toBe(403); + expect(await res.json()).toEqual({ detail: 'forbidden' }); + }); + + it('returns 503 when the platform API is not configured', async () => { + bff.baseURL = undefined; + const fetchMock = vi.fn(); + vi.stubGlobal('fetch', fetchMock); + + const res = await call('t1'); + + expect(res.status).toBe(503); + expect(fetchMock).not.toHaveBeenCalled(); + }); +}); diff --git a/agentex-ui/app/api/traces/[traceId]/spans/route.ts b/agentex-ui/app/api/traces/[traceId]/spans/route.ts new file mode 100644 index 00000000..2300a035 --- /dev/null +++ b/agentex-ui/app/api/traces/[traceId]/spans/route.ts @@ -0,0 +1,47 @@ +import { NextResponse } from 'next/server'; + +import { applyBffCredentials, SGP_BASE_URL } from '@/app/api/_lib/bff'; + +/** + * Scoped BFF proxy for one trace's spans, read from the platform's span search with the + * credentials attached server-side. Only this path is exposed, not a catch-all, so the + * browser can't reach arbitrary platform endpoints with those credentials. + */ +export const dynamic = 'force-dynamic'; + +// The platform caps a search page at this size, and the sidebar shows the first page. +const PAGE_SIZE = 100; + +export async function GET( + request: Request, + ctx: { params: Promise<{ traceId: string }> } +): Promise { + if (!SGP_BASE_URL) { + return NextResponse.json( + { error: 'SGP traces are not configured. Set SGP_API_URL.' }, + { status: 503 } + ); + } + + const { traceId } = await ctx.params; + const headers = new Headers({ + 'Content-Type': 'application/json', + accept: 'application/json', + }); + await applyBffCredentials(request, headers); + + const query = new URLSearchParams({ + limit: String(PAGE_SIZE), + sort_by: 'start_timestamp', + sort_order: 'asc', + }); + const upstream = await fetch(`${SGP_BASE_URL}/v5/spans/search?${query}`, { + method: 'POST', + headers, + body: JSON.stringify({ trace_ids: [traceId] }), + }); + return new Response(upstream.body, { + status: upstream.status, + headers: { 'content-type': 'application/json' }, + }); +} diff --git a/agentex-ui/components/traces-sidebar/traces-sidebar.tsx b/agentex-ui/components/traces-sidebar/traces-sidebar.tsx index d056a32a..4d8cf703 100644 --- a/agentex-ui/components/traces-sidebar/traces-sidebar.tsx +++ b/agentex-ui/components/traces-sidebar/traces-sidebar.tsx @@ -20,7 +20,7 @@ type TracesSidebarProps = { export function TracesSidebar({ isOpen }: TracesSidebarProps) { const { taskID } = useSafeSearchParams(); - const { spans, isLoading, error } = useSpans(taskID); + const { spans, hasMore, isLoading, error } = useSpans(taskID); return ( @@ -76,8 +76,15 @@ export function TracesSidebar({ isOpen }: TracesSidebarProps) { )} + {hasMore && ( +
+ Showing the first {spans.length} spans. Use Investigate + traces for the full trace. +
+ )} + {spans.map(span => { - const startTime = new Date(span.start_time); + const startTime = new Date(span.start_timestamp); return (
diff --git a/agentex-ui/example.env.development b/agentex-ui/example.env.development index 37547eee..ef6d2123 100644 --- a/agentex-ui/example.env.development +++ b/agentex-ui/example.env.development @@ -1,6 +1,6 @@ AGENTEX_API_URL=http://localhost:5003 # /api/agentex → agentex API # ENABLE_AGENT_RUN_SCHEDULES=true # enables scheduled tasks in both API and UI -# SGP_API_URL= # optional: /api/feedback & /api/user-info → SGP API +# SGP_API_URL= # optional: /api/feedback, /api/user-info & /api/traces → SGP API # NEXT_PUBLIC_SGP_APP_URL= # optional: links to SGP traces #---- OIDC login (opt-in) — set AGENTEX_UI_AUTH_PROVIDER_ID to enable login. Its value must diff --git a/agentex-ui/hooks/use-spans.test.tsx b/agentex-ui/hooks/use-spans.test.tsx new file mode 100644 index 00000000..0b01f6bb --- /dev/null +++ b/agentex-ui/hooks/use-spans.test.tsx @@ -0,0 +1,114 @@ +import type { ReactNode } from 'react'; + +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import { renderHook, waitFor } from '@testing-library/react'; +import { afterEach, describe, expect, it, vi } from 'vitest'; + +import { useSpans } from './use-spans'; + +vi.mock('@/hooks/use-safe-search-params', () => ({ + useSafeSearchParams: () => ({ sgpAccountID: 'acct-1' }), +})); + +function createWrapper() { + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + return function Wrapper({ children }: { children: ReactNode }) { + return ( + {children} + ); + }; +} + +function jsonResponse(body: unknown, status = 200) { + return new Response(JSON.stringify(body), { + status, + headers: { 'content-type': 'application/json' }, + }); +} + +const span = { + id: 'span-1', + trace_id: 'task-1', + parent_id: null, + name: 'run_agent', + start_timestamp: '2026-01-01T00:00:00Z', + end_timestamp: null, +}; + +describe('useSpans', () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it('reads the task trace through the BFF with the selected account', async () => { + const fetchMock = vi + .fn() + .mockResolvedValue(jsonResponse({ items: [span], has_more: false })); + vi.stubGlobal('fetch', fetchMock); + + const { result } = renderHook(() => useSpans('task-1'), { + wrapper: createWrapper(), + }); + + await waitFor(() => expect(result.current.isLoading).toBe(false)); + + expect(result.current.spans).toEqual([span]); + expect(result.current.hasMore).toBe(false); + expect(result.current.error).toBeNull(); + expect(fetchMock).toHaveBeenCalledTimes(1); + const [url, init] = fetchMock.mock.calls[0]!; + expect(url).toBe('/api/traces/task-1/spans'); + expect(init.credentials).toBe('include'); + expect(init.headers).toEqual({ 'x-selected-account-id': 'acct-1' }); + }); + + it('reports when the trace has more spans than the page', async () => { + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue(jsonResponse({ items: [span], has_more: true })) + ); + + const { result } = renderHook(() => useSpans('task-1'), { + wrapper: createWrapper(), + }); + + await waitFor(() => expect(result.current.isLoading).toBe(false)); + + expect(result.current.spans).toEqual([span]); + expect(result.current.hasMore).toBe(true); + }); + + it('surfaces the BFF error message when the platform is not configured', async () => { + vi.stubGlobal( + 'fetch', + vi + .fn() + .mockResolvedValue( + jsonResponse({ error: 'SGP traces are not configured.' }, 503) + ) + ); + + const { result } = renderHook(() => useSpans('task-1'), { + wrapper: createWrapper(), + }); + + await waitFor(() => + expect(result.current.error).toBe('SGP traces are not configured.') + ); + expect(result.current.spans).toEqual([]); + }); + + it('does not fetch without a task', () => { + const fetchMock = vi.fn(); + vi.stubGlobal('fetch', fetchMock); + + const { result } = renderHook(() => useSpans(null), { + wrapper: createWrapper(), + }); + + expect(result.current.spans).toEqual([]); + expect(fetchMock).not.toHaveBeenCalled(); + }); +}); diff --git a/agentex-ui/hooks/use-spans.ts b/agentex-ui/hooks/use-spans.ts index f2213186..15a43667 100644 --- a/agentex-ui/hooks/use-spans.ts +++ b/agentex-ui/hooks/use-spans.ts @@ -2,9 +2,7 @@ import { useQuery } from '@tanstack/react-query'; -import { useAgentexClient } from '@/components/providers'; - -import type { Span } from 'agentex/resources'; +import { useSafeSearchParams } from '@/hooks/use-safe-search-params'; export const spansKeys = { all: ['spans'] as const, @@ -12,50 +10,88 @@ export const spansKeys = { taskId ? ([...spansKeys.all, 'task', taskId] as const) : spansKeys.all, }; +/** A platform span as the traces BFF route returns it. */ +export type TraceSpan = { + id: string; + trace_id: string; + parent_id: string | null; + name: string; + start_timestamp: string; + end_timestamp: string | null; + status?: string | null; + type?: string | null; + input?: Record | null; + output?: Record | null; + metadata?: Record | null; +}; + +type SpansPage = { + items: TraceSpan[]; + has_more?: boolean; +}; + +type SpansResult = { + items: TraceSpan[]; + hasMore: boolean; +}; + type UseSpansState = { - spans: Span[]; + spans: TraceSpan[]; + // True when the trace has more spans than the one page the sidebar shows. + hasMore: boolean; isLoading: boolean; error: string | null; }; /** - * Fetches execution spans for observability and debugging of task execution. - * - * Queries by task_id first. Falls back to trace_id=taskId for backward - * compatibility with spans created before the task_id column was added. + * Fetches a task's execution spans from Scale GenAI Platform, where agents trace under the + * task id, through the same-origin BFF route that attaches credentials server-side. * - * @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 + * @param taskId - The task ID to fetch spans for, or null to disable the query + * @returns The first page of spans in start order, whether more exist, the loading state, and any error message */ export function useSpans(taskId: string | null): UseSpansState { - const { agentexClient } = useAgentexClient(); + const { sgpAccountID } = useSafeSearchParams(); - const { data, isLoading, error } = useQuery({ + const { data, isLoading, error } = useQuery({ queryKey: spansKeys.byTaskId(taskId), - queryFn: async ({ signal }) => { + queryFn: async ({ signal }): Promise => { if (!taskId) { - return []; + return { items: [], hasMore: false }; } - // 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[0], - { signal } + const response = await fetch( + `/api/traces/${encodeURIComponent(taskId)}/spans`, + { + credentials: 'include', + // Selected account, same source as the SDK, forwarded by the BFF. + headers: sgpAccountID + ? { 'x-selected-account-id': sgpAccountID } + : {}, + signal, + } ); - if (spansByTaskId.length > 0) { - return spansByTaskId; + if (!response.ok) { + const body = await response.json().catch(() => ({})); + const message = + typeof body.error === 'string' + ? body.error + : typeof body.detail === 'string' + ? body.detail + : `Request failed with status ${response.status}`; + throw new Error(message); } - // Fallback: query by trace_id=taskId for backward compat with old spans - return await agentexClient.spans.list({ trace_id: taskId }, { signal }); + const page: SpansPage = await response.json(); + return { items: page.items ?? [], hasMore: page.has_more ?? false }; }, enabled: taskId !== null, }); return { - spans: data ?? [], + spans: data?.items ?? [], + hasMore: data?.hasMore ?? false, isLoading, error: error?.message ?? null, }; diff --git a/agentex/database/migrations/alembic/versions/2026_09_11_1239_drop_spans_78384970fed5.py b/agentex/database/migrations/alembic/versions/2026_09_11_1239_drop_spans_78384970fed5.py new file mode 100644 index 00000000..1c2a9948 --- /dev/null +++ b/agentex/database/migrations/alembic/versions/2026_09_11_1239_drop_spans_78384970fed5.py @@ -0,0 +1,63 @@ +"""drop spans + +Revision ID: 78384970fed5 +Revises: c4e8b2a7f91d +Create Date: 2026-09-11 12:39:10.000000 + +Drops the legacy Postgres-backed spans table. Agent spans are written to the +platform's tracing service by the SDK's SGP tracing processor, and the +/spans API that fed this table is removed in the same change, so nothing +reads or writes it any more. + +Safety: +- DROP TABLE is metadata-only in PostgreSQL (the files are unlinked), so it + completes well inside the statement timeout regardless of table size. It + needs an AccessExclusiveLock; a writer still holding the table (an old pod + mid-rollout) makes the lock wait hit lock_timeout, and the pod retries the + migration on its next start. +- IF EXISTS keeps re-runs idempotent. The indexes and the foreign key to + tasks go with the table. +- Downgrade recreates the empty table with the shape the ORM last declared. + Indexes on a table created in the same migration are built plain, since + there are no writers to block. +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "78384970fed5" +down_revision: str | None = "c4e8b2a7f91d" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.execute("DROP TABLE IF EXISTS spans") + + +def downgrade() -> None: + op.create_table( + "spans", + sa.Column("id", sa.String(), primary_key=True), + sa.Column("trace_id", sa.String(), nullable=False), + sa.Column( + "task_id", + sa.String(), + sa.ForeignKey("tasks.id", ondelete="SET NULL", name="fk_spans_task_id_tasks"), + nullable=True, + ), + sa.Column("parent_id", sa.String(), nullable=True), + sa.Column("name", sa.String(), nullable=False), + sa.Column("start_time", sa.DateTime(timezone=True), nullable=False), + sa.Column("end_time", sa.DateTime(timezone=True), nullable=True), + sa.Column("input", sa.JSON(), nullable=True), + sa.Column("output", sa.JSON(), nullable=True), + sa.Column("data", sa.JSON(), nullable=True), + ) + op.create_index("ix_spans_trace_id", "spans", ["trace_id"]) + op.create_index("ix_spans_trace_id_start_time", "spans", ["trace_id", "start_time"]) + op.create_index("ix_spans_parent_id", "spans", ["parent_id"]) + op.create_index("ix_spans_task_id", "spans", ["task_id"]) diff --git a/agentex/docs/runbooks/spans-task-id-backfill.md b/agentex/docs/runbooks/spans-task-id-backfill.md deleted file mode 100644 index a2171b63..00000000 --- a/agentex/docs/runbooks/spans-task-id-backfill.md +++ /dev/null @@ -1,232 +0,0 @@ -# Runbook: backfill `spans.task_id` from `trace_id` - -## Purpose - -Populate the `spans.task_id` column on historical rows with the value from -`trace_id`, where `trace_id` matches an existing `tasks.id`. The column is -written natively by current clients; existing rows still have -`task_id = NULL`. - -**This backfill is optional.** `SpanRepository.list` ORs on `trace_id` when -filtering by `task_id`, so application reads are correct without it. Run this -runbook only when: - -- We want to drop the `trace_id` fallback in the application (cleanup goal). -- A downstream consumer specifically needs `WHERE task_id IS NOT NULL` to - identify task-scoped spans for historical data. - -## Background - -An earlier in-band attempt to backfill this column via Alembic ran a single -large `UPDATE` inside Alembic's transaction on a multi-tens-of-GB `spans` -table. It held row locks and exhausted the application connection pool while -concurrent span writes piled up, taking the service offline until the -migration was killed. - -This runbook does the same work in **batched, committed-per-batch chunks** -with explicit operator control over timing. - -## Pre-flight - -> Run all queries against the agentex database. - -1. Confirm the migration that adds the column has been applied: - - ```sql - SELECT version_num FROM alembic_version; - ``` - - Expected: `a9959ebcbe98` (or later). If the value is `4a9b7787ccd7` or - earlier, **stop** — deploy the migration first. - -2. Confirm the column exists and the FK + index are in place: - - ```sql - \d spans - ``` - - Expected: `task_id` column, `fk_spans_task_id_tasks` constraint, - `ix_spans_task_id` index. - -3. Measure the size of the backlog: - - ```sql - SELECT - count(*) FILTER (WHERE s.task_id IS NULL AND t.id IS NOT NULL) AS to_backfill, - count(*) FILTER (WHERE s.task_id IS NULL AND t.id IS NULL) AS untouched_non_task_spans, - count(*) FILTER (WHERE s.task_id IS NOT NULL) AS already_set, - count(*) AS total - FROM spans s - LEFT JOIN tasks t ON t.id = s.trace_id; - ``` - - Note the `to_backfill` number. With 10 000 rows per batch and 100 ms - sleep between batches, expect roughly 1 000 batches per minute. A - ~25M-row backlog is therefore ~25 minutes of wall-clock time at the - default pace. - -4. Confirm there is no other long-running write or maintenance activity on - the `spans` table: - - ```sql - SELECT pid, now() - xact_start AS duration, state, wait_event, query - FROM pg_stat_activity - WHERE query ILIKE '%spans%' - AND state <> 'idle' - ORDER BY duration DESC; - ``` - -## Coordination requirements - -- Notify the operating team before starting and post the `to_backfill` - count from step 3. -- Run during a confirmed low-traffic window, or shift traffic away from the - service for the duration. -- Get sign-off from infra/platform. -- Have a rollback plan ready — the script below is safely cancellable - (Ctrl-C in psql, then run the cancel snippet); no schema change happens. - -## Execution - -Connect to the agentex DB with `psql` (not via the application), and run: - -```sql --- 1. Apply per-session timeouts. lock_timeout fails fast if any other --- session holds AccessExclusiveLock on spans (e.g. a stuck migration). --- statement_timeout caps the *per-batch* runtime; total runtime is --- unbounded because we loop and commit between batches. -SET lock_timeout = '3s'; -SET statement_timeout = '60s'; - --- 2. Loop in 10 000-row chunks. ROW_COUNT is checked between iterations to --- detect when the backlog is drained. Each iteration commits before the --- next starts, so autovacuum can reclaim dead tuples and the table --- does not bloat. -DO $$ -DECLARE - rows_updated INT := 1; - total_updated BIGINT := 0; -BEGIN - WHILE rows_updated > 0 LOOP - WITH batch AS ( - SELECT s.ctid - FROM spans s - JOIN tasks t ON t.id = s.trace_id - WHERE s.task_id IS NULL - LIMIT 10000 - ) - UPDATE spans - SET task_id = trace_id - WHERE ctid IN (SELECT ctid FROM batch); - - GET DIAGNOSTICS rows_updated = ROW_COUNT; - total_updated := total_updated + rows_updated; - - RAISE NOTICE 'updated batch: % rows (running total: %)', - rows_updated, total_updated; - - COMMIT; - PERFORM pg_sleep(0.1); - END LOOP; -END$$; -``` - -Notes on the SQL choices: - -- `ctid` selection in a CTE means each batch operates on a fixed set of rows - selected at the start of the batch — we don't read and write the same row - in a single statement, and we don't hold locks across batches. -- `JOIN tasks` filters out spans whose `trace_id` is not actually a task id - (system or framework spans). Those rows stay with `task_id = NULL`, which - matches the application's existing semantics. -- The `COMMIT` inside the `DO` block requires PostgreSQL ≥ 11. Confirm the - server version with `SELECT version();` if running this in an older env. -- `pg_sleep(0.1)` gives autovacuum and concurrent writes breathing room. - -### Monitoring while it runs - -In a separate `psql` session: - -```sql --- Live progress (rerun periodically) -SELECT count(*) AS remaining -FROM spans s JOIN tasks t ON t.id = s.trace_id -WHERE s.task_id IS NULL; - --- Active sessions on spans -SELECT pid, now() - xact_start AS duration, state, wait_event, left(query, 80) -FROM pg_stat_activity -WHERE query ILIKE '%spans%' AND state <> 'idle' -ORDER BY duration DESC; - --- Lock contention -SELECT pg_class.relname, pg_locks.mode, pg_locks.granted, pg_locks.pid -FROM pg_locks -JOIN pg_class ON pg_class.oid = pg_locks.relation -WHERE pg_class.relname = 'spans'; -``` - -## Cancellation - -The script is safe to interrupt at any batch boundary — already-committed -batches are durable; in-flight batches roll back cleanly. - -- **From the same psql session**: Ctrl-C cancels the current statement. -- **From another session**, if the runner is stuck: - - ```sql - -- Cancel the runner gracefully (releases locks at end of current batch) - SELECT pg_cancel_backend(); - - -- If pg_cancel_backend doesn't return control, escalate: - SELECT pg_terminate_backend(); - ``` - - Prefer `pg_cancel_backend` for this batched runbook; - `pg_terminate_backend` should be reserved for situations where graceful - cancellation has already failed. - -## Exit criteria - -The backfill is complete when both of the following return `0`: - -```sql -SELECT count(*) AS remaining_to_backfill -FROM spans s JOIN tasks t ON t.id = s.trace_id -WHERE s.task_id IS NULL; -``` - -```sql --- Sanity: no orphaned task_ids (FK is NOT VALID, so this is the only --- way to verify referential cleanliness for backfilled rows). NOT EXISTS --- with a correlated subquery uses the tasks(id) primary-key index and --- avoids materialising every tasks.id (which NOT IN would force). -SELECT count(*) AS orphaned_task_ids -FROM spans s -WHERE s.task_id IS NOT NULL - AND NOT EXISTS (SELECT 1 FROM tasks t WHERE t.id = s.task_id); -``` - -Once both are zero, notify the operating team and restore normal traffic -(if it was diverted). - -## Follow-ups after a successful backfill - -- Open a PR to drop the `trace_id` OR-fallback in - `SpanRepository.list` — task-scoped spans can now be queried purely by - `task_id`. -- Optionally run `ALTER TABLE spans VALIDATE CONSTRAINT fk_spans_task_id_tasks` - to convert the FK from `NOT VALID` to `VALID`. This takes a - `ShareUpdateExclusiveLock` (does **not** block reads/writes) and scans the - table once. Coordinate with infra; a multi-tens-of-GB scan is non-trivial - even when non-blocking. - -## What this runbook deliberately does *not* do - -- It does **not** modify any schema. The column, FK, and index are managed - by the alembic migrations. -- It does **not** populate `task_id` for spans whose `trace_id` is not a - task id. Those rows correctly remain `NULL`. -- It does **not** run from inside a deploy or pod-startup path. Migrations - run on agentex pod startup; a multi-minute backfill in that path would - recreate the original failure mode. diff --git a/agentex/openapi.yaml b/agentex/openapi.yaml index 88dd2a78..13fa0fa2 100644 --- a/agentex/openapi.yaml +++ b/agentex/openapi.yaml @@ -2098,162 +2098,6 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' - /spans: - post: - tags: - - Spans - summary: Create Span - description: Create a new span with the provided parameters - operationId: create_span_spans_post - requestBody: - required: true - content: - application/json: - schema: - $ref: '#/components/schemas/CreateSpanRequest' - responses: - '200': - description: Successful Response - content: - application/json: - schema: - $ref: '#/components/schemas/Span' - '422': - description: Validation Error - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPValidationError' - get: - tags: - - Spans - summary: List Spans - description: List spans, optionally filtered by trace_id and/or task_id - operationId: list_spans_spans_get - parameters: - - name: trace_id - in: query - required: false - schema: - anyOf: - - type: string - - type: 'null' - title: Trace Id - - name: task_id - in: query - required: false - schema: - anyOf: - - type: string - - type: 'null' - title: Task Id - - name: limit - in: query - required: false - schema: - type: integer - maximum: 1000 - minimum: 1 - default: 50 - title: Limit - - name: page_number - in: query - required: false - schema: - type: integer - minimum: 1 - default: 1 - title: Page Number - - name: order_by - in: query - required: false - schema: - anyOf: - - type: string - - type: 'null' - title: Order By - - name: order_direction - in: query - required: false - schema: - type: string - default: desc - title: Order Direction - responses: - '200': - description: Successful Response - content: - application/json: - schema: - type: array - items: - $ref: '#/components/schemas/Span' - title: Response List Spans Spans Get - '422': - description: Validation Error - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPValidationError' - /spans/{span_id}: - patch: - tags: - - Spans - summary: Partial Update Span - description: Update a span with the provided output data and mark it as complete - operationId: partial_update_span_spans__span_id__patch - parameters: - - name: span_id - in: path - required: true - schema: - type: string - title: Span Id - requestBody: - required: true - content: - application/json: - schema: - $ref: '#/components/schemas/UpdateSpanRequest' - responses: - '200': - description: Successful Response - content: - application/json: - schema: - $ref: '#/components/schemas/Span' - '422': - description: Validation Error - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPValidationError' - get: - tags: - - Spans - summary: Get Span - description: Get a span by ID - operationId: get_span_spans__span_id__get - parameters: - - name: span_id - in: path - required: true - schema: - type: string - title: Span Id - responses: - '200': - description: Successful Response - content: - application/json: - schema: - $ref: '#/components/schemas/Span' - '422': - description: Validation Error - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPValidationError' /states: post: tags: @@ -5102,86 +4946,6 @@ components: required: - docker_image title: CreateDeploymentRequest - CreateSpanRequest: - properties: - id: - anyOf: - - type: string - - type: 'null' - title: Unique Span ID - description: Unique identifier for the span. If not provided, an ID will - be generated. - trace_id: - type: string - title: The trace ID for this span - description: Unique identifier for the trace this span belongs to - task_id: - anyOf: - - type: string - - type: 'null' - title: The task ID this span is associated with - description: ID of the task this span belongs to - parent_id: - anyOf: - - type: string - - type: 'null' - title: The parent span ID if this is a child span - description: ID of the parent span if this is a child span in a trace - name: - type: string - title: The name of the span - description: Name that describes what operation this span represents - start_time: - type: string - format: date-time - title: The start time of the span - description: The time the span started - end_time: - anyOf: - - type: string - format: date-time - - type: 'null' - title: The end time of the span - description: The time the span ended - input: - anyOf: - - additionalProperties: true - type: object - - items: - additionalProperties: true - type: object - type: array - - type: 'null' - title: The input data for the span - description: Input parameters or data for the operation - output: - anyOf: - - additionalProperties: true - type: object - - items: - additionalProperties: true - type: object - type: array - - type: 'null' - title: The output data from the span - description: Output data resulting from the operation - data: - anyOf: - - additionalProperties: true - type: object - - items: - additionalProperties: true - type: object - type: array - - type: 'null' - title: Additional data associated with the span - description: Any additional metadata or context for the span - type: object - required: - - trace_id - - name - - start_time - title: CreateSpanRequest CreateStateRequest: properties: task_id: @@ -6379,83 +6143,6 @@ components: required: - scheduled_time title: SkipRunScheduleRequest - Span: - properties: - id: - type: string - title: Unique Span ID - trace_id: - type: string - title: The trace ID for this span - description: Unique identifier for the trace this span belongs to - task_id: - anyOf: - - type: string - - type: 'null' - title: The task ID this span is associated with - description: ID of the task this span belongs to - parent_id: - anyOf: - - type: string - - type: 'null' - title: The parent span ID if this is a child span - description: ID of the parent span if this is a child span in a trace - name: - type: string - title: The name of the span - description: Name that describes what operation this span represents - start_time: - type: string - format: date-time - title: The start time of the span - description: The time the span started - end_time: - anyOf: - - type: string - format: date-time - - type: 'null' - title: The end time of the span - description: The time the span ended - input: - anyOf: - - additionalProperties: true - type: object - - items: - additionalProperties: true - type: object - type: array - - type: 'null' - title: The input data for the span - description: Input parameters or data for the operation - output: - anyOf: - - additionalProperties: true - type: object - - items: - additionalProperties: true - type: object - type: array - - type: 'null' - title: The output data from the span - description: Output data resulting from the operation - data: - anyOf: - - additionalProperties: true - type: object - - items: - additionalProperties: true - type: object - type: array - - type: 'null' - title: Additional data associated with the span - description: Any additional metadata or context for the span - type: object - required: - - id - - trace_id - - name - - start_time - title: Span State: properties: task_id: @@ -7445,81 +7132,6 @@ components: type: object title: UpdateAgentTaskTrackerRequest description: Request model for updating an agent task tracker. - UpdateSpanRequest: - properties: - trace_id: - anyOf: - - type: string - - type: 'null' - title: The trace ID for this span - description: Unique identifier for the trace this span belongs to - task_id: - anyOf: - - type: string - - type: 'null' - title: The task ID this span is associated with - description: ID of the task this span belongs to - parent_id: - anyOf: - - type: string - - type: 'null' - title: The parent span ID if this is a child span - description: ID of the parent span if this is a child span in a trace - name: - anyOf: - - type: string - - type: 'null' - title: The name of the span - description: Name that describes what operation this span represents - start_time: - anyOf: - - type: string - format: date-time - - type: 'null' - title: The start time of the span - description: The time the span started - end_time: - anyOf: - - type: string - format: date-time - - type: 'null' - title: The end time of the span - description: The time the span ended - input: - anyOf: - - additionalProperties: true - type: object - - items: - additionalProperties: true - type: object - type: array - - type: 'null' - title: The input data for the span - description: Input parameters or data for the operation - output: - anyOf: - - additionalProperties: true - type: object - - items: - additionalProperties: true - type: object - type: array - - type: 'null' - title: The output data from the span - description: Output data resulting from the operation - data: - anyOf: - - additionalProperties: true - type: object - - items: - additionalProperties: true - type: object - type: array - - type: 'null' - title: Additional data associated with the span - description: Any additional metadata or context for the span - type: object - title: UpdateSpanRequest UpdateStateRequest: properties: state: diff --git a/agentex/src/adapters/orm.py b/agentex/src/adapters/orm.py index a53718f8..6a00a7de 100644 --- a/agentex/src/adapters/orm.py +++ b/agentex/src/adapters/orm.py @@ -152,32 +152,6 @@ class AgentTaskTrackerORM(BaseORM): ) -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) - end_time = Column(DateTime(timezone=True), nullable=True) - input = Column(JSON, nullable=True) - output = Column(JSON, nullable=True) - data = Column(JSON, nullable=True) - - # Indexes for efficient querying - __table_args__ = ( - # Index for filtering spans by trace_id - Index("ix_spans_trace_id", "trace_id"), - # Composite index for filtering by trace_id and ordering by start_time - 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"), - ) - - class AgentAPIKeyORM(BaseORM): __tablename__ = "agent_api_keys" id = Column(String, primary_key=True, default=orm_id) diff --git a/agentex/src/api/app.py b/agentex/src/api/app.py index 1d197d34..3be0008b 100644 --- a/agentex/src/api/app.py +++ b/agentex/src/api/app.py @@ -41,7 +41,6 @@ linear, messages, slack, - spans, states, task_retention, tasks, @@ -200,7 +199,6 @@ async def handle_unexpected(request, exc): fastapi_app.include_router(agents.router) fastapi_app.include_router(tasks.router) fastapi_app.include_router(messages.router) -fastapi_app.include_router(spans.router) fastapi_app.include_router(states.router) fastapi_app.include_router(events.router) fastapi_app.include_router(slack.router) diff --git a/agentex/src/api/routes/spans.py b/agentex/src/api/routes/spans.py deleted file mode 100644 index 888b71ac..00000000 --- a/agentex/src/api/routes/spans.py +++ /dev/null @@ -1,103 +0,0 @@ -from fastapi import APIRouter, Query - -from src.api.schemas.spans import CreateSpanRequest, Span, UpdateSpanRequest -from src.domain.use_cases.spans_use_case import DSpanUseCase -from src.utils.logging import make_logger - -logger = make_logger(__name__) - -router = APIRouter(prefix="/spans", tags=["Spans"]) - - -@router.post( - "", - response_model=Span, -) -async def create_span( - request: CreateSpanRequest, - span_use_case: DSpanUseCase, -) -> Span: - """ - Create a new span with the provided parameters - """ - 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, - end_time=request.end_time, - input_data=request.input, - output_data=request.output, - data=request.data, - ) - - -@router.patch( - "/{span_id}", - response_model=Span, -) -async def partial_update_span( - span_id: str, - request: UpdateSpanRequest, - span_use_case: DSpanUseCase, -) -> Span: - """ - Update a span with the provided output data and mark it as complete - """ - 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, - end_time=request.end_time, - input_data=request.input, - output_data=request.output, - data=request.data, - ) - - -@router.get( - "/{span_id}", - response_model=Span, -) -async def get_span( - span_id: str, - span_use_case: DSpanUseCase, -) -> Span: - """ - Get a span by ID - """ - span = await span_use_case.get(span_id=span_id) - return span - - -@router.get( - "", - response_model=list[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 spans, optionally filtered by trace_id and/or task_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, - order_direction=order_direction, - ) - return [Span.model_validate(span) for span in spans] diff --git a/agentex/src/api/schemas/spans.py b/agentex/src/api/schemas/spans.py deleted file mode 100644 index 6f493e7b..00000000 --- a/agentex/src/api/schemas/spans.py +++ /dev/null @@ -1,108 +0,0 @@ -from datetime import datetime -from typing import Any - -from pydantic import Field - -from src.utils.model_utils import BaseModel - - -class CreateSpanRequest(BaseModel): - id: str | None = Field( - None, - title="Unique Span ID", - description="Unique identifier for the span. If not provided, an ID will be generated.", - ) - trace_id: str = Field( - ..., - 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", - description="ID of the parent span if this is a child span in a trace", - ) - name: str = Field( - ..., - title="The name of the span", - description="Name that describes what operation this span represents", - ) - start_time: datetime = Field( - ..., title="The start time of the span", description="The time the span started" - ) - end_time: datetime | None = Field( - None, title="The end time of the span", description="The time the span ended" - ) - input: dict[str, Any] | list[dict[str, Any]] | None = Field( - None, - title="The input data for the span", - description="Input parameters or data for the operation", - ) - output: dict[str, Any] | list[dict[str, Any]] | None = Field( - None, - title="The output data from the span", - description="Output data resulting from the operation", - ) - data: dict[str, Any] | list[dict[str, Any]] | None = Field( - None, - title="Additional data associated with the span", - description="Any additional metadata or context for the span", - ) - - -class UpdateSpanRequest(BaseModel): - trace_id: str | None = Field( - None, - 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", - description="ID of the parent span if this is a child span in a trace", - ) - name: str | None = Field( - None, - title="The name of the span", - description="Name that describes what operation this span represents", - ) - start_time: datetime | None = Field( - None, - title="The start time of the span", - description="The time the span started", - ) - end_time: datetime | None = Field( - None, title="The end time of the span", description="The time the span ended" - ) - input: dict[str, Any] | list[dict[str, Any]] | None = Field( - None, - title="The input data for the span", - description="Input parameters or data for the operation", - ) - output: dict[str, Any] | list[dict[str, Any]] | None = Field( - None, - title="The output data from the span", - description="Output data resulting from the operation", - ) - data: dict[str, Any] | list[dict[str, Any]] | None = Field( - None, - title="Additional data associated with the span", - description="Any additional metadata or context for the span", - ) - - -class Span(CreateSpanRequest): - id: str = Field( - ..., - title="Unique Span ID", - ) diff --git a/agentex/src/domain/entities/spans.py b/agentex/src/domain/entities/spans.py deleted file mode 100644 index 76243704..00000000 --- a/agentex/src/domain/entities/spans.py +++ /dev/null @@ -1,49 +0,0 @@ -from datetime import datetime -from typing import Any - -from pydantic import Field - -from src.utils.model_utils import BaseModel - - -class SpanEntity(BaseModel): - id: str = Field( - ..., - title="Unique Span ID", - ) - trace_id: str = Field( - ..., - 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", - ) - name: str = Field( - ..., - title="The name of the span", - ) - start_time: datetime = Field( - ..., - title="The time the span started", - ) - end_time: datetime | None = Field( - None, - title="The time the span ended", - ) - input: dict[str, Any] | list[dict[str, Any]] | None = Field( - None, - title="The input data for the span", - ) - output: dict[str, Any] | list[dict[str, Any]] | None = Field( - None, - title="The output data from the span", - ) - data: dict[str, Any] | list[dict[str, Any]] | None = Field( - None, - title="Additional data associated with the span", - ) diff --git a/agentex/src/domain/repositories/span_repository.py b/agentex/src/domain/repositories/span_repository.py deleted file mode 100644 index fb8e4c2b..00000000 --- a/agentex/src/domain/repositories/span_repository.py +++ /dev/null @@ -1,78 +0,0 @@ -from typing import Annotated, Any - -from fastapi import Depends -from sqlalchemy import or_, select -from src.adapters.crud_store.adapter_postgres import PostgresCRUDRepository -from src.adapters.orm import SpanORM -from src.config.dependencies import ( - DDatabaseAsyncReadOnlySessionMaker, - DDatabaseAsyncReadWriteSessionMaker, -) -from src.domain.entities.spans import SpanEntity -from src.utils.logging import make_logger - -logger = make_logger(__name__) - - -class SpanRepository(PostgresCRUDRepository[SpanORM, SpanEntity]): - def __init__( - self, - async_read_write_session_maker: DDatabaseAsyncReadWriteSessionMaker, - async_read_only_session_maker: DDatabaseAsyncReadOnlySessionMaker, - ): - super().__init__( - async_read_write_session_maker, - async_read_only_session_maker, - SpanORM, - SpanEntity, - ) - - async def list( - self, - filters: dict[str, Any] | None = None, - limit: int | None = None, - page_number: int | None = None, - order_by: str | None = None, - order_direction: str | None = None, - ) -> list[SpanEntity]: - # Default to start_time if no order_by specified - effective_order_by = order_by or "start_time" - - # Filtering by task_id matches both the new task_id column and historical - # rows where the value was stored in trace_id. The task_id column was - # added late in the table's life and the prod backfill is run out-of-band - # rather than via migration (see docs/runbooks/spans-task-id-backfill.md), - # so old rows can have task_id NULL even when they belong to a task. For - # task-scoped spans, trace_id holds the task id, so we OR the two columns - # at read time. Both columns are indexed. - # - # The OR fallback is skipped when task_id is None — applying it would - # expand to (task_id IS NULL OR trace_id IS NULL), which on a large - # spans table where virtually all historical rows have task_id NULL - # would return an enormous, unintended result set. A None task_id - # filter falls through to the parent's normal IS NULL handling. - if filters and filters.get("task_id") is not None: - remaining_filters = {k: v for k, v in filters.items() if k != "task_id"} - task_id_value = filters["task_id"] - query = select(self.orm).where( - or_(SpanORM.task_id == task_id_value, SpanORM.trace_id == task_id_value) - ) - return await super().list( - filters=remaining_filters or None, - query=query, - order_by=effective_order_by, - order_direction=order_direction, - limit=limit, - page_number=page_number, - ) - - return await super().list( - filters=filters, - order_by=effective_order_by, - order_direction=order_direction, - limit=limit, - page_number=page_number, - ) - - -DSpanRepository = Annotated[SpanRepository, Depends(SpanRepository)] diff --git a/agentex/src/domain/use_cases/spans_use_case.py b/agentex/src/domain/use_cases/spans_use_case.py deleted file mode 100644 index 8faff82c..00000000 --- a/agentex/src/domain/use_cases/spans_use_case.py +++ /dev/null @@ -1,140 +0,0 @@ -from datetime import datetime -from typing import Annotated, Any - -from fastapi import Depends - -from src.domain.entities.spans import SpanEntity -from src.domain.repositories.span_repository import DSpanRepository -from src.utils.ids import orm_id -from src.utils.logging import make_logger - -logger = make_logger(__name__) - - -class SpanUseCase: - def __init__(self, span_repository: DSpanRepository): - logger.info("Initializing SpanUseCase") - self.span_repo = span_repository - - async def create( - self, - 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, - input_data: dict[str, Any] | None = None, - output_data: dict[str, Any] | None = None, - data: dict[str, Any] | None = None, - ) -> SpanEntity: - """ - Create a new span with the given parameters - """ - # Generate ID if not provided - if id is None: - id = orm_id() - - span = SpanEntity( - id=id, - trace_id=trace_id, - task_id=task_id, - parent_id=parent_id, - name=name, - start_time=start_time, - end_time=end_time, - input=input_data, - output=output_data, - data=data, - ) - return await self.span_repo.create(span) - - 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, - end_time: datetime | None = None, - input_data: dict[str, Any] | None = None, - output_data: dict[str, Any] | None = None, - data: dict[str, Any] | None = None, - ) -> SpanEntity: - """ - Update an existing span with partial data - """ - # Get the existing span - span = await self.span_repo.get(id=id) - - # Apply partial updates for all fields - 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 - - if parent_id is not None: - span.parent_id = parent_id - - if start_time is not None: - span.start_time = start_time - - if end_time is not None: - span.end_time = end_time - - if input_data is not None: - span.input = input_data - - if output_data is not None: - span.output = output_data - - if data is not None: - # Merge with existing data if present - if span.data: - span.data.update(data) - else: - span.data = data - - return await self.span_repo.update(span) - - async def get(self, span_id: str) -> SpanEntity: - """ - Get a span by ID - """ - return await self.span_repo.get(id=span_id) - - async def list( - self, - 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 spans, optionally filtered by trace_id and/or task_id - """ - 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, - page_number=page_number, - order_by=order_by, - order_direction=order_direction, - ) - - -DSpanUseCase = Annotated[SpanUseCase, Depends(SpanUseCase)] diff --git a/agentex/tests/fixtures/repositories.py b/agentex/tests/fixtures/repositories.py index 5f2b5e26..6f725cc7 100644 --- a/agentex/tests/fixtures/repositories.py +++ b/agentex/tests/fixtures/repositories.py @@ -60,22 +60,6 @@ async def session_maker(): ) -def create_span_repository(postgres_session): - """Factory function to create SpanRepository with given PostgreSQL session""" - from contextlib import asynccontextmanager - - from src.domain.repositories.span_repository import SpanRepository - - @asynccontextmanager - async def session_maker(): - yield postgres_session - - return SpanRepository( - async_read_write_session_maker=session_maker, - async_read_only_session_maker=session_maker, - ) - - def create_task_state_repository(mongodb_database): """Factory function to create TaskStateRepository with given MongoDB database""" from src.domain.repositories.task_state_repository import TaskStateRepository @@ -173,12 +157,6 @@ def event_repository(unit_db_session): return create_event_repository(unit_db_session) -@pytest.fixture -def span_repository(unit_db_session): - """Span repository for unit tests""" - return create_span_repository(unit_db_session) - - @pytest.fixture def task_state_repository(unit_mongodb_database): """Task state repository for unit tests""" diff --git a/agentex/tests/integration/api/spans/test_spans_api.py b/agentex/tests/integration/api/spans/test_spans_api.py deleted file mode 100644 index d606c756..00000000 --- a/agentex/tests/integration/api/spans/test_spans_api.py +++ /dev/null @@ -1,428 +0,0 @@ -""" -Integration tests for span endpoints following FastAPI async testing best practices. -Tests the full HTTP request → FastAPI → response cycle with API-first validation. -""" - -from datetime import UTC, datetime - -import pytest -import pytest_asyncio -from src.domain.entities.agents import ACPType, AgentEntity -from src.domain.entities.spans import SpanEntity -from src.domain.entities.tasks import TaskEntity -from src.utils.ids import orm_id - - -@pytest.mark.asyncio -class TestSpansAPIIntegration: - """Integration tests for span endpoints using API-first validation""" - - @pytest_asyncio.fixture - async def test_agent(self, isolated_repositories): - """Create a test agent for task creation.""" - agent_repo = isolated_repositories["agent_repository"] - return await agent_repo.create( - AgentEntity( - id=orm_id(), - name="spans-test-agent", - description="Agent for span integration tests", - acp_url="http://test:8000", - acp_type=ACPType.SYNC, - ) - ) - - @pytest_asyncio.fixture - async def test_tasks(self, isolated_repositories, test_agent): - """Create test tasks that can be referenced by spans via FK.""" - task_repo = isolated_repositories["task_repository"] - tasks = {} - for name in [ - "task-a", - "task-b", - "task-x", - "task-y", - "task-create", - "task-update", - ]: - task = await task_repo.create( - agent_id=test_agent.id, - task=TaskEntity(id=orm_id(), name=name), - ) - tasks[name] = task - return tasks - - @pytest_asyncio.fixture - async def test_pagination_spans(self, isolated_repositories): - """Create spans for pagination tests""" - span_repo = isolated_repositories["span_repository"] - spans = [] - for i in range(60): - span = SpanEntity( - id=orm_id(), - trace_id=orm_id(), - name=f"test-span-{i}", - start_time=datetime.now(UTC), - ) - spans.append(await span_repo.create(span)) - return spans - - async def test_create_and_retrieve_span_consistency( - self, isolated_client, test_tasks - ): - """Test span creation and validate POST → GET consistency (API-first)""" - task_id = test_tasks["task-create"].id - - # Given - Span creation data - span_data = { - "trace_id": "test-trace-123", - "task_id": task_id, - "name": "test-operation", - "start_time": "2024-01-01T10:00:00Z", - "end_time": "2024-01-01T10:00:05Z", - "input": {"key": "value"}, - "output": {"result": "success"}, - "metadata": {"test": True}, - } - - # When - Create span via POST - create_response = await isolated_client.post("/spans", json=span_data) - - # Then - Should succeed and return created span - assert create_response.status_code == 200 - created_span = create_response.json() - - # Validate response has required fields - assert "id" in created_span - assert created_span["trace_id"] == span_data["trace_id"] - assert created_span["task_id"] == task_id - assert created_span["name"] == span_data["name"] - span_id = created_span["id"] - - # API-first validation: GET the created span - get_response = await isolated_client.get(f"/spans/{span_id}") - assert get_response.status_code == 200 - retrieved_span = get_response.json() - - # Validate POST/GET consistency - assert retrieved_span["id"] == span_id - assert retrieved_span["trace_id"] == span_data["trace_id"] - assert retrieved_span["task_id"] == task_id - assert retrieved_span["name"] == span_data["name"] - assert retrieved_span["input"] == span_data["input"] - assert retrieved_span["output"] == span_data["output"] - - async def test_create_span_without_task_id(self, isolated_client): - """Test span creation without task_id (should default to null)""" - span_data = { - "trace_id": "test-trace-no-task", - "name": "test-no-task", - "start_time": "2024-01-01T10:00:00Z", - } - - create_response = await isolated_client.post("/spans", json=span_data) - assert create_response.status_code == 200 - created_span = create_response.json() - assert created_span["task_id"] is None - - async def test_update_span_and_validate_changes(self, isolated_client, test_tasks): - """Test span update and validate PATCH → GET consistency""" - task_id = test_tasks["task-update"].id - - # Given - Create a span first - initial_data = { - "trace_id": "update-trace-456", - "name": "initial-name", - "start_time": "2024-01-01T10:00:00Z", - } - create_response = await isolated_client.post("/spans", json=initial_data) - assert create_response.status_code == 200 - span_id = create_response.json()["id"] - - # When - Update the span including task_id - update_data = { - "name": "updated-name", - "task_id": task_id, - "parent_id": "parent-id", - "start_time": "2024-01-01T10:10:00Z", - "end_time": "2024-01-01T10:10:05Z", - "input": {"key": "value"}, - "output": {"status": "completed"}, - "data": {"test": True}, - } - patch_response = await isolated_client.patch( - f"/spans/{span_id}", json=update_data - ) - - # Then - Should succeed - assert patch_response.status_code == 200 - - # API-first validation: GET updated span - get_response = await isolated_client.get(f"/spans/{span_id}") - assert get_response.status_code == 200 - updated_span = get_response.json() - - # Validate changes were applied - assert updated_span["name"] == "updated-name" - assert updated_span["task_id"] == task_id - assert updated_span["output"]["status"] == "completed" - assert updated_span["parent_id"] == "parent-id" - assert updated_span["start_time"] == "2024-01-01T10:10:00Z" - assert updated_span["end_time"] == "2024-01-01T10:10:05Z" - assert updated_span["input"] == {"key": "value"} - assert updated_span["data"] == {"test": True} - assert updated_span["trace_id"] == initial_data["trace_id"] # Unchanged - - # We can also update trace ID and add values into metadata - patch_response = await isolated_client.patch( - f"/spans/{span_id}", - json={ - "trace_id": "updated-trace-789", - "data": {"version": "2.0.0"}, - }, - ) - assert patch_response.status_code == 200 - updated_span = patch_response.json() - assert updated_span["name"] == "updated-name" - assert updated_span["task_id"] == task_id # Still set from prior update - assert updated_span["output"]["status"] == "completed" - assert updated_span["parent_id"] == "parent-id" - assert updated_span["start_time"] == "2024-01-01T10:10:00Z" - assert updated_span["end_time"] == "2024-01-01T10:10:05Z" - assert updated_span["input"] == {"key": "value"} - assert updated_span["trace_id"] == "updated-trace-789" - assert updated_span["data"] == {"test": True, "version": "2.0.0"} - - async def test_list_spans_with_trace_id_filtering(self, isolated_client): - """Test list spans endpoint with trace_id filtering""" - # Given - Create spans with different trace_ids - trace_id_1 = "list-trace-001" - trace_id_2 = "list-trace-002" - - span1_data = { - "trace_id": trace_id_1, - "name": "span-1", - "start_time": "2024-01-01T10:00:00Z", - } - span2_data = { - "trace_id": trace_id_2, - "name": "span-2", - "start_time": "2024-01-01T10:00:00Z", - } - - create1 = await isolated_client.post("/spans", json=span1_data) - create2 = await isolated_client.post("/spans", json=span2_data) - assert create1.status_code == 200 - assert create2.status_code == 200 - - all_spans = await isolated_client.get("/spans") - assert all_spans.status_code == 200 - all_spans_data = all_spans.json() - assert isinstance(all_spans_data, list) - assert len(all_spans_data) == 2 - assert all_spans_data[0]["trace_id"] == trace_id_1 - assert all_spans_data[1]["trace_id"] == trace_id_2 - - # When - List spans filtered by trace_id - list_response = await isolated_client.get(f"/spans?trace_id={trace_id_1}") - - # Then - Should return only spans with matching trace_id - assert list_response.status_code == 200 - spans = list_response.json() - assert isinstance(spans, list) - - # Validate filtering worked - for span in spans: - assert span["trace_id"] == trace_id_1 - - async def test_list_spans_with_task_id_filtering(self, isolated_client, test_tasks): - """Test list spans endpoint with task_id filtering""" - task_id_a = test_tasks["task-a"].id - task_id_b = test_tasks["task-b"].id - - for i in range(3): - resp = await isolated_client.post( - "/spans", - json={ - "trace_id": f"trace-task-filter-{i}", - "task_id": task_id_a, - "name": f"span-task-a-{i}", - "start_time": "2024-01-01T10:00:00Z", - }, - ) - assert resp.status_code == 200 - - for i in range(2): - resp = await isolated_client.post( - "/spans", - json={ - "trace_id": f"trace-task-filter-b-{i}", - "task_id": task_id_b, - "name": f"span-task-b-{i}", - "start_time": "2024-01-01T10:00:00Z", - }, - ) - assert resp.status_code == 200 - - # One span with no task_id - resp = await isolated_client.post( - "/spans", - json={ - "trace_id": "trace-no-task", - "name": "span-no-task", - "start_time": "2024-01-01T10:00:00Z", - }, - ) - assert resp.status_code == 200 - - # When - Filter by task_id_a - response = await isolated_client.get(f"/spans?task_id={task_id_a}") - assert response.status_code == 200 - spans = response.json() - assert len(spans) == 3 - for span in spans: - assert span["task_id"] == task_id_a - - # When - Filter by task_id_b - response = await isolated_client.get(f"/spans?task_id={task_id_b}") - assert response.status_code == 200 - spans = response.json() - assert len(spans) == 2 - for span in spans: - assert span["task_id"] == task_id_b - - # When - No filter returns all 6 - response = await isolated_client.get("/spans") - assert response.status_code == 200 - assert len(response.json()) == 6 - - async def test_list_spans_with_combined_trace_and_task_filtering( - self, isolated_client, test_tasks - ): - """Test list spans with both trace_id and task_id filters""" - shared_trace = "combined-trace" - task_id_x = test_tasks["task-x"].id - task_id_y = test_tasks["task-y"].id - - await isolated_client.post( - "/spans", - json={ - "trace_id": shared_trace, - "task_id": task_id_x, - "name": "span-match", - "start_time": "2024-01-01T10:00:00Z", - }, - ) - await isolated_client.post( - "/spans", - json={ - "trace_id": shared_trace, - "task_id": task_id_y, - "name": "span-same-trace-diff-task", - "start_time": "2024-01-01T10:00:00Z", - }, - ) - await isolated_client.post( - "/spans", - json={ - "trace_id": "other-trace", - "task_id": task_id_x, - "name": "span-diff-trace-same-task", - "start_time": "2024-01-01T10:00:00Z", - }, - ) - - # When - Filter by both trace_id and task_id - response = await isolated_client.get( - f"/spans?trace_id={shared_trace}&task_id={task_id_x}" - ) - assert response.status_code == 200 - spans = response.json() - assert len(spans) == 1 - assert spans[0]["name"] == "span-match" - assert spans[0]["trace_id"] == shared_trace - assert spans[0]["task_id"] == task_id_x - - async def test_get_span_non_existent(self, isolated_client): - """Test getting a non-existent span returns 404""" - # When - Get a non-existent span - response = await isolated_client.get("/spans/non-existent-id") - # Then - Should return 404 - assert response.status_code == 404 - - async def test_list_spans_pagination(self, isolated_client, test_pagination_spans): - """Test GET /spans/ endpoint with pagination.""" - # Given - A span record exists - # (created by test_pagination_spans fixture) - - # When - List all spans with pagination - response = await isolated_client.get("/spans") - assert response.status_code == 200 - response_data = response.json() - # Default limit if none specified - assert len(response_data) == 50 - - page_number = 1 - paginated_spans = [] - while True: - response = await isolated_client.get( - "/spans", params={"limit": 7, "page_number": page_number} - ) - assert response.status_code == 200 - spans_data = response.json() - paginated_spans.extend(spans_data) - if len(spans_data) < 1: - break - page_number += 1 - assert len(paginated_spans) == len(test_pagination_spans) - assert {(d["id"], d["name"]) for d in paginated_spans} == { - (d.id, d.name) for d in test_pagination_spans - } - - async def test_list_spans_with_order_by(self, isolated_client): - """Test that list spans endpoint supports order_by parameter""" - # Given - Create multiple spans with different start times - trace_id = "order-by-trace" - spans_data = [ - { - "trace_id": trace_id, - "name": f"order-span-{i}", - "start_time": f"2024-01-01T10:0{i}:00Z", - } - for i in range(3) - ] - - for span_data in spans_data: - response = await isolated_client.post("/spans", json=span_data) - assert response.status_code == 200 - - # When - Request spans with order_by=start_time and order_direction=asc - response_asc = await isolated_client.get( - f"/spans?trace_id={trace_id}&order_by=start_time&order_direction=asc" - ) - - # Then - Should return spans in ascending order - assert response_asc.status_code == 200 - spans_asc = response_asc.json() - assert len(spans_asc) == 3 - - # Verify ascending order - for i in range(len(spans_asc) - 1): - assert spans_asc[i]["start_time"] <= spans_asc[i + 1]["start_time"] - - # When - Request spans with order_by=start_time and order_direction=desc - response_desc = await isolated_client.get( - f"/spans?trace_id={trace_id}&order_by=start_time&order_direction=desc" - ) - - # Then - Should return spans in descending order - assert response_desc.status_code == 200 - spans_desc = response_desc.json() - assert len(spans_desc) == 3 - - # Verify descending order - for i in range(len(spans_desc) - 1): - assert spans_desc[i]["start_time"] >= spans_desc[i + 1]["start_time"] - - # Verify the order is actually reversed - assert spans_asc[0]["id"] == spans_desc[-1]["id"] - assert spans_asc[-1]["id"] == spans_desc[0]["id"] diff --git a/agentex/tests/integration/fixtures/integration_client.py b/agentex/tests/integration/fixtures/integration_client.py index 1e173506..29f76430 100644 --- a/agentex/tests/integration/fixtures/integration_client.py +++ b/agentex/tests/integration/fixtures/integration_client.py @@ -267,7 +267,6 @@ async def __aenter__(self): ) from src.domain.repositories.deployment_repository import DeploymentRepository from src.domain.repositories.event_repository import EventRepository - from src.domain.repositories.span_repository import SpanRepository from src.domain.repositories.task_message_repository import TaskMessageRepository from src.domain.repositories.task_repository import TaskRepository from src.domain.repositories.task_state_repository import TaskStateRepository @@ -305,9 +304,6 @@ def __init__(self, redis_url): "event_repository": EventRepository( async_rw_session_factory, async_ro_session_factory ), - "span_repository": SpanRepository( - async_rw_session_factory, async_ro_session_factory - ), "agent_task_tracker_repository": AgentTaskTrackerRepository( async_rw_session_factory, async_ro_session_factory ), @@ -383,7 +379,6 @@ async def isolated_integration_app( ) from src.domain.use_cases.events_use_case import EventUseCase from src.domain.use_cases.messages_use_case import MessagesUseCase - from src.domain.use_cases.spans_use_case import SpanUseCase from src.domain.use_cases.states_use_case import StatesUseCase from src.domain.use_cases.task_retention_use_case import TaskRetentionUseCase from src.domain.use_cases.tasks_use_case import TasksUseCase @@ -424,9 +419,6 @@ def create_deployment_history_use_case(): def create_events_use_case(): return EventUseCase(event_repository=isolated_repositories["event_repository"]) - def create_spans_use_case(): - return SpanUseCase(span_repository=isolated_repositories["span_repository"]) - def create_states_use_case(): return StatesUseCase( task_state_repository=isolated_repositories["task_state_repository"] @@ -519,7 +511,6 @@ def create_task_retention_use_case(): ) from src.domain.repositories.deployment_repository import DeploymentRepository from src.domain.repositories.event_repository import EventRepository - from src.domain.repositories.span_repository import SpanRepository from src.domain.repositories.task_message_repository import TaskMessageRepository from src.domain.repositories.task_repository import TaskRepository from src.domain.repositories.task_state_repository import TaskStateRepository @@ -540,7 +531,6 @@ def create_task_retention_use_case(): CheckpointsUseCase: create_checkpoints_use_case, AgentsUseCase: create_agents_use_case, EventUseCase: create_events_use_case, - SpanUseCase: create_spans_use_case, StatesUseCase: create_states_use_case, AgentTaskTrackerUseCase: create_agent_task_tracker_use_case, TasksUseCase: create_tasks_use_case, @@ -562,7 +552,6 @@ def create_task_retention_use_case(): ], TaskRepository: lambda: isolated_repositories["task_repository"], EventRepository: lambda: isolated_repositories["event_repository"], - SpanRepository: lambda: isolated_repositories["span_repository"], AgentTaskTrackerRepository: lambda: isolated_repositories[ "agent_task_tracker_repository" ], diff --git a/agentex/tests/unit/repositories/test_span_repository.py b/agentex/tests/unit/repositories/test_span_repository.py deleted file mode 100644 index 0276f391..00000000 --- a/agentex/tests/unit/repositories/test_span_repository.py +++ /dev/null @@ -1,393 +0,0 @@ -import asyncio -import os - -# Import the repository and entities we need to test -import sys -from datetime import UTC, datetime - -import pytest -from sqlalchemy import text -from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine - -sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..", "..", "src")) - -from adapters.orm import BaseORM, TaskORM -from domain.entities.spans import SpanEntity -from domain.repositories.span_repository import SpanRepository -from utils.ids import orm_id - - -@pytest.mark.asyncio -@pytest.mark.unit -async def test_span_repository_crud_operations(postgres_url): - """Test SpanRepository CRUD operations with JSON fields and time ordering""" - - # URL conversion for SQLAlchemy async - sqlalchemy_asyncpg_url = postgres_url.replace( - "postgresql+psycopg2://", "postgresql+asyncpg://" - ) - - # Wait for database readiness - for attempt in range(10): - try: - engine = create_async_engine(sqlalchemy_asyncpg_url, echo=True) - async with engine.begin() as conn: - await conn.run_sync(BaseORM.metadata.create_all) - await conn.execute(text("SELECT 1")) - break - except Exception as e: - if attempt < 9: - print( - f"Database not ready (attempt {attempt + 1}), retrying... Error: {e}" - ) - await asyncio.sleep(2) - continue - raise - - # Create async session maker and repository - async_session_maker = async_sessionmaker(engine, expire_on_commit=False) - span_repo = SpanRepository(async_session_maker, async_session_maker) - - # Create a task row to satisfy the FK constraint on spans.task_id - task_id = orm_id() - async with async_session_maker() as session: - session.add(TaskORM(id=task_id, name="test-task")) - await session.commit() - - # Test CREATE operation with JSON fields - now = datetime.now(UTC) - span_id = orm_id() - trace_id = orm_id() - - span = SpanEntity( - id=span_id, - trace_id=trace_id, - task_id=task_id, - parent_id=None, - name="test-span-operation", - start_time=now, - end_time=None, # Still running - input={"operation": "test", "parameters": {"limit": 10}}, - output=None, # Not finished yet - data={"metadata": {"version": "1.0", "environment": "test"}}, - ) - - created_span = await span_repo.create(span) - assert created_span.id == span_id - assert created_span.trace_id == trace_id - assert created_span.task_id == task_id - assert created_span.name == "test-span-operation" - assert created_span.input["operation"] == "test" - assert created_span.data["metadata"]["version"] == "1.0" - print("✅ CREATE operation successful with JSON fields") - - # Test UPDATE operation (complete the span) - end_time = datetime.now(UTC) - updated_span = SpanEntity( - id=span_id, - trace_id=trace_id, - task_id=task_id, - parent_id=None, - name="test-span-operation", - start_time=now, - end_time=end_time, - input={"operation": "test", "parameters": {"limit": 10}}, - output={"result": "success", "processed": 5}, - data={ - "metadata": {"version": "1.0", "environment": "test", "duration_ms": 150} - }, - ) - - result_span = await span_repo.update(updated_span) - assert result_span.end_time is not None - assert result_span.output["result"] == "success" - assert result_span.data["metadata"]["duration_ms"] == 150 - print("✅ UPDATE operation successful") - - # Test GET operation - retrieved_span = await span_repo.get(id=span_id) - assert retrieved_span.id == span_id - assert retrieved_span.output["processed"] == 5 - print("✅ GET operation successful") - - # Create a child span to test ordering - child_span_id = orm_id() - child_start_time = datetime.now(UTC) - - child_span = SpanEntity( - id=child_span_id, - trace_id=trace_id, - task_id=task_id, - parent_id=span_id, # Child of the first span - name="child-span-operation", - start_time=child_start_time, - end_time=None, - input={"operation": "child_task"}, - output=None, - data={"parent_context": "inherited"}, - ) - - await span_repo.create(child_span) - print("✅ Child span created for ordering test") - - # Test LIST operation with time-based ordering - all_spans = await span_repo.list() - assert len(all_spans) >= 2 - span_ids = [s.id for s in all_spans] - assert span_id in span_ids - assert child_span_id in span_ids - - # Should be ordered by start_time (parent should come first) - parent_index = next(i for i, s in enumerate(all_spans) if s.id == span_id) - child_index = next(i for i, s in enumerate(all_spans) if s.id == child_span_id) - assert parent_index < child_index, "Parent span should come before child span" - print("✅ LIST operation successful with time ordering") - - # Test DELETE operation - await span_repo.delete(child_span_id) - - # Verify deletion - all_spans_after_delete = await span_repo.list() - span_ids_after_delete = [s.id for s in all_spans_after_delete] - assert child_span_id not in span_ids_after_delete - assert span_id in span_ids_after_delete - print("✅ DELETE operation successful") - - print("✅ Test isolation provided by session-scoped PostgreSQL container") - print("🎉 ALL SPAN REPOSITORY TESTS PASSED!") - - -@pytest.mark.asyncio -@pytest.mark.unit -async def test_span_task_id_set_null_on_task_delete(postgres_url): - """Deleting a referenced task should null out spans.task_id, not fail with FK violation.""" - - sqlalchemy_asyncpg_url = postgres_url.replace( - "postgresql+psycopg2://", "postgresql+asyncpg://" - ) - - engine = create_async_engine(sqlalchemy_asyncpg_url, echo=False) - async with engine.begin() as conn: - await conn.run_sync(BaseORM.metadata.create_all) - - async_session_maker = async_sessionmaker(engine, expire_on_commit=False) - span_repo = SpanRepository(async_session_maker, async_session_maker) - - # Seed a task and a span referencing it - task_id = orm_id() - span_id = orm_id() - async with async_session_maker() as session: - session.add(TaskORM(id=task_id, name="task-to-delete")) - await session.commit() - - await span_repo.create( - SpanEntity( - id=span_id, - trace_id=orm_id(), - task_id=task_id, - parent_id=None, - name="span-with-task-fk", - start_time=datetime.now(UTC), - ) - ) - - # Delete the task — should succeed, not raise a FK violation - async with async_session_maker() as session: - task = await session.get(TaskORM, task_id) - await session.delete(task) - await session.commit() - - # Span should survive with task_id set to NULL - retrieved = await span_repo.get(id=span_id) - assert retrieved is not None - assert retrieved.task_id is None - - -@pytest.mark.asyncio -@pytest.mark.unit -async def test_list_by_task_id_falls_back_to_trace_id(postgres_url): - """Listing by task_id should also match historical rows that have the value - in trace_id but a NULL task_id (pre-backfill state).""" - - sqlalchemy_asyncpg_url = postgres_url.replace( - "postgresql+psycopg2://", "postgresql+asyncpg://" - ) - - engine = create_async_engine(sqlalchemy_asyncpg_url, echo=False) - async with engine.begin() as conn: - await conn.run_sync(BaseORM.metadata.create_all) - - async_session_maker = async_sessionmaker(engine, expire_on_commit=False) - span_repo = SpanRepository(async_session_maker, async_session_maker) - - task_id = orm_id() - async with async_session_maker() as session: - session.add(TaskORM(id=task_id, name="task-or-fallback")) - await session.commit() - - # Historical span: task_id NULL, trace_id holds the task id (pre-backfill) - historical_id = orm_id() - await span_repo.create( - SpanEntity( - id=historical_id, - trace_id=task_id, - task_id=None, - parent_id=None, - name="historical", - start_time=datetime.now(UTC), - ) - ) - - # New-style span: task_id set explicitly, trace_id is unrelated - new_id = orm_id() - await span_repo.create( - SpanEntity( - id=new_id, - trace_id=orm_id(), - task_id=task_id, - parent_id=None, - name="new-style", - start_time=datetime.now(UTC), - ) - ) - - # Unrelated span: should not match - unrelated_id = orm_id() - await span_repo.create( - SpanEntity( - id=unrelated_id, - trace_id=orm_id(), - task_id=None, - parent_id=None, - name="unrelated", - start_time=datetime.now(UTC), - ) - ) - - matched = await span_repo.list(filters={"task_id": task_id}) - matched_ids = {s.id for s in matched} - assert historical_id in matched_ids - assert new_id in matched_ids - assert unrelated_id not in matched_ids - - -@pytest.mark.asyncio -@pytest.mark.unit -async def test_list_with_none_task_id_does_not_or_on_trace_id(postgres_url): - """A None task_id filter must NOT trigger the trace_id OR fallback, - otherwise the predicate expands to (task_id IS NULL OR trace_id IS NULL) - and returns nearly every row on a partially backfilled table.""" - - sqlalchemy_asyncpg_url = postgres_url.replace( - "postgresql+psycopg2://", "postgresql+asyncpg://" - ) - - engine = create_async_engine(sqlalchemy_asyncpg_url, echo=False) - async with engine.begin() as conn: - await conn.run_sync(BaseORM.metadata.create_all) - - async_session_maker = async_sessionmaker(engine, expire_on_commit=False) - span_repo = SpanRepository(async_session_maker, async_session_maker) - - # Span with non-null trace_id and null task_id (pre-backfill historical row). - # If the OR fallback were applied to a None task_id filter, this row would - # incorrectly match because trace_id IS NOT NULL but task_id IS NULL — the - # generated predicate (task_id IS NULL OR trace_id IS NULL) would be true. - # Wait: with this row trace_id IS NOT NULL, so trace_id IS NULL is false. - # The bug is the *other* direction: task_id IS NULL is true → row matches - # the (incorrectly) ORed predicate, even though the caller asked for - # task_id IS NULL only. - historical_id = orm_id() - await span_repo.create( - SpanEntity( - id=historical_id, - trace_id=orm_id(), - task_id=None, - parent_id=None, - name="historical-null-task", - start_time=datetime.now(UTC), - ) - ) - - # Span where both task_id and trace_id are non-null. This row should NOT - # match a "task_id IS NULL" filter under either correct or incorrect - # behavior — included as a sanity check. - populated_id = orm_id() - task_id = orm_id() - async with async_session_maker() as session: - session.add(TaskORM(id=task_id, name="task-for-populated-span")) - await session.commit() - await span_repo.create( - SpanEntity( - id=populated_id, - trace_id=orm_id(), - task_id=task_id, - parent_id=None, - name="populated", - start_time=datetime.now(UTC), - ) - ) - - # Filtering by task_id=None should match only the historical (NULL task_id) - # row, NOT trigger the OR fallback against trace_id. - matched = await span_repo.list(filters={"task_id": None}) - matched_ids = {s.id for s in matched} - assert historical_id in matched_ids - assert populated_id not in matched_ids - - -@pytest.mark.asyncio -@pytest.mark.unit -async def test_list_combines_task_id_and_trace_id_filters(postgres_url): - """When both task_id and trace_id are passed, the trace_id filter still - applies on top of the task_id OR-fallback (logical AND between filters).""" - - sqlalchemy_asyncpg_url = postgres_url.replace( - "postgresql+psycopg2://", "postgresql+asyncpg://" - ) - - engine = create_async_engine(sqlalchemy_asyncpg_url, echo=False) - async with engine.begin() as conn: - await conn.run_sync(BaseORM.metadata.create_all) - - async_session_maker = async_sessionmaker(engine, expire_on_commit=False) - span_repo = SpanRepository(async_session_maker, async_session_maker) - - task_id = orm_id() - other_trace_id = orm_id() - async with async_session_maker() as session: - session.add(TaskORM(id=task_id, name="task-and")) - await session.commit() - - # Span matches task_id but not the requested trace_id — should be excluded - excluded_id = orm_id() - await span_repo.create( - SpanEntity( - id=excluded_id, - trace_id=orm_id(), - task_id=task_id, - parent_id=None, - name="excluded", - start_time=datetime.now(UTC), - ) - ) - - # Span matches both — should be included - included_id = orm_id() - await span_repo.create( - SpanEntity( - id=included_id, - trace_id=other_trace_id, - task_id=task_id, - parent_id=None, - name="included", - start_time=datetime.now(UTC), - ) - ) - - matched = await span_repo.list( - filters={"task_id": task_id, "trace_id": other_trace_id} - ) - matched_ids = {s.id for s in matched} - assert included_id in matched_ids - assert excluded_id not in matched_ids From 2e9516160631432b9efed04f0a1645ce97fd1302 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Fri, 11 Sep 2026 09:45:04 -0400 Subject: [PATCH 2/4] fix(agentex-ui): scope the spans cache to the account and abort the platform search with the browser The traces query keyed on the task id alone, so a deep link that switched the account while keeping the task could show the other account's cached spans. The BFF route also let a cancelled browser request keep the platform search running, unlike the Agentex proxy next to it. Co-Authored-By: Claude Fable 5.1 --- .../api/traces/[traceId]/spans/route.test.ts | 25 +++++++++++++++++-- .../app/api/traces/[traceId]/spans/route.ts | 21 ++++++++++++---- agentex-ui/hooks/use-spans.test.tsx | 10 +++++++- agentex-ui/hooks/use-spans.ts | 9 ++++--- 4 files changed, 54 insertions(+), 11 deletions(-) diff --git a/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts b/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts index 40a1d9da..d37e0c6f 100644 --- a/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts +++ b/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts @@ -16,8 +16,8 @@ vi.mock('@/app/api/_lib/bff', () => ({ applyBffCredentials: bff.applyBffCredentials, })); -function call(traceId: string) { - return GET(new Request(`http://ui.local/api/traces/${traceId}/spans`), { +function call(traceId: string, init?: RequestInit) { + return GET(new Request(`http://ui.local/api/traces/${traceId}/spans`, init), { params: Promise.resolve({ traceId }), }); } @@ -52,6 +52,27 @@ describe('GET /api/traces/[traceId]/spans', () => { expect(new Headers(init.headers).get('authorization')).toBe( 'Bearer server-side' ); + expect(init.signal).toBeInstanceOf(AbortSignal); + }); + + it('answers 499 when the browser aborts before the platform replies', async () => { + const controller = new AbortController(); + vi.stubGlobal( + 'fetch', + vi.fn((_url: string, init: RequestInit) => { + const signal = init.signal as AbortSignal; + return new Promise((_resolve, reject) => { + if (signal.aborted) reject(signal.reason); + signal.addEventListener('abort', () => reject(signal.reason)); + }); + }) + ); + + const pending = call('t1', { signal: controller.signal }); + controller.abort(); + const res = await pending; + + expect(res.status).toBe(499); }); it('passes the upstream status through', async () => { diff --git a/agentex-ui/app/api/traces/[traceId]/spans/route.ts b/agentex-ui/app/api/traces/[traceId]/spans/route.ts index 2300a035..3f14e532 100644 --- a/agentex-ui/app/api/traces/[traceId]/spans/route.ts +++ b/agentex-ui/app/api/traces/[traceId]/spans/route.ts @@ -35,11 +35,22 @@ export async function GET( sort_by: 'start_timestamp', sort_order: 'asc', }); - const upstream = await fetch(`${SGP_BASE_URL}/v5/spans/search?${query}`, { - method: 'POST', - headers, - body: JSON.stringify({ trace_ids: [traceId] }), - }); + let upstream: Response; + try { + upstream = await fetch(`${SGP_BASE_URL}/v5/spans/search?${query}`, { + method: 'POST', + headers, + body: JSON.stringify({ trace_ids: [traceId] }), + signal: request.signal, + }); + } catch (error) { + // A browser that navigated away aborts the request, and the abort reason is what the + // fetch rejects with, so the same identity check the Agentex proxy uses applies here. + if (error === request.signal.reason) { + return new Response(null, { status: 499 }); + } + throw error; + } return new Response(upstream.body, { status: upstream.status, headers: { 'content-type': 'application/json' }, diff --git a/agentex-ui/hooks/use-spans.test.tsx b/agentex-ui/hooks/use-spans.test.tsx index 0b01f6bb..6c5a3949 100644 --- a/agentex-ui/hooks/use-spans.test.tsx +++ b/agentex-ui/hooks/use-spans.test.tsx @@ -4,7 +4,7 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; import { renderHook, waitFor } from '@testing-library/react'; import { afterEach, describe, expect, it, vi } from 'vitest'; -import { useSpans } from './use-spans'; +import { spansKeys, useSpans } from './use-spans'; vi.mock('@/hooks/use-safe-search-params', () => ({ useSafeSearchParams: () => ({ sgpAccountID: 'acct-1' }), @@ -37,6 +37,14 @@ const span = { end_timestamp: null, }; +describe('spansKeys', () => { + it('scopes a task query to the selected account', () => { + expect(spansKeys.byTaskId('task-1', 'acct-1')).not.toEqual( + spansKeys.byTaskId('task-1', 'acct-2') + ); + }); +}); + describe('useSpans', () => { afterEach(() => { vi.unstubAllGlobals(); diff --git a/agentex-ui/hooks/use-spans.ts b/agentex-ui/hooks/use-spans.ts index 15a43667..446ea90e 100644 --- a/agentex-ui/hooks/use-spans.ts +++ b/agentex-ui/hooks/use-spans.ts @@ -6,8 +6,11 @@ import { useSafeSearchParams } from '@/hooks/use-safe-search-params'; export const spansKeys = { all: ['spans'] as const, - byTaskId: (taskId: string | null) => - taskId ? ([...spansKeys.all, 'task', taskId] as const) : spansKeys.all, + // The account scopes the platform read, so it scopes the cache entry too. + byTaskId: (taskId: string | null, accountId: string | null) => + taskId + ? ([...spansKeys.all, 'task', taskId, accountId ?? ''] as const) + : spansKeys.all, }; /** A platform span as the traces BFF route returns it. */ @@ -54,7 +57,7 @@ export function useSpans(taskId: string | null): UseSpansState { const { sgpAccountID } = useSafeSearchParams(); const { data, isLoading, error } = useQuery({ - queryKey: spansKeys.byTaskId(taskId), + queryKey: spansKeys.byTaskId(taskId, sgpAccountID), queryFn: async ({ signal }): Promise => { if (!taskId) { return { items: [], hasMore: false }; From cfe5f6c3d08da2e9d9ba96b3469214a6499b9115 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Fri, 11 Sep 2026 10:03:44 -0400 Subject: [PATCH 3/4] fix(agentex-ui): anchor the platform span search on the task and accept short pages The platform defaults an omitted search window to the last 90 days and refuses a page that exceeds its byte budget unless the caller opts into short pages, so an old or a heavy task read as empty or as an error. The sidebar now passes the task's creation time, which the proxy turns into a 90-day window, and opts into short pages. The header no longer loads spans to learn the trace id, which is the task id by construction. The truncation notice names the Investigate link only when it renders, and CI now runs the UI unit tests. Co-Authored-By: Claude Fable 5.1 --- .../workflows/agentex-ui-lint-typecheck.yml | 4 ++ agentex-ui/README.md | 2 +- .../api/traces/[traceId]/spans/route.test.ts | 67 +++++++++++++++++-- .../app/api/traces/[traceId]/spans/route.ts | 28 +++++++- .../components/task-header/task-header.tsx | 5 +- .../traces-sidebar/traces-sidebar.tsx | 17 ++++- agentex-ui/hooks/use-spans.test.tsx | 47 ++++++++++--- agentex-ui/hooks/use-spans.ts | 32 +++++++-- 8 files changed, 173 insertions(+), 29 deletions(-) diff --git a/.github/workflows/agentex-ui-lint-typecheck.yml b/.github/workflows/agentex-ui-lint-typecheck.yml index a5db0e44..cfc346e5 100644 --- a/.github/workflows/agentex-ui-lint-typecheck.yml +++ b/.github/workflows/agentex-ui-lint-typecheck.yml @@ -40,3 +40,7 @@ jobs: - name: Run lint run: npm run lint working-directory: ./agentex-ui + + - name: Run unit tests + run: npm run test:run + working-directory: ./agentex-ui diff --git a/agentex-ui/README.md b/agentex-ui/README.md index 90e9c0db..707ceed9 100644 --- a/agentex-ui/README.md +++ b/agentex-ui/README.md @@ -33,7 +33,7 @@ A modern web interface for building, testing, and monitoring intelligent agents. ### Observability -- **Execution Traces** - View a task's spans from Scale GenAI Platform (needs `SGP_API_URL`) +- **Execution Traces** - View a task's spans from Scale GenAI Platform (needs `SGP_API_URL` or `NEXT_PUBLIC_SGP_APP_URL`) - **Span Visualization** - Hierarchical view of execution flow - **Performance Metrics** - Timing and duration information for each execution step - **Error Tracking** - Detailed error information when tasks fail diff --git a/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts b/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts index d37e0c6f..d564b7f4 100644 --- a/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts +++ b/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts @@ -16,10 +16,15 @@ vi.mock('@/app/api/_lib/bff', () => ({ applyBffCredentials: bff.applyBffCredentials, })); -function call(traceId: string, init?: RequestInit) { - return GET(new Request(`http://ui.local/api/traces/${traceId}/spans`, init), { - params: Promise.resolve({ traceId }), - }); +function call(traceId: string, init?: RequestInit, search = '') { + return GET( + new Request(`http://ui.local/api/traces/${traceId}/spans${search}`, init), + { params: Promise.resolve({ traceId }) } + ); +} + +function upstreamURL(fetchMock: ReturnType) { + return new URL(fetchMock.mock.calls[0]![0] as string); } describe('GET /api/traces/[traceId]/spans', () => { @@ -45,7 +50,7 @@ describe('GET /api/traces/[traceId]/spans', () => { expect(bff.applyBffCredentials).toHaveBeenCalledTimes(1); const [url, init] = fetchMock.mock.calls[0]!; expect(url).toBe( - 'https://sgp.example/api/v5/spans/search?limit=100&sort_by=start_timestamp&sort_order=asc' + 'https://sgp.example/api/v5/spans/search?limit=100&sort_by=start_timestamp&sort_order=asc&allow_short_pages=true' ); expect(init.method).toBe('POST'); expect(JSON.parse(init.body)).toEqual({ trace_ids: ['t1'] }); @@ -75,6 +80,58 @@ describe('GET /api/traces/[traceId]/spans', () => { expect(res.status).toBe(499); }); + it('anchors the search window on the task creation time', async () => { + const fetchMock = vi + .fn() + .mockResolvedValue(new Response('{"items":[]}', { status: 200 })); + vi.stubGlobal('fetch', fetchMock); + const from = '2026-01-01T00:00:00.000Z'; + + await call('t1', undefined, `?from=${encodeURIComponent(from)}`); + + const params = upstreamURL(fetchMock).searchParams; + const fromTs = Date.parse(params.get('from_ts')!); + const toTs = Date.parse(params.get('to_ts')!); + expect(fromTs).toBe(Date.parse(from) - 5 * 60 * 1000); + expect(toTs - fromTs).toBe(90 * 24 * 60 * 60 * 1000 - 60 * 1000); + }); + + it('caps the window at now for a recent task', async () => { + const fetchMock = vi + .fn() + .mockResolvedValue(new Response('{"items":[]}', { status: 200 })); + vi.stubGlobal('fetch', fetchMock); + const from = new Date(Date.now() - 60 * 60 * 1000).toISOString(); + + await call('t1', undefined, `?from=${encodeURIComponent(from)}`); + + const toTs = Date.parse(upstreamURL(fetchMock).searchParams.get('to_ts')!); + expect(Date.now() - toTs).toBeLessThan(5000); + }); + + it('sends no window without a creation time', async () => { + const fetchMock = vi + .fn() + .mockResolvedValue(new Response('{"items":[]}', { status: 200 })); + vi.stubGlobal('fetch', fetchMock); + + await call('t1'); + + const params = upstreamURL(fetchMock).searchParams; + expect(params.has('from_ts')).toBe(false); + expect(params.has('to_ts')).toBe(false); + }); + + it('rejects a creation time that is not a timestamp', async () => { + const fetchMock = vi.fn(); + vi.stubGlobal('fetch', fetchMock); + + const res = await call('t1', undefined, '?from=yesterday'); + + expect(res.status).toBe(400); + expect(fetchMock).not.toHaveBeenCalled(); + }); + it('passes the upstream status through', async () => { vi.stubGlobal( 'fetch', diff --git a/agentex-ui/app/api/traces/[traceId]/spans/route.ts b/agentex-ui/app/api/traces/[traceId]/spans/route.ts index 3f14e532..e058cb8a 100644 --- a/agentex-ui/app/api/traces/[traceId]/spans/route.ts +++ b/agentex-ui/app/api/traces/[traceId]/spans/route.ts @@ -9,8 +9,24 @@ import { applyBffCredentials, SGP_BASE_URL } from '@/app/api/_lib/bff'; */ export const dynamic = 'force-dynamic'; -// The platform caps a search page at this size, and the sidebar shows the first page. +// The sidebar shows one page and reports the rest through has_more. const PAGE_SIZE = 100; +// The platform refuses a window wider than 90 days, and it defaults an omitted window to the +// last 90 days, which hides older tasks. Anchor the window on the task's creation instead. +const WINDOW_MS = 90 * 24 * 60 * 60 * 1000 - 60 * 1000; +const SKEW_MS = 5 * 60 * 1000; + +function searchWindow(from: string | null): Record | null { + if (from === null) return {}; + const start = Date.parse(from); + if (Number.isNaN(start)) return null; + const fromTs = start - SKEW_MS; + const toTs = Math.min(Date.now(), fromTs + WINDOW_MS); + return { + from_ts: new Date(fromTs).toISOString(), + to_ts: new Date(toTs).toISOString(), + }; +} export async function GET( request: Request, @@ -24,6 +40,13 @@ export async function GET( } const { traceId } = await ctx.params; + const window = searchWindow(new URL(request.url).searchParams.get('from')); + if (window === null) { + return NextResponse.json( + { error: 'from must be an ISO timestamp' }, + { status: 400 } + ); + } const headers = new Headers({ 'Content-Type': 'application/json', accept: 'application/json', @@ -34,6 +57,9 @@ export async function GET( limit: String(PAGE_SIZE), sort_by: 'start_timestamp', sort_order: 'asc', + // Over the byte budget the platform shortens the page instead of refusing it. + allow_short_pages: 'true', + ...window, }); let upstream: Response; try { diff --git a/agentex-ui/components/task-header/task-header.tsx b/agentex-ui/components/task-header/task-header.tsx index b8b3aa2b..c1af5b0b 100644 --- a/agentex-ui/components/task-header/task-header.tsx +++ b/agentex-ui/components/task-header/task-header.tsx @@ -13,7 +13,6 @@ 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'; @@ -37,8 +36,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; + // Agents trace under the task id, so the task is the trace. + const traceId = taskId; const copyTaskId = async () => { if (taskId) { diff --git a/agentex-ui/components/traces-sidebar/traces-sidebar.tsx b/agentex-ui/components/traces-sidebar/traces-sidebar.tsx index 4d8cf703..05749490 100644 --- a/agentex-ui/components/traces-sidebar/traces-sidebar.tsx +++ b/agentex-ui/components/traces-sidebar/traces-sidebar.tsx @@ -1,5 +1,6 @@ import { AnimatePresence, motion } from 'framer-motion'; +import { useAgentexClient } from '@/components/providers'; import { JsonViewer } from '@/components/ui/json-viewer'; import { ResizableSidebar } from '@/components/ui/resizable-sidebar'; import { @@ -10,6 +11,7 @@ import { } from '@/components/ui/tooltip'; import { useSafeSearchParams } from '@/hooks/use-safe-search-params'; import { useSpans } from '@/hooks/use-spans'; +import { useTask } from '@/hooks/use-tasks'; const MIN_SIDEBAR_WIDTH = 350; const DEFAULT_SIDEBAR_WIDTH = 350; @@ -20,7 +22,15 @@ type TracesSidebarProps = { export function TracesSidebar({ isOpen }: TracesSidebarProps) { const { taskID } = useSafeSearchParams(); - const { spans, hasMore, isLoading, error } = useSpans(taskID); + const { agentexClient, sgpAppURL } = useAgentexClient(); + const { data: task, isError: taskUnavailable } = useTask({ + agentexClient, + taskId: taskID ?? '', + }); + // The task's creation time anchors the search window. Without it the query waits, unless + // the task itself cannot be read, in which case the platform's default window is used. + const createdAt = task?.created_at ?? (taskUnavailable ? null : undefined); + const { spans, hasMore, isLoading, error } = useSpans(taskID, createdAt); return ( @@ -78,8 +88,9 @@ export function TracesSidebar({ isOpen }: TracesSidebarProps) { {hasMore && (
- Showing the first {spans.length} spans. Use Investigate - traces for the full trace. + Showing the first {spans.length} spans. + {sgpAppURL && + ' Use Investigate traces for the full trace.'}
)} diff --git a/agentex-ui/hooks/use-spans.test.tsx b/agentex-ui/hooks/use-spans.test.tsx index 6c5a3949..fb577de3 100644 --- a/agentex-ui/hooks/use-spans.test.tsx +++ b/agentex-ui/hooks/use-spans.test.tsx @@ -39,8 +39,8 @@ const span = { describe('spansKeys', () => { it('scopes a task query to the selected account', () => { - expect(spansKeys.byTaskId('task-1', 'acct-1')).not.toEqual( - spansKeys.byTaskId('task-1', 'acct-2') + expect(spansKeys.byTaskId('task-1', 'acct-1', null)).not.toEqual( + spansKeys.byTaskId('task-1', 'acct-2', null) ); }); }); @@ -56,9 +56,10 @@ describe('useSpans', () => { .mockResolvedValue(jsonResponse({ items: [span], has_more: false })); vi.stubGlobal('fetch', fetchMock); - const { result } = renderHook(() => useSpans('task-1'), { - wrapper: createWrapper(), - }); + const { result } = renderHook( + () => useSpans('task-1', '2026-01-01T00:00:00.000Z'), + { wrapper: createWrapper() } + ); await waitFor(() => expect(result.current.isLoading).toBe(false)); @@ -67,7 +68,9 @@ describe('useSpans', () => { expect(result.current.error).toBeNull(); expect(fetchMock).toHaveBeenCalledTimes(1); const [url, init] = fetchMock.mock.calls[0]!; - expect(url).toBe('/api/traces/task-1/spans'); + expect(url).toBe( + '/api/traces/task-1/spans?from=2026-01-01T00%3A00%3A00.000Z' + ); expect(init.credentials).toBe('include'); expect(init.headers).toEqual({ 'x-selected-account-id': 'acct-1' }); }); @@ -78,7 +81,7 @@ describe('useSpans', () => { vi.fn().mockResolvedValue(jsonResponse({ items: [span], has_more: true })) ); - const { result } = renderHook(() => useSpans('task-1'), { + const { result } = renderHook(() => useSpans('task-1', null), { wrapper: createWrapper(), }); @@ -98,7 +101,7 @@ describe('useSpans', () => { ) ); - const { result } = renderHook(() => useSpans('task-1'), { + const { result } = renderHook(() => useSpans('task-1', null), { wrapper: createWrapper(), }); @@ -112,11 +115,37 @@ describe('useSpans', () => { const fetchMock = vi.fn(); vi.stubGlobal('fetch', fetchMock); - const { result } = renderHook(() => useSpans(null), { + const { result } = renderHook(() => useSpans(null, null), { wrapper: createWrapper(), }); expect(result.current.spans).toEqual([]); expect(fetchMock).not.toHaveBeenCalled(); }); + + it('waits until the task creation time is known', () => { + const fetchMock = vi.fn(); + vi.stubGlobal('fetch', fetchMock); + + renderHook(() => useSpans('task-1', undefined), { + wrapper: createWrapper(), + }); + + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('reads without a window when the task cannot be loaded', async () => { + const fetchMock = vi + .fn() + .mockResolvedValue(jsonResponse({ items: [], has_more: false })); + vi.stubGlobal('fetch', fetchMock); + + const { result } = renderHook(() => useSpans('task-1', null), { + wrapper: createWrapper(), + }); + + await waitFor(() => expect(result.current.isLoading).toBe(false)); + + expect(fetchMock.mock.calls[0]![0]).toBe('/api/traces/task-1/spans'); + }); }); diff --git a/agentex-ui/hooks/use-spans.ts b/agentex-ui/hooks/use-spans.ts index 446ea90e..20f262d4 100644 --- a/agentex-ui/hooks/use-spans.ts +++ b/agentex-ui/hooks/use-spans.ts @@ -6,10 +6,20 @@ import { useSafeSearchParams } from '@/hooks/use-safe-search-params'; export const spansKeys = { all: ['spans'] as const, - // The account scopes the platform read, so it scopes the cache entry too. - byTaskId: (taskId: string | null, accountId: string | null) => + // The account and the window anchor scope the platform read, so they scope the cache too. + byTaskId: ( + taskId: string | null, + accountId: string | null, + createdAt: string | null | undefined + ) => taskId - ? ([...spansKeys.all, 'task', taskId, accountId ?? ''] as const) + ? ([ + ...spansKeys.all, + 'task', + taskId, + accountId ?? '', + createdAt ?? '', + ] as const) : spansKeys.all, }; @@ -51,20 +61,28 @@ type UseSpansState = { * task id, through the same-origin BFF route that attaches credentials server-side. * * @param taskId - The task ID to fetch spans for, or null to disable the query + * @param createdAt - The task's creation time, which anchors the platform's search window. + * Undefined means not known yet (the query waits), null means unknown (no window is sent). * @returns The first page of spans in start order, whether more exist, the loading state, and any error message */ -export function useSpans(taskId: string | null): UseSpansState { +export function useSpans( + taskId: string | null, + createdAt: string | null | undefined +): UseSpansState { const { sgpAccountID } = useSafeSearchParams(); const { data, isLoading, error } = useQuery({ - queryKey: spansKeys.byTaskId(taskId, sgpAccountID), + queryKey: spansKeys.byTaskId(taskId, sgpAccountID, createdAt), queryFn: async ({ signal }): Promise => { if (!taskId) { return { items: [], hasMore: false }; } + const search = createdAt + ? `?${new URLSearchParams({ from: createdAt })}` + : ''; const response = await fetch( - `/api/traces/${encodeURIComponent(taskId)}/spans`, + `/api/traces/${encodeURIComponent(taskId)}/spans${search}`, { credentials: 'include', // Selected account, same source as the SDK, forwarded by the BFF. @@ -89,7 +107,7 @@ export function useSpans(taskId: string | null): UseSpansState { const page: SpansPage = await response.json(); return { items: page.items ?? [], hasMore: page.has_more ?? false }; }, - enabled: taskId !== null, + enabled: taskId !== null && createdAt !== undefined, }); return { From caeb229c5dd6f44f5de3b24327aa0ac2cb28e1c2 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Fri, 11 Sep 2026 10:04:21 -0400 Subject: [PATCH 4/4] fix(agentex-ui): retry a span read once after a session refresh The raw fetch behind the traces sidebar returned a 401 to the user when the access token had expired, where the SDK client refreshes the session and retries. The hook now does the same, only when login is enabled. The migration docstring also stops claiming that nothing writes the table and that the drop always fits the statement timeout. Co-Authored-By: Claude Fable 5.1 --- .../components/providers/agentex-provider.tsx | 2 +- agentex-ui/components/providers/index.ts | 6 ++- agentex-ui/hooks/use-spans.test.tsx | 45 +++++++++++++++++++ agentex-ui/hooks/use-spans.ts | 26 ++++++----- ...2026_09_11_1239_drop_spans_78384970fed5.py | 15 ++++--- 5 files changed, 74 insertions(+), 20 deletions(-) diff --git a/agentex-ui/components/providers/agentex-provider.tsx b/agentex-ui/components/providers/agentex-provider.tsx index 04b43bbb..522b80aa 100644 --- a/agentex-ui/components/providers/agentex-provider.tsx +++ b/agentex-ui/components/providers/agentex-provider.tsx @@ -20,7 +20,7 @@ import { // Hitting /api/auth/session runs the jwt-callback refresh and rotates the cookie. Deduped // so a burst of 401s (e.g. a refocused tab) shares one refresh. let sessionRefresh: Promise | null = null; -function refreshSession(): Promise { +export function refreshSession(): Promise { sessionRefresh ??= fetch('/api/auth/session', { credentials: 'include' }) .catch(() => {}) .finally(() => { diff --git a/agentex-ui/components/providers/index.ts b/agentex-ui/components/providers/index.ts index 670f12df..2b3e2a45 100644 --- a/agentex-ui/components/providers/index.ts +++ b/agentex-ui/components/providers/index.ts @@ -1,4 +1,8 @@ -export { AgentexProvider, useAgentexClient } from './agentex-provider'; +export { + AgentexProvider, + refreshSession, + useAgentexClient, +} from './agentex-provider'; export { TaskProvider } from './task-provider'; export { ThemeProvider } from './theme-provider'; export { QueryProvider } from './query-provider'; diff --git a/agentex-ui/hooks/use-spans.test.tsx b/agentex-ui/hooks/use-spans.test.tsx index fb577de3..c2e55004 100644 --- a/agentex-ui/hooks/use-spans.test.tsx +++ b/agentex-ui/hooks/use-spans.test.tsx @@ -10,6 +10,16 @@ vi.mock('@/hooks/use-safe-search-params', () => ({ useSafeSearchParams: () => ({ sgpAccountID: 'acct-1' }), })); +const providers = vi.hoisted(() => ({ + authEnabled: false, + refreshSession: vi.fn(async () => undefined), +})); + +vi.mock('@/components/providers', () => ({ + useAgentexClient: () => ({ authEnabled: providers.authEnabled }), + refreshSession: providers.refreshSession, +})); + function createWrapper() { const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } }, @@ -48,6 +58,41 @@ describe('spansKeys', () => { describe('useSpans', () => { afterEach(() => { vi.unstubAllGlobals(); + providers.authEnabled = false; + providers.refreshSession.mockClear(); + }); + + it('refreshes the session once and retries a 401 when login is enabled', async () => { + providers.authEnabled = true; + const fetchMock = vi + .fn() + .mockResolvedValueOnce(jsonResponse({ detail: 'expired' }, 401)) + .mockResolvedValueOnce(jsonResponse({ items: [span], has_more: false })); + vi.stubGlobal('fetch', fetchMock); + + const { result } = renderHook(() => useSpans('task-1', null), { + wrapper: createWrapper(), + }); + + await waitFor(() => expect(result.current.isLoading).toBe(false)); + + expect(result.current.spans).toEqual([span]); + expect(providers.refreshSession).toHaveBeenCalledTimes(1); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); + + it('does not refresh on a 401 when login is disabled', async () => { + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue(jsonResponse({ detail: 'expired' }, 401)) + ); + + const { result } = renderHook(() => useSpans('task-1', null), { + wrapper: createWrapper(), + }); + + await waitFor(() => expect(result.current.error).toBe('expired')); + expect(providers.refreshSession).not.toHaveBeenCalled(); }); it('reads the task trace through the BFF with the selected account', async () => { diff --git a/agentex-ui/hooks/use-spans.ts b/agentex-ui/hooks/use-spans.ts index 20f262d4..2778da48 100644 --- a/agentex-ui/hooks/use-spans.ts +++ b/agentex-ui/hooks/use-spans.ts @@ -2,6 +2,7 @@ import { useQuery } from '@tanstack/react-query'; +import { refreshSession, useAgentexClient } from '@/components/providers'; import { useSafeSearchParams } from '@/hooks/use-safe-search-params'; export const spansKeys = { @@ -70,6 +71,7 @@ export function useSpans( createdAt: string | null | undefined ): UseSpansState { const { sgpAccountID } = useSafeSearchParams(); + const { authEnabled } = useAgentexClient(); const { data, isLoading, error } = useQuery({ queryKey: spansKeys.byTaskId(taskId, sgpAccountID, createdAt), @@ -81,17 +83,19 @@ export function useSpans( const search = createdAt ? `?${new URLSearchParams({ from: createdAt })}` : ''; - const response = await fetch( - `/api/traces/${encodeURIComponent(taskId)}/spans${search}`, - { - credentials: 'include', - // Selected account, same source as the SDK, forwarded by the BFF. - headers: sgpAccountID - ? { 'x-selected-account-id': sgpAccountID } - : {}, - signal, - } - ); + const url = `/api/traces/${encodeURIComponent(taskId)}/spans${search}`; + const init: RequestInit = { + credentials: 'include', + // Selected account, same source as the SDK, forwarded by the BFF. + headers: sgpAccountID ? { 'x-selected-account-id': sgpAccountID } : {}, + signal, + }; + let response = await fetch(url, init); + if (response.status === 401 && authEnabled) { + // The access token expired between refreshes, so refresh once and retry like the SDK client. + await refreshSession(); + response = await fetch(url, init); + } if (!response.ok) { const body = await response.json().catch(() => ({})); diff --git a/agentex/database/migrations/alembic/versions/2026_09_11_1239_drop_spans_78384970fed5.py b/agentex/database/migrations/alembic/versions/2026_09_11_1239_drop_spans_78384970fed5.py index 1c2a9948..62253f6f 100644 --- a/agentex/database/migrations/alembic/versions/2026_09_11_1239_drop_spans_78384970fed5.py +++ b/agentex/database/migrations/alembic/versions/2026_09_11_1239_drop_spans_78384970fed5.py @@ -6,15 +6,16 @@ Drops the legacy Postgres-backed spans table. Agent spans are written to the platform's tracing service by the SDK's SGP tracing processor, and the -/spans API that fed this table is removed in the same change, so nothing -reads or writes it any more. +/spans API that fed this table is removed in the same change, so the +server neither reads nor writes it any more. Clients on SDK versions that +still default to that API get 404s from here on. Safety: -- DROP TABLE is metadata-only in PostgreSQL (the files are unlinked), so it - completes well inside the statement timeout regardless of table size. It - needs an AccessExclusiveLock; a writer still holding the table (an old pod - mid-rollout) makes the lock wait hit lock_timeout, and the pod retries the - migration on its next start. +- DROP TABLE unlinks the table's files rather than scanning rows, so its + cost does not grow with the row count. It needs an AccessExclusiveLock, + so a writer holding a conflicting lock for longer than lock_timeout (an + old pod mid-rollout) aborts the migration, and the pod retries it on its + next start. - IF EXISTS keeps re-runs idempotent. The indexes and the foreign key to tasks go with the table. - Downgrade recreates the empty table with the shape the ORM last declared.