@earendil-works/pi-durable 1.4.0 → 1.5.0

package/dist/harness/generation.js

@@ -4,13 +4,14 @@ import { isRetryableAssistantError, retryDelayMs } from "@earendil-works/pi-ai/u
 import { getCurrentTools } from "@earendil-works/pi-ai/utils/transcript";
 import { AssistantEntry, ResetEntry, SystemEntry, UserEntry } from "../entries.js";
 import { defineTask } from "../tasks.js";
-import { addTools, sectionExtension } from "./agent.js";
+import { addTools, agentHooks, sectionExtension } from "./agent.js";
 import { compactionPolicy, createCompaction, estimateContext, selectCut } from "./compaction.js";
-import { applyBoundary, prepareBoundary, withdrawQueuedInputs } from "./inbox.js";
-import { assignJson } from "./json.js";
-import { endRun, LiveDoc } from "./live.js";
+import { applyBoundary, hasQueuedReset, prepareBoundary, withdrawQueuedInputs } from "./inbox.js";
+import { cutOffCount, cutOffTooOften, endRun, LiveDoc, PROGRESS_COMMIT, salvageable, } from "./live.js";
+import { PartialChanges } from "./partial-changes.js";
 import { planSystemEntries, renderSections, replaySections } from "./prompt.js";
 import { ensureProviderSessionId } from "./provider.js";
+import { endedByCrashes } from "./scheduler.js";
 import { appendToolResult, harnessError, ToolTask } from "./tool.js";
 import { recordUsage } from "./usage.js";
 import { optionalExtension } from "./util.js";
@@ -21,8 +22,15 @@ const DEFAULT_POLL_AFTER_MS = 5000;
  */
 export const GenerationTask = defineTask({
     name: "pi.generation",
-    version: 1,
+    version: 2,
     initial: () => ({ phase: "prepare", attempt: 1 }),
+    // Version 2 leaves calls to tools not offered in a sequential round's `pending`, flagged `unoffered` in their slots.
+    // A version 1 Harness would start their tool tasks unflagged, resolving the name, so it leaves version 2 blocked
+    // instead. A version 1 checkpoint has no such calls and is kept as it is.
+    migrate: (input, checkpoint) => ({
+        input: input,
+        checkpoint: checkpoint,
+    }),
     phases: {
         /**
          * Render the system prompt and tool loadout and append the positional `pi.system` entries they need, then move to
@@ -31,6 +39,8 @@ export const GenerationTask = defineTask({
          * may land meanwhile.
          */
         prepare: async (task, runtime, context) => {
+            if (cutOffTooOften(task, runtime))
+                return endCutOff(runtime, task, context);
             const { conversationId } = runtime;
             const agent = await runtime.agent(context);
             const settings = runtime.settings;
@@ -54,6 +64,12 @@ export const GenerationTask = defineTask({
                 }
             }
             const view = await runtime.context(conversationId, context);
+            // The model owes an answer only to a context ending in a user message or tool results; system messages say
+            // nothing it answers. Otherwise, as for a continue after an answer or in an empty conversation, end the run
+            // before any request.
+            const last = view.messages.findLast((message) => message.role !== "system");
+            if (last?.role !== "user" && last?.role !== "toolResult")
+                return failNothingToAnswer(runtime, context);
             const shown = replaySections(view.messages);
             let env;
             try {
@@ -62,11 +78,11 @@ export const GenerationTask = defineTask({
             catch (error) {
                 if (context.abortSignal?.aborted)
                     throw error;
-                runtime.report(error, { type: "host", callback: "env" });
+                runtime.report(error, { kind: "host", callback: "env" });
             }
             const input = { conversationId, agent, env, shown: Object.fromEntries(shown), read: runtime };
             const desired = await renderSections(agent.sections, input, shown, (error, section) => runtime.report(error, {
-                type: "extension",
+                kind: "extension",
                 ...optionalExtension(sectionExtension(agent, section.key)),
             }), context);
             const entries = planSystemEntries(view, desired, agent.tools, runtime.now());
@@ -83,11 +99,10 @@ export const GenerationTask = defineTask({
                 return;
             }
             await runtime.commit(async (tx) => {
-                let cutoff = (await tx.scanEntries({ conversationId }, 1)).items[0]?.id;
+                // The context ends in a message the model answers, so the conversation has an entry.
+                let cutoff = (await tx.latestEntryId(conversationId));
                 for (const entry of entries)
                     cutoff = (await tx.appendEntry(SystemEntry, conversationId, entry)).id;
-                if (cutoff === undefined)
-                    throw new Error(`Conversation ${conversationId} has no entries to send`);
                 // Checked in this commit, so a compaction admitted during preparation counts.
                 if (threshold === "background" && (await tx.doc(LiveDoc, conversationId)).compactions === undefined) {
                     await createCompaction(tx, conversationId, { reason: "threshold" });
@@ -104,14 +119,22 @@ export const GenerationTask = defineTask({
             }, context);
         },
         request: async (task, runtime, context) => {
+            // Charging the attempt may itself be what keeps failing, such as converting a huge partial: the limit ends it.
+            if (cutOffTooOften(task, runtime))
+                return endCutOff(runtime, task, context);
+            // The host stopped while this request was in flight: a failed attempt, not a resend (spec §8.3).
+            if (task.restarts !== undefined)
+                return interrupted(runtime, task.state.checkpoint, context);
             const { attempt, compacted, model: ref, thinkingLevel, streamOptions, cutoff } = task.state.checkpoint;
             const conversationId = runtime.conversationId;
+            // The attempt shown is presentation, stored like progress (`PROGRESS_COMMIT`); converting a partial left by an
+            // interrupted attempt appends an entry, which makes the commit immediate.
             await runtime.commit(async (tx) => {
                 const live = await tx.doc(LiveDoc, conversationId);
                 await convertPartial(tx, live, conversationId);
                 live.generation = { attempt };
                 return undefined;
-            }, context);
+            }, context, PROGRESS_COMMIT);
             const model = runtime.models.getModel(ref.provider, ref.modelId);
             if (model === undefined)
                 return failNoModel(runtime, ref, context);
@@ -128,29 +151,44 @@ export const GenerationTask = defineTask({
                 telemetry: context.telemetry,
                 sessionId: await ensureProviderSessionId(runtime, context),
                 ...(thinkingLevel === "off" ? {} : { reasoning: thinkingLevel }),
+                onPayload: (payload) => providerPayload(runtime, model, payload, context),
             };
             const message = await streamResponse(runtime, model, messages, options, attempt, context);
-            const request = { attempt, compacted, model: ref, cutoff, messages: view.messages };
+            const request = { attempt, compacted, model: ref, thinkingLevel, cutoff, messages: view.messages };
             await classify(runtime, request, message, context);
         },
         /** After the backoff's sleep; a run before `until` (an early wake, or a record of an earlier version) sleeps again. */
         retry: async (task, runtime, context) => {
-            const { attempt, compacted, until } = task.state.checkpoint;
+            const { attempt, compacted, until, poll } = task.state.checkpoint;
+            if (cutOffTooOften(task, runtime))
+                return endCutOff(runtime, task, context);
             if (until > runtime.now())
                 return runtime.sleepUntil(until, task.state.checkpoint, context);
             await runtime.commit(async (tx) => {
-                (await tx.doc(LiveDoc, runtime.conversationId)).generation = { attempt: attempt + 1 };
-                const checkpoint = {
-                    phase: "prepare",
-                    attempt: attempt + 1,
-                    ...(compacted === undefined ? {} : { compacted }),
-                };
-                return { status: "running", checkpoint };
+                const live = await tx.doc(LiveDoc, runtime.conversationId);
+                const kept = compacted === undefined ? {} : { compacted };
+                // A poll the host died during polls its response again, as the next attempt.
+                if (poll !== undefined) {
+                    const pollAt = runtime.now();
+                    live.generation = { attempt: attempt + 1, deferred: { pollAt } };
+                    return {
+                        status: "running",
+                        checkpoint: { phase: "poll", attempt: attempt + 1, ...kept, ...poll, pollAt },
+                    };
+                }
+                live.generation = { attempt: attempt + 1 };
+                return { status: "running", checkpoint: { phase: "prepare", attempt: attempt + 1, ...kept } };
             }, context);
         },
         /** After the sleep until `pollAt`; a run before it (an early wake, or a record of an earlier version) sleeps again. */
         poll: async (task, runtime, context) => {
-            const { attempt, compacted, model: ref, cutoff, handle, pollAt } = task.state.checkpoint;
+            const { attempt, compacted, model: ref, thinkingLevel, cutoff, handle, pollAt } = task.state.checkpoint;
+            if (cutOffTooOften(task, runtime))
+                return endCutOff(runtime, task, context);
+            // The host stopped while the poll's fetch was in flight: a failed attempt, as for a request. Checked before the
+            // time, which may have moved backwards on a replacement host; a sleeping poll is never cut off.
+            if (task.restarts !== undefined)
+                return interrupted(runtime, task.state.checkpoint, context);
             if (pollAt > runtime.now())
                 return runtime.sleepUntil(pollAt, task.state.checkpoint, context);
             const model = runtime.models.getModel(ref.provider, ref.modelId);
@@ -160,10 +198,12 @@ export const GenerationTask = defineTask({
                 signal: runtime.signal,
                 telemetry: context.telemetry,
             });
-            const request = { attempt, compacted, model: ref, cutoff, pollAt };
+            const request = { attempt, compacted, model: ref, thinkingLevel, cutoff, pollAt };
             await classify(runtime, request, message, context);
         },
         tools: async (task, runtime, context) => {
+            if (cutOffTooOften(task, runtime))
+                return endCutOff(runtime, task, context);
             const { assistant, tools, pending } = task.state.checkpoint;
             const [next, ...rest] = pending;
             if (next === undefined)
@@ -171,8 +211,8 @@ export const GenerationTask = defineTask({
             // Sequential round: start the next call and wait for it.
             await runtime.commit(async (tx) => {
                 const live = await tx.doc(LiveDoc, runtime.conversationId);
-                const taskId = await createToolTask(tx, runtime, assistant, next);
                 const slot = live.tools?.find((slot) => slot.callId === next && slot.taskId === undefined);
+                const taskId = await createToolTask(tx, runtime, assistant, next, slot?.unoffered);
                 if (slot !== undefined)
                     slot.taskId = taskId;
                 const checkpoint = {
@@ -184,25 +224,50 @@ export const GenerationTask = defineTask({
                 return { status: "waiting", checkpoint, on: [taskId], policy: "allSettled" };
             }, context);
         },
+        /**
+         * Ask the `afterRun` hooks for tasks and create them, as the generation's children, in the commit that moves to
+         * `drain`; without any, end the run now.
+         */
+        finish: async (task, runtime, context) => {
+            const { run } = task.state.checkpoint;
+            if (cutOffTooOften(task, runtime))
+                return endAnsweredCutOff(runtime, task, run.answer, context);
+            const prepared = [];
+            await runtime.hooks.each("afterRun", async (hook) => {
+                const asked = await hook(run, runtime, context);
+                if (asked !== undefined)
+                    prepareRunEndTasks(runtime, asked.tasks, prepared);
+            });
+            if (prepared.length === 0)
+                return endAnsweredRun(runtime, run.answer, context);
+            await runtime.commit(async (tx) => {
+                const tasks = [];
+                for (const { task, input } of prepared) {
+                    tasks.push(await tx.createTask(task, input, { ownership: { kind: "task", taskId: runtime.taskId } }));
+                }
+                const live = await tx.doc(LiveDoc, runtime.conversationId);
+                if (live.run?.taskId === runtime.taskId && live.run.afterRun !== undefined)
+                    live.run.afterRun.tasks = tasks;
+                const checkpoint = { phase: "drain", answer: run.answer, tasks };
+                return { status: "waiting", checkpoint, on: tasks, policy: "allSettled" };
+            }, context);
+        },
+        /** The after-run tasks are terminal, whatever their outcomes: end the run. */
+        drain: async (task, runtime, context) => {
+            const { answer } = task.state.checkpoint;
+            if (cutOffTooOften(task, runtime))
+                return endAnsweredCutOff(runtime, task, answer, context);
+            return endAnsweredRun(runtime, answer, context);
+        },
     },
     abort: async (task, runtime, context) => {
         const checkpoint = task.state.checkpoint;
-        if (checkpoint.phase === "poll") {
-            const model = runtime.models.getModel(checkpoint.model.provider, checkpoint.model.modelId);
-            if (model !== undefined) {
-                try {
-                    await runtime.models.cancelDeferred(model, checkpoint.handle, {
-                        signal: runtime.signal,
-                        telemetry: context.telemetry,
-                    });
-                }
-                catch (error) {
-                    runtime.report(error, { type: "host", callback: "models" });
-                }
-            }
-        }
+        const polled = deferredOf(checkpoint);
+        if (polled !== undefined)
+            await cancelDeferred(runtime, polled, context);
         const conversationId = runtime.conversationId;
-        // Runs after the round's tool tasks are terminal; calls never started get `aborted` results (spec §8.5).
+        // Runs after the round's tool tasks are terminal; calls never started get `aborted` results (spec §8.5), unoffered
+        // ones `tool_unavailable`.
         const unstarted = checkpoint.phase === "tools"
             ? await readCalls(runtime, checkpoint.assistant, checkpoint.pending, context)
             : [];
@@ -210,9 +275,14 @@ export const GenerationTask = defineTask({
             const live = await tx.doc(LiveDoc, conversationId);
             await convertPartial(tx, live, conversationId);
             for (const call of unstarted) {
-                const result = harnessError("aborted", `Tool ${call.name} was aborted`);
+                // A call its request did not offer could not have run either way.
+                const unoffered = live.tools?.find((slot) => slot.callId === call.id)?.unoffered === true;
+                const result = unoffered
+                    ? harnessError("tool_unavailable", `Tool ${call.name} is not available`)
+                    : harnessError("aborted", `Tool ${call.name} was aborted`);
                 await appendToolResult(tx, conversationId, call, result, { timestamp: runtime.now() });
             }
+            // In `finish` and `drain` the run answered already: it ends without a failure, its inputs stay `done`.
             endRun(tx, live, runtime.taskId, { status: "unanswered", reason: "aborted" });
             return { status: "terminal", outcome: { status: "aborted" } };
         }, context);
@@ -224,9 +294,9 @@ async function readCalls(runtime, assistant, callIds, context) {
     const calls = message?.role === "assistant" ? message.content.filter((content) => content.type === "toolCall") : [];
     return callIds.flatMap((id) => calls.filter((call) => call.id === id).slice(0, 1));
 }
-/** A tool task for call `callId`, owned by the generation. */
-function createToolTask(tx, runtime, assistant, callId) {
-    return tx.createTask(ToolTask, { kind: "model", assistant, callId }, { ownership: { kind: "task", taskId: runtime.taskId } });
+/** A tool task for call `callId`, owned by the generation; `unoffered` for a call to a tool its request did not offer. */
+function createToolTask(tx, runtime, assistant, callId, unoffered) {
+    return tx.createTask(ToolTask, { kind: "model", assistant, callId, ...(unoffered === undefined ? {} : { unoffered }) }, { ownership: { kind: "task", taskId: runtime.taskId } });
 }
 /**
  * Which threshold compaction preparation starts before its request (spec §8.3): `blocking` above
@@ -268,16 +338,121 @@ async function failModelError(runtime, text, context) {
         };
     }, context);
 }
+/**
+ * The host stopped while this attempt's request or poll was in flight (`TaskRecord.restarts`): a failed attempt under
+ * `settings.retry` (spec §8.3). The committed partial becomes an aborted entry; the next attempt follows the backoff, a
+ * poll's by polling its deferred response again, or, with no retry left, the run ends `interrupted`.
+ */
+async function interrupted(runtime, checkpoint, context) {
+    const { attempt, compacted } = checkpoint;
+    const policy = runtime.settings.retry;
+    const retry = policy.enabled && attempt <= policy.maxRetries;
+    const until = retry ? runtime.now() + retryDelayMs(policy, attempt) : 0;
+    const polled = checkpoint.phase === "poll" ? checkpoint : undefined;
+    const text = polled === undefined
+        ? `The host stopped while the model was answering, on attempt ${attempt}`
+        : `The host stopped while the deferred answer was fetched, on attempt ${attempt}`;
+    // A run that ends does not leave the provider's job running.
+    if (!retry && polled !== undefined)
+        await cancelDeferred(runtime, polled, context);
+    await runtime.commit(async (tx) => {
+        // Read before the first table write; only a failure uses it.
+        const failure = retry ? undefined : await prepareFailure(tx, runtime);
+        const live = await tx.doc(LiveDoc, runtime.conversationId);
+        await convertPartial(tx, live, runtime.conversationId);
+        if (failure === undefined) {
+            live.generation = { attempt, retry: { at: until, error: text } };
+            const poll = polled === undefined
+                ? {}
+                : {
+                    poll: {
+                        model: polled.model,
+                        ...(polled.thinkingLevel === undefined ? {} : { thinkingLevel: polled.thinkingLevel }),
+                        cutoff: polled.cutoff,
+                        handle: polled.handle,
+                    },
+                };
+            const next = { phase: "retry", attempt, ...(compacted === undefined ? {} : { compacted }), until, ...poll };
+            return { status: "sleeping", checkpoint: next, until };
+        }
+        await endFailedRun(tx, runtime, live, failure, { status: "unanswered", reason: "interrupted", detail: text });
+        return {
+            status: "terminal",
+            outcome: { status: "failed", error: { message: text, detail: { reason: "interrupted" } } },
+        };
+    }, context);
+}
+/**
+ * End a run whose generation was cut off `settings.restartLimit` times in a row (spec §5.4). At the limit it still
+ * keeps a committed partial as an aborted entry and cancels a deferred response it polled; past the limit that cleanup
+ * may be what cut it off, so it ends without it.
+ */
+async function endCutOff(runtime, task, context) {
+    const checkpoint = task.state.checkpoint;
+    const text = `The generation was interrupted ${cutOffCount(task)} in its ${checkpoint.phase} phase`;
+    const salvage = salvageable(task, runtime);
+    const polled = deferredOf(checkpoint);
+    if (salvage && polled !== undefined)
+        await cancelDeferred(runtime, polled, context);
+    await runtime.commit(async (tx) => {
+        const failure = await prepareFailure(tx, runtime);
+        const live = await tx.doc(LiveDoc, runtime.conversationId);
+        if (salvage)
+            await convertPartial(tx, live, runtime.conversationId);
+        await endFailedRun(tx, runtime, live, failure, { status: "unanswered", reason: "restart_limit", detail: text });
+        return {
+            status: "terminal",
+            outcome: { status: "failed", error: { message: text, detail: { reason: "restart_limit" } } },
+        };
+    }, context);
+}
+/** The deferred response a checkpoint polls, or will poll again after a retry's backoff. */
+function deferredOf(checkpoint) {
+    if (checkpoint.phase === "poll")
+        return checkpoint;
+    return checkpoint.phase === "retry" ? checkpoint.poll : undefined;
+}
+/** Cancel a deferred response no attempt will poll again; a failure is reported, never thrown. */
+async function cancelDeferred(runtime, checkpoint, context) {
+    const model = runtime.models.getModel(checkpoint.model.provider, checkpoint.model.modelId);
+    if (model === undefined)
+        return;
+    try {
+        await runtime.models.cancelDeferred(model, checkpoint.handle, {
+            signal: runtime.signal,
+            telemetry: context.telemetry,
+        });
+    }
+    catch (error) {
+        runtime.report(error, { kind: "host", callback: "models" });
+    }
+}
 /** Settle the run's inputs `unanswered` with `no_model` and fail. */
 async function failNoModel(runtime, ref, context) {
     const message = ref === undefined ? "No model is configured" : `Model ${ref.provider}/${ref.modelId} is not available`;
     await runtime.commit(async (tx) => {
         const failure = await prepareFailure(tx, runtime);
         const live = await tx.doc(LiveDoc, runtime.conversationId);
-        await endFailedRun(tx, runtime, live, failure, { status: "unanswered", reason: "no_model" });
+        await endFailedRun(tx, runtime, live, failure, {
+            status: "unanswered",
+            reason: "no_model",
+            detail: message,
+            settlement: { status: "unanswered", reason: "no_model" },
+        });
         return { status: "terminal", outcome: { status: "failed", error: { message, detail: { reason: "no_model" } } } };
     }, context);
 }
+/** End the run without a request: its context has nothing the model would answer. */
+async function failNothingToAnswer(runtime, context) {
+    const message = "The conversation has nothing to answer";
+    await runtime.commit(async (tx) => {
+        const failure = await prepareFailure(tx, runtime);
+        const live = await tx.doc(LiveDoc, runtime.conversationId);
+        await endFailedRun(tx, runtime, live, failure, { status: "unanswered", reason: "nothing_to_answer" });
+        const error = { message, detail: { reason: "nothing_to_answer" } };
+        return { status: "terminal", outcome: { status: "failed", error } };
+    }, context);
+}
 async function prepareFailure(tx, runtime) {
     const settings = runtime.settings;
     const policy = settings.afterFailedRun;
@@ -289,16 +464,16 @@ async function prepareFailure(tx, runtime) {
  * End a failed run, settling its inputs with `settlement`, and do with the inputs queued behind it what
  * `settings.afterFailedRun` says (spec §6): keep them, start the next run with them, or withdraw them.
  */
-async function endFailedRun(tx, runtime, live, failure, settlement) {
+async function endFailedRun(tx, runtime, live, failure, end) {
     const conversationId = runtime.conversationId;
-    endRun(tx, live, runtime.taskId, settlement);
+    endRun(tx, live, runtime.taskId, end);
     if (failure.policy === "withdraw")
         await withdrawQueuedInputs(tx, conversationId, "run_failed");
     if (failure.boundary === undefined)
         return;
-    const { users } = await applyBoundary(tx, failure.boundary, "final", runtime.now());
+    const { users, first } = await applyBoundary(tx, failure.boundary, "final", runtime.now());
     if (users.length > 0)
-        await startRun(tx, conversationId, live, users);
+        await startRun(tx, conversationId, live, { inputs: users, first }, { keepFailure: true });
 }
 /**
  * Append a committed partial left by an interrupted, aborted, faulted, or orphaned attempt as an aborted assistant
@@ -311,10 +486,22 @@ export async function convertPartial(tx, live, conversationId) {
     const message = copyJson(partial);
     await appendAssistant(tx, conversationId, { ...message, stopReason: "aborted" });
 }
+/** The `beforeProviderRequest` replacement chain over one provider request body. */
+async function providerPayload(runtime, model, payload, context) {
+    let current = payload;
+    await runtime.hooks.each("beforeProviderRequest", async (hook) => {
+        const replaced = await hook({ model, payload: current }, runtime, context);
+        if (replaced !== undefined)
+            current = replaced.payload;
+    });
+    return current;
+}
 /**
  * Stream one request and return the terminal message. Partials commit as trailing writes at most every
- * `progress.partialIntervalMs` (default 100 ms) with one commit in flight; `finally` stops the throttle and awaits that
- * commit, so no stale partial lands after the outcome.
+ * `progress.partialIntervalMs` (default 100 ms) with one commit in flight, deferred (`PROGRESS_COMMIT`); `finally` stops
+ * the throttle and awaits that commit, so no stale partial lands after the outcome. A commit appends what text and
+ * thinking grew by and applies the argument changes since the last one (`PartialChanges`), so its cost follows the new
+ * text, not the strings' length.
  */
 async function streamResponse(runtime, model, messages, options, attempt, context) {
     const interval = runtime.settings.progress.partialIntervalMs;
@@ -322,6 +509,7 @@ async function streamResponse(runtime, model, messages, options, attempt, contex
     let timer;
     let inFlight;
     let stopped = false;
+    const changes = new PartialChanges();
     const flush = () => {
         timer = undefined;
         const partial = pending;
@@ -329,14 +517,22 @@ async function streamResponse(runtime, model, messages, options, attempt, contex
         if (partial === undefined || stopped)
             return;
         inFlight = (async () => {
-            // Copy synchronously: the provider keeps mutating its partial.
+            // Copy synchronously: the provider keeps mutating its partial. The changes since the last commit belong to it.
+            // `copyJson()` returns untyped JSON; a copy of a message is that message's JSON representation.
             const message = copyJson(partial, { omitUndefinedProperties: true });
-            await runtime.commit(async (tx) => {
-                const live = await tx.doc(LiveDoc, runtime.conversationId);
-                live.generation ??= { attempt };
-                assignJson(live.generation, "message", message);
-                return undefined;
-            }, context);
+            const taken = changes.take(message);
+            try {
+                await runtime.commit(async (tx) => {
+                    const live = await tx.doc(LiveDoc, runtime.conversationId);
+                    live.generation ??= { attempt };
+                    changes.commit(live.generation, message, taken);
+                    return undefined;
+                }, context, PROGRESS_COMMIT);
+            }
+            catch (error) {
+                changes.restore(taken);
+                throw error;
+            }
         })()
             .catch((error) => {
             // Rejections after an abort mark or close are expected; the committed state stays consistent.
@@ -356,6 +552,7 @@ async function streamResponse(runtime, model, messages, options, attempt, contex
             // never gets past it, so it never leaves a partial.
             if (event.type === "done" || event.type === "error" || event.partial.content.length === 0)
                 continue;
+            changes.observe(event);
             pending = event.partial;
             if (timer === undefined && inFlight === undefined)
                 timer = runtime.setTimer(flush, interval);
@@ -373,7 +570,10 @@ async function classify(runtime, request, message, context) {
     // An abort mark or close: the abort invocation or the reopened run handles the committed state.
     runtime.signal.throwIfAborted();
     const conversationId = runtime.conversationId;
-    const { attempt, compacted, model: ref, cutoff } = request;
+    const { attempt, compacted, model: ref, thinkingLevel, cutoff } = request;
+    // Record the requested level, as pi-agent-core's agent loop does.
+    if (thinkingLevel !== undefined)
+        message.thinkingLevel = thinkingLevel;
     if (message.stopReason === "deferred" && message.deferred !== undefined) {
         const handle = message.deferred;
         const pollAt = Math.max(runtime.now() + (handle.pollAfterMs ?? DEFAULT_POLL_AFTER_MS), request.pollAt === undefined ? Number.NEGATIVE_INFINITY : request.pollAt + 1);
@@ -384,6 +584,7 @@ async function classify(runtime, request, message, context) {
                 attempt,
                 ...(compacted === undefined ? {} : { compacted }),
                 model: ref,
+                ...(thinkingLevel === undefined ? {} : { thinkingLevel }),
                 cutoff,
                 handle,
                 pollAt,
@@ -452,8 +653,9 @@ async function classify(runtime, request, message, context) {
 }
 /**
  * A final answer; the final boundary places queued items (spec §6). The first `onYield` continuation appends a user
- * message and hands the run to a successor generation, but only when the boundary selected no user item and no reset.
- * Otherwise the run's inputs settle `done`, and selected user items start the next run.
+ * message and hands the run to a successor generation, but only when the inbox holds no user item and no reset.
+ * Otherwise the run's inputs settle `done`; with `afterRun` hooks selected the run waits for the tasks they ask for,
+ * in `finish` and `drain`, before the final boundary; without, selected user items start the next run now.
  */
 async function answer(runtime, message, context) {
     let continuation;
@@ -462,6 +664,7 @@ async function answer(runtime, message, context) {
             return;
         continuation = (await hook(message, runtime, context))?.continue;
     });
+    const afterRun = await asksAfterRun(runtime, context);
     const conversationId = runtime.conversationId;
     await runtime.commit(async (tx) => {
         // Queue modes are read on the Session line, when the boundary is decided.
@@ -469,24 +672,135 @@ async function answer(runtime, message, context) {
         const live = await tx.doc(LiveDoc, conversationId);
         const entry = await appendAssistant(tx, conversationId, message);
         const result = { status: "terminal", outcome: { status: "completed", result: { entryId: entry.id } } };
-        const { users, reset } = await applyBoundary(tx, boundary, "final", runtime.now());
-        if (continuation !== undefined && users.length === 0 && !reset) {
-            const user = { role: "user", content: continuation, timestamp: runtime.now() };
+        const continues = takesContinuation(continuation, boundary) ? continuation : undefined;
+        if (continues === undefined && afterRun) {
+            const held = holdRun(tx, live, runtime, { reason: "answer", answer: entry.id, last: entry.id });
+            if (held !== undefined)
+                return held;
+        }
+        const { users, first } = await applyBoundary(tx, boundary, "final", runtime.now());
+        if (continues !== undefined) {
+            const user = { role: "user", content: continues, timestamp: runtime.now() };
             await tx.appendEntry(UserEntry, conversationId, { model: [user] });
+            // The continuation starts the run's turn count over: bounding `onYield` loops is the extension's job.
+            if (live.run?.taskId === runtime.taskId)
+                live.run.turns = 0;
             handOver(live, runtime.taskId, await createGeneration(tx, conversationId));
             delete live.generation;
             return result;
         }
         endRun(tx, live, runtime.taskId, { status: "done", answer: entry.id });
         if (users.length > 0)
-            await startRun(tx, conversationId, live, users);
+            await startRun(tx, conversationId, live, { inputs: users, first });
         return result;
     }, context);
 }
+/** Whether the phase's agent selects an `afterRun` hook, so a run that answers waits for its tasks (spec §8.3). */
+async function asksAfterRun(runtime, context) {
+    const hooks = agentHooks(await runtime.agent(context), GenerationTask.definition.name);
+    return hooks.some(({ handlers }) => typeof handlers.afterRun === "function");
+}
+/**
+ * Commit A of a run that answered with `afterRun` hooks selected (spec §8.3): settle the run's inputs `done` with the
+ * answer and keep `pi.live.run` for the after-run tasks, then ask the hooks in `finish`. The final boundary waits for
+ * the run's end, so queued items stay queued.
+ */
+function holdRun(tx, live, runtime, end) {
+    const run = live.run;
+    if (run?.taskId !== runtime.taskId)
+        return undefined;
+    for (const id of run.inputs)
+        tx.settleSubmission(id, { status: "done", answer: end.answer });
+    run.inputs = [];
+    run.afterRun = {};
+    delete live.generation;
+    delete live.tools;
+    delete live.nestedTools;
+    // A run started before `first` was recorded has its answer as its first entry.
+    const summary = { ...end, first: run.first ?? end.answer };
+    return { status: "running", checkpoint: { phase: "finish", run: summary } };
+}
+/**
+ * Check the tasks one `afterRun` hook asked for, off the line, and add those that pass to `prepared`, so a broken request
+ * is skipped instead of failing the creating commit: its kind must be an installed task that is not built-in, its input
+ * JSON, and its definition's `initial()` must return a JSON checkpoint of one of its phases. The task is created from
+ * the installed definition, as every task runs under it, with that checkpoint, so `initial()` runs once. Then throws
+ * the rejections, which `hooks.each` reports against the hook's extension.
+ */
+function prepareRunEndTasks(runtime, requests, prepared) {
+    const rejected = [];
+    for (const request of requests) {
+        try {
+            const name = request.task.definition.name;
+            if (name.startsWith("pi."))
+                throw new Error(`An afterRun hook cannot start the built-in task ${name}`);
+            const installed = runtime.registry.task(name);
+            if (installed === undefined)
+                throw new Error(`Task ${name} is not installed`);
+            const input = copyJson(request.input, { omitUndefinedProperties: true });
+            // The registry erases definitions; an installed one is a task definition of its kind.
+            const definition = installed.definition;
+            const checkpoint = copyJson(definition.initial(input), { omitUndefinedProperties: true });
+            if (!Object.hasOwn(definition.phases, checkpoint.phase)) {
+                throw new Error(`Task ${name} starts at the unknown phase ${String(checkpoint.phase)}`);
+            }
+            prepared.push({ task: { definition: { ...definition, initial: () => checkpoint } }, input });
+        }
+        catch (error) {
+            rejected.push(error);
+        }
+    }
+    if (rejected.length === 1)
+        throw rejected[0];
+    if (rejected.length > 1)
+        throw new AggregateError(rejected, `${rejected.length} afterRun task requests were rejected`);
+}
+/**
+ * Commit B of a run that answered (spec §8.3): its after-run tasks are terminal, or there were none. Place what the
+ * final boundary selects, end the run, and start the next one for selected user items.
+ */
+async function endAnsweredRun(runtime, answer, context) {
+    const conversationId = runtime.conversationId;
+    await runtime.commit(async (tx) => {
+        const boundary = await prepareBoundary(tx, conversationId, runtime.settings);
+        const live = await tx.doc(LiveDoc, conversationId);
+        const { users, first } = await applyBoundary(tx, boundary, "final", runtime.now());
+        endRun(tx, live, runtime.taskId, { status: "done", answer });
+        if (users.length > 0)
+            await startRun(tx, conversationId, live, { inputs: users, first });
+        return { status: "terminal", outcome: { status: "completed", result: { entryId: answer } } };
+    }, context);
+}
+/**
+ * End a run that answered once its `finish` or `drain` phase was cut off `settings.restartLimit` times (spec §5.4):
+ * the run answered, so it ends `completed`, without a failure and without the tasks it would still have waited for,
+ * marked so the scheduler reports it `task_interrupted`. Queued items stay queued, as after a failed run.
+ */
+async function endAnsweredCutOff(runtime, task, answer, context) {
+    const text = `The generation was interrupted ${cutOffCount(task)} in its ${task.state.checkpoint.phase} phase`;
+    await runtime.commit(async (tx) => {
+        const live = await tx.doc(LiveDoc, runtime.conversationId);
+        endRun(tx, live, runtime.taskId, { status: "done", answer });
+        const completed = { status: "terminal", outcome: { status: "completed", result: { entryId: answer } } };
+        return endedByCrashes(completed, text);
+    }, context);
+}
 /**
- * Append the tool-calling answer and start its tool round in one commit (spec §8.3). A call to a tool the request did
- * not offer gets its `tool_unavailable` result here; every other call gets a tool task owned by the generation, only the
- * first one now when the round is sequential. The generation then waits for them in its `tools` phase, keeping the run.
+ * Whether an answer's `onYield` continuation is taken (spec §8.3), decided from the inbox before the final boundary
+ * places anything: only when it holds no user item, of which the boundary would select one, and no reset, which would
+ * otherwise run first. A continuation starts the turn count over.
+ */
+function takesContinuation(continuation, boundary) {
+    const items = boundary.inbox.items;
+    return (continuation !== undefined &&
+        !items.some((item) => item.mode === "steer" || item.mode === "followUp") &&
+        !hasQueuedReset(items));
+}
+/**
+ * Append the tool-calling answer and start its tool round in one commit (spec §8.3). Every call gets a tool task owned by
+ * the generation, in call order, only the first one now when the round is sequential; a call to a tool the request did
+ * not offer is flagged, and its task settles `tool_unavailable`. The generation then waits for them in its `tools`
+ * phase, keeping the run.
  */
 async function startToolRound(runtime, request, message, calls, context) {
     const conversationId = runtime.conversationId;
@@ -503,30 +817,33 @@ async function startToolRound(runtime, request, message, calls, context) {
         const tools = [];
         const pending = [];
         for (const call of calls) {
-            if (!offered.has(call.name)) {
-                const unavailable = harnessError("tool_unavailable", `Tool ${call.name} is not available`);
-                const result = await appendToolResult(tx, conversationId, call, unavailable, { timestamp: runtime.now() });
-                slots.push({ callId: call.id, name: call.name, status: "done", entry: result.id });
-                continue;
-            }
+            const unoffered = offered.has(call.name) ? {} : { unoffered: true };
             if (sequential && tools.length > 0) {
                 pending.push(call.id);
-                slots.push({ callId: call.id, name: call.name, status: "pending" });
+                slots.push({ callId: call.id, name: call.name, status: "pending", ...unoffered });
                 continue;
             }
-            const taskId = await createToolTask(tx, runtime, entry.id, call.id);
+            const taskId = await createToolTask(tx, runtime, entry.id, call.id, unoffered.unoffered);
             tools.push(taskId);
-            slots.push({ callId: call.id, name: call.name, taskId, status: "pending" });
+            slots.push({ callId: call.id, name: call.name, taskId, status: "pending", ...unoffered });
         }
         delete live.generation;
         live.tools = slots;
+        if (live.run?.taskId === runtime.taskId) {
+            // A tool round is one turn of the run (`AgentState.maxTurns`).
+            live.run.turns += 1;
+            // A run its boundary started without placing anything, as a continue after an answer, begins here.
+            live.run.first ??= entry.id;
+        }
         const checkpoint = { phase: "tools", assistant: entry.id, tools, pending };
         return { status: "waiting", checkpoint, on: tools, policy: "allSettled" };
     }, context);
 }
 /**
- * The round's tools are terminal: apply their controls and either end the run at the final boundary (`terminate`,
- * `handoff`, or a queued reset) or hand it to the next generation at the `postTools` boundary (spec §8.5).
+ * The round's tools are terminal: apply their controls, append a `handoff`'s reset, and then end the run at the final
+ * boundary (`terminate`), end it like a failure at the conversation's turn limit (`turn_limit`), end it at a queued
+ * reset, or hand it to the next generation at the `postTools` boundary, in the handoff's new context when there is one
+ * (spec §8.5).
  */
 async function finishToolRound(runtime, assistant, tools, context) {
     const conversationId = runtime.conversationId;
@@ -540,36 +857,60 @@ async function finishToolRound(runtime, assistant, tools, context) {
     const slots = (await runtime.snapshot(LiveDoc, conversationId, context))?.tools ?? [];
     const results = slots.flatMap((slot) => (slot.entry === undefined ? [] : [slot.entry]));
     await runtime.hooks.each("afterTools", (hook) => hook(assistant, results, runtime, context));
-    // Every call of the round, including those answered without a task, must ask to terminate.
-    const terminate = slots.length > 0 &&
-        slots.every((slot) => slot.taskId !== undefined && controls.get(slot.taskId)?.terminate === true);
-    const added = [...controls.values()].flatMap((control) => control?.addTools ?? []);
+    // In call order, the slots' order; a slot without a task (faulted or orphaned before it got one) has no control.
+    const ordered = slots.map((slot) => (slot.taskId === undefined ? undefined : controls.get(slot.taskId)));
+    // Every call of the round must ask to terminate.
+    const terminate = slots.length > 0 && ordered.every((control) => control?.terminate === true);
+    const added = ordered.flatMap((control) => control?.addTools ?? []);
     // The last handoff in call order wins.
-    const handoff = [...controls.values()].findLast((control) => control?.handoff !== undefined)?.handoff;
+    const handoff = ordered.findLast((control) => control?.handoff !== undefined)?.handoff;
+    // Read from the phase's agent, as every agent field; a `configure()` during the round applies at its end.
+    const maxTurns = (await runtime.agent(context)).maxTurns;
+    const afterRun = terminate && (await asksAfterRun(runtime, context));
     await runtime.commit(async (tx) => {
         const boundary = await prepareBoundary(tx, conversationId, runtime.settings);
         if (added.length > 0)
             await addTools(tx, conversationId, added);
         const live = await tx.doc(LiveDoc, conversationId);
         const now = runtime.now();
-        if (terminate || handoff !== undefined) {
-            if (handoff !== undefined) {
-                const message = { role: "user", content: handoff, timestamp: now };
-                const entry = await tx.appendEntry(ResetEntry, conversationId, { head: "self", model: [message] });
-                boundary.head = entry.id;
-            }
-            const { users } = await applyBoundary(tx, boundary, "final", now);
+        // The run's last entry so far: the newest result, or the assistant when no result has an entry.
+        let last = results.reduce((newest, id) => (id > newest ? id : newest), assistant);
+        if (handoff !== undefined) {
+            // The handoff's reset precedes every queued item, so queued writes and selected steers land in its context.
+            const message = { role: "user", content: handoff, timestamp: now };
+            const entry = await tx.appendEntry(ResetEntry, conversationId, { head: "self", model: [message] });
+            boundary.head = entry.id;
+            last = entry.id;
+        }
+        const turns = live.run?.taskId === runtime.taskId ? live.run.turns : 0;
+        const held = afterRun ? holdRun(tx, live, runtime, { reason: "terminate", answer: assistant, last }) : undefined;
+        if (held !== undefined)
+            return held;
+        if (terminate) {
+            const { users, first } = await applyBoundary(tx, boundary, "final", now);
             endRun(tx, live, runtime.taskId, { status: "done", answer: assistant });
             if (users.length > 0)
-                await startRun(tx, conversationId, live, users);
+                await startRun(tx, conversationId, live, { inputs: users, first });
+        }
+        else if (maxTurns !== undefined && turns >= maxTurns) {
+            // The turn limit ends the run like a failure: `afterFailedRun` decides the queued items, and a continue resumes
+            // it on the round's results. The table was read for the boundary already, before any write.
+            const policy = runtime.settings.afterFailedRun;
+            const failure = { policy, ...(policy === "continue" ? { boundary } : {}) };
+            await endFailedRun(tx, runtime, live, failure, {
+                status: "unanswered",
+                reason: "turn_limit",
+                detail: turns === 1 ? "1 turn" : `${turns} turns`,
+            });
         }
         else {
-            const { users, reset } = await applyBoundary(tx, boundary, "postTools", now);
+            const { users, first, reset } = await applyBoundary(tx, boundary, "postTools", now);
             if (reset) {
                 // The queued reset cut the run's context before an answer.
                 endRun(tx, live, runtime.taskId, { status: "unanswered", reason: "reset" });
-                if (users.length > 0)
-                    await startRun(tx, conversationId, live, users);
+                if (users.length > 0) {
+                    await startRun(tx, conversationId, live, { inputs: users, first }, { keepFailure: true });
+                }
             }
             else {
                 delete live.tools;
@@ -590,9 +931,17 @@ async function appendAssistant(tx, conversationId, message) {
     await recordUsage(tx, conversationId, "models", `${message.provider}/${message.model}`, message.usage);
     return tx.appendEntry(AssistantEntry, conversationId, { model: [message] });
 }
-/** Start a run for `inputs`, placed input submissions: a new generation takes `pi.live.run`. */
-export async function startRun(tx, conversationId, live, inputs) {
-    live.run = { taskId: await createGeneration(tx, conversationId, { submissionIds: inputs }), inputs };
+/**
+ * Start a run for `inputs`, placed input submissions, whose first entry is `first` (`LiveRun.first`): a new generation
+ * takes `pi.live.run`. Removes `runFailure`, unless `keepFailure`: a run started in the commit that ended the failed one
+ * keeps its failure.
+ */
+export async function startRun(tx, conversationId, live, run, options) {
+    const { inputs, first } = run;
+    const taskId = await createGeneration(tx, conversationId, { submissionIds: inputs });
+    live.run = { taskId, inputs, turns: 0, ...(first === undefined ? {} : { first }) };
+    if (options?.keepFailure !== true)
+        delete live.runFailure;
 }
 /** A generation owned by its conversation. */
 function createGeneration(tx, conversationId, telemetry = {}) {