@holokai/neuron-sdk 0.1.0

package/dist/index.mjs

@@ -0,0 +1,1108 @@
+import { z } from 'zod';
+import { EventEmitter } from 'eventemitter3';
+import { zodToJsonSchema } from 'zod-to-json-schema';
+import AjvDefault from 'ajv';
+import addFormatsDefault from 'ajv-formats';
+
+// ../neuron-protocol/src/capability.ts
+var CapabilityTypeSchema = z.string().min(1).regex(
+  /^[A-Za-z0-9_.\-]+(\/[A-Za-z0-9_.\-]+)+$/,
+  "capability type must be namespaced: {owner}/{package}[/{name}][.{variant}]"
+);
+var UserIdSchema = z.string().min(1);
+var CapabilityScopeSchema = z.union([
+  z.literal("any"),
+  z.object({ onBehalfOf: UserIdSchema }).strict()
+]);
+var CapabilitySchema = z.object({
+  type: CapabilityTypeSchema,
+  scope: CapabilityScopeSchema,
+  concurrency: z.number().int().positive()
+}).strict();
+var JsonSchemaSchema = z.unknown();
+var LeaseIdSchema = z.string().min(1);
+var TaskPayloadSchema = z.object({
+  taskId: z.string().min(1),
+  workflowId: z.string().min(1),
+  capability: CapabilityTypeSchema,
+  onBehalfOf: UserIdSchema.optional(),
+  input: z.unknown(),
+  attempt: z.number().int().nonnegative(),
+  metadata: z.record(z.string(), z.unknown()).optional()
+}).strict();
+var LeasePayloadSchema = z.object({
+  leaseId: LeaseIdSchema,
+  expiresAt: z.string().datetime(),
+  task: TaskPayloadSchema,
+  inputSchema: JsonSchemaSchema,
+  outputSchema: JsonSchemaSchema
+}).strict();
+var NeuronIdSchema = z.string().min(1);
+var RegisterFrameSchema = z.object({
+  type: z.literal("register"),
+  neuronId: NeuronIdSchema,
+  // Free-form by design — the gateway negotiates compatibility against its
+  // own PROTOCOL_VERSION (semver), so a v1.1 server can still accept a v1.0
+  // neuron. Do not pin this to the exported constant.
+  protocolVersion: z.string().min(1),
+  capabilities: z.array(CapabilitySchema).min(1).refine(
+    (caps) => new Set(caps.map((c) => c.type)).size === caps.length,
+    "duplicate capability types are not allowed in a single register frame"
+  )
+}).strict();
+var UpdateCapabilityFrameSchema = z.object({
+  type: z.literal("update-capability"),
+  neuronId: NeuronIdSchema,
+  capability: CapabilitySchema
+}).strict();
+var HeartbeatInFlightEntrySchema = z.object({
+  leaseId: LeaseIdSchema,
+  taskId: z.string().min(1),
+  startedAt: z.string().datetime()
+}).strict();
+var HeartbeatFrameSchema = z.object({
+  type: z.literal("heartbeat"),
+  neuronId: NeuronIdSchema,
+  inFlight: z.array(HeartbeatInFlightEntrySchema),
+  sentAt: z.string().datetime()
+}).strict();
+var AckFrameSchema = z.object({
+  type: z.literal("ack"),
+  leaseId: LeaseIdSchema,
+  output: z.unknown(),
+  completedAt: z.string().datetime()
+}).strict();
+var NackKindSchema = z.enum([
+  "retryable",
+  "terminal",
+  "schema-mismatch",
+  "cancelled"
+]);
+var NackFrameSchema = z.object({
+  type: z.literal("nack"),
+  leaseId: LeaseIdSchema,
+  kind: NackKindSchema,
+  reason: z.string().min(1),
+  code: z.string().optional(),
+  details: z.unknown().optional(),
+  failedAt: z.string().datetime()
+}).strict();
+var ProgressFrameSchema = z.object({
+  type: z.literal("progress"),
+  leaseId: LeaseIdSchema,
+  percent: z.number().min(0).max(100).optional(),
+  message: z.string().optional(),
+  data: z.unknown().optional(),
+  sentAt: z.string().datetime()
+}).strict();
+z.discriminatedUnion("type", [
+  RegisterFrameSchema,
+  UpdateCapabilityFrameSchema,
+  HeartbeatFrameSchema,
+  AckFrameSchema,
+  NackFrameSchema,
+  ProgressFrameSchema
+]);
+var RegisteredFrameSchema = z.object({
+  type: z.literal("registered"),
+  sessionId: z.string().min(1),
+  // Effective capability bindings after auth + policy filters; may be a
+  // subset of what the neuron advertised.
+  capabilities: z.array(CapabilitySchema),
+  acceptedAt: z.string().datetime()
+}).strict();
+var LeaseFrameSchema = z.object({
+  type: z.literal("lease"),
+  lease: LeasePayloadSchema,
+  sentAt: z.string().datetime()
+}).strict();
+var CancelFrameSchema = z.object({
+  type: z.literal("cancel"),
+  leaseId: LeaseIdSchema,
+  reason: z.string().min(1),
+  cancelledAt: z.string().datetime()
+}).strict();
+var CapabilityRevokedFrameSchema = z.object({
+  type: z.literal("capability-revoked"),
+  capability: CapabilityTypeSchema,
+  reason: z.string().min(1),
+  revokedAt: z.string().datetime()
+}).strict();
+var DisconnectFrameSchema = z.object({
+  type: z.literal("disconnect"),
+  reason: z.string().min(1),
+  code: z.string().optional(),
+  reconnectable: z.boolean(),
+  sentAt: z.string().datetime()
+}).strict();
+var ServerToNeuronFrameSchema = z.discriminatedUnion("type", [
+  RegisteredFrameSchema,
+  LeaseFrameSchema,
+  CancelFrameSchema,
+  CapabilityRevokedFrameSchema,
+  DisconnectFrameSchema
+]);
+z.discriminatedUnion("type", [
+  RegisterFrameSchema,
+  UpdateCapabilityFrameSchema,
+  HeartbeatFrameSchema,
+  AckFrameSchema,
+  NackFrameSchema,
+  ProgressFrameSchema,
+  RegisteredFrameSchema,
+  LeaseFrameSchema,
+  CancelFrameSchema,
+  CapabilityRevokedFrameSchema,
+  DisconnectFrameSchema
+]);
+
+// ../neuron-protocol/src/transport.ts
+var PROTOCOL_VERSION = "1.0.0";
+var SSE_PATH = "/neuron/events";
+var SSE_EVENT = {
+  registered: "registered",
+  lease: "lease",
+  cancel: "cancel",
+  capabilityRevoked: "capability-revoked",
+  disconnect: "disconnect"
+};
+var POST_PATH = {
+  register: "/neuron/register",
+  updateCapability: "/neuron/update-capability",
+  heartbeat: "/neuron/heartbeat",
+  ack: "/neuron/ack",
+  nack: "/neuron/nack",
+  progress: "/neuron/progress"
+};
+function defineCapability(init) {
+  return {
+    type: init.type,
+    description: init.description,
+    inputSchema: zodToJsonSchema(init.input, { $refStrategy: "none" }),
+    outputSchema: zodToJsonSchema(init.output, { $refStrategy: "none" }),
+    scope: init.scope,
+    concurrency: init.concurrency,
+    _zod: { input: init.input, output: init.output }
+  };
+}
+function toPlanShape(def) {
+  return {
+    type: def.type,
+    description: def.description,
+    input: def.inputSchema,
+    output: def.outputSchema
+  };
+}
+function toRegistrationCapability(def, opts) {
+  const scope = opts?.scope ?? def.scope;
+  const concurrency = opts?.concurrency ?? def.concurrency;
+  if (!scope) {
+    throw new Error(
+      `capability ${def.type}: scope is required for registration \u2014 set it on defineCapability(...) or at neuron.handle(def, { scope }, fn)`
+    );
+  }
+  if (typeof concurrency !== "number") {
+    throw new Error(
+      `capability ${def.type}: concurrency is required for registration \u2014 set it on defineCapability(...) or at neuron.handle(def, { concurrency }, fn)`
+    );
+  }
+  return { type: def.type, scope, concurrency };
+}
+
+// src/heartbeat.ts
+var DEFAULT_INTERVAL_MS = 1e4;
+var HeartbeatLoop = class {
+  constructor(opts) {
+    this.opts = opts;
+    this.intervalMs = opts.intervalMs ?? DEFAULT_INTERVAL_MS;
+  }
+  timer = null;
+  running = false;
+  intervalMs;
+  start() {
+    if (this.timer) return;
+    this.timer = setInterval(() => {
+      void this.tick();
+    }, this.intervalMs);
+    if (typeof this.timer.unref === "function") this.timer.unref();
+  }
+  stop() {
+    if (this.timer) {
+      clearInterval(this.timer);
+      this.timer = null;
+    }
+  }
+  /** Public for tests so they don't have to wait for the timer. */
+  async tick() {
+    if (this.running) return;
+    if (!this.opts.transport.isConnected()) return;
+    this.running = true;
+    try {
+      await this.opts.transport.post(POST_PATH.heartbeat, {
+        type: "heartbeat",
+        neuronId: this.opts.neuronId,
+        inFlight: this.opts.inFlight(),
+        sentAt: (/* @__PURE__ */ new Date()).toISOString()
+      });
+    } catch (err) {
+      this.opts.logger?.warn({ err }, "heartbeat post failed; will retry next tick");
+    } finally {
+      this.running = false;
+    }
+  }
+};
+
+// src/lease-runner.ts
+var TerminalFailure = class extends Error {
+  constructor(reason, code) {
+    super(reason);
+    this.reason = reason;
+    this.code = code;
+    this.name = "TerminalFailure";
+  }
+};
+function validateInput(zod, validator, jsonSchema, data) {
+  if (zod) {
+    const result = zod.safeParse(data);
+    return result.success ? { ok: true, value: result.data } : { ok: false, errors: result.error.issues };
+  }
+  const errors = validator.validate(
+    jsonSchema,
+    data
+  );
+  return errors === null ? { ok: true, value: data } : { ok: false, errors };
+}
+async function runLease(opts) {
+  const { lease, handler, zodInput, zodOutput, transport, validator, logger, signal } = opts;
+  const inputCheck = validateInput(zodInput, validator, lease.inputSchema, lease.task.input);
+  if (!inputCheck.ok) {
+    await sendNack(
+      transport,
+      lease.leaseId,
+      "schema-mismatch",
+      "lease input does not match inputSchema",
+      { code: "INPUT_SCHEMA_MISMATCH", details: { errors: inputCheck.errors } },
+      logger
+    );
+    return;
+  }
+  const parsedInput = inputCheck.value;
+  const progressFn = makeProgressFn(transport, lease.leaseId, logger);
+  const ctx = {
+    task: lease.task,
+    leaseId: lease.leaseId,
+    signal,
+    log: scopedLogger(logger, lease),
+    progress: progressFn,
+    fail: (reason, opts2) => {
+      throw new TerminalFailure(reason, opts2?.code);
+    }
+  };
+  let output;
+  try {
+    output = await handler(parsedInput, ctx);
+  } catch (err) {
+    if (signal.aborted || err.name === "AbortError") {
+      await sendNack(
+        transport,
+        lease.leaseId,
+        "cancelled",
+        "handler aborted",
+        { code: "CANCELLED" },
+        logger
+      );
+      return;
+    }
+    if (err instanceof TerminalFailure) {
+      await sendNack(
+        transport,
+        lease.leaseId,
+        "terminal",
+        err.reason,
+        { code: err.code ?? "HANDLER_TERMINAL" },
+        logger
+      );
+      return;
+    }
+    const message = err instanceof Error ? err.message : String(err);
+    logger.error({ err, leaseId: lease.leaseId }, "handler threw; nacking retryable");
+    await sendNack(
+      transport,
+      lease.leaseId,
+      "retryable",
+      message,
+      { code: "HANDLER_THREW" },
+      logger
+    );
+    return;
+  }
+  const outputCheck = validateInput(zodOutput, validator, lease.outputSchema, output);
+  if (!outputCheck.ok) {
+    await sendNack(
+      transport,
+      lease.leaseId,
+      "schema-mismatch",
+      "handler output does not match outputSchema",
+      { code: "OUTPUT_SCHEMA_MISMATCH", details: { errors: outputCheck.errors } },
+      logger
+    );
+    return;
+  }
+  const parsedOutput = outputCheck.value;
+  await transport.post(POST_PATH.ack, {
+    type: "ack",
+    leaseId: lease.leaseId,
+    output: parsedOutput,
+    completedAt: (/* @__PURE__ */ new Date()).toISOString()
+  }).catch((err) => {
+    logger.warn(
+      { err, leaseId: lease.leaseId },
+      "ack POST rejected by gateway; lease will be reclaimed by reaper"
+    );
+  });
+}
+async function sendNack(transport, leaseId, kind, reason, extras = {}, logger) {
+  await transport.post(POST_PATH.nack, {
+    type: "nack",
+    leaseId,
+    kind,
+    reason,
+    code: extras.code,
+    details: extras.details,
+    failedAt: (/* @__PURE__ */ new Date()).toISOString()
+  }).catch((err) => {
+    logger?.warn(
+      { err, leaseId, kind },
+      "nack POST rejected by gateway; lease will be reclaimed by reaper"
+    );
+  });
+}
+function makeProgressFn(transport, leaseId, logger) {
+  return (update) => {
+    void transport.post(POST_PATH.progress, {
+      type: "progress",
+      leaseId,
+      percent: update.percent,
+      message: update.message,
+      data: update.data,
+      sentAt: (/* @__PURE__ */ new Date()).toISOString()
+    }).catch((err) => {
+      logger.warn({ err, leaseId }, "progress POST failed; continuing");
+    });
+  };
+}
+function scopedLogger(logger, lease) {
+  const scope = { leaseId: lease.leaseId, taskId: lease.task.taskId };
+  return {
+    debug: (meta, msg) => logger.debug({ ...scope, ...meta }, msg),
+    info: (meta, msg) => logger.info({ ...scope, ...meta }, msg),
+    warn: (meta, msg) => logger.warn({ ...scope, ...meta }, msg),
+    error: (meta, msg) => logger.error({ ...scope, ...meta }, msg)
+  };
+}
+
+// src/logger.ts
+var noopLogger = {
+  debug() {
+  },
+  info() {
+  },
+  warn() {
+  },
+  error() {
+  }
+};
+var Ajv = AjvDefault.default ?? AjvDefault;
+var addFormats = addFormatsDefault.default ?? addFormatsDefault;
+var SchemaValidator = class {
+  ajv;
+  cache = /* @__PURE__ */ new WeakMap();
+  constructor() {
+    this.ajv = new Ajv({ allErrors: true, strict: false });
+    addFormats(this.ajv);
+  }
+  /**
+   * Returns null if the value matches the schema (or no schema given);
+   * otherwise an array of structured errors.
+   */
+  validate(schema, data) {
+    if (!schema || typeof schema !== "object") return null;
+    let validator = this.cache.get(schema);
+    if (!validator) {
+      try {
+        validator = this.ajv.compile(schema);
+      } catch (err) {
+        return [
+          {
+            instancePath: "",
+            schemaPath: "",
+            keyword: "schema",
+            params: {},
+            message: `schema failed to compile: ${err instanceof Error ? err.message : String(err)}`
+          }
+        ];
+      }
+      this.cache.set(schema, validator);
+    }
+    const v = validator;
+    const ok = v(data);
+    return ok ? null : v.errors ?? [];
+  }
+};
+
+// src/sse-parser.ts
+async function* parseSSEStream(stream, signal) {
+  const reader = stream.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+  const onAbort = () => {
+    void reader.cancel().catch(() => void 0);
+  };
+  signal?.addEventListener("abort", onAbort);
+  try {
+    while (true) {
+      if (signal?.aborted) return;
+      const { value, done } = await reader.read();
+      if (done) break;
+      buffer += decoder.decode(value, { stream: true }).replace(/\r\n/g, "\n").replace(/\r/g, "\n");
+      let blockEnd;
+      while ((blockEnd = buffer.indexOf("\n\n")) !== -1) {
+        const block = buffer.slice(0, blockEnd);
+        buffer = buffer.slice(blockEnd + 2);
+        const event = parseBlock(block);
+        if (event) yield event;
+      }
+    }
+  } finally {
+    signal?.removeEventListener("abort", onAbort);
+    try {
+      reader.releaseLock();
+    } catch {
+    }
+  }
+}
+function parseBlock(block) {
+  if (!block.trim()) return null;
+  let event = "message";
+  const dataLines = [];
+  for (const rawLine of block.split("\n")) {
+    const line = rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine;
+    if (line.startsWith(":")) continue;
+    if (line.startsWith("event:")) {
+      event = line.slice("event:".length).trim();
+    } else if (line.startsWith("data:")) {
+      dataLines.push(line.slice("data:".length).replace(/^ /, ""));
+    }
+  }
+  if (dataLines.length === 0) return null;
+  return { event, data: dataLines.join("\n") };
+}
+
+// src/transport.ts
+var Transport = class {
+  opts;
+  callbacks;
+  abortController = null;
+  stopped = false;
+  connected = false;
+  reconnectAttempts = 0;
+  reconnectTimer = null;
+  constructor(opts) {
+    this.opts = {
+      gatewayUrl: opts.gatewayUrl.replace(/\/+$/, ""),
+      neuronId: opts.neuronId,
+      auth: opts.auth,
+      initialReconnectDelayMs: opts.initialReconnectDelayMs ?? 500,
+      maxReconnectDelayMs: opts.maxReconnectDelayMs ?? 3e4,
+      logger: opts.logger ?? noopLogger,
+      fetch: opts.fetch ?? globalThis.fetch.bind(globalThis)
+    };
+    this.callbacks = {
+      onOpen: opts.onOpen,
+      onClose: opts.onClose,
+      onFrame: opts.onFrame,
+      onError: opts.onError,
+      onReconnecting: opts.onReconnecting
+    };
+  }
+  isConnected() {
+    return this.connected;
+  }
+  /** Open the SSE stream. Resolves after the connection is established. */
+  async start() {
+    this.stopped = false;
+    await this.openStream();
+  }
+  /** Tear down the SSE stream and prevent further reconnect attempts. */
+  async stop() {
+    this.stopped = true;
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer);
+      this.reconnectTimer = null;
+    }
+    if (this.abortController) {
+      this.abortController.abort();
+      this.abortController = null;
+    }
+    this.connected = false;
+  }
+  /**
+   * POST a client frame to the gateway. Refreshes the token via the auth
+   * callback and retries once on 401 — that's the contract: callers
+   * implement refresh inside their callback.
+   */
+  async post(path, frame) {
+    return this.postRaw(path, frame, true);
+  }
+  async postRaw(path, frame, allowRetry) {
+    const token = await this.opts.auth();
+    const url = `${this.opts.gatewayUrl}${path}`;
+    const res = await this.opts.fetch(url, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Authorization: `Bearer ${token}`
+      },
+      body: JSON.stringify(frame)
+    });
+    if (res.status === 401 && allowRetry) {
+      this.opts.logger.warn({ path }, "gateway returned 401; retrying after auth refresh");
+      return this.postRaw(path, frame, false);
+    }
+    if (!res.ok) {
+      const text = await res.text().catch(() => "");
+      throw new TransportError(
+        `${path} failed: ${res.status} ${res.statusText}${text ? ` \u2014 ${text}` : ""}`,
+        res.status
+      );
+    }
+    return res.headers.get("content-type")?.includes("application/json") ? await res.json() : await res.text();
+  }
+  async openStream(allowAuthRetry = true) {
+    if (this.stopped) return;
+    const url = `${this.opts.gatewayUrl}${SSE_PATH}?neuronId=${encodeURIComponent(
+      this.opts.neuronId
+    )}`;
+    this.abortController = new AbortController();
+    const signal = this.abortController.signal;
+    let res;
+    try {
+      const token = await this.opts.auth();
+      res = await this.opts.fetch(url, {
+        method: "GET",
+        headers: {
+          Accept: "text/event-stream",
+          Authorization: `Bearer ${token}`
+        },
+        signal
+      });
+    } catch (err) {
+      this.scheduleReconnect(`fetch error: ${err.message}`);
+      return;
+    }
+    if (res.status === 401 && allowAuthRetry) {
+      this.opts.logger.warn(
+        { reason: "SSE 401" },
+        "SSE open returned 401; retrying after auth refresh"
+      );
+      this.abortController = null;
+      void res.body?.cancel().catch(() => void 0);
+      return this.openStream(false);
+    }
+    if (!res.ok || !res.body) {
+      const reason = `SSE open failed: ${res.status} ${res.statusText}`;
+      this.scheduleReconnect(reason);
+      return;
+    }
+    this.connected = true;
+    this.reconnectAttempts = 0;
+    this.callbacks.onOpen?.();
+    void this.consumeStream(res.body, signal);
+  }
+  async consumeStream(body, signal) {
+    try {
+      for await (const event of parseSSEStream(body, signal)) {
+        let parsed;
+        try {
+          const data = JSON.parse(event.data);
+          parsed = ServerToNeuronFrameSchema.parse(data);
+        } catch (err) {
+          this.opts.logger.warn(
+            { err, event: event.event, data: event.data.slice(0, 200) },
+            "dropped malformed SSE frame"
+          );
+          continue;
+        }
+        this.callbacks.onFrame?.(parsed);
+      }
+      this.handleStreamEnd("clean close");
+    } catch (err) {
+      this.handleStreamEnd(`read error: ${err.message}`);
+    }
+  }
+  handleStreamEnd(reason) {
+    if (!this.connected) return;
+    this.connected = false;
+    this.abortController = null;
+    this.callbacks.onClose?.(reason);
+    if (!this.stopped) {
+      this.scheduleReconnect(reason);
+    }
+  }
+  scheduleReconnect(reason) {
+    if (this.stopped) return;
+    const attempt = this.reconnectAttempts++;
+    const delay = Math.min(
+      this.opts.maxReconnectDelayMs,
+      this.opts.initialReconnectDelayMs * Math.pow(2, attempt)
+    );
+    this.opts.logger.warn(
+      { reason, attempt: attempt + 1, delayMs: delay },
+      "transport reconnecting"
+    );
+    this.callbacks.onReconnecting?.({ attempt: attempt + 1, delayMs: delay, reason });
+    this.reconnectTimer = setTimeout(() => {
+      this.reconnectTimer = null;
+      void this.openStream();
+    }, delay);
+    if (typeof this.reconnectTimer === "object" && this.reconnectTimer !== null && "unref" in this.reconnectTimer && typeof this.reconnectTimer.unref === "function") {
+      this.reconnectTimer.unref();
+    }
+  }
+};
+var TransportError = class extends Error {
+  constructor(message, statusCode) {
+    super(message);
+    this.statusCode = statusCode;
+    this.name = "TransportError";
+  }
+};
+
+// src/worker-pool.ts
+var WorkerPool = class {
+  constructor(maxConcurrency) {
+    this.maxConcurrency = maxConcurrency;
+    if (maxConcurrency <= 0) {
+      throw new Error(
+        `WorkerPool requires positive concurrency, got ${maxConcurrency}`
+      );
+    }
+  }
+  inFlight = 0;
+  queue = [];
+  draining = false;
+  drained = null;
+  /** Currently-running task count. */
+  size() {
+    return this.inFlight;
+  }
+  /** Tasks waiting for a slot. */
+  pending() {
+    return this.queue.length;
+  }
+  /**
+   * Run `fn` once a concurrency slot is available. Resolves with the
+   * function's return value (or rejects with its error). Throws
+   * ShutdownError if the pool is already draining.
+   */
+  async run(fn) {
+    if (this.draining) {
+      throw new ShutdownError("worker pool is draining");
+    }
+    if (this.inFlight >= this.maxConcurrency) {
+      await new Promise((resolve) => this.queue.push(resolve));
+    }
+    this.inFlight++;
+    try {
+      return await fn();
+    } finally {
+      this.inFlight--;
+      const next = this.queue.shift();
+      if (next) {
+        next();
+      } else if (this.draining && this.inFlight === 0 && this.drained) {
+        this.drained.resolve();
+      }
+    }
+  }
+  /**
+   * Update the max concurrency. Lowering doesn't kill in-flight tasks —
+   * they finish naturally — but it does prevent new ones from starting
+   * until the in-flight count drops below the new limit.
+   */
+  setConcurrency(newMax) {
+    if (newMax <= 0) {
+      throw new Error(
+        `WorkerPool concurrency must be positive, got ${newMax}`
+      );
+    }
+    const oldMax = this.maxConcurrency;
+    this.maxConcurrency = newMax;
+    if (newMax > oldMax) {
+      const headroom = newMax - this.inFlight;
+      for (let i = 0; i < headroom && this.queue.length > 0; i++) {
+        const next = this.queue.shift();
+        next?.();
+      }
+    }
+  }
+  /**
+   * Stop accepting new submissions. Returns a promise that resolves once
+   * every in-flight AND queued task has finished — queued waiters from
+   * before drain are released as in-flight slots free, run normally, and
+   * are awaited as part of drain. Only fresh `run()` calls received after
+   * `drain()` reject with ShutdownError. Callers that need a hard stop
+   * should race this promise against a timeout and abort handlers.
+   */
+  drain() {
+    if (!this.draining) {
+      this.draining = true;
+      let resolveFn;
+      const promise = new Promise((resolve) => {
+        resolveFn = resolve;
+      });
+      this.drained = { resolve: resolveFn, promise };
+      if (this.inFlight === 0) {
+        resolveFn();
+      }
+    }
+    return this.drained.promise;
+  }
+};
+var ShutdownError = class extends Error {
+  constructor(message) {
+    super(message);
+    this.name = "ShutdownError";
+  }
+};
+
+// src/neuron.ts
+var Neuron = class extends EventEmitter {
+  transport;
+  handlers = /* @__PURE__ */ new Map();
+  inFlightLeases = /* @__PURE__ */ new Map();
+  /**
+   * Leases admitted by `dispatchLease` but still waiting for a worker-pool
+   * slot. Tracked separately from `inFlightLeases` so the heartbeat snapshot
+   * doesn't lie — only running leases extend the gateway's lease TTL.
+   */
+  queuedLeases = /* @__PURE__ */ new Map();
+  heartbeat;
+  validator;
+  logger;
+  opts;
+  started = false;
+  stopping = false;
+  hasOpenedOnce = false;
+  constructor(opts) {
+    super();
+    this.opts = opts;
+    this.logger = opts.logger ?? noopLogger;
+    this.validator = new SchemaValidator();
+    this.transport = new Transport({
+      gatewayUrl: opts.gatewayUrl,
+      neuronId: opts.neuronId,
+      auth: opts.auth,
+      logger: this.logger,
+      initialReconnectDelayMs: opts.initialReconnectDelayMs,
+      maxReconnectDelayMs: opts.maxReconnectDelayMs,
+      fetch: opts.fetch,
+      // Internal callbacks — Transport has a single subscriber (this
+      // Neuron) so a callbacks-on-options surface is simpler than an
+      // event emitter. The Neuron is the one that re-publishes lifecycle
+      // signals to its own consumers via the public emitter.
+      onOpen: () => {
+        this.emit("connection:open");
+        void this.handleConnectionOpen();
+      },
+      onClose: (reason) => {
+        this.handleConnectionClose(reason);
+        this.emit("connection:close", reason);
+      },
+      onFrame: (frame) => this.handleFrame(frame),
+      onError: (err) => {
+        this.logger.error({ err }, "transport error");
+        this.emit("error", err);
+      },
+      onReconnecting: (info) => this.emit("connection:reconnecting", info)
+    });
+    this.heartbeat = new HeartbeatLoop({
+      transport: this.transport,
+      neuronId: opts.neuronId,
+      inFlight: () => this.snapshotInFlight(),
+      intervalMs: opts.heartbeatIntervalMs,
+      logger: this.logger
+    });
+  }
+  handle(defOrReg, optsOrHandler, maybeHandler) {
+    if (this.started) {
+      throw new Error(
+        "Neuron.handle() must be called before start(). Use updateCapability(...) for runtime changes."
+      );
+    }
+    const isCapDef = "_zod" in defOrReg;
+    let capability;
+    let inputSchema;
+    let outputSchema;
+    let zodInput;
+    let zodOutput;
+    let handler;
+    if (isCapDef) {
+      const def = defOrReg;
+      let opts;
+      if (typeof optsOrHandler === "function") {
+        handler = optsOrHandler;
+      } else {
+        opts = optsOrHandler;
+        if (!maybeHandler) {
+          throw new Error("Neuron.handle(def, opts, fn): handler is required");
+        }
+        handler = maybeHandler;
+      }
+      capability = toRegistrationCapability(def, opts);
+      inputSchema = def.inputSchema;
+      outputSchema = def.outputSchema;
+      zodInput = def._zod.input;
+      zodOutput = def._zod.output;
+    } else {
+      const reg = defOrReg;
+      if (typeof optsOrHandler !== "function") {
+        throw new Error(
+          "Neuron.handle(registration, fn): expected handler function as second argument"
+        );
+      }
+      handler = optsOrHandler;
+      capability = { type: reg.type, scope: reg.scope, concurrency: reg.concurrency };
+      inputSchema = reg.inputSchema;
+      outputSchema = reg.outputSchema;
+    }
+    this.handlers.set(capability.type, {
+      capability,
+      inputSchema,
+      outputSchema,
+      zodInput,
+      zodOutput,
+      handler,
+      pool: new WorkerPool(capability.concurrency)
+    });
+  }
+  /**
+   * Adjust an already-registered capability's concurrency at runtime and
+   * notify the gateway via an `update-capability` frame so prefetch can be
+   * tuned. Schema and handler are not replaced — re-register before start
+   * if those need to change.
+   */
+  async updateCapability(reg) {
+    const entry = this.handlers.get(reg.type);
+    if (!entry) {
+      throw new Error(
+        `updateCapability: no handler registered for ${reg.type}`
+      );
+    }
+    entry.capability = { ...entry.capability, concurrency: reg.concurrency };
+    entry.pool.setConcurrency(reg.concurrency);
+    if (this.started && this.transport.isConnected()) {
+      await this.transport.post(POST_PATH.updateCapability, {
+        type: "update-capability",
+        neuronId: this.opts.neuronId,
+        capability: entry.capability
+      }).catch((err) => {
+        this.logger.warn(
+          { err, capability: reg.type },
+          "update-capability POST failed; gateway will pick up new value on next reconnect"
+        );
+      });
+    }
+  }
+  /**
+   * Open the SSE stream, post register with all advertised capabilities,
+   * and start the heartbeat loop. Resolves once the gateway accepts the
+   * register frame.
+   */
+  async start() {
+    if (this.started) return;
+    if (this.handlers.size === 0) {
+      throw new Error(
+        "Neuron.start() requires at least one handler \u2014 call Neuron.handle(...) first."
+      );
+    }
+    this.started = true;
+    await this.transport.start();
+  }
+  /**
+   * Refuse new leases, drain in-flight handlers up to `timeoutMs`, then
+   * tear down the connection. Pending leases (queued but not started)
+   * are abandoned — the gateway's lease TTL will reclaim them.
+   */
+  async stop(opts = {}) {
+    if (this.stopping) return;
+    this.stopping = true;
+    const drain = opts.drain ?? true;
+    const timeoutMs = opts.timeoutMs ?? 3e4;
+    this.heartbeat.stop();
+    if (drain) {
+      const drainPromises = Array.from(this.handlers.values()).map(
+        (h) => h.pool.drain()
+      );
+      const allDrained = Promise.all(drainPromises);
+      const timeout = new Promise((resolve) => {
+        const t = setTimeout(() => {
+          this.logger.warn({ timeoutMs }, "drain timeout exceeded; aborting in-flight leases");
+          for (const inFlight of this.inFlightLeases.values()) {
+            inFlight.abort.abort();
+          }
+          for (const queued of this.queuedLeases.values()) {
+            queued.abort.abort();
+          }
+          resolve();
+        }, timeoutMs);
+        if (typeof t.unref === "function") t.unref();
+      });
+      await Promise.race([allDrained, timeout]);
+    } else {
+      for (const inFlight of this.inFlightLeases.values()) {
+        inFlight.abort.abort();
+      }
+      for (const queued of this.queuedLeases.values()) {
+        queued.abort.abort();
+      }
+    }
+    await this.transport.stop();
+  }
+  // ----- internals ----------------------------------------------------
+  snapshotInFlight() {
+    return Array.from(this.inFlightLeases.entries()).map(([leaseId, info]) => ({
+      leaseId,
+      taskId: info.taskId,
+      startedAt: info.startedAt
+    }));
+  }
+  async handleConnectionOpen() {
+    try {
+      await this.postRegister();
+      this.hasOpenedOnce = true;
+      this.heartbeat.start();
+    } catch (err) {
+      this.logger.error({ err }, "register POST failed; transport will retry on reconnect");
+    }
+  }
+  handleConnectionClose(reason) {
+    this.logger.warn({ reason }, "transport disconnected");
+    this.heartbeat.stop();
+  }
+  async postRegister() {
+    await this.transport.post(POST_PATH.register, {
+      type: "register",
+      neuronId: this.opts.neuronId,
+      protocolVersion: PROTOCOL_VERSION,
+      capabilities: Array.from(this.handlers.values()).map((h) => h.capability)
+    });
+  }
+  handleFrame(frame) {
+    this.emit("frame", frame);
+    switch (frame.type) {
+      case "registered":
+        this.logger.info(
+          { sessionId: frame.sessionId, capabilities: frame.capabilities.length },
+          "gateway accepted registration"
+        );
+        this.emit("connection:registered", {
+          sessionId: frame.sessionId,
+          capabilities: frame.capabilities.length
+        });
+        return;
+      case "lease":
+        void this.dispatchLease(frame.lease);
+        return;
+      case "cancel":
+        this.handleCancel(frame.leaseId, frame.reason);
+        return;
+      case "capability-revoked":
+        this.logger.warn(
+          { capability: frame.capability, reason: frame.reason },
+          "capability revoked by gateway"
+        );
+        return;
+      case "disconnect":
+        this.logger.info({ reason: frame.reason, code: frame.code }, "gateway disconnect");
+        return;
+    }
+  }
+  async dispatchLease(lease) {
+    const handler = this.handlers.get(lease.task.capability);
+    if (!handler) {
+      this.logger.warn(
+        { capability: lease.task.capability, leaseId: lease.leaseId },
+        "no handler registered for lease capability; nacking retryable"
+      );
+      await this.transport.post(POST_PATH.nack, {
+        type: "nack",
+        leaseId: lease.leaseId,
+        kind: "retryable",
+        reason: "no handler registered for capability",
+        failedAt: (/* @__PURE__ */ new Date()).toISOString()
+      }).catch(() => void 0);
+      return;
+    }
+    const abort = new AbortController();
+    this.queuedLeases.set(lease.leaseId, { abort });
+    try {
+      await handler.pool.run(async () => {
+        this.queuedLeases.delete(lease.leaseId);
+        this.inFlightLeases.set(lease.leaseId, {
+          taskId: lease.task.taskId,
+          startedAt: (/* @__PURE__ */ new Date()).toISOString(),
+          abort
+        });
+        try {
+          await runLease({
+            lease,
+            handler: handler.handler,
+            zodInput: handler.zodInput,
+            zodOutput: handler.zodOutput,
+            transport: this.transport,
+            validator: this.validator,
+            logger: this.logger,
+            signal: abort.signal
+          });
+        } finally {
+          this.inFlightLeases.delete(lease.leaseId);
+        }
+      });
+    } catch (err) {
+      this.queuedLeases.delete(lease.leaseId);
+      if (err instanceof ShutdownError) {
+        this.logger.info(
+          { leaseId: lease.leaseId },
+          "lease refused during shutdown; nacking retryable"
+        );
+        await this.transport.post(POST_PATH.nack, {
+          type: "nack",
+          leaseId: lease.leaseId,
+          kind: "retryable",
+          reason: "neuron shutting down",
+          failedAt: (/* @__PURE__ */ new Date()).toISOString()
+        }).catch(() => void 0);
+      } else {
+        this.logger.error({ err, leaseId: lease.leaseId }, "unhandled lease runner error");
+      }
+    }
+  }
+  handleCancel(leaseId, reason) {
+    const tracked = this.inFlightLeases.get(leaseId) ?? this.queuedLeases.get(leaseId);
+    if (!tracked) {
+      this.logger.warn(
+        { leaseId, reason },
+        "cancel for unknown lease; ignoring (already completed?)"
+      );
+      return;
+    }
+    this.logger.info({ leaseId, reason }, "cancelling lease");
+    tracked.abort.abort();
+  }
+};
+var SSE_EVENT_NAMES = SSE_EVENT;
+
+export { Neuron, SSE_EVENT_NAMES, SchemaValidator, ShutdownError, Transport, TransportError, WorkerPool, defineCapability, toPlanShape, toRegistrationCapability };
+//# sourceMappingURL=index.mjs.map
+//# sourceMappingURL=index.mjs.map