@ilya-lesikov/pi-pi 0.4.0 → 0.6.0

package/extensions/orchestrator/transition-controller.test.ts

@@ -0,0 +1,207 @@
+import { describe, expect, it, vi } from "vitest";
+import { TransitionController, type TransitionHost, type MainSession } from "./transition-controller.js";
+
+// Fake host + session: records outbound primitives and lets tests control idle +
+// compaction. `controller` is pre-wired with both.
+function makeHost(opts: { idle?: boolean; step?: string | null; canCompact?: boolean } = {}) {
+  let compactCb: { onComplete?: () => void; onError?: (e: Error) => void } | null = null;
+  const calls = {
+    userMessages: [] as Array<{ text: string; deliverAs: string }>,
+    customMessages: [] as Array<{ customType: string; deliverAs: string }>,
+    compactStarted: 0,
+  };
+  const session: MainSession = {
+    sendUserMessage: (text, options) => calls.userMessages.push({ text, deliverAs: options?.deliverAs ?? "" }),
+    sendMessage: (message, options) => calls.customMessages.push({ customType: message.customType, deliverAs: options?.deliverAs ?? "" }),
+  };
+  const host: TransitionHost = {
+    compact: (options) => {
+      if (opts.canCompact === false) return false;
+      calls.compactStarted++;
+      compactCb = options;
+      return true;
+    },
+    isIdle: () => opts.idle ?? false,
+    currentStep: () => opts.step ?? "llm_work",
+  };
+  return {
+    host,
+    session,
+    calls,
+    controller: new TransitionController(host, session),
+    // Simulate the SDK firing the compaction completion event.
+    completeCompaction: () => compactCb?.onComplete?.(),
+    failCompaction: (msg: string) => compactCb?.onError?.(new Error(msg)),
+  };
+}
+
+describe("TransitionController.send", () => {
+  it("delivers instruction as followUp (always starts a turn)", () => {
+    const { host, session, calls } = makeHost();
+    const c = new TransitionController(host, session);
+    c.send("go", "instruction");
+    expect(calls.userMessages).toEqual([{ text: "go", deliverAs: "followUp" }]);
+  });
+
+  it("rejects context role (plain user messages always start a turn)", () => {
+    const { host, session, calls } = makeHost();
+    const c = new TransitionController(host, session);
+    expect(() => c.send("ctx", "context")).toThrow(/sendCustom/);
+    expect(calls.userMessages).toHaveLength(0);
+  });
+
+  it("sendCustom maps roles to delivery modes", () => {
+    const { host, session, calls } = makeHost();
+    const c = new TransitionController(host, session);
+    c.sendCustom({ customType: "pp-context", content: "x", display: false }, "context");
+    c.sendCustom({ customType: "pp-artifact", content: "y", display: false }, "instruction");
+    expect(calls.customMessages).toEqual([
+      { customType: "pp-context", deliverAs: "steer" },
+      { customType: "pp-artifact", deliverAs: "followUp" },
+    ]);
+  });
+});
+
+describe("TransitionController.isRunning / shouldBlockAgentStart", () => {
+  it("running with llm_work step is running and does not block", () => {
+    const { host, session } = makeHost({ step: "llm_work" });
+    const c = new TransitionController(host, session);
+    expect(c.isRunning()).toBe(true);
+    expect(c.shouldBlockAgentStart()).toBe(false);
+  });
+
+  it("await_planners and await_reviewers steps are not running (block)", () => {
+    for (const step of ["await_planners", "await_reviewers"]) {
+      const { host, session } = makeHost({ step });
+      const c = new TransitionController(host, session);
+      expect(c.isRunning()).toBe(false);
+      expect(c.shouldBlockAgentStart()).toBe(true);
+    }
+  });
+
+  it("gateAgentStart aborts and reports when not running, no-op when running", () => {
+    const running = makeHost({ step: "llm_work" });
+    const rc = new TransitionController(running.host, running.session);
+    const abort1 = vi.fn();
+    expect(rc.gateAgentStart(abort1)).toBe(false);
+    expect(abort1).not.toHaveBeenCalled();
+
+    const waiting = makeHost({ step: "await_planners" });
+    const wc = new TransitionController(waiting.host, waiting.session);
+    const abort2 = vi.fn();
+    expect(wc.gateAgentStart(abort2)).toBe(true);
+    expect(abort2).toHaveBeenCalledOnce();
+  });
+
+  it("abortMainAgent issues the abort (controller-owned cleanup abort)", () => {
+    const { controller } = makeHost();
+    const abort = vi.fn();
+    controller.abortMainAgent(abort);
+    expect(abort).toHaveBeenCalledOnce();
+    expect(() => controller.abortMainAgent(undefined)).not.toThrow();
+  });
+});
+
+describe("TransitionController phase transition flow", () => {
+  it("waits for agent_end when not idle, then compacts and resumes", async () => {
+    const { host, session, calls, completeCompaction } = makeHost({ idle: false });
+    const c = new TransitionController(host, session);
+    const onResume = vi.fn();
+    const p = c.requestTransition({ kind: "phase", summary: "s", onResume, instruction: "Begin working." });
+    expect(c.getState()).toBe("pending");
+    expect(calls.compactStarted).toBe(0); // not idle -> waits
+
+    c.onAgentEnd();
+    expect(c.getState()).toBe("compacting");
+    expect(calls.compactStarted).toBe(1);
+
+    completeCompaction();
+    await p;
+    expect(onResume).toHaveBeenCalledOnce();
+    expect(c.getState()).toBe("running");
+    expect(calls.userMessages).toEqual([{ text: "Begin working.", deliverAs: "followUp" }]);
+  });
+
+  it("already-idle path compacts immediately without agent_end", async () => {
+    const { host, session, calls, completeCompaction } = makeHost({ idle: true });
+    const c = new TransitionController(host, session);
+    const p = c.requestTransition({ kind: "done", summary: "done" });
+    expect(c.getState()).toBe("compacting");
+    expect(calls.compactStarted).toBe(1);
+    completeCompaction();
+    await p;
+    expect(c.getState()).toBe("running");
+  });
+
+  it("treats compact no-op errors as a clean resume", async () => {
+    for (const msg of ["Nothing to compact (session too small)", "Already compacted"]) {
+      const { host, session, failCompaction } = makeHost({ idle: true });
+      const c = new TransitionController(host, session);
+      const onResume = vi.fn();
+      const p = c.requestTransition({ kind: "phase", onResume, instruction: "go" });
+      failCompaction(msg);
+      await p;
+      expect(onResume).toHaveBeenCalledOnce();
+      expect(c.getState()).toBe("running");
+    }
+  });
+
+  it("resumes via session_compact event", async () => {
+    const { host, session } = makeHost({ idle: false });
+    const c = new TransitionController(host, session);
+    const onResume = vi.fn();
+    const p = c.requestTransition({ kind: "phase", onResume, instruction: "go" });
+    c.onAgentEnd();
+    c.onSessionCompact();
+    await p;
+    expect(onResume).toHaveBeenCalledOnce();
+    expect(c.getState()).toBe("running");
+  });
+
+  it("ignores agent_end while running (no self-trigger loop)", () => {
+    const { host, session, calls } = makeHost({ idle: false });
+    const c = new TransitionController(host, session);
+    c.onAgentEnd();
+    c.onAgentEnd();
+    expect(calls.compactStarted).toBe(0);
+    expect(c.getState()).toBe("running");
+  });
+
+  it("ignores session_compact while running (no spurious resume)", () => {
+    const { host, session, calls } = makeHost();
+    const c = new TransitionController(host, session);
+    c.onSessionCompact();
+    expect(calls.userMessages).toHaveLength(0);
+    expect(c.getState()).toBe("running");
+  });
+
+  it("resolves the awaitable even when no live ctx can compact", async () => {
+    const { host, session, calls } = makeHost({ idle: true, canCompact: false });
+    const c = new TransitionController(host, session);
+    const onResume = vi.fn();
+    await c.requestTransition({ kind: "done", onResume });
+    expect(onResume).toHaveBeenCalledOnce();
+    expect(calls.compactStarted).toBe(0);
+    expect(c.getState()).toBe("running");
+  });
+
+  it("does not send an instruction when none is provided (await_planners notify-only)", async () => {
+    const { host, session, calls, completeCompaction } = makeHost({ idle: true });
+    const c = new TransitionController(host, session);
+    const p = c.requestTransition({ kind: "phase", onResume: () => {} });
+    completeCompaction();
+    await p;
+    expect(calls.userMessages).toHaveLength(0);
+  });
+
+  it("compaction completion is idempotent across onComplete and session_compact", async () => {
+    const { host, session, completeCompaction } = makeHost({ idle: true });
+    const c = new TransitionController(host, session);
+    const onResume = vi.fn();
+    const p = c.requestTransition({ kind: "phase", onResume, instruction: "go" });
+    completeCompaction();
+    c.onSessionCompact(); // second terminus — must be a no-op
+    await p;
+    expect(onResume).toHaveBeenCalledOnce();
+  });
+});

package/extensions/orchestrator/transition-controller.ts

@@ -0,0 +1,259 @@
+import { getLogger } from "./log.js";
+
+// Role of an outbound message. Drives the SDK delivery mode:
+//   - "context": injected background context/artifacts. Delivered as "steer" so
+//     it never starts a turn (the agent is or will be working).
+//   - "instruction": a handoff that should make the agent act. Delivered as
+//     "followUp" so it queues and triggers a turn once the current one settles
+//     (or runs immediately when idle). Never throws "Agent is already processing".
+export type SendRole = "context" | "instruction";
+
+// Dependency passed to standalone phase-module spawn functions so they can emit
+// status/context messages WITHOUT importing the orchestrator (avoids circular
+// imports) while still routing every send through the controller. Bound to
+// TransitionController.sendCustom by the caller.
+export type PhaseSend = (
+  message: { customType: string; content: string; display: boolean; details?: unknown },
+  role: SendRole,
+) => void;
+
+// The kind of transition the controller is coordinating. Phase transitions resume
+// the main loop with a "Begin working" instruction; done/stop transitions finish
+// the task (cleanup already performed by the caller) and resolve their awaitable.
+export type TransitionKind = "phase" | "done";
+
+export interface TransitionRequest {
+  kind: TransitionKind;
+  // Compaction summary used by the session_before_compact handler.
+  summary?: string;
+  // Runs after compaction completes (or is skipped), before the resume message.
+  // Used for phase transitions to switch model + inject context/artifacts and to
+  // spawn planners at the right moment. Throwing here is logged, not fatal.
+  onResume?: () => void | Promise<void>;
+  // For phase transitions: the instruction sent after onResume. When omitted
+  // (e.g. plan/await_planners which only notifies, or done transitions), no
+  // instruction is sent.
+  instruction?: string;
+}
+
+export type ControllerState = "running" | "pending" | "compacting" | "resuming";
+
+// The raw main-session messaging surface (satisfied by ExtensionAPI / pi). The
+// controller calls these directly so it is literally the only code that calls
+// sendUserMessage/sendMessage — a whole-tree grep finds zero raw sends elsewhere.
+export interface MainSession {
+  sendUserMessage(content: string, options?: { deliverAs?: "steer" | "followUp" }): void;
+  sendMessage(
+    message: { customType: string; content: string; display: boolean; details?: unknown },
+    options?: { deliverAs?: "steer" | "followUp" | "nextTurn" },
+  ): void;
+}
+
+// Live-session surface the controller needs from the Orchestrator (the bits that
+// depend on the current ExtensionContext / persisted state). Kept as an interface
+// so the controller is unit-testable with a fake host.
+export interface TransitionHost {
+  // Compaction + idle probe come from the live ExtensionContext (lastCtx).
+  compact(options: { customInstructions?: string; onComplete?: () => void; onError?: (err: Error) => void }): boolean;
+  isIdle(): boolean;
+  // Persisted step, used to derive the await_* "not running" predicate.
+  currentStep(): string | null;
+}
+
+const COMPACT_NOOP_MESSAGES = [
+  "Nothing to compact (session too small)",
+  "Already compacted",
+];
+
+function isCompactNoop(err: unknown): boolean {
+  const msg = err instanceof Error ? err.message : String(err ?? "");
+  return COMPACT_NOOP_MESSAGES.some((m) => msg.includes(m));
+}
+
+// Single event-driven owner of phase/task transitions and main-session
+// compaction + outbound messaging. See transition-controller flow in the plan:
+//   requestTransition -> pending; agent_end (or already-idle) -> compact ->
+//   compacting; session_compact / compact no-op / error -> resuming -> onResume
+//   -> send(instruction) -> running.
+export class TransitionController {
+  private state: ControllerState = "running";
+  private active: TransitionRequest | null = null;
+  // Resolvers for callers that await the transition (done/stop/new-task paths).
+  private waiters: Array<() => void> = [];
+
+  constructor(
+    private readonly host: TransitionHost,
+    // The raw main-session messaging surface (pi). The controller calls it
+    // directly so no other code issues raw sendUserMessage/sendMessage.
+    private readonly session: MainSession,
+  ) {}
+
+  // Pre-bound PhaseSend handed to standalone spawn functions so they route every
+  // outbound message through the controller without importing the orchestrator.
+  readonly phaseSend: PhaseSend = (message, role) => this.sendCustom(message, role);
+
+  getState(): ControllerState {
+    return this.state;
+  }
+
+  // True only when the controller is idle (running) AND not waiting on subagents.
+  // The await_* states live in persisted step, not the enum. Every consumer of
+  // "may the agent loop start / may we nudge / may the menu proceed" reads THIS.
+  isRunning(): boolean {
+    if (this.state !== "running") return false;
+    const step = this.host.currentStep();
+    return step !== "await_planners" && step !== "await_reviewers";
+  }
+
+  // Inverse predicate used by in-tool / menu code to decide whether to abort.
+  shouldBlockAgentStart(): boolean {
+    return !this.isRunning();
+  }
+
+  // The transition/await abort gate, owned by the controller. before_agent_start
+  // delegates here: the controller decides AND issues the abort (via the passed
+  // abort fn from the live ctx) so it is the sole owner of transition aborts.
+  // Returns true if the agent start was aborted.
+  gateAgentStart(abort: () => void): boolean {
+    if (this.isRunning()) return false;
+    getLogger().debug({ s: "controller", state: this.state }, "aborting agent start (not running)");
+    abort();
+    return true;
+  }
+
+  // Abort the in-flight main-agent turn ahead of an interactive pause/stop/finish
+  // cleanup. Routed through the controller so it is the sole owner of main-session
+  // aborts; the caller then performs cleanup and requests a done transition.
+  abortMainAgent(abort: (() => void) | undefined): void {
+    abort?.();
+  }
+
+  // True while a transition is mid-flight (pending/compacting/resuming) — i.e.
+  // the controller initiated the current compaction. Used by session_before_compact
+  // to decide between supplying the transition summary vs. re-injecting artifacts
+  // after a natural (user-triggered) compaction.
+  isTransitioning(): boolean {
+    return this.state !== "running";
+  }
+
+  // Compaction summary for the in-flight transition (empty when not transitioning).
+  currentSummary(): string {
+    return this.active?.summary ?? "";
+  }
+
+  // Outbound plain user-message. The ONLY path for main-session user messaging.
+  // Plain user messages (pi.sendUserMessage) ALWAYS start a turn — even with
+  // deliverAs:"steer" when idle (SDK prompt()) — so this only serves the
+  // "instruction" role (queue + start a turn via followUp). A non-turn-starting
+  // "context" payload must use sendCustom (a custom message that only appends
+  // when idle), so passing "context" here is a misuse and is rejected.
+  send(text: string, role: SendRole): void {
+    if (role === "context") {
+      throw new Error("send(context) would start a turn; use sendCustom for non-turn-starting context");
+    }
+    this.session.sendUserMessage(text, { deliverAs: "followUp" });
+  }
+
+  // Custom (non-LLM) messages (pp-context / pp-artifact). "context" => steer
+  // (append-only when idle, never starts a turn); "instruction" => followUp
+  // (queues + starts a turn while streaming; callers that need a guaranteed turn
+  // when idle pair this with a send(...) instruction).
+  sendCustom(message: { customType: string; content: string; display: boolean; details?: unknown }, role: SendRole): void {
+    this.session.sendMessage(message, { deliverAs: role === "context" ? "steer" : "followUp" });
+  }
+
+  // Request a transition. Returns a promise that resolves once the transition
+  // has resumed (or finished, for done/stop). If the agent is already idle the
+  // controller compacts immediately; otherwise it waits for agent_end.
+  requestTransition(req: TransitionRequest): Promise<void> {
+    const log = getLogger();
+    if (this.state !== "running") {
+      // A transition is already in flight. Coalesce: ignore the new request but
+      // still attach the waiter so the caller is released when this settles.
+      if (req.kind === "phase") {
+        log.warn({ s: "controller", state: this.state }, "phase transition coalesced while not running — onResume/instruction dropped");
+      } else {
+        log.debug({ s: "controller", state: this.state, kind: req.kind }, "requestTransition while not running — coalescing");
+      }
+      return new Promise<void>((resolve) => this.waiters.push(resolve));
+    }
+    this.active = req;
+    this.state = "pending";
+    log.debug({ s: "controller", kind: req.kind }, "transition pending");
+    const promise = new Promise<void>((resolve) => this.waiters.push(resolve));
+    // Already-idle path: no agent turn is in flight, so no agent_end will fire.
+    // Compact immediately rather than waiting forever.
+    if (this.host.isIdle()) {
+      this.beginCompaction();
+    }
+    return promise;
+  }
+
+  // agent_end fires when the main loop goes idle. If a transition is pending,
+  // this is the moment to compact. Repeated agent_end while running is a no-op
+  // (prevents the self-trigger loop from injected followUp messages).
+  onAgentEnd(): void {
+    if (this.state !== "pending") return;
+    this.beginCompaction();
+  }
+
+  // session_compact fires after compaction completes. Only advance the transition
+  // we initiated (state must be compacting) — an unrelated/manual compaction
+  // while running must not resume the wrong transition.
+  onSessionCompact(): void {
+    if (this.state !== "compacting") return;
+    void this.resume();
+  }
+
+  private beginCompaction(): void {
+    const req = this.active;
+    if (!req) {
+      this.state = "running";
+      return;
+    }
+    this.state = "compacting";
+    getLogger().debug({ s: "controller", kind: req.kind }, "compacting");
+    const started = this.host.compact({
+      customInstructions: req.summary,
+      // onComplete is a fallback resume terminus; the authoritative one is the
+      // session_compact event. Whichever lands first wins (resume is idempotent).
+      onComplete: () => {
+        if (this.state === "compacting") void this.resume();
+      },
+      onError: (err: Error) => {
+        if (this.state !== "compacting") return;
+        if (isCompactNoop(err)) {
+          getLogger().debug({ s: "controller", err: err.message }, "compact no-op — resuming");
+          void this.resume();
+        } else {
+          getLogger().error({ s: "controller", err: err.message }, "compact error — resuming anyway");
+          void this.resume();
+        }
+      },
+    });
+    // host.compact returns false when no ctx is available (e.g. no live session);
+    // treat as a no-op and resume so awaiting callers don't hang.
+    if (!started) {
+      void this.resume();
+    }
+  }
+
+  private async resume(): Promise<void> {
+    const req = this.active;
+    this.state = "resuming";
+    getLogger().debug({ s: "controller", kind: req?.kind }, "resuming");
+    try {
+      await req?.onResume?.();
+    } catch (err: any) {
+      getLogger().error({ s: "controller", err: err?.message ?? String(err) }, "onResume failed");
+    }
+    if (req?.instruction) {
+      this.send(req.instruction, "instruction");
+    }
+    this.active = null;
+    this.state = "running";
+    const waiters = this.waiters;
+    this.waiters = [];
+    for (const w of waiters) w();
+  }
+}