@cursor/july 0.1.1

package/src/ab.ts

@@ -0,0 +1,675 @@
+/**
+ * Live A/B metrics plug-in for agent-serve.
+ *
+ * Author `agent/ab.ts` (or `agent/ab/<name>.ts`) with {@link defineAB}. At
+ * session creation the engine runs {@link ABConfig.split} and appends durable
+ * `ab.assigned` events; arms are exposed on {@link SessionInfo.abs} for the
+ * agent loop (tools / optional per-variant `instructions`). {@link ABCollector}
+ * folds the event stream into metrics and calls {@link ABConfig.onSample}.
+ * There is no `agent-serve ab` test runner.
+ *
+ * Assignments and foldable metrics live in each session's `events.ndjson`
+ * (survive serve restart). Optional project defaults in `agent/ab.config.ts`
+ * ({@link defineABConfig}) let authors retain the playground fold window and
+ * persist samples / aggregate snapshots wherever they like — same idea as
+ * eval `persistRuns`.
+ *
+ * ```ts
+ * import { defineAB, splitBySessionHash } from "@anysphere/agent-serve/ab";
+ *
+ * export default defineAB({
+ *   name: "concise-instructions",
+ *   variants: {
+ *     control: { label: "Baseline" },
+ *     treatment: { label: "Shorter" },
+ *   },
+ *   split: splitBySessionHash({ holdout: 0.1 }),
+ *   onSample(sample) {
+ *     console.log(sample.variant, sample.metrics);
+ *   },
+ * });
+ * ```
+ */
+
+import { randomUUID } from "node:crypto";
+import { appendFile, mkdir } from "node:fs/promises";
+import { join } from "node:path";
+import type { ABSnapshot } from "./internal/ab-snapshot.js";
+import { normalizePositiveInt } from "./internal/bounded-int.js";
+import { brandDefinition } from "./internal/brand.js";
+import {
+  assertNonEmptyDir,
+  deleteFileInDir,
+  listJsonBasenames,
+  readJsonDir,
+  resolveProjectRelativeDir,
+  writeJsonFile,
+} from "./internal/json-dir-store.js";
+import type { SessionEvent, SessionInfo } from "./types.js";
+
+export type { ABSnapshot } from "./internal/ab-snapshot.js";
+
+export interface ABVariantMeta {
+  label?: string;
+  description?: string;
+  /**
+   * Optional prompt overlay for this arm. Appended to the agent's
+   * instructions for the session (local `AGENTS.md` + cloud first-turn
+   * preamble) after sticky enrollment.
+   */
+  instructions?: string;
+}
+
+export interface ABSplitContext {
+  agent: { name: string };
+  channel: { id: string };
+  /**
+   * Live sessions only — the engine never calls `split` for
+   * `purpose: "eval"` harness sessions.
+   */
+  session: SessionInfo;
+  experiment: string;
+  /** Variant keys from {@link ABConfig.variants}, in declaration order. */
+  variants: string[];
+}
+
+export interface ABCollectContext {
+  agent: { name: string };
+  channel: { id: string };
+  session: SessionInfo;
+  experiment: string;
+  variant: string;
+}
+
+/** Cumulative per-session metrics for one experiment arm. */
+export interface ABMetrics {
+  turns: number;
+  turnFailures: number;
+  toolCalls: number;
+  toolErrors: number;
+  inputTokens: number;
+  outputTokens: number;
+  cacheReadTokens: number;
+  cacheWriteTokens: number;
+  /** Sum of turn wall times (turn.started → turn.completed/failed). */
+  wallTimeMs: number;
+  /** Custom counters from {@link ABConfig.derive}. */
+  custom: Record<string, number>;
+}
+
+export type ABSampleReason = "turn.completed" | "turn.failed";
+
+export interface ABMetricSample {
+  experiment: string;
+  variant: string;
+  variantLabel?: string;
+  sessionId: string;
+  channelId: string;
+  metrics: ABMetrics;
+  reason: ABSampleReason;
+  at: string;
+}
+
+/** Variant key to enroll, or `null` to leave the session out (sticky). */
+export type ABSplitResult = string | null;
+
+/**
+ * Sticky enrollment decision. Return a variant key, or `null` to skip.
+ * Prefer helpers ({@link splitBySessionHash}, {@link splitIf}, …).
+ */
+export type ABSplit = (
+  ctx: ABSplitContext
+) => ABSplitResult | Promise<ABSplitResult>;
+
+export interface ABSplitOptions {
+  /** Relative weights per variant key (default: equal). */
+  weights?: Record<string, number>;
+  /** Fraction of sessions left out (`0`–`1`, sticky). */
+  holdout?: number;
+  /** Salt so a new rollout reshuffles without renaming the experiment. */
+  salt?: string;
+}
+
+/** Sticky hash of session id → arm (recommended default). */
+export function splitBySessionHash(options?: ABSplitOptions): ABSplit {
+  assertHoldout(options?.holdout);
+  return (ctx) =>
+    pickArm(
+      ctx.variants,
+      options?.weights,
+      hashToUnitInterval(`${options?.salt ?? ""}\0${ctx.session.id}`),
+      options?.holdout
+    );
+}
+
+/** Random draw once per new session (then sticky via persistence). */
+export function splitByRandom(options?: Omit<ABSplitOptions, "salt">): ABSplit {
+  assertHoldout(options?.holdout);
+  return (ctx) =>
+    pickArm(ctx.variants, options?.weights, Math.random(), options?.holdout);
+}
+
+/** Pin every new session to one arm. */
+export function splitAlways(variant: string): ABSplit {
+  return () => variant;
+}
+
+/** Skip every session (pause without deleting the experiment). */
+export function splitNone(): ABSplit {
+  return () => null;
+}
+
+/** Run `inner` only when `predicate` is true; otherwise skip. */
+export function splitIf(
+  predicate: (ctx: ABSplitContext) => boolean | Promise<boolean>,
+  inner: ABSplit
+): ABSplit {
+  return async (ctx) => ((await predicate(ctx)) ? inner(ctx) : null);
+}
+
+export interface ABConfig {
+  /** Experiment id. Defaults to the file stem (`ab/concise.ts` → `concise`). */
+  name?: string;
+  /**
+   * Named arms (at least two). Prefer `control` / `treatment` over `a` / `b`.
+   */
+  variants: Record<string, ABVariantMeta>;
+  /**
+   * Choose an arm at session creation, or return `null` to skip. The engine
+   * appends a durable `ab.assigned` event; the collector folds that stream.
+   */
+  split: ABSplit;
+  /**
+   * Optional custom metric extractors. Return a numeric delta, or
+   * `null`/`undefined` to skip. Called on **every** session event (including
+   * high-frequency `message.appended` deltas) — keep these synchronous and
+   * cheap; do not await I/O.
+   */
+  derive?: Record<
+    string,
+    (
+      event: SessionEvent,
+      ctx: ABCollectContext
+    ) => number | null | undefined | Promise<number | null | undefined>
+  >;
+  /**
+   * Called on turn complete/fail for enrolled sessions only. Observe-only;
+   * errors are logged and never fatal.
+   */
+  onSample: (sample: ABMetricSample) => void | Promise<void>;
+}
+
+export type ABDefinition = ABConfig & { readonly __agentServe: "ab" };
+
+/** Author a live A/B metrics experiment under `agent/ab`. */
+export function defineAB(config: ABConfig): ABDefinition {
+  const keys = Object.keys(config.variants ?? {});
+  if (keys.length < 2) {
+    throw new Error(
+      `defineAB: variants needs at least two arms (got ${keys.length})`
+    );
+  }
+  for (const key of keys) {
+    if (key.length === 0 || key.includes("/") || key.includes("\\")) {
+      throw new Error(
+        `defineAB: invalid variant key ${JSON.stringify(key)} (single segment)`
+      );
+    }
+  }
+  if (typeof config.split !== "function") {
+    throw new Error(
+      "defineAB: split must be a function (splitBySessionHash(), splitIf(...), or custom)"
+    );
+  }
+  if (typeof config.onSample !== "function") {
+    throw new Error("defineAB: onSample(sample) is required");
+  }
+  if (config.derive !== undefined) {
+    for (const [key, fn] of Object.entries(config.derive)) {
+      if (typeof fn !== "function") {
+        throw new Error(`defineAB: derive.${key} must be a function`);
+      }
+    }
+  }
+  return brandDefinition("ab", config);
+}
+
+/** Resolve {@link ABConfig.split} to a variant key or `null`. */
+export async function resolveABSplit(
+  split: ABSplit,
+  ctx: ABSplitContext
+): Promise<ABSplitResult> {
+  const result = await split(ctx);
+  if (result === null || result === undefined) {
+    return null;
+  }
+  if (typeof result === "string") {
+    return result;
+  }
+  throw new Error(
+    `defineAB: split must return a variant key or null (got ${typeof result})`
+  );
+}
+
+function assertHoldout(holdout: number | undefined): void {
+  if (holdout === undefined) {
+    return;
+  }
+  if (!Number.isFinite(holdout) || holdout < 0 || holdout > 1) {
+    throw new Error("defineAB: holdout must be a finite number in [0, 1]");
+  }
+}
+
+/** Pick a variant (or `null` via holdout) from a unit interval in `[0, 1)`. */
+function pickArm(
+  variantKeys: readonly string[],
+  weights: Record<string, number> | undefined,
+  unitInterval: number,
+  holdout: number | undefined
+): ABSplitResult {
+  const u = clampUnit(unitInterval);
+  const out = holdout ?? 0;
+  if (out >= 1 || (out > 0 && u < out)) {
+    return null;
+  }
+  const remapped = out > 0 ? (u - out) / (1 - out) : u;
+  return pickWeightedVariant(variantKeys, weights, remapped);
+}
+
+/** FNV-1a → [0, 1). */
+export function hashToUnitInterval(input: string): number {
+  let h = 2166136261;
+  for (let i = 0; i < input.length; i++) {
+    h ^= input.charCodeAt(i);
+    h = Math.imul(h, 16777619);
+  }
+  return (h >>> 0) / 0x1_0000_0000;
+}
+
+export function pickWeightedVariant(
+  variantKeys: readonly string[],
+  weights: Record<string, number> | undefined,
+  unitInterval: number
+): string {
+  const u = clampUnit(unitInterval);
+  const rows: Array<{ key: string; weight: number }> = [];
+  let total = 0;
+  for (const key of variantKeys) {
+    const weight = weights?.[key] ?? 1;
+    if (weight > 0 && Number.isFinite(weight)) {
+      rows.push({ key, weight });
+      total += weight;
+    }
+  }
+  if (rows.length === 0 || total <= 0) {
+    throw new Error("defineAB: no positive-weight variants to split across");
+  }
+  let cursor = u * total;
+  for (const row of rows) {
+    cursor -= row.weight;
+    if (cursor < 0) {
+      return row.key;
+    }
+  }
+  return rows[rows.length - 1]!.key;
+}
+
+function clampUnit(unitInterval: number): number {
+  return Math.min(Math.max(unitInterval, 0), 0.999999999);
+}
+
+export function emptyABMetrics(): ABMetrics {
+  return {
+    turns: 0,
+    turnFailures: 0,
+    toolCalls: 0,
+    toolErrors: 0,
+    inputTokens: 0,
+    outputTokens: 0,
+    cacheReadTokens: 0,
+    cacheWriteTokens: 0,
+    wallTimeMs: 0,
+    custom: {},
+  };
+}
+
+export function cloneABMetrics(metrics: ABMetrics): ABMetrics {
+  return { ...metrics, custom: { ...metrics.custom } };
+}
+
+/**
+ * Append enrolled per-variant `instructions` overlays onto the base agent
+ * prompt. Used by the session engine before materializing the harness.
+ */
+export function mergeABInstructions(
+  base: string | undefined,
+  abs: Readonly<Record<string, string | null>> | undefined,
+  experiments: ReadonlyArray<{ name: string; definition: ABDefinition }>
+): string | undefined {
+  if (abs === undefined) {
+    return base;
+  }
+  const extras: string[] = [];
+  for (const experiment of experiments) {
+    const variant = abs[experiment.name];
+    if (variant === undefined || variant === null) {
+      continue;
+    }
+    const extra = experiment.definition.variants[variant]?.instructions;
+    if (typeof extra === "string" && extra.trim() !== "") {
+      extras.push(extra.trim());
+    }
+  }
+  if (extras.length === 0) {
+    return base;
+  }
+  const baseText = base?.trim() ?? "";
+  return baseText === ""
+    ? extras.join("\n\n")
+    : `${baseText}\n\n${extras.join("\n\n")}`;
+}
+
+// ---------------------------------------------------------------------------
+// Project config (`agent/ab.config.ts`) — fold window + optional persistence
+// ---------------------------------------------------------------------------
+
+/** Default sessions folded into `GET /v1/abs` / the A/Bs tab. */
+export const AB_DEFAULT_MAX_PLAYGROUND_SESSIONS = 200;
+
+/** Hard ceiling for {@link ABConfigFile.maxPlaygroundSessions}. */
+export const AB_MAX_PLAYGROUND_SESSIONS = 10_000;
+
+/** Default snapshot files kept when {@link ABConfigFile.persistSnapshots} is set. */
+export const AB_DEFAULT_RETAIN_SNAPSHOTS = 20;
+
+/** Hard ceiling for {@link ABConfigFile.retainSnapshots}. */
+export const AB_MAX_RETAIN_SNAPSHOTS = 500;
+
+/** Context passed to A/B persistence hooks. */
+export interface ABPersistContext {
+  /** Absolute agent project root (directory that contains `agent/`). */
+  projectRoot: string;
+}
+
+/**
+ * Durable sink for {@link ABMetricSample} rows (same payload as
+ * {@link ABConfig.onSample}). Session `events.ndjson` remains the source of
+ * truth for assignment + fold; this is for analytics export.
+ */
+export interface ABSamplePersistence {
+  save(sample: ABMetricSample, ctx: ABPersistContext): void | Promise<void>;
+}
+
+/**
+ * Durable store for playground-style aggregate snapshots (`GET /v1/abs`).
+ * Without this, aggregates are always recomputed from session logs (nothing
+ * is kept after sessions are pruned).
+ *
+ * **Runtime note:** `serve` currently only calls {@link save} (after each
+ * `GET /v1/abs`). {@link load} / {@link delete} are for author tooling or a
+ * custom store; they are not wired into the playground hydrate path.
+ */
+export interface ABSnapshotPersistence {
+  load(ctx: ABPersistContext): ABSnapshot[] | Promise<ABSnapshot[]>;
+  save(snapshot: ABSnapshot, ctx: ABPersistContext): void | Promise<void>;
+  delete?(snapshotId: string, ctx: ABPersistContext): void | Promise<void>;
+}
+
+/** Lets {@link defineABConfig} push `retainSnapshots` into {@link persistABSnapshotsToDir}. */
+const abSnapshotRetainState = new WeakMap<
+  ABSnapshotPersistence,
+  { retain: number }
+>();
+
+/**
+ * Project-wide A/B defaults (`agent/ab.config.ts`).
+ *
+ * **Assignments and foldable metrics** already persist in each session's
+ * `events.ndjson` under `--state-root`. Use this config to:
+ * - cap how many sessions the playground / `GET /v1/abs` folds
+ * - export samples and/or aggregate snapshots to author-chosen storage
+ */
+export interface ABConfigFile {
+  /**
+   * Optional. Max newest sessions folded into **`GET /v1/abs` / the A/Bs
+   * tab** (and into {@link persistSnapshots} payloads). Does not prune
+   * session logs, change assignment, or limit {@link persistSamples}.
+   *
+   * Default {@link AB_DEFAULT_MAX_PLAYGROUND_SESSIONS}; hard-capped at
+   * {@link AB_MAX_PLAYGROUND_SESSIONS}.
+   */
+  maxPlaygroundSessions?: number;
+  /**
+   * Snapshot-file budget when {@link persistSnapshots} is set.
+   * Default {@link AB_DEFAULT_RETAIN_SNAPSHOTS}; hard-capped at
+   * {@link AB_MAX_RETAIN_SNAPSHOTS}.
+   *
+   * When using {@link persistABSnapshotsToDir}, this value controls on-disk
+   * pruning (overriding the helper's `{ retain }` default). Custom stores
+   * should apply the same budget themselves.
+   */
+  retainSnapshots?: number;
+  /**
+   * Optional sink for cumulative metric samples (in addition to each
+   * experiment's {@link ABConfig.onSample}).
+   *
+   * ```ts
+   * persistSamples: persistABSamplesToDir(".agent-serve/ab-samples"),
+   * ```
+   */
+  persistSamples?: ABSamplePersistence;
+  /**
+   * Optional store for aggregate snapshots written when `GET /v1/abs` runs.
+   *
+   * ```ts
+   * persistSnapshots: persistABSnapshotsToDir(".agent-serve/ab-snapshots"),
+   * ```
+   */
+  persistSnapshots?: ABSnapshotPersistence;
+}
+
+/**
+ * Defaults for every experiment under `agent/ab*` (`agent/ab.config.ts`).
+ *
+ * Session event logs remain the assignment/metrics source of truth; this
+ * file controls the playground fold window and optional durable exports.
+ */
+export function defineABConfig(config: ABConfigFile): ABConfigFile {
+  const persistSamples = normalizeSamplePersistence(config.persistSamples);
+  const persistSnapshots = normalizeSnapshotPersistence(
+    config.persistSnapshots
+  );
+  const helperRetain =
+    persistSnapshots !== undefined
+      ? abSnapshotRetainState.get(persistSnapshots)
+      : undefined;
+  let retainSnapshots: number | undefined;
+  if (config.retainSnapshots !== undefined) {
+    retainSnapshots = normalizeABRetainSnapshots(config.retainSnapshots);
+    if (helperRetain !== undefined) {
+      helperRetain.retain = retainSnapshots;
+    }
+  } else if (persistSnapshots !== undefined) {
+    retainSnapshots =
+      helperRetain?.retain ??
+      normalizeABRetainSnapshots(AB_DEFAULT_RETAIN_SNAPSHOTS);
+  }
+  return {
+    maxPlaygroundSessions: normalizeABMaxPlaygroundSessions(
+      config.maxPlaygroundSessions ?? AB_DEFAULT_MAX_PLAYGROUND_SESSIONS
+    ),
+    ...(retainSnapshots === undefined ? {} : { retainSnapshots }),
+    ...(persistSamples === undefined ? {} : { persistSamples }),
+    ...(persistSnapshots === undefined ? {} : { persistSnapshots }),
+  };
+}
+
+/** Validate {@link ABConfigFile.maxPlaygroundSessions}. */
+export function normalizeABMaxPlaygroundSessions(value: unknown): number {
+  return normalizePositiveInt(value, {
+    label: "ab.config.ts: maxPlaygroundSessions",
+    max: AB_MAX_PLAYGROUND_SESSIONS,
+  });
+}
+
+/**
+ * Effective `GET /v1/abs` fold window. Invalid values fall back to
+ * {@link AB_DEFAULT_MAX_PLAYGROUND_SESSIONS}.
+ */
+export function resolveABMaxPlaygroundSessions(value: unknown): number {
+  if (value === undefined) {
+    return AB_DEFAULT_MAX_PLAYGROUND_SESSIONS;
+  }
+  try {
+    return normalizeABMaxPlaygroundSessions(value);
+  } catch {
+    return AB_DEFAULT_MAX_PLAYGROUND_SESSIONS;
+  }
+}
+
+/** Validate {@link ABConfigFile.retainSnapshots}. */
+export function normalizeABRetainSnapshots(value: unknown): number {
+  return normalizePositiveInt(value, {
+    label: "ab.config.ts: retainSnapshots",
+    max: AB_MAX_RETAIN_SNAPSHOTS,
+  });
+}
+
+function normalizeSamplePersistence(
+  value: ABSamplePersistence | undefined
+): ABSamplePersistence | undefined {
+  if (value === undefined) {
+    return undefined;
+  }
+  if (
+    typeof value !== "object" ||
+    value === null ||
+    typeof value.save !== "function"
+  ) {
+    throw new Error(
+      "ab.config.ts: persistSamples must define save(sample, ctx)"
+    );
+  }
+  return value;
+}
+
+function normalizeSnapshotPersistence(
+  value: ABSnapshotPersistence | undefined
+): ABSnapshotPersistence | undefined {
+  if (value === undefined) {
+    return undefined;
+  }
+  if (
+    typeof value !== "object" ||
+    value === null ||
+    typeof value.load !== "function" ||
+    typeof value.save !== "function"
+  ) {
+    throw new Error(
+      "ab.config.ts: persistSnapshots must define load(ctx) and save(snapshot, ctx)"
+    );
+  }
+  if (value.delete !== undefined && typeof value.delete !== "function") {
+    throw new Error("ab.config.ts: persistSnapshots.delete must be a function");
+  }
+  return value;
+}
+
+/**
+ * Append-only NDJSON {@link ABSamplePersistence} under `dir` (relative to
+ * the agent project root).
+ *
+ * ```ts
+ * // agent/ab.config.ts
+ * export default defineABConfig({
+ *   persistSamples: persistABSamplesToDir(".agent-serve/ab-samples"),
+ * });
+ * ```
+ */
+export function persistABSamplesToDir(dir: string): ABSamplePersistence {
+  const relativeDir = assertNonEmptyDir(dir, "persistABSamplesToDir");
+  const resolveDir = (projectRoot: string): string =>
+    resolveProjectRelativeDir({ dir: relativeDir, projectRoot });
+
+  return {
+    async save(sample, ctx) {
+      const root = resolveDir(ctx.projectRoot);
+      await mkdir(root, { recursive: true });
+      await appendFile(
+        join(root, "samples.ndjson"),
+        `${JSON.stringify(sample)}\n`,
+        "utf8"
+      );
+    },
+  };
+}
+
+/**
+ * Filesystem {@link ABSnapshotPersistence}: one JSON file per snapshot under
+ * `dir`, pruned to `retain` files (default {@link AB_DEFAULT_RETAIN_SNAPSHOTS}).
+ *
+ * Prefer setting {@link ABConfigFile.retainSnapshots} in `defineABConfig` —
+ * that value is applied to this helper's prune budget. Pass `{ retain }` here
+ * only when constructing the helper outside `defineABConfig`, or as a default
+ * that `retainSnapshots` may override.
+ */
+export function persistABSnapshotsToDir(
+  dir: string,
+  options?: { retain?: number }
+): ABSnapshotPersistence {
+  const relativeDir = assertNonEmptyDir(dir, "persistABSnapshotsToDir");
+  const retainState = {
+    retain: normalizeABRetainSnapshots(
+      options?.retain ?? AB_DEFAULT_RETAIN_SNAPSHOTS
+    ),
+  };
+  const resolveDir = (projectRoot: string): string =>
+    resolveProjectRelativeDir({ dir: relativeDir, projectRoot });
+  let saveChain: Promise<void> = Promise.resolve();
+
+  const persist: ABSnapshotPersistence = {
+    async load(ctx) {
+      return readJsonDir<ABSnapshot>(resolveDir(ctx.projectRoot), {
+        sortBasenames: "desc",
+      });
+    },
+    async save(snapshot, ctx) {
+      const run = async (): Promise<void> => {
+        const root = resolveDir(ctx.projectRoot);
+        const snapshotId =
+          snapshot.config?.snapshotId ??
+          `absnap_${randomUUID().replace(/-/g, "")}`;
+        const toWrite: ABSnapshot = {
+          ...snapshot,
+          config: {
+            ...snapshot.config,
+            snapshotId,
+          },
+        };
+        await writeJsonFile({
+          root,
+          fileName: `${snapshotId}.json`,
+          value: toWrite,
+        });
+        const names = (await listJsonBasenames(root)).sort().reverse();
+        for (const stale of names.slice(retainState.retain)) {
+          await deleteFileInDir({ root, fileName: stale });
+        }
+      };
+      const next = saveChain.then(run, run);
+      saveChain = next.then(
+        () => undefined,
+        () => undefined
+      );
+      await next;
+    },
+    async delete(snapshotId, ctx) {
+      await deleteFileInDir({
+        root: resolveDir(ctx.projectRoot),
+        fileName: `${snapshotId}.json`,
+      });
+    },
+  };
+  abSnapshotRetainState.set(persist, retainState);
+  return persist;
+}