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/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..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 OpenTelemetry-style spans for task execution +- **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 @@ -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..d564b7f4 --- /dev/null +++ b/agentex-ui/app/api/traces/[traceId]/spans/route.test.ts @@ -0,0 +1,161 @@ +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, 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', () => { + 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&allow_short_pages=true' + ); + 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' + ); + 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('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', + 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..e058cb8a --- /dev/null +++ b/agentex-ui/app/api/traces/[traceId]/spans/route.ts @@ -0,0 +1,84 @@ +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 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, + 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 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', + }); + await applyBffCredentials(request, headers); + + const query = new URLSearchParams({ + 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 { + 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/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/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 d056a32a..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, 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 ( @@ -76,8 +86,16 @@ export function TracesSidebar({ isOpen }: TracesSidebarProps) { )} + {hasMore && ( +
+ Showing the first {spans.length} spans. + {sgpAppURL && + ' 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..c2e55004 --- /dev/null +++ b/agentex-ui/hooks/use-spans.test.tsx @@ -0,0 +1,196 @@ +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 { spansKeys, useSpans } from './use-spans'; + +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 } }, + }); + 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('spansKeys', () => { + it('scopes a task query to the selected account', () => { + expect(spansKeys.byTaskId('task-1', 'acct-1', null)).not.toEqual( + spansKeys.byTaskId('task-1', 'acct-2', null) + ); + }); +}); + +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 () => { + const fetchMock = vi + .fn() + .mockResolvedValue(jsonResponse({ items: [span], has_more: false })); + vi.stubGlobal('fetch', fetchMock); + + const { result } = renderHook( + () => useSpans('task-1', '2026-01-01T00:00:00.000Z'), + { 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?from=2026-01-01T00%3A00%3A00.000Z' + ); + 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', null), { + 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', null), { + 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, 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 f2213186..2778da48 100644 --- a/agentex-ui/hooks/use-spans.ts +++ b/agentex-ui/hooks/use-spans.ts @@ -2,60 +2,121 @@ import { useQuery } from '@tanstack/react-query'; -import { useAgentexClient } from '@/components/providers'; - -import type { Span } from 'agentex/resources'; +import { refreshSession, useAgentexClient } from '@/components/providers'; +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 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 ?? '', + createdAt ?? '', + ] 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. + * 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. * - * Queries by task_id first. Falls back to trace_id=taskId for backward - * compatibility with spans created before the task_id column was added. - * - * @param 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 + * @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 { - const { agentexClient } = useAgentexClient(); +export function useSpans( + taskId: string | null, + createdAt: string | null | undefined +): UseSpansState { + const { sgpAccountID } = useSafeSearchParams(); + const { authEnabled } = useAgentexClient(); - const { data, isLoading, error } = useQuery({ - queryKey: spansKeys.byTaskId(taskId), - queryFn: async ({ signal }) => { + const { data, isLoading, error } = useQuery({ + queryKey: spansKeys.byTaskId(taskId, sgpAccountID, createdAt), + 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 search = createdAt + ? `?${new URLSearchParams({ from: createdAt })}` + : ''; + 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 (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, + enabled: taskId !== null && createdAt !== undefined, }); 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..62253f6f --- /dev/null +++ b/agentex/database/migrations/alembic/versions/2026_09_11_1239_drop_spans_78384970fed5.py @@ -0,0 +1,64 @@ +"""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 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 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. + 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