@beignet/core/src/agent-capabilities/index.ts

Core framework primitives for Beignet

import type { StandardSchemaV1 } from "@standard-schema/spec";
import {
  createProviderInstrumentation,
  type ProviderInstrumentationTarget,
} from "../providers/index.js";
import { runWithResolvedTracingContext } from "../tracing/execution.js";
import type { TracingPort } from "../tracing/index.js";

/** Any Standard Schema compatible validator. */
export type AgentCapabilitySchema = StandardSchemaV1<unknown, unknown>;

/** Value or promise of that value. */
export type MaybePromise<T> = T | Promise<T>;

/** Infer the input accepted by a capability schema. */
export type InferAgentCapabilitySchemaInput<Schema extends StandardSchemaV1> =
  StandardSchemaV1.InferInput<Schema>;

/** Infer the parsed output produced by a capability schema. */
export type InferAgentCapabilitySchemaOutput<Schema extends StandardSchemaV1> =
  StandardSchemaV1.InferOutput<Schema>;

/** A typed action that an authenticated agent transport may expose. */
export interface AgentCapabilityDef<
  Name extends string = string,
  Input extends AgentCapabilitySchema = AgentCapabilitySchema,
  Output extends AgentCapabilitySchema = AgentCapabilitySchema,
  Ctx = unknown,
  Principal = unknown,
> {
  readonly kind: "agentCapability";
  readonly name: Name;
  readonly description: string;
  readonly input: Input;
  readonly output: Output;
  readonly handle: (args: {
    capability: AgentCapabilityDef<Name, Input, Output, Ctx, Principal>;
    ctx: Ctx;
    principal: Principal;
    input: InferAgentCapabilitySchemaOutput<Input>;
  }) => MaybePromise<InferAgentCapabilitySchemaInput<Output>>;
}

/** Broad capability definition shape tied to one context and principal. */
export interface AgentCapabilityFor<Ctx, Principal> {
  readonly kind: "agentCapability";
  readonly name: string;
  readonly description: string;
  readonly input: AgentCapabilitySchema;
  readonly output: AgentCapabilitySchema;
  readonly handle: (args: {
    capability: never;
    ctx: Ctx;
    principal: Principal;
    input: never;
  }) => MaybePromise<unknown>;
}

/** Broad capability definition shape used by registries and adapters. */
export type AnyAgentCapabilityDef = AgentCapabilityFor<never, never>;

/** Infer a capability's application context. */
export type AgentCapabilityContext<Capability extends AnyAgentCapabilityDef> =
  Capability extends { handle(args: { ctx: infer Ctx }): unknown }
    ? Ctx
    : never;

/** Infer a capability's verified principal. */
export type AgentCapabilityPrincipal<Capability extends AnyAgentCapabilityDef> =
  Capability extends { handle(args: { principal: infer Principal }): unknown }
    ? Principal
    : never;

/** Infer a capability's raw input. */
export type AgentCapabilityInput<Capability extends AnyAgentCapabilityDef> =
  InferAgentCapabilitySchemaInput<Capability["input"]>;

/** Infer the parsed input passed to capability context and execution. */
export type AgentCapabilityParsedInput<
  Capability extends AnyAgentCapabilityDef,
> = InferAgentCapabilitySchemaOutput<Capability["input"]>;

/** Infer a capability's validated output. */
export type AgentCapabilityOutput<Capability extends AnyAgentCapabilityDef> =
  InferAgentCapabilitySchemaOutput<Capability["output"]>;

/** Options accepted by `defineAgentCapability(...)`. */
export interface DefineAgentCapabilityOptions<
  Name extends string,
  Input extends AgentCapabilitySchema,
  Output extends AgentCapabilitySchema,
  Ctx,
  Principal,
> {
  description: string;
  input: Input;
  output: Output;
  handle(args: {
    capability: AgentCapabilityDef<Name, Input, Output, Ctx, Principal>;
    ctx: Ctx;
    principal: Principal;
    input: InferAgentCapabilitySchemaOutput<Input>;
  }): MaybePromise<InferAgentCapabilitySchemaInput<Output>>;
}

/** Context-bound capability definition helpers. */
export interface AgentCapabilities<Ctx, Principal> {
  defineAgentCapability<
    const Name extends string,
    Input extends AgentCapabilitySchema,
    Output extends AgentCapabilitySchema,
  >(
    name: Name,
    options: DefineAgentCapabilityOptions<Name, Input, Output, Ctx, Principal>,
  ): AgentCapabilityDef<Name, Input, Output, Ctx, Principal>;
  defineAgentCapabilityRegistry<
    const Definitions extends readonly AgentCapabilityFor<Ctx, Principal>[],
  >(definitions: Definitions): AgentCapabilityRegistry<Definitions>;
}

/** Bind application context and principal types once for capability definitions. */
export function createAgentCapabilities<Ctx, Principal>(): AgentCapabilities<
  Ctx,
  Principal
> {
  return {
    defineAgentCapability(name, options) {
      return {
        kind: "agentCapability",
        name,
        description: options.description,
        input: options.input,
        output: options.output,
        handle: options.handle,
      };
    },
    defineAgentCapabilityRegistry,
  };
}

/** Error thrown when a capability registry is ambiguous. */
export class AgentCapabilityRegistryError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "AgentCapabilityRegistryError";
  }
}

/** Registered capability definitions available to execution adapters. */
export interface AgentCapabilityRegistry<
  Definitions extends
    readonly AnyAgentCapabilityDef[] = readonly AnyAgentCapabilityDef[],
> {
  readonly definitions: Definitions;
  get(name: string): Definitions[number] | undefined;
}

/** Build a capability registry and reject duplicate external names. */
function defineAgentCapabilityRegistry<
  const Definitions extends readonly AnyAgentCapabilityDef[],
>(definitions: Definitions): AgentCapabilityRegistry<Definitions> {
  const byName = new Map<string, Definitions[number]>();

  for (const capability of definitions) {
    if (byName.has(capability.name)) {
      throw new AgentCapabilityRegistryError(
        `Duplicate agent capability definition "${capability.name}" in registry.`,
      );
    }
    byName.set(capability.name, capability);
  }

  return {
    definitions: [...definitions] as unknown as Definitions,
    get(name) {
      return byName.get(name);
    },
  };
}

/** Stable capability execution failure categories. */
export type AgentCapabilityErrorCode =
  | "unknown_capability"
  | "invalid_input"
  | "invalid_output"
  | "execution_failed";

/** Error raised by capability lookup, validation, or execution. */
export class AgentCapabilityError extends Error {
  readonly code: AgentCapabilityErrorCode;
  readonly capabilityName: string;
  readonly issues?: readonly StandardSchemaV1.Issue[];
  override readonly cause?: unknown;

  constructor(args: {
    code: AgentCapabilityErrorCode;
    capabilityName: string;
    message: string;
    issues?: readonly StandardSchemaV1.Issue[];
    cause?: unknown;
  }) {
    super(args.message);
    this.name = "AgentCapabilityError";
    this.code = args.code;
    this.capabilityName = args.capabilityName;
    this.issues = args.issues;
    this.cause = args.cause;
  }
}

/** A parsed invocation passed to the app-owned context resolver. */
export type AgentCapabilityContextRequest<
  Definitions extends readonly AnyAgentCapabilityDef[],
  Principal,
> = {
  [Capability in Definitions[number] as Capability["name"]]: {
    capability: Capability;
    name: Capability["name"];
    principal: Principal;
    input: InferAgentCapabilitySchemaOutput<Capability["input"]>;
  };
}[Definitions[number]["name"]];

/** Stage reached by an agent capability attempt. */
export type AgentCapabilityRunStage =
  | "lookup"
  | "input"
  | "authorization"
  | "context"
  | "handler"
  | "output";

/** Successful capability lifecycle event with validated application data. */
export type AgentCapabilityCompletedRunEvent<
  Ctx,
  Principal,
  Definitions extends
    readonly AnyAgentCapabilityDef[] = readonly AnyAgentCapabilityDef[],
  Capability extends Definitions[number] = Definitions[number],
> = Capability extends Definitions[number]
  ? {
      name: Capability["name"];
      phase: "end";
      stage: "output";
      capability: Capability;
      ctx: Ctx;
      principal: Principal;
      input: AgentCapabilityParsedInput<Capability>;
      output: AgentCapabilityOutput<Capability>;
      durationMs: number;
      error?: undefined;
    }
  : never;

/** Lifecycle event observed around a capability attempt. */
export type AgentCapabilityRunEvent<
  Ctx,
  Principal,
  Definitions extends
    readonly AnyAgentCapabilityDef[] = readonly AnyAgentCapabilityDef[],
> =
  | {
      name: string;
      phase: "start";
      stage: "lookup";
      capability?: Definitions[number];
      ctx?: undefined;
      principal: Principal;
      input?: undefined;
      output?: undefined;
      durationMs?: undefined;
      error?: undefined;
    }
  | AgentCapabilityCompletedRunEvent<Ctx, Principal, Definitions>
  | {
      name: string;
      phase: "error";
      stage: AgentCapabilityRunStage;
      capability?: Definitions[number];
      ctx?: Ctx;
      principal: Principal;
      /** Present only when input validation completed successfully. */
      input?: unknown;
      output?: undefined;
      durationMs: number;
      error: unknown;
    };

/** Best-effort lifecycle observer for capability execution. */
export type AgentCapabilityHook<
  Ctx,
  Principal,
  Definitions extends
    readonly AnyAgentCapabilityDef[] = readonly AnyAgentCapabilityDef[],
> = (
  event: AgentCapabilityRunEvent<Ctx, Principal, Definitions>,
) => MaybePromise<void>;

/** Options for a capability executor. */
export interface CreateAgentCapabilityExecutorOptions<
  Ctx,
  Principal,
  Definitions extends readonly AnyAgentCapabilityDef[],
> {
  registry: AgentCapabilityRegistry<Definitions>;
  createContext(
    request: AgentCapabilityContextRequest<Definitions, Principal>,
  ): MaybePromise<Ctx>;
  hooks?: readonly AgentCapabilityHook<Ctx, Principal, Definitions>[];
  /** Independent target used to observe failures before context exists. */
  instrumentation?: ProviderInstrumentationTarget;
  tracing?: TracingPort;
}

/** Raw dynamic invocation accepted by protocol adapters. */
export interface DynamicAgentCapabilityInvocation<Principal> {
  name: string;
  principal: Principal;
  input: unknown;
  /**
   * Transport-owned authorization against the exact parsed input that will be
   * passed to context construction and the capability handler.
   */
  authorize?(
    request: DynamicAgentCapabilityAuthorization<Principal>,
  ): MaybePromise<void>;
}

/** Parsed invocation exposed to a dynamic transport authorization callback. */
export interface DynamicAgentCapabilityAuthorization<Principal> {
  capability: AnyAgentCapabilityDef;
  name: string;
  principal: Principal;
  input: unknown;
}

/** Typed and dynamic execution surface for one capability registry. */
export interface AgentCapabilityExecutor<
  Principal,
  Definitions extends readonly AnyAgentCapabilityDef[],
> {
  execute<Name extends Definitions[number]["name"]>(args: {
    name: Name;
    principal: Principal;
    input: AgentCapabilityInput<Extract<Definitions[number], { name: Name }>>;
  }): Promise<
    AgentCapabilityOutput<Extract<Definitions[number], { name: Name }>>
  >;
  executeDynamic(
    args: DynamicAgentCapabilityInvocation<Principal>,
  ): Promise<unknown>;
}

function formatPath(path: StandardSchemaV1.Issue["path"]): string {
  if (!path?.length) return "";
  return path
    .map((segment) =>
      typeof segment === "object" && segment !== null && "key" in segment
        ? String(segment.key)
        : String(segment),
    )
    .join(".");
}

function formatIssues(issues: readonly StandardSchemaV1.Issue[]): string {
  return issues
    .map((issue) => {
      const path = formatPath(issue.path);
      return path ? `${path}: ${issue.message}` : issue.message;
    })
    .join("; ");
}

async function parseCapabilitySchema<Schema extends AgentCapabilitySchema>(
  schema: Schema,
  input: unknown,
  args: {
    capabilityName: string;
    code: "invalid_input" | "invalid_output";
  },
): Promise<InferAgentCapabilitySchemaOutput<Schema>> {
  const result = await schema["~standard"].validate(input);
  if (result.issues?.length) {
    throw new AgentCapabilityError({
      code: args.code,
      capabilityName: args.capabilityName,
      message: `Agent capability "${args.capabilityName}" ${args.code === "invalid_input" ? "input" : "output"} validation failed: ${formatIssues(result.issues)}`,
      issues: result.issues,
    });
  }
  if ("value" in result) {
    return result.value as InferAgentCapabilitySchemaOutput<Schema>;
  }
  throw new AgentCapabilityError({
    code: args.code,
    capabilityName: args.capabilityName,
    message: `Agent capability "${args.capabilityName}" schema returned no value.`,
  });
}

function instrumentationTarget(ctx: unknown): ProviderInstrumentationTarget {
  if (!ctx || typeof ctx !== "object") return undefined;
  if ("ports" in ctx) {
    return (ctx as { ports?: ProviderInstrumentationTarget }).ports;
  }
  return ctx as ProviderInstrumentationTarget;
}

async function notifyHooks<
  Ctx,
  Principal,
  Definitions extends readonly AnyAgentCapabilityDef[],
>(
  hooks: readonly AgentCapabilityHook<Ctx, Principal, Definitions>[],
  event: AgentCapabilityRunEvent<Ctx, Principal, Definitions>,
): Promise<void> {
  await Promise.all(
    hooks.map(async (hook) => {
      try {
        await hook(event);
      } catch {
        // Observers must not alter application behavior.
      }
    }),
  );
}

/** Create a validated, traced executor for a capability registry. */
export function createAgentCapabilityExecutor<
  const Definitions extends readonly AnyAgentCapabilityDef[],
>(
  options: CreateAgentCapabilityExecutorOptions<
    AgentCapabilityContext<Definitions[number]>,
    AgentCapabilityPrincipal<Definitions[number]>,
    Definitions
  >,
): AgentCapabilityExecutor<
  AgentCapabilityPrincipal<Definitions[number]>,
  Definitions
> {
  type Ctx = AgentCapabilityContext<Definitions[number]>;
  type Principal = AgentCapabilityPrincipal<Definitions[number]>;
  const hooks = options.hooks ?? [];
  const instrumentationConfig = {
    providerName: "agent-capabilities",
    watcher: "agentCapabilities",
  } as const;

  async function executeDynamic(
    invocation: DynamicAgentCapabilityInvocation<Principal>,
  ): Promise<unknown> {
    const startedAt = Date.now();
    let stage: AgentCapabilityRunStage = "lookup";
    let ctx: Ctx | undefined;
    let capability: Definitions[number] | undefined;
    let input: unknown;
    let hasValidatedInput = false;
    let instrumentation = createProviderInstrumentation(
      options.instrumentation,
      instrumentationConfig,
    );

    try {
      capability = options.registry.get(invocation.name);
      const observedName = capability?.name ?? "unknown";
      instrumentation.custom({
        name: "agentCapability.started",
        label: `Agent capability ${observedName}`,
        summary: "Capability execution started",
        details: { capabilityName: observedName, phase: "start" },
      });
      await notifyHooks(hooks, {
        name: invocation.name,
        phase: "start",
        stage,
        principal: invocation.principal,
      });

      if (!capability) {
        throw new AgentCapabilityError({
          code: "unknown_capability",
          capabilityName: invocation.name,
          message: `Unknown agent capability "${invocation.name}".`,
        });
      }

      const registeredCapability = capability;
      const runtimeCapability =
        registeredCapability as unknown as AgentCapabilityDef<
          string,
          AgentCapabilitySchema,
          AgentCapabilitySchema,
          Ctx,
          Principal
        >;
      return await runWithResolvedTracingContext({
        tracing: options.tracing,
        ctx: async () => {
          stage = "input";
          input = await parseCapabilitySchema(
            runtimeCapability.input,
            invocation.input,
            {
              capabilityName: runtimeCapability.name,
              code: "invalid_input",
            },
          );
          hasValidatedInput = true;
          if (invocation.authorize) {
            stage = "authorization";
            await invocation.authorize({
              capability: registeredCapability,
              name: registeredCapability.name,
              principal: invocation.principal,
              input,
            });
          }
          stage = "context";
          ctx = await options.createContext({
            capability: registeredCapability,
            name: registeredCapability.name,
            principal: invocation.principal,
            input,
          } as AgentCapabilityContextRequest<Definitions, Principal>);
          return ctx;
        },
        operation: {
          name: `beignet.agent_capability ${registeredCapability.name}`,
          type: "agentCapability",
          kind: "internal",
          attributes: {
            "beignet.agent_capability.name": registeredCapability.name,
          },
          metricAttributes: {
            "beignet.agent_capability.name": registeredCapability.name,
          },
        },
        async run(resolvedCtx) {
          ctx = resolvedCtx;
          if (options.instrumentation === undefined) {
            instrumentation = createProviderInstrumentation(
              instrumentationTarget(ctx),
              instrumentationConfig,
            );
            instrumentation.custom({
              name: "agentCapability.started",
              label: `Agent capability ${registeredCapability.name}`,
              summary: "Capability execution started",
              details: {
                capabilityName: registeredCapability.name,
                phase: "start",
              },
            });
          }

          stage = "handler";
          const rawOutput = await runtimeCapability.handle({
            capability: runtimeCapability,
            ctx,
            principal: invocation.principal,
            input,
          });
          stage = "output";
          const output = await parseCapabilitySchema(
            runtimeCapability.output,
            rawOutput,
            {
              capabilityName: registeredCapability.name,
              code: "invalid_output",
            },
          );
          const durationMs = Date.now() - startedAt;
          instrumentation.custom({
            name: "agentCapability.completed",
            label: `Agent capability ${registeredCapability.name}`,
            summary: "Capability execution completed",
            details: {
              capabilityName: registeredCapability.name,
              phase: "end",
              durationMs,
            },
          });
          await notifyHooks(hooks, {
            name: registeredCapability.name,
            phase: "end",
            stage,
            capability: registeredCapability,
            ctx,
            principal: invocation.principal,
            input,
            output,
            durationMs,
          } as AgentCapabilityRunEvent<Ctx, Principal, Definitions>);
          return output;
        },
      });
    } catch (error) {
      const wrapped =
        error instanceof AgentCapabilityError
          ? error
          : new AgentCapabilityError({
              code: "execution_failed",
              capabilityName: capability?.name ?? invocation.name,
              message: `Agent capability "${capability?.name ?? invocation.name}" execution failed.`,
              cause: error,
            });
      const durationMs = Date.now() - startedAt;
      const observedName = capability?.name ?? "unknown";
      instrumentation.custom({
        name: "agentCapability.failed",
        label: `Agent capability ${observedName}`,
        summary: "Capability execution failed",
        details: {
          capabilityName: observedName,
          phase: "error",
          stage,
          durationMs,
          code: wrapped.code,
        },
      });
      await notifyHooks(hooks, {
        name: invocation.name,
        phase: "error",
        stage,
        ...(capability ? { capability } : {}),
        ctx,
        principal: invocation.principal,
        durationMs,
        error: wrapped,
        ...(hasValidatedInput ? { input } : {}),
      } as AgentCapabilityRunEvent<Ctx, Principal, Definitions>);
      throw wrapped;
    }
  }

  return {
    execute: executeDynamic as AgentCapabilityExecutor<
      Principal,
      Definitions
    >["execute"],
    executeDynamic,
  };
}