@gmickel/gno 1.46.0 → 2.1.0

package/src/llm/native-worker/client.ts

@@ -0,0 +1,541 @@
+import type { ModelLifecycleStats } from "../nodeLlamaCpp/lifecycle";
+import type { InferenceOptions } from "../types";
+import type { Owner, Pending } from "./owner";
+import type { ApprovedModel, NativeRequest, NativeResponse } from "./protocol";
+import type { NativeRuntimeConfig } from "./runtime-config";
+
+import { InferenceSettlement } from "../inference-cancellation";
+import {
+  inferenceOptions,
+  recordInferenceTimeout,
+  isBackgroundInference,
+} from "../inference-scope";
+import { NativeWorkerError } from "./errors";
+import { retireOwnedChild, type NativeDisposeOptions } from "./owned-exit";
+import {
+  cancelPending,
+  wireResult,
+  waitForQuarantine,
+  releaseQuarantine,
+  selectNext,
+  recordCompletion,
+} from "./owner";
+import {
+  frameNativeMessage,
+  NativeFrameDecoder,
+  NativeRequestLedger,
+  parseNativeRequest,
+  NativeExecutionStartedSchema,
+} from "./protocol";
+import {
+  nativeWorkerEnvironment,
+  NativeRuntimeConfigSchema,
+  NativeRegistrationSchema,
+  wireError,
+} from "./runtime-config";
+
+type Input = NativeRequest extends infer R
+  ? R extends NativeRequest
+    ? Omit<R, "version" | "generation" | "requestId">
+    : never
+  : never;
+export interface NativeWorkerClientOptions {
+  models: readonly ApprovedModel[];
+  loadTimeout: number;
+  inferenceTimeout: number;
+  warmModelTtl?: number;
+  /** Internal fault-test launch seam; never populated from caller configuration. */
+  entryPath?: string;
+}
+
+/** One process generation, one active operation and 64 waiting logical calls. */
+export class NativeWorkerClient {
+  private owner?: Owner;
+  private lifecycle?: ModelLifecycleStats;
+  private generation = 0;
+  private requestId = 0;
+  private closed = false;
+  private leases = 0;
+  private idleOwner?: Owner;
+  private registration: Promise<void> = Promise.resolve();
+  private readonly options: NativeWorkerClientOptions;
+
+  constructor(options: NativeWorkerClientOptions) {
+    this.options = { ...options, models: structuredClone(options.models) };
+    process.once("exit", this.onParentExit);
+  }
+
+  private readonly onParentExit = (): void => {
+    this.owner?.child.kill("SIGKILL");
+  };
+
+  get processId(): number | undefined {
+    return this.owner?.child.pid;
+  }
+  get currentGeneration(): number {
+    return this.generation;
+  }
+
+  /** Last child-reported snapshot; reading it sends no IPC and cannot extend idle. */
+  getLifecycleStats(): ModelLifecycleStats {
+    const snapshot = this.lifecycle ?? {
+      activeLeases: 0,
+      leaseAcquisitions: 0,
+      leaseReleases: 0,
+      loadedModels: 0,
+      loadAttempts: 0,
+      loadSuccesses: 0,
+      loadFailures: 0,
+      inflightLoads: 0,
+    };
+    return this.owner
+      ? { ...snapshot }
+      : {
+          ...snapshot,
+          activeLeases: 0,
+          loadedModels: 0,
+          inflightLoads: 0,
+        };
+  }
+
+  acquireLease(): { release(): void } {
+    if (this.closed) throw new NativeWorkerError("exited");
+    this.leases++;
+    let released = false;
+    return {
+      release: () => {
+        if (released) return;
+        released = true;
+        this.leases--;
+        const owner = this.idleOwner;
+        if (
+          !this.leases &&
+          owner &&
+          this.owner === owner &&
+          !owner.pending.length &&
+          !owner.busy
+        )
+          void this.retire(owner).catch((error: unknown) =>
+            console.error(error)
+          );
+      },
+    };
+  }
+
+  async disposeModel(uri: string): Promise<void> {
+    const operation = this.registration.then(async () => {
+      const owner = this.owner;
+      if (
+        !owner ||
+        !this.options.models.some((entry) => entry.modelUri === uri)
+      )
+        return;
+      if (owner.retiring) {
+        await owner.retirement;
+        return;
+      }
+      if (owner.pending.length || owner.busy)
+        await new Promise<void>((resolve) => owner.drain.add(resolve));
+      // Explicit disposal reclaims the complete process allocation. Individual
+      // model idle timers may already have evicted weights within this owner.
+      await this.retire(owner);
+    });
+    this.registration = operation.catch(() => {});
+    await operation;
+  }
+
+  /** Only the parent model-selection/policy layer may approve these descriptors. */
+  registerModel(model: ApprovedModel): Promise<void> {
+    const descriptor = NativeRegistrationSchema.parse({
+      register: model,
+    }).register;
+    const operation = this.registration.then(async () => {
+      if (this.closed) throw new NativeWorkerError("exited");
+      const existing = this.options.models.find(
+        (entry) => entry.id === descriptor.id
+      );
+      if (existing && JSON.stringify(existing) === JSON.stringify(descriptor))
+        return;
+      const models = [
+        ...this.options.models.filter((entry) => entry.id !== descriptor.id),
+        descriptor,
+      ];
+      NativeRuntimeConfigSchema.parse({
+        generation: Math.max(1, this.generation),
+        models,
+        loadTimeout: this.options.loadTimeout,
+        inferenceTimeout: this.options.inferenceTimeout,
+        warmModelTtl: this.options.warmModelTtl ?? 300_000,
+      });
+      const owner = this.owner;
+      if (owner?.retiring) await owner.retirement;
+      if (existing && owner) {
+        if (owner.pending.length || owner.busy)
+          await new Promise<void>((resolve) => owner.drain.add(resolve));
+        await this.retire(owner);
+      } else if (owner && !owner.retiring) {
+        owner.child.send({ register: descriptor });
+      }
+      this.options.models = models;
+    });
+    this.registration = operation.catch(() => {});
+    return operation;
+  }
+
+  async request(
+    input: Input,
+    expectedGeneration?: number,
+    options?: InferenceOptions
+  ): Promise<NativeResponse["result"]> {
+    const operational = inferenceOptions({
+      ...options,
+      deadlineAt: options?.deadlineAt ?? input.deadlineAt,
+    });
+    const settlement = new InferenceSettlement<
+      Extract<NativeResponse["result"], { ok: true }>["value"]
+    >(operational, this.options.inferenceTimeout);
+    void this.enqueue(
+      { ...input, deadlineAt: operational.deadlineAt },
+      expectedGeneration,
+      settlement
+    ).catch(() => {
+      settlement.fail(new NativeWorkerError("exited").detail);
+    });
+    return settlement.completion.then((result) => {
+      if (!result.ok && result.error.code === "TIMEOUT")
+        recordInferenceTimeout();
+      return wireResult(result);
+    });
+  }
+
+  private async enqueue(
+    input: Input,
+    expectedGeneration: number | undefined,
+    settlement: Pending["settlement"]
+  ): Promise<void> {
+    const result = await this.admit(input, expectedGeneration, settlement);
+    if (settlement.phase === "queued") settlement.startNative();
+    settlement.nativeSettled(result);
+    settlement.publish();
+  }
+
+  private async admit(
+    input: Input,
+    expectedGeneration: number | undefined,
+    settlement: Pending["settlement"]
+  ): Promise<NativeResponse["result"]> {
+    await this.registration;
+    if (settlement.signal.aborted)
+      return settlement.completion.then(wireResult);
+    if (this.closed)
+      return {
+        ok: false,
+        error: wireError(new NativeWorkerError("exited").detail),
+      };
+    while (this.owner?.quarantined && !this.owner.retiring) {
+      try {
+        await waitForQuarantine(this.owner, settlement);
+      } catch (cause) {
+        return {
+          ok: false,
+          error: wireError(
+            (cause instanceof NativeWorkerError
+              ? cause
+              : new NativeWorkerError("exited")
+            ).detail
+          ),
+        };
+      }
+      if (settlement.signal.aborted)
+        return settlement.completion.then(wireResult);
+    }
+    if (this.owner?.retiring) await this.owner.retirement;
+    if (this.closed)
+      return {
+        ok: false,
+        error: wireError(new NativeWorkerError("exited").detail),
+      };
+    if (settlement.signal.aborted)
+      return settlement.completion.then(wireResult);
+    if (input.op === "dispose" && !this.owner) return { ok: true, value: null };
+    try {
+      // Metadata-dependent inference must never start/replay on a replacement
+      // owner. All asynchronous registration/retirement waits precede this fence.
+      if (
+        expectedGeneration !== undefined &&
+        (!this.owner ||
+          this.owner.retiring ||
+          this.owner.generation !== expectedGeneration)
+      )
+        throw new NativeWorkerError("stale_generation");
+      const owner = this.owner ?? this.start();
+      const request = parseNativeRequest(
+        {
+          ...input,
+          deadlineAt: input.deadlineAt,
+          version: 1,
+          generation: owner.generation,
+          requestId: ++this.requestId,
+        },
+        owner.generation,
+        this.options.models
+      );
+      owner.ledger.admit(request);
+      return new Promise((resolve) => {
+        const pending: Pending = {
+          request,
+          resolve,
+          settlement,
+          background: isBackgroundInference(),
+        };
+        owner.pending.push(pending);
+        settlement.signal.addEventListener(
+          "abort",
+          () => this.cancel(owner, pending),
+          { once: true }
+        );
+        if (settlement.signal.aborted) this.cancel(owner, pending);
+        if (owner.pending.length === 1 && owner.ready) this.sendNext(owner);
+      });
+    } catch (cause) {
+      return {
+        ok: false,
+        error: wireError(
+          (cause instanceof NativeWorkerError
+            ? cause
+            : new NativeWorkerError("exited")
+          ).detail
+        ),
+      };
+    }
+  }
+
+  private start(): Owner {
+    this.lifecycle = undefined;
+    // A compiled executable cannot interpret an external TS entry: invoking it
+    // here would recursively launch the CLI. npm/desktop ship the source runtime.
+    if (import.meta.dir.includes("$bunfs"))
+      throw new NativeWorkerError("exited");
+    const config: NativeRuntimeConfig = NativeRuntimeConfigSchema.parse({
+      generation: ++this.generation,
+      models: this.options.models,
+      loadTimeout: this.options.loadTimeout,
+      inferenceTimeout: this.options.inferenceTimeout,
+      warmModelTtl: this.options.warmModelTtl ?? 300_000,
+    });
+    // IPC callbacks execute after this synchronous ownership publication.
+    const child = Bun.spawn({
+      cmd: [
+        process.execPath,
+        "--no-env-file",
+        this.options.entryPath ?? `${import.meta.dir}/entry.ts`,
+        JSON.stringify(config),
+      ],
+      env: nativeWorkerEnvironment(),
+      stdin: "ignore",
+      stdout: "ignore",
+      stderr: "inherit",
+      serialization: "advanced",
+      ipc: (message) => this.receive(owner, message),
+    });
+    const owner: Owner = {
+      generation: config.generation,
+      child,
+      ledger: new NativeRequestLedger(config.generation),
+      decoder: new NativeFrameDecoder(config.generation),
+      pending: [],
+      foregroundCompletions: 0,
+      ready: false,
+      busy: false,
+      quarantined: false,
+      waiters: new Set(),
+      retiring: false,
+      drain: new Set(),
+    };
+    this.owner = owner;
+    owner.timer = setTimeout(
+      () => this.fail(owner, new NativeWorkerError("timeout")),
+      this.options.loadTimeout
+    );
+    void child.exited.then(() => {
+      this.fail(owner, new NativeWorkerError("exited"));
+      this.finishOwner(owner, new NativeWorkerError("exited"));
+      if (this.owner === owner) this.owner = undefined;
+      releaseQuarantine(owner);
+    });
+    return owner;
+  }
+
+  private receive(owner: Owner, message: unknown): void {
+    if (owner.retiring || this.owner !== owner) return;
+    try {
+      if (message === "ready") {
+        if (owner.ready) throw new NativeWorkerError("protocol");
+        owner.ready = true;
+        clearTimeout(owner.timer);
+        this.sendNext(owner);
+        return;
+      }
+      if (message === "idle") {
+        this.idleOwner = owner;
+        if (!this.leases && !owner.pending.length && !owner.busy)
+          void this.retire(owner).catch((error: unknown) =>
+            console.error(error)
+          );
+        return;
+      }
+      if (
+        typeof message === "object" &&
+        message !== null &&
+        "executionStarted" in message
+      ) {
+        const parsed = NativeExecutionStartedSchema.safeParse(message);
+        const pending = owner.pending[0];
+        if (
+          !parsed.success ||
+          !owner.busy ||
+          !pending ||
+          pending.executionStarted ||
+          ["init", "dispose"].includes(pending.request.op) ||
+          parsed.data.generation !== owner.generation ||
+          parsed.data.requestId !== pending.request.requestId
+        )
+          throw new NativeWorkerError("protocol");
+        if (
+          pending.loadDeadline !== undefined &&
+          performance.now() >= pending.loadDeadline
+        )
+          pending.settlement.cancel("timeout");
+        pending.loadDeadline = undefined;
+        pending.executionStarted = true;
+        clearTimeout(owner.timer);
+        pending.settlement.startExecution();
+        return;
+      }
+      if (!(message instanceof Uint8Array))
+        throw new NativeWorkerError("protocol");
+      const decoded = owner.decoder.push(message);
+      if (decoded === undefined) return;
+      const response = owner.ledger.settle(decoded);
+      const pending = owner.pending[0];
+      if (pending?.request.requestId !== response.requestId)
+        throw new NativeWorkerError("protocol");
+      clearTimeout(owner.timer);
+      if (
+        pending.loadDeadline !== undefined &&
+        performance.now() >= pending.loadDeadline
+      )
+        pending.settlement.cancel("timeout");
+      clearTimeout(pending.cancelTimer);
+      owner.pending.shift();
+      recordCompletion(owner, pending);
+      releaseQuarantine(owner);
+      if (response.lifecycle) this.lifecycle = response.lifecycle;
+      pending.resolve(response.result);
+      // Promise delivery has been queued before acknowledging settlement.
+      queueMicrotask(() => {
+        if (owner.retiring) return;
+        try {
+          owner.child.send({ ack: response.requestId });
+          owner.busy = false;
+          if (!owner.pending.length) {
+            for (const resolve of owner.drain) resolve();
+            owner.drain.clear();
+          }
+          this.sendNext(owner);
+        } catch {
+          this.fail(owner, new NativeWorkerError("exited"));
+        }
+      });
+    } catch (cause) {
+      this.fail(
+        owner,
+        cause instanceof NativeWorkerError
+          ? cause
+          : new NativeWorkerError("protocol")
+      );
+    }
+  }
+
+  private sendNext(owner: Owner): void {
+    if (!owner.ready || owner.retiring || owner.busy) return;
+    selectNext(owner);
+    const pending = owner.pending[0];
+    if (!pending) return;
+    this.idleOwner = undefined;
+    if (!pending.settlement.startNative()) {
+      this.cancel(owner, pending);
+      return;
+    }
+    owner.busy = true;
+    pending.loadDeadline = performance.now() + this.options.loadTimeout;
+    clearTimeout(owner.timer);
+    owner.timer = setTimeout(
+      () => pending.settlement.cancel("timeout"),
+      this.options.loadTimeout
+    );
+    try {
+      for (const frame of frameNativeMessage(pending.request))
+        owner.child.send(frame);
+    } catch {
+      this.fail(owner, new NativeWorkerError("exited"));
+    }
+  }
+
+  private cancel(owner: Owner, pending: Pending): void {
+    cancelPending(
+      owner,
+      pending,
+      (current, error) => this.fail(current, error),
+      (current) => this.sendNext(current)
+    );
+  }
+
+  private fail(owner: Owner, error: NativeWorkerError): void {
+    if (owner.retiring) return;
+    // Delivery can fail now; ownership is acknowledged only after child exit.
+    for (const pending of owner.pending) {
+      clearTimeout(pending.cancelTimer);
+      pending.settlement.fail(error.detail);
+    }
+    void this.retire(owner, { force: true }).catch((cause: unknown) =>
+      console.error(cause)
+    );
+  }
+
+  private retire(
+    owner: Owner,
+    options: NativeDisposeOptions = {}
+  ): Promise<void> {
+    return retireOwnedChild(owner, options).then(() => {
+      if (this.owner === owner) this.owner = undefined;
+      releaseQuarantine(owner);
+    });
+  }
+
+  async dispose(options: NativeDisposeOptions = {}): Promise<void> {
+    this.closed = true;
+    const owner = this.owner;
+    if (!owner) {
+      process.removeListener("exit", this.onParentExit);
+      return;
+    }
+    const failure = new NativeWorkerError("exited");
+    for (const pending of owner.pending) {
+      clearTimeout(pending.cancelTimer);
+      pending.settlement.fail(failure.detail);
+    }
+    await this.retire(owner, options);
+    process.removeListener("exit", this.onParentExit);
+    this.finishOwner(owner, failure);
+  }
+
+  private finishOwner(owner: Owner, failure: NativeWorkerError): void {
+    for (const response of owner.ledger.failAll(failure))
+      owner.pending
+        .find((entry) => entry.request.requestId === response.requestId)
+        ?.resolve(response.result);
+    owner.pending.length = 0;
+    owner.busy = false;
+  }
+}