feat(webapp,core): page large run traces from the trace API

The run trace API (`GET /api/v1/runs/:runId/trace`) can page large
traces. Pass `page[size]` and `page[after]` to read every span of a run
at the root of its trace. Requests without these parameters are
unchanged.

Mono-RevId: 3a811d1a61118cfe089ba692b4edd56282674c65

Author: Mariefay

apps/webapp/app/routes/api.v1.runs.$runId.trace.ts

@@ -1,10 +1,23 @@
+import { trace } from "@opentelemetry/api";
 import { json } from "@remix-run/server-runtime";
+import type { RetrieveRunTracePageResponseBody } from "@trigger.dev/core/v3";
 import { BatchId } from "@trigger.dev/core/v3/isomorphic";
 import { z } from "zod";
 import { $replica } from "~/db.server";
+import { env } from "~/env.server";
 import { anyResource, createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server";
+import { clampToEmergencySpanCap } from "~/v3/eventRepository/emergencySpanCap.server";
 import { getEventRepositoryForStore } from "~/v3/eventRepository/index.server";
 import { getTraceInsertedAtEnd } from "~/v3/eventRepository/traceInsertedAtBound";
+import { getRunTracePage } from "~/v3/eventRepository/runTracePage.server";
+import { FEATURE_FLAG } from "~/v3/featureFlags";
+import { makeFlag } from "~/v3/featureFlags.server";
+import { toPageSpan } from "~/v3/eventRepository/tracePage";
+import {
+  DEFAULT_TRACE_PAGE_SIZE,
+  parseTracePageRequest,
+  type TracePageRequest,
+} from "~/v3/eventRepository/tracePageRequest";
 import { getTaskEventStoreTableForRun } from "~/v3/taskEventStore.server";
 import { findRunByIdWithMollifierFallback } from "~/v3/mollifier/readFallback.server";
 import { buildSyntheticTraceBody } from "~/v3/mollifier/syntheticApiResponses.server";
@@ -78,7 +91,22 @@ export const loader = createLoaderApiRoute(
       },
     },
   },
-  async ({ resource: resolved, authentication }) => {
+  async ({ resource: resolved, authentication, request }) => {
+    const searchParams = new URL(request.url).searchParams;
+
+    if (
+      hasPageParams(searchParams) &&
+      (await isTracePagingEnabled(authentication.environment.organization.featureFlags))
+    ) {
+      const pageRequest = parseTracePageRequest(searchParams, DEFAULT_TRACE_PAGE_SIZE);
+      if (pageRequest.isErr()) {
+        return json({ error: PAGE_REQUEST_ERRORS[pageRequest.error] }, { status: 400 });
+      }
+      if (pageRequest.value) {
+        return tracePageResponse(resolved, authentication.environment, pageRequest.value);
+      }
+    }
+
     if (resolved.source === "buffer") {
       // Buffered runs have no events ingested yet — the drainer hasn't
       // materialised the PG row and the worker hasn't started executing.
@@ -104,6 +132,11 @@ export const loader = createLoaderApiRoute(
       { insertedAtEnd: getTraceInsertedAtEnd(run) }
     );
 
+    trace.getActiveSpan()?.setAttributes({
+      "run.depth": run.depth,
+      "trace.truncated": traceSummary?.isTruncated ?? false,
+    });
+
     if (!traceSummary) {
       return json({ error: "Trace not found" }, { status: 404 });
     }
@@ -116,3 +149,88 @@ export const loader = createLoaderApiRoute(
     );
   }
 );
+
+function hasPageParams(searchParams: URLSearchParams): boolean {
+  return searchParams.has("page[size]") || searchParams.has("page[after]");
+}
+
+// Off: page parameters are ignored and the request gets today's tree response.
+function isTracePagingEnabled(orgFeatureFlags: unknown): Promise<boolean> {
+  return makeFlag()({
+    key: FEATURE_FLAG.publicTracePagingEnabled,
+    defaultValue: false,
+    overrides: (orgFeatureFlags as Record<string, unknown>) ?? {},
+  });
+}
+
+const PAGE_REQUEST_ERRORS = {
+  invalid_page_size: "page[size] must be a positive integer",
+  invalid_cursor: "page[after] is not a valid cursor; pass back pagination.next unchanged",
+} as const;
+
+const CHILD_RUN_ERROR =
+  "Paging is only supported for runs at the root of their trace. Request this run's trace without page parameters.";
+
+// The same synthetic root span the unpaged response returns for a run still in the trigger buffer.
+function bufferedRunPage(
+  run: Parameters<typeof buildSyntheticTraceBody>[0]
+): RetrieveRunTracePageResponseBody {
+  if (!run.spanId) {
+    return { data: [], attemptFailures: [], pagination: {} };
+  }
+
+  const { rootSpan } = buildSyntheticTraceBody(run).trace;
+  return { data: [toPageSpan(rootSpan)], attemptFailures: [], pagination: {} };
+}
+
+async function tracePageResponse(
+  resolved: ResolvedRun,
+  environment: { id: string; organization: { id: string } },
+  pageRequest: TracePageRequest
+) {
+  // Paging reads the whole trace, so the run's span must be the trace root (no parent run or span).
+  if (resolved.run.parentTaskRunId || resolved.run.parentSpanId) {
+    return json({ error: CHILD_RUN_ERROR }, { status: 400 });
+  }
+
+  if (resolved.source === "buffer") {
+    return json(bufferedRunPage(resolved.run), { status: 200 });
+  }
+
+  const capped = env.TRACE_VIEW_EMERGENCY_SPAN_CAP !== undefined;
+  if (capped && pageRequest.after) {
+    return json(
+      { error: "Trace paging is temporarily limited to the first page" },
+      { status: 503, headers: { "Retry-After": "60", "x-should-retry": "false" } }
+    );
+  }
+
+  const run = resolved.run;
+  const page = await getRunTracePage({
+    repository: await getEventRepositoryForStore(run.taskEventStore, environment.organization.id),
+    storeTable: getTaskEventStoreTableForRun(run),
+    environmentId: environment.id,
+    run,
+    pageRequest: { ...pageRequest, limit: clampToEmergencySpanCap(pageRequest.limit) },
+    // Like the dashboard, stop paging under the emergency cap instead of reading the whole trace.
+    stopAfterPage: capped,
+  });
+
+  if (page.isOk()) {
+    return json(page.value, { status: 200 });
+  }
+
+  switch (page.error) {
+    case "store_unavailable":
+      // The SDK retries 5xx; tell it not to, so a struggling store isn't hit harder.
+      return json(
+        { error: "The trace store is temporarily unavailable" },
+        { status: 503, headers: { "Retry-After": "5", "x-should-retry": "false" } }
+      );
+    case "paging_unsupported":
+      return json(
+        { error: "Trace paging isn't supported by this deployment's trace store" },
+        { status: 501 }
+      );
+  }
+}

apps/webapp/app/v3/eventRepository/clickhouseEventRepository.server.ts

@@ -75,6 +75,7 @@ import type {
   TraceChunk,
   TraceChunkCursor,
   TraceChunkEvent,
+  TraceChunkScopeOptions,
   TraceDetailedSummary,
   TraceErrorEvents,
   TraceEventOptions,
@@ -1757,11 +1758,11 @@ export class ClickhouseEventRepository implements IEventRepository {
     startCreatedAt: Date,
     endCreatedAt: Date | undefined,
     cursor: TraceChunkCursor | undefined,
-    options?: { includeDebugLogs?: boolean; limit?: number; tailInsertedAtSinceMs?: number }
+    options?: TraceChunkScopeOptions & { limit?: number }
   ): Promise<TraceChunk | undefined> {
     const limit = options?.limit ?? this.maximumTraceChunkSize;
 
-    const { events, nextCursor, hasMore } = await this.#fetchTraceChunkRecords({
+    const { events, nextCursor, hasMore, droppedKeyRows } = await this.#fetchTraceChunkRecords({
       environmentId,
       traceId,
       startCreatedAt,
@@ -1775,37 +1776,10 @@ export class ClickhouseEventRepository implements IEventRepository {
       events: events.map((record) => this.#toTraceChunkEvent(record)),
       nextCursor,
       hasMore,
+      droppedKeyRows,
     };
   }
 
-  async getTraceSpanCount(
-    storeTable: TaskEventStoreTable,
-    environmentId: string,
-    traceId: string,
-    startCreatedAt: Date,
-    endCreatedAt: Date | undefined,
-    options?: { includeDebugLogs?: boolean }
-  ): Promise<number | undefined> {
-    const queryBuilder = this.#createTraceSpanCountQueryBuilder();
-    this.#applyTraceScopeWhere(queryBuilder, {
-      environmentId,
-      traceId,
-      startCreatedAt,
-      endCreatedAt,
-      options,
-    });
-
-    const [queryError, records] = await queryBuilder.execute();
-
-    if (queryError) {
-      logger.error("getTraceSpanCount failed", { error: queryError, traceId });
-      return undefined;
-    }
-
-    const count = records?.[0]?.count;
-    return count === undefined ? undefined : Number(count);
-  }
-
   #createTraceChunkQueryBuilder() {
     return this._version === "v2"
       ? this._clickhouse.taskEventsV2.traceChunkQueryBuilder()
@@ -1827,11 +1801,12 @@ export class ClickhouseEventRepository implements IEventRepository {
     endCreatedAt?: Date;
     cursor?: TraceChunkCursor;
     limit: number;
-    options?: { includeDebugLogs?: boolean; tailInsertedAtSinceMs?: number };
+    options?: TraceChunkScopeOptions;
   }): Promise<{
     events: TaskEventChunkV2Result[];
     nextCursor: TraceChunkCursor | null;
     hasMore: boolean;
+    droppedKeyRows: boolean;
   }> {
     const queryBuilder = this.#applyTraceChunkScope({
       environmentId,
@@ -1869,21 +1844,28 @@ export class ClickhouseEventRepository implements IEventRepository {
       groupBuilder.where(clause, params);
       groupBuilder.orderBy(TRACE_CHUNK_ORDER_BY);
       // The next cursor skips past this key, so rows beyond the cap are dropped.
-      groupBuilder.limit(this.maximumKeyRows);
+      groupBuilder.limit(this.maximumKeyRows + 1);
 
       const [groupError, groupRecords] = await groupBuilder.execute();
       if (groupError) {
         throw groupError;
       }
 
+      const rows = groupRecords ?? [];
       return {
-        events: groupRecords ?? [],
+        events: rows.slice(0, this.maximumKeyRows),
         nextCursor: slice.nextCursor,
         hasMore: slice.hasMore,
+        droppedKeyRows: rows.length > this.maximumKeyRows,
       };
     }
 
-    return { events: slice.events, nextCursor: slice.nextCursor, hasMore: slice.hasMore };
+    return {
+      events: slice.events,
+      nextCursor: slice.nextCursor,
+      hasMore: slice.hasMore,
+      droppedKeyRows: false,
+    };
   }
 
   #applyTraceChunkScope({
@@ -1897,7 +1879,7 @@ export class ClickhouseEventRepository implements IEventRepository {
     traceId: string;
     startCreatedAt: Date;
     endCreatedAt?: Date;
-    options?: { includeDebugLogs?: boolean; tailInsertedAtSinceMs?: number };
+    options?: TraceChunkScopeOptions;
   }) {
     const queryBuilder = this.#createTraceChunkQueryBuilder();
     this.#applyTraceScopeWhere(queryBuilder, {
@@ -1910,12 +1892,6 @@ export class ClickhouseEventRepository implements IEventRepository {
     return queryBuilder;
   }
 
-  #createTraceSpanCountQueryBuilder() {
-    return this._version === "v2"
-      ? this._clickhouse.taskEventsV2.traceSpanCountQueryBuilder()
-      : this._clickhouse.taskEvents.traceSpanCountQueryBuilder();
-  }
-
   #applyTraceScopeWhere(
     queryBuilder: { where: (clause: string, params?: any) => unknown },
     {
@@ -1929,7 +1905,7 @@ export class ClickhouseEventRepository implements IEventRepository {
       traceId: string;
       startCreatedAt: Date;
       endCreatedAt?: Date;
-      options?: { includeDebugLogs?: boolean; tailInsertedAtSinceMs?: number };
+      options?: TraceChunkScopeOptions;
     }
   ) {
     queryBuilder.where("environment_id = {environmentId: String}", { environmentId });
@@ -1964,6 +1940,12 @@ export class ClickhouseEventRepository implements IEventRepository {
           ),
         });
       }
+
+      if (options?.insertedAtEnd) {
+        queryBuilder.where("inserted_at <= {insertedAtEnd: DateTime64(3)}", {
+          insertedAtEnd: convertDateToClickhouseDateTime(options.insertedAtEnd),
+        });
+      }
     }
 
     if (options?.includeDebugLogs === false) {

apps/webapp/app/v3/eventRepository/eventRepository.server.ts

@@ -71,6 +71,7 @@ import type {
   TraceChunk,
   TraceChunkCursor,
   TraceChunkEvent,
+  TraceChunkScopeOptions,
   TraceDetailedSummary,
   TraceErrorEvents,
   TraceEventOptions,
@@ -487,22 +488,11 @@ export class EventRepository implements IEventRepository {
     _startCreatedAt: Date,
     _endCreatedAt: Date | undefined,
     _cursor: TraceChunkCursor | undefined,
-    _options?: { includeDebugLogs?: boolean; limit?: number }
+    _options?: TraceChunkScopeOptions & { limit?: number }
   ): Promise<TraceChunk | undefined> {
     return undefined;
   }
 
-  public async getTraceSpanCount(
-    _storeTable: TaskEventStoreTable,
-    _environmentId: string,
-    _traceId: string,
-    _startCreatedAt: Date,
-    _endCreatedAt: Date | undefined,
-    _options?: { includeDebugLogs?: boolean }
-  ): Promise<number | undefined> {
-    return undefined;
-  }
-
   public async getTraceErrorEvents(
     _storeTable: TaskEventStoreTable,
     _environmentId: string,

apps/webapp/app/v3/eventRepository/eventRepository.types.ts

@@ -311,10 +311,20 @@ export type TraceChunkEvent = {
 
 export type { TraceChunkCursor };
 
+export type TraceChunkScopeOptions = {
+  includeDebugLogs?: boolean;
+  // Live tail: only rows written at or after this time (ms since epoch).
+  tailInsertedAtSinceMs?: number;
+  // Only rows written at or before this time.
+  insertedAtEnd?: Date;
+};
+
 export type TraceChunk = {
   events: Array<TraceChunkEvent>;
   nextCursor: TraceChunkCursor | null;
   hasMore: boolean;
+  // A span had more rows at one timestamp than one read allows; the extra rows were skipped.
+  droppedKeyRows?: boolean;
 };
 
 export type TraceErrorEvents = {
@@ -449,18 +459,9 @@ export interface IEventRepository {
     startCreatedAt: Date,
     endCreatedAt: Date | undefined,
     cursor: TraceChunkCursor | undefined,
-    options?: { includeDebugLogs?: boolean; limit?: number; tailInsertedAtSinceMs?: number }
+    options?: TraceChunkScopeOptions & { limit?: number }
   ): Promise<TraceChunk | undefined>;
 
-  getTraceSpanCount(
-    storeTable: TaskEventStoreTable,
-    environmentId: string,
-    traceId: string,
-    startCreatedAt: Date,
-    endCreatedAt: Date | undefined,
-    options?: { includeDebugLogs?: boolean }
-  ): Promise<number | undefined>;
-
   getTraceErrorEvents(
     storeTable: TaskEventStoreTable,
     environmentId: string,