camino-package-runtime/src/index.ts

import http from "node:http";
import { AsyncLocalStorage } from "node:async_hooks";
import { create } from "@bufbuild/protobuf";
import { createClient, type Client } from "@connectrpc/connect";
import {
  connectNodeAdapter,
  createConnectTransport,
} from "@connectrpc/connect-node";
import {
  CaminoService,
  CrdtValueSchema,
  ListValueSchema,
  NullValueSchema,
  ObjectValueSchema,
  RefValueSchema,
  ValueSchema,
  type Value,
} from "./camino/api_pb.js";
import {
  HandshakeResponseSchema,
  InvokeResponseSchema,
  PackageRuntime,
  type InvokeRequest,
} from "./quixos/runtime_pb.js";
import { FunctionRefSchema } from "./quixos/refs_pb.js";

declare const objectRefBrand: unique symbol;

export type ObjectRef<ClassId extends string> = string & {
  readonly [objectRefBrand]: ClassId;
};

export type CaminoClient = Client<typeof CaminoService>;

export type { InvokeRequest, Value };

export type DerivedDependency =
  | {
      kind: "object";
      objectId: string;
    }
  | {
      kind: "field";
      objectId: string;
      fieldName: string;
    }
  | {
      kind: "edges";
      objectId: string;
      sourceField?: string;
    };

type DerivedDependencyTracker = {
  record(dependency: DerivedDependency): void;
  dependencies(): DerivedDependency[];
};

const dependencyKey = (dependency: DerivedDependency) => {
  switch (dependency.kind) {
    case "object":
      return `object:${dependency.objectId}`;
    case "field":
      return `field:${dependency.objectId}:${dependency.fieldName}`;
    case "edges":
      return `edges:${dependency.objectId}:${dependency.sourceField ?? ""}`;
  }
};

const createDependencyTracker = (): DerivedDependencyTracker => {
  const dependencies = new Map<string, DerivedDependency>();
  return {
    record(dependency) {
      dependencies.set(dependencyKey(dependency), dependency);
    },
    dependencies() {
      return [...dependencies.values()];
    },
  };
};

const dependencyStorage = new AsyncLocalStorage<DerivedDependencyTracker>();

export const recordDerivedDependency = (dependency: DerivedDependency) => {
  dependencyStorage.getStore()?.record(dependency);
};

export const getActiveDerivedDependencies = (): DerivedDependency[] =>
  dependencyStorage.getStore()?.dependencies() ?? [];

export const runWithDerivedDependencyTracking = async <T>(
  fn: () => Promise<T> | T,
): Promise<{ value: T; dependencies: DerivedDependency[] }> => {
  const tracker = createDependencyTracker();
  const value = await dependencyStorage.run(tracker, fn);
  return {
    value,
    dependencies: tracker.dependencies(),
  };
};

const logDerivedDependencies = (dependencies: DerivedDependency[]) => {
  if (process.env.QUIXOS_DERIVED_DEPS_LOG !== "1") {
    return;
  }
  process.stderr.write(
    `[derived-deps] ${JSON.stringify({ dependencies })}\n`,
  );
};

export type Field<T> = {
  get(): Promise<T | undefined>;
  set(value: T): Promise<void>;
};

export type RuntimeHandler<Context> = (
  context: Context,
) => Promise<unknown> | unknown;

export type FieldOperation<Context> = {
  get?: RuntimeHandler<Context>;
  set?: RuntimeHandler<Context>;
  watchStart?: RuntimeHandler<Context>;
  watchStop?: RuntimeHandler<Context>;
};

export const fieldOps = <Context>(
  operation: FieldOperation<Context>,
): FieldOperation<Context> => operation;

export const derivedWithDeps = <Context>(operation: {
  get: RuntimeHandler<Context>;
}): FieldOperation<Context> =>
  fieldOps({
    async get(context) {
      const { value, dependencies } = await runWithDerivedDependencyTracking(
        () => operation.get(context),
      );
      logDerivedDependencies(dependencies);
      return value;
    },
  });

export type RuntimeFunctionSpec<ExportName extends string = string> = {
  exportName: ExportName;
  symbol: string;
  operation?: string;
};

const isRecord = (value: unknown): value is Record<string, unknown> =>
  Boolean(value) && typeof value === "object" && !Array.isArray(value);

const bytesToBase64 = (value: Uint8Array) =>
  Buffer.from(value).toString("base64");

const base64ToBytes = (value: string) => Buffer.from(value, "base64");

export const jsToProtoValue = (value: unknown): Value => {
  if (value === null || value === undefined) {
    return create(ValueSchema, {
      kind: {
        case: "nullValue",
        value: create(NullValueSchema, {}),
      },
    });
  }
  if (typeof value === "boolean") {
    return create(ValueSchema, {
      kind: { case: "boolValue", value },
    });
  }
  if (typeof value === "number") {
    return create(ValueSchema, {
      kind: { case: "numberValue", value },
    });
  }
  if (typeof value === "bigint") {
    return create(ValueSchema, {
      kind: { case: "integerValue", value: value.toString() },
    });
  }
  if (typeof value === "string") {
    return create(ValueSchema, {
      kind: { case: "stringValue", value },
    });
  }
  if (value instanceof Uint8Array) {
    return create(ValueSchema, {
      kind: { case: "bytesValue", value },
    });
  }
  if (Array.isArray(value)) {
    return create(ValueSchema, {
      kind: {
        case: "listValue",
        value: create(ListValueSchema, {
          values: value.map(jsToProtoValue),
        }),
      },
    });
  }
  if (isRecord(value)) {
    if (typeof value.$caminoRef === "string") {
      return create(ValueSchema, {
        kind: {
          case: "refValue",
          value: create(RefValueSchema, {
            objectId: value.$caminoRef,
          }),
        },
      });
    }
    if (
      typeof value.$caminoBytes === "string" &&
      Object.keys(value).every((key) => key === "$caminoBytes")
    ) {
      return create(ValueSchema, {
        kind: {
          case: "bytesValue",
          value: base64ToBytes(value.$caminoBytes),
        },
      });
    }
    if (
      typeof value.$caminoCrdtType === "string" &&
      typeof value.$caminoCrdtPayload === "string"
    ) {
      return create(ValueSchema, {
        kind: {
          case: "crdtValue",
          value: create(CrdtValueSchema, {
            type: value.$caminoCrdtType,
            encoding:
              typeof value.$caminoCrdtEncoding === "string"
                ? value.$caminoCrdtEncoding
                : "base64",
            payload: base64ToBytes(value.$caminoCrdtPayload),
          }),
        },
      });
    }
    return create(ValueSchema, {
      kind: {
        case: "objectValue",
        value: create(ObjectValueSchema, {
          fields: jsObjectToProtoFields(value),
        }),
      },
    });
  }
  throw new Error(`Unsupported Camino RPC value: ${typeof value}`);
};

export const protoValueToJs = (value: Value | undefined): unknown => {
  switch (value?.kind.case) {
    case "nullValue":
    case undefined:
      return null;
    case "boolValue":
    case "numberValue":
    case "stringValue":
      return value.kind.value;
    case "integerValue":
      return value.kind.value;
    case "bytesValue":
      return bytesToBase64(value.kind.value);
    case "listValue":
      return value.kind.value.values.map(protoValueToJs);
    case "objectValue":
      return protoFieldsToJs(value.kind.value.fields);
    case "refValue":
      return value.kind.value.objectId;
    case "crdtValue":
      return {
        $caminoCrdtType: value.kind.value.type,
        $caminoCrdtEncoding: value.kind.value.encoding,
        $caminoCrdtPayload: bytesToBase64(value.kind.value.payload),
      };
  }
};

export const jsObjectToProtoFields = (value: Record<string, unknown>) =>
  Object.fromEntries(
    Object.entries(value)
      .filter(([, fieldValue]) => fieldValue !== undefined)
      .map(([key, fieldValue]) => [key, jsToProtoValue(fieldValue)]),
  );

export const protoFieldsToJs = (
  fields: Record<string, Value> | undefined,
): Record<string, unknown> =>
  Object.fromEntries(
    Object.entries(fields ?? {}).map(([key, value]) => [
      key,
      protoValueToJs(value),
    ]),
  );

export const objectRef = <ClassId extends string>(objectId: string) =>
  objectId as ObjectRef<ClassId>;

export const createField = <T>(
  camino: CaminoClient,
  objectId: string,
  fieldName: string,
): Field<T> => ({
  async get() {
    recordDerivedDependency({ kind: "field", objectId, fieldName });
    const response = await camino.getObject({ objectId });
    return protoValueToJs(response.object?.fields[fieldName]) as T | undefined;
  },
  async set(value) {
    await camino.setField({
      objectId,
      fieldName,
      value: jsToProtoValue(value),
    });
  },
});

const createDependencyTrackingCaminoClient = (
  camino: CaminoClient,
): CaminoClient =>
  new Proxy(camino, {
    get(target, property, receiver) {
      if (property === "getObject") {
        return async (request: { objectId: string }, ...args: unknown[]) => {
          if (request.objectId) {
            recordDerivedDependency({
              kind: "object",
              objectId: request.objectId,
            });
          }
          return await (
            Reflect.get(target, property, receiver) as (
              request: { objectId: string },
              ...args: unknown[]
            ) => Promise<unknown>
          )(request, ...args);
        };
      }
      if (property === "listEdges") {
        return async (
          request: { objectId: string; sourceField?: string },
          ...args: unknown[]
        ) => {
          if (request.objectId) {
            recordDerivedDependency({
              kind: "edges",
              objectId: request.objectId,
              ...(request.sourceField ? { sourceField: request.sourceField } : {}),
            });
          }
          return await (
            Reflect.get(target, property, receiver) as (
              request: { objectId: string; sourceField?: string },
              ...args: unknown[]
            ) => Promise<unknown>
          )(request, ...args);
        };
      }
      return Reflect.get(target, property, receiver);
    },
  });

export const serveQuixosPackageRuntime = <
  Context,
  Spec extends readonly RuntimeFunctionSpec[],
>(params: {
  packageNamespace: string;
  packageName: string;
  runtimeProtocolVersion: string;
  functions: Spec;
  implementation: {
    [Name in Spec[number]["exportName"]]:
      | RuntimeHandler<Context>
      | FieldOperation<Context>;
  };
  createContext: (camino: CaminoClient, request: InvokeRequest) => Context;
}) => {
  const host = process.env.QUIXOS_RUNTIME_HOST ?? "127.0.0.1";
  const port = Number(process.env.QUIXOS_RUNTIME_PORT ?? "0");
  const caminoUrl = process.env.CAMINO_URL ?? "http://127.0.0.1:7310";
  const camino = createClient(
    CaminoService,
    createConnectTransport({ baseUrl: caminoUrl, httpVersion: "1.1" }),
  );
  const trackingCamino = createDependencyTrackingCaminoClient(camino);

  if (!Number.isInteger(port) || port <= 0) {
    throw new Error("QUIXOS_RUNTIME_PORT must be a positive integer");
  }

  const functionRefs = params.functions.map((entry) =>
    create(FunctionRefSchema, {
      packageNamespace: params.packageNamespace,
      packageName: params.packageName,
      symbol: entry.symbol,
      operation: entry.operation ?? "",
    }),
  );

  const handler = connectNodeAdapter({
    routes: (router) => {
      router.service(PackageRuntime, {
        async handshake() {
          return create(HandshakeResponseSchema, {
            packageNamespace: params.packageNamespace,
            packageName: params.packageName,
            runtimeProtocolVersion: params.runtimeProtocolVersion,
            functions: functionRefs,
          });
        },

        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}` : ""
                }`,
              );
            }
            return create(InvokeResponseSchema, {
              ok: true,
              result: jsToProtoValue(
                await runtimeFunction(params.createContext(trackingCamino, request)),
              ),
            });
          } catch (error) {
            return create(InvokeResponseSchema, {
              ok: false,
              result: jsToProtoValue(null),
              error: error instanceof Error ? error.message : String(error),
            });
          }
        },
      });
    },
  });

  const server = http.createServer((request, response) => {
    void handler(request, response);
  });
  const close = () => server.close(() => process.exit(0));

  server.listen(port, host, () => {
    process.stderr.write(
      `${params.packageName} listening on ${host}:${port}\n`,
    );
  });
  process.on("SIGTERM", close);
  process.on("SIGINT", close);
};