feat(agent): share session benchmark infrastructure

packages/agent/benchmark/session/benchmark.ts

@@ -0,0 +1,101 @@
+import { strictEqual } from "node:assert/strict";
+import { bench, describe } from "vitest";
+
+export interface BenchmarkTarget<TFixture extends AsyncDisposable> {
+	readonly name: string;
+	createFixture(): Promise<TFixture>;
+}
+
+const READ_BENCHMARK_OPTIONS = {
+	time: 500,
+	iterations: 10,
+	warmupTime: 100,
+	warmupIterations: 5,
+} as const;
+
+export const WRITE_BENCHMARK_ITERATIONS = 30;
+export const WRITE_BENCHMARK_WARMUP_ITERATIONS = 5;
+
+export const WRITE_BENCHMARK_OPTIONS = {
+	time: 0,
+	iterations: WRITE_BENCHMARK_ITERATIONS,
+	warmupTime: 0,
+	warmupIterations: WRITE_BENCHMARK_WARMUP_ITERATIONS,
+} as const;
+
+interface NamedBenchmarkDataset {
+	readonly name: string;
+}
+
+interface ReadBenchmarkScenario<TSubject, TDataset> {
+	readonly name: string;
+	expectedResult(dataset: TDataset): number;
+	run(subject: TSubject, dataset: TDataset): Promise<number>;
+}
+
+interface RegisterReadBenchmarksOptions<
+	TDataset extends NamedBenchmarkDataset,
+	TFixture extends AsyncDisposable,
+	TSubject,
+> {
+	readonly datasets: readonly TDataset[];
+	readonly targets: readonly BenchmarkTarget<TFixture>[];
+	readonly scenarios: readonly ReadBenchmarkScenario<TSubject, TDataset>[];
+	prepare(fixture: TFixture, dataset: TDataset): Promise<void>;
+	getSubject(fixture: TFixture): TSubject;
+}
+
+interface PreparedReadFixture<TFixture> {
+	readonly datasetName: string;
+	readonly targetName: string;
+	readonly fixture: TFixture;
+}
+
+/** Prepares, validates, registers, and disposes one immutable fixture per target and dataset. */
+export async function registerReadBenchmarks<
+	TDataset extends NamedBenchmarkDataset,
+	TFixture extends AsyncDisposable,
+	TSubject,
+>(options: RegisterReadBenchmarksOptions<TDataset, TFixture, TSubject>): Promise<void> {
+	const fixtures: TFixture[] = [];
+	const preparedFixtures: PreparedReadFixture<TFixture>[] = [];
+
+	process.once("beforeExit", async () => {
+		await Promise.all(fixtures.map((fixture) => fixture[Symbol.asyncDispose]()));
+	});
+
+	for (const dataset of options.datasets) {
+		for (const target of options.targets) {
+			const fixture = await target.createFixture();
+			fixtures.push(fixture);
+			await options.prepare(fixture, dataset);
+			const subject = options.getSubject(fixture);
+			for (const scenario of options.scenarios) {
+				strictEqual(await scenario.run(subject, dataset), scenario.expectedResult(dataset));
+			}
+			preparedFixtures.push({ datasetName: dataset.name, targetName: target.name, fixture });
+		}
+	}
+
+	for (const dataset of options.datasets) {
+		for (const scenario of options.scenarios) {
+			describe(`${scenario.name} (${dataset.name})`, () => {
+				for (const target of options.targets) {
+					const prepared = preparedFixtures.find(
+						(candidate) => candidate.datasetName === dataset.name && candidate.targetName === target.name,
+					);
+					if (prepared === undefined) throw new Error("Benchmark fixture was not initialized");
+					const subject = options.getSubject(prepared.fixture);
+
+					bench(
+						target.name,
+						async () => {
+							await scenario.run(subject, dataset);
+						},
+						READ_BENCHMARK_OPTIONS,
+					);
+				}
+			});
+		}
+	}
+}

packages/agent/benchmark/session/session-repo-targets.ts

@@ -1,15 +1,11 @@
 import { MemorySessionRepo } from "../../src/harness/session/index.ts";
 import type { SessionRepo } from "../../src/harness/session/types.ts";
+import type { BenchmarkTarget } from "./benchmark.ts";
 
 export interface SessionRepoBenchmarkFixture extends AsyncDisposable {
 	readonly repo: SessionRepo;
 }
 
-interface SessionRepoBenchmarkTarget {
-	readonly name: string;
-	createFixture(): Promise<SessionRepoBenchmarkFixture>;
-}
-
 const NOW = 1_700_000_000_000;
 
 export const sessionRepoBenchmarkTargets = [
@@ -23,4 +19,4 @@ export const sessionRepoBenchmarkTargets = [
 			});
 		},
 	},
-] satisfies readonly SessionRepoBenchmarkTarget[];
+] satisfies readonly BenchmarkTarget<SessionRepoBenchmarkFixture>[];

packages/agent/benchmark/session/session-repo.bench.ts

@@ -1,67 +1,19 @@
-import { strictEqual } from "node:assert/strict";
-import { bench, describe } from "vitest";
 import {
 	SESSION_REPO_CATALOG_BENCHMARK_DATASETS,
 	SESSION_REPO_CATALOG_READ_BENCHMARK_SCENARIOS,
 	seedSessionRepoCatalogBenchmark,
 } from "../../src/harness/session/testing/index.ts";
-import { type SessionRepoBenchmarkFixture, sessionRepoBenchmarkTargets } from "./session-repo-targets.ts";
+import { registerReadBenchmarks } from "./benchmark.ts";
+import { sessionRepoBenchmarkTargets } from "./session-repo-targets.ts";
 
-const BENCHMARK_TIME_MS = 500;
-const BENCHMARK_WARMUP_TIME_MS = 100;
-
-interface PreparedCatalogFixture {
-	readonly datasetName: string;
-	readonly targetName: string;
-	readonly fixture: SessionRepoBenchmarkFixture;
-}
-
-let benchmarkSink = 0;
-const allFixtures: SessionRepoBenchmarkFixture[] = [];
-const catalogFixtures: PreparedCatalogFixture[] = [];
-
-process.once("beforeExit", () => {
-	void Promise.all(allFixtures.map((fixture) => fixture[Symbol.asyncDispose]()));
-});
-
-// Vitest 4's benchmark runner does not execute normal test lifecycle hooks.
-// Prepare one immutable catalog per backend/dataset before registration.
-for (const dataset of SESSION_REPO_CATALOG_BENCHMARK_DATASETS) {
-	for (const target of sessionRepoBenchmarkTargets) {
-		const fixture = await target.createFixture();
-		allFixtures.push(fixture);
+await registerReadBenchmarks({
+	datasets: SESSION_REPO_CATALOG_BENCHMARK_DATASETS,
+	targets: sessionRepoBenchmarkTargets,
+	scenarios: SESSION_REPO_CATALOG_READ_BENCHMARK_SCENARIOS,
+	async prepare(fixture, dataset) {
 		await seedSessionRepoCatalogBenchmark(fixture.repo, dataset);
-		for (const scenario of SESSION_REPO_CATALOG_READ_BENCHMARK_SCENARIOS) {
-			strictEqual(await scenario.run(fixture.repo), scenario.expectedResult(dataset));
-		}
-		catalogFixtures.push({ datasetName: dataset.name, targetName: target.name, fixture });
-	}
-}
-
-for (const dataset of SESSION_REPO_CATALOG_BENCHMARK_DATASETS) {
-	for (const scenario of SESSION_REPO_CATALOG_READ_BENCHMARK_SCENARIOS) {
-		describe(`${scenario.name} (${dataset.name})`, () => {
-			for (const target of sessionRepoBenchmarkTargets) {
-				const prepared = catalogFixtures.find(
-					(candidate) => candidate.datasetName === dataset.name && candidate.targetName === target.name,
-				);
-				if (prepared === undefined) throw new Error("Benchmark fixture was not initialized");
-
-				bench(
-					target.name,
-					async () => {
-						benchmarkSink = await scenario.run(prepared.fixture.repo);
-					},
-					{
-						time: BENCHMARK_TIME_MS,
-						iterations: 10,
-						warmupTime: BENCHMARK_WARMUP_TIME_MS,
-						warmupIterations: 5,
-					},
-				);
-			}
-		});
-	}
-}
-
-void benchmarkSink;
+	},
+	getSubject(fixture) {
+		return fixture.repo;
+	},
+});

packages/agent/benchmark/session/storage-targets.ts

@@ -1,10 +1,6 @@
 import { MemoryStorage } from "../../src/harness/session/memory.ts";
 import type { StorageFixture } from "../../src/harness/session/testing/index.ts";
-
-interface StorageBenchmarkTarget {
-	readonly name: string;
-	createFixture(): Promise<StorageFixture>;
-}
+import type { BenchmarkTarget } from "./benchmark.ts";
 
 const NOW = 1_700_000_000_000;
 
@@ -19,4 +15,4 @@ export const storageBenchmarkTargets = [
 			});
 		},
 	},
-] satisfies readonly StorageBenchmarkTarget[];
+] satisfies readonly BenchmarkTarget<StorageFixture>[];

packages/agent/benchmark/session/storage.bench.ts

@@ -7,50 +7,42 @@ import {
 	seedStorageBenchmark,
 	type StorageFixture,
 } from "../../src/harness/session/testing/index.ts";
+import {
+	registerReadBenchmarks,
+	WRITE_BENCHMARK_ITERATIONS,
+	WRITE_BENCHMARK_OPTIONS,
+	WRITE_BENCHMARK_WARMUP_ITERATIONS,
+} from "./benchmark.ts";
 import { storageBenchmarkTargets } from "./storage-targets.ts";
 
-const BENCHMARK_TIME_MS = 500;
-const BENCHMARK_WARMUP_TIME_MS = 100;
-const WRITE_ITERATIONS = 30;
-const WRITE_WARMUP_ITERATIONS = 5;
-
-interface PreparedReadFixture {
-	readonly datasetName: string;
-	readonly targetName: string;
-	readonly fixture: StorageFixture;
-}
-
 interface PreparedWriteFixtures {
 	readonly scenarioName: string;
 	readonly targetName: string;
 	readonly pendingFixtures: StorageFixture[];
 }
 
-let benchmarkSink = 0;
 const allFixtures: StorageFixture[] = [];
-const readFixtures: PreparedReadFixture[] = [];
 const writeFixtures: PreparedWriteFixtures[] = [];
 
-process.once("beforeExit", () => {
-	void Promise.all(allFixtures.map((fixture) => fixture[Symbol.asyncDispose]()));
+process.once("beforeExit", async () => {
+	await Promise.all(allFixtures.map((fixture) => fixture[Symbol.asyncDispose]()));
+});
+
+await registerReadBenchmarks({
+	datasets: STORAGE_BENCHMARK_DATASETS,
+	targets: storageBenchmarkTargets,
+	scenarios: STORAGE_READ_BENCHMARK_SCENARIOS,
+	prepare(fixture, dataset) {
+		return seedStorageBenchmark(fixture.storage, dataset);
+	},
+	getSubject(fixture) {
+		return fixture.storage;
+	},
 });
 
 // Vitest 4's benchmark runner does not execute normal test lifecycle hooks.
-// Prepare read fixtures and one write fixture per warmup/measured iteration
-// before registration. Vitest 5 can express the same lifecycle with each
-// registration's beforeAll/beforeEach/afterEach/afterAll hooks.
-for (const dataset of STORAGE_BENCHMARK_DATASETS) {
-	for (const target of storageBenchmarkTargets) {
-		const fixture = await target.createFixture();
-		allFixtures.push(fixture);
-		await seedStorageBenchmark(fixture.storage, dataset);
-		for (const scenario of STORAGE_READ_BENCHMARK_SCENARIOS) {
-			strictEqual(await scenario.run(fixture.storage, dataset), scenario.expectedResult(dataset));
-		}
-		readFixtures.push({ datasetName: dataset.name, targetName: target.name, fixture });
-	}
-}
-
+// Prepare one write fixture per warmup/measured iteration before registration.
+// Vitest 5 can express the same lifecycle with each registration's hooks.
 for (const scenario of STORAGE_WRITE_BENCHMARK_SCENARIOS) {
 	for (const target of storageBenchmarkTargets) {
 		const validationFixture = await target.createFixture();
@@ -59,7 +51,11 @@ for (const scenario of STORAGE_WRITE_BENCHMARK_SCENARIOS) {
 		strictEqual(await scenario.run(validationFixture.storage), scenario.writeCount);
 
 		const pendingFixtures: StorageFixture[] = [];
-		for (let index = 0; index < WRITE_WARMUP_ITERATIONS + WRITE_ITERATIONS; index++) {
+		for (
+			let index = 0;
+			index < WRITE_BENCHMARK_WARMUP_ITERATIONS + WRITE_BENCHMARK_ITERATIONS;
+			index++
+		) {
 			const fixture = await target.createFixture();
 			allFixtures.push(fixture);
 			await scenario.prepare?.(fixture.storage);
@@ -69,32 +65,6 @@ for (const scenario of STORAGE_WRITE_BENCHMARK_SCENARIOS) {
 	}
 }
 
-for (const dataset of STORAGE_BENCHMARK_DATASETS) {
-	for (const scenario of STORAGE_READ_BENCHMARK_SCENARIOS) {
-		describe(`${scenario.name} (${dataset.name})`, () => {
-			for (const target of storageBenchmarkTargets) {
-				const prepared = readFixtures.find(
-					(candidate) => candidate.datasetName === dataset.name && candidate.targetName === target.name,
-				);
-				if (prepared === undefined) throw new Error("Benchmark fixture was not initialized");
-
-				bench(
-					target.name,
-					async () => {
-						benchmarkSink = await scenario.run(prepared.fixture.storage, dataset);
-					},
-					{
-						time: BENCHMARK_TIME_MS,
-						iterations: 10,
-						warmupTime: BENCHMARK_WARMUP_TIME_MS,
-						warmupIterations: 5,
-					},
-				);
-			}
-		});
-	}
-}
-
 for (const scenario of STORAGE_WRITE_BENCHMARK_SCENARIOS) {
 	describe(scenario.name, () => {
 		for (const target of storageBenchmarkTargets) {
@@ -108,17 +78,10 @@ for (const scenario of STORAGE_WRITE_BENCHMARK_SCENARIOS) {
 				async () => {
 					const fixture = prepared.pendingFixtures.shift();
 					if (fixture === undefined) throw new Error("Benchmark fixture pool was exhausted");
-					benchmarkSink = await scenario.run(fixture.storage);
-				},
-				{
-					time: 0,
-					iterations: WRITE_ITERATIONS,
-					warmupTime: 0,
-					warmupIterations: WRITE_WARMUP_ITERATIONS,
+					await scenario.run(fixture.storage);
 				},
+				WRITE_BENCHMARK_OPTIONS,
 			);
 		}
 	});
 }
-
-void benchmarkSink;