@modelprofile.com/flexharness 11.1.0 → 11.2.0

package/ts/classes.codesandbox.ts

@@ -1,5 +1,5 @@
 import * as plugins from './plugins.js';
-import { FlexHarnessValidationError } from './errors.js';
+import { FlexHarnessClosedError, FlexHarnessValidationError } from './errors.js';
 import type { IFlexCodeModeLimits, IFlexCodeRunFailure, TJsonValue } from './interfaces.js';
 import type { TCodeWorkerPhase, TCodeWorkerReport, TCodeWorkerRequest } from './classes.codesandbox.worker.js';
 import { maxTransferMetadataBytes, truncateUtf8 } from './utils.transfer.js';
@@ -27,11 +27,44 @@ export type TFlexCodeRunResult =
   | { ok: true; value: TJsonValue; calls: number; logs?: string }
   | { ok: false; failure: IFlexCodeRunFailure };
 
+/**
+ * When a sandbox replaces its workers. A worker runs programs one after another, each in a QuickJS
+ * runtime of its own that it frees after the program, and is ended after a limit stopped a program,
+ * when it fails, after `maxPrograms`, once its WebAssembly memory grew past `maxMemoryBytes`, or after
+ * waiting `idleMs` for a program.
+ */
+export interface IFlexCodeSandboxPolicy {
+  /** Programs one worker runs before it is replaced. */
+  maxPrograms: number;
+  /** The WebAssembly memory a worker may have grown to and still run another program; it never shrinks. */
+  maxMemoryBytes: number;
+  /** How long a worker waits for its next program before it ends. */
+  idleMs: number;
+  /** Workers kept waiting for programs, across scopes; the one that waited longest goes first. */
+  maxWaiting: number;
+}
+
+export const FLEX_CODE_SANDBOX_DEFAULT_POLICY: Readonly<IFlexCodeSandboxPolicy> = Object.freeze({
+  maxPrograms: 100,
+  maxMemoryBytes: 64 * 1024 * 1024,
+  idleMs: 60_000,
+  maxWaiting: 4,
+});
+
 type TStop =
   | { kind: 'limit'; limit: NonNullable<IFlexCodeRunFailure['limit']>; message: string }
   | { kind: 'aborted'; message: string }
   | { kind: 'host'; error: unknown };
 
+/** What a worker hands the program it serves: everything but its own bookkeeping. */
+type TCodeRunReport = Exclude<TCodeWorkerReport, { kind: 'released' }>;
+
+interface ICodeWorkerClient {
+  report(report: TCodeRunReport): void;
+  /** The worker failed or exited by itself; it is ended. */
+  fail(error: unknown): void;
+}
+
 /** How long calls in flight get to settle after their program stopped and their signal aborted. */
 const stopGraceMs = 5_000;
 /**
@@ -40,6 +73,8 @@ const stopGraceMs = 5_000;
  * array many times over, can run far longer.
  */
 const cpuGraceMs = 250;
+/** How long a worker gets to free the runtime of a program that ended by itself before it is ended. */
+const releaseGraceMs = 1_000;
 
 const programPrefix = 'globalThis.__flexMain = async function () {\n';
 const programSuffix = '\n};';
@@ -50,16 +85,32 @@ const workerUrl = new URL(`./classes.codesandbox.worker${plugins.path.extname(im
 let enginePromise: Promise<plugins.ICodeSandboxEngine> | undefined;
 
 /**
- * Runs `run_code` programs in QuickJS-ng compiled to WebAssembly: every program in a worker thread and
- * WebAssembly instance of its own, with no globals beyond the language, `api`, `ApiError` and
- * `console`, and every limit of `IFlexCodeModeLimits` enforced. Program types are stripped with amaro
- * first.
+ * Runs `run_code` programs in QuickJS-ng compiled to WebAssembly, with no globals beyond the language,
+ * `api`, `ApiError` and `console`, and every limit of `IFlexCodeModeLimits` enforced. Program types are
+ * stripped with amaro first.
+ *
+ * The sandbox owns its worker threads. Each runs one program at a time, in a QuickJS runtime and
+ * context of its own that it frees once the program ended, so nothing of a program reaches the next,
+ * and serves only programs of one scope: programs of different scopes never share a WebAssembly
+ * instance. `IFlexCodeSandboxPolicy` says when a worker is replaced; `dispose()` stops the programs
+ * running and ends every worker.
  */
 export class FlexCodeSandbox {
-  private constructor(private readonly engine: plugins.ICodeSandboxEngine) {}
+  private readonly workers = new Set<CodeWorker>();
+  /** Workers waiting for a program, the one that waited longest first. */
+  private readonly waiting: CodeWorker[] = [];
+  private readonly runs = new Map<CodeRun, Promise<unknown>>();
+  private started = 0;
+  private closed = false;
+  private disposal?: Promise<void>;
+
+  private constructor(
+    private readonly engine: plugins.ICodeSandboxEngine,
+    private readonly policy: Readonly<IFlexCodeSandboxPolicy>,
+  ) {}
 
-  /** Loads QuickJS and amaro, once per process. */
-  public static async load(): Promise<FlexCodeSandbox> {
+  /** Loads QuickJS and amaro, once per process, into a sandbox with workers of its own. */
+  public static async load(policy: Partial<IFlexCodeSandboxPolicy> = {}): Promise<FlexCodeSandbox> {
     enginePromise ??= plugins.importCodeSandboxEngine().catch((error: unknown) => {
       enginePromise = undefined;
       throw new FlexHarnessValidationError(
@@ -68,22 +119,224 @@ export class FlexCodeSandbox {
         { cause: error },
       );
     });
-    return new FlexCodeSandbox(await enginePromise);
+    return new FlexCodeSandbox(await enginePromise, { ...FLEX_CODE_SANDBOX_DEFAULT_POLICY, ...policy });
   }
 
-  public run(
+  /** Workers started so far, alive now, and waiting for a program. */
+  public get workerCounts(): { started: number; live: number; waiting: number } {
+    return { started: this.started, live: this.workers.size, waiting: this.waiting.length };
+  }
+
+  /** Starts a worker for the scope's next program, unless one waits for it already. */
+  public prewarm(scope?: string): void {
+    if (this.closed || this.policy.maxWaiting === 0) return;
+    if (this.waiting.some((worker) => worker.programs === 0 || worker.scope === scope)) return;
+    this.park(this.spawn());
+  }
+
+  /** Runs a program on a worker that waits for the scope, or on a new one. */
+  public async run(
     source: string,
     host: IFlexCodeHost,
     limits: Readonly<IFlexCodeModeLimits>,
     signal?: AbortSignal,
+    scope?: string,
   ): Promise<TFlexCodeRunResult> {
-    return new CodeRun(this.engine, source, host, limits, signal).start();
+    if (this.closed) throw new FlexHarnessClosedError();
+    const run = new CodeRun(this.engine, source, host, limits, signal, () => this.lease(scope));
+    const outcome = run.start();
+    this.runs.set(run, outcome);
+    try {
+      return await outcome;
+    } finally {
+      this.runs.delete(run);
+    }
+  }
+
+  /** Stops the programs running, and ends every worker once they settled. */
+  public dispose(): Promise<void> {
+    if (!this.disposal) {
+      this.closed = true;
+      this.disposal = this.endAll();
+    }
+    return this.disposal;
+  }
+
+  private async endAll(): Promise<void> {
+    for (const run of this.runs.keys()) run.abort('The code sandbox was disposed.');
+    await Promise.allSettled(this.runs.values());
+    await Promise.all([...this.workers].map((worker) => worker.end(false)));
+  }
+
+  /** A waiting worker of the scope, the one that waited least; else one that has served no scope yet. */
+  private lease(scope: string | undefined): CodeWorker {
+    let index = this.lastWaiting((worker) => worker.programs > 0 && worker.scope === scope);
+    if (index === -1) index = this.lastWaiting((worker) => worker.programs === 0);
+    const worker = index === -1 ? this.spawn() : this.waiting.splice(index, 1)[0]!;
+    worker.claim(scope);
+    return worker;
+  }
+
+  private lastWaiting(matches: (worker: CodeWorker) => boolean): number {
+    for (let index = this.waiting.length - 1; index >= 0; index -= 1) {
+      if (matches(this.waiting[index]!)) return index;
+    }
+    return -1;
+  }
+
+  private spawn(): CodeWorker {
+    const worker = new CodeWorker({
+      returned: (returned) => this.returned(returned),
+      ended: (ended, replace) => this.ended(ended, replace),
+    });
+    this.workers.add(worker);
+    this.started += 1;
+    return worker;
+  }
+
+  /** A worker freed its program's runtime: it waits for the next program, or is replaced. */
+  private returned(worker: CodeWorker): void {
+    if (this.closed) {
+      void worker.end(false);
+      return;
+    }
+    if (worker.programs >= this.policy.maxPrograms || worker.memoryBytes > this.policy.maxMemoryBytes) {
+      void worker.end(true);
+      return;
+    }
+    this.park(worker);
+  }
+
+  private park(worker: CodeWorker): void {
+    worker.wait(this.policy.idleMs);
+    this.waiting.push(worker);
+    while (this.waiting.length > this.policy.maxWaiting) void this.waiting.shift()!.end(false);
+  }
+
+  /** A worker that served programs and was ended before its time is replaced in the background. */
+  private ended(worker: CodeWorker, replace: boolean): void {
+    this.workers.delete(worker);
+    const index = this.waiting.indexOf(worker);
+    if (index !== -1) this.waiting.splice(index, 1);
+    if (replace) this.prewarm(worker.scope);
   }
 }
 
-class CodeRun {
-  private worker?: plugins.Worker;
+/**
+ * A worker thread of a sandbox, and the WebAssembly instance it loads QuickJS into. It is `waiting`
+ * for a program, `serving` one, `releasing` the runtime of a program that ended by itself, or `ended`.
+ * Once it served a program, it serves only programs of that program's scope.
+ */
+class CodeWorker {
+  public state: 'waiting' | 'serving' | 'releasing' | 'ended' = 'waiting';
+  public programs = 0;
+  public scope?: string;
+  /** Its WebAssembly memory after the last program, which never shrinks. */
+  public memoryBytes = 0;
+  private ready = false;
+  private readonly thread: plugins.Worker;
+  private client?: ICodeWorkerClient;
+  private idleTimer?: ReturnType<typeof setTimeout>;
+  private releaseTimer?: ReturnType<typeof setTimeout>;
+  /** Settles the `leave()` of the program it releases, once it is returned or ended. */
+  private handedBack?: () => void;
   private terminated?: Promise<unknown>;
+
+  constructor(private readonly owner: {
+    returned(worker: CodeWorker): void;
+    ended(worker: CodeWorker, replace: boolean): void;
+  }) {
+    this.thread = new plugins.Worker(workerUrl);
+    this.thread.on('message', (report: TCodeWorkerReport) => this.receive(report));
+    this.thread.on('error', (error: unknown) => this.fail(error));
+    this.thread.on('exit', (code: number) => this.fail(new Error(`The code mode worker exited with code ${code}.`)));
+  }
+
+  /** Takes the worker out of waiting for a program of the scope. */
+  public claim(scope: string | undefined): void {
+    clearTimeout(this.idleTimer);
+    this.thread.ref();
+    this.state = 'serving';
+    this.programs += 1;
+    this.scope = scope;
+  }
+
+  /** Hands the worker's reports to the program it serves, starting with `ready` once QuickJS is loaded. */
+  public serve(client: ICodeWorkerClient): void {
+    this.client = client;
+    if (this.ready) client.report({ kind: 'ready' });
+  }
+
+  public send(request: TCodeWorkerRequest): void {
+    this.thread.postMessage(request);
+  }
+
+  /**
+   * The program is over. A worker whose program ended by itself frees its runtime and is returned to
+   * the sandbox, or ended if that takes `releaseGraceMs`; any other is ended now, whatever it computes.
+   * Settles once the worker is returned or ended, so the program's next one finds it waiting.
+   */
+  public leave(reusable: boolean): Promise<unknown> {
+    this.client = undefined;
+    if (!reusable || this.state !== 'serving') return this.end(true);
+    this.state = 'releasing';
+    this.releaseTimer = setTimeout(() => void this.end(true), releaseGraceMs);
+    return new Promise<void>((resolve) => {
+      this.handedBack = resolve;
+    });
+  }
+
+  /** Waits for the next program, keeping the process alive no longer than its own work does. */
+  public wait(idleMs: number): void {
+    this.state = 'waiting';
+    this.thread.unref();
+    this.idleTimer = setTimeout(() => void this.end(false), idleMs);
+    this.idleTimer.unref();
+  }
+
+  /** Terminates the thread, which frees its WebAssembly instance; `replace` asks for a successor. */
+  public end(replace: boolean): Promise<unknown> {
+    if (!this.terminated) {
+      this.state = 'ended';
+      this.client = undefined;
+      clearTimeout(this.idleTimer);
+      clearTimeout(this.releaseTimer);
+      this.terminated = this.thread.terminate();
+      const handedBack = this.handedBack;
+      this.handedBack = undefined;
+      if (handedBack) void this.terminated.then(handedBack);
+      this.owner.ended(this, replace && this.programs > 0);
+    }
+    return this.terminated;
+  }
+
+  private receive(report: TCodeWorkerReport): void {
+    if (report.kind === 'released') {
+      if (this.state !== 'releasing') return;
+      clearTimeout(this.releaseTimer);
+      this.memoryBytes = report.memoryBytes;
+      this.owner.returned(this);
+      this.handedBack?.();
+      this.handedBack = undefined;
+      return;
+    }
+    if (report.kind === 'ready') this.ready = true;
+    this.client?.report(report);
+  }
+
+  private fail(error: unknown): void {
+    if (this.terminated) return;
+    const client = this.client;
+    void this.end(true);
+    client?.fail(error);
+  }
+}
+
+class CodeRun {
+  private worker?: CodeWorker;
+  private left?: Promise<unknown>;
+  /** The program ended by itself, so its worker can run the next one. */
+  private reusable = false;
   private program = '';
   private readonly calls = new AbortController();
   private readonly queue: Array<() => void> = [];
@@ -111,6 +364,7 @@ class CodeRun {
     private readonly host: IFlexCodeHost,
     private readonly limits: Readonly<IFlexCodeModeLimits>,
     private readonly signal: AbortSignal | undefined,
+    private readonly lease: () => CodeWorker,
   ) {}
 
   public async start(): Promise<TFlexCodeRunResult> {
@@ -139,7 +393,7 @@ class CodeRun {
     this.signal?.addEventListener('abort', onAbort, { once: true });
     try {
       if (this.signal?.aborted) onAbort();
-      else this.spawn();
+      else this.attach();
       return await outcome;
     } finally {
       this.signal?.removeEventListener('abort', onAbort);
@@ -147,17 +401,20 @@ class CodeRun {
     }
   }
 
-  private spawn(): void {
-    const worker = new plugins.Worker(workerUrl);
-    this.worker = worker;
-    worker.on('message', (report: TCodeWorkerReport) => this.receive(report));
-    worker.on('error', (error: unknown) => this.halt({ kind: 'host', error }));
-    worker.on('exit', (code: number) => {
-      this.halt({ kind: 'host', error: new Error(`The code mode worker exited with code ${code}.`) });
+  /** Stops the program from outside, as an abort. */
+  public abort(message: string): void {
+    this.halt({ kind: 'aborted', message });
+  }
+
+  private attach(): void {
+    this.worker = this.lease();
+    this.worker.serve({
+      report: (report) => this.receive(report),
+      fail: (error) => this.halt({ kind: 'host', error }),
     });
   }
 
-  private receive(report: TCodeWorkerReport): void {
+  private receive(report: TCodeRunReport): void {
     if (this.closing) return;
     switch (report.kind) {
       case 'ready':
@@ -190,7 +447,7 @@ class CodeRun {
           this.halt({ kind: 'host', error: new Error(`The code mode prelude failed: ${report.dumped}`) });
           return;
         }
-        this.done(this.failure(this.classify(report, report.phase)));
+        this.finish(this.failure(this.classify(report, report.phase)));
         return;
       case 'stopped':
         this.halt(report.limit === 'cpu'
@@ -203,7 +460,7 @@ class CodeRun {
   /** Has the worker compute: start the program, or go on with an answer. */
   private send(request: TCodeWorkerRequest): void {
     this.busy += 1;
-    this.worker!.postMessage(request);
+    this.worker!.send(request);
     this.armWatchdog();
   }
 
@@ -295,13 +552,13 @@ class CodeRun {
     if (outcome.ok) {
       if (unawaited > 0) {
         // Calls the program did not await are not part of its result: it fails, and proposes nothing.
-        this.done(this.failure({
+        this.finish(this.failure({
           error: 'program_error',
           message: `The program returned with ${unawaited} call${unawaited === 1 ? '' : 's'} it did not await; await every call.`,
         }));
         return;
       }
-      this.done({
+      this.finish({
         ok: true,
         value: outcome.value ?? null,
         calls: Math.min(this.callCount, this.limits.maxCalls),
@@ -309,7 +566,7 @@ class CodeRun {
       });
       return;
     }
-    this.done(this.failure(this.classify(outcome, 'program')));
+    this.finish(this.failure(this.classify(outcome, 'program')));
   }
 
   private classify(
@@ -363,12 +620,21 @@ class CodeRun {
       : { error: 'aborted', message: reason.message }));
   }
 
+  /**
+   * The program ended by itself with this outcome. Its worker runs the next program, unless the
+   * program reached a limit: such a worker is not trusted with another.
+   */
+  private finish(result: TFlexCodeRunResult): void {
+    this.reusable = result.ok || result.failure.error !== 'limit_exceeded';
+    this.done(result);
+  }
+
   private done(result: TFlexCodeRunResult): void {
     this.close(() => this.settle(result));
   }
 
   /**
-   * Decides the run's outcome and ends the program's worker. Calls not started yet never start. Calls
+   * Decides the run's outcome and leaves the program's worker. Calls not started yet never start. Calls
    * in flight are told to stop through their signal, and the outcome is delivered once they settled,
    * so no call runs on after its program; a host call that ignores its signal gets `stopGraceMs`, and
    * its answer is dropped.
@@ -377,7 +643,7 @@ class CodeRun {
     if (this.closing) return;
     this.closing = true;
     this.queue.length = 0;
-    void this.terminate();
+    void this.leave();
     if (this.inFlight.size === 0) {
       deliver();
       return;
@@ -408,11 +674,14 @@ class CodeRun {
     this.wallTimer = setTimeout(check, Math.max(10, Math.min(250, this.limits.wallMs)));
   }
 
-  /** Ends the program's worker, whatever it computes, which frees its WebAssembly instance with it. */
-  private terminate(): Promise<unknown> {
+  /**
+   * Hands the worker back: for the next program if this one ended by itself, ended otherwise, whatever
+   * it computes, which frees its WebAssembly instance with it.
+   */
+  private leave(): Promise<unknown> {
     clearTimeout(this.watchdog);
-    this.terminated ??= this.worker?.terminate() ?? Promise.resolve();
-    return this.terminated;
+    this.left ??= this.worker?.leave(this.reusable) ?? Promise.resolve();
+    return this.left;
   }
 
   private async dispose(): Promise<void> {
@@ -420,6 +689,6 @@ class CodeRun {
     if (this.wallTimer) clearTimeout(this.wallTimer);
     this.queue.length = 0;
     this.calls.abort();
-    await this.terminate();
+    await this.leave();
   }
 }

package/ts/classes.codesandbox.worker.ts

@@ -1,8 +1,9 @@
 /**
- * The QuickJS side of a `run_code` program. `FlexCodeSandbox` runs every program in a worker thread of
- * its own with this module, and ends it by terminating the worker, which frees its WebAssembly instance
- * too: a native operation that never reaches QuickJS's interrupt handler cannot outlast its limits, and
- * a program never blocks the host's event loop.
+ * The QuickJS side of `run_code` programs. A worker thread of `FlexCodeSandbox` loads QuickJS once and
+ * runs one program at a time, each in a QuickJS runtime and context of its own, which it frees once the
+ * program ended by itself. A program stopped at a limit ends with its worker: terminating the worker
+ * frees its WebAssembly instance too, so a native operation that never reaches QuickJS's interrupt
+ * handler cannot outlast its limits, and a program never blocks the host's event loop.
  *
  * Written in erasable TypeScript that imports nothing of the package at run time, so the module loads
  * uncompiled as well.
@@ -29,10 +30,13 @@ export type TCodeWorkerRequest =
 /**
  * What a program's worker reports. `ready` comes once QuickJS is loaded; every `start` and `answer`
  * ends in `idle`, with the milliseconds computed so far, unless the program stopped. `done` counts the
- * calls the program had no answer to when it returned.
+ * calls the program had no answer to when it returned. `released` follows a program that ended by
+ * itself (`done`, or `failed` short of a limit): its runtime is freed, and the worker's WebAssembly
+ * memory, which never shrinks, has `memoryBytes`.
  */
 export type TCodeWorkerReport =
   | { kind: 'ready' }
+  | { kind: 'released'; memoryBytes: number }
   | { kind: 'call'; id: number; method: string; params: string }
   | { kind: 'log'; text: string }
   | { kind: 'idle'; cpuMs: number }
@@ -134,7 +138,8 @@ const prelude = `'use strict';
 /**
  * One program in its QuickJS runtime. It enforces the limits QuickJS can see: memory, stack, console
  * output, and computing time at every interrupt point. The host enforces the rest, and ends the worker
- * when the program computes past its time between two interrupt points.
+ * when the program computes past its time between two interrupt points. Once the program's outcome is
+ * reported, nothing more of it runs: work it left pending is interrupted and dropped with its runtime.
  */
 class CodeWorkerRun {
   private readonly report: (report: TCodeWorkerReport) => void;
@@ -177,6 +182,19 @@ class CodeWorkerRun {
     this.idle();
   }
 
+  /** The program ended by itself, so its runtime can be freed and the worker run the next program. */
+  public get finished(): boolean {
+    return this.ended && this.stopped === undefined;
+  }
+
+  /** Frees the program's runtime with everything in it, and the calls it left unanswered. */
+  public release(): void {
+    for (const deferred of this.deferreds.values()) deferred.dispose();
+    this.deferreds.clear();
+    this.vm.dispose();
+    this.runtime.dispose();
+  }
+
   /** Settles a call of the program with the host's outcome, and runs what it unblocked. */
   public answer(id: number, text: string): void {
     const deferred = this.deferreds.get(id);
@@ -233,7 +251,7 @@ class CodeWorkerRun {
       return false;
     }
     result.value.dispose();
-    return true;
+    return !this.ended;
   }
 
   /** Runs the program's pending jobs after a call settled or the program started. */
@@ -259,6 +277,11 @@ class CodeWorkerRun {
   }
 
   private failWith(handle: plugins.IQuickJsHandle, phase: TCodeWorkerPhase): void {
+    if (this.ended) {
+      // The interruption of work the program left pending once its outcome was out.
+      handle.dispose();
+      return;
+    }
     let dumped: unknown;
     // Reading the thrown value can run program code, such as a toJSON: it counts as computing.
     this.enter();
@@ -268,6 +291,7 @@ class CodeWorkerRun {
       dumped = undefined;
     } finally {
       this.leave();
+      handle.dispose();
     }
     if (this.stopped) {
       this.stop(this.stopped);
@@ -313,7 +337,7 @@ class CodeWorkerRun {
   }
 
   private interrupted(): boolean {
-    if (this.stopped) return true;
+    if (this.stopped || this.ended) return true;
     if (this.enteredAt !== undefined && this.cpuMs + performance.now() - this.enteredAt > this.limits.cpuMs) {
       this.stopped = 'cpu';
       return true;
@@ -340,6 +364,12 @@ if (port) {
     } else {
       run?.answer(request.id, request.text);
     }
+    // A failure to free the runtime throws: the worker fails, and the host ends it.
+    if (run?.finished) {
+      run.release();
+      run = undefined;
+      report({ kind: 'released', memoryBytes: module.getWasmMemory().buffer.byteLength });
+    }
   });
   report({ kind: 'ready' });
 }

package/ts/classes.flexharness.callbacks.ts

@@ -39,7 +39,7 @@ export class FlexHarnessRunCallbacks<TScope> {
     const run = this.callbackRun(state, sessionId);
     if (!run || !delta) return;
     const lastPart = run.callbackParts.at(-1);
-    const openPart = lastPart?.type === 'text' && !run.completedTextPartIds.has(lastPart.partId)
+    const openPart = lastPart?.type === 'text' && run.stepPartIds.has(lastPart.partId)
       ? lastPart
       : undefined;
     const created = !openPart;
@@ -54,6 +54,7 @@ export class FlexHarnessRunCallbacks<TScope> {
     )) return;
     if (created) {
       run.callbackParts.push(part);
+      run.stepPartIds.add(part.partId);
       run.partTextUtf8Bytes.set(part.partId, 0);
       this.core.emitPartEvent(run, 'part.started', part);
     }
@@ -75,6 +76,7 @@ export class FlexHarnessRunCallbacks<TScope> {
     };
     run.reasoningPartIds.set(id, part.partId);
     run.callbackParts.push(part);
+    run.stepPartIds.add(part.partId);
     run.partTextUtf8Bytes.set(part.partId, 0);
     this.core.emitPartEvent(run, 'part.started', part);
   }
@@ -115,6 +117,9 @@ export class FlexHarnessRunCallbacks<TScope> {
     }
     part.status = 'completed';
     part.completedAt = new Date().toISOString();
+    // Model calls number their reasoning summaries afresh, so a later call's summary with the
+    // same ID is another one.
+    run.reasoningPartIds.delete(id);
     this.core.emitPartEvent(run, 'part.completed', part);
   }
 
@@ -128,10 +133,10 @@ export class FlexHarnessRunCallbacks<TScope> {
     const run = this.callbackRun(state, sessionId);
     if (!run) return;
     for (const part of run.callbackParts) {
-      if (part.type !== 'text' || run.completedTextPartIds.has(part.partId)) continue;
-      run.completedTextPartIds.add(part.partId);
+      if (part.type !== 'text' || !run.stepPartIds.has(part.partId)) continue;
       this.core.emitPartEvent(run, 'part.completed', part);
     }
+    run.stepPartIds.clear();
   }
 
   private textAppendAccounting(
@@ -372,9 +377,23 @@ export class FlexHarnessRunCallbacks<TScope> {
     }
   }
 
+  /**
+   * Announces a retry. The failed model call's text and reasoning are not in the session, so their
+   * parts are removed first, last part first, each with a `part.removed` event; its tool calls
+   * stay. The removals are outside the callback limits, like completions: their number is bounded
+   * by the run's parts. What the removed parts took of the limits stays taken.
+   */
   public onRetry(state: IStorageState, sessionId: string, retry: plugins.IAgentRetryEvent): void {
     const run = this.callbackRun(state, sessionId);
     if (!run || !this.reserveCallbackCapacity(run, 1, 0, 0)) return;
+    for (let partIndex = run.callbackParts.length - 1; partIndex >= 0; partIndex--) {
+      const part = run.callbackParts[partIndex]!;
+      if (!run.stepPartIds.has(part.partId)) continue;
+      this.core.emitPartRemovedEvent(run, part);
+      run.callbackParts.splice(partIndex, 1);
+      run.partTextUtf8Bytes.delete(part.partId);
+    }
+    run.stepPartIds.clear();
     this.core.emitEvent(run.scopeId, run.stored.session, {
       type: 'run.retrying',
       runId: run.runId,