Lifecycle Hooks

// server/lifecycle.ts
import type { AgentLifecycle } from "@pai/core";
import type { LanguageModel } from "ai";
import { z } from "zod";
import type { AssistantContext } from "./agent";
import { redactTextDeltas } from "./stream-transforms";

function latestUserText(
  messages: readonly { role: string; parts: readonly { type: string; text?: string }[] }[],
) {
  const message = [...messages].reverse().find((item) => item.role === "user");
  return message?.parts
    .filter((part) => part.type === "text")
    .map((part) => part.text ?? "")
    .join("\n")
    .trim() || null;
}

export function withRag(): AgentLifecycle<AssistantContext> {
  return {
    name: "workspace-rag",

    async prepareModelStep({ input, runtimeContext }) {
      const query = latestUserText(input.messages);
      if (!query) return;

      const docs = await runtimeContext.search.similar(query);
      if (docs.length === 0) return;

      return {
        action: "continue",
        patch: {
          appendModelMessages: [
            {
              role: "system",
              parts: [{ type: "text", text: `Workspace context:\n${docs.join("\n\n")}` }],
            },
          ],
        },
      };
    },
  };
}

export const blockUnsafeRequests: AgentLifecycle<AssistantContext> = {
  name: "block-unsafe-requests",

  async prepareRun({ latestUserMessage, runtimeContext }) {
    if (!latestUserMessage) return;
    if (!(await runtimeContext.policy.isUnsafe([latestUserMessage]))) return;

    return {
      action: "respond",
      outcome: "blocked",
      reason: "unsafe_request",
      response: {
        parts: [{ type: "text", text: "I cannot help with that request." }],
      },
    };
  },
};

export const redactSecrets: AgentLifecycle<AssistantContext> = {
  name: "redact-secrets",

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

export const blockUnsafeOutput: AgentLifecycle<AssistantContext> = {
  name: "block-unsafe-output",

  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." }],
              },
            });
          }
        },
      }),
    );
  },
};

export function withModelFallback(
  fallbackModel: LanguageModel,
): AgentLifecycle<AssistantContext> {
  return {
    name: "fallback-model",

    aroundModelCall(input, next) {
      if (modelHealth.isUnavailable(input.model)) {
        return next({ model: fallbackModel });
      }

      return next();
    },
  };
}

export const emailPolicy: AgentLifecycle<AssistantContext> = {
  name: "email-policy",

  beforeToolCall({ tool, input, runtimeContext }) {
    if (tool.name !== "sendEmail") return;

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

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

// server/agent.ts
import { composeLifecycle, defineAgent, type InferAgentContract } from "@pai/core";
import {
  blockUnsafeRequests,
  blockUnsafeOutput,
  emailPolicy,
  redactSecrets,
  withModelFallback,
  withRag,
} from "./lifecycle";
import { fallbackModel } from "./models";

export const agent = defineAgent({
  name: "report",
  model: ({ runtimeContext }) => runtimeContext.models.primary,
  instructions: ({ runtimeContext }) =>
    `Draft concise reports for ${runtimeContext.workspace.name}.`,
  runtimeContext: buildAssistantContext,
  tools: {
    sendEmail,
    updateReport,
  },
  lifecycle: composeLifecycle(
    blockUnsafeRequests,
    withRag(),
    withModelFallback(fallbackModel),
    emailPolicy,
    blockUnsafeOutput,
    redactSecrets,
  ),
});

export type AssistantAgentContract = InferAgentContract<typeof agent>;

import { createTestRuntime, mockModel, textResponse } from "@pai/test-utils";
import { expect, it } from "vitest";
import { agent } from "../server/agent";

it("redacts secrets before persistence", async () => {
  const runtime = createTestRuntime({
    agent,
    model: mockModel([textResponse("Token: sk-test-secret")]),
  });

  try {
    const pai = runtime.client({ identity });
    const thread = pai.thread(pai.newThreadId());

    await (await thread.send("Show token")).waitUntilIdle();
    await thread.refresh();

    const text = thread
      .getState()
      .messages.flatMap((message) => message.parts)
      .filter((part) => part.type === "text")
      .map((part) => part.text);

    expect(text).toContain("Token: [redacted secret]");
  } finally {
    await runtime.close();
  }
});