Add derived watch lifecycle helper

src/index.ts

@@ -1,5 +1,6 @@
 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 {
@@ -107,6 +108,48 @@ const logDerivedDependencies = (dependencies: DerivedDependency[]) => {
   );
 };
 
+const logDerivedWatch = (event: string, detail: Record<string, unknown>) => {
+  if (process.env.QUIXOS_DERIVED_WATCH_LOG !== "1") {
+    return;
+  }
+  process.stderr.write(
+    `[derived-watch] ${JSON.stringify({ event, ...detail })}\n`,
+  );
+};
+
+type WatchableContext = {
+  camino: CaminoClient;
+  request: InvokeRequest;
+};
+
+type ActiveDerivedWatch<Context extends WatchableContext> = {
+  watchId: string;
+  context: Context;
+  get: RuntimeHandler<Context>;
+  controllers: Map<string, AbortController>;
+  running: boolean;
+  pending: boolean;
+  stopped: boolean;
+  lastValue: unknown;
+};
+
+const dependencyObjectIds = (dependencies: DerivedDependency[]) =>
+  new Set(dependencies.map((dependency) => dependency.objectId));
+
+const stopDerivedWatch = <Context extends WatchableContext>(
+  watch: ActiveDerivedWatch<Context>,
+) => {
+  if (watch.stopped) {
+    return;
+  }
+  watch.stopped = true;
+  for (const controller of watch.controllers.values()) {
+    controller.abort();
+  }
+  watch.controllers.clear();
+  logDerivedWatch("stopped", { watchId: watch.watchId });
+};
+
 export type Field<T> = {
   get(): Promise<T | undefined>;
   set(value: T): Promise<void>;
@@ -129,8 +172,116 @@ export const fieldOps = <Context>(
 
 export const derivedWithDeps = <Context>(operation: {
   get: RuntimeHandler<Context>;
-}): FieldOperation<Context> =>
+}): FieldOperation<Context> => {
-  fieldOps({
+  type DerivedContext = Context & WatchableContext;
+  const watches = new Map<string, ActiveDerivedWatch<DerivedContext>>();
+
+  const rerun = async (watch: ActiveDerivedWatch<DerivedContext>) => {
+    if (watch.stopped) {
+      return;
+    }
+    if (watch.running) {
+      watch.pending = true;
+      return;
+    }
+    watch.running = true;
+    try {
+      do {
+        watch.pending = false;
+        const { value, dependencies } = await runWithDerivedDependencyTracking(
+          () => watch.get(watch.context),
+        );
+        watch.lastValue = value;
+        logDerivedDependencies(dependencies);
+        logDerivedWatch("evaluated", {
+          watchId: watch.watchId,
+          value,
+          dependencies,
+        });
+        reconcileSubscriptions(watch, dependencies);
+      } while (watch.pending && !watch.stopped);
+    } catch (error) {
+      logDerivedWatch("error", {
+        watchId: watch.watchId,
+        error: error instanceof Error ? error.message : String(error),
+      });
+    } finally {
+      watch.running = false;
+    }
+  };
+
+  const scheduleRerun = (
+    watch: ActiveDerivedWatch<DerivedContext>,
+    reason: string,
+  ) => {
+    if (watch.stopped) {
+      return;
+    }
+    logDerivedWatch("scheduled", { watchId: watch.watchId, reason });
+    void rerun(watch);
+  };
+
+  const startObjectWatch = (
+    watch: ActiveDerivedWatch<DerivedContext>,
+    objectId: string,
+  ) => {
+    const controller = new AbortController();
+    watch.controllers.set(objectId, controller);
+    logDerivedWatch("subscribed", { watchId: watch.watchId, objectId });
+    void (async () => {
+      try {
+        for await (const event of watch.context.camino.watchObject(
+          {
+            objectId,
+            includeSnapshot: false,
+          },
+          { signal: controller.signal },
+        )) {
+          if (controller.signal.aborted || watch.stopped) {
+            break;
+          }
+          scheduleRerun(
+            watch,
+            event.op?.id ? `op:${event.op.id}` : `object:${objectId}`,
+          );
+        }
+      } catch (error) {
+        if (!controller.signal.aborted && !watch.stopped) {
+          logDerivedWatch("watch-error", {
+            watchId: watch.watchId,
+            objectId,
+            error: error instanceof Error ? error.message : String(error),
+          });
+        }
+      } finally {
+        if (watch.controllers.get(objectId) === controller) {
+          watch.controllers.delete(objectId);
+          logDerivedWatch("unsubscribed", { watchId: watch.watchId, objectId });
+        }
+      }
+    })();
+  };
+
+  const reconcileSubscriptions = (
+    watch: ActiveDerivedWatch<DerivedContext>,
+    dependencies: DerivedDependency[],
+  ) => {
+    const objectIds = dependencyObjectIds(dependencies);
+    for (const [objectId, controller] of watch.controllers.entries()) {
+      if (!objectIds.has(objectId)) {
+        controller.abort();
+        watch.controllers.delete(objectId);
+        logDerivedWatch("unsubscribed", { watchId: watch.watchId, objectId });
+      }
+    }
+    for (const objectId of objectIds) {
+      if (!watch.controllers.has(objectId)) {
+        startObjectWatch(watch, objectId);
+      }
+    }
+  };
+
+  return fieldOps({
     async get(context) {
       const { value, dependencies } = await runWithDerivedDependencyTracking(
         () => operation.get(context),
@@ -138,7 +289,43 @@ export const derivedWithDeps = <Context>(operation: {
       logDerivedDependencies(dependencies);
       return value;
     },
+    async watchStart(context) {
+      const watchId = `watch:${randomUUID()}`;
+      const watch: ActiveDerivedWatch<DerivedContext> = {
+        watchId,
+        context: context as DerivedContext,
+        get: operation.get as RuntimeHandler<DerivedContext>,
+        controllers: new Map(),
+        running: false,
+        pending: false,
+        stopped: false,
+        lastValue: undefined,
+      };
+      watches.set(watchId, watch);
+      logDerivedWatch("started", { watchId });
+      await rerun(watch);
+      return {
+        watch_id: watchId,
+      };
+    },
+    watchStop(context) {
+      const input = protoFieldsToJs((context as DerivedContext).request.input);
+      const watchId =
+        typeof input.watch_id === "string"
+          ? input.watch_id
+          : typeof input.watchId === "string"
+            ? input.watchId
+            : "";
+      const watch = watches.get(watchId);
+      if (!watch) {
+        throw new Error(`Unknown derived watch: ${watchId}`);
+      }
+      stopDerivedWatch(watch);
+      watches.delete(watchId);
+      return {};
+    },
   });
+};
 
 export type RuntimeFunctionSpec<ExportName extends string = string> = {
   exportName: ExportName;