Refactor the metadata streams stuff to be better

Author: ericallam

packages/core/src/v3/realtimeStreams/manager.ts

@@ -10,8 +10,8 @@ import {
 } from "./types.js";
 import { taskContext } from "../task-context-api.js";
 import { ApiClient } from "../apiClient/index.js";
-import { MetadataStream } from "../runMetadata/metadataStream.js";
-import { S2MetadataStream } from "../runMetadata/s2MetadataStream.js";
+import { StreamsWriterV1 } from "./streamsWriterV1.js";
+import { StreamsWriterV2 } from "./streamsWriterV2.js";
 
 export class StandardRealtimeStreamsManager implements RealtimeStreamsManager {
   constructor(
@@ -53,7 +53,7 @@ export class StandardRealtimeStreamsManager implements RealtimeStreamsManager {
 
     const streamInstance =
       parsedResponse.version === "v1"
-        ? new MetadataStream({
+        ? new StreamsWriterV1({
             key,
             runId,
             source: asyncIterableSource,
@@ -63,7 +63,7 @@ export class StandardRealtimeStreamsManager implements RealtimeStreamsManager {
             version,
             target: "self",
           })
-        : new S2MetadataStream({
+        : new StreamsWriterV2({
             basin: parsedResponse.basin,
             stream: key,
             accessToken: parsedResponse.accessToken,

packages/core/src/v3/runMetadata/metadataStream.ts renamed to packages/core/src/v3/realtimeStreams/streamsWriterV1.ts

@@ -2,9 +2,9 @@ import { request as httpsRequest } from "node:https";
 import { request as httpRequest } from "node:http";
 import { URL } from "node:url";
 import { randomBytes } from "node:crypto";
-import type { StreamInstance } from "./types.js";
+import { StreamsWriter } from "./types.js";
 
-export type MetadataOptions<T> = {
+export type StreamsWriterV1Options<T> = {
   baseUrl: string;
   runId: string;
   key: string;
@@ -23,7 +23,7 @@ interface BufferedChunk<T> {
   data: T;
 }
 
-export class MetadataStream<T> implements StreamInstance {
+export class StreamsWriterV1<T> implements StreamsWriter {
   private controller = new AbortController();
   private serverStream: ReadableStream<T>;
   private consumerStream: ReadableStream<T>;
@@ -42,7 +42,7 @@ export class MetadataStream<T> implements StreamInstance {
   private bufferReaderTask: Promise<void> | null = null;
   private streamComplete = false;
 
-  constructor(private options: MetadataOptions<T>) {
+  constructor(private options: StreamsWriterV1Options<T>) {
     const [serverStream, consumerStream] = this.createTeeStreams();
     this.serverStream = serverStream;
     this.consumerStream = consumerStream;

packages/core/src/v3/runMetadata/s2MetadataStream.ts renamed to packages/core/src/v3/realtimeStreams/streamsWriterV2.ts

@@ -1,5 +1,5 @@
 import { S2 } from "@s2-dev/streamstore";
-import type { StreamInstance } from "./types.js";
+import { StreamsWriter } from "./types.js";
 
 type LimitFunction = {
   readonly activeCount: number;
@@ -11,7 +11,7 @@ type LimitFunction = {
   ): Promise<ReturnType>;
 };
 
-export type S2MetadataStreamOptions<T = any> = {
+export type StreamsWriterV2Options<T = any> = {
   basin: string;
   stream: string;
   accessToken: string;
@@ -24,7 +24,7 @@ export type S2MetadataStreamOptions<T = any> = {
 };
 
 /**
- * S2MetadataStream writes metadata stream data directly to S2 (https://s2.dev).
+ * StreamsWriterV2 writes metadata stream data directly to S2 (https://s2.dev).
  *
  * Features:
  * - Batching: Reads chunks as fast as possible and buffers them
@@ -53,7 +53,7 @@ export type S2MetadataStreamOptions<T = any> = {
  * }
  * ```
  */
-export class S2MetadataStream<T = any> implements StreamInstance {
+export class StreamsWriterV2<T = any> implements StreamsWriter {
   private s2Client: S2;
   private serverStream: ReadableStream<T>;
   private consumerStream: ReadableStream<T>;
@@ -76,7 +76,7 @@ export class S2MetadataStream<T = any> implements StreamInstance {
   private readonly baseDelayMs = 1000;
   private readonly maxDelayMs = 30000;
 
-  constructor(private options: S2MetadataStreamOptions<T>) {
+  constructor(private options: StreamsWriterV2Options<T>) {
     this.limiter = options.limiter(1);
     this.debug = options.debug ?? false;
 

packages/core/src/v3/realtimeStreams/types.ts

@@ -19,3 +19,7 @@ export interface RealtimeStreamInstance<T> {
   wait(): Promise<void>;
   get stream(): AsyncIterableStream<T>;
 }
+
+export interface StreamsWriter {
+  wait(): Promise<void>;
+}

packages/core/src/v3/runMetadata/types.ts

@@ -29,7 +29,3 @@ export interface RunMetadataManager extends RunMetadataUpdater {
   get parent(): RunMetadataUpdater;
   get root(): RunMetadataUpdater;
 }
-
-export interface StreamInstance {
-  wait(): Promise<void>;
-}

packages/core/test/metadataStream.test.ts renamed to packages/core/test/streamsWriterV1.test.ts

@@ -1,11 +1,11 @@
 import { describe, it, expect, beforeEach, afterEach } from "vitest";
 import { createServer, Server, IncomingMessage, ServerResponse } from "node:http";
 import { AddressInfo } from "node:net";
-import { MetadataStream } from "../src/v3/runMetadata/metadataStream.js";
+import { StreamsWriterV1 } from "../src/v3/realtimeStreams/streamsWriterV1.js";
 
 type RequestHandler = (req: IncomingMessage, res: ServerResponse) => void;
 
-describe("MetadataStream", () => {
+describe("StreamsWriterV1", () => {
   let server: Server;
   let baseUrl: string;
   let requestHandler: RequestHandler | null = null;
@@ -67,7 +67,7 @@ describe("MetadataStream", () => {
       yield { chunk: 2, data: "chunk 2" };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -95,7 +95,7 @@ describe("MetadataStream", () => {
       yield { chunk: 0 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -138,7 +138,7 @@ describe("MetadataStream", () => {
       yield { chunk: 2 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -187,7 +187,7 @@ describe("MetadataStream", () => {
       yield { chunk: 0 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -214,7 +214,7 @@ describe("MetadataStream", () => {
 
       if (requestCount === 1) {
         // First request - don't respond, let it timeout
-        // (timeout is set to 15 minutes in MetadataStream, so we can't actually test this easily)
+        // (timeout is set to 15 minutes in StreamsWriterV1, so we can't actually test this easily)
         // Instead we'll just delay and then respond
         setTimeout(() => {
           res.writeHead(200);
@@ -231,7 +231,7 @@ describe("MetadataStream", () => {
       yield { chunk: 0 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -274,7 +274,7 @@ describe("MetadataStream", () => {
       }
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -319,7 +319,7 @@ describe("MetadataStream", () => {
       yield { chunk: 0 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -366,7 +366,7 @@ describe("MetadataStream", () => {
         yield { chunk: 1 };
       }
 
-      const metadataStream = new MetadataStream({
+      const metadataStream = new StreamsWriterV1({
         baseUrl,
         runId: "run_123",
         key: "test-stream",
@@ -414,7 +414,7 @@ describe("MetadataStream", () => {
       yield { chunk: 0 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -455,7 +455,7 @@ describe("MetadataStream", () => {
       yield { chunk: 0 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -478,7 +478,7 @@ describe("MetadataStream", () => {
       }
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -528,7 +528,7 @@ describe("MetadataStream", () => {
       }
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -587,7 +587,7 @@ describe("MetadataStream", () => {
       yield { chunk: 0 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -637,7 +637,7 @@ describe("MetadataStream", () => {
       }
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -659,7 +659,7 @@ describe("MetadataStream", () => {
       yield { chunk: 2, data: "data 2" };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -697,7 +697,7 @@ describe("MetadataStream", () => {
       yield { chunk: 0 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -739,7 +739,7 @@ describe("MetadataStream", () => {
       yield { chunk: 0 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -770,7 +770,7 @@ describe("MetadataStream", () => {
       yield { chunk: 1 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -794,7 +794,7 @@ describe("MetadataStream", () => {
       return;
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -820,7 +820,7 @@ describe("MetadataStream", () => {
       // Note: Throwing here would test error handling, but causes test infrastructure issues
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -861,7 +861,7 @@ describe("MetadataStream", () => {
       yield { chunk: 1 };
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",
@@ -911,7 +911,7 @@ describe("MetadataStream", () => {
         yield { chunk: 0 };
       }
 
-      const metadataStream = new MetadataStream({
+      const metadataStream = new StreamsWriterV1({
         baseUrl,
         runId: "run_123",
         key: "test-stream",
@@ -957,7 +957,7 @@ describe("MetadataStream", () => {
       }
     }
 
-    const metadataStream = new MetadataStream({
+    const metadataStream = new StreamsWriterV1({
       baseUrl,
       runId: "run_123",
       key: "test-stream",