Keep a running build going through a Chopin restart

apps/connector/src/connection.test.ts

@@ -585,8 +585,8 @@ test("a connection failure's cause never shows credentials from the error", () =
 	expect(text).not.toContain("abcdef0123456789xyz");
 });
 
-test("machine, Planner and pairing key tokens are masked wherever chopin prints text", () => {
-	for (let prefix of ["chopin_m_", "chopin_p_", "chopin_k_"]) {
+test("machine, Planner, pairing key and build run tokens are masked wherever chopin prints text", () => {
+	for (let prefix of ["chopin_m_", "chopin_p_", "chopin_k_", "chopin_r_"]) {
 		let token = `${prefix}${"Zx9_-".repeat(9).slice(0, 43)}`;
 		expect(redacted(`ran chopin connect --key ${token} -- agent`)).toBe(
 			"ran chopin connect --key [redacted] -- agent",

apps/connector/src/connection.ts

@@ -49,7 +49,7 @@ function stopped(reason: Exclude<ErrorClass, "retry" | "forgotten">, error: unkn
 
 const USERINFO = /\b([a-z][a-z0-9+.-]*:\/\/)[^\s/@]+@/gi;
 const TOKEN =
-	/\b(?:gh[opsur]_|github_pat_|chopin_[kmp]_|sk-)[A-Za-z0-9_-]{8,}|\b(Bearer\s+)\S{8,}/g;
+	/\b(?:gh[opsur]_|github_pat_|chopin_[kmpr]_|sk-)[A-Za-z0-9_-]{8,}|\b(Bearer\s+)\S{8,}/g;
 
 /**
  * `text` with URL credentials, token-shaped strings and each of the `secrets` it knows, such as

apps/connector/src/run.test.ts

@@ -598,6 +598,37 @@ test("a network drop mid-run keeps the agent running and delivers its completion
 	expect(stopped).not.toBe(running);
 });
 
+test("a build keeps its agent through a server restart and its heartbeat is heard soon after", async () => {
+	let abort = new AbortController();
+	let renewals: Array<Record<string, unknown>> = [];
+	let started = performance.now();
+	let { documentId, calls } = await run(
+		{},
+		["--hang"],
+		(name, args) => {
+			if (name !== "renew_implementation_build") return {};
+			renewals.push(args);
+			// Chopin is down for eight heartbeats, then forgets the connection the connector had.
+			if (renewals.length <= 8) throw unreachable();
+			if (renewals.length === 9) {
+				throw new ConnectorError("connection-unavailable", "connection-unavailable");
+			}
+			// It carries on with the build; the person then stops the agent from the terminal.
+			abort.abort();
+			return { accepted: true };
+		},
+		abort.signal,
+		{ random: () => 0.999 },
+	);
+	// Retries wait no longer than a heartbeat apart, rather than backing off for minutes, so a
+	// restarted Chopin hears the build again within its window.
+	expect(performance.now() - started).toBeLessThan(4000);
+	expect(renewals).toHaveLength(10);
+	// Each heartbeat names the document, so a restarted Chopin can find the build again.
+	expect(renewals.every(args => args.documentId === documentId)).toBe(true);
+	expect(calls.at(-1)?.[1]).toMatchObject({ state: "failed", error: STOPPED });
+});
+
 test("a run out of reach for its whole grace period stops its agent", async () => {
 	let { id, calls } = await run(
 		{},

apps/connector/src/run.ts

@@ -141,8 +141,11 @@ export async function runWork(
 			abort.abort();
 		},
 	});
-	// A heartbeat that fails while Chopin is out of reach retries on the reconnect backoff and
-	// leaves the agent running; only Chopin ending the run, or the grace period, stops it.
+	// A heartbeat that fails while Chopin is out of reach retries on the reconnect backoff, never
+	// waiting longer than a heartbeat apart, and leaves the agent running; only Chopin ending the
+	// run, or the grace period, stops it. A Chopin back from a restart hears the build again soon
+	// enough to let it carry on.
+	let heartbeatMs = options.heartbeatMs ?? 10_000;
 	let beating = new AbortController();
 	// Ctrl-C, the run ending and the grace period running out each abandon a heartbeat in flight.
 	let beat = AbortSignal.any([beating.signal, options.signal, reports.signal]);
@@ -151,8 +154,8 @@ export async function runWork(
 		while (!beat.aborted) {
 			await delay(
 				failures
-					? Math.min(backoff(failures - 1, options.random), reports.remaining())
-					: options.heartbeatMs ?? 10_000,
+					? Math.min(backoff(failures - 1, options.random), heartbeatMs, reports.remaining())
+					: heartbeatMs,
 				beat,
 			);
 			if (beat.aborted) return;

apps/connector/src/testing/fake-implementer.ts

@@ -1,7 +1,7 @@
 import { AgentApp, ndJsonStream, PROTOCOL_VERSION } from "@agentclientprotocol/sdk";
 import type { McpServer } from "@agentclientprotocol/sdk";
 import { spawn, spawnSync } from "node:child_process";
-import { writeFile } from "node:fs/promises";
+import { access, writeFile } from "node:fs/promises";
 import { Readable, Writable } from "node:stream";
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
@@ -198,11 +198,31 @@ let app = new AgentApp()
 					},
 				}),
 		);
-		let call = async (name: string, args: Record<string, unknown> = {}) => {
+		let once = async (name: string, args: Record<string, unknown>) => {
 			let result = await client.callTool({ name, arguments: args });
 			if (result.isError) throw new Error(JSON.stringify(result.content));
 			return result;
 		};
+		// With CHOPIN_TEST_RESUME, an agent that retries a failed call for a minute, as one does
+		// while Chopin restarts.
+		let call = async (name: string, args: Record<string, unknown> = {}) => {
+			for (let started = Date.now();;) {
+				try {
+					return await once(name, args);
+				} catch (error) {
+					if (!process.env.CHOPIN_TEST_RESUME || Date.now() - started > 60_000) throw error;
+					await new Promise(resolve => setTimeout(resolve, 500));
+				}
+			}
+		};
+		/** With CHOPIN_TEST_RESUME, say the build is under way and wait until told to go on. */
+		let pause = async () => {
+			if (!process.env.CHOPIN_TEST_RESUME) return;
+			await writeFile(process.env.CHOPIN_TEST_PAUSED!, "");
+			while (!await access(process.env.CHOPIN_TEST_RESUME).then(() => true, () => false)) {
+				await new Promise(resolve => setTimeout(resolve, 100));
+			}
+		};
 		// A completing fixture lands a rebuild with one commit on the first pull request.
 		if (rebuilding) {
 			try {
@@ -227,6 +247,7 @@ let app = new AgentApp()
 			if (process.argv.includes("--complete")) {
 				for (let [index, task] of context.graph.definition.tasks.entries()) {
 					await call("start_task", { taskId: task.id, idempotencyKey: `start-${task.id}` });
+					if (index === 0) await pause();
 					await call("report_pr", {
 						taskId: task.id,
 						idempotencyKey: `pr-${task.id}`,

apps/server/src/experiments/connections.ts

@@ -1,4 +1,4 @@
-import { randomBytes } from "node:crypto";
+import { createHmac, hkdfSync, randomBytes, timingSafeEqual } from "node:crypto";
 import type { Source } from "@chopin/experiment";
 import { fail, fingerprint } from "./service";
 
@@ -79,6 +79,18 @@ const FRESH_MS = 45_000;
 
 /** The first connector protocol whose connection is the machine rather than one checkout. */
 export const ANY_REPOSITORY_PROTOCOL = 5;
+/**
+ * The first connector protocol whose build heartbeats and reports name the document, so a server
+ * that restarted can find the build again and its machine can carry on with it.
+ */
+export const RESUME_PROTOCOL = 6;
+
+/** A build's run, as its durable credential names it. */
+export type DurableRun = { id: string; documentId: string; kind: "implementation" | "rebuild" };
+const RUN_CONTEXT = "chopin:run-credential:v1";
+const RUN_PREFIX = "chopin_r_";
+/** An HMAC-SHA256 signature in unpadded base64url. */
+const SIGNATURE_LENGTH = 43;
 
 /** Whether a connection may take work in a repository. */
 export function serves(connection: Connection, repositoryId: string) {
@@ -135,7 +147,68 @@ export class Connections {
 	#about = new Map<string, About>();
 	/** Called once with each connection that ends: disconnected, expired, or its machine revoked. */
 	onEnded?: (connection: Connection) => void;
-	constructor(private now = () => Date.now()) {}
+	/** Signs build run credentials so they can be checked again after a restart; none without a key. */
+	#runKey: Buffer | undefined;
+	constructor(
+		private now = () => Date.now(),
+		/** The deployment's SESSION_ENCRYPTION_KEY, from which build run credentials are derived. */
+		key?: Uint8Array,
+	) {
+		this.#runKey = key
+			? Buffer.from(hkdfSync("sha256", key, "", RUN_CONTEXT, 32))
+			: undefined;
+	}
+	/** Whether build run credentials outlive this process, so a build can resume after a restart. */
+	get durable() {
+		return !!this.#runKey;
+	}
+	#sign(payload: string) {
+		return createHmac("sha256", this.#runKey!).update(payload).digest();
+	}
+	/**
+	 * A build's run credential: its run, signed with a key derived from the deployment's own, so a
+	 * restarted server can check it without having stored it. The build's state still decides
+	 * whether it is accepted: it ends with the build.
+	 */
+	#durableToken(run: DurableRun) {
+		let payload = Buffer.from([run.documentId, run.id, run.kind].join("\n")).toString("base64url");
+		// The signature first, at its fixed length, keeps the whole token one base64url word.
+		return `${RUN_PREFIX}${this.#sign(payload).toString("base64url")}${payload}`;
+	}
+	/** The run a durable build credential names, when this deployment signed it. */
+	durableRun(token: string): DurableRun | undefined {
+		if (!this.#runKey || !token.startsWith(RUN_PREFIX) || token.length > 512) return undefined;
+		let body = token.slice(RUN_PREFIX.length);
+		let signature = body.slice(0, SIGNATURE_LENGTH);
+		let payload = body.slice(SIGNATURE_LENGTH);
+		if (!payload || !/^[A-Za-z0-9_-]+$/.test(body)) return undefined;
+		let given = Buffer.from(signature, "base64url");
+		let expected = this.#sign(payload);
+		if (given.length !== expected.length || !timingSafeEqual(given, expected)) return undefined;
+		let [documentId, id, kind, rest] = Buffer.from(payload, "base64url").toString("utf8")
+			.split("\n");
+		if (!documentId || !id || rest !== undefined) return undefined;
+		if (kind !== "implementation" && kind !== "rebuild") return undefined;
+		return { documentId, id, kind };
+	}
+	/**
+	 * Accept a durable build credential this process has not seen, as after a restart, for the
+	 * connection that now holds its build.
+	 */
+	adopt(token: string, connectionId: string, run: DurableRun) {
+		if (!this.#connections.has(connectionId)) return;
+		this.#runTokens.set(`${connectionId}:${run.id}:1:${run.kind}`, token);
+		this.#tokens.set(fingerprint(token), {
+			connectionId,
+			run: { id: run.id, documentId: run.documentId, generation: 1, kind: run.kind },
+		});
+		this.#remember(run.id, run.documentId);
+	}
+	/** Whether this process already holds a credential, or refused it. */
+	holds(token: string) {
+		let key = fingerprint(token);
+		return this.#tokens.has(key) || this.#revoked.has(key);
+	}
 	create(input: PairingInput) {
 		this.sweep();
 		if (this.#pending.size >= 100) fail("pairing-capacity");
@@ -236,7 +309,7 @@ export class Connections {
 	}
 	/** Whether a token is a run's, which the coding agent itself presents. */
 	run(token: string) {
-		return !!this.#tokens.get(fingerprint(token))?.run;
+		return !!this.#tokens.get(fingerprint(token))?.run || !!this.durableRun(token);
 	}
 	lookup(token: string) {
 		let refused = this.#revoked.get(fingerprint(token));
@@ -411,19 +484,25 @@ export class Connections {
 		let key = `${connectionId}:${id}:${generation}:${kind ?? "experiment"}`;
 		let existing = this.#runTokens.get(key);
 		if (existing) return existing;
-		let token = randomBytes(32).toString("base64url");
+		// A build's credential can be checked again after a restart; an investigation's cannot.
+		let token = kind && this.#runKey
+			? this.#durableToken({ id, documentId, kind })
+			: randomBytes(32).toString("base64url");
 		this.#runTokens.set(key, token);
 		this.#tokens.set(fingerprint(token), {
 			connectionId,
 			run: { id, documentId, generation, ...(kind ? { kind } : {}) },
 		});
+		this.#remember(id, documentId);
+		return token;
+	}
+	#remember(id: string, documentId: string) {
 		this.#runs.delete(id);
 		this.#runs.set(id, documentId);
 		for (let old of this.#runs.keys()) {
 			if (this.#runs.size <= RUNS) break;
 			this.#runs.delete(old);
 		}
-		return token;
 	}
 	/** The document of a run this process handed out, whichever connection holds it now. */
 	runDocument(id: string) {

apps/server/src/experiments/routes.ts

@@ -257,7 +257,7 @@ function described(live: Connection[]) {
 }
 
 export function registerExperimentRoutes(router: Router, auth: HostedAuth, options: Options) {
-	let connections = new Connections();
+	let connections = new Connections(undefined, auth.config.encryptionKey);
 	/**
 	 * The stored machine choices, read before any work is routed. A read that fails is retried a
 	 * few times, then refuses the work with `machine-choice-unavailable`; it is never taken to mean
@@ -1396,6 +1396,10 @@ export function registerExperimentRoutes(router: Router, auth: HostedAuth, optio
 	 * A connection that serves any repository is checked per document, as its work is claimed.
 	 */
 	async function connector(token: string) {
+		// A build's run credential from before a restart is accepted again for the connection that
+		// now holds its build, once its machine has come back.
+		let durable = !connections.holds(token) ? connections.durableRun(token) : undefined;
+		if (durable) await options.implementations?.()?.adopt(token, durable);
 		let found = connections.lookup(token);
 		let machine: StoredMachine;
 		try {

apps/server/src/tasks/builds.ts

@@ -57,6 +57,7 @@ let buildSchema = z.object({
 		at: z.string().datetime(),
 	}).strict().optional(),
 	reportKey: z.string().uuid().optional(),
+	resumable: z.literal(true).optional(),
 	kind: z.literal("rebuild").optional(),
 	baseRevision: z.number().int().nonnegative().optional(),
 	targetRevision: z.number().int().nonnegative().optional(),
@@ -470,6 +471,8 @@ export function pickBuild(
 	plan: Plan,
 	user: string,
 	connectionId: string,
+	/** Its run credential outlives a restart and its connector names the document as it reports. */
+	resumable = false,
 ): Promise<BuildRequest | undefined> {
 	return exclusive(plan, async () => {
 		let build = plan.builds.at(-1);
@@ -480,7 +483,12 @@ export function pickBuild(
 			return;
 		}
 		let previous = plan.builds;
-		let next: BuildRequest = { ...build, state: "starting", expiresAt: Date.now() + 60_000 };
+		let next: BuildRequest = {
+			...build,
+			state: "starting",
+			expiresAt: Date.now() + 60_000,
+			...(resumable ? { resumable: true as const } : {}),
+		};
 		plan.builds = [...previous.slice(0, -1), next];
 		try {
 			await persistExclusive(plan, true);
@@ -507,7 +515,12 @@ export function rebindBuild(
 		let build = plan.builds.at(-1);
 		if (!build || build.id !== id || build.connectionId !== from) return false;
 		let previous = plan.builds;
-		plan.builds = [...previous.slice(0, -1), { ...build, connectionId: to }];
+		// A build its machine carries on with after a restart may have gone longer than a heartbeat.
+		plan.builds = [...previous.slice(0, -1), {
+			...build,
+			connectionId: to,
+			expiresAt: Math.max(build.expiresAt, Date.now() + 60_000),
+		}];
 		try {
 			await persistExclusive(plan, true);
 		} catch (error) {