diff --git a/packages/durable/CHANGELOG.md b/packages/durable/CHANGELOG.md index 1b3515e89..83a9c5bf5 100644 --- a/packages/durable/CHANGELOG.md +++ b/packages/durable/CHANGELOG.md @@ -152,6 +152,7 @@ - Run-end tasks: a `GenerationHooks.afterRun(run, api)` hook returns `{ tasks }`, built with the typed `runEndTask(task, input)`, after a run answered (a final answer or a `terminate` round; not when an `onYield` continuation is taken, not after a failure or an abort). The answer is committed and the run's inputs settle `done` first, so `Submission.wait()` resolves at the answer; the generation then creates the tasks as its children and waits for them, keeping `pi.live.run` (with `run.afterRun`) and the conversation busy, so idle waits include them, Stop aborts them, and a crash resumes them. Input queued meanwhile, steers included, is placed when they are terminal and starts the next run; a task admits a model-visible write or a fix-run follow-up in its own commit with `submit(tx, draft)`. A crash before the tasks exist asks the hooks again; a broken request (an uninstalled or built-in task, a non-JSON input, a throwing `initial()`) is reported and skipped. `run` names the run's entries: `first` (`pi.live.run.first`, carried into `runFailure.first` so a continue covers a failed run's history), `answer`, and `last` ([#11141](https://github.com/earendil-works/pi/issues/11141)) - `registryReader(value, context)`: a `RegistryReader` over a Chord v2 registry value whose composition is a registry snapshot (`RegistryValue`): `snapshot()` reads the current composition, `subscribe` observes its compositions synchronously inside Chord's batch, so a Harness opened in a Chord object takes its extensions from a Chord registry. +- `GenerationHooks.beforeProviderRequest({ model, payload }, api)` replaces the provider request body of a generation request, for that request only, e.g. to set a provider-specific field. The body is in the provider API's format; the hook returns `{ payload }` or `undefined`. Hooks chain in extension order and run each time pi-ai builds a body; a throw is reported and the chain continues. Compaction requests do not use it ([#10804](https://github.com/earendil-works/pi/issues/10804)) ### Changed diff --git a/packages/durable/README.md b/packages/durable/README.md index 5c2b28b68..95e9c3085 100644 --- a/packages/durable/README.md +++ b/packages/durable/README.md @@ -659,7 +659,7 @@ const Guard = defineExtension({ }); ``` -- **Generation:** `beforeRequest` (replace the messages of one request), `afterResponse`, `onYield` (continue the run with another user message), `afterTools` (runs once a round's tools are done), and `afterRun` (tasks to run before the run ends, below). +- **Generation:** `beforeRequest` (replace the messages of one request), `beforeProviderRequest` (replace the provider request body of one request, e.g. to set a provider-specific field), `afterResponse`, `onYield` (continue the run with another user message), `afterTools` (runs once a round's tools are done), and `afterRun` (tasks to run before the run ends, below). - **Tools:** `beforeTool` (block or rewrite arguments) and `afterTool` (replace the result). `afterTool` sees every result a tool task writes, also the error results the Harness writes itself (unavailable, invalid, blocked, interrupted, restart limit, aborted), told apart by their diagnostic code; past `settings.restartLimit` it no longer runs. On an abort it runs inside the abort handler, so a slow hook delays `abort()`. To limit a hook to some conversations, select its extension only there, for example with `configure({ extensions: { add: [Guard] } })`. diff --git a/packages/durable/docs/spec.md b/packages/durable/docs/spec.md index dbbd2e492..fa9d90e23 100644 --- a/packages/durable/docs/spec.md +++ b/packages/durable/docs/spec.md @@ -3339,6 +3339,7 @@ consuming commit may rerun them. Abort errors always propagate. | hook | composition | ordinary throw | |---|---|---| | `beforeRequest` | replacement chain | report, continue | +| `beforeProviderRequest` | replacement chain | report, continue | | `afterResponse` | all observers | report, continue | | `onYield` | first continuation wins | report, continue | | `beforeTool` | argument replacement chain; first block wins | block tool with error text | @@ -3380,6 +3381,12 @@ interface GenerationHooks { api: HookApi, context: Context, ): HookResult<{ readonly messages: readonly Message[] }>; + /** After pi-ai builds the provider request body, before it is sent; replaces the body for that request only. */ + beforeProviderRequest( + request: { readonly model: Model; readonly payload: unknown }, + api: HookApi, + context: Context, + ): HookResult<{ readonly payload: unknown }>; /** Every terminal provider message, before classification. */ afterResponse(message: AssistantMessage, api: HookApi, context: Context): void | Promise; /** A final answer; `continue` appends a user message and continues the run. */ @@ -4900,7 +4907,9 @@ The run's inputs live in `pi.live.run`, not in the task input. (section 4.1), so Storage gets them at the next commit that reads `pi.live` or within `settings.progress.persistWithinMs`. Before streaming, the `beforeRequest` chain may - replace the messages for this request only. After a close, recovery resends + replace the messages for this request only. Once pi-ai has built the provider + request body, the `beforeProviderRequest` chain may replace it, again for this + request only; it runs each time pi-ai builds a body. After a close, recovery resends the same committed messages with the same pinned model, thinking level, and stream options, and reruns `beforeRequest`; after a crash the attempt is charged instead, as above. diff --git a/packages/durable/src/harness/generation.ts b/packages/durable/src/harness/generation.ts index 6fc2311b3..17c28db47 100644 --- a/packages/durable/src/harness/generation.ts +++ b/packages/durable/src/harness/generation.ts @@ -292,6 +292,7 @@ export const GenerationTask = defineTask providerPayload(runtime, model, payload, context), }; const message = await streamResponse(runtime, model, messages, options, attempt, context); const request = { attempt, compacted, model: ref, thinkingLevel, cutoff, messages: view.messages }; @@ -664,6 +665,21 @@ export async function convertPartial(tx: Tx, live: Draft, conversatio await appendAssistant(tx, conversationId, { ...message, stopReason: "aborted" }); } +/** The `beforeProviderRequest` replacement chain over one provider request body. */ +async function providerPayload( + runtime: Runtime, + model: Model, + payload: unknown, + context: Context, +): Promise { + 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, deferred (`PROGRESS_COMMIT`); `finally` stops diff --git a/packages/durable/src/harness/types.ts b/packages/durable/src/harness/types.ts index d7f4a05ec..d114bf6cc 100644 --- a/packages/durable/src/harness/types.ts +++ b/packages/durable/src/harness/types.ts @@ -1,9 +1,11 @@ import type { Context, JsonValue, WatchHandle } from "@earendil-works/chord"; import type { + Api, AssistantMessage, CacheRetention, ClientIdentity, Message, + Model, Models, ModelThinkingLevel, Static, @@ -970,6 +972,16 @@ export interface GenerationHooks { api: HookApi, context: Context, ): HookResult<{ readonly messages: readonly Message[] }>; + /** + * After pi-ai builds the provider request body, before it is sent; the result replaces the body for that request only. + * Runs again whenever pi-ai rebuilds the body, such as a retry within one attempt. The body is in the provider API's + * own format (`request.model.api`), so a hook checks the API before changing it. + */ + beforeProviderRequest( + request: { readonly model: Model; readonly payload: unknown }, + api: HookApi, + context: Context, + ): HookResult<{ readonly payload: unknown }>; /** Every terminal provider message, before classification. */ afterResponse(message: AssistantMessage, api: HookApi, context: Context): void | Promise; /** A final answer; the first `continue` appends a user message and continues the run. */ diff --git a/packages/durable/test/harness-tools.test.ts b/packages/durable/test/harness-tools.test.ts index 6849861b8..d2937fd77 100644 --- a/packages/durable/test/harness-tools.test.ts +++ b/packages/durable/test/harness-tools.test.ts @@ -752,6 +752,75 @@ describe("generation hooks", () => { await harness.close(context); }); + it("chains beforeProviderRequest over every provider request body", async () => { + const setup = chatSetup(); + const seen: { hook: string; model: string; payload: unknown }[] = []; + addHooks(setup.registry, GenerationTask, { + beforeProviderRequest: ({ model, payload }) => { + seen.push({ hook: "first", model: `${model.provider}/${model.id}`, payload }); + return { payload: { ...(payload as object), first: true } }; + }, + }); + addHooks(setup.registry, GenerationTask, { + beforeProviderRequest: ({ model, payload }) => { + seen.push({ hook: "keep", model: `${model.provider}/${model.id}`, payload }); + return undefined; + }, + }); + addHooks(setup.registry, GenerationTask, { + beforeProviderRequest: ({ payload }) => ({ payload: { ...(payload as object), last: true } }), + }); + const sent: unknown[] = []; + // pi-ai builds a body per provider call, and again when it resends within one attempt. + const respond: FauxResponseStep = async (_context, options, _state, model) => { + sent.push(await options!.onPayload!({ body: 1 }, model)); + sent.push(await options!.onPayload!({ body: 2 }, model)); + return DONE; + }; + const { harness, status } = await run(setup, [respond]); + expect(status).toBe("done"); + expect(sent).toEqual([ + { body: 1, first: true, last: true }, + { body: 2, first: true, last: true }, + ]); + expect(seen).toEqual([ + { hook: "first", model: "faux/faux-1", payload: { body: 1 } }, + { hook: "keep", model: "faux/faux-1", payload: { body: 1, first: true } }, + { hook: "first", model: "faux/faux-1", payload: { body: 2 } }, + { hook: "keep", model: "faux/faux-1", payload: { body: 2, first: true } }, + ]); + await harness.close(context); + }); + + it("reports a throwing beforeProviderRequest hook and continues the chain", async () => { + const setup = chatSetup(); + addHooks( + setup.registry, + GenerationTask, + { + beforeProviderRequest: () => { + throw new Error("broken hook"); + }, + }, + "broken", + ); + addHooks(setup.registry, GenerationTask, { + beforeProviderRequest: ({ payload }) => ({ payload: { ...(payload as object), after: true } }), + }); + let sent: unknown; + const respond: FauxResponseStep = async (_context, options, _state, model) => { + sent = await options!.onPayload!({ body: 1 }, model); + return DONE; + }; + const { harness, status } = await run(setup, [respond]); + expect(status).toBe("done"); + expect(sent).toEqual({ body: 1, after: true }); + expect(setup.infos.filter((info) => info.kind === "extension_error")).toEqual([ + expect.objectContaining({ extension: "broken" }), + ]); + await harness.close(context); + }); + it("keeps the run's input open across an onYield continuation and answers it with the final answer", async () => { const setup = chatSetup(); let yields = 0;