AkurAI Build
Menu

popagent

public

Latest change 7cecf6a89b6f39dae8c5678f5b369014199aeb3c - Add self-hosted Mastra observability by AkurAI Build

import { z } from "zod";
import { appLogger } from "./observability";

const HOOK_EVENTS = [
  "SessionStart",
  "TaskStart",
  "UserPromptSubmit",
  "SubagentStart",
  "PreToolUse",
  "PostToolUse",
  "PostToolUseFailure",
  "SubagentStop",
  "TaskStop",
  "Stop",
  "StopFailure",
  "SessionEnd",
] as const;

type HookEventName = (typeof HOOK_EVENTS)[number];
export type HookEvent = {
  schemaVersion: 1;
  eventId: string;
  eventName: HookEventName;
  occurredAt: string;
  sessionId: string;
  turnId?: string;
  toolCallId?: string;
  model: string;
  detail: Record<string, unknown>;
};

type HttpHookHandler = {
  id: string;
  type: "http";
  url: string;
  matcher?: string[];
  timeoutMs?: number;
  failureMode?: "open" | "closed";
};
type HookConfig = { version: 1; hooks: Partial<Record<HookEventName, HttpHookHandler[]>> };
type HookHandlerResult =
  | { schemaVersion: 1; outcome: "pass"; additionalContext?: string }
  | { schemaVersion: 1; outcome: "deny"; reason: string }
  | { schemaVersion: 1; outcome: "modify"; reason: string; replacement: unknown; additionalContext?: string }
  | { schemaVersion: 1; outcome: "continue"; reason: string; additionalContext: string };
type HookDispatchResult = {
  replacement?: unknown;
  additionalContext: string[];
  continue?: { reason: string; additionalContext: string };
};
export type HookRunRecord = {
  id: string;
  eventId: string;
  sessionId: string;
  turnId?: string;
  toolCallId?: string;
  eventName: HookEventName;
  handlerId: string;
  startedAt: string;
  completedAt: string;
  durationMs: number;
  status: "passed" | "modified" | "denied" | "continued" | "failed-open" | "failed-closed";
  reason?: string;
  error?: string;
};

type Transport = (handler: HttpHookHandler, event: HookEvent, signal: AbortSignal) => Promise<unknown>;
type RuntimeOptions = {
  config: HookConfig;
  transport?: Transport;
  audit?: (record: HookRunRecord) => Promise<void>;
};

const handlerSchema = z.object({
  id: z.string().min(1).max(128),
  type: z.literal("http"),
  url: z.url().refine((value) => new URL(value).protocol === "https:", "HTTP hooks require HTTPS"),
  matcher: z.array(z.string().min(1).max(128)).max(128).optional(),
  timeoutMs: z.number().int().min(50).max(30_000).optional(),
  failureMode: z.enum(["open", "closed"]).optional(),
}).strict();
const hooksShape = Object.fromEntries(
  HOOK_EVENTS.map((event) => [event, z.array(handlerSchema).max(64).optional()]),
) as Record<HookEventName, z.ZodOptional<z.ZodArray<typeof handlerSchema>>>;
const openOnlyEvents = new Set<HookEventName>([
  "SessionStart",
  "TaskStart",
  "PostToolUse",
  "PostToolUseFailure",
  "SubagentStop",
  "TaskStop",
  "StopFailure",
  "SessionEnd",
]);
const configSchema = z.object({
  version: z.literal(1),
  hooks: z.object(hooksShape).strict(),
}).strict().superRefine((config, refinement) => {
  for (const event of openOnlyEvents) {
    config.hooks[event]?.forEach((handler, index) => {
      if (handler.failureMode === "closed") {
        refinement.addIssue({
          code: "custom",
          message: `${event} hooks are observational and must fail open`,
          path: ["hooks", event, index, "failureMode"],
        });
      }
    });
  }
});
const utf8String = (maximum: number, label: string) => z.string().min(1).refine(
  (value) => Buffer.byteLength(value, "utf8") <= maximum,
  `${label} exceeds ${maximum} UTF-8 bytes`,
);
const MAX_ADDITIONAL_CONTEXT_BYTES = 16_384;
const context = utf8String(MAX_ADDITIONAL_CONTEXT_BYTES, "additionalContext");
const reason = utf8String(2_048, "reason");
const resultSchema = z.discriminatedUnion("outcome", [
  z.object({ schemaVersion: z.literal(1), outcome: z.literal("pass"), additionalContext: context.optional() }).strict(),
  z.object({ schemaVersion: z.literal(1), outcome: z.literal("deny"), reason }).strict(),
  z.object({
    schemaVersion: z.literal(1), outcome: z.literal("modify"), reason,
    replacement: z.unknown(), additionalContext: context.optional(),
  }).strict(),
  z.object({
    schemaVersion: z.literal(1), outcome: z.literal("continue"), reason,
    additionalContext: context,
  }).strict(),
]);
const allowed: Record<HookEventName, HookHandlerResult["outcome"][]> = {
  SessionStart: ["pass"],
  TaskStart: ["pass"],
  UserPromptSubmit: ["pass", "deny", "modify"],
  SubagentStart: ["pass", "deny", "modify"],
  PreToolUse: ["pass", "deny", "modify"],
  PostToolUse: ["pass"],
  PostToolUseFailure: ["pass"],
  SubagentStop: ["pass"],
  TaskStop: ["pass"],
  Stop: ["pass", "continue"],
  StopFailure: ["pass"],
  SessionEnd: ["pass"],
};
const closedByDefault: Partial<Record<HookEventName, true>> = {
  UserPromptSubmit: true,
  SubagentStart: true,
  PreToolUse: true,
  Stop: true,
};

export class HookBlockedError extends Error {
  override readonly name = "HookBlockedError";
  constructor(readonly reason: string, readonly handlerId: string) {
    super(reason);
  }
}

class HookProtocolError extends Error {}
const messageOf = (error: unknown) => error instanceof Error ? error.message : String(error);

const MAX_HOOK_RESPONSE_BYTES = 65_536;

async function boundedJSON(response: Response): Promise<unknown> {
  if (!response.ok) throw new Error(`HTTP ${response.status}`);
  const contentLength = response.headers.get("content-length");
  if (contentLength !== null && Number(contentLength) > MAX_HOOK_RESPONSE_BYTES) {
    throw new Error(`response exceeds ${MAX_HOOK_RESPONSE_BYTES} bytes`);
  }
  const chunks: Uint8Array[] = [];
  let size = 0;
  const reader = response.body?.getReader();
  if (reader) {
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      size += value.byteLength;
      if (size > MAX_HOOK_RESPONSE_BYTES) {
        await reader.cancel().catch(() => undefined);
        throw new Error(`response exceeds ${MAX_HOOK_RESPONSE_BYTES} bytes`);
      }
      chunks.push(value);
    }
  }
  const bytes = new Uint8Array(size);
  let offset = 0;
  for (const chunk of chunks) {
    bytes.set(chunk, offset);
    offset += chunk.byteLength;
  }
  return JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(bytes));
}

const httpTransport: Transport = async (handler, event, signal) => boundedJSON(await fetch(handler.url, {
  method: "POST",
  redirect: "error",
  headers: { "content-type": "application/json" },
  body: JSON.stringify(event),
  signal,
}));

function eventWithReplacement(event: HookEvent, replacement: unknown): HookEvent {
  const detail = structuredClone(event.detail);
  if (event.eventName === "UserPromptSubmit") detail.prompt = replacement;
  if (event.eventName === "SubagentStart") detail.prompt = replacement;
  if (event.eventName === "PreToolUse") (detail.tool as { input: unknown }).input = replacement;
  return { ...event, detail };
}

function resultViolation(event: HookEventName, result: HookHandlerResult): string | undefined {
  if (!allowed[event].includes(result.outcome)) return `Invalid ${event} hook outcome`;
  if ("additionalContext" in result && result.additionalContext !== undefined) {
    const supported = event === "UserPromptSubmit" || (event === "Stop" && result.outcome === "continue");
    if (!supported) return `${event} hooks cannot add context`;
  }
  if (result.outcome === "modify") {
    if (
      (event === "UserPromptSubmit" || event === "SubagentStart")
      && typeof result.replacement !== "string"
    ) {
      return `${event} replacements must be strings`;
    }
    if (
      event === "PreToolUse"
      && (typeof result.replacement !== "object" || result.replacement === null || Array.isArray(result.replacement))
    ) {
      return "PreToolUse replacements must be argument objects";
    }
  }
}

function failureModeFor(event: HookEventName, handler: HttpHookHandler): "open" | "closed" {
  return handler.failureMode ?? (closedByDefault[event] ? "closed" : "open");
}

export class HookRuntime {
  readonly config: HookConfig;
  private readonly transport: Transport;
  private readonly audit?: RuntimeOptions["audit"];

  constructor(options: RuntimeOptions) {
    const parsed = configSchema.safeParse(options.config);
    if (!parsed.success) throw new Error(`Invalid hook configuration: ${z.prettifyError(parsed.error)}`);
    this.config = parsed.data as HookConfig;
    this.transport = options.transport ?? httpTransport;
    this.audit = options.audit;
  }

  static fromJSON(json: string, options: Omit<RuntimeOptions, "config"> = {}): HookRuntime {
    try {
      return new HookRuntime({ ...options, config: JSON.parse(json) });
    } catch (error) {
      if (error instanceof SyntaxError) throw new Error(`Invalid hook configuration JSON: ${error.message}`);
      throw error;
    }
  }

  async dispatch(initial: HookEvent): Promise<HookDispatchResult> {
    let event = initial;
    let replacement: unknown;
    const additionalContext: string[] = [];
    let additionalContextBytes = 0;
    for (const handler of this.config.hooks[event.eventName] ?? []) {
      if (!this.matches(handler, event)) continue;
      const startedAt = new Date().toISOString();
      const started = performance.now();
      let result: HookHandlerResult;
      try {
        const parsed = resultSchema.safeParse(
          await this.transport(handler, event, AbortSignal.timeout(handler.timeoutMs ?? 2_000)),
        );
        if (!parsed.success) throw new HookProtocolError(`Invalid ${event.eventName} hook result`);
        const violation = resultViolation(event.eventName, parsed.data);
        if (violation) throw new HookProtocolError(violation);
        const nextContext = "additionalContext" in parsed.data ? parsed.data.additionalContext : undefined;
        if (
          nextContext !== undefined
          && additionalContextBytes + Buffer.byteLength(nextContext, "utf8") > MAX_ADDITIONAL_CONTEXT_BYTES
        ) {
          throw new HookProtocolError(
            `Combined additionalContext exceeds ${MAX_ADDITIONAL_CONTEXT_BYTES} UTF-8 bytes`,
          );
        }
        if (nextContext !== undefined) {
          additionalContextBytes += Buffer.byteLength(nextContext, "utf8");
        }
        result = parsed.data;
      } catch (error) {
        const failureMode = failureModeFor(event.eventName, handler);
        const errorText = messageOf(error);
        await this.record(event, handler, startedAt, started, `failed-${failureMode}`, undefined, errorText);
        if (failureMode === "closed") throw new HookBlockedError(`Hook ${handler.id} failed: ${errorText}`, handler.id);
        continue;
      }
      if (result.outcome === "deny") {
        await this.record(event, handler, startedAt, started, "denied", result.reason);
        throw new HookBlockedError(result.reason, handler.id);
      }
      if (result.outcome === "modify") {
        replacement = result.replacement;
        event = eventWithReplacement(event, replacement);
        if (result.additionalContext) additionalContext.push(result.additionalContext);
        await this.record(event, handler, startedAt, started, "modified", result.reason);
      } else if (result.outcome === "continue") {
        await this.record(event, handler, startedAt, started, "continued", result.reason);
        return { replacement, additionalContext, continue: result };
      } else {
        if (result.additionalContext) additionalContext.push(result.additionalContext);
        await this.record(event, handler, startedAt, started, "passed");
      }
    }
    return { ...(replacement === undefined ? {} : { replacement }), additionalContext };
  }

  private matches(handler: HttpHookHandler, event: HookEvent): boolean {
    if (!handler.matcher?.length) return true;
    const subject = (
      event.detail.tool as { name?: unknown } | undefined
    )?.name ?? (
      event.detail.agent as { id?: unknown } | undefined
    )?.id;
    return typeof subject === "string" && handler.matcher.includes(subject);
  }

  private async record(
    event: HookEvent, handler: HttpHookHandler, startedAt: string, started: number,
    status: HookRunRecord["status"], reason?: string, error?: string,
  ) {
    if (!this.audit) return;
    try {
      await this.audit({
        id: crypto.randomUUID(), eventId: event.eventId, sessionId: event.sessionId,
        turnId: event.turnId, toolCallId: event.toolCallId, eventName: event.eventName,
        handlerId: handler.id, startedAt, completedAt: new Date().toISOString(),
        durationMs: performance.now() - started, status, reason, error,
      });
    } catch (auditError) {
      appLogger().error("hook.audit.persistence_failed", { error: messageOf(auditError) });
    }
  }
}

export function createHookEvent(
  eventName: HookEventName,
  input: Omit<HookEvent, "schemaVersion" | "eventId" | "eventName" | "occurredAt">,
): HookEvent {
  return { schemaVersion: 1, eventId: crypto.randomUUID(), eventName, occurredAt: new Date().toISOString(), ...input };
}

export function loadHookRuntime(audit?: RuntimeOptions["audit"]): HookRuntime {
  const source = process.env.POPAGENT_HOOKS?.trim();
  const runtime = source
    ? HookRuntime.fromJSON(source, { audit })
    : new HookRuntime({ config: { version: 1, hooks: {} }, audit });
  const hosts = (process.env.POPAGENT_HOOK_ALLOWED_HOSTS ?? "").split(",").map((host) => host.trim()).filter(Boolean);
  for (const handlers of Object.values(runtime.config.hooks)) {
    for (const handler of handlers ?? []) {
      if (!hosts.includes(new URL(handler.url).hostname)) throw new Error(`Hook host is not allowed: ${new URL(handler.url).hostname}`);
    }
  }
  return runtime;
}