The stream.read() span now works better when specifying a startIndex

Author: ericallam

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

@@ -19,6 +19,7 @@ import { WaitpointPresenter } from "./WaitpointPresenter.server";
 import { engine } from "~/v3/runEngine.server";
 import { resolveEventRepositoryForStore } from "~/v3/eventRepository/index.server";
 import { IEventRepository, SpanDetail } from "~/v3/eventRepository/eventRepository.types";
+import { safeJsonParse } from "~/utils/json";
 
 type Result = Awaited<ReturnType<SpanPresenter["call"]>>;
 export type Span = NonNullable<NonNullable<Result>["span"]>;
@@ -570,13 +571,18 @@ export class SpanPresenter extends BasePresenter {
           return { ...data, entity: null };
         }
 
+        const metadata = span.entity.metadata
+          ? (safeJsonParse(span.entity.metadata) as Record<string, unknown> | undefined)
+          : undefined;
+
         return {
           ...data,
           entity: {
             type: "realtime-stream" as const,
             object: {
               runId,
               streamKey,
+              metadata,
             },
           },
         };

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

@@ -1161,6 +1161,7 @@ function SpanEntity({ span }: { span: Span }) {
           <RealtimeStreamViewer
             runId={span.entity.object.runId}
             streamKey={span.entity.object.streamKey}
+            metadata={span.entity.object.metadata}
           />
         </div>
       );

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

@@ -71,23 +71,28 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => {
     run.realtimeStreamsVersion
   );
 
-  return realtimeStream.streamResponse(
-    request,
-    run.friendlyId,
-    streamKey,
-    request.signal,
-    lastEventId
-  );
+  return realtimeStream.streamResponse(request, run.friendlyId, streamKey, request.signal, {
+    lastEventId,
+  });
 };
 
-export function RealtimeStreamViewer({ runId, streamKey }: { runId: string; streamKey: string }) {
+export function RealtimeStreamViewer({
+  runId,
+  streamKey,
+  metadata,
+}: {
+  runId: string;
+  streamKey: string;
+  metadata: Record<string, unknown> | undefined;
+}) {
   const organization = useOrganization();
   const project = useProject();
   const environment = useEnvironment();
 
   const resourcePath = `/resources/orgs/${organization.slug}/projects/${project.slug}/env/${environment.slug}/runs/${runId}/streams/${streamKey}`;
 
-  const { chunks, error, isConnected } = useRealtimeStream(resourcePath);
+  const startIndex = typeof metadata?.startIndex === "number" ? metadata.startIndex : undefined;
+  const { chunks, error, isConnected } = useRealtimeStream(resourcePath, startIndex);
   const scrollRef = useRef<HTMLDivElement>(null);
   const bottomRef = useRef<HTMLDivElement>(null);
   const [isAtBottom, setIsAtBottom] = useState(true);
@@ -124,7 +129,10 @@ export function RealtimeStreamViewer({ runId, streamKey }: { runId: string; stre
     }
   }, [chunks, isAtBottom]);
 
-  const maxLineNumberWidth = chunks.length.toString().length;
+  const firstLineNumber = startIndex ?? 0;
+  const lastLineNumber = firstLineNumber + chunks.length - 1;
+  const maxLineNumberWidth = (chunks.length > 0 ? lastLineNumber : firstLineNumber).toString()
+    .length;
 
   return (
     <div className="flex h-full flex-col overflow-hidden border-t border-grid-bright">
@@ -178,7 +186,7 @@ export function RealtimeStreamViewer({ runId, streamKey }: { runId: string; stre
               <StreamChunkLine
                 key={index}
                 chunk={chunk}
-                lineNumber={index + 1}
+                lineNumber={firstLineNumber + index}
                 maxLineNumberWidth={maxLineNumberWidth}
               />
             ))}
@@ -246,7 +254,7 @@ function StreamChunkLine({
   );
 }
 
-function useRealtimeStream(resourcePath: string) {
+function useRealtimeStream(resourcePath: string, startIndex?: number) {
   const [chunks, setChunks] = useState<StreamChunk[]>([]);
   const [error, setError] = useState<Error | null>(null);
   const [isConnected, setIsConnected] = useState(false);
@@ -259,6 +267,8 @@ function useRealtimeStream(resourcePath: string) {
       try {
         const sseSubscription = new SSEStreamSubscription(resourcePath, {
           signal: abortController.signal,
+          lastEventId: startIndex ? (startIndex - 1).toString() : undefined,
+          timeoutInSeconds: 30,
         });
 
         const stream = await sseSubscription.subscribe();
@@ -300,7 +310,7 @@ function useRealtimeStream(resourcePath: string) {
       abortController.abort();
       reader?.cancel();
     };
-  }, [resourcePath]);
+  }, [resourcePath, startIndex]);
 
   return { chunks, error, isConnected };
 }

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

@@ -424,19 +424,24 @@ export class ClickhouseEventRepository implements IEventRepository {
 
   private extractEntityFromAttributes(
     attributes: Attributes
-  ): { entityType: string; entityId?: string } | undefined {
+  ): { entityType: string; entityId?: string; entityMetadata?: string } | undefined {
     if (!attributes || typeof attributes !== "object") {
       return undefined;
     }
 
     const entityType = attributes[SemanticInternalAttributes.ENTITY_TYPE];
     const entityId = attributes[SemanticInternalAttributes.ENTITY_ID];
+    const entityMetadata = attributes[SemanticInternalAttributes.ENTITY_METADATA];
 
     if (typeof entityType !== "string") {
       return undefined;
     }
 
-    return { entityType, entityId: entityId as string | undefined };
+    return {
+      entityType,
+      entityId: entityId as string | undefined,
+      entityMetadata: entityMetadata as string | undefined,
+    };
   }
 
   private addToBatch(events: TaskEventV1Input[] | TaskEventV1Input) {
@@ -1101,6 +1106,7 @@ export class ClickhouseEventRepository implements IEventRepository {
           entity: {
             type: undefined,
             id: undefined,
+            metadata: undefined,
           },
           metadata: {},
         };
@@ -1140,6 +1146,12 @@ export class ClickhouseEventRepository implements IEventRepository {
         span.entity = {
           id: parsedMetadata.entity.entityId,
           type: parsedMetadata.entity.entityType,
+          metadata:
+            "entityMetadata" in parsedMetadata.entity &&
+            parsedMetadata.entity.entityMetadata &&
+            typeof parsedMetadata.entity.entityMetadata === "string"
+              ? parsedMetadata.entity.entityMetadata
+              : undefined,
         };
       }
 

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

@@ -783,6 +783,7 @@ export class EventRepository implements IEventRepository {
           SemanticInternalAttributes.ENTITY_TYPE
         ),
         id: rehydrateAttribute<string>(spanEvent.properties, SemanticInternalAttributes.ENTITY_ID),
+        metadata: undefined,
       };
 
       return {

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

@@ -217,6 +217,7 @@ export type SpanDetail = {
     // Used for entity type switching in SpanEntity
     type: string | undefined;
     id: string | undefined;
+    metadata: string | undefined;
   };
 
   metadata: any; // Used by SpanPresenter for entity processing

packages/core/src/v3/semanticInternalAttributes.ts

@@ -29,6 +29,7 @@ export const SemanticInternalAttributes = {
   SPAN: "$span",
   ENTITY_TYPE: "$entity.type",
   ENTITY_ID: "$entity.id",
+  ENTITY_METADATA: "$entity.metadata",
   OUTPUT: "$output",
   OUTPUT_TYPE: "$mime_type_output",
   STYLE: "$style",

packages/trigger-sdk/src/v3/streams.ts

@@ -115,6 +115,9 @@ async function readStream<T>(
       runId,
       [SemanticInternalAttributes.ENTITY_TYPE]: "realtime-stream",
       [SemanticInternalAttributes.ENTITY_ID]: `${runId}:${key}`,
+      [SemanticInternalAttributes.ENTITY_METADATA]: JSON.stringify({
+        startIndex: options?.startIndex,
+      }),
       [SemanticInternalAttributes.STYLE_ICON]: "streams",
       ...accessoryAttributes({
         items: [