@m64/nats-pi-bridge 0.0.1

package/dist/server.js

@@ -0,0 +1,689 @@
+#!/usr/bin/env node
+
+// server.ts
+import { existsSync, readFileSync, statSync } from "node:fs";
+import { homedir } from "node:os";
+import { join, resolve } from "node:path";
+import {
+  connect,
+  credsAuthenticator,
+  jwtAuthenticator,
+  nkeyAuthenticator,
+  tokenAuthenticator,
+  usernamePasswordAuthenticator
+} from "@nats-io/transport-node";
+import { Svcm } from "@nats-io/services";
+import {
+  AuthStorage,
+  ModelRegistry,
+  SessionManager,
+  createAgentSession
+} from "@mariozechner/pi-coding-agent";
+var STATE_DIR = join(homedir(), ".pi-exec");
+var CONFIG_FILE = join(STATE_DIR, "config.json");
+var NATS_CONTEXT_DIR = join(homedir(), ".config", "nats", "context");
+var SERVICE_NAME = "pi-exec";
+var SERVICE_VERSION = "0.0.1";
+var DEFAULT_MAX_LIFETIME_SECONDS = 1800;
+var STALE_REQUEST_CUTOFF_MS = 30 * 60 * 1e3;
+var STALE_PRUNE_INTERVAL_MS = 6e4;
+var LIFETIME_CHECK_INTERVAL_MS = 3e4;
+var SHUTDOWN_FORCED_EXIT_MS = 2e3;
+var VALID_THINKING_LEVELS = ["off", "minimal", "low", "medium", "high", "xhigh"];
+var DEFAULT_CONTEXT = {
+  url: "demo.nats.io",
+  description: "NATS demo server (no auth)"
+};
+function intakeSubject(owner2) {
+  return `agents.pi-exec.${owner2}`;
+}
+function sessionSubject(owner2, sessionId) {
+  return `agents.pi-exec.${owner2}.${sessionId}`;
+}
+function sessionInspectSubject(owner2, sessionId) {
+  return `${sessionSubject(owner2, sessionId)}.inspect`;
+}
+function sanitizeSessionName(s) {
+  return s.replace(/[^a-zA-Z0-9_-]/g, "-").toLowerCase().replace(/^-+|-+$/g, "");
+}
+function loadNatsContext(name) {
+  const contextFile = join(NATS_CONTEXT_DIR, `${name}.json`);
+  try {
+    return JSON.parse(readFileSync(contextFile, "utf8"));
+  } catch (err) {
+    throw new Error(
+      `NATS context "${name}" not found at ${contextFile} (${err.message})`
+    );
+  }
+}
+function contextToConnectOpts(ctx) {
+  const opts = {
+    name: "pi-exec"
+  };
+  if (ctx.url) {
+    opts.servers = ctx.url;
+  }
+  if (ctx.creds) {
+    opts.authenticator = credsAuthenticator(readFileSync(ctx.creds));
+  } else if (ctx.nkey) {
+    opts.authenticator = nkeyAuthenticator(readFileSync(ctx.nkey));
+  } else if (ctx.user_jwt && ctx.user_seed) {
+    const seed = new TextEncoder().encode(ctx.user_seed);
+    opts.authenticator = jwtAuthenticator(ctx.user_jwt, seed);
+  } else if (ctx.token) {
+    opts.authenticator = tokenAuthenticator(ctx.token);
+  } else if (ctx.user) {
+    opts.authenticator = usernamePasswordAuthenticator(ctx.user, ctx.password ?? "");
+  }
+  if (ctx.cert || ctx.key || ctx.ca) {
+    opts.tls = {
+      certFile: ctx.cert || void 0,
+      keyFile: ctx.key || void 0,
+      caFile: ctx.ca || void 0,
+      handshakeFirst: ctx.tls_first || void 0
+    };
+  }
+  if (ctx.inbox_prefix) {
+    opts.inboxPrefix = ctx.inbox_prefix;
+  }
+  return opts;
+}
+function loadConfig() {
+  try {
+    return JSON.parse(readFileSync(CONFIG_FILE, "utf8"));
+  } catch {
+    return {};
+  }
+}
+function parseEnvelope(data) {
+  const text = typeof data === "string" ? data : new TextDecoder().decode(data);
+  let parsed;
+  try {
+    parsed = JSON.parse(text);
+  } catch {
+    return { from: "anonymous", body: text };
+  }
+  if (typeof parsed === "object" && parsed !== null && typeof parsed.body === "string") {
+    const rec = parsed;
+    return {
+      from: typeof rec.from === "string" ? rec.from : "anonymous",
+      body: rec.body
+    };
+  }
+  return { from: "anonymous", body: text };
+}
+function publishChunked(nc2, replySubject, text) {
+  const maxPayload = nc2.info?.max_payload ?? 1024 * 1024;
+  const encoded = new TextEncoder().encode(text);
+  if (encoded.byteLength <= maxPayload) {
+    nc2.publish(replySubject, text);
+    return;
+  }
+  let offset = 0;
+  while (offset < encoded.byteLength) {
+    const chunk = encoded.subarray(offset, offset + maxPayload);
+    nc2.publish(replySubject, chunk);
+    offset += maxPayload;
+  }
+}
+function buildInspectResponse(sessionName, subject, cwd) {
+  return {
+    name: sessionName,
+    description: `PI agent in ${cwd}`,
+    capabilities: [
+      {
+        name: "prompt",
+        endpoint: subject,
+        description: "Send a prompt (plain text or {from,body} JSON) and receive a streamed response. Empty payload signals end-of-stream.",
+        pattern: "request/reply"
+      }
+    ]
+  };
+}
+async function startStatusLoop(conn) {
+  try {
+    for await (const s of conn.status()) {
+      switch (s.type) {
+        case "disconnect":
+          process.stderr.write(`pi-exec: NATS disconnected from ${s.server} \u2014 retrying
+`);
+          break;
+        case "reconnect":
+          process.stderr.write(`pi-exec: NATS reconnected to ${s.server}
+`);
+          break;
+        case "error": {
+          const err = s.error;
+          const data = s.data;
+          process.stderr.write(
+            `pi-exec: NATS error: ${err?.message ?? String(data ?? "(unknown)")}
+`
+          );
+          break;
+        }
+      }
+    }
+  } catch {
+  }
+}
+var shuttingDown = false;
+var requestCounter = 0;
+var sessions = /* @__PURE__ */ new Map();
+var creating = /* @__PURE__ */ new Set();
+var nc;
+var intakeService;
+var owner;
+var config;
+var authStorage;
+var modelRegistry;
+var lifetimeInterval;
+var pruneInterval;
+function resolveModel(modelSpec) {
+  if (!modelSpec) return void 0;
+  const slash = modelSpec.indexOf("/");
+  if (slash < 0) return void 0;
+  return modelRegistry.find(modelSpec.slice(0, slash), modelSpec.slice(slash + 1));
+}
+function validateThinkingLevel(level) {
+  if (level === void 0 || level === "") return void 0;
+  if (VALID_THINKING_LEVELS.includes(level)) {
+    return level;
+  }
+  throw new Error(
+    `invalid thinkingLevel: ${level} (must be one of ${VALID_THINKING_LEVELS.join(", ")})`
+  );
+}
+async function createPiSession(intake) {
+  if (!intake.cwd) throw new Error("cwd is required");
+  const absCwd = resolve(intake.cwd);
+  if (!existsSync(absCwd) || !statSync(absCwd).isDirectory()) {
+    throw new Error(`cwd not found or not a directory: ${absCwd}`);
+  }
+  const modelSpec = intake.model ?? config.defaultModel;
+  const model = resolveModel(modelSpec);
+  if (modelSpec && !model) throw new Error(`unknown model: ${modelSpec}`);
+  const thinkingLevel = validateThinkingLevel(
+    intake.thinkingLevel ?? config.defaultThinkingLevel
+  );
+  const { session } = await createAgentSession({
+    cwd: absCwd,
+    sessionManager: SessionManager.inMemory(),
+    authStorage,
+    modelRegistry,
+    model,
+    thinkingLevel
+  });
+  return {
+    session,
+    cwd: absCwd,
+    resolvedModel: model ? `${model.provider}/${model.id}` : void 0,
+    thinkingLevel
+  };
+}
+async function createPerSessionService(sessionId, subject, cwd) {
+  const svcm2 = new Svcm(nc);
+  const service = await svcm2.add({
+    name: SERVICE_NAME,
+    version: SERVICE_VERSION,
+    description: `${sessionId} \u2014 ${cwd}`,
+    metadata: {
+      type: "session",
+      platform: "pi",
+      sessionId,
+      cwd,
+      owner
+    },
+    queue: ""
+  });
+  service.addEndpoint("prompt", {
+    subject,
+    handler: (err, msg) => handlePerSessionMessage(sessionId, err, msg)
+  });
+  const inspectResp = JSON.stringify(buildInspectResponse(sessionId, subject, cwd));
+  service.addEndpoint("inspect", {
+    subject: sessionInspectSubject(owner, sessionId),
+    handler: (_err, msg) => {
+      try {
+        msg.respond(inspectResp);
+      } catch {
+      }
+    }
+  });
+  return service;
+}
+function respondJson(msg, payload) {
+  try {
+    msg.respond(JSON.stringify(payload));
+  } catch {
+  }
+}
+function handleIntakeMessage(err, msg) {
+  if (err) {
+    process.stderr.write(`pi-exec: intake error: ${err.message}
+`);
+    return;
+  }
+  if (shuttingDown) {
+    respondJson(msg, { error: "shutting_down" });
+    return;
+  }
+  let intake;
+  try {
+    intake = JSON.parse(msg.string());
+  } catch (e) {
+    respondJson(msg, { error: "invalid_json", message: e.message });
+    return;
+  }
+  if (!intake || typeof intake.sessionMode !== "string") {
+    respondJson(msg, { error: "invalid_request", message: "sessionMode required" });
+    return;
+  }
+  switch (intake.sessionMode) {
+    case "run":
+      void handleRunMode(msg, intake);
+      break;
+    case "session":
+      void handleSessionMode(msg, intake);
+      break;
+    case "stop":
+      void handleStopMode(msg, intake);
+      break;
+    case "list":
+      handleListMode(msg);
+      break;
+    default:
+      respondJson(msg, {
+        error: "invalid_sessionMode",
+        sessionMode: intake.sessionMode
+      });
+  }
+}
+async function handleRunMode(msg, intake) {
+  const reply = msg.reply;
+  if (!reply) {
+    respondJson(msg, { error: "no_reply_subject" });
+    return;
+  }
+  if (!intake.body || !intake.cwd) {
+    respondJson(msg, { error: "missing_fields", required: ["body", "cwd"] });
+    return;
+  }
+  let session;
+  let unsubscribe;
+  try {
+    const created = await createPiSession(intake);
+    session = created.session;
+    unsubscribe = session.subscribe((ev) => {
+      if (ev.type === "message_update" && ev.assistantMessageEvent.type === "text_delta" && ev.assistantMessageEvent.delta) {
+        publishChunked(nc, reply, ev.assistantMessageEvent.delta);
+      }
+    });
+    await session.prompt(intake.body);
+  } catch (e) {
+    try {
+      nc.publish(reply, `error: ${e.message}`);
+    } catch {
+    }
+    process.stderr.write(`pi-exec: run error: ${e.message}
+`);
+  } finally {
+    try {
+      unsubscribe?.();
+    } catch {
+    }
+    try {
+      session?.dispose();
+    } catch {
+    }
+    try {
+      nc.publish(reply, "");
+    } catch {
+    }
+    try {
+      await nc.flush();
+    } catch {
+    }
+  }
+}
+async function handleSessionMode(msg, intake) {
+  const reply = msg.reply;
+  if (!reply) {
+    respondJson(msg, { error: "no_reply_subject" });
+    return;
+  }
+  if (!intake.body || !intake.cwd || !intake.sessionId) {
+    respondJson(msg, {
+      error: "missing_fields",
+      required: ["body", "cwd", "sessionId"]
+    });
+    return;
+  }
+  const sessionId = sanitizeSessionName(intake.sessionId);
+  if (!sessionId) {
+    respondJson(msg, {
+      error: "invalid_sessionId",
+      message: "empty after sanitize"
+    });
+    return;
+  }
+  if (sessionId !== intake.sessionId) {
+    respondJson(msg, {
+      error: "invalid_sessionId",
+      sessionId: intake.sessionId,
+      message: `sessionId must match [a-z0-9_-]+; suggested: ${sessionId}`
+    });
+    return;
+  }
+  if (sessions.has(sessionId) || creating.has(sessionId)) {
+    respondJson(msg, {
+      error: "session_exists",
+      sessionId,
+      subject: sessionSubject(owner, sessionId),
+      message: "Session already exists. Send follow-up prompts to the session subject."
+    });
+    return;
+  }
+  creating.add(sessionId);
+  let created;
+  try {
+    created = await createPiSession(intake);
+  } catch (e) {
+    creating.delete(sessionId);
+    respondJson(msg, {
+      error: "session_create_failed",
+      message: e.message
+    });
+    return;
+  }
+  const subject = sessionSubject(owner, sessionId);
+  let service;
+  try {
+    service = await createPerSessionService(sessionId, subject, created.cwd);
+  } catch (e) {
+    creating.delete(sessionId);
+    try {
+      created.session.dispose();
+    } catch {
+    }
+    respondJson(msg, {
+      error: "service_register_failed",
+      message: e.message
+    });
+    return;
+  }
+  const managed = {
+    sessionId,
+    session: created.session,
+    service,
+    cwd: created.cwd,
+    model: created.resolvedModel,
+    thinkingLevel: created.thinkingLevel,
+    createdAt: Date.now(),
+    lastActivity: Date.now(),
+    maxLifetime: intake.maxLifetime ?? config.defaultMaxLifetime ?? DEFAULT_MAX_LIFETIME_SECONDS,
+    subject,
+    pendingRequests: /* @__PURE__ */ new Map(),
+    requestQueue: [],
+    activeRequestId: null,
+    disposed: false
+  };
+  sessions.set(sessionId, managed);
+  creating.delete(sessionId);
+  process.stderr.write(`pi-exec: session created ${sessionId} cwd=${created.cwd}
+`);
+  const initialId = `init-${++requestCounter}`;
+  managed.pendingRequests.set(initialId, {
+    requestId: initialId,
+    replySubject: reply,
+    from: intake.from ?? "anonymous",
+    body: intake.body,
+    createdAt: Date.now()
+  });
+  managed.requestQueue.push(initialId);
+  void drainSession(
+    managed,
+    () => {
+      process.stderr.write(`pi-exec: initial prompt failed for ${sessionId}; disposing
+`);
+      void disposeSession(managed);
+    },
+    initialId
+  );
+}
+async function handleStopMode(msg, intake) {
+  if (!intake.sessionId) {
+    respondJson(msg, { error: "missing_fields", required: ["sessionId"] });
+    return;
+  }
+  const sid = sanitizeSessionName(intake.sessionId);
+  const managed = sessions.get(sid);
+  if (!managed) {
+    respondJson(msg, { error: "not_found", sessionId: sid });
+    return;
+  }
+  await disposeSession(managed);
+  respondJson(msg, { ok: true, sessionId: sid });
+}
+function handleListMode(msg) {
+  const now = Date.now();
+  const list = Array.from(sessions.values()).map((m) => {
+    const elapsedSec = Math.floor((now - m.createdAt) / 1e3);
+    const remaining = m.maxLifetime > 0 ? Math.max(0, m.maxLifetime - elapsedSec) : 0;
+    return {
+      sessionId: m.sessionId,
+      subject: m.subject,
+      cwd: m.cwd,
+      model: m.model,
+      thinkingLevel: m.thinkingLevel,
+      createdAt: new Date(m.createdAt).toISOString(),
+      lastActivity: new Date(m.lastActivity).toISOString(),
+      maxLifetime: m.maxLifetime,
+      remainingLifetime: remaining,
+      activeRequest: m.activeRequestId !== null,
+      queuedRequests: m.requestQueue.length
+    };
+  });
+  respondJson(msg, { sessions: list });
+}
+function handlePerSessionMessage(sessionId, err, msg) {
+  if (err) {
+    process.stderr.write(`pi-exec: ${sessionId} handler err: ${err.message}
+`);
+    return;
+  }
+  if (!msg.reply) return;
+  const managed = sessions.get(sessionId);
+  if (!managed || managed.disposed) return;
+  const env = parseEnvelope(msg.data);
+  if (!env.body) return;
+  const rid = `${sessionId}-${++requestCounter}`;
+  managed.pendingRequests.set(rid, {
+    requestId: rid,
+    replySubject: msg.reply,
+    from: env.from,
+    body: env.body,
+    createdAt: Date.now()
+  });
+  managed.requestQueue.push(rid);
+  managed.lastActivity = Date.now();
+  void drainSession(managed);
+}
+async function drainSession(m, onInitialFailure, initialRequestId) {
+  if (m.activeRequestId) return;
+  if (m.disposed) return;
+  const next = m.requestQueue.shift();
+  if (!next) return;
+  const pr = m.pendingRequests.get(next);
+  if (!pr) {
+    void drainSession(m);
+    return;
+  }
+  m.activeRequestId = next;
+  m.lastActivity = Date.now();
+  let unsubscribe;
+  let failed = false;
+  try {
+    unsubscribe = m.session.subscribe((ev) => {
+      if (ev.type === "message_update" && ev.assistantMessageEvent.type === "text_delta" && ev.assistantMessageEvent.delta) {
+        publishChunked(nc, pr.replySubject, ev.assistantMessageEvent.delta);
+      }
+    });
+    await m.session.prompt(pr.body);
+  } catch (e) {
+    failed = true;
+    try {
+      nc.publish(pr.replySubject, `error: ${e.message}`);
+    } catch {
+    }
+    process.stderr.write(`pi-exec: ${m.sessionId} prompt error: ${e.message}
+`);
+  } finally {
+    try {
+      unsubscribe?.();
+    } catch {
+    }
+    try {
+      nc.publish(pr.replySubject, "");
+    } catch {
+    }
+    try {
+      await nc.flush();
+    } catch {
+    }
+    m.pendingRequests.delete(next);
+    m.activeRequestId = null;
+    m.lastActivity = Date.now();
+    if (failed && onInitialFailure && next === initialRequestId) {
+      onInitialFailure();
+      return;
+    }
+    if (!m.disposed && m.requestQueue.length > 0) {
+      setImmediate(() => void drainSession(m));
+    }
+  }
+}
+async function disposeSession(m) {
+  if (m.disposed) return;
+  m.disposed = true;
+  for (const pr of m.pendingRequests.values()) {
+    try {
+      nc.publish(pr.replySubject, "");
+    } catch {
+    }
+  }
+  m.pendingRequests.clear();
+  m.requestQueue.length = 0;
+  try {
+    await m.service.stop();
+  } catch {
+  }
+  try {
+    m.session.dispose();
+  } catch {
+  }
+  sessions.delete(m.sessionId);
+  process.stderr.write(`pi-exec: session disposed ${m.sessionId}
+`);
+}
+function checkLifetimes() {
+  const now = Date.now();
+  for (const m of sessions.values()) {
+    if (m.maxLifetime <= 0) continue;
+    if (now - m.createdAt > m.maxLifetime * 1e3) {
+      process.stderr.write(`pi-exec: session ${m.sessionId} expired (${m.maxLifetime}s)
+`);
+      void disposeSession(m);
+    }
+  }
+}
+function pruneStaleRequests() {
+  const cutoff = Date.now() - STALE_REQUEST_CUTOFF_MS;
+  for (const m of sessions.values()) {
+    for (const [id, pr] of m.pendingRequests) {
+      if (id === m.activeRequestId) continue;
+      if (pr.createdAt < cutoff) {
+        m.pendingRequests.delete(id);
+        const qi = m.requestQueue.indexOf(id);
+        if (qi >= 0) m.requestQueue.splice(qi, 1);
+      }
+    }
+  }
+}
+async function shutdown(signal) {
+  if (shuttingDown) return;
+  shuttingDown = true;
+  process.stderr.write(`pi-exec: shutdown (${signal})
+`);
+  const forceTimer = setTimeout(() => {
+    process.stderr.write("pi-exec: forced exit\n");
+    process.exit(0);
+  }, SHUTDOWN_FORCED_EXIT_MS);
+  forceTimer.unref();
+  clearInterval(lifetimeInterval);
+  clearInterval(pruneInterval);
+  await Promise.allSettled(Array.from(sessions.values()).map((m) => disposeSession(m)));
+  try {
+    await intakeService.stop();
+  } catch {
+  }
+  try {
+    await nc.drain();
+  } catch {
+  }
+  clearTimeout(forceTimer);
+  process.exit(0);
+}
+process.on("unhandledRejection", (err) => {
+  process.stderr.write(`pi-exec: unhandledRejection: ${err}
+`);
+});
+process.on("uncaughtException", (err) => {
+  process.stderr.write(`pi-exec: uncaughtException: ${err}
+`);
+});
+config = loadConfig();
+var contextName = process.env.NATS_CONTEXT ?? config.context;
+config.defaultModel = process.env.PI_EXEC_DEFAULT_MODEL ?? config.defaultModel;
+var envMaxLifetime = process.env.PI_EXEC_DEFAULT_MAX_LIFETIME;
+config.defaultMaxLifetime = envMaxLifetime ? Number(envMaxLifetime) : config.defaultMaxLifetime ?? DEFAULT_MAX_LIFETIME_SECONDS;
+var natsCtx;
+try {
+  natsCtx = contextName ? loadNatsContext(contextName) : DEFAULT_CONTEXT;
+} catch (e) {
+  process.stderr.write(`pi-exec: ${e.message}
+`);
+  process.exit(1);
+}
+process.stderr.write(
+  `pi-exec: connecting to ${natsCtx.url ?? "default"} (context: ${contextName ?? "default"})
+`
+);
+var connectOpts = contextToConnectOpts(natsCtx);
+connectOpts.name = "pi-exec";
+nc = await connect(connectOpts);
+process.stderr.write(`pi-exec: connected
+`);
+owner = sanitizeSessionName(process.env.USER ?? "unknown") || "unknown";
+authStorage = AuthStorage.create();
+modelRegistry = ModelRegistry.create(authStorage);
+var svcm = new Svcm(nc);
+intakeService = await svcm.add({
+  name: SERVICE_NAME,
+  version: SERVICE_VERSION,
+  description: "intake",
+  metadata: { type: "intake", platform: "pi", owner },
+  queue: ""
+});
+intakeService.addEndpoint("intake", {
+  subject: intakeSubject(owner),
+  handler: handleIntakeMessage
+});
+process.stderr.write(`pi-exec: intake registered on ${intakeSubject(owner)}
+`);
+lifetimeInterval = setInterval(checkLifetimes, LIFETIME_CHECK_INTERVAL_MS);
+lifetimeInterval.unref();
+pruneInterval = setInterval(pruneStaleRequests, STALE_PRUNE_INTERVAL_MS);
+pruneInterval.unref();
+void startStatusLoop(nc);
+process.on("SIGTERM", () => void shutdown("SIGTERM"));
+process.on("SIGINT", () => void shutdown("SIGINT"));