Stream derived watch events from package runtimes

src/index.ts

@@ -2,7 +2,11 @@ import http from "node:http";
 import { AsyncLocalStorage } from "node:async_hooks";
 import { randomUUID } from "node:crypto";
 import { create } from "@bufbuild/protobuf";
-import { createClient, type Client } from "@connectrpc/connect";
+import {
+  createClient,
+  type Client,
+  type HandlerContext,
+} from "@connectrpc/connect";
 import {
   connectNodeAdapter,
   createConnectTransport,
@@ -18,10 +22,15 @@ import {
   type Value,
 } from "./camino/api_pb.js";
 import {
+  DerivedDependencySchema,
   HandshakeResponseSchema,
   InvokeResponseSchema,
   PackageRuntime,
+  WatchEventSchema,
+  WatchRequestSchema,
   type InvokeRequest,
+  type WatchEvent,
+  type WatchRequest,
 } from "./quixos/runtime_pb.js";
 import { FunctionRefSchema } from "./quixos/refs_pb.js";
 
@@ -122,14 +131,30 @@ type WatchableContext = {
   request: InvokeRequest;
 };
 
+type DerivedWatchEmission = {
+  watchId: string;
+  value?: unknown;
+  dependencies: DerivedDependency[];
+  error?: string;
+  initial: boolean;
+};
+
+type DerivedWatchSink = {
+  publish(event: DerivedWatchEmission): void;
+};
+
+const watchSinkStorage = new AsyncLocalStorage<DerivedWatchSink>();
+
 type ActiveDerivedWatch<Context extends WatchableContext> = {
   watchId: string;
   context: Context;
   get: RuntimeHandler<Context>;
+  sink?: DerivedWatchSink;
   controllers: Map<string, AbortController>;
   running: boolean;
   pending: boolean;
   stopped: boolean;
+  emitted: boolean;
   lastValue: unknown;
 };
 
@@ -198,13 +223,28 @@ export const derivedWithDeps = <Context>(operation: {
           value,
           dependencies,
         });
+        watch.sink?.publish({
+          watchId: watch.watchId,
+          value,
+          dependencies,
+          initial: !watch.emitted,
+        });
+        watch.emitted = true;
         reconcileSubscriptions(watch, dependencies);
       } while (watch.pending && !watch.stopped);
     } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
       logDerivedWatch("error", {
         watchId: watch.watchId,
-        error: error instanceof Error ? error.message : String(error),
+        error: message,
       });
+      watch.sink?.publish({
+        watchId: watch.watchId,
+        dependencies: [],
+        error: message,
+        initial: !watch.emitted,
+      });
+      watch.emitted = true;
     } finally {
       watch.running = false;
     }
@@ -295,10 +335,12 @@ export const derivedWithDeps = <Context>(operation: {
         watchId,
         context: context as DerivedContext,
         get: operation.get as RuntimeHandler<DerivedContext>,
+        sink: watchSinkStorage.getStore(),
         controllers: new Map(),
         running: false,
         pending: false,
         stopped: false,
+        emitted: false,
         lastValue: undefined,
       };
       watches.set(watchId, watch);
@@ -333,6 +375,51 @@ export type RuntimeFunctionSpec<ExportName extends string = string> = {
   operation?: string;
 };
 
+class AsyncEventQueue<T> implements AsyncIterableIterator<T> {
+  private readonly values: T[] = [];
+  private readonly waiters: Array<(result: IteratorResult<T>) => void> = [];
+  private closed = false;
+
+  push(value: T) {
+    if (this.closed) {
+      return;
+    }
+    const waiter = this.waiters.shift();
+    if (waiter) {
+      waiter({ value, done: false });
+      return;
+    }
+    this.values.push(value);
+  }
+
+  close() {
+    if (this.closed) {
+      return;
+    }
+    this.closed = true;
+    for (const waiter of this.waiters.splice(0)) {
+      waiter({ value: undefined, done: true });
+    }
+  }
+
+  async next(): Promise<IteratorResult<T>> {
+    const value = this.values.shift();
+    if (value !== undefined) {
+      return { value, done: false };
+    }
+    if (this.closed) {
+      return { value: undefined, done: true };
+    }
+    return await new Promise<IteratorResult<T>>((resolve) => {
+      this.waiters.push(resolve);
+    });
+  }
+
+  [Symbol.asyncIterator]() {
+    return this;
+  }
+}
+
 const isRecord = (value: unknown): value is Record<string, unknown> =>
   Boolean(value) && typeof value === "object" && !Array.isArray(value);
 
@@ -482,6 +569,31 @@ export const protoFieldsToJs = (
     ]),
   );
 
+const derivedDependencyToProto = (dependency: DerivedDependency) =>
+  create(DerivedDependencySchema, {
+    kind: dependency.kind,
+    objectId: dependency.objectId,
+    fieldName: dependency.kind === "field" ? dependency.fieldName : "",
+    sourceField: dependency.kind === "edges" ? dependency.sourceField ?? "" : "",
+  });
+
+const derivedEmissionToProto = (event: DerivedWatchEmission): WatchEvent =>
+  create(WatchEventSchema, {
+    watchId: event.watchId,
+    value: jsToProtoValue(event.value ?? null),
+    dependencies: event.dependencies.map(derivedDependencyToProto),
+    error: event.error ?? "",
+    initial: event.initial,
+  });
+
+const watchIdFromResult = (value: unknown) => {
+  if (!isRecord(value)) {
+    return "";
+  }
+  const watchId = value.watch_id ?? value.watchId;
+  return typeof watchId === "string" ? watchId : "";
+};
+
 export const objectRef = <ClassId extends string>(objectId: string) =>
   objectId as ObjectRef<ClassId>;
 
@@ -586,6 +698,52 @@ export const serveQuixosPackageRuntime = <
     }),
   );
 
+  const resolveRuntimeFunction = (
+    fn: InvokeRequest["function"],
+  ): {
+    spec: RuntimeFunctionSpec;
+    runtimeFunction: RuntimeHandler<Context>;
+  } => {
+    const symbol = fn?.symbol ?? "";
+    const spec = params.functions.find(
+      (entry) =>
+        entry.symbol === symbol &&
+        (entry.operation ?? "") === (fn?.operation ?? ""),
+    );
+    if (!spec) {
+      throw new Error(
+        `Unknown ${params.packageName} function: ${symbol}${
+          fn?.operation ? `.${fn.operation}` : ""
+        }`,
+      );
+    }
+    const implementationExport = (
+      params.implementation as Record<
+        string,
+        RuntimeHandler<Context> | FieldOperation<Context>
+      >
+    )[spec.exportName];
+    const runtimeFunction =
+      spec.operation && typeof implementationExport === "object"
+        ? implementationExport[spec.operation as keyof FieldOperation<Context>]
+        : implementationExport;
+    if (!runtimeFunction) {
+      throw new Error(
+        `Missing ${params.packageName} implementation export: ${spec.exportName}${
+          spec.operation ? `.${spec.operation}` : ""
+        }`,
+      );
+    }
+    if (typeof runtimeFunction !== "function") {
+      throw new Error(
+        `${params.packageName} implementation export is not callable: ${spec.exportName}${
+          spec.operation ? `.${spec.operation}` : ""
+        }`,
+      );
+    }
+    return { spec, runtimeFunction };
+  };
+
   const handler = connectNodeAdapter({
     routes: (router) => {
       router.service(PackageRuntime, {
@@ -600,45 +758,7 @@ export const serveQuixosPackageRuntime = <
 
         async invoke(request) {
           try {
-            const symbol = request.function?.symbol ?? "";
-            const spec = params.functions.find(
-              (entry) =>
-                entry.symbol === symbol &&
-                (entry.operation ?? "") === (request.function?.operation ?? ""),
-            );
-            if (!spec) {
-              throw new Error(
-                `Unknown ${params.packageName} function: ${symbol}${
-                  request.function?.operation
-                    ? `.${request.function.operation}`
-                    : ""
-                }`,
-              );
-            }
-            const implementationExport = (
-              params.implementation as Record<
-                string,
-                RuntimeHandler<Context> | FieldOperation<Context>
-              >
-            )[spec.exportName];
-            const runtimeFunction =
-              spec.operation && typeof implementationExport === "object"
-                ? implementationExport[spec.operation as keyof FieldOperation<Context>]
-                : implementationExport;
-            if (!runtimeFunction) {
-              throw new Error(
-                `Missing ${params.packageName} implementation export: ${spec.exportName}${
-                  spec.operation ? `.${spec.operation}` : ""
-                }`,
-              );
-            }
-            if (typeof runtimeFunction !== "function") {
-              throw new Error(
-                `${params.packageName} implementation export is not callable: ${spec.exportName}${
-                  spec.operation ? `.${spec.operation}` : ""
-                }`,
-              );
-            }
+            const { runtimeFunction } = resolveRuntimeFunction(request.function);
             return create(InvokeResponseSchema, {
               ok: true,
               result: jsToProtoValue(
@@ -653,6 +773,85 @@ export const serveQuixosPackageRuntime = <
             });
           }
         },
+
+        async *watch(request: WatchRequest, context: HandlerContext) {
+          const queue = new AsyncEventQueue<WatchEvent>();
+          let watchId = "";
+          const closeQueue = () => queue.close();
+          context.signal.addEventListener("abort", closeQueue, { once: true });
+          try {
+            const { spec, runtimeFunction } = resolveRuntimeFunction(
+              request.function,
+            );
+            if (spec.operation !== "watchStart") {
+              throw new Error(
+                `PackageRuntime.watch requires a watchStart operation, got ${spec.symbol}${
+                  spec.operation ? `.${spec.operation}` : ""
+                }`,
+              );
+            }
+            const sink: DerivedWatchSink = {
+              publish(event) {
+                queue.push(derivedEmissionToProto(event));
+              },
+            };
+            const result = await watchSinkStorage.run(sink, async () =>
+              runtimeFunction(
+                params.createContext(
+                  trackingCamino,
+                  request as unknown as InvokeRequest,
+                ),
+              ),
+            );
+            watchId = watchIdFromResult(result);
+            if (!watchId) {
+              throw new Error(
+                `Watch start did not return a watch_id for ${spec.symbol}`,
+              );
+            }
+            for await (const event of queue) {
+              yield event;
+            }
+          } catch (error) {
+            yield create(WatchEventSchema, {
+              watchId,
+              value: jsToProtoValue(null),
+              dependencies: [],
+              error: error instanceof Error ? error.message : String(error),
+              initial: !watchId,
+            });
+          } finally {
+            context.signal.removeEventListener("abort", closeQueue);
+            queue.close();
+            if (watchId) {
+              try {
+                const stopFunction = create(FunctionRefSchema, {
+                  ...(request.function ?? create(FunctionRefSchema, {})),
+                  operation: "watchStop",
+                });
+                const { runtimeFunction } = resolveRuntimeFunction(stopFunction);
+                await runtimeFunction(
+                  params.createContext(
+                    trackingCamino,
+                    create(WatchRequestSchema, {
+                      ...request,
+                      function: stopFunction,
+                      input: {
+                        ...request.input,
+                        watch_id: jsToProtoValue(watchId),
+                      },
+                    }) as unknown as InvokeRequest,
+                  ),
+                );
+              } catch (error) {
+                logDerivedWatch("stop-error", {
+                  watchId,
+                  error: error instanceof Error ? error.message : String(error),
+                });
+              }
+            }
+          }
+        },
       });
     },
   });