Lifecycle Hooks

Lifecycle hooks are the advanced extension point for agent behavior.

They live on the agent definition because they describe what the agent does. The same hooks run when the agent is used through HTTP, an in-process CLI client, tests, workers, or another transport.

import { composeLifecycle, defineAgent } from "@pai/core";

export const agent = defineAgent({
  name: "report",
  model,
  instructions,
  tools,
  lifecycle: composeLifecycle(
    withWorkspaceRag(),
    redactSecrets(),
    {
      prepareRun() {
        return {
          action: "continue",
          patch: {
            appendModelMessages: [
              {
                role: "system",
                content: "Prefer concise answers.",
              },
            ],
          },
        };
      },
    },
  ),
});

Why Lifecycle Instead Of Middleware

Generic middleware usually wraps a whole request and can mutate anything. Agent runtimes need stricter boundaries:

• model input changes affect what the LLM sees;
• canonical stream transforms affect persisted transcript state;
• tool hooks affect side effects;
• resolveModelContext chooses durable/base context for each model step;
• afterStep can affect only a later execution-local model step;
• live stream transforms should not change stored state.

PAI uses named phases so the API makes those differences explicit. See Lifecycle Phases for the full execution order.

Hooks return structured decisions such as continue, respond, nextStep, complete, replaceInput, replaceOutput, or reject. They do not receive the raw storage provider. When a hook needs durable thread changes, it uses the runtime-owned thread facade or a deliberate trusted/admin API. That keeps scope checks, revisions, active-run leases, realtime invalidation, and provider portability in one place.

Phase Guide

Use runtimeContext to load product state and dependencies. Use lifecycle hooks to change how that state is applied to the agent loop.

RAG And Model Context

Use prepareModelStep when each model step should see additional context that is not part of the user's transcript.

The example uses a local latestUserText() helper so later tool or assistant steps do not accidentally become the retrieval query.

function withWorkspaceRag(): AgentLifecycle<AssistantContext> {
  return {
    async prepareModelStep({ input, runtimeContext }) {
      const query = latestUserText(input.messages);
      if (!query) return;

      const docs = await runtimeContext.search.similar(query);

      return {
        action: "continue",
        patch: {
          appendModelMessages: [
            {
              role: "system",
              content: `Relevant workspace context:\n${docs.join("\n\n")}`,
            },
          ],
        },
      };
    },
  };
}

The appended message is model-only context. It is not added to the transcript.

prepareRun runs once per execution episode, including resume. Use it for context that can be reconstructed whenever the run acquires execution. Use a hidden lifecycle context message when model context must survive waiting and resume.

Guardrails

Use prepareRun to stop before the first model step and commit a normal assistant response.

const blockPromptInjection: AgentLifecycle<AssistantContext> = {
  async prepareRun({ latestUserMessage, runtimeContext }) {
    const latest = latestUserMessage ? latestUserText([latestUserMessage]) : null;

    if (latest && (await runtimeContext.policy.isPromptInjection(latest))) {
      return {
        action: "respond",
        outcome: "blocked",
        reason: "prompt_injection",
        response: {
          parts: [
            {
              type: "text",
              text: "I cannot follow instructions that bypass workspace policy.",
            },
          ],
        },
      };
    }
  },
};

Throw from a hook when the run should fail instead of producing an assistant response.

action: "respond" is terminal: it commits that assistant response, skips the model/tool phases, and ends the run. The response still goes through canonical output policy before commit.

Model Routing

prepareModelStep can patch only the current model step. The next step starts from the current execution episode's defaults and model-only overlay again.

const routeLargeDrafts: AgentLifecycle<AssistantContext> = {
  prepareModelStep({ input, runtimeContext }) {
    if (
      runtimeContext.workspace.plan === "enterprise" &&
      looksLikeLargeDraft(input)
    ) {
      return {
        action: "continue",
        patch: {
          model: runtimeContext.models.large,
          modelSettings: {
            maxOutputTokens: 8_000,
            temperature: 0.2,
          },
          providerOptions: {
            openai: {
              reasoningEffort: "medium",
            },
          },
        },
      };
    }
  },
};

A sparse run or step patch that supplies model starts a fresh route, even if it supplies the same model object. Inherited modelSettings, providerOptions, providerTools, and toolChoice are cleared before values from that patch are applied. Registered PAI tools remain available unless the patch also replaces tools. Patches without model merge model settings by field and provider options by provider key; provider tools and tool choice replace as whole values.

This is the replacement for old onBeforeStep and input processor use cases.

Canonical Redaction

Use transformStepStream when the transformed output should be persisted, shown after refresh, and reused as future model context.

const redactSecrets: AgentLifecycle = {
  transformStepStream({ stream }) {
    return stream.pipeThrough(
      redactTextDeltas({
        patterns: [/sk-[A-Za-z0-9_-]+/g],
        replacement: "[redacted secret]",
      }),
    );
  },
};

Do not use live presentation code for redaction that matters. Live presentation code is not the durable source of truth.

Use structured aborts when the stream should stop because of a policy decision.

const blockUnsafeOutput: AgentLifecycle<AssistantContext> = {
  transformStepStream({ stream, abort }) {
    return stream.pipeThrough(
      inspectTextDeltas({
        async onText(text) {
          if (await isUnsafeOutput(text)) {
            await abort({
              reason: "unsafe_output",
              outcome: "blocked",
              replacement: {
                parts: [{ type: "text", text: "I cannot continue with that output." }],
              },
            });
          }
        },
      }),
    );
  },
};

Already-emitted chunks cannot be unsent. Strict guardrails should buffer until content is safe to release.

Tool Policy

Use tool hooks for policy around model-requested tool calls and final tool results.

const protectEmailTools: AgentLifecycle<AssistantContext> = {
  beforeToolCall({ tool, input, runtimeContext }) {
    if (tool.name !== "sendEmail") return;

    const to = z.object({ to: z.string().email() }).parse(input).to;

    if (!runtimeContext.policy.canEmail(to)) {
      return {
        action: "reject",
        message: `Email to ${to} is not allowed in this workspace.`,
      };
    }
  },

  afterToolResult({ tool, result }) {
    if (tool.name === "lookupCustomer" && result.status === "output-available") {
      return {
        action: "replaceOutput",
        output: removePrivateFields(result.output),
      };
    }
  },
};

For tool calls in the same model step, PAI runs beforeToolCall for every tool call before starting any ready backend execution or client routing. Ready tools then run in parallel. This keeps policy deterministic without making independent tool calls wait on each other.

afterToolResult runs before any final tool result returns to the model, including skipped outputs, client-submitted results, resumed backend tools, cancellations, and failures that are converted into model-visible tool errors.

For backend errors, result.source distinguishes an exception-based failure ("throw") from a deliberate ctx.fail(failure) result ("fail"). The distinction is also what the model sees: a deliberate failure reaches it verbatim, while a thrown one is redacted.

afterToolResult is agent-wide, so it sees each result as name: string with output: unknown. To keep agent state from one specific tool's result, register onOutput beside that tool instead: it is typed against the tool's own schemas, so no name matching or casting is needed and a rename becomes a compile error. Use afterToolResult when the behavior really is agent-wide, or when you need to observe failures as well as outputs.

Human approval is still a tool feature, not a lifecycle hook. Use Suspend And Resume when execution must pause for external input.

Conditional Model Routing

Use aroundModelCall when a last-mile condition should route one provider attempt differently. The hook receives the effective request, while next() accepts only a sparse patch.

The request is a shallow read-only view, and input.runtimeContext has the runtime-context type inferred from the agent. Top-level fields cannot be reassigned. Nested application and provider objects are deliberately not made recursively read-only and should be treated as observational inputs.

import type { LanguageModel } from "ai";

function withModelFallback(
  fallbackModel: LanguageModel,
): AgentLifecycle<AssistantContext> {
  return {
    aroundModelCall(input, next) {
      if (modelHealth.isUnavailable(input.model)) {
        return next({ model: fallbackModel });
      }

      return next();
    },
  };
}

Supplying model applies the same atomic route rule as run and step patches: inherited settings, provider options, provider-native tools, and tool choice are cleared before values from the same patch are applied. Registered PAI tools are not cleared. Do not pass or spread input into next(); it is a complete request and includes runtime-owned fields as well as the previous route's configuration.

next() resolves when the provider stream has been created, not when that stream has finished. A normal try/catch around await next() therefore does not catch failures that arrive while the stream is consumed. Stream-time fallback needs an explicit stream-aware policy and is safe only before output has been released to clients.

An attempt-specific signal can be supplied as a patch. PAI combines it with the runtime signal, so stop and shutdown remain authoritative:

return next({
  signal: AbortSignal.timeout(10_000),
});

Use provider-specific model middleware when the behavior belongs to a model provider. Use PAI lifecycle when the behavior depends on agent runtime context, tools, or thread state.

After Step

afterStep runs after a canonical assistant step has been assembled and before the final step/run transition is committed. It cannot rewrite the assembled message. It can request another model step or complete the run.

const continueAfterDraft: AgentLifecycle = {
  afterStep({ step }) {
    if (!needsVerification(step.message)) return;
    return { action: "nextStep" };
  },
};

Return undefined to keep the runtime's natural transition, nextStep to request another eligible model call, or complete to finish deliberately. Waiting and steering take precedence over nextStep. If complete is returned for a step with waiting tools, the runtime cancels those pending parts before completing the run.

Queue admission, leases, and wakeups are internal runtime interceptors. They do not depend on user lifecycle hook ordering.

appendModelMessages and replaceModelMessages are execution-episode-local. Waiting/resume, steering, completion, cancellation, or loss of execution ownership discards them. Persist durable application state through typed run metadata, your app database, or tools.

Model Context Resolution

resolveModelContext is the lifecycle boundary for moving windows, retrieval, compaction, and other durable context policies. Its next() function resolves the downstream strategy or the default chronological set of context-visible messages:

const recentContext: AgentLifecycle = {
  async resolveModelContext(input, next) {
    const context = await next();
    return context.slice(-20);
  },
};

next() is memoized. next(base) replaces the downstream leaf while still running inner middleware, which lets a cursor-backed source preserve redaction and policy layers. Both consume and return the AI SDK's ModelMessage[] shape: model-facing messages use content, not UI-message parts.

Durable history uses a different trusted boundary. input.history.collect() returns LifecycleMessage[]: the canonical native-first PAI message plus visibility and private metadata. input.history.modelMessages() converts the same fixed-anchor range directly from its private native values into ModelMessage[] for provider context. PAI never round-trips the public LifecycleMessage projection back through the SDK.

A strategy can bypass composition and collect an inclusive anchored range with input.history.collect(). Raw exclusive input.history.before() pages remain available when a strategy needs to own pagination and handle history.through itself. Every read is provider-backed and fenced to the owning agent and current thread incarnation. It is also bounded by the newest physical message observed at step start, so context does not require an eager full-thread snapshot and concurrent appends cannot enter the current step. The active run lease prevents regeneration and administrative history rewrites from overlapping this work. Resolver hooks must not rewrite message visibility while selecting context, so their phase-specific thread facade omits messages.updateVisibility() directly and inside mutate().

LifecycleMessage.id and LifecycleMessage.metadata.pai.updatedAt retain durable provenance. ModelMessage is deliberately provider-facing and carries only the SDK context shape. Model-only messages from run preparation are appended after the durable resolver returns.

Thread History Changes

Use the server-only thread facade for strategy state, never raw storage. Typed thread.privateMetadata is suitable for cursors that clients do not need to render, while thread.metadata is server-written and client-readable.

Writes staged inside thread.mutate() commit atomically after its callback returns. A complete expectedVersion CAS is available when every root change must invalidate the work. A slow strategy can instead compare its private cursor without rejecting unrelated usage or lease revisions:

await thread.mutate(
  async (tx) => {
    await tx.messages.append(summaryMessage);
    await tx.privateMetadata.merge({ cursor: throughMessageId });
  },
  {
    expectedPrivateMetadata: privateMetadata,
  },
);

The facade is scoped to the current thread and execution lease so the runtime preserves revisions, realtime behavior, provider portability, and scope isolation. Reads inside a mutation callback observe persisted state and do not include earlier staged writes.

For the common summary strategy, use createContextCompaction() instead of rewriting original message visibility yourself. It installs context resolution, transcript-visible/context-hidden summary records, public run progress, a private cursor, atomic mutation, and regeneration repair as one lifecycle helper. See Context compaction.

When regeneration or retry truncates physical history, or trusted code changes stored message visibility, prepareHistoryRewrite can return thread metadata patches that commit in the same transaction. Its rewrite input discriminates kind: "regenerate" from kind: "visibility"; the latter includes complete before/after visibility for every changed message. historyRewritten is a best-effort post-commit fact for idempotent external reconciliation. Hooks are serialized per thread and concurrency-bounded across threads; pending facts may be coalesced. It cannot roll back the rewrite or guarantee delivery, so an external atomic invariant needs an application outbox or coordinated store.

Testing Hooks

Test lifecycle behavior through an in-process runtime.

import { createTestRuntime, mockModel, textResponse } from "@pai/test-utils";

const runtime = createTestRuntime({
  agent,
  model: mockModel([textResponse("The API key is sk-test.")]),
});

const pai = runtime.client({ identity });
const thread = pai.thread(pai.newThreadId());
const state = await (await thread.send("Show a secret")).waitUntilIdle();

const assistant = state.messages.find(
  (message) =>
    message.role === "assistant" && message.metadata.pai.producer === "model",
);

expect(assistant?.parts).toContainEqual(
  expect.objectContaining({
    type: "text",
    text: "The API key is [redacted secret].",
  }),
);

For canonical transforms, assert against refreshed thread state. For client or transport presentation transforms, assert against the watch/request stream separately.