@justjoshn/pi-pstack 0.0.0-stage → 0.15.36-pi.11

package/src/subscriptions/supervisor-main.ts

@@ -0,0 +1,134 @@
+#!/usr/bin/env node
+import { spawn } from 'node:child_process';
+import { resolve } from 'node:path';
+import { StringDecoder } from 'node:string_decoder';
+import { fileURLToPath } from 'node:url';
+
+import { PACKAGE_ROOT } from '../paths.ts';
+import { createSupervisor, type FreshMeta } from './supervisor.ts';
+
+function usage(): never {
+  console.error('Usage: pstack-routines-supervisor --workspace <cwd>');
+  process.exit(2);
+}
+
+function parseArgs(argv: string[]): { workspace: string } {
+  let workspace: string | undefined;
+  for (let index = 0; index < argv.length; index += 1) {
+    const arg = argv[index];
+    if (arg === '--workspace') {
+      workspace = argv[index + 1];
+      index += 1;
+      continue;
+    }
+    if (arg === '--help' || arg === '-h') usage();
+    console.error(`Unknown argument: ${arg}`);
+    usage();
+  }
+  if (workspace === undefined || workspace.trim() === '') usage();
+  return { workspace: resolve(workspace) };
+}
+
+/** The model, trust, and print arguments a detached new_session fire uses. */
+export function freshArgs(meta: FreshMeta): string[] {
+  const model = meta.model === undefined ? [] : ['--model', `${meta.model.provider}/${meta.model.id}`, ...(meta.model.thinking === undefined ? [] : ['--thinking', meta.model.thinking])];
+  const trust = meta.trusted === undefined ? [] : [meta.trusted ? '--approve' : '--no-approve'];
+  return [...trust, ...model, '-p', '--'];
+}
+
+/** The text of one JSON-mode assistant message, or undefined for another record. */
+export function assistantText(event: unknown): string | undefined {
+  const value = event as { type?: unknown; message?: { role?: unknown; content?: unknown } } | null;
+  if (value?.type !== 'message_end' || value.message?.role !== 'assistant' || !Array.isArray(value.message.content)) return undefined;
+  const text = (value.message.content as { type?: unknown; text?: unknown }[])
+    .filter((part) => part.type === 'text' && typeof part.text === 'string')
+    .map((part) => part.text as string)
+    .join('');
+  return text.trim() === '' ? undefined : text;
+}
+
+const STDERR_LIMIT = 2_000;
+
+/** Spawns a headless Pi print run; resolves with its final assistant message once the child exits 0. */
+function runFresh(prompt: string, name: string, meta: FreshMeta): Promise<string | undefined> {
+  return new Promise((resolvePromise, reject) => {
+    const raw: unknown = JSON.parse(process.env['PI_PSTACK_PI_COMMAND'] ?? '["pi"]');
+    if (!Array.isArray(raw) || raw.length === 0 || !raw.every((part) => typeof part === 'string' && part !== '')) {
+      reject(new Error('Invalid Pi command supplied to the subscription supervisor.'));
+      return;
+    }
+    const [executable, ...prefix] = raw as string[];
+    if (executable === undefined) {
+      reject(new Error(`Could not spawn fire for "${name}" (${meta.subscriptionId}): no pi command`));
+      return;
+    }
+    const child = spawn(executable, [...prefix, '--mode', 'json', '-e', PACKAGE_ROOT, ...freshArgs(meta), prompt], {
+      cwd: meta.cwd,
+      env: process.env,
+      stdio: ['ignore', 'pipe', 'pipe'],
+    });
+    let outcome: string | undefined;
+    let buffer = '';
+    let stderr = '';
+    const decoder = new StringDecoder('utf8');
+    // JSONL framing: split only on LF, and keep reading so a full pipe never stalls the child.
+    child.stdout.on('data', (chunk: Buffer) => {
+      buffer += decoder.write(chunk);
+      for (let index = buffer.indexOf('\n'); index !== -1; index = buffer.indexOf('\n')) {
+        const line = buffer.slice(0, index).trim();
+        buffer = buffer.slice(index + 1);
+        if (line === '') continue;
+        try {
+          outcome = assistantText(JSON.parse(line)) ?? outcome;
+        } catch {
+          // A child extension can print non-JSON to stdout; it is not a session event.
+        }
+      }
+    });
+    child.stderr.on('data', (chunk: Buffer) => {
+      stderr = (stderr + chunk.toString('utf8')).slice(-STDERR_LIMIT);
+    });
+    child.once('error', (error) => reject(new Error(`Could not spawn fire for "${name}" (${meta.subscriptionId}): ${error.message}`)));
+    child.once('exit', (code, signal) => {
+      const detail = stderr.trim() === '' ? '' : ` ${stderr.trim()}`;
+      if (signal != null) {
+        reject(new Error(`Fire for "${name}" (${meta.subscriptionId}) exited from signal ${signal}.${detail}`));
+        return;
+      }
+      if (code !== 0) {
+        reject(new Error(`Fire for "${name}" (${meta.subscriptionId}) exited with code ${code ?? 'null'}.${detail}`));
+        return;
+      }
+      resolvePromise(outcome);
+    });
+  });
+}
+
+async function main(): Promise<void> {
+  const { workspace } = parseArgs(process.argv.slice(2));
+  const supervisor = createSupervisor({ workspaceCwd: workspace, runFresh, isDaemon: true });
+  const result = await supervisor.start();
+  if (!result.started) {
+    console.log(`pstack-routines-supervisor already running for ${workspace} (pid ${result.pid}).`);
+    process.exit(0);
+  }
+  console.log(`pstack-routines-supervisor started for ${workspace} (pid ${process.pid}).`);
+  await new Promise<void>((resolve) => {
+    const shutdown = () => {
+      supervisor.stop();
+      resolve();
+    };
+    process.on('SIGINT', shutdown);
+    process.on('SIGTERM', shutdown);
+    process.on('SIGHUP', () => {
+      void supervisor.reconcile();
+    });
+  });
+}
+
+if (process.argv[1] !== undefined && resolve(process.argv[1]) === fileURLToPath(import.meta.url)) {
+  main().catch((error: unknown) => {
+    console.error(error instanceof Error ? error.message : error);
+    process.exit(1);
+  });
+}

package/src/subscriptions/supervisor.ts

@@ -0,0 +1,325 @@
+import { randomUUID } from 'node:crypto';
+
+import { routinesDir } from '../paths.ts';
+import { cronMatches, parseCron } from './cron.ts';
+import {
+  acceptEvent,
+  acquireSupervisorLock,
+  alive,
+  canonicalPath,
+  clearError,
+  expireDesired,
+  finishFresh,
+  isDesired,
+  readLock,
+  readRegistry,
+  readRoutinesPort,
+  recordError,
+  releaseSupervisorLock,
+  runsDetached,
+  startFresh,
+  writeArm,
+  writeRoutinesPort,
+} from './durable-store.ts';
+import type { DurableSubscription, EventAdmission, RunModel, StoredSubscription } from './durable-types.ts';
+import { defaultGhRunner, fetchPrs, type GhRunner, type PrSnapshot, pollChanges } from './github.ts';
+import { credentialFile, DEFAULT_ROUTINES_PORT, ensureCredentials, RoutineServer, type WebhookEvent, webhookEventBlock } from './webhook.ts';
+
+const MAX_STEP_MS = 24 * 24 * 3_600_000;
+const MINUTE_MS = 60_000;
+const ADMISSION_FULL = 'The delivery queue is full; the event was not accepted and will be offered again.';
+export interface FreshMeta {
+  subscriptionId: string;
+  cwd: string;
+  /** The owner's model when it created a new_session timer. */
+  model?: RunModel;
+  /** The owner's project-trust decision when it created the timer. */
+  trusted?: boolean;
+}
+export type RunFresh = (prompt: string, name: string, meta: FreshMeta) => Promise<string | undefined> | string | undefined;
+export interface SupervisorOptions {
+  workspaceCwd: string;
+  runFresh: RunFresh;
+  now?: () => number;
+  routinesPort?: number;
+  routinesDir?: string;
+  pollIntervalMs?: number;
+  gh?: GhRunner;
+  reconcileIntervalMs?: number;
+  onAccepted?: () => void;
+  isDaemon?: boolean;
+  ownerSessionId?: string;
+}
+type TimerConfig = Extract<DurableSubscription['config'], { kind: 'timer' }>;
+type Armed = { sub: StoredSubscription; timers: Set<ReturnType<typeof setTimeout>>; baseline?: Map<number, PrSnapshot>; polling?: boolean };
+
+function envPort(): number | undefined {
+  const raw = process.env['PSTACK_ROUTINES_PORT'];
+  if (raw === undefined || raw === '') return undefined;
+  const port = Number(raw);
+  return Number.isInteger(port) && port >= 0 && port < 65536 ? port : undefined;
+}
+export function delayFireIndex(createdAt: number, delaySeconds: number, at: number): number {
+  return Math.floor(Math.max(0, at - createdAt) / (Math.max(1, delaySeconds) * 1_000));
+}
+export function timerClaimKey(sub: DurableSubscription, config: TimerConfig, at: number): string {
+  if (config.once) return `${sub.id}:once`;
+  return `${sub.id}:${config.cron === undefined ? delayFireIndex(sub.createdAt, config.delaySeconds ?? 1, at) : Math.floor(at / MINUTE_MS)}`;
+}
+
+export function createSupervisor(options: SupervisorOptions) {
+  const cwd = canonicalPath(options.workspaceCwd);
+  const now = options.now ?? Date.now;
+  const gh = options.gh ?? defaultGhRunner;
+  const pollMs = options.pollIntervalMs ?? MINUTE_MS;
+  const dir = options.routinesDir ?? routinesDir();
+  const explicitPort = options.routinesPort ?? envPort();
+  const server = new RoutineServer(explicitPort ?? readRoutinesPort(cwd) ?? DEFAULT_ROUTINES_PORT, explicitPort === undefined);
+  const armed = new Map<string, Armed>();
+  const isDaemon = options.isDaemon ?? true;
+  let active = false;
+  let epoch: string | undefined;
+  let watchTimer: ReturnType<typeof setTimeout> | undefined;
+  let reconciling: Promise<void> | undefined;
+
+  const owns = (entry: Armed) => active && armed.get(entry.sub.id) === entry;
+  const at = (entry: Armed, when: number, fn: () => void) => {
+    const handle = setTimeout(
+      () => {
+        entry.timers.delete(handle);
+        if (!owns(entry)) return;
+        if (now() < when) at(entry, when, fn);
+        else fn();
+      },
+      Math.min(Math.max(0, when - now()), MAX_STEP_MS),
+    );
+    handle.unref?.();
+    entry.timers.add(handle);
+  };
+  const disarm = (entry: Armed) => {
+    for (const handle of entry.timers) clearTimeout(handle);
+    entry.timers.clear();
+    if (entry.sub.config.kind === 'webhook') {
+      server.routes.delete(entry.sub.config.slug);
+      if (server.routes.size === 0) server.close(false);
+    }
+  };
+  const error = (id: string, cause: unknown) => recordError(cwd, id, cause);
+  /** Persists the port the listener actually bound, so a restart reuses it and sender URLs keep working. */
+  const syncPort = () => {
+    const port = server.boundPort();
+    if (port !== undefined) writeRoutinesPort(cwd, port);
+  };
+  const offlineFresh = () => {
+    if (!active) return;
+    const registry = readRegistry(cwd);
+    const eligible = registry.pending.filter((event) => {
+      if (event.claim !== undefined && alive(event.claim.pid)) return false;
+      if (event.retryAt !== undefined && now() < event.retryAt) return false;
+      const sub = registry.subscriptions.find((s) => s.id === event.subscriptionId);
+      if (sub === undefined || sub.lifecycle === 'cancelled' || !runsDetached(sub)) return false;
+      if (!isDaemon && (options.ownerSessionId === undefined || sub.ownerSessionId !== options.ownerSessionId)) return false;
+      return true;
+    });
+    for (const event of eligible) {
+      const claim = startFresh(
+        cwd,
+        event.id,
+        undefined,
+        () => active,
+        (sub, pending) =>
+          options.runFresh(pending.content, sub.name, {
+            subscriptionId: sub.id,
+            cwd,
+            ...(sub.config.kind !== 'github' && sub.config.model !== undefined ? { model: sub.config.model } : {}),
+            ...(sub.config.kind !== 'github' && sub.config.trusted !== undefined ? { trusted: sub.config.trusted } : {}),
+          }),
+      );
+      if (claim !== undefined)
+        void claim.result
+          .then(
+            (outcome) => finishFresh(cwd, event.id, claim.token, undefined, typeof outcome === 'string' ? outcome : undefined),
+            (cause: unknown) => finishFresh(cwd, event.id, claim.token, cause),
+          )
+          .catch((cause: unknown) => {
+            console.error(`Subscription delivery status failed. ${String(cause)}`);
+          });
+    }
+  };
+  const accepted = () => {
+    options.onAccepted?.();
+    offlineFresh();
+  };
+  const fireTimer = (entry: Armed, config: TimerConfig) => {
+    if (!owns(entry)) return;
+    try {
+      const prompt = config.strategy === 'wake_self' ? `[timer] ${entry.sub.name}\n\n${config.prompt}` : config.prompt;
+      const admission = acceptEvent(cwd, entry.sub.id, timerClaimKey(entry.sub, config, now()), prompt, now());
+      if (admission === 'accepted') accepted();
+      else if (admission === 'full') error(entry.sub.id, new Error(ADMISSION_FULL));
+    } catch (cause) {
+      error(entry.sub.id, cause);
+    }
+  };
+  const armTimer = (entry: Armed, config: TimerConfig) => {
+    if (config.cron !== undefined) {
+      const schedule = parseCron(config.cron);
+      const tick = () => {
+        if (cronMatches(schedule, new Date(now()))) fireTimer(entry, config);
+        const current = readRegistry(cwd).subscriptions.find((sub) => sub.id === entry.sub.id);
+        if (current !== undefined && isDesired(current, now())) at(entry, (Math.floor(now() / MINUTE_MS) + 1) * MINUTE_MS, tick);
+      };
+      at(entry, (Math.floor(now() / MINUTE_MS) + 1) * MINUTE_MS, tick);
+    } else {
+      const period = (config.delaySeconds ?? 1) * 1_000;
+      const next = () => (config.once ? entry.sub.createdAt + period : entry.sub.createdAt + (Math.floor(Math.max(0, now() - entry.sub.createdAt) / period) + 1) * period);
+      const fire = () => {
+        fireTimer(entry, config);
+        const current = readRegistry(cwd).subscriptions.find((sub) => sub.id === entry.sub.id);
+        if (current !== undefined && isDesired(current, now())) at(entry, config.once ? now() + 1_000 : next(), fire);
+      };
+      at(entry, next(), fire);
+    }
+  };
+  const onWebhook =
+    (entry: Armed, prompt: string) =>
+    async (event: WebhookEvent): Promise<EventAdmission> => {
+      if (!owns(entry)) return 'inactive';
+      const content = [`[routine] ${entry.sub.name}`, prompt, webhookEventBlock(event)].join('\n\n');
+      const result: EventAdmission = acceptEvent(cwd, entry.sub.id, event.eventId, content, now());
+      if (result === 'accepted') accepted();
+      else if (result === 'full') error(entry.sub.id, new Error(ADMISSION_FULL));
+      return result;
+    };
+  const poll = async (entry: Armed, target: Extract<DurableSubscription['config'], { kind: 'github' }>['target']) => {
+    if (entry.polling || !owns(entry)) return;
+    entry.polling = true;
+    try {
+      const after = await fetchPrs(gh, target);
+      if (!owns(entry)) return;
+      let blocked = false;
+      if (entry.baseline !== undefined) {
+        const blocks = await pollChanges(gh, target, entry.baseline, after);
+        if (!owns(entry)) return;
+        if (blocks.length > 0) {
+          const admission = acceptEvent(cwd, entry.sub.id, randomUUID(), [`[github] ${entry.sub.name}`, ...blocks].join('\n\n'), now());
+          if (admission === 'full') {
+            error(entry.sub.id, new Error(ADMISSION_FULL));
+            blocked = true;
+          } else if (admission === 'accepted') accepted();
+        }
+      }
+      if (!blocked) {
+        entry.baseline = after;
+        clearError(cwd, entry.sub.id);
+      }
+    } finally {
+      entry.polling = false;
+    }
+  };
+  const arm = async (entry: Armed): Promise<string | undefined> => {
+    if (entry.sub.ownerSessionId === undefined) throw new Error('Owner session ID unavailable. Recreate this subscription.');
+    if (entry.sub.expiresAt !== 0)
+      at(entry, entry.sub.expiresAt, () => {
+        expireDesired(cwd, now());
+        void reconcile().catch((cause: unknown) => error(entry.sub.id, cause));
+      });
+    const config = entry.sub.config;
+    if (config.kind === 'timer') armTimer(entry, config);
+    else if (config.kind === 'webhook') {
+      server.routes.set(config.slug, { slug: config.slug, credentialFile: ensureCredentials(dir, config.slug), onEvent: onWebhook(entry, config.prompt) });
+      return await server.url(config.slug);
+    } else {
+      const loop = () => {
+        void poll(entry, config.target).then(
+          () => {
+            if (owns(entry)) at(entry, now() + pollMs, loop);
+          },
+          (cause: unknown) => {
+            error(entry.sub.id, cause);
+            if (owns(entry)) at(entry, now() + pollMs, loop);
+          },
+        );
+      };
+      loop();
+    }
+    return undefined;
+  };
+  const reconcile = (): Promise<void> => {
+    if (!active) return Promise.resolve();
+    if (reconciling !== undefined) return reconciling;
+    reconciling = (async () => {
+      const registry = readRegistry(cwd);
+      if (registry.subscriptions.some((sub) => sub.lifecycle === 'active' && sub.expiresAt !== 0 && sub.expiresAt <= now())) expireDesired(cwd, now());
+      const desired = registry.subscriptions.filter((sub) => isDesired(sub, now()));
+      for (const [id, entry] of armed)
+        if (!desired.some((sub) => sub.id === id && sub.revision === entry.sub.revision)) {
+          disarm(entry);
+          armed.delete(id);
+        }
+      for (const sub of desired) {
+        if (!active) return;
+        if (armed.has(sub.id)) continue;
+        const entry: Armed = { sub, timers: new Set() };
+        armed.set(sub.id, entry);
+        try {
+          const url = await arm(entry);
+          if (!owns(entry) || epoch === undefined) continue;
+          clearError(cwd, sub.id);
+          writeArm(cwd, { state: 'ready', subscriptionId: sub.id, revision: sub.revision, epoch, ...(url === undefined ? {} : { url }) });
+        } catch (cause) {
+          disarm(entry);
+          armed.delete(sub.id);
+          if (!active || epoch === undefined) return;
+          const message = cause instanceof Error ? cause.message : String(cause);
+          error(sub.id, cause);
+          writeArm(cwd, { state: 'error', subscriptionId: sub.id, revision: sub.revision, epoch, message });
+        }
+      }
+      offlineFresh();
+      syncPort();
+    })().finally(() => {
+      reconciling = undefined;
+    });
+    return reconciling;
+  };
+  const watch = () => {
+    if (!active) return;
+    watchTimer = setTimeout(() => {
+      void reconcile()
+        .catch((cause: unknown) => console.error(`Subscription supervisor reconcile failed. ${String(cause)}`))
+        .finally(watch);
+    }, options.reconcileIntervalMs ?? 250);
+  };
+  return {
+    start: async (): Promise<{ started: boolean; alreadyRunning?: boolean; pid?: number }> => {
+      const lock = acquireSupervisorLock(cwd);
+      if (!lock.acquired) return { started: false, alreadyRunning: true, pid: lock.pid };
+      epoch = readLock(cwd)?.epoch;
+      active = true;
+      try {
+        await reconcile();
+        syncPort();
+        watch();
+        return { started: true };
+      } catch (cause) {
+        active = false;
+        releaseSupervisorLock(cwd, process.pid, epoch);
+        throw cause;
+      }
+    },
+    reconcile,
+    stop: () => {
+      active = false;
+      if (watchTimer !== undefined) clearTimeout(watchTimer);
+      for (const entry of armed.values()) disarm(entry);
+      armed.clear();
+      server.close();
+      if (epoch !== undefined) releaseSupervisorLock(cwd, process.pid, epoch);
+      epoch = undefined;
+    },
+    isListening: () => server.isListening(),
+    credentialFile: (slug: string) => credentialFile(dir, slug),
+  };
+}
+export type Supervisor = ReturnType<typeof createSupervisor>;