@tt-a1i/openpi 0.5.0 → 0.6.0

package/web/runtime/pi-runtime.ts

@@ -15,11 +15,15 @@ import {
   hasTrustRequiringProjectResources,
 } from "@earendil-works/pi-coding-agent";
 import {
+  type WebActiveTurn,
   type WebModelSelectionOptions,
   type WebPromptOptions,
+  type WebPromptAdmissionReceipt,
   type WebRuntimeController,
   type WebRuntimeEvent,
   type WebSessionCreationOptions,
+  type WebTurnCancellationOptions,
+  type WebTurnCancellationResult,
   WebRuntimeRequestError,
 } from "./types.ts";
 import { projectMessage } from "../protocol/types.ts";
@@ -35,6 +39,7 @@ import {
 } from "./web-host-lease.ts";
 
 const STARTUP_TIMEOUT_MS = 15_000;
+const TURN_CANCELLATION_SETTLEMENT_TIMEOUT_MS = 10_000;
 const BOOTSTRAP_WORKSPACE_DIRECTORY = ".bootstrap-workspace";
 
 type PromptTrace = {
@@ -43,6 +48,13 @@ type PromptTrace = {
   startedAt: number;
   started: boolean;
   queued: boolean;
+  userMessageObserved: boolean;
+  epoch?: number;
+  outcome?: "completed" | "cancelled" | "failed" | "uncertain";
+};
+
+type TurnSettlement = WebActiveTurn & {
+  outcome: "completed" | "cancelled" | "failed" | "uncertain";
 };
 
 function errorText(error: unknown) {
@@ -77,6 +89,14 @@ export class PiWebRuntime implements WebRuntimeController {
   private promptAdmission: Promise<void> = Promise.resolve();
   private activePromptTrace?: PromptTrace;
   private readonly pendingPromptTraces: PromptTrace[] = [];
+  private nextTurnEpoch = 0;
+  private readonly terminalTurnKeys = new Set<string>();
+  private readonly turnSettlementWaiters = new Map<
+    string,
+    Set<(settlement: TurnSettlement) => void>
+  >();
+  /** Native aborts remain owned by Pi until its agent_settled event arrives. */
+  private readonly turnAbortOperations = new Map<string, Promise<unknown>>();
   private liveMessageKey?: string;
   private liveMessageSequence = 0;
   private readonly webSessionDirectory: string;
@@ -124,6 +144,7 @@ export class PiWebRuntime implements WebRuntimeController {
     const webSessionDirectory = join(getAgentDir(), "web-sessions");
     const webHostLease = await acquireWebHostLease(webSessionDirectory);
     let runtime: PiWebRuntime | undefined;
+    let timeoutHandle: ReturnType<typeof setTimeout> | undefined;
     try {
       const created = await PiWebRuntime.createRuntime(
         canonicalCwd,
@@ -174,6 +195,109 @@ export class PiWebRuntime implements WebRuntimeController {
     return !this.runtime.session.isStreaming;
   }
 
+  getActiveTurn() {
+    return this.activeTurnFromTrace(this.activePromptTrace);
+  }
+
+  cancelTurn(options: WebTurnCancellationOptions) {
+    return this.serializeControllerMutation(() =>
+      this.cancelActiveTurn(options),
+    );
+  }
+
+  private async cancelActiveTurn(
+    options: WebTurnCancellationOptions,
+  ): Promise<WebTurnCancellationResult> {
+    this.assertActive();
+    this.assertWorkspaceSelected();
+    const activeSessionId = this.runtime.session.sessionManager.getSessionId();
+    if (options.sessionId !== activeSessionId) {
+      return { ...options, state: "stale-session" };
+    }
+    const key = this.turnKey(options);
+    if (this.terminalTurnKeys.has(key)) {
+      return { ...options, state: "already-settled" };
+    }
+    const activeTurn = this.getActiveTurn();
+    if (
+      !activeTurn ||
+      activeTurn.commandId !== options.commandId ||
+      activeTurn.epoch !== options.epoch
+    ) {
+      return { ...options, state: "stale-turn" };
+    }
+    if (this.turnAbortOperations.has(key)) {
+      return {
+        ...options,
+        state: "failed",
+        error: "Cancellation is already waiting for Pi to settle this turn",
+      };
+    }
+
+    let ownWaiter: ((settlement: TurnSettlement) => void) | undefined;
+    const settlement = new Promise<TurnSettlement>((resolveSettlement) => {
+      ownWaiter = resolveSettlement;
+      const waiters = this.turnSettlementWaiters.get(key) ?? new Set();
+      waiters.add(resolveSettlement);
+      this.turnSettlementWaiters.set(key, waiters);
+    });
+    let timeoutHandle: ReturnType<typeof setTimeout> | undefined;
+    try {
+      const abortOperation = this.runtime.session.abort();
+      this.turnAbortOperations.set(key, abortOperation);
+      void abortOperation.catch(() => {
+        if (this.turnAbortOperations.get(key) === abortOperation) {
+          this.turnAbortOperations.delete(key);
+        }
+      });
+      const abortFailure = new Promise<never>((_, reject) => {
+        void abortOperation.catch(reject);
+      });
+      const settlementTimeout = new Promise<never>((_, reject) => {
+        timeoutHandle = setTimeout(
+          () =>
+            reject(
+              new Error(
+                "Cancellation did not settle within the bounded wait window",
+              ),
+            ),
+          TURN_CANCELLATION_SETTLEMENT_TIMEOUT_MS,
+        );
+      });
+      const terminal = await Promise.race([
+        settlement,
+        abortFailure,
+        settlementTimeout,
+      ]);
+      return {
+        ...options,
+        state:
+          terminal.outcome === "cancelled"
+            ? "accepted"
+            : terminal.outcome === "completed"
+              ? "already-settled"
+              : "failed",
+        ...(terminal.outcome === "failed"
+          ? { error: "The active turn failed while cancellation was requested" }
+          : terminal.outcome === "uncertain"
+            ? {
+                error:
+                  "Pi settled without a terminal assistant outcome for this cancellation",
+              }
+            : {}),
+      };
+    } catch (error) {
+      return { ...options, state: "failed", error: errorText(error) };
+    } finally {
+      if (timeoutHandle) clearTimeout(timeoutHandle);
+      const waiters = this.turnSettlementWaiters.get(key);
+      if (waiters && ownWaiter) {
+        waiters.delete(ownWaiter);
+        if (waiters.size === 0) this.turnSettlementWaiters.delete(key);
+      }
+    }
+  }
+
   listModels() {
     const current = this.runtime.session.model;
     const available = [...this.runtime.services.modelRuntime.getAvailableSnapshot()];
@@ -287,23 +411,25 @@ export class PiWebRuntime implements WebRuntimeController {
       releaseAdmission = resolveAdmission;
     });
     const startedAt = performance.now();
-    const queued = session.isStreaming;
     const promptTrace: PromptTrace | undefined = options?.commandId
       ? {
           commandId: options.commandId,
           sessionId,
           startedAt,
           started: false,
-          queued,
+          queued: false,
+          userMessageObserved: false,
         }
       : undefined;
     this.retainRuntimeReference(agentRuntime);
-    let resolveRequest: () => void = () => undefined;
+    let resolveRequest: (receipt: WebPromptAdmissionReceipt) => void = () => undefined;
     let rejectRequest: (error: unknown) => void = () => undefined;
-    const requestAdmission = new Promise<void>((resolveRequestAdmission, reject) => {
-      resolveRequest = resolveRequestAdmission;
-      rejectRequest = reject;
-    });
+    const requestAdmission = new Promise<WebPromptAdmissionReceipt>(
+      (resolveRequestAdmission, reject) => {
+        resolveRequest = resolveRequestAdmission;
+        rejectRequest = reject;
+      },
+    );
     const operation = (async () => {
       let preflightObserved = false;
       let admitted = false;
@@ -331,14 +457,15 @@ export class PiWebRuntime implements WebRuntimeController {
             elapsedMs: elapsed(startedAt),
           });
         }
-        const pendingMessagesBefore = session.pendingMessageCount;
+        let followUpMessages = session.getFollowUpMessages().length;
         unsubscribePromptLifecycle = session.subscribe((event) => {
           if (event.type === "agent_start") agentLifecycleStarted = true;
-          if (
-            event.type === "queue_update" &&
-            event.steering.length + event.followUp.length > pendingMessagesBefore
-          ) {
-            queuedForAgent = true;
+          if (event.type === "queue_update") {
+            if (event.followUp.length > followUpMessages) {
+              queuedForAgent = true;
+              if (promptTrace) promptTrace.queued = true;
+            }
+            followUpMessages = event.followUp.length;
           }
         });
         await session.prompt(content, {
@@ -363,7 +490,9 @@ export class PiWebRuntime implements WebRuntimeController {
               );
             }
             if (accepted) {
-              resolveRequest();
+              resolveRequest({
+                pendingFollowUps: session.getFollowUpMessages().length,
+              });
             } else {
               rejectRequest(
                 new WebRuntimeRequestError(
@@ -432,6 +561,7 @@ export class PiWebRuntime implements WebRuntimeController {
           });
         }
         if (promptTrace) {
+          promptTrace.outcome = "failed";
           traceWeb("prompt_operation_failed", {
             commandId: promptTrace.commandId,
             sessionId,
@@ -451,7 +581,7 @@ export class PiWebRuntime implements WebRuntimeController {
       () => this.promptOperations.delete(operation),
       () => this.promptOperations.delete(operation),
     );
-    await requestAdmission;
+    return await requestAdmission;
   }
 
   newSession(workspacePath: string, options?: WebSessionCreationOptions) {
@@ -725,13 +855,17 @@ export class PiWebRuntime implements WebRuntimeController {
 
   private projectEvent(session: AgentSession, event: AgentSessionEvent) {
     if (session !== this.runtime.session) return;
+    if (event.type === "agent_start" && this.activePromptTrace) {
+      this.startPromptTrace(this.activePromptTrace);
+    }
     if (event.type === "message_start" && event.message.role === "user") {
       if (!this.activePromptTrace) {
         this.activePromptTrace = this.pendingPromptTraces.shift();
-      } else if (this.activePromptTrace.started && this.pendingPromptTraces.length > 0) {
-        this.activePromptTrace = this.pendingPromptTraces.shift();
       }
-      if (this.activePromptTrace) this.activePromptTrace.started = true;
+      if (this.activePromptTrace) {
+        this.startPromptTrace(this.activePromptTrace);
+        this.activePromptTrace.userMessageObserved = true;
+      }
     }
     const promptTrace = this.activePromptTrace;
     if (promptTrace) {
@@ -770,15 +904,25 @@ export class PiWebRuntime implements WebRuntimeController {
     }
     switch (event.type) {
       case "agent_start":
+        this.emit(event.type, {
+          sessionId: session.sessionManager.getSessionId(),
+          ...(this.getActiveTurn()
+            ? { activeTurn: this.getActiveTurn() }
+            : {}),
+        });
+        break;
       case "agent_settled":
-        this.emit(event.type);
-        if (
-          event.type === "agent_settled" &&
-          this.activePromptTrace?.started &&
-          this.pendingPromptTraces.length === 0
-        ) {
-          this.activePromptTrace = undefined;
+        // Pi emits this only after the whole agent run (including tool loops
+        // and admitted follow-ups) has reached a terminal state. A
+        // message_end is only one model response and must not settle a turn.
+        if (this.activePromptTrace?.started) {
+          this.settlePromptTrace(this.activePromptTrace);
         }
+        this.activePromptTrace = undefined;
+        this.pendingPromptTraces.length = 0;
+        this.emit(event.type, {
+          sessionId: session.sessionManager.getSessionId(),
+        });
         break;
       case "auto_retry_start":
         this.emit(event.type, {
@@ -796,6 +940,30 @@ export class PiWebRuntime implements WebRuntimeController {
         break;
       case "message_update":
       case "message_end":
+        if (
+          event.type === "message_end" &&
+          event.message.role === "assistant" &&
+          this.activePromptTrace
+        ) {
+          // Preserve the terminal model result for classification, but defer
+          // publication until Pi confirms the entire run is settled.
+          const outcome =
+            event.message.stopReason === "aborted"
+              ? "cancelled"
+              : event.message.stopReason === "error"
+                ? "failed"
+                : event.message.stopReason === "stop" ||
+                    event.message.stopReason === "length"
+                  ? "completed"
+                  : undefined;
+          // A later queued continuation must not erase proof that the
+          // provider result targeted by Stop was aborted. The control remains
+          // owned until agent_settled; this outcome does not claim that every
+          // queued follow-up in the same Pi execution was cancelled.
+          if (outcome && this.activePromptTrace.outcome !== "cancelled") {
+            this.activePromptTrace.outcome = outcome;
+          }
+        }
         this.emit(event.type, {
           message: projectMessage(event.message),
           ...(this.liveMessageKey ? { messageKey: this.liveMessageKey } : {}),
@@ -819,10 +987,56 @@ export class PiWebRuntime implements WebRuntimeController {
     for (const listener of this.listeners) listener({ type, detail });
   }
 
+  private activeTurnFromTrace(trace?: PromptTrace): WebActiveTurn | undefined {
+    if (!trace?.started || trace.epoch === undefined) return undefined;
+    return {
+      sessionId: trace.sessionId,
+      commandId: trace.commandId,
+      epoch: trace.epoch,
+    };
+  }
+
+  private startPromptTrace(trace: PromptTrace) {
+    if (trace.started) return;
+    trace.started = true;
+    trace.epoch = ++this.nextTurnEpoch;
+    const activeTurn = this.activeTurnFromTrace(trace);
+    if (activeTurn) this.emit("turn_started", { ...activeTurn });
+  }
+
+  private settlePromptTrace(trace: PromptTrace) {
+    const activeTurn = this.activeTurnFromTrace(trace);
+    if (!activeTurn) return;
+    const settlement: TurnSettlement = {
+      ...activeTurn,
+      outcome:
+        trace.outcome ?? "uncertain",
+    };
+    const key = this.turnKey(activeTurn);
+    if (this.terminalTurnKeys.has(key)) return;
+    this.terminalTurnKeys.add(key);
+    this.turnAbortOperations.delete(key);
+    while (this.terminalTurnKeys.size > 64) {
+      const oldest = this.terminalTurnKeys.values().next().value;
+      if (typeof oldest === "string") this.terminalTurnKeys.delete(oldest);
+    }
+    this.emit("turn_settled", { ...settlement });
+    for (const resolveSettlement of this.turnSettlementWaiters.get(key) ?? []) {
+      resolveSettlement(settlement);
+    }
+    this.turnSettlementWaiters.delete(key);
+  }
+
+  private turnKey(turn: WebActiveTurn) {
+    return `${turn.sessionId}\u0000${turn.commandId}\u0000${turn.epoch}`;
+  }
+
   private removePromptTrace(trace: PromptTrace) {
     const pendingIndex = this.pendingPromptTraces.indexOf(trace);
     if (pendingIndex !== -1) this.pendingPromptTraces.splice(pendingIndex, 1);
     if (this.activePromptTrace !== trace) return;
+    // A started trace can only be terminally projected by agent_settled.
+    if (trace.started) return;
     this.activePromptTrace = this.pendingPromptTraces.shift();
   }
 
@@ -987,5 +1201,6 @@ export class PiWebRuntime implements WebRuntimeController {
   private resetPromptTraces() {
     this.activePromptTrace = undefined;
     this.pendingPromptTraces.length = 0;
+    this.turnAbortOperations.clear();
   }
 }

package/web/runtime/types.ts

@@ -33,6 +33,30 @@ export interface WebPromptOptions {
   expectedSessionId?: string;
 }
 
+export interface WebPromptAdmissionReceipt {
+  pendingFollowUps: number;
+}
+
+export interface WebActiveTurn {
+  sessionId: string;
+  commandId: string;
+  epoch: number;
+}
+
+export interface WebTurnCancellationOptions extends WebActiveTurn {}
+
+export type WebTurnCancellationState =
+  | "accepted"
+  | "already-settled"
+  | "stale-session"
+  | "stale-turn"
+  | "failed";
+
+export interface WebTurnCancellationResult extends WebActiveTurn {
+  state: WebTurnCancellationState;
+  error?: string;
+}
+
 export interface WebModelSelectionOptions {
   expectedSessionId?: string;
 }
@@ -54,7 +78,14 @@ export interface WebRuntimeController {
   readonly sessionDirectory: string;
   readonly sessionManager: SessionManager;
   isIdle(): boolean;
-  sendPrompt(content: string, options?: WebPromptOptions): Promise<void>;
+  getActiveTurn(): WebActiveTurn | undefined;
+  sendPrompt(
+    content: string,
+    options?: WebPromptOptions,
+  ): Promise<WebPromptAdmissionReceipt>;
+  cancelTurn(
+    options: WebTurnCancellationOptions,
+  ): Promise<WebTurnCancellationResult>;
   newSession(
     workspacePath: string,
     options?: WebSessionCreationOptions,