ember-bench

1import type {
2	EngineInterface,
3	ModelForkResult,
4	ProcessRunResult,
5	Register,
6	Timer,
7} from "claude-code";
8
9import type { Ember } from "./bench.ts";
10import {
11	ARM_PREFIX,
12	armKey,
13	bunCandidates,
14	EMBER,
15	emberOf,
16	forkRecord,
17	INGEST_EVERY_MS,
18	limitsOf,
19	OFF,
20	OFF_FORK,
21	offRefusal,
22	sessionRecord,
23	UNKNOWN_EMBER,
24	USAGE,
25} from "./bench.ts";
26import type { BenchRecord } from "./record.ts";
27import { isRecord, isSessionId } from "./record.ts";
28
29type Live = { timer?: Timer; ember: Ember; isSeen: boolean };
30
31const scriptOf = ($: EngineInterface) => `${$.plugin.root}/scripts/bench.ts`;
32
33const debug = ($: EngineInterface, text: string) =>
34	$.ui.log(`ember-bench: ${text}`, { to: "debug" });
35
36/** Starts the script and does not wait for it; tries each bun in turn until one starts. */
37function spawnScript($: EngineInterface, args: readonly string[]): void {
38	void (async () => {
39		for (const bun of bunCandidates(await $.env.get("HOME"))) {
40			try {
41				for await (const _ of $.process.spawn({
42					argv: [bun, scriptOf($), ...args],
43				})) {
44					// Drained so the child never blocks on a full pipe.
45				}
46				return;
47			} catch {
48				// Not startable here: try the next bun.
49			}
50		}
51		debug($, "could not start bun");
52	})().catch(() => debug($, "could not start bun"));
53}
54
55/** Runs the script and waits for its output; tries each bun in turn until one starts. */
56async function runScript(
57	$: EngineInterface,
58	args: readonly string[],
59): Promise<ProcessRunResult> {
60	let failure: unknown;
61	for (const bun of bunCandidates(await $.env.get("HOME"))) {
62		try {
63			return await $.process.run([bun, scriptOf($), ...args], {
64				timeoutMs: 60_000,
65			});
66		} catch (e) {
67			failure = e;
68		}
69	}
70	throw failure;
71}
72
73function send($: EngineInterface, r: BenchRecord): void {
74	if (isRecord(r)) spawnScript($, ["record", JSON.stringify(r)]);
75	else debug($, "refused an invalid record");
76}
77
78const isOffNow = async ($: EngineInterface, session: string) =>
79	(await $.store.get(armKey(session))) === OFF;
80
81async function recordSession($: EngineInterface, live: Live): Promise<void> {
82	const session = await $.session.id();
83	if (!isSessionId(session)) return debug($, "the session id is no UUID");
84	const [isOff, version, usage, pluginConfigs] = await Promise.all([
85		isOffNow($, session),
86		$.session.version(),
87		$.session.usage(),
88		$.settings.read({ key: "pluginConfigs" }).catch(() => undefined),
89	]);
90	send(
91		$,
92		sessionRecord({
93			session,
94			startedMs: usage.startedAt,
95			isOff,
96			ccVersion: version.version,
97			ember: live.ember,
98			pluginConfigs,
99		}),
100	);
101}
102
103function start($: EngineInterface, live: Live): void {
104	live.timer?.cancel();
105	spawnScript($, ["ingest"]);
106	live.timer = $.clock.every(INGEST_EVERY_MS, () => spawnScript($, ["ingest"]));
107	recordSession($, live).catch(() => debug($, "could not record the session"));
108}
109
110/** Marks, once per load, that ember's calls pass through ember-bench: the proof it runs before ember. */
111async function noteSeen($: EngineInterface, live: Live): Promise<void> {
112	if (live.isSeen) return;
113	live.isSeen = true;
114	debug($, "saw a call from ember");
115	try {
116		const session = await $.session.id();
117		send($, { kind: "seen", session, atMs: await $.clock.now() });
118	} catch {
119		debug($, "could not record a call from ember");
120	}
121}
122
123const limitsNow = async ($: EngineInterface) =>
124	limitsOf((await $.session.usage().catch(() => undefined))?.rateLimits);
125
126async function benchOff($: EngineInterface, live: Live): Promise<string> {
127	const session = await $.session.id();
128	if (!isSessionId(session)) return "ember-bench: this session has no id yet.";
129	const refusal = offRefusal(
130		await $.session.turns(),
131		await isOffNow($, session),
132	);
133	if (refusal !== undefined) return refusal;
134	await $.store.set(armKey(session), OFF);
135	await recordSession($, live);
136	return "ember-bench: this session is off. Ember's cache refreshes are refused here, and no refresh request is sent.";
137}
138
139async function benchStatus($: EngineInterface): Promise<string> {
140	const session = await $.session.id();
141	if (!isSessionId(session)) return "ember-bench: this session has no id yet.";
142	const arm = `arm: ${(await isOffNow($, session)) ? "off" : "on"}`;
143	try {
144		const out = await runScript($, ["status", session]);
145		return out.exitCode === 0
146			? `ember-bench\n${arm}\n${out.stdout.trimEnd()}`
147			: `ember-bench\n${arm}\nThe database could not be read.`;
148	} catch {
149		return `ember-bench\n${arm}\nbun could not be started.`;
150	}
151}
152
153async function benchDelete($: EngineInterface): Promise<string> {
154	let answer: string;
155	try {
156		answer = await $.ui.ask("Delete all ember-bench data?", ["Delete", "Keep"]);
157	} catch {
158		answer = "Keep";
159	}
160	if (answer !== "Delete") return "ember-bench: Nothing was deleted.";
161	try {
162		const out = await runScript($, ["delete"]);
163		if (out.exitCode !== 0) return "ember-bench: the delete failed.";
164	} catch {
165		return "ember-bench: bun could not be started; nothing was deleted.";
166	}
167	// The current session keeps its arm: it cannot be chosen again after the first prompt.
168	const keep = armKey(await $.session.id());
169	for (const key of await $.store.keys())
170		if (key.startsWith(ARM_PREFIX) && key !== keep) await $.store.delete(key);
171	return "ember-bench: all collected data was deleted. The next ingest collects the transcripts again.";
172}
173
174export const register: Register = (on) => {
175	const live: Live = { ember: UNKNOWN_EMBER, isSeen: false };
176
177	on("plugin.register", { name: "ember" }, (_, e, next) => {
178		live.ember = emberOf(e);
179		return next(e);
180	});
181
182	on("session.start", async ($, e, next) => {
183		start($, live);
184		await $.command
185			.register({
186				name: "bench",
187				description:
188					"ember-bench: switch this session off before its first prompt, show status, or delete all data",
189				argumentHint: "off | status | delete",
190				immediate: true,
191			})
192			.catch(() => debug($, "/bench is taken; it is not available"));
193		return next(e);
194	});
195
196	// session.start does not fire after /clear, /resume or /branch; this does.
197	on(
198		"classic.SessionStart",
199		{ source: ["clear", "resume", "fork"] },
200		($, e, next) => {
201			start($, live);
202			return next(e);
203		},
204	).catch((_, e, next) => next(e));
205
206	on("model.fork", async ($, e, next) => {
207		if (next.origin.plugin !== EMBER) return next(e);
208		void noteSeen($, live);
209		const session = await $.session.id();
210		const firedMs = await $.clock.now();
211		if (await isOffNow($, session)) {
212			const none = limitsOf(undefined);
213			send(
214				$,
215				forkRecord({
216					session,
217					firedMs,
218					durationMs: 0,
219					result: OFF_FORK,
220					isOff: true,
221					before: none,
222					after: none,
223				}),
224			);
225			return { value: OFF_FORK };
226		}
227		const before = await limitsNow($);
228		const r = await next(e);
229		if ("value" in r && r.value !== undefined) {
230			const result: ModelForkResult = r.value;
231			void (async () => {
232				const [endMs, after] = await Promise.all([$.clock.now(), limitsNow($)]);
233				send(
234					$,
235					forkRecord({
236						session,
237						firedMs,
238						durationMs: endMs - firedMs,
239						result,
240						isOff: false,
241						before,
242						after,
243					}),
244				);
245			})().catch(() => debug($, "could not record a fork"));
246		}
247		return r;
248		// A failure before next refuses ember's fork: an off session must never send one, and ember retries in 30 s.
249	}).catch((_, e, next) =>
250		next.called || next.origin.plugin !== EMBER ? next(e) : { value: OFF_FORK },
251	);
252
253	on("ui.open", ($, e, next) => {
254		if (next.origin.plugin === EMBER) void noteSeen($, live);
255		return next(e);
256	});
257
258	on("command.run", { command: "ember-probe" }, async ($, e, next) => {
259		void (async () => {
260			const session = await $.session.id();
261			send($, { kind: "probe", session, atMs: await $.clock.now() });
262		})().catch(() => debug($, "could not record a probe"));
263		return next(e);
264	});
265
266	on("command.run", { command: "bench" }, async ($, e) => {
267		const arg = e.args.trim();
268		const text =
269			arg === "off"
270				? await benchOff($, live)
271				: arg === "status"
272					? await benchStatus($)
273					: arg === "delete"
274						? await benchDelete($)
275						: USAGE;
276		return { text };
277	});
278};
279

1import type { ModelForkResult, SessionRateLimit } from "claude-code";
2
3import type { ForkRecord, SessionRecord } from "./record.ts";
4
5export const EMBER = "ember";
6/** Store key prefix of a session's arm; the value is 0 for off. Unset is on. */
7export const ARM_PREFIX = "arm:";
8export const OFF = 0;
9export const INGEST_EVERY_MS = 30 * 60_000;
10
11export const armKey = (session: string) => `${ARM_PREFIX}${session}`;
12
13/** What an off session answers ember's fork with: no request was made, and ember's judgeFork reads `aborted` as refused. */
14export const OFF_FORK: ModelForkResult = {
15	isAnswered: false,
16	reason: "aborted",
17	usage: {
18		input_tokens: 0,
19		output_tokens: 0,
20		cache_read_input_tokens: 0,
21		cache_creation_input_tokens: 0,
22	},
23};
24
25export type Ember = { version: string | null; source: string };
26export const UNKNOWN_EMBER: Ember = { version: null, source: "unknown" };
27
28export const emberOf = (e: {
29	version?: string;
30	provenance: string;
31}): Ember => ({
32	version: e.version ?? null,
33	source: e.provenance,
34});
35
36/** Where bun may be: on PATH first, then the installer's default. */
37export const bunCandidates = (home: string | undefined): string[] =>
38	home === undefined ? ["bun"] : ["bun", `${home}/.bun/bin/bun`];
39
40export const percentOf = (
41	limits: readonly SessionRateLimit[] | undefined,
42	kind: string,
43): number | null => limits?.find((l) => l.kind === kind)?.percentUsed ?? null;
44
45const objOf = (v: unknown): Record<string, unknown> | undefined =>
46	typeof v === "object" && v !== null && !Array.isArray(v)
47		? (v as Record<string, unknown>)
48		: undefined;
49const numberOr = (v: unknown) =>
50	typeof v === "number" && Number.isFinite(v) ? v : null;
51const flagOr = (v: unknown): 0 | 1 | null =>
52	typeof v === "boolean" ? (v ? 1 : 0) : null;
53
54/**
55 * Ember's option values as settings hold them (`pluginConfigs[<plugin id>].options`), the installed id first and
56 * the --plugin-dir spellings after. A field settings do not hold is null: ember's manifest default applies.
57 */
58export function optionsOf(pluginConfigs: unknown, source: string) {
59	const all = objOf(pluginConfigs);
60	const key = [source, EMBER, `${EMBER}@inline`].find(
61		(k) => objOf(all?.[k]) !== undefined,
62	);
63	const o = objOf(objOf(key === undefined ? undefined : all?.[key])?.options);
64	return {
65		ttlMinutes: numberOr(o?.ttlMinutes),
66		refresh: flagOr(o?.refresh),
67		refreshShortTtl: flagOr(o?.refreshShortTtl),
68		refreshLeadSeconds: numberOr(o?.refreshLeadSeconds),
69	};
70}
71
72export function sessionRecord(p: {
73	session: string;
74	startedMs: number;
75	isOff: boolean;
76	ccVersion: string | null;
77	ember: Ember;
78	pluginConfigs: unknown;
79}): SessionRecord {
80	return {
81		kind: "session",
82		session: p.session,
83		startedMs: p.startedMs,
84		arm: p.isOff ? 0 : 1,
85		ccVersion: p.ccVersion,
86		emberVersion: p.ember.version,
87		emberSource: p.ember.source,
88		...optionsOf(p.pluginConfigs, p.ember.source),
89	};
90}
91
92export type Limits = { fiveHour: number | null; sevenDay: number | null };
93export const limitsOf = (
94	limits: readonly SessionRateLimit[] | undefined,
95): Limits => ({
96	fiveHour: percentOf(limits, "five_hour"),
97	sevenDay: percentOf(limits, "seven_day"),
98});
99
100/** One fork as a row: when it fired, how long it took, how it ended, what it cost, and the plan's use around it. */
101export function forkRecord(p: {
102	session: string;
103	firedMs: number;
104	durationMs: number;
105	result: ModelForkResult;
106	isOff: boolean;
107	before: Limits;
108	after: Limits;
109}): ForkRecord {
110	const r = p.result;
111	const u = "usage" in r ? r.usage : undefined;
112	return {
113		kind: "fork",
114		session: p.session,
115		firedMs: p.firedMs,
116		durationMs: p.durationMs,
117		isAnswered: r.isAnswered ? 1 : 0,
118		reason: p.isOff ? "off" : r.isAnswered ? "answered" : r.reason,
119		input: u?.input_tokens ?? 0,
120		output: u?.output_tokens ?? 0,
121		read: u?.cache_read_input_tokens ?? 0,
122		creation: u?.cache_creation_input_tokens ?? 0,
123		fiveHourBefore: p.before.fiveHour,
124		fiveHourAfter: p.after.fiveHour,
125		sevenDayBefore: p.before.sevenDay,
126		sevenDayAfter: p.after.sevenDay,
127	};
128}
129
130export const USAGE = "Usage: /bench off | status | delete";
131
132/** The answer to `/bench off`, or undefined when the session may switch off now. */
133export function offRefusal(turns: number, isOff: boolean): string | undefined {
134	if (isOff) return "ember-bench: this session is already off.";
135	if (turns > 0)
136		return "ember-bench: /bench off works only before the first prompt, so the choice cannot depend on how the session goes. This session stays on.";
137	return undefined;
138}
139

1/** The rows the mod hands the script. Shared by both sides: the mod checks before it sends, the script before it writes. */
2
3export const FORK_REASONS = [
4	"answered",
5	"api-error",
6	"empty-reply",
7	"aborted",
8	"nothing-to-fork",
9	"off",
10] as const;
11export type ForkReason = (typeof FORK_REASONS)[number];
12
13export type SessionRecord = {
14	kind: "session";
15	session: string;
16	startedMs: number;
17	/** 1 records refreshes as usual, 0 refuses ember's refresh forks. */
18	arm: 0 | 1;
19	ccVersion: string | null;
20	emberVersion: string | null;
21	/** `<name>@<marketplace>` installed, `<name>@inline` by --plugin-dir, or `unknown`. */
22	emberSource: string;
23	ttlMinutes: number | null;
24	refresh: 0 | 1 | null;
25	refreshShortTtl: 0 | 1 | null;
26	refreshLeadSeconds: number | null;
27};
28
29export type ForkRecord = {
30	kind: "fork";
31	session: string;
32	firedMs: number;
33	durationMs: number;
34	isAnswered: 0 | 1;
35	reason: ForkReason;
36	input: number;
37	output: number;
38	read: number;
39	creation: number;
40	fiveHourBefore: number | null;
41	fiveHourAfter: number | null;
42	sevenDayBefore: number | null;
43	sevenDayAfter: number | null;
44};
45
46/** `probe`: /ember-probe ran (its fork is no refresh). `seen`: ember-bench saw a call ember made. */
47export type MarkRecord = {
48	kind: "probe" | "seen";
49	session: string;
50	atMs: number;
51};
52
53export type BenchRecord = SessionRecord | ForkRecord | MarkRecord;
54
55const UUID = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;
56const VERSION = /^[\w.+-]{1,64}$/;
57const SOURCE = /^[\w.-]{1,64}@[\w.-]{1,64}$/;
58
59export const isSessionId = (v: unknown): v is string =>
60	typeof v === "string" && UUID.test(v);
61
62const isCount = (v: unknown): v is number =>
63	typeof v === "number" && Number.isFinite(v) && v >= 0;
64const isCountOrNull = (v: unknown) => v === null || isCount(v);
65const isFlag = (v: unknown) => v === 0 || v === 1;
66const isFlagOrNull = (v: unknown) => v === null || isFlag(v);
67const isVersionOrNull = (v: unknown) =>
68	v === null || (typeof v === "string" && VERSION.test(v));
69
70const KEYS: Record<BenchRecord["kind"], readonly string[]> = {
71	session: [
72		"kind",
73		"session",
74		"startedMs",
75		"arm",
76		"ccVersion",
77		"emberVersion",
78		"emberSource",
79		"ttlMinutes",
80		"refresh",
81		"refreshShortTtl",
82		"refreshLeadSeconds",
83	],
84	fork: [
85		"kind",
86		"session",
87		"firedMs",
88		"durationMs",
89		"isAnswered",
90		"reason",
91		"input",
92		"output",
93		"read",
94		"creation",
95		"fiveHourBefore",
96		"fiveHourAfter",
97		"sevenDayBefore",
98		"sevenDayAfter",
99	],
100	probe: ["kind", "session", "atMs"],
101	seen: ["kind", "session", "atMs"],
102};
103
104/** True when `v` is exactly one record: known keys only, every number finite, the session a UUID, every text from a closed form. */
105export function isRecord(v: unknown): v is BenchRecord {
106	if (typeof v !== "object" || v === null || Array.isArray(v)) return false;
107	const r = v as Record<string, unknown>;
108	const keys = KEYS[r.kind as BenchRecord["kind"]];
109	if (keys === undefined || !isSessionId(r.session)) return false;
110	const own = Object.keys(r);
111	if (own.length !== keys.length || own.some((k) => !keys.includes(k)))
112		return false;
113	switch (r.kind) {
114		case "session":
115			return (
116				isCount(r.startedMs) &&
117				isFlag(r.arm) &&
118				isVersionOrNull(r.ccVersion) &&
119				isVersionOrNull(r.emberVersion) &&
120				(r.emberSource === "unknown" ||
121					(typeof r.emberSource === "string" && SOURCE.test(r.emberSource))) &&
122				isCountOrNull(r.ttlMinutes) &&
123				isFlagOrNull(r.refresh) &&
124				isFlagOrNull(r.refreshShortTtl) &&
125				isCountOrNull(r.refreshLeadSeconds)
126			);
127		case "fork":
128			return (
129				isCount(r.firedMs) &&
130				isCount(r.durationMs) &&
131				isFlag(r.isAnswered) &&
132				FORK_REASONS.includes(r.reason as ForkReason) &&
133				isCount(r.input) &&
134				isCount(r.output) &&
135				isCount(r.read) &&
136				isCount(r.creation) &&
137				isCountOrNull(r.fiveHourBefore) &&
138				isCountOrNull(r.fiveHourAfter) &&
139				isCountOrNull(r.sevenDayBefore) &&
140				isCountOrNull(r.sevenDayAfter)
141			);
142		default:
143			return isCount(r.atMs);
144	}
145}
146