Merge branch 'main' into limits-page-2

Author: samejr

apps/webapp/app/components/primitives/TreeView/TreeView.tsx

@@ -527,6 +527,7 @@ export function useTree<TData, TFilterValue>({
 /** An actual tree structure with custom data */
 export type Tree<TData> = {
   id: string;
+  runId?: string;
   children?: Tree<TData>[];
   data: TData;
 };
@@ -535,6 +536,7 @@ export type Tree<TData> = {
 export type FlatTreeItem<TData> = {
   id: string;
   parentId?: string | undefined;
+  runId?: string;
   children: string[];
   hasChildren: boolean;
   /** The indentation level, the root is 0 */
@@ -552,6 +554,7 @@ export function flattenTree<TData>(tree: Tree<TData>): FlatTree<TData> {
     flatTree.push({
       id: node.id,
       parentId,
+      runId: node.runId,
       children,
       hasChildren: children.length > 0,
       level,
@@ -571,6 +574,7 @@ export function flattenTree<TData>(tree: Tree<TData>): FlatTree<TData> {
 type FlatTreeWithoutChildren<TData> = {
   id: string;
   parentId: string | undefined;
+  runId?: string;
   data: TData;
 };
 
@@ -580,7 +584,7 @@ export function createTreeFromFlatItems<TData>(
 ): Tree<TData> | undefined {
   // Index items by id
   const indexedItems: { [id: string]: Tree<TData> } = withoutChildren.reduce((acc, item) => {
-    acc[item.id] = { id: item.id, data: item.data, children: [] };
+    acc[item.id] = { id: item.id, runId: item.runId, data: item.data, children: [] };
     return acc;
   }, {} as { [id: string]: Tree<TData> });
 

apps/webapp/app/presenters/v3/RunPresenter.server.ts

@@ -209,6 +209,10 @@ export class RunPresenter {
     //we need the start offset for each item, and the total duration of the entire tree
     const treeRootStartTimeMs = tree ? tree?.data.startTime.getTime() : 0;
     let totalDuration = tree?.data.duration ?? 0;
+
+    // Build the linkedRunIdBySpanId map during the same walk
+    const linkedRunIdBySpanId: Record<string, string> = {};
+
     const events = tree
       ? flattenTree(tree).map((n) => {
           const offset = millisecondsToNanoseconds(
@@ -218,6 +222,12 @@ export class RunPresenter {
           if (!n.data.isDebug) {
             totalDuration = Math.max(totalDuration, offset + n.data.duration);
           }
+
+          // For cached spans, store the mapping from spanId to the linked run's ID
+          if (n.data.style?.icon === "task-cached" && n.runId) {
+            linkedRunIdBySpanId[n.id] = n.runId;
+          }
+
           return {
             ...n,
             data: {
@@ -260,6 +270,7 @@ export class RunPresenter {
           ? millisecondsToNanoseconds(run.startedAt.getTime() - run.createdAt.getTime())
           : undefined,
         overridesBySpanId: traceSummary.overridesBySpanId,
+        linkedRunIdBySpanId,
       },
       maximumLiveReloadingSetting: eventRepository.maximumLiveReloadingSetting,
     };

apps/webapp/app/presenters/v3/SpanPresenter.server.ts

@@ -35,11 +35,13 @@ export class SpanPresenter extends BasePresenter {
     projectSlug,
     spanId,
     runFriendlyId,
+    linkedRunId,
   }: {
     userId: string;
     projectSlug: string;
     spanId: string;
     runFriendlyId: string;
+    linkedRunId?: string;
   }) {
     const project = await this._replica.project.findFirst({
       where: {
@@ -88,6 +90,7 @@ export class SpanPresenter extends BasePresenter {
       traceId,
       eventRepository,
       spanId,
+      linkedRunId,
       createdAt: parentRun.createdAt,
       completedAt: parentRun.completedAt,
       environmentId: parentRun.runtimeEnvironmentId,
@@ -126,6 +129,7 @@ export class SpanPresenter extends BasePresenter {
     traceId,
     eventRepository,
     spanId,
+    linkedRunId,
     createdAt,
     completedAt,
   }: {
@@ -134,19 +138,12 @@ export class SpanPresenter extends BasePresenter {
     traceId: string;
     eventRepository: IEventRepository;
     spanId: string;
+    linkedRunId?: string;
     createdAt: Date;
     completedAt: Date | null;
   }) {
-    const originalRunId = await eventRepository.getSpanOriginalRunId(
-      eventStore,
-      environmentId,
-      spanId,
-      traceId,
-      createdAt,
-      completedAt ?? undefined
-    );
-
-    const run = await this.findRun({ originalRunId, spanId, environmentId });
+    // Use linkedRunId if provided (for cached spans), otherwise look up by spanId
+    const run = await this.findRun({ originalRunId: linkedRunId, spanId, environmentId });
 
     if (!run) {
       return;
@@ -272,7 +269,7 @@ export class SpanPresenter extends BasePresenter {
       workerQueue: run.workerQueue,
       traceId: run.traceId,
       spanId: run.spanId,
-      isCached: !!originalRunId,
+      isCached: !!linkedRunId,
       machinePreset: machine?.name,
       taskEventStore: run.taskEventStore,
       externalTraceId,

apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam/route.tsx

@@ -476,6 +476,12 @@ function TraceView({
 
   const spanOverrides = selectedSpanId ? overridesBySpanId?.[selectedSpanId] : undefined;
 
+  // Get the linked run ID for cached spans (map built during RunPresenter walk)
+  const { linkedRunIdBySpanId } = trace;
+  const selectedSpanLinkedRunId = selectedSpanId
+    ? linkedRunIdBySpanId?.[selectedSpanId]
+    : undefined;
+
   return (
     <div className={cn("grid h-full max-h-full grid-cols-1 overflow-hidden")}>
       <ResizablePanelGroup
@@ -526,6 +532,7 @@ function TraceView({
               spanId={selectedSpanId}
               spanOverrides={spanOverrides as SpanOverride | undefined}
               closePanel={() => replaceSearchParam("span")}
+              linkedRunId={selectedSpanLinkedRunId}
             />
           </ResizablePanel>
         )}

apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.spans.$spanParam/route.tsx

@@ -93,6 +93,9 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => {
   const { projectParam, organizationSlug, envParam, runParam, spanParam } =
     v3SpanParamsSchema.parse(params);
 
+  const url = new URL(request.url);
+  const linkedRunId = url.searchParams.get("linkedRunId") ?? undefined;
+
   const presenter = new SpanPresenter();
 
   try {
@@ -101,6 +104,7 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => {
       spanId: spanParam,
       runFriendlyId: runParam,
       userId,
+      linkedRunId,
     });
 
     return typedjson(result);
@@ -130,11 +134,13 @@ export function SpanView({
   spanId,
   spanOverrides,
   closePanel,
+  linkedRunId,
 }: {
   runParam: string;
   spanId: string | undefined;
   spanOverrides?: SpanOverride;
   closePanel?: () => void;
+  linkedRunId?: string;
 }) {
   const organization = useOrganization();
   const project = useProject();
@@ -143,10 +149,11 @@ export function SpanView({
 
   useEffect(() => {
     if (spanId === undefined) return;
-    fetcher.load(
-      `/resources/orgs/${organization.slug}/projects/${project.slug}/env/${environment.slug}/runs/${runParam}/spans/${spanId}`
-    );
-  }, [organization.slug, project.slug, environment.slug, runParam, spanId]);
+    const url = `/resources/orgs/${organization.slug}/projects/${project.slug}/env/${
+      environment.slug
+    }/runs/${runParam}/spans/${spanId}${linkedRunId ? `?linkedRunId=${linkedRunId}` : ""}`;
+    fetcher.load(url);
+  }, [organization.slug, project.slug, environment.slug, runParam, spanId, linkedRunId]);
 
   if (spanId === undefined) {
     return null;
@@ -305,7 +312,12 @@ function RunBody({
   useEffect(() => {
     if (resetFetcher.data && resetFetcher.state === "idle") {
       // Check if the response indicates success
-      if (resetFetcher.data && typeof resetFetcher.data === "object" && "success" in resetFetcher.data && resetFetcher.data.success === true) {
+      if (
+        resetFetcher.data &&
+        typeof resetFetcher.data === "object" &&
+        "success" in resetFetcher.data &&
+        resetFetcher.data.success === true
+      ) {
         toast.custom(
           (t) => (
             <ToastUI
@@ -574,7 +586,11 @@ function RunBody({
                     <div className="flex items-start justify-between gap-2">
                       <div className="flex-1">
                         {run.idempotencyKey ? (
-                          <CopyableText value={run.idempotencyKey} copyValue={run.idempotencyKey} asChild />
+                          <CopyableText
+                            value={run.idempotencyKey}
+                            copyValue={run.idempotencyKey}
+                            asChild
+                          />
                         ) : (
                           <div className="break-all">–</div>
                         )}
@@ -951,7 +967,9 @@ function RunBody({
           {run.logsDeletedAt === null ? (
             <>
               <LinkButton
-                to={`${v3LogsPath(organization, project, environment)}?runId=${runParam}&from=${new Date(run.createdAt).getTime() - 60000}`}
+                to={`${v3LogsPath(organization, project, environment)}?runId=${runParam}&from=${
+                  new Date(run.createdAt).getTime() - 60000
+                }`}
                 variant="secondary/medium"
               >
                 View logs
@@ -968,8 +986,6 @@ function RunBody({
               </LinkButton>
             </>
           ) : null}
-
-
         </div>
       </div>
     </div>

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

@@ -59,7 +59,6 @@ import type {
   TraceEventOptions,
   TraceSummary,
 } from "./eventRepository.types";
-import { originalRunIdCache } from "./originalRunIdCache.server";
 
 export type ClickhouseEventRepositoryConfig = {
   clickhouse: ClickHouse;
@@ -705,13 +704,6 @@ export class ClickhouseEventRepository implements IEventRepository {
       expires_at: convertDateToClickhouseDateTime(new Date(Date.now() + 365 * 24 * 60 * 60 * 1000)),
     };
 
-    const originalRunId =
-      options.attributes.properties?.[SemanticInternalAttributes.ORIGINAL_RUN_ID];
-
-    if (typeof originalRunId === "string") {
-      await originalRunIdCache.set(traceId, spanId, originalRunId);
-    }
-
     const events = [event];
 
     if (failedWithError) {
@@ -1197,17 +1189,6 @@ export class ClickhouseEventRepository implements IEventRepository {
     return span;
   }
 
-  async getSpanOriginalRunId(
-    storeTable: TaskEventStoreTable,
-    environmentId: string,
-    spanId: string,
-    traceId: string,
-    startCreatedAt: Date,
-    endCreatedAt?: Date
-  ): Promise<string | undefined> {
-    return await originalRunIdCache.lookup(traceId, spanId);
-  }
-
   #mergeRecordsIntoSpanDetail(
     spanId: string,
     records: TaskEventDetailsV1Result[],

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

@@ -64,7 +64,6 @@ import type {
   TraceEventOptions,
   TraceSummary,
 } from "./eventRepository.types";
-import { originalRunIdCache } from "./originalRunIdCache.server";
 
 const MAX_FLUSH_DEPTH = 5;
 
@@ -817,40 +816,6 @@ export class EventRepository implements IEventRepository {
     });
   }
 
-  async getSpanOriginalRunId(
-    storeTable: TaskEventStoreTable,
-    environmentId: string,
-    spanId: string,
-    traceId: string,
-    startCreatedAt: Date,
-    endCreatedAt?: Date
-  ): Promise<string | undefined> {
-    return await startActiveSpan("getSpanOriginalRunId", async (s) => {
-      return await originalRunIdCache.swr(traceId, spanId, async () => {
-        const spanEvent = await this.#getSpanEvent({
-          storeTable,
-          spanId,
-          environmentId,
-          startCreatedAt,
-          endCreatedAt,
-          options: { includeDebugLogs: false },
-        });
-
-        if (!spanEvent) {
-          return;
-        }
-        // This is used when the span is a cached run (because of idempotency key)
-        // so this span isn't the actual run span, but points to the original run
-        const originalRun = rehydrateAttribute<string>(
-          spanEvent.properties,
-          SemanticInternalAttributes.ORIGINAL_RUN_ID
-        );
-
-        return originalRun;
-      });
-    });
-  }
-
   async #createSpanFromEvent(
     storeTable: TaskEventStoreTable,
     event: PreparedEvent,

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

@@ -393,15 +393,6 @@ export interface IEventRepository {
     options?: { includeDebugLogs?: boolean }
   ): Promise<SpanDetail | undefined>;
 
-  getSpanOriginalRunId(
-    storeTable: TaskEventStoreTable,
-    environmentId: string,
-    spanId: string,
-    traceId: string,
-    startCreatedAt: Date,
-    endCreatedAt?: Date
-  ): Promise<string | undefined>;
-
   // Event recording methods
   recordEvent(
     message: string,