pi-amq 0.1.1 → 0.1.2

package/extension/worker-runtime.ts

@@ -2,7 +2,7 @@ import { spawn as spawnChild } from "node:child_process";
 import { createHash } from "node:crypto";
 import { access, readFile } from "node:fs/promises";
 import { constants } from "node:fs";
-import { delimiter, isAbsolute, join } from "node:path";
+import { delimiter, isAbsolute, join, normalize, resolve } from "node:path";
 import { fileURLToPath } from "node:url";
 
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
@@ -55,7 +55,12 @@ import {
   NodePiProcessFactory,
   WorkerProcessSupervisor,
 } from "../workflow/worker-process-supervisor.js";
-import { WorkflowLeaseManager } from "../workflow/workflow-lease.js";
+import {
+  WorkflowLeaseManager,
+  type ProcessIdentity,
+  type ProcessIdentityInspector,
+  type WorkflowLeaseOwnerV1,
+} from "../workflow/workflow-lease.js";
 import type {
   RunProjectionV2,
   ModelSpec,
@@ -94,6 +99,37 @@ const CAPABILITY_TOOLS: Readonly<Record<string, readonly string[]>> = {
   message_worker: ["message_worker"],
 };
 
+const WORKER_RUNTIME_REGISTRY_SYMBOL = Symbol.for("pi-amq.worker-runtime-registry.v1");
+
+interface WorkerRuntimeCoreV1 {
+  readonly schema: 1;
+  readonly key: string;
+  readonly workspacePath: string;
+  readonly workspaceId: string;
+  readonly namespaceDirectory: string;
+  readonly processIdentity: ProcessIdentity;
+  readonly supervisorGeneration: string;
+  readonly leaseManager: WorkflowLeaseManager;
+  readonly store: WorkflowRunStore;
+  readonly supervisor: WorkerProcessSupervisor;
+  readonly cleanup: WorkflowRunCleanup;
+  readonly doctor: WorkflowDoctor;
+  readonly artifactDirectory: string;
+  readonly requestKey: string;
+}
+
+interface WorkerRuntimeRegistryV1 {
+  readonly schema: 1;
+  readonly entries: Map<string, unknown>;
+}
+
+interface WorkerRuntimeAttachment {
+  readonly core: WorkerRuntimeCoreV1;
+  readonly notifier: WorkerCompletionNotifier;
+  active: boolean;
+  unsubscribe: () => void;
+}
+
 export class PiWorkerRuntime implements WorkerToolProvider {
   readonly #publishFleet: (records: readonly FleetAgentRecord[]) => void;
   readonly #unlimitedDeclaration: () => unknown;
@@ -111,15 +147,7 @@ export class PiWorkerRuntime implements WorkerToolProvider {
   #requestKey: string | undefined;
   #error: string | undefined;
   #guidance = DRIVER_DOCTRINE;
-  readonly #runtimes = new Map<string, {
-    readonly store: WorkflowRunStore;
-    readonly supervisor: WorkerProcessSupervisor;
-    readonly cleanup: WorkflowRunCleanup;
-    readonly doctor: WorkflowDoctor;
-    readonly artifactDirectory: string;
-    readonly requestKey: string;
-    readonly unsubscribe: () => void;
-  }>();
+  #attachment: WorkerRuntimeAttachment | undefined;
 
   constructor(
     publishFleet: (records: readonly FleetAgentRecord[]) => void = () => {},
@@ -134,36 +162,44 @@ export class PiWorkerRuntime implements WorkerToolProvider {
   }
 
   async start(ctx: ExtensionContext): Promise<void> {
+    this.#detach();
     await this.refreshGuidance(ctx);
-    const existing = this.#runtimes.get(ctx.cwd);
-    if (existing) {
-      this.#cwd = ctx.cwd;
-      this.#store = existing.store;
-      this.#supervisor = existing.supervisor;
-      this.#cleanup = existing.cleanup;
-      this.#doctor = existing.doctor;
-      this.#artifactDirectory = existing.artifactDirectory;
-      this.#requestKey = existing.requestKey;
-      existing.supervisor.updateLaunchAuthority(await launchAuthority(ctx));
-      this.#publish(await requiredRun(existing.store, existing.requestKey));
-      void this.#publishWithConversation(await requiredRun(existing.store, existing.requestKey));
-      return;
-    }
     try {
-      const runBudget = await confirmedRunBudget(ctx, this.#unlimitedDeclaration());
-      const authority = await launchAuthority(ctx);
-      const namespaceDirectory = workflowNamespaceDirectory(defaultPiStateDirectory(), ctx.cwd);
+      const workspacePath = canonicalWorkerWorkspace(ctx.cwd);
+      const workspaceId = workflowWorkspaceId(workspacePath);
+      const namespaceDirectory = workflowNamespaceDirectory(defaultPiStateDirectory(), workspacePath);
+      const key = workerRuntimeRegistryKey(namespaceDirectory, workspaceId);
+      const registry = processWorkerRuntimeRegistry();
       const inspector = new SystemProcessIdentityInspector();
       const leaseManager = new WorkflowLeaseManager({
         namespaceDirectory,
-        workspaceId: workflowWorkspaceId(ctx.cwd),
+        workspaceId,
         inspector,
       });
+      const registered = registry.entries.get(key);
+      if (registered !== undefined) {
+        const existing = parseWorkerRuntimeCore(registered, key);
+        await validateRuntimeAdoption(existing, {
+          workspacePath,
+          workspaceId,
+          namespaceDirectory,
+          inspector,
+          leaseManager,
+        });
+        existing.supervisor.updateLaunchAuthority(await launchAuthority(ctx));
+        const current = await requiredRun(existing.store, existing.requestKey);
+        this.#attach(existing, current);
+        this.#error = undefined;
+        return;
+      }
+
+      const runBudget = await confirmedRunBudget(ctx, this.#unlimitedDeclaration());
+      const authority = await launchAuthority(ctx);
       const acquired = await leaseManager.acquireSupervisor();
       if (acquired.kind !== "acquired") {
         throw new Error(
           acquired.kind === "held"
-            ? `worker supervisor is held by PID ${acquired.owner.pid}`
+            ? `worker supervisor is held by PID ${acquired.owner.pid}; fully restart Pi without deleting the supervisor lease`
             : acquired.diagnostic,
         );
       }
@@ -173,7 +209,7 @@ export class PiWorkerRuntime implements WorkerToolProvider {
         supervisorGeneration: acquired.owner.generation,
       });
       const sessionIdentity = ctx.sessionManager.getSessionFile() ?? `process:${process.pid}`;
-      const requestKey = `pi-session:${hash(`${ctx.cwd}\0${sessionIdentity}`)}`;
+      const requestKey = `pi-session:${hash(`${workspacePath}\0${sessionIdentity}`)}`;
       await store.startRun({
         requestKey,
         source: {
@@ -212,53 +248,37 @@ export class PiWorkerRuntime implements WorkerToolProvider {
         leases: leaseManager,
         inspector,
       });
-      const initial = await requiredRun(store, requestKey);
-      this.#baselineLifecycle(initial);
-      const notifier = new WorkerCompletionNotifier({ publish: this.#publishNotice });
-      notifier.baseline(initial);
-      const unsubscribeStore = store.subscribe((projection) => {
-        this.#publish(projection);
-        void this.#publishWithConversation(projection);
-        void this.#publishLifecycleChanges(projection).catch(() => {});
-        notifier.observe(projection);
-      });
-      const unsubscribe = (): void => {
-        unsubscribeStore();
-        notifier.dispose();
-      };
       const reconciled = await supervisor.reconcileRun(requestKey);
-      this.#cwd = ctx.cwd;
-      this.#store = store;
-      this.#supervisor = supervisor;
-      this.#cleanup = cleanup;
-      this.#doctor = doctor;
-      this.#artifactDirectory = artifactDirectory;
-      this.#requestKey = requestKey;
-      this.#runtimes.set(ctx.cwd, {
+      const core: WorkerRuntimeCoreV1 = {
+        schema: 1,
+        key,
+        workspacePath,
+        workspaceId,
+        namespaceDirectory,
+        processIdentity: processIdentityFromOwner(acquired.owner),
+        supervisorGeneration: acquired.owner.generation,
+        leaseManager,
         store,
         supervisor,
         cleanup,
         doctor,
         artifactDirectory,
         requestKey,
-        unsubscribe,
-      });
+      };
+      registry.entries.set(key, core);
+      this.#attach(core, reconciled);
       this.#error = undefined;
-      this.#publish(reconciled);
     } catch (error) {
+      this.#detach();
       this.#error = toError(error).message;
-      this.#store = undefined;
-      this.#supervisor = undefined;
-      this.#cleanup = undefined;
-      this.#doctor = undefined;
-      this.#artifactDirectory = undefined;
-      this.#requestKey = undefined;
     }
   }
 
   async shutdown(): Promise<void> {
     // Workers own a parent-liveness descriptor. Session changes do not imply cancellation,
-    // so the runtime intentionally retains the supervisor instead of inventing a stop policy.
+    // so only this extension instance's advisory subscribers are detached. The process-global
+    // core and its supervisor lease remain intact for a replacement extension instance.
+    this.#detach();
   }
 
   async refreshGuidance(ctx: ExtensionContext): Promise<void> {
@@ -278,7 +298,7 @@ export class PiWorkerRuntime implements WorkerToolProvider {
       this.#guidance = `${DRIVER_DOCTRINE}\n\n${guidance.text}${
         guidance.notices.length > 0 ? `\n${guidance.notices.join("\n")}` : ""
       }`;
-      if (this.#cwd === ctx.cwd) {
+      if (this.#cwd === canonicalWorkerWorkspace(ctx.cwd)) {
         this.#supervisor?.updateLaunchAuthority(await launchAuthority(ctx));
       }
     } catch (error) {
@@ -704,11 +724,67 @@ export class PiWorkerRuntime implements WorkerToolProvider {
     };
   }
 
+  #attach(core: WorkerRuntimeCoreV1, projection: RunProjectionV2): void {
+    this.#detach();
+    this.#cwd = core.workspacePath;
+    this.#store = core.store;
+    this.#supervisor = core.supervisor;
+    this.#cleanup = core.cleanup;
+    this.#doctor = core.doctor;
+    this.#artifactDirectory = core.artifactDirectory;
+    this.#requestKey = core.requestKey;
+    this.#lifecycleStates.clear();
+    this.#baselineLifecycle(projection);
+    const notifier = new WorkerCompletionNotifier({ publish: this.#publishNotice });
+    notifier.baseline(projection);
+    const attachment: WorkerRuntimeAttachment = {
+      core,
+      notifier,
+      active: true,
+      unsubscribe: () => {},
+    };
+    this.#attachment = attachment;
+    attachment.unsubscribe = core.store.subscribe((run) => {
+      if (!this.#isAttached(attachment)) return;
+      this.#publish(run);
+      void this.#publishWithConversation(run, attachment);
+      void this.#publishLifecycleChanges(run, attachment).catch(() => {});
+      notifier.observe(run);
+    });
+    this.#publish(projection);
+    void this.#publishWithConversation(projection, attachment);
+  }
+
+  #detach(): void {
+    const attachment = this.#attachment;
+    this.#attachment = undefined;
+    if (attachment) {
+      attachment.active = false;
+      attachment.unsubscribe();
+      attachment.notifier.dispose();
+    }
+    this.#lifecycleStates.clear();
+    this.#cwd = undefined;
+    this.#store = undefined;
+    this.#supervisor = undefined;
+    this.#cleanup = undefined;
+    this.#doctor = undefined;
+    this.#artifactDirectory = undefined;
+    this.#requestKey = undefined;
+  }
+
+  #isAttached(attachment: WorkerRuntimeAttachment): boolean {
+    return attachment.active && this.#attachment === attachment;
+  }
+
   #publish(run: RunProjectionV2): void {
     this.#publishFleet(projectFleetRecords(run));
   }
 
-  async #publishWithConversation(run: RunProjectionV2): Promise<void> {
+  async #publishWithConversation(
+    run: RunProjectionV2,
+    attachment: WorkerRuntimeAttachment,
+  ): Promise<void> {
     const records = await Promise.all(Object.values(run.workers).map(async (worker): Promise<FleetAgentRecord> => {
       const attempt = worker.attempts.at(-1);
       const contract = effectiveWorkerContract(worker);
@@ -739,7 +815,9 @@ export class PiWorkerRuntime implements WorkerToolProvider {
         ],
       };
     }));
-    if (this.#requestKey === run.requestKey) this.#publishFleet(records);
+    if (this.#isAttached(attachment) && attachment.core.requestKey === run.requestKey) {
+      this.#publishFleet(records);
+    }
   }
 
   #baselineLifecycle(run: RunProjectionV2): void {
@@ -748,16 +826,160 @@ export class PiWorkerRuntime implements WorkerToolProvider {
     }
   }
 
-  async #publishLifecycleChanges(run: RunProjectionV2): Promise<void> {
+  async #publishLifecycleChanges(
+    run: RunProjectionV2,
+    attachment: WorkerRuntimeAttachment,
+  ): Promise<void> {
     for (const worker of Object.values(run.workers)) {
+      if (!this.#isAttached(attachment)) return;
       const key = `${run.requestKey}\0${worker.workerId}`;
       if (this.#lifecycleStates.get(key) === worker.state) continue;
       this.#lifecycleStates.set(key, worker.state);
+      if (!this.#isAttached(attachment)) return;
       await this.#publishLifecycle(projectWorkerLifecycle(worker.workerId, worker.state)).catch(() => {});
     }
   }
 }
 
+function processWorkerRuntimeRegistry(): WorkerRuntimeRegistryV1 {
+  const existing = Reflect.get(globalThis, WORKER_RUNTIME_REGISTRY_SYMBOL) as unknown;
+  if (existing === undefined) {
+    const created: WorkerRuntimeRegistryV1 = { schema: 1, entries: new Map() };
+    Reflect.set(globalThis, WORKER_RUNTIME_REGISTRY_SYMBOL, created);
+    return created;
+  }
+  try {
+    if (!isRecord(existing) || existing["schema"] !== 1 || !(existing["entries"] instanceof Map)) {
+      throw new Error("registry envelope is invalid");
+    }
+    const registry = existing as unknown as WorkerRuntimeRegistryV1;
+    for (const [key, value] of registry.entries) {
+      if (typeof key !== "string" || key.length === 0) throw new Error("registry key is invalid");
+      parseWorkerRuntimeCore(value, key);
+    }
+    return registry;
+  } catch (error) {
+    throw reloadAdoptionRefused(`process-global registry is corrupt: ${toError(error).message}`);
+  }
+}
+
+function parseWorkerRuntimeCore(value: unknown, expectedKey: string): WorkerRuntimeCoreV1 {
+  if (!isRecord(value) || value["schema"] !== 1) {
+    throw reloadAdoptionRefused("process-global registry entry is corrupt");
+  }
+  const core = value as unknown as WorkerRuntimeCoreV1;
+  if (core.key !== expectedKey
+    || !nonEmptyString(core.workspacePath)
+    || !nonEmptyString(core.workspaceId)
+    || !nonEmptyString(core.namespaceDirectory)
+    || !nonEmptyString(core.supervisorGeneration)
+    || !nonEmptyString(core.artifactDirectory)
+    || !nonEmptyString(core.requestKey)
+    || !runtimeProcessIdentity(core.processIdentity)
+    || !hasMethods(core.leaseManager, ["inspectSupervisor"])
+    || !hasMethods(core.store, ["readRun", "subscribe"])
+    || !hasMethods(core.supervisor, ["updateLaunchAuthority", "steerWorker", "cancelWorker"])
+    || !hasMethods(core.cleanup, ["plan", "execute"])
+    || !hasMethods(core.doctor, ["inspect"])) {
+    throw reloadAdoptionRefused("process-global registry entry is corrupt");
+  }
+  return core;
+}
+
+async function validateRuntimeAdoption(
+  core: WorkerRuntimeCoreV1,
+  expected: {
+    readonly workspacePath: string;
+    readonly workspaceId: string;
+    readonly namespaceDirectory: string;
+    readonly inspector: ProcessIdentityInspector;
+    readonly leaseManager: WorkflowLeaseManager;
+  },
+): Promise<void> {
+  if (core.workspacePath !== expected.workspacePath
+    || core.workspaceId !== expected.workspaceId
+    || core.namespaceDirectory !== expected.namespaceDirectory
+    || core.key !== workerRuntimeRegistryKey(expected.namespaceDirectory, expected.workspaceId)) {
+    throw reloadAdoptionRefused("workspace or namespace identity changed");
+  }
+  const current = await expected.inspector.current();
+  if (!sameProcessIdentity(core.processIdentity, current)) {
+    throw reloadAdoptionRefused("process-start identity changed");
+  }
+  const lease = await expected.leaseManager.inspectSupervisor();
+  if (lease.kind !== "owned") {
+    throw reloadAdoptionRefused(
+      lease.kind === "corrupt" ? `supervisor lease is corrupt: ${lease.diagnostic}` : "supervisor lease is missing",
+    );
+  }
+  if (lease.owner.kind !== "supervisor"
+    || lease.owner.workspaceId !== expected.workspaceId
+    || lease.owner.generation !== core.supervisorGeneration) {
+    throw reloadAdoptionRefused("supervisor generation or workspace does not match the registry");
+  }
+  const leaseIdentity = processIdentityFromOwner(lease.owner);
+  if (!sameProcessIdentity(leaseIdentity, current)
+    || !sameProcessIdentity(leaseIdentity, core.processIdentity)) {
+    throw reloadAdoptionRefused("supervisor lease process identity does not match the registry");
+  }
+}
+
+function canonicalWorkerWorkspace(workspacePath: string): string {
+  if (!isAbsolute(workspacePath)) throw new TypeError("Worker workspace must be absolute");
+  return normalize(resolve(workspacePath));
+}
+
+function workerRuntimeRegistryKey(namespaceDirectory: string, workspaceId: string): string {
+  return `${normalize(resolve(namespaceDirectory))}\0${workspaceId}`;
+}
+
+function processIdentityFromOwner(owner: WorkflowLeaseOwnerV1): ProcessIdentity {
+  return {
+    pid: owner.pid,
+    processStartIdentity: owner.processStartIdentity,
+    bootIdentity: owner.bootIdentity,
+    executable: owner.executable,
+    hostname: owner.hostname,
+  };
+}
+
+function sameProcessIdentity(left: ProcessIdentity, right: ProcessIdentity): boolean {
+  return left.pid === right.pid
+    && left.processStartIdentity === right.processStartIdentity
+    && left.bootIdentity === right.bootIdentity
+    && normalize(left.executable) === normalize(right.executable)
+    && left.hostname === right.hostname;
+}
+
+function runtimeProcessIdentity(value: unknown): value is ProcessIdentity {
+  return isRecord(value)
+    && Number.isSafeInteger(value["pid"])
+    && (value["pid"] as number) > 0
+    && nonEmptyString(value["processStartIdentity"])
+    && (value["bootIdentity"] === null || nonEmptyString(value["bootIdentity"]))
+    && nonEmptyString(value["executable"])
+    && isAbsolute(value["executable"] as string)
+    && nonEmptyString(value["hostname"]);
+}
+
+function hasMethods(value: unknown, methods: readonly string[]): boolean {
+  return isRecord(value) && methods.every((method) => typeof value[method] === "function");
+}
+
+function nonEmptyString(value: unknown): value is string {
+  return typeof value === "string" && value.length > 0;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function reloadAdoptionRefused(reason: string): Error {
+  return new Error(
+    `worker runtime reload adoption refused: ${reason}; fully restart Pi without deleting the supervisor lease`,
+  );
+}
+
 export function projectFleetRecords(run: RunProjectionV2): readonly FleetAgentRecord[] {
   return Object.values(run.workers).map((worker): FleetAgentRecord => {
     const attempt = worker.attempts.at(-1);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-amq",
-  "version": "0.1.1",
+  "version": "0.1.2",
   "description": "A daemonless AMQ bridge and governed worker runtime for pi",
   "type": "module",
   "license": "MIT",
@@ -39,6 +39,7 @@
     "typecheck": "tsc --project tsconfig.json",
     "test:compile": "tsc --project tsconfig.test.json",
     "test": "npm run test:compile && node --test .test-dist/test/*.test.js",
+    "test:package": "npm run test:compile && node --test .test-dist/test/package/*.test.js",
     "prepack": "npm run build"
   },
   "peerDependencies": {