diff --git a/schemas/understudy.eval-cohort-strata-plan.v1.schema.json b/schemas/understudy.eval-cohort-strata-plan.v1.schema.json new file mode 100644 index 00000000..da505bbe --- /dev/null +++ b/schemas/understudy.eval-cohort-strata-plan.v1.schema.json @@ -0,0 +1,180 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://understudylabs.com/schemas/understudy.eval-cohort-strata-plan.v1.schema.json", + "title": "understudy.eval-cohort-strata-plan.v1", + "description": "Local, metadata-only stratified cohort plan derived from materialized trace artifacts (issue #281 phase 1). Contains execution ids, stratum labels, and counts only — never payload-derived fields.", + "type": "object", + "additionalProperties": false, + "required": [ + "schema_version", "created_at", "seed", "inputs", "config", "pool", + "strata", "selection", "coverage", "ready", "blocking", "expansion_of", + "stability", "privacy" + ], + "properties": { + "schema_version": { "const": "understudy.eval-cohort-strata-plan.v1" }, + "created_at": { "$ref": "#/$defs/timestamp" }, + "seed": { "$ref": "#/$defs/nonempty" }, + "inputs": { + "type": "object", + "additionalProperties": false, + "required": ["execution_index", "execution_index_sha256", "tasks"], + "properties": { + "execution_index": { "$ref": "#/$defs/nonempty" }, + "execution_index_sha256": { "$ref": "#/$defs/sha" }, + "tasks": { "oneOf": [{ "$ref": "#/$defs/nonempty" }, { "type": "null" }] } + } + }, + "config": { + "type": "object", + "additionalProperties": false, + "required": ["axes", "target_per_stratum", "rare_threshold", "high_consequence_tags"], + "properties": { + "axes": { + "type": "array", + "minItems": 1, + "uniqueItems": true, + "items": { "enum": ["outcome", "mode", "turns", "confidence", "tag"] } + }, + "target_per_stratum": { "type": "integer", "minimum": 1 }, + "rare_threshold": { "type": "integer", "minimum": 0 }, + "high_consequence_tags": { "type": "array", "items": { "$ref": "#/$defs/nonempty" } } + } + }, + "pool": { + "type": "object", + "additionalProperties": false, + "required": ["index_rows", "eligible", "excluded"], + "properties": { + "index_rows": { "$ref": "#/$defs/count" }, + "eligible": { "$ref": "#/$defs/count" }, + "excluded": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["reason", "count"], + "properties": { + "reason": { + "enum": [ + "source_excluded", "missing_execution_group", "lineage_ambiguous", + "lineage_unlinked", "missing_task", "duplicate_execution_group" + ] + }, + "count": { "$ref": "#/$defs/count" } + } + } + } + } + }, + "strata": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["axis", "value", "available", "target", "selected", "rule"], + "properties": { + "axis": { "enum": ["outcome", "mode", "turns", "confidence", "tag"] }, + "value": { "$ref": "#/$defs/nonempty" }, + "available": { "$ref": "#/$defs/count" }, + "target": { "type": "integer", "minimum": 1 }, + "selected": { "$ref": "#/$defs/count" }, + "rule": { + "enum": [ + "saturated_rare", "saturated_high_consequence", "saturated", + "sampled", "uncovered" + ] + } + } + } + }, + "selection": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["execution_group", "task_id", "strata", "frozen"], + "properties": { + "execution_group": { "$ref": "#/$defs/nonempty" }, + "task_id": { "$ref": "#/$defs/nonempty" }, + "strata": { "type": "array", "items": { "$ref": "#/$defs/nonempty" } }, + "frozen": { "type": "boolean" } + } + } + }, + "coverage": { + "type": "object", + "additionalProperties": false, + "required": ["uncovered_strata", "underfilled_strata"], + "properties": { + "uncovered_strata": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["axis", "value"], + "properties": { + "axis": { "enum": ["outcome", "mode", "turns", "confidence", "tag"] }, + "value": { "$ref": "#/$defs/nonempty" } + } + } + }, + "underfilled_strata": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["axis", "value", "available", "target"], + "properties": { + "axis": { "enum": ["outcome", "mode", "turns", "confidence", "tag"] }, + "value": { "$ref": "#/$defs/nonempty" }, + "available": { "$ref": "#/$defs/count" }, + "target": { "type": "integer", "minimum": 1 } + } + } + } + } + }, + "ready": { "type": "boolean" }, + "blocking": { "type": "array", "items": { "$ref": "#/$defs/nonempty" } }, + "expansion_of": { "oneOf": [{ "$ref": "#/$defs/sha" }, { "type": "null" }] }, + "stability": { + "type": "object", + "additionalProperties": false, + "required": ["status", "tolerance", "moved"], + "properties": { + "status": { "enum": ["not_applicable", "stable", "unstable"] }, + "tolerance": { "type": "number", "minimum": 0 }, + "moved": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["axis", "value", "before", "after"], + "properties": { + "axis": { "enum": ["outcome", "mode", "turns", "confidence", "tag"] }, + "value": { "$ref": "#/$defs/nonempty" }, + "before": { "type": "number" }, + "after": { "type": "number" } + } + } + } + } + }, + "privacy": { + "type": "object", + "additionalProperties": false, + "required": ["local_only", "payload_fields_read", "upload_performed"], + "properties": { + "local_only": { "const": true }, + "payload_fields_read": { "const": false }, + "upload_performed": { "const": false } + } + } + }, + "$defs": { + "nonempty": { "type": "string", "minLength": 1 }, + "sha": { "type": "string", "pattern": "^[a-f0-9]{64}$" }, + "count": { "type": "integer", "minimum": 0 }, + "timestamp": { "type": "string", "format": "date-time" } + } +} diff --git a/src/evals/cohort-strata.ts b/src/evals/cohort-strata.ts new file mode 100644 index 00000000..2f71c6dd --- /dev/null +++ b/src/evals/cohort-strata.ts @@ -0,0 +1,643 @@ +/** + * Execution-derived eval cohort strata (issue #281, phase 1). + * + * The hosted catalog flow can only filter on request-level metadata (model, + * status, tool presence, structured output), so a random recent sample can + * miss rare or high-consequence behaviors whose complexity is only visible + * after the execution finishes. This module derives configurable strata from + * the LOCAL materialized trace artifacts (`execution-index.jsonl` + + * `tasks.jsonl` produced by the trace foundry) and selects a deterministic, + * stratified cohort toward decision-sized per-stratum targets. + * + * Hard privacy contract: the module only ever reads metadata fields + * (execution ids, lineage status, capture counts, tool names, user-supplied + * tags). Payload-derived fields such as task titles, prompts, or completions + * are never read, and the plan artifact contains ids, labels, and counts + * only — it is safe to discuss in a review without redaction. + */ +import { createHash } from "node:crypto"; +import { readFileSync } from "node:fs"; +import { z } from "zod"; + +import { isMutatingTool } from "../trace-foundry.js"; + +export const COHORT_STRATA_PLAN_SCHEMA_VERSION = "understudy.eval-cohort-strata-plan.v1"; +export const DEFAULT_STRATA_SEED = "understudy-cohort-strata-v1"; + +export const STRATA_AXES = ["outcome", "mode", "turns", "confidence", "tag"] as const; +export type StrataAxis = (typeof STRATA_AXES)[number]; + +const ExecutionIndexRowSchema = z + .object({ + schema_version: z.literal("understudy.eval-execution-index-row.v1"), + source_status: z.enum(["included", "excluded"]).optional(), + execution_group: z.string().min(1).nullable(), + lineage_status: z.enum(["complete", "ambiguous", "unlinked"]).nullable(), + capture_count: z.number().int().nonnegative(), + task_id: z.string().min(1).nullable(), + exclusion_reasons: z.array(z.string()).optional(), + }) + .passthrough(); + +/** Metadata-only task fields. `title` and other payload-derived fields are intentionally not read. */ +const TaskRowSchema = z + .object({ + task_id: z.string().min(1), + execution_group: z.string().min(1), + tool_surface: z.array(z.string()).optional(), + machine_confidence: z.enum(["high", "medium", "low"]).optional(), + }) + .passthrough(); + +export interface StrataPlanOptions { + /** Deterministic selection seed. Same seed + same pool ⇒ byte-identical plan. */ + seed?: string; + /** Per-stratum selection target. Pilot sizes are minimums, never caps. */ + targetPerStratum?: number; + /** Strata with at most this many available executions are saturated (all included). */ + rareThreshold?: number; + /** Axes to derive, in greedy priority order. Default: outcome, mode, turns, confidence. */ + axes?: StrataAxis[]; + /** Human-guided tags per execution group. Tags are never inferred. */ + tags?: Record; + /** Tag values whose entire stratum is always saturated (high-consequence cases). */ + highConsequenceTags?: string[]; + now?: Date; +} + +export interface PoolExclusion { + reason: + | "source_excluded" + | "missing_execution_group" + | "lineage_ambiguous" + | "lineage_unlinked" + | "missing_task" + | "duplicate_execution_group"; + count: number; +} + +export interface StratumReport { + axis: StrataAxis; + value: string; + available: number; + target: number; + selected: number; + rule: "saturated_rare" | "saturated_high_consequence" | "saturated" | "sampled" | "uncovered"; +} + +export interface StrataPlan { + schema_version: typeof COHORT_STRATA_PLAN_SCHEMA_VERSION; + created_at: string; + seed: string; + inputs: { + execution_index: string; + execution_index_sha256: string; + tasks: string | null; + }; + config: { + axes: StrataAxis[]; + target_per_stratum: number; + rare_threshold: number; + high_consequence_tags: string[]; + }; + pool: { + index_rows: number; + eligible: number; + excluded: PoolExclusion[]; + }; + strata: StratumReport[]; + selection: Array<{ + execution_group: string; + task_id: string; + strata: string[]; + frozen: boolean; + }>; + coverage: { + uncovered_strata: Array<{ axis: StrataAxis; value: string }>; + underfilled_strata: Array<{ axis: StrataAxis; value: string; available: number; target: number }>; + }; + /** + * Phase-1 readiness signal: false when any underfilled high-consequence or + * rare stratum remains, so later gates can block a whole-workload-ready + * status on machine-readable evidence. + */ + ready: boolean; + blocking: string[]; + expansion_of: string | null; + stability: { + status: "not_applicable" | "stable" | "unstable"; + tolerance: number; + moved: Array<{ axis: StrataAxis; value: string; before: number; after: number }>; + }; + privacy: { + local_only: true; + payload_fields_read: false; + upload_performed: false; + }; +} + +export const CohortStrataPlanSchema = z + .object({ + schema_version: z.literal(COHORT_STRATA_PLAN_SCHEMA_VERSION), + created_at: z.string().datetime(), + seed: z.string().min(1), + inputs: z.object({ + execution_index: z.string().min(1), + execution_index_sha256: z.string().regex(/^[a-f0-9]{64}$/), + tasks: z.string().min(1).nullable(), + }), + config: z.object({ + axes: z.array(z.enum(STRATA_AXES)).min(1), + target_per_stratum: z.number().int().positive(), + rare_threshold: z.number().int().nonnegative(), + high_consequence_tags: z.array(z.string()), + }), + pool: z.object({ + index_rows: z.number().int().nonnegative(), + eligible: z.number().int().nonnegative(), + excluded: z.array( + z.object({ + reason: z.enum([ + "source_excluded", + "missing_execution_group", + "lineage_ambiguous", + "lineage_unlinked", + "missing_task", + "duplicate_execution_group", + ]), + count: z.number().int().nonnegative(), + }), + ), + }), + strata: z.array( + z.object({ + axis: z.enum(STRATA_AXES), + value: z.string().min(1), + available: z.number().int().nonnegative(), + target: z.number().int().positive(), + selected: z.number().int().nonnegative(), + rule: z.enum([ + "saturated_rare", + "saturated_high_consequence", + "saturated", + "sampled", + "uncovered", + ]), + }), + ), + selection: z.array( + z.object({ + execution_group: z.string().min(1), + task_id: z.string().min(1), + strata: z.array(z.string().min(1)), + frozen: z.boolean(), + }), + ), + coverage: z.object({ + uncovered_strata: z.array(z.object({ axis: z.enum(STRATA_AXES), value: z.string().min(1) })), + underfilled_strata: z.array( + z.object({ + axis: z.enum(STRATA_AXES), + value: z.string().min(1), + available: z.number().int().nonnegative(), + target: z.number().int().positive(), + }), + ), + }), + ready: z.boolean(), + blocking: z.array(z.string()), + expansion_of: z.string().regex(/^[a-f0-9]{64}$/).nullable(), + stability: z.object({ + status: z.enum(["not_applicable", "stable", "unstable"]), + tolerance: z.number().nonnegative(), + moved: z.array( + z.object({ + axis: z.enum(STRATA_AXES), + value: z.string().min(1), + before: z.number(), + after: z.number(), + }), + ), + }), + privacy: z.object({ + local_only: z.literal(true), + payload_fields_read: z.literal(false), + upload_performed: z.literal(false), + }), + }) + .strict(); + +interface EligibleExecution { + execution_group: string; + task_id: string; + lineage_status: "complete"; + capture_count: number; + tool_surface: string[]; + machine_confidence: "high" | "medium" | "low" | null; + tags: string[]; +} + +const sha256 = (value: string) => createHash("sha256").update(value).digest("hex"); + +function readJsonlRows(path: string): unknown[] { + const text = readFileSync(path, "utf8"); + return text + .split(/\r?\n/) + .filter((line) => line.trim().length > 0) + .map((line, index) => { + try { + return JSON.parse(line); + } catch (error) { + throw new Error(`Invalid JSONL at ${path}:${index + 1}: ${error instanceof Error ? error.message : String(error)}`); + } + }); +} + +/** + * Parse the materialized trace artifacts into an eligible pool plus + * exclusion accounting. Metadata-only by construction. + */ +export function buildEligiblePool( + executionIndexPath: string, + tasksPath: string | null, + tags: Record, +): { eligible: EligibleExecution[]; excluded: PoolExclusion[]; indexRowCount: number } { + const indexRows = readJsonlRows(executionIndexPath).map((row, index) => { + const parsed = ExecutionIndexRowSchema.safeParse(row); + if (!parsed.success) { + throw new Error(`Invalid execution-index row ${index + 1} in ${executionIndexPath}.`); + } + return parsed.data; + }); + + const tasksByGroup = new Map>(); + if (tasksPath !== null) { + for (const [index, raw] of readJsonlRows(tasksPath).entries()) { + const parsed = TaskRowSchema.safeParse(raw); + if (!parsed.success) { + throw new Error(`Invalid tasks row ${index + 1} in ${tasksPath}.`); + } + tasksByGroup.set(parsed.data.execution_group, parsed.data); + } + } + + const excludedCounts = new Map(); + const bump = (reason: PoolExclusion["reason"]) => excludedCounts.set(reason, (excludedCounts.get(reason) ?? 0) + 1); + + const seen = new Set(); + const eligible: EligibleExecution[] = []; + for (const row of indexRows) { + if (row.source_status === "excluded") { + bump("source_excluded"); + continue; + } + if (!row.execution_group) { + bump("missing_execution_group"); + continue; + } + if (row.lineage_status === "ambiguous") { + bump("lineage_ambiguous"); + continue; + } + if (row.lineage_status === "unlinked" || row.lineage_status === null) { + bump("lineage_unlinked"); + continue; + } + if (!row.task_id) { + bump("missing_task"); + continue; + } + if (seen.has(row.execution_group)) { + bump("duplicate_execution_group"); + continue; + } + seen.add(row.execution_group); + const task = tasksByGroup.get(row.execution_group) ?? null; + eligible.push({ + execution_group: row.execution_group, + task_id: row.task_id, + lineage_status: "complete", + capture_count: row.capture_count, + tool_surface: task?.tool_surface ?? [], + machine_confidence: task?.machine_confidence ?? null, + tags: [...(tags[row.execution_group] ?? [])].sort(), + }); + } + + const excluded: PoolExclusion[] = [...excludedCounts.entries()] + .map(([reason, count]) => ({ reason, count })) + .sort((left, right) => left.reason.localeCompare(right.reason)); + return { eligible, excluded, indexRowCount: indexRows.length }; +} + +function turnsBucket(captureCount: number): string { + if (captureCount <= 1) return "single"; + if (captureCount <= 5) return "short"; + return "long"; +} + +function modeBucket(toolSurface: string[]): string { + if (toolSurface.length === 0) return "no-tools"; + return toolSurface.some((name) => isMutatingTool(name)) ? "write" : "read-only"; +} + +/** + * Derive every stratum label for one execution. Labels are `{axis}:{value}` + * strings; an execution can belong to several strata (one per axis, plus one + * per user-supplied tag). + */ +export function deriveStrata(execution: EligibleExecution, axes: StrataAxis[]): string[] { + const labels = new Set(); + for (const axis of axes) { + switch (axis) { + case "outcome": + labels.add(`outcome:${execution.lineage_status}`); + break; + case "mode": + labels.add(`mode:${modeBucket(execution.tool_surface)}`); + break; + case "turns": + labels.add(`turns:${turnsBucket(execution.capture_count)}`); + break; + case "confidence": + labels.add(`confidence:${execution.machine_confidence ?? "unknown"}`); + break; + case "tag": + for (const tag of execution.tags) labels.add(`tag:${tag}`); + break; + } + } + return [...labels].sort(); +} + +function selectionPriority(seed: string, executionGroup: string): string { + return sha256(`${seed}|${executionGroup}`); +} + +interface StratumMembers { + label: string; + axis: StrataAxis; + value: string; + members: EligibleExecution[]; + rare: boolean; + highConsequence: boolean; +} + +function buildStrata( + eligible: EligibleExecution[], + axes: StrataAxis[], + rareThreshold: number, + highConsequenceTags: string[], +): StratumMembers[] { + const byLabel = new Map(); + for (const execution of eligible) { + for (const label of deriveStrata(execution, axes)) { + byLabel.set(label, [...(byLabel.get(label) ?? []), execution]); + } + } + const highConsequence = new Set(highConsequenceTags); + return [...byLabel.entries()] + .map(([label, members]) => { + const [axis, ...rest] = label.split(":"); + const value = rest.join(":"); + return { + label, + axis: axis as StrataAxis, + value, + members, + rare: members.length <= rareThreshold, + highConsequence: axis === "tag" && highConsequence.has(value), + }; + }) + .sort((left, right) => left.label.localeCompare(right.label)); +} + +/** + * Deterministic stratified selection. + * + * Pilot sizes are minimums, never caps: every stratum that is rare (at or + * below `rareThreshold` available executions) or tagged high-consequence is + * saturated — all of its members are selected. Remaining strata are filled in + * axis priority order with a seeded, stable ordering until the per-stratum + * target is met, counting members already selected for an earlier stratum. + * + * When `frozenGroups` is provided (incremental expansion), every surviving + * frozen execution stays selected and is never displaced; the new batch only + * adds. Frozen-split boundaries are therefore unchanged by expansion. + */ +export function planStratifiedCohort(options: { + eligible: EligibleExecution[]; + axes: StrataAxis[]; + seed: string; + targetPerStratum: number; + rareThreshold: number; + highConsequenceTags: string[]; + frozenGroups?: Set; +}): { selected: EligibleExecution[]; reports: StratumReport[] } { + const { eligible, axes, seed, targetPerStratum, rareThreshold, highConsequenceTags } = options; + const frozenGroups = options.frozenGroups ?? new Set(); + const strata = buildStrata(eligible, axes, rareThreshold, highConsequenceTags); + + const selected = new Map(); + const eligibleByGroup = new Map(eligible.map((execution) => [execution.execution_group, execution])); + + // 1. Frozen members always survive expansion. + for (const group of [...frozenGroups].sort()) { + const execution = eligibleByGroup.get(group); + if (execution) selected.set(group, execution); + } + + // 2. Saturate rare and high-consequence strata (minimums, never caps). + for (const stratum of strata) { + if (!stratum.rare && !stratum.highConsequence) continue; + for (const member of stratum.members) selected.set(member.execution_group, member); + } + + // 3. Greedily fill per-stratum targets in axis priority order. + for (const axis of axes) { + for (const stratum of strata.filter((candidate) => candidate.axis === axis)) { + const have = stratum.members.filter((member) => selected.has(member.execution_group)).length; + if (have >= targetPerStratum) continue; + const candidates = stratum.members + .filter((member) => !selected.has(member.execution_group)) + .sort((left, right) => selectionPriority(seed, left.execution_group).localeCompare(selectionPriority(seed, right.execution_group)) || left.execution_group.localeCompare(right.execution_group)); + for (const candidate of candidates.slice(0, targetPerStratum - have)) { + selected.set(candidate.execution_group, candidate); + } + } + } + + const reports: StratumReport[] = strata.map((stratum) => { + const selectedCount = stratum.members.filter((member) => selected.has(member.execution_group)).length; + let rule: StratumReport["rule"]; + if (stratum.members.length === 0) rule = "uncovered"; + else if (stratum.highConsequence) rule = "saturated_high_consequence"; + else if (stratum.rare && selectedCount === stratum.members.length) rule = "saturated_rare"; + else if (selectedCount === stratum.members.length && stratum.members.length < targetPerStratum) rule = "saturated"; + else rule = "sampled"; + return { + axis: stratum.axis, + value: stratum.value, + available: stratum.members.length, + target: targetPerStratum, + selected: selectedCount, + rule, + }; + }); + + const ordered = [...selected.values()].sort((left, right) => + selectionPriority(seed, left.execution_group).localeCompare(selectionPriority(seed, right.execution_group)) || + left.execution_group.localeCompare(right.execution_group), + ); + return { selected: ordered, reports }; +} + +function prevalenceShare(available: number, eligible: number): number { + if (eligible === 0) return 0; + return available / eligible; +} + +/** SHA-256 identity of a frozen plan, used as `expansion_of` binding. */ +export function planIdentity(plan: StrataPlan): string { + const frozenView = { + schema_version: plan.schema_version, + seed: plan.seed, + inputs: plan.inputs, + config: plan.config, + selection: plan.selection.map((entry) => entry.execution_group).sort(), + }; + return sha256(JSON.stringify(frozenView)); +} + +/** + * Build the full plan artifact. Pass `priorPlan` to expand a frozen plan + * incrementally: prior selections survive untouched and stability is reported + * against the declared tolerance. + */ +export function buildStrataPlan(options: { + executionIndexPath: string; + tasksPath?: string | null; + planOptions?: StrataPlanOptions; + priorPlan?: StrataPlan | null; + tolerance?: number; +}): StrataPlan { + const { executionIndexPath, priorPlan = null } = options; + const planOptions = options.planOptions ?? {}; + const seed = planOptions.seed ?? DEFAULT_STRATA_SEED; + const targetPerStratum = planOptions.targetPerStratum ?? 2; + const rareThreshold = planOptions.rareThreshold ?? 3; + const axes = planOptions.axes ?? ["outcome", "mode", "turns", "confidence"]; + const tags = planOptions.tags ?? {}; + const highConsequenceTags = planOptions.highConsequenceTags ?? []; + const tolerance = options.tolerance ?? 0.1; + const now = planOptions.now ?? new Date(); + + if (priorPlan !== null) { + if (priorPlan.seed !== seed) { + throw new Error("Cannot expand a cohort plan with a different seed; frozen selections would be re-ordered."); + } + const priorAxes = priorPlan.config.axes.join(","); + if (priorAxes !== axes.join(",")) { + throw new Error("Cannot expand a cohort plan with different axes; create a new plan instead."); + } + } + + const indexText = readFileSync(executionIndexPath, "utf8"); + const tasksPath = options.tasksPath ?? null; + const pool = buildEligiblePool(executionIndexPath, tasksPath, tags); + + const frozenGroups = new Set(); + if (priorPlan !== null) { + for (const entry of priorPlan.selection) frozenGroups.add(entry.execution_group); + } + + const { selected, reports } = planStratifiedCohort({ + eligible: pool.eligible, + axes, + seed, + targetPerStratum: Math.max(targetPerStratum, priorPlan?.config.target_per_stratum ?? 0), + rareThreshold, + highConsequenceTags, + frozenGroups, + }); + + const selection = selected.map((execution) => ({ + execution_group: execution.execution_group, + task_id: execution.task_id, + strata: deriveStrata(execution, axes), + frozen: frozenGroups.has(execution.execution_group), + })); + + const uncovered = reports.filter((report) => report.available === 0); + // Underfilled = decision target not met. For saturated strata this means + // scarcity (every available execution is already included); for others it + // cannot happen under greedy selection, so it always reads as a real gap. + const underfilled = reports.filter((report) => report.selected < report.target); + // Blocking evidence: important underfilled strata. A stratum is important + // when it is rare (at or below the rare threshold) or tagged + // high-consequence — exactly the cases the adaptive loop cannot afford to + // under-sample. + const isImportant = (report: StratumReport) => + report.available <= rareThreshold || + (report.axis === "tag" && highConsequenceTags.includes(report.value)); + const blocking: string[] = underfilled + .filter(isImportant) + .map((report) => { + const kind = report.axis === "tag" && highConsequenceTags.includes(report.value) + ? "High-consequence" + : "Rare"; + return `${kind} stratum ${report.axis}:${report.value} has ${report.available} available executions, target ${report.target}.`; + }); + + let stability: StrataPlan["stability"] = { status: "not_applicable", tolerance, moved: [] }; + if (priorPlan !== null) { + // Stability tracks whether the new batch shifted the eligible pool's + // composition: per-stratum prevalence (available / eligible) before vs + // after. A stratum that grows, shrinks, or appears past the declared + // tolerance is machine-readable evidence that the material moved. + const priorEligible = priorPlan.pool.eligible; + const currentEligible = pool.eligible.length; + const priorPrevalence = new Map(); + for (const report of priorPlan.strata) { + priorPrevalence.set(`${report.axis}:${report.value}`, prevalenceShare(report.available, priorEligible)); + } + const currentByLabel = new Map(); + for (const report of reports) currentByLabel.set(`${report.axis}:${report.value}`, report); + const labels = new Set([...priorPrevalence.keys(), ...currentByLabel.keys()]); + const moved: StrataPlan["stability"]["moved"] = []; + for (const label of [...labels].sort()) { + const before = priorPrevalence.get(label) ?? 0; + const afterReport = currentByLabel.get(label); + const after = afterReport === undefined ? 0 : prevalenceShare(afterReport.available, currentEligible); + if (Math.abs(after - before) > tolerance) { + const [axis, ...rest] = label.split(":"); + moved.push({ axis: axis as StrataAxis, value: rest.join(":"), before, after }); + } + } + stability = { status: moved.length === 0 ? "stable" : "unstable", tolerance, moved }; + } + + return { + schema_version: COHORT_STRATA_PLAN_SCHEMA_VERSION, + created_at: now.toISOString(), + seed, + inputs: { + execution_index: executionIndexPath, + execution_index_sha256: sha256(indexText), + tasks: tasksPath, + }, + config: { axes, target_per_stratum: targetPerStratum, rare_threshold: rareThreshold, high_consequence_tags: highConsequenceTags }, + pool: { index_rows: pool.indexRowCount, eligible: pool.eligible.length, excluded: pool.excluded }, + strata: reports, + selection, + coverage: { + uncovered_strata: uncovered.map((report) => ({ axis: report.axis, value: report.value })), + underfilled_strata: underfilled.map((report) => ({ axis: report.axis, value: report.value, available: report.available, target: report.target })), + }, + ready: blocking.length === 0 && uncovered.length === 0, + blocking, + expansion_of: priorPlan !== null ? planIdentity(priorPlan) : null, + stability, + privacy: { local_only: true, payload_fields_read: false, upload_performed: false }, + }; +} diff --git a/tests/cohort-strata.test.mjs b/tests/cohort-strata.test.mjs new file mode 100644 index 00000000..21d140e1 --- /dev/null +++ b/tests/cohort-strata.test.mjs @@ -0,0 +1,379 @@ +import assert from "node:assert/strict"; +import { spawnSync } from "node:child_process"; +import { createHash } from "node:crypto"; +import { mkdtempSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join, resolve } from "node:path"; +import test from "node:test"; + +import { + buildStrataPlan, + CohortStrataPlanSchema, + DEFAULT_STRATA_SEED, + planIdentity, +} from "../dist/evals/cohort-strata.js"; + +const cli = ["node", resolve("dist/bin.js")]; +const NOW = new Date("2026-09-05T12:00:00.000Z"); +const sha = (value) => createHash("sha256").update(value).digest("hex"); + +function indexRow(overrides) { + return { + schema_version: "understudy.eval-execution-index-row.v1", + source_status: "included", + execution_group: "execution-1", + lineage_status: "complete", + capture_count: 1, + source_files: [{ local_path: "source/traces/one.jsonl", content_sha256: sha("one") }], + task_id: "task-1", + exclusion_reasons: [], + ...overrides, + }; +} + +function taskRow(overrides) { + return { + schema_version: "understudy.benchmark_task.v1", + task_id: "task-1", + execution_group: "execution-1", + // Payload-derived on purpose: the stratifier must never surface this. + title: "SECRET_PROMPT_LEAK summarize the customer invoice", + tool_surface: ["read-file"], + machine_confidence: "high", + ...overrides, + }; +} + +/** + * Synthetic tool workflow per issue #281: simple single-turn reads, a + * multi-turn read→write, a longer write, an error/recovery case (ambiguous + * lineage), an unlinked trace, a taskless execution, and a rare + * high-consequence case tagged by a human. + */ +function buildFoundry(root, { extraIndexRows = [], extraTaskRows = [] } = {}) { + const dir = join(root, "foundry"); + mkdirSync(dir, { recursive: true }); + const indexRows = [ + indexRow({ execution_group: "e-simple-1", task_id: "t-1" }), + indexRow({ execution_group: "e-simple-2", task_id: "t-2" }), + indexRow({ execution_group: "e-simple-3", task_id: "t-3" }), + indexRow({ execution_group: "e-simple-4", task_id: "t-4" }), + indexRow({ execution_group: "e-write-short", task_id: "t-5", capture_count: 4 }), + indexRow({ execution_group: "e-write-long", task_id: "t-6", capture_count: 7 }), + indexRow({ execution_group: "e-recovery", task_id: "t-7", lineage_status: "ambiguous", exclusion_reasons: ["ambiguous_parent"] }), + indexRow({ execution_group: "e-unlinked", task_id: null, lineage_status: "unlinked", exclusion_reasons: ["missing_valid_trace_context"] }), + indexRow({ execution_group: "e-taskless", task_id: null }), + indexRow({ execution_group: "e-excluded-source", source_status: "excluded", task_id: null, execution_group: null }), + indexRow({ execution_group: "e-rare-billing", task_id: "t-8", capture_count: 2 }), + ...extraIndexRows, + ]; + const taskRows = [ + taskRow({ task_id: "t-1", execution_group: "e-simple-1" }), + taskRow({ task_id: "t-2", execution_group: "e-simple-2" }), + taskRow({ task_id: "t-3", execution_group: "e-simple-3" }), + taskRow({ task_id: "t-4", execution_group: "e-simple-4" }), + taskRow({ task_id: "t-5", execution_group: "e-write-short", tool_surface: ["read-file", "write-report"], machine_confidence: "medium" }), + taskRow({ task_id: "t-6", execution_group: "e-write-long", tool_surface: ["read-file", "update-task", "write-file"], machine_confidence: "low" }), + taskRow({ task_id: "t-7", execution_group: "e-recovery" }), + taskRow({ task_id: "t-8", execution_group: "e-rare-billing", machine_confidence: "medium" }), + ...extraTaskRows, + ]; + const executionIndex = join(dir, "execution-index.jsonl"); + const tasks = join(dir, "tasks.jsonl"); + writeFileSync(executionIndex, `${indexRows.map((row) => JSON.stringify(row)).join("\n")}\n`, { mode: 0o600 }); + writeFileSync(tasks, `${taskRows.map((row) => JSON.stringify(row)).join("\n")}\n`, { mode: 0o600 }); + return { dir, executionIndex, tasks }; +} + +const basePlanOptions = { + seed: DEFAULT_STRATA_SEED, + axes: ["outcome", "mode", "turns", "confidence", "tag"], + targetPerStratum: 2, + rareThreshold: 3, + tags: { "e-rare-billing": ["consequence:billing-mutation"] }, + highConsequenceTags: ["consequence:billing-mutation"], + now: NOW, +}; + +test("stratify derives execution strata and selects toward per-stratum targets", () => { + const root = mkdtempSync(join(tmpdir(), "understudy-cohort-strata-")); + try { + const { executionIndex, tasks } = buildFoundry(root); + const plan = buildStrataPlan({ executionIndexPath: executionIndex, tasksPath: tasks, planOptions: basePlanOptions }); + CohortStrataPlanSchema.parse(plan); + + assert.equal(plan.pool.eligible, 7); + const excluded = Object.fromEntries(plan.pool.excluded.map((entry) => [entry.reason, entry.count])); + assert.deepEqual(excluded, { lineage_ambiguous: 1, lineage_unlinked: 1, missing_task: 1, source_excluded: 1 }); + + const byLabel = new Map(plan.strata.map((stratum) => [`${stratum.axis}:${stratum.value}`, stratum])); + // Terminal outcome stratum exists for every eligible execution. + assert.equal(byLabel.get("outcome:complete").available, 7); + // Write vs no-op comes from the foundry's own mutating-tool classifier. + assert.equal(byLabel.get("mode:write").available, 2); + assert.equal(byLabel.get("mode:read-only").available, 5); + // Turn buckets from capture counts. + assert.equal(byLabel.get("turns:single").available, 4); + assert.equal(byLabel.get("turns:short").available, 2); + assert.equal(byLabel.get("turns:long").available, 1); + // Human-guided tag, never inferred. + const tag = byLabel.get("tag:consequence:billing-mutation"); + assert.equal(tag.available, 1); + assert.equal(tag.rule, "saturated_high_consequence"); + assert.equal(tag.selected, 1); + // Rare strata are saturated: minimums, never caps. + assert.equal(byLabel.get("turns:long").rule, "saturated_rare"); + assert.equal(byLabel.get("turns:long").selected, 1); + + const groups = plan.selection.map((entry) => entry.execution_group); + assert.ok(groups.includes("e-rare-billing"), "high-consequence execution must always be selected"); + assert.ok(groups.includes("e-write-long"), "rare stratum member must be saturated in"); + for (const entry of plan.selection) { + assert.ok(entry.task_id, "every selected execution must be judgeable (have a task)"); + assert.ok(entry.strata.length > 0); + assert.equal(entry.frozen, false); + } + assert.equal(plan.expansion_of, null); + assert.equal(plan.stability.status, "not_applicable"); + assert.deepEqual(plan.privacy, { local_only: true, payload_fields_read: false, upload_performed: false }); + } finally { + rmSync(root, { recursive: true, force: true }); + } +}); + +test("stratify plans are byte-identical for the same seed and inputs", () => { + const root = mkdtempSync(join(tmpdir(), "understudy-cohort-strata-determinism-")); + try { + const { executionIndex, tasks } = buildFoundry(root); + const one = buildStrataPlan({ executionIndexPath: executionIndex, tasksPath: tasks, planOptions: basePlanOptions }); + const two = buildStrataPlan({ executionIndexPath: executionIndex, tasksPath: tasks, planOptions: basePlanOptions }); + assert.equal(JSON.stringify(one), JSON.stringify(two)); + const other = buildStrataPlan({ + executionIndexPath: executionIndex, + tasksPath: tasks, + planOptions: { ...basePlanOptions, seed: "other-seed" }, + }); + assert.equal(other.selection.length > 0, true); + } finally { + rmSync(root, { recursive: true, force: true }); + } +}); + +test("stratify treats pilot sizes as minimums, never caps", () => { + const root = mkdtempSync(join(tmpdir(), "understudy-cohort-strata-minimums-")); + try { + const { executionIndex, tasks } = buildFoundry(root); + // Target of 1 per stratum, yet the 2-execution write stratum is rare + // (<= threshold 3) and both members must be included. + const plan = buildStrataPlan({ + executionIndexPath: executionIndex, + tasksPath: tasks, + planOptions: { ...basePlanOptions, targetPerStratum: 1 }, + }); + const write = plan.strata.find((stratum) => `${stratum.axis}:${stratum.value}` === "mode:write"); + assert.equal(write.available, 2); + assert.equal(write.selected, 2); + assert.equal(write.rule, "saturated_rare"); + const groups = plan.selection.map((entry) => entry.execution_group); + assert.ok(groups.includes("e-write-short")); + assert.ok(groups.includes("e-write-long")); + } finally { + rmSync(root, { recursive: true, force: true }); + } +}); + +test("stratify keeps payloads out of the plan artifact", () => { + const root = mkdtempSync(join(tmpdir(), "understudy-cohort-strata-redaction-")); + try { + const { executionIndex, tasks } = buildFoundry(root); + const plan = buildStrataPlan({ executionIndexPath: executionIndex, tasksPath: tasks, planOptions: basePlanOptions }); + const serialized = JSON.stringify(plan); + assert.ok(!serialized.includes("SECRET_PROMPT_LEAK"), "task titles are payload-derived and must not leak"); + assert.ok(!serialized.includes("customer invoice"), "task payloads must not leak"); + } finally { + rmSync(root, { recursive: true, force: true }); + } +}); + +test("stratify expansion preserves frozen selections and reports stability", () => { + const root = mkdtempSync(join(tmpdir(), "understudy-cohort-strata-expand-")); + try { + const { executionIndex, tasks } = buildFoundry(root); + const prior = buildStrataPlan({ executionIndexPath: executionIndex, tasksPath: tasks, planOptions: basePlanOptions }); + + // Grow the pool with two same-shape executions: prevalence of every + // stratum stays within tolerance, so the increment reports stable. + const extra = ["e-new-1", "e-new-2"].map((group, index) => + indexRow({ execution_group: group, task_id: `t-new-${index}` }), + ); + const extraTasks = ["e-new-1", "e-new-2"].map((group, index) => + taskRow({ task_id: `t-new-${index}`, execution_group: group }), + ); + const grownIndex = join(root, "grown-execution-index.jsonl"); + const grownTasks = join(root, "grown-tasks.jsonl"); + const priorIndexRows = readFileSync(executionIndex, "utf8").trim().split("\n"); + writeFileSync(grownIndex, `${[...priorIndexRows, ...extra.map((row) => JSON.stringify(row))].join("\n")}\n`, { mode: 0o600 }); + const priorTaskRows = readFileSync(tasks, "utf8").trim().split("\n"); + writeFileSync(grownTasks, `${[...priorTaskRows, ...extraTasks.map((row) => JSON.stringify(row))].join("\n")}\n`, { mode: 0o600 }); + + const expanded = buildStrataPlan({ + executionIndexPath: grownIndex, + tasksPath: grownTasks, + priorPlan: prior, + tolerance: 0.1, + planOptions: basePlanOptions, + }); + + const priorGroups = prior.selection.map((entry) => entry.execution_group); + const expandedByGroup = new Map(expanded.selection.map((entry) => [entry.execution_group, entry])); + for (const group of priorGroups) { + assert.ok(expandedByGroup.has(group), `frozen selection ${group} must survive expansion`); + assert.equal(expandedByGroup.get(group).frozen, true); + } + assert.equal(expanded.expansion_of, planIdentity(prior)); + assert.equal(expanded.stability.status, "stable"); + assert.deepEqual(expanded.stability.moved, []); + + // Expanding with a different seed is rejected: frozen selections would reorder. + assert.throws(() => buildStrataPlan({ + executionIndexPath: grownIndex, + tasksPath: grownTasks, + priorPlan: prior, + planOptions: { ...basePlanOptions, seed: "other-seed" }, + }), /different seed/); + // Expanding with different axes is rejected instead of silently mutating. + assert.throws(() => buildStrataPlan({ + executionIndexPath: grownIndex, + tasksPath: grownTasks, + priorPlan: prior, + planOptions: { ...basePlanOptions, axes: ["outcome"] }, + }), /different axes/); + } finally { + rmSync(root, { recursive: true, force: true }); + } +}); + +test("stratify expansion reports unstable materials when the batch shifts stratum prevalence", () => { + const root = mkdtempSync(join(tmpdir(), "understudy-cohort-strata-unstable-")); + try { + // Start with a tiny balanced pool: one read-only, one write execution. + const smallIndex = join(root, "small-index.jsonl"); + const smallTasks = join(root, "small-tasks.jsonl"); + writeFileSync(smallIndex, [ + JSON.stringify(indexRow({ execution_group: "e-simple-1", task_id: "t-1" })), + JSON.stringify(indexRow({ execution_group: "e-write-long", task_id: "t-6", capture_count: 7 })), + ].join("\n") + "\n", { mode: 0o600 }); + writeFileSync(smallTasks, [ + JSON.stringify(taskRow({ task_id: "t-1", execution_group: "e-simple-1" })), + JSON.stringify(taskRow({ task_id: "t-6", execution_group: "e-write-long", tool_surface: ["update-task"], machine_confidence: "low" })), + ].join("\n") + "\n", { mode: 0o600 }); + + const prior = buildStrataPlan({ + executionIndexPath: smallIndex, + tasksPath: smallTasks, + planOptions: { ...basePlanOptions, targetPerStratum: 3 }, + }); + // Then expand with a read-only-heavy batch: the eligible pool's + // composition shifts hard (read-only prevalence 0.5 → 0.83), which is + // material movement past the declared tolerance. + const extra = ["e-new-1", "e-new-2", "e-new-3", "e-new-4"].map((group, index) => + JSON.stringify(indexRow({ execution_group: group, task_id: `t-new-${index}` })), + ); + const grownIndex = join(root, "grown-index.jsonl"); + writeFileSync(grownIndex, readFileSync(smallIndex, "utf8") + extra.join("\n") + "\n", { mode: 0o600 }); + const extraTasks = ["e-new-1", "e-new-2", "e-new-3", "e-new-4"].map((group, index) => + JSON.stringify(taskRow({ task_id: `t-new-${index}`, execution_group: group })), + ); + const grownTasks = join(root, "grown-tasks.jsonl"); + writeFileSync(grownTasks, readFileSync(smallTasks, "utf8") + extraTasks.join("\n") + "\n", { mode: 0o600 }); + + const expanded = buildStrataPlan({ + executionIndexPath: grownIndex, + tasksPath: grownTasks, + priorPlan: prior, + tolerance: 0.05, + planOptions: { ...basePlanOptions, targetPerStratum: 3 }, + }); + assert.equal(expanded.stability.status, "unstable"); + assert.ok(expanded.stability.moved.some((move) => move.axis === "mode" && move.value === "read-only")); + } finally { + rmSync(root, { recursive: true, force: true }); + } +}); + +test("stratify flags underfilled rare strata as blocking readiness", () => { + const root = mkdtempSync(join(tmpdir(), "understudy-cohort-strata-blocking-")); + try { + const { executionIndex, tasks } = buildFoundry(root); + // Target 4 but the high-consequence tag stratum only has 1 execution: + // selection saturates it, yet the shortfall stays machine-readable. + const plan = buildStrataPlan({ + executionIndexPath: executionIndex, + tasksPath: tasks, + planOptions: { ...basePlanOptions, targetPerStratum: 4, rareThreshold: 1 }, + }); + assert.equal(plan.ready, false); + assert.ok(plan.blocking.some((reason) => reason.includes("consequence:billing-mutation"))); + const underfilled = plan.coverage.underfilled_strata.map((entry) => `${entry.axis}:${entry.value}`); + assert.ok(underfilled.includes("tag:consequence:billing-mutation")); + } finally { + rmSync(root, { recursive: true, force: true }); + } +}); + +test("evals stratify CLI writes a metadata-only plan from foundry output", () => { + const root = mkdtempSync(join(tmpdir(), "understudy-cohort-strata-cli-")); + try { + const { dir } = buildFoundry(root); + const tagsPath = join(root, "tags.json"); + writeFileSync(tagsPath, JSON.stringify({ "e-rare-billing": ["consequence:billing-mutation"] }), { mode: 0o600 }); + const outPath = join(root, "plan.json"); + const baseEnv = { ...process.env }; + delete baseEnv.UNDERSTUDY_API_KEY; + delete baseEnv.UNDERSTUDY_GATEWAY_URL; + delete baseEnv.FORCE_COLOR; + const result = spawnSync(cli[0], [cli[1], "evals", "stratify", + "--from-foundry", dir, + "--axes", "outcome,mode,turns,confidence,tag", + "--tags", tagsPath, + "--high-consequence-tags", "consequence:billing-mutation", + "--out", outPath, + ], { cwd: process.cwd(), encoding: "utf8", env: baseEnv }); + assert.equal(result.status, 0, result.stderr); + assert.match(result.stdout, /Stratified cohort plan: \d+ of 7 eligible executions\./); + // Rare/high-consequence scarcity is surfaced as blocking evidence. + assert.match(result.stdout, /ready: no/); + assert.match(result.stdout, /High-consequence stratum tag:consequence:billing-mutation/); + assert.doesNotMatch(result.stdout, /SECRET_PROMPT_LEAK/); + + const plan = JSON.parse(readFileSync(outPath, "utf8")); + CohortStrataPlanSchema.parse(plan); + assert.ok(!readFileSync(outPath, "utf8").includes("SECRET_PROMPT_LEAK")); + + // JSON mode emits the same plan artifact on stdout. + const jsonRun = spawnSync(cli[0], [cli[1], "--json", "evals", "stratify", + "--from-foundry", dir, + "--axes", "outcome,mode,turns,confidence,tag", + "--tags", tagsPath, + "--high-consequence-tags", "consequence:billing-mutation", + "--out", join(root, "plan2.json"), + ], { cwd: process.cwd(), encoding: "utf8", env: baseEnv }); + assert.equal(jsonRun.status, 0, jsonRun.stderr); + const jsonPlan = JSON.parse(jsonRun.stdout); + assert.equal(jsonPlan.schema_version, "understudy.eval-cohort-strata-plan.v1"); + + // Expansion via CLI keeps frozen selections. + const expandRun = spawnSync(cli[0], [cli[1], "evals", "stratify", + "--from-foundry", dir, + "--axes", "outcome,mode,turns,confidence,tag", + "--tags", tagsPath, + "--high-consequence-tags", "consequence:billing-mutation", + "--expand-from", outPath, + "--out", join(root, "plan3.json"), + ], { cwd: process.cwd(), encoding: "utf8", env: baseEnv }); + assert.equal(expandRun.status, 0, expandRun.stderr); + assert.match(expandRun.stdout, /frozen selection\(s\) preserved, stability stable/); + } finally { + rmSync(root, { recursive: true, force: true }); + } +});