@zachwill/pi-orchestrate 0.7.2 → 0.9.0

package/extension/host.ts

@@ -1,27 +1,89 @@
-import { DeliveryCoordinator } from "./delivery.js";
+import { Cause, Effect, Exit, Layer, ManagedRuntime } from "effect";
 import {
-  createOrchestratorRuntime,
-  type OrchestratorRuntime,
+  Delivery,
+  deliveryLayer,
+  type DeliveryService,
+} from "./delivery.js";
+import {
+  Orchestration,
+  orchestrationLayer,
+  type AbortTarget,
+  type AcceptedRun,
+  type CompletedRun,
+  type OrchestrationContext,
+  type OrchestrationService,
+  SHUTDOWN_CLEANUP_GRACE_MS,
+  type RuntimeSnapshot,
+  type SettlementListener,
+  type UnsubscribeSettlement,
 } from "./runtime.js";
-import { createWorkerSessionFactory } from "./worker-session.js";
+import { createChildSessionsLayer } from "./worker-session.js";
+import type { OrchestrateTaskInput, RunMode } from "./domain.js";
 
 const PROCESS_HOST_KEY = Symbol.for("@zachwill/pi-orchestrate/process-host/v3");
 
+type RunResult<M extends RunMode> = M extends "async"
+  ? AcceptedRun
+  : CompletedRun;
+
+export interface OrchestratorRuntime {
+  orchestrate<M extends RunMode>(
+    context: OrchestrationContext,
+    task: OrchestrateTaskInput,
+    mode: M,
+    signal?: AbortSignal,
+    onSettlement?: SettlementListener,
+  ): Promise<RunResult<M>>;
+  sendInteractive<M extends RunMode>(
+    context: OrchestrationContext,
+    workerId: string,
+    instructions: string,
+    mode: M,
+    signal?: AbortSignal,
+    onSettlement?: SettlementListener,
+  ): Promise<RunResult<M>>;
+  abort(ownerSessionId: string, target: AbortTarget): Promise<void>;
+  closeInteractive(ownerSessionId: string, workerId: string): Promise<void>;
+  snapshot(ownerSessionId: string): Promise<RuntimeSnapshot>;
+  subscribeSettlement(listener: SettlementListener): UnsubscribeSettlement;
+  subscribeState(
+    ownerSessionId: string,
+    listener: (snapshot: RuntimeSnapshot) => void,
+  ): () => void;
+  shutdown(): Promise<void>;
+}
+
 export interface ProcessHost {
   readonly runtime: OrchestratorRuntime;
-  readonly delivery: DeliveryCoordinator;
+  readonly delivery: DeliveryService;
 }
 
 export interface ProcessHostAttachment {
   readonly host: ProcessHost;
 }
 
+export interface ProcessHostDestructionOptions {
+  /** Process-boundary deadline capability; production uses the orchestration cleanup grace. */
+  readonly awaitShutdown?: (
+    shutdown: Promise<void>,
+    graceMs: number,
+  ) => Promise<void>;
+  /** Process-boundary deadline capability; production uses the orchestration cleanup grace. */
+  readonly awaitRootDisposal?: (
+    disposal: Promise<void>,
+    graceMs: number,
+  ) => Promise<void>;
+}
+
 interface AttachmentAwareProcessHost extends ProcessHost {
   attachments?: Set<ProcessHostAttachment>;
 }
 
+type ProcessHostLifecycle = "active" | "destroying" | "destroyed";
+
 interface OwnedProcessHost extends AttachmentAwareProcessHost {
-  readonly unsubscribeSettlement?: () => void;
+  readonly effectRuntime?: ManagedRuntime.ManagedRuntime<Orchestration | Delivery, never>;
+  lifecycle?: ProcessHostLifecycle;
   destroyPromise?: Promise<void>;
 }
 
@@ -37,32 +99,190 @@ export function getProcessHost(): ProcessHost | undefined {
   return processGlobal()[PROCESS_HOST_KEY];
 }
 
+/** Pi-facing adapter. Every Promise operation executes one complete Orchestration Effect. */
+export class ProcessHostRuntimeAdapter<R = never> implements OrchestratorRuntime {
+  constructor(
+    private readonly effectRuntime: ManagedRuntime.ManagedRuntime<Orchestration | R, never>,
+    private readonly orchestration: OrchestrationService,
+  ) {}
+
+  orchestrate(
+    context: OrchestrationContext,
+    task: OrchestrateTaskInput,
+    mode: "async",
+    signal?: AbortSignal,
+    onSettlement?: SettlementListener,
+  ): Promise<AcceptedRun>;
+  orchestrate(
+    context: OrchestrationContext,
+    task: OrchestrateTaskInput,
+    mode: "inline",
+    signal?: AbortSignal,
+    onSettlement?: SettlementListener,
+  ): Promise<CompletedRun>;
+  orchestrate(
+    context: OrchestrationContext,
+    task: OrchestrateTaskInput,
+    mode: RunMode,
+    signal?: AbortSignal,
+    onSettlement?: SettlementListener,
+  ): Promise<AcceptedRun | CompletedRun>;
+  orchestrate(
+    context: OrchestrationContext,
+    task: OrchestrateTaskInput,
+    mode: RunMode,
+    signal?: AbortSignal,
+    onSettlement?: SettlementListener,
+  ): Promise<AcceptedRun | CompletedRun> {
+    return this.run(
+      this.orchestration.orchestrate(context, task, mode, onSettlement),
+      signal,
+    );
+  }
+
+  sendInteractive(
+    context: OrchestrationContext,
+    workerId: string,
+    instructions: string,
+    mode: "async",
+    signal?: AbortSignal,
+    onSettlement?: SettlementListener,
+  ): Promise<AcceptedRun>;
+  sendInteractive(
+    context: OrchestrationContext,
+    workerId: string,
+    instructions: string,
+    mode: "inline",
+    signal?: AbortSignal,
+    onSettlement?: SettlementListener,
+  ): Promise<CompletedRun>;
+  sendInteractive(
+    context: OrchestrationContext,
+    workerId: string,
+    instructions: string,
+    mode: RunMode,
+    signal?: AbortSignal,
+    onSettlement?: SettlementListener,
+  ): Promise<AcceptedRun | CompletedRun>;
+  sendInteractive(
+    context: OrchestrationContext,
+    workerId: string,
+    instructions: string,
+    mode: RunMode,
+    signal?: AbortSignal,
+    onSettlement?: SettlementListener,
+  ): Promise<AcceptedRun | CompletedRun> {
+    return this.run(
+      this.orchestration.sendInteractive(
+        context,
+        workerId,
+        instructions,
+        mode,
+        onSettlement,
+      ),
+      signal,
+    );
+  }
+
+  abort(ownerSessionId: string, target: AbortTarget): Promise<void> {
+    return this.run(this.orchestration.abort(ownerSessionId, target));
+  }
+
+  closeInteractive(ownerSessionId: string, workerId: string): Promise<void> {
+    return this.run(this.orchestration.closeInteractive(ownerSessionId, workerId));
+  }
+
+  snapshot(ownerSessionId: string): Promise<RuntimeSnapshot> {
+    // Snapshot is dependency-free and remains readable from a retained host reference
+    // after the process root has been disposed.
+    return Effect.runPromise(this.orchestration.snapshot(ownerSessionId));
+  }
+
+  subscribeSettlement(listener: SettlementListener): UnsubscribeSettlement {
+    return this.orchestration.subscribeSettlement(listener);
+  }
+
+  subscribeState(
+    ownerSessionId: string,
+    listener: (snapshot: RuntimeSnapshot) => void,
+  ): () => void {
+    return this.orchestration.subscribeState(ownerSessionId, listener);
+  }
+
+  shutdown(): Promise<void> {
+    return this.run(this.orchestration.shutdown());
+  }
+
+  private async run<A, E>(
+    effect: Effect.Effect<A, E>,
+    signal?: AbortSignal,
+  ): Promise<A> {
+    if (signal?.aborted) throw abortSignalReason(signal);
+    if (!signal) return this.effectRuntime.runPromise(effect);
+
+    const signalInterruption = {};
+    const exit = await this.effectRuntime.runPromiseExit(
+      Effect.raceFirst(effect, abortSignalEffect(signal, signalInterruption)),
+    );
+    if (Exit.isSuccess(exit)) return exit.value;
+
+    const error = Cause.squash(exit.cause);
+    if (error === signalInterruption) throw abortSignalReason(signal);
+    throw error;
+  }
+}
+
+export function createProcessHostRuntimeAdapter<R>(
+  effectRuntime: ManagedRuntime.ManagedRuntime<Orchestration | R, never>,
+): OrchestratorRuntime {
+  // Orchestration acquisition is synchronous; subscriptions must remain reentrant.
+  const orchestration = effectRuntime.runSync(Orchestration);
+  return new ProcessHostRuntimeAdapter(effectRuntime, orchestration);
+}
+
+export function createProcessApplicationLayer(): Layer.Layer<Orchestration | Delivery> {
+  const orchestration = orchestrationLayer().pipe(
+    Layer.provide(createChildSessionsLayer()),
+  );
+  return deliveryLayer.pipe(Layer.provideMerge(orchestration));
+}
+
 export function createProcessHost(): ProcessHost {
   const global = processGlobal();
   const existing = global[PROCESS_HOST_KEY];
-  if (existing) return existing;
+  if (existing?.lifecycle === "destroying") {
+    throw new Error("Cannot create a process host while the current host is being destroyed");
+  }
+  if (existing?.lifecycle === "destroyed") {
+    delete global[PROCESS_HOST_KEY];
+  } else if (existing) {
+    existing.lifecycle = "active";
+    return existing;
+  }
 
-  const runtime = createOrchestratorRuntime({
-    workerSessionFactory: createWorkerSessionFactory(),
-  });
-  const delivery = new DeliveryCoordinator();
+  const effectRuntime = ManagedRuntime.make(createProcessApplicationLayer());
+  const runtime = createProcessHostRuntimeAdapter(effectRuntime);
+  const delivery = effectRuntime.runSync(Delivery);
   const host: OwnedProcessHost = {
     runtime,
     delivery,
+    effectRuntime,
     attachments: new Set(),
-    unsubscribeSettlement: runtime.subscribeSettlement((settlement) => {
-      delivery.accept(settlement);
-    }),
+    lifecycle: "active",
   };
   global[PROCESS_HOST_KEY] = host;
   return host;
 }
 
 export function attachProcessHost(host: ProcessHost): ProcessHostAttachment {
+  const ownedHost = host as OwnedProcessHost;
+  if (ownedHost.lifecycle === "destroying" || ownedHost.lifecycle === "destroyed") {
+    throw new Error(`Cannot attach to a ${ownedHost.lifecycle} process host`);
+  }
+
   const attachment: ProcessHostAttachment = { host };
-  const attachmentAwareHost = host as AttachmentAwareProcessHost;
-  attachmentAwareHost.attachments ??= new Set();
-  attachmentAwareHost.attachments.add(attachment);
+  ownedHost.attachments ??= new Set();
+  ownedHost.attachments.add(attachment);
   return attachment;
 }
 
@@ -77,27 +297,65 @@ export function detachProcessHost(
   return attachments.size === 0;
 }
 
-export async function destroyProcessHost(host: ProcessHost): Promise<void> {
+export function destroyProcessHost(
+  host: ProcessHost,
+  options: ProcessHostDestructionOptions = {},
+): Promise<void> {
   const ownedHost = host as OwnedProcessHost;
-  if ((ownedHost.attachments?.size ?? 0) > 0) return;
-  if (ownedHost.destroyPromise) {
-    await ownedHost.destroyPromise;
-    return;
-  }
+  if (ownedHost.destroyPromise) return ownedHost.destroyPromise;
+  if ((ownedHost.attachments?.size ?? 0) > 0) return Promise.resolve();
+  if (ownedHost.lifecycle === "destroyed") return Promise.resolve();
+
+  ownedHost.lifecycle = "destroying";
+  let resolveDestruction!: () => void;
+  let rejectDestruction!: (error: unknown) => void;
+  const destroyPromise = new Promise<void>((resolve, reject) => {
+    resolveDestruction = resolve;
+    rejectDestruction = reject;
+  });
+  // Publish the exact shared result before shutdown can synchronously reenter destruction.
+  ownedHost.destroyPromise = destroyPromise;
 
-  ownedHost.destroyPromise = (async () => {
+  let shutdown: Promise<void>;
+  try {
+    // shutdown() closes Orchestration admission before returning its bounded teardown Promise.
+    shutdown = ownedHost.runtime.shutdown();
+  } catch (error) {
+    shutdown = Promise.reject(error);
+  }
+  const awaitShutdown =
+    options.awaitShutdown ?? awaitPromiseWithinCleanupGrace;
+  const awaitRootDisposal =
+    options.awaitRootDisposal ?? awaitPromiseWithinCleanupGrace;
+  const teardown = (async () => {
     try {
-      await ownedHost.runtime.shutdown();
+      // A timed-out shutdown keeps physical best-effort finalizers alive. Observe
+      // either late outcome while proceeding to the separately bounded root disposal.
+      void shutdown.catch(() => {});
+      await awaitShutdown(shutdown, SHUTDOWN_CLEANUP_GRACE_MS);
     } finally {
-      ownedHost.delivery.clear();
-      ownedHost.unsubscribeSettlement?.();
-      const global = processGlobal();
-      if (global[PROCESS_HOST_KEY] === ownedHost) {
-        delete global[PROCESS_HOST_KEY];
+      try {
+        const disposal = ownedHost.effectRuntime?.dispose();
+        if (disposal) {
+          // ManagedRuntime disposal keeps ownership of uninterruptible finalizers after timeout.
+          // Observe its eventual rejection while host destruction proceeds best-effort.
+          void disposal.catch(() => {});
+          await awaitRootDisposal(disposal, SHUTDOWN_CLEANUP_GRACE_MS);
+        }
+      } finally {
+        // After either deadline the host is logically destroyed and detached even
+        // though abandoned physical finalizers may still settle. A replacement may
+        // overlap only that cleanup; identity guards keep stale completion harmless.
+        ownedHost.lifecycle = "destroyed";
+        const global = processGlobal();
+        if (global[PROCESS_HOST_KEY] === ownedHost) {
+          delete global[PROCESS_HOST_KEY];
+        }
       }
     }
   })();
-  await ownedHost.destroyPromise;
+  void teardown.then(resolveDestruction, rejectDestruction);
+  return destroyPromise;
 }
 
 export async function quitProcessHost(): Promise<void> {
@@ -105,3 +363,35 @@ export async function quitProcessHost(): Promise<void> {
   if (!host) return;
   await destroyProcessHost(host);
 }
+
+function awaitPromiseWithinCleanupGrace(
+  promise: Promise<void>,
+  graceMs: number,
+): Promise<void> {
+  return Effect.runPromise(
+    Effect.promise(() => promise).pipe(
+      Effect.timeoutOption(graceMs),
+      Effect.asVoid,
+    ),
+  );
+}
+
+function abortSignalEffect(
+  signal: AbortSignal,
+  interruption: object,
+): Effect.Effect<never, object> {
+  return Effect.callback((resume) => {
+    const onAbort = () => resume(Effect.fail(interruption));
+    signal.addEventListener("abort", onAbort, { once: true });
+    if (signal.aborted) {
+      signal.removeEventListener("abort", onAbort);
+      onAbort();
+    }
+    return Effect.sync(() => signal.removeEventListener("abort", onAbort));
+  });
+}
+
+function abortSignalReason(signal: AbortSignal): unknown {
+  if ("reason" in signal) return signal.reason;
+  return new DOMException("This operation was aborted", "AbortError");
+}

package/extension/index.ts

@@ -52,36 +52,31 @@ export function createOrchestrationExtension(
   dependencies: OrchestrationExtensionDependencies = {},
 ): ExtensionFactory {
   return (pi) => {
-    const host = dependencies.getHost?.() ?? createProcessHost();
     const discoverCatalog = dependencies.discoverCatalog ?? discoverWorkerCatalog;
-    const statusController =
-      dependencies.createStatusController?.(host.runtime) ??
-      createStatusController(host.runtime);
     const dispatchDecisions = new Map<string, StoredDispatchDecision>();
+    let host: ProcessHost | undefined;
     let hostAttachment: ProcessHostAttachment | undefined;
+    let statusController: StatusController | undefined;
+    let toolsRegistered = false;
     let activeBinding: OwnerBinding | undefined;
     let cachedCatalog: WorkerCatalog | undefined;
 
-    const catalogFor = (ctx: ExtensionContext): WorkerCatalog => {
-      if (cachedCatalog) return cachedCatalog;
-      cachedCatalog = discoverCatalog({
+    const discoverCatalogFor = (ctx: ExtensionContext): WorkerCatalog =>
+      discoverCatalog({
         cwd: ctx.cwd,
         projectTrusted: ctx.isProjectTrusted(),
       });
+
+    const catalogFor = (ctx: ExtensionContext): WorkerCatalog => {
+      if (cachedCatalog) return cachedCatalog;
+      cachedCatalog = discoverCatalogFor(ctx);
       return cachedCatalog;
     };
 
-    registerOrchestrationTools(pi, {
-      runtime: host.runtime,
-      getCatalog: catalogFor,
-      getDispatchDecision: (toolCallId) =>
-        dispatchDecisions.get(toolCallId) ?? { mode: "inline" },
-    });
     registerOrchestrationPresentation(pi);
 
     pi.on("session_start", (_event, ctx) => {
-      hostAttachment ??= attachProcessHost(host);
-      if (activeBinding) {
+      if (activeBinding && host && statusController) {
         host.delivery.unbind(
           activeBinding.ownerSessionId,
           activeBinding.generation,
@@ -89,6 +84,21 @@ export function createOrchestrationExtension(
         statusController.unbind(activeBinding.ownerSessionId);
       }
 
+      host ??= dependencies.getHost?.() ?? createProcessHost();
+      statusController ??=
+        dependencies.createStatusController?.(host.runtime) ??
+        createStatusController(host.runtime);
+      if (!toolsRegistered) {
+        registerOrchestrationTools(pi, {
+          runtime: host.runtime,
+          getCatalog: catalogFor,
+          getDispatchDecision: (toolCallId) =>
+            dispatchDecisions.get(toolCallId) ?? { mode: "inline" },
+        });
+        toolsRegistered = true;
+      }
+      hostAttachment ??= attachProcessHost(host);
+
       dispatchDecisions.clear();
       cachedCatalog = undefined;
       const binding: OwnerBinding = {
@@ -106,10 +116,7 @@ export function createOrchestrationExtension(
     });
 
     pi.on("before_agent_start", (event, ctx) => {
-      cachedCatalog = discoverCatalog({
-        cwd: ctx.cwd,
-        projectTrusted: ctx.isProjectTrusted(),
-      });
+      cachedCatalog = discoverCatalogFor(ctx);
       return {
         systemPrompt: appendOrchestratorContract(
           event.systemPrompt,
@@ -125,16 +132,16 @@ export function createOrchestrationExtension(
       );
       const ownerSessionId = activeBinding?.ownerSessionId;
       if (!ownerSessionId) return;
-      const groupedOrchestration =
+      const isOrchestrateGroup =
         toolCalls.length > 1 &&
         toolCalls.every((toolCall) => toolCall.name === "orchestrate");
-      const synthesisGroup = groupedOrchestration
+      const synthesisGroup = isOrchestrateGroup
         ? { id: `orchestrate:${toolCalls[0]?.id ?? "group"}`, size: toolCalls.length }
         : undefined;
 
       for (const toolCall of toolCalls) {
         if (!DISPATCH_TOOL_NAMES.has(toolCall.name)) continue;
-        const mode = groupedOrchestration || toolCalls.length === 1
+        const mode = isOrchestrateGroup || toolCalls.length === 1
           ? "async"
           : "inline";
         dispatchDecisions.set(toolCall.id, {
@@ -151,7 +158,7 @@ export function createOrchestrationExtension(
       const decision = dispatchDecisions.get(event.toolCallId);
       dispatchDecisions.delete(event.toolCallId);
       if (!event.isError || !decision?.synthesisGroup) return;
-      host.delivery.skipSynthesisGroupMember(
+      host?.delivery.skipSynthesisGroupMember(
         decision.ownerSessionId,
         decision.synthesisGroup.id,
         decision.synthesisGroup.size,
@@ -161,7 +168,7 @@ export function createOrchestrationExtension(
     pi.on("agent_start", () => {
       const binding = activeBinding;
       if (!binding) return;
-      host.delivery.markAgentStarted(
+      host?.delivery.markAgentStarted(
         binding.ownerSessionId,
         binding.generation,
       );
@@ -170,7 +177,7 @@ export function createOrchestrationExtension(
     pi.on("agent_settled", () => {
       const binding = activeBinding;
       if (!binding) return;
-      host.delivery.markAgentSettled(
+      host?.delivery.markAgentSettled(
         binding.ownerSessionId,
         binding.generation,
       );
@@ -182,17 +189,17 @@ export function createOrchestrationExtension(
       cachedCatalog = undefined;
       dispatchDecisions.clear();
 
-      if (binding) {
+      if (binding && host) {
         host.delivery.unbind(binding.ownerSessionId, binding.generation);
       }
-      statusController.dispose();
+      statusController?.dispose();
 
       const attachment = hostAttachment;
       hostAttachment = undefined;
-      const wasLastAttachment = attachment
+      const wasLastAttachment = host && attachment
         ? detachProcessHost(host, attachment)
         : false;
-      if (event.reason === "quit" && wasLastAttachment) {
+      if (event.reason === "quit" && wasLastAttachment && host) {
         await (dependencies.destroyHost ?? destroyProcessHost)(host);
       }
     });