Spaces:

Mike0021
/

chat-ui-agent-demo

Paused

chat-ui-agent-demo / src /lib /server /agent /WorkspaceProvisioner.ts
deploy: production agent journey fixes (part 2)
77374bd verified 2 months ago
55.5 kB
import { randomUUID } from "node:crypto";
import type { ObjectId } from "mongodb";
import { collections } from "$lib/server/database";
import { config } from "$lib/server/config";
import {
	issueAgentProxyCredential,
	type IssuedAgentProxyCredential,
} from "$lib/server/agent/llmProxy";
import {
	SandboxClient,
	type CancelMatchingSandboxJobsOptions,
	type CreateSandboxOptions,
	type SandboxConnection,
	type SandboxExecResult,
	type SandboxJobInfo,
	type SandboxProcess,
} from "$lib/server/agent/SandboxClient";
import type {
	Workspace,
	WorkspaceLlmProxy,
	WorkspacePiSession,
	WorkspaceSandbox,
	WorkspaceSandboxCleanup,
	WorkspaceSandboxStatus,
} from "$lib/types/Workspace";

export const AGENT_RUNTIME_IMAGE = "hf.co/spaces/Mike0021/chat-ui-agent-runtime";
export const AGENT_BRIDGE_PORT = 4_777;
export const AGENT_PROVIDER_ID = "chat-ui-agent";
export const AGENT_WORKSPACE_DIR = "/workspace";
export const AGENT_POLICY_PATH = "/run/chat-ui-agent/policy";

const AGENT_CONFIG_DIR = "/root/.pi/agent";
const AGENT_MODELS_PATH = `${AGENT_CONFIG_DIR}/models.json`;
const AGENT_SESSION_DIR = `${AGENT_CONFIG_DIR}/sessions`;
const AGENT_BRIDGE_PATH = "/opt/chat-ui-agent/bridge.mjs";
const AGENT_APPROVAL_EXTENSION_PATH = "/opt/chat-ui-agent/approval-gate.mjs";
const DEFAULT_HUB_URL = "https://huggingface.co";
const DEFAULT_BRANCH = "main";
const BRIDGE_START_TIMEOUT_MS = 20_000;
const MAX_STORED_ERROR_CHARS = 1_000;
const PROVISIONING_LEASE_MS = 15 * 60 * 1_000;
const TERMINAL_SANDBOX_STAGES = new Set(["COMPLETED", "ERROR", "DELETED", "CANCELED"]);
export const BLANK_WORKSPACE_GITIGNORE = `# Local dependency trees and generated caches
node_modules/
.venv/
venv/
__pypackages__/
__pycache__/
*.pyc
*.pyo
.pytest_cache/
.mypy_cache/
.ruff_cache/
.coverage
# Ignore conventional local credential files as defense in depth
.env
.env.*
!.env.example
!.env.sample
!.env.template
`;

export type WorkspaceProvisioningPhase =
	"claiming" | "sandbox" | "repository" | "credentials" | "bridge" | "ready" | "cleanup";

export interface WorkspaceProvisioningProgress {
	phase: WorkspaceProvisioningPhase;
	message: string;
}

export type WorkspaceProvisioningCallback = (
	progress: WorkspaceProvisioningProgress
) => void | Promise<void>;

export interface ProvisionWorkspaceOptions {
	workspace: Workspace;
	/** Models exposed to pi and allowed through this workspace's proxy. */
	availableModels?: AgentRuntimeModelDefinition[];
	/**
	 * Current request's HF credential may be used only as a server-side fallback for
	 * the Jobs lifecycle. Source checkout auth is resolved later by the Git proxy and
	 * no raw Hugging Face credential is ever copied into the sandbox.
	 */
	requestHfToken?: string;
	/** Optional override; otherwise AGENT_SANDBOX_NAMESPACE or token whoami is used. */
	namespace?: string;
	/** Public chat-ui origin reachable from the sandbox. */
	publicOrigin?: string;
	runtimeImage?: string;
	hubUrl?: string;
	onProgress?: WorkspaceProvisioningCallback;
}

export interface ReconnectWorkspaceOptions {
	workspace: Workspace;
	requestHfToken?: string;
	hubUrl?: string;
	onProgress?: WorkspaceProvisioningCallback;
	/** Health-check only: never rewrite lifecycle state or top-level revision. */
	observational?: boolean;
}

export interface KillWorkspaceOptions {
	workspace: Workspace;
	requestHfToken?: string;
	hubUrl?: string;
	onProgress?: WorkspaceProvisioningCallback;
}

export interface RestartWorkspaceOptions extends ProvisionWorkspaceOptions {}

export interface AgentRuntimeModelDefinition {
	id: string;
	name?: string;
	reasoning?: boolean;
	multimodal?: boolean;
	contextWindow?: number;
	maxTokens?: number;
	cost?: { input: number; output: number; cacheRead?: number; cacheWrite?: number };
}

/** Server-only runtime handle. Its SandboxClient contains the Jobs credential. */
export interface WorkspaceRuntimeConnection {
	workspace: Workspace;
	sandbox: SandboxClient;
	bridgeUrl: string;
}

interface RuntimePatch {
	sandbox?: WorkspaceSandbox;
	piSession?: WorkspacePiSession | null;
	llmProxy?: WorkspaceLlmProxy | null;
	updatedAt: Date;
}

export interface WorkspaceSandboxExpectation {
	statuses: readonly WorkspaceSandboxStatus[];
	id?: string;
	nonce?: string;
	generation?: string;
	/** Undefined ignores cleanup; null requires none; an object requires exact identity. */
	cleanup?: WorkspaceSandboxCleanup | null;
	/** Optional document revision fence used by provisioning claims. */
	updatedAt?: Date;
	/** Undefined ignores retirement; null requires no lease; an id requires exact ownership. */
	retirement?: { id: string } | null;
}

export interface WorkspaceProvisionerStore {
	claim(
		workspaceId: ObjectId,
		expected: WorkspaceSandboxExpectation,
		llmProxy: WorkspaceLlmProxy,
		generation: string,
		cleanup: WorkspaceSandboxCleanup | undefined,
		now: Date,
		staleBefore: Date
	): Promise<Workspace | null>;
	claimStop(
		workspaceId: ObjectId,
		expected: WorkspaceSandboxExpectation,
		generation: string,
		cleanup: WorkspaceSandboxCleanup | undefined,
		now: Date
	): Promise<Workspace | null>;
	get(workspaceId: ObjectId): Promise<Workspace | null>;
	update(
		workspaceId: ObjectId,
		patch: RuntimePatch,
		expected: WorkspaceSandboxExpectation
	): Promise<boolean>;
}

interface SandboxRuntime {
	readonly connection: SandboxConnection;
	readonly proxyHeaders: Record<string, string>;
	health(): Promise<boolean>;
	inspect(): Promise<SandboxJobInfo>;
	exec(
		cmd: string | string[],
		options?: {
			cwd?: string;
			environment?: Record<string, string>;
			timeoutSeconds?: number;
			stdin?: string;
		}
	): Promise<SandboxExecResult>;
	spawn(
		cmd: string | string[],
		options?: { cwd?: string; environment?: Record<string, string>; tag?: string }
	): Promise<SandboxProcess>;
	listProcesses(): Promise<SandboxProcess[]>;
	writeFile(path: string, content: string | ArrayBuffer, mode?: string): Promise<void>;
	proxy(port: number, path: string, init?: RequestInit): Promise<Response>;
	proxyUrlFor(port: number, path?: string, scheme?: "https" | "wss"): string;
	kill(): Promise<void>;
}

export interface WorkspaceProvisionerDependencies {
	store?: WorkspaceProvisionerStore;
	createSandbox?: (options: CreateSandboxOptions) => Promise<SandboxRuntime>;
	cancelMatchingSandboxJobs?: (
		options: CancelMatchingSandboxJobsOptions
	) => Promise<readonly string[]>;
	connectSandbox?: (
		connection: SandboxConnection,
		hfToken: string,
		options?: { hubUrl?: string; fetch?: typeof fetch }
	) => SandboxRuntime;
	issueCredential?: typeof issueAgentProxyCredential;
	fetch?: typeof fetch;
	now?: () => Date;
	randomUuid?: () => string;
	delay?: (milliseconds: number) => Promise<void>;
	configValue?: (name: string) => string;
}

export class WorkspaceProvisioningError extends Error {
	constructor(
		message: string,
		readonly code:
			| "invalid_configuration"
			| "invalid_repository"
			| "already_provisioning"
			| "sandbox_unavailable"
			| "workspace_missing" = "sandbox_unavailable"
	) {
		super(message);
		this.name = "WorkspaceProvisioningError";
	}
}

export function isRetryableWorkspaceProvisioningFailure(cause: unknown): boolean {
	return !(cause instanceof WorkspaceProvisioningError) || cause.code === "sandbox_unavailable";
}

/**
 * Owns the credential boundary and lifecycle of one dedicated HF Sandbox.
 * The Jobs token is only ever passed to the Hub/SandboxClient. The sandbox
 * receives a short-lived agent proxy credential whose hash alone is durable.
 */
export class WorkspaceProvisioner {
	private readonly store: WorkspaceProvisionerStore;
	private readonly createSandbox: (options: CreateSandboxOptions) => Promise<SandboxRuntime>;
	private readonly cancelMatchingSandboxJobs: NonNullable<
		WorkspaceProvisionerDependencies["cancelMatchingSandboxJobs"]
	>;
	private readonly connectSandbox: NonNullable<WorkspaceProvisionerDependencies["connectSandbox"]>;
	private readonly issueCredential: typeof issueAgentProxyCredential;
	private readonly requestFetch: typeof fetch;
	private readonly now: () => Date;
	private readonly randomUuid: () => string;
	private readonly wait: (milliseconds: number) => Promise<void>;
	private readonly configValue: (name: string) => string;

	constructor(dependencies: WorkspaceProvisionerDependencies = {}) {
		this.store = dependencies.store ?? mongoWorkspaceStore;
		this.createSandbox = dependencies.createSandbox ?? ((options) => SandboxClient.create(options));
		this.cancelMatchingSandboxJobs =
			dependencies.cancelMatchingSandboxJobs ??
			((options) => SandboxClient.cancelMatchingJobs(options));
		this.connectSandbox =
			dependencies.connectSandbox ??
			((connection, hfToken, options) => SandboxClient.connect(connection, hfToken, options));
		this.issueCredential = dependencies.issueCredential ?? issueAgentProxyCredential;
		this.requestFetch = dependencies.fetch ?? fetch;
		this.now = dependencies.now ?? (() => new Date());
		this.randomUuid = dependencies.randomUuid ?? randomUUID;
		this.wait = dependencies.delay ?? delay;
		this.configValue = dependencies.configValue ?? readConfigValue;
	}

	async provision(options: ProvisionWorkspaceOptions): Promise<WorkspaceRuntimeConnection> {
		if (options.workspace.retirement) {
			throw new WorkspaceProvisioningError(
				"Agent workspace is being retired",
				"sandbox_unavailable"
			);
		}
		if (
			options.workspace.sandbox.status === "ready" ||
			options.workspace.sandbox.status === "idle"
		) {
			return this.reconnect(options);
		}
		const claimNow = this.now();
		if (
			options.workspace.sandbox.status === "provisioning" &&
			options.workspace.sandbox.claimedAt &&
			options.workspace.sandbox.claimedAt.getTime() > claimNow.getTime() - PROVISIONING_LEASE_MS
		) {
			throw new WorkspaceProvisioningError(
				"Another request is already provisioning this workspace",
				"already_provisioning"
			);
		}

		const jobsToken = this.resolveJobsToken(options.requestHfToken);
		const hubUrl = trimTrailingSlash(options.hubUrl || DEFAULT_HUB_URL);
		const publicOrigin = this.resolvePublicOrigin(options.publicOrigin);
		const runtimeImage =
			options.runtimeImage || this.configValue("AGENT_RUNTIME_IMAGE") || AGENT_RUNTIME_IMAGE;
		validateRepository(options.workspace);
		const requestedRuntimeModels = normalizeRuntimeModels(
			options.workspace,
			options.availableModels
		);
		let namespace = options.namespace || this.configValue("AGENT_SANDBOX_NAMESPACE") || undefined;
		const requireNamespace = async (): Promise<string> =>
			(namespace ??= await this.resolveNamespace(jobsToken, hubUrl));
		const abandonedGeneration = abandonedUnpersistedGeneration(
			options.workspace,
			claimNow,
			PROVISIONING_LEASE_MS
		);
		const cleanup =
			options.workspace.sandbox.cleanup ??
			(abandonedGeneration
				? {
						generation: abandonedGeneration,
						namespace: options.workspace.sandbox.namespace || (await requireNamespace()),
						hubUrl,
					}
				: undefined);
		const credential = await this.issueCredential(requestedRuntimeModels.map((model) => model.id));
		const persistedCredential = toPersistedCredential(credential);

		await emit(options.onProgress, "claiming", "Reserving agent workspace");
		const generation = this.randomUuid();
		let claimed = await this.store.claim(
			options.workspace._id,
			provisionClaimExpectation(options.workspace),
			persistedCredential,
			generation,
			cleanup,
			claimNow,
			new Date(claimNow.getTime() - PROVISIONING_LEASE_MS)
		);
		if (!claimed) {
			const current = await this.store.get(options.workspace._id);
			if (!current) {
				throw new WorkspaceProvisioningError(
					"Agent workspace no longer exists",
					"workspace_missing"
				);
			}
			if (
				(current.sandbox.status === "ready" || current.sandbox.status === "idle") &&
				current.sandbox.id &&
				!current.retirement
			) {
				return this.reconnect({ ...options, workspace: current });
			}
			throw new WorkspaceProvisioningError(
				"Another request is already provisioning this workspace",
				"already_provisioning"
			);
		}
		let expectedSandbox: WorkspaceSandboxExpectation = {
			...sandboxExpectation(claimed.sandbox, ["provisioning"]),
			retirement: null,
		};
		if (claimed.sandbox.cleanup) {
			const cleanupTarget = claimed.sandbox.cleanup;
			await emit(options.onProgress, "cleanup", "Stopping abandoned HF Sandbox");
			try {
				await this.cancelMatchingSandboxJobs({
					hfToken: jobsToken,
					namespace: cleanupTarget.namespace,
					labels: sandboxOwnershipLabels(claimed._id, cleanupTarget.generation),
					hubUrl: cleanupTarget.hubUrl,
					fetch: this.requestFetch,
				});
			} catch (error) {
				const failedAt = this.now();
				const safeMessage = redactError(error, [jobsToken, options.requestHfToken ?? ""]);
				await this.store
					.update(
						claimed._id,
						{
							sandbox: {
								...claimed.sandbox,
								status: "error",
								flavor: claimed.flavor,
								lastSeenAt: failedAt,
								error: safeMessage,
							},
							llmProxy: null,
							updatedAt: failedAt,
						},
						expectedSandbox
					)
					.catch(() => undefined);
				throw new WorkspaceProvisioningError(safeMessage, "sandbox_unavailable");
			}
			const cleanupClearedAt = this.now();
			const cleanupClearedSandbox = { ...claimed.sandbox };
			delete cleanupClearedSandbox.cleanup;
			await this.updateRuntime(
				claimed._id,
				{ sandbox: cleanupClearedSandbox, updatedAt: cleanupClearedAt },
				expectedSandbox
			);
			claimed = { ...claimed, sandbox: cleanupClearedSandbox, updatedAt: cleanupClearedAt };
			expectedSandbox = {
				...sandboxExpectation(cleanupClearedSandbox, ["provisioning"]),
				retirement: null,
			};
		}

		if (hasCompleteConnection(claimed.sandbox)) {
			await emit(options.onProgress, "cleanup", "Stopping previous HF Sandbox");
			try {
				const previousSandbox = this.connectSandbox(connectionFromWorkspace(claimed), jobsToken, {
					hubUrl,
					fetch: this.requestFetch,
				});
				await previousSandbox.kill();
			} catch (error) {
				const failedAt = this.now();
				const safeMessage = redactError(error, [
					jobsToken,
					options.requestHfToken ?? "",
					credential.token,
				]);
				await this.updateRuntime(
					claimed._id,
					{
						sandbox: {
							...claimed.sandbox,
							status: "error",
							lastSeenAt: failedAt,
							error: safeMessage,
						},
						llmProxy: null,
						updatedAt: failedAt,
					},
					expectedSandbox
				);
				throw new WorkspaceProvisioningError(safeMessage, "sandbox_unavailable");
			}
		}

		let sandbox: SandboxRuntime | undefined;
		let createdConnection: SandboxConnection | undefined;
		let connectionPersisted = false;
		try {
			// Source and model configuration must come from the record protected by
			// the successful claim, never from the caller's potentially stale copy.
			const repository = validateRepository(claimed);
			const runtimeModels = normalizeRuntimeModels(claimed, options.availableModels);
			const persistCreatedConnection = async (connection: SandboxConnection): Promise<void> => {
				if (createdConnection && !sameSandboxConnection(createdConnection, connection)) {
					throw new WorkspaceProvisioningError(
						"HF Sandbox creation returned inconsistent connection metadata",
						"sandbox_unavailable"
					);
				}
				createdConnection = connection;
				if (connectionPersisted) return;
				const startedAt = this.now();
				const provisioningSandbox = workspaceSandboxFromConnection(
					connection,
					"provisioning",
					startedAt,
					generation,
					claimNow
				);
				await this.updateRuntime(
					claimed._id,
					{
						sandbox: provisioningSandbox,
						updatedAt: startedAt,
					},
					expectedSandbox
				);
				expectedSandbox = {
					...sandboxExpectation(provisioningSandbox, ["provisioning"]),
					retirement: null,
				};
				connectionPersisted = true;
			};
			await emit(options.onProgress, "sandbox", "Starting dedicated HF Sandbox");
			sandbox = await this.createSandbox({
				hfToken: jobsToken,
				namespace: await requireNamespace(),
				image: runtimeImage,
				flavor: claimed.flavor,
				labels: sandboxOwnershipLabels(claimed._id, generation),
				hubUrl,
				onStatus: async (message) => emit(options.onProgress, "sandbox", message),
				onCreated: persistCreatedConnection,
			});
			// Keep injected implementations compatible, while the production client
			// persists from onCreated before it begins readiness polling.
			await persistCreatedConnection(sandbox.connection);

			await emit(options.onProgress, "repository", "Preparing repository");
			await this.prepareRepository(
				sandbox,
				repository,
				claimed.branch,
				publicOrigin,
				credential.token
			);

			const piSession: WorkspacePiSession = {
				id: claimed.piSession?.id || this.randomUuid(),
				dir: AGENT_SESSION_DIR,
				// A new bridge owns a new in-memory replay buffer whose first seq is 1.
				lastSequence: 0,
			};
			await emit(options.onProgress, "credentials", "Configuring inference proxy");
			await this.configurePi(sandbox, claimed, credential, publicOrigin, runtimeModels);

			await emit(options.onProgress, "bridge", "Starting coding agent");
			const process = await sandbox.spawn(["node", AGENT_BRIDGE_PATH], {
				cwd: AGENT_WORKSPACE_DIR,
				tag: `chat-ui-agent:${claimed._id.toString()}`,
				environment: bridgeEnvironment(claimed, piSession),
			});
			if (!Number.isSafeInteger(process.pid) || process.pid <= 0) {
				throw new WorkspaceProvisioningError(
					"Sandbox did not return a coding-agent process id",
					"sandbox_unavailable"
				);
			}
			piSession.bridgePid = process.pid;
			await this.waitForBridge(sandbox);

			const readyAt = this.now();
			const readySandbox = workspaceSandboxFromConnection(
				sandbox.connection,
				"ready",
				readyAt,
				generation
			);
			const readyExpectation: WorkspaceSandboxExpectation = {
				...liveSandboxExpectation(readySandbox),
				retirement: null,
			};
			await this.updateRuntime(
				claimed._id,
				{
					sandbox: readySandbox,
					piSession,
					updatedAt: readyAt,
				},
				expectedSandbox
			);
			const readyWorkspace = await this.requireWorkspace(claimed._id);
			if (!matchesWorkspaceExpectation(readyWorkspace, readyExpectation)) {
				await this.throwWorkspaceChanged(claimed._id);
			}
			await emit(options.onProgress, "ready", "Agent workspace ready");
			return runtimeConnection(readyWorkspace, sandbox);
		} catch (error) {
			await emit(options.onProgress, "cleanup", "Cleaning up failed sandbox");
			let cleanupError: unknown;
			try {
				const cleanupSandbox =
					sandbox ??
					(createdConnection
						? this.connectSandbox(createdConnection, jobsToken, {
								hubUrl,
								fetch: this.requestFetch,
							})
						: undefined);
				await cleanupSandbox?.kill();
			} catch (killError) {
				cleanupError = killError;
			}
			const failedAt = this.now();
			const originalMessage = redactError(error, [
				jobsToken,
				options.requestHfToken ?? "",
				credential.token,
			]);
			const cleanupMessage = cleanupError
				? redactError(cleanupError, [jobsToken, options.requestHfToken ?? "", credential.token])
				: undefined;
			const safeMessage = cleanupMessage
				? `${originalMessage}. Automatic sandbox cleanup failed: ${cleanupMessage}`.slice(
						0,
						MAX_STORED_ERROR_CHARS
					)
				: originalMessage;
			await this.store
				.update(
					claimed._id,
					{
						sandbox:
							createdConnection && cleanupError
								? {
										...workspaceSandboxFromConnection(
											createdConnection,
											"provisioning",
											failedAt,
											generation,
											claimNow
										),
										status: "error",
										error: safeMessage,
									}
								: {
										status: "error",
										generation,
										claimedAt: claimNow,
										flavor: claimed.flavor,
										image: runtimeImage,
										lastSeenAt: failedAt,
										error: safeMessage,
									},
						llmProxy: null,
						updatedAt: failedAt,
					},
					expectedSandbox
				)
				.catch(() => undefined);
			throw new WorkspaceProvisioningError(
				safeMessage,
				error instanceof WorkspaceProvisioningError ? error.code : "sandbox_unavailable"
			);
		}
	}

	async reconnect(options: ReconnectWorkspaceOptions): Promise<WorkspaceRuntimeConnection> {
		if (
			options.workspace.sandbox.status !== "ready" &&
			options.workspace.sandbox.status !== "idle"
		) {
			throw new WorkspaceProvisioningError(
				`Workspace sandbox is not reconnectable (${options.workspace.sandbox.status})`,
				options.workspace.sandbox.status === "provisioning"
					? "already_provisioning"
					: "sandbox_unavailable"
			);
		}
		const jobsToken = this.resolveJobsToken(options.requestHfToken);
		const connection = connectionFromWorkspace(options.workspace);
		await emit(options.onProgress, "sandbox", "Reconnecting to HF Sandbox");
		const sandbox = this.connectSandbox(connection, jobsToken, {
			hubUrl: trimTrailingSlash(options.hubUrl || DEFAULT_HUB_URL),
			fetch: this.requestFetch,
		});
		const job = await sandbox.inspect().catch(() => undefined);
		if (job && TERMINAL_SANDBOX_STAGES.has(job.stage.toUpperCase())) {
			if (!options.observational) {
				await this.retireSandbox(
					options.workspace,
					sandbox,
					jobsToken,
					"HF Sandbox has stopped",
					liveWorkspaceExpectation(options.workspace),
					options.onProgress
				);
			}
			throw new WorkspaceProvisioningError("HF Sandbox has stopped", "sandbox_unavailable");
		}

		// HF edge routes can briefly return 404 independently. Only treat a
		// missing bridge as authoritative when the control route is healthy and
		// its process list confirms the persisted PID is gone. Otherwise let the
		// actual operation use its own bounded transport reconnect.
		if (await sandbox.health()) {
			const bridgePid = options.workspace.piSession?.bridgePid;
			if (Number.isSafeInteger(bridgePid) && bridgePid && bridgePid > 0) {
				const processes = await sandbox.listProcesses().catch(() => undefined);
				const bridge = processes?.find((process) => process.pid === bridgePid);
				if (processes && (!bridge || bridge.running === false || bridge.exitCode != null)) {
					if (!options.observational) {
						await this.retireSandbox(
							options.workspace,
							sandbox,
							jobsToken,
							"Coding-agent bridge has stopped",
							liveWorkspaceExpectation(options.workspace),
							options.onProgress
						);
					}
					throw new WorkspaceProvisioningError(
						"Coding-agent bridge is not running",
						"sandbox_unavailable"
					);
				}
			}
		}
		if (options.observational) {
			const current = await this.requireWorkspace(options.workspace._id);
			if (!matchesWorkspaceExpectation(current, liveWorkspaceExpectation(options.workspace))) {
				await this.throwWorkspaceChanged(options.workspace._id);
			}
			return runtimeConnection(current, sandbox);
		}
		const readyAt = this.now();
		const healthySandbox = { ...options.workspace.sandbox };
		delete healthySandbox.error;
		const reconnectExpectation = liveWorkspaceExpectation(options.workspace);
		await this.updateRuntime(
			options.workspace._id,
			{
				sandbox: {
					...healthySandbox,
					status: "ready",
					lastSeenAt: readyAt,
				},
				updatedAt: readyAt,
			},
			reconnectExpectation
		);
		const workspace = await this.requireWorkspace(options.workspace._id);
		if (!matchesWorkspaceExpectation(workspace, reconnectExpectation)) {
			await this.throwWorkspaceChanged(options.workspace._id);
		}
		await emit(options.onProgress, "ready", "Agent workspace reconnected");
		return runtimeConnection(workspace, sandbox);
	}

	async kill(options: KillWorkspaceOptions): Promise<Workspace> {
		let stopRuntime: { jobsToken: string; sandbox: SandboxRuntime } | undefined;
		let cleanup = options.workspace.sandbox.cleanup;
		const sandboxStatus = options.workspace.sandbox.status;
		const hasAttachedJobMetadata = Boolean(
			options.workspace.sandbox.id ||
			options.workspace.sandbox.namespace ||
			options.workspace.sandbox.baseUrl ||
			options.workspace.sandbox.nonce
		);
		const hubUrl = trimTrailingSlash(options.hubUrl || DEFAULT_HUB_URL);
		if (
			sandboxStatus === "ready" ||
			sandboxStatus === "idle" ||
			(sandboxStatus !== "dead" && hasAttachedJobMetadata)
		) {
			const connection = connectionFromWorkspace(options.workspace);
			const jobsToken = this.resolveJobsToken(options.requestHfToken);
			const sandbox = this.connectSandbox(connection, jobsToken, {
				hubUrl,
				fetch: this.requestFetch,
			});
			stopRuntime = { jobsToken, sandbox };
		} else if (
			sandboxStatus !== "dead" &&
			sandboxStatus !== "unprovisioned" &&
			options.workspace.sandbox.generation
		) {
			// A process can die after Jobs accepts creation but before onCreated
			// persists the connection. Exact ownership labels are the only durable
			// handle in that interval, so preflight credentials/namespace before
			// claiming and cancel the old generation only after the CAS succeeds.
			const jobsToken = this.resolveJobsToken(options.requestHfToken);
			const namespace =
				options.workspace.sandbox.namespace ||
				this.configValue("AGENT_SANDBOX_NAMESPACE") ||
				(await this.resolveNamespace(jobsToken, hubUrl));
			cleanup ??= {
				namespace,
				generation: options.workspace.sandbox.generation,
				hubUrl,
			};
		}
		const cleanupJobsToken = cleanup
			? (stopRuntime?.jobsToken ?? this.resolveJobsToken(options.requestHfToken))
			: undefined;

		const claimedAt = this.now();
		const generation = this.randomUuid();
		let claimed = await this.store.claimStop(
			options.workspace._id,
			stopClaimExpectation(options.workspace),
			generation,
			cleanup,
			claimedAt
		);
		if (!claimed) {
			return this.throwWorkspaceChanged(options.workspace._id);
		}
		let stopExpectation = stopClaimExpectation(claimed, ["provisioning"]);
		if (claimed.sandbox.cleanup) {
			const cleanupTarget = claimed.sandbox.cleanup;
			const cleanupToken = cleanupJobsToken ?? this.resolveJobsToken(options.requestHfToken);
			await emit(options.onProgress, "cleanup", "Stopping unattached HF Sandbox");
			const stoppedAt = this.now();
			try {
				await this.cancelMatchingSandboxJobs({
					hfToken: cleanupToken,
					namespace: cleanupTarget.namespace,
					labels: sandboxOwnershipLabels(claimed._id, cleanupTarget.generation),
					hubUrl: cleanupTarget.hubUrl,
					fetch: this.requestFetch,
				});
			} catch (error) {
				const safeMessage = redactError(error, [cleanupToken, options.requestHfToken ?? ""]);
				await this.updateRuntime(
					claimed._id,
					{
						sandbox: {
							...claimed.sandbox,
							status: "error",
							flavor: claimed.flavor,
							lastSeenAt: stoppedAt,
							error: `HF Sandbox stop failed: ${safeMessage}`.slice(0, MAX_STORED_ERROR_CHARS),
						},
						piSession: claimed.piSession ? { ...claimed.piSession, bridgePid: undefined } : null,
						llmProxy: null,
						updatedAt: stoppedAt,
					},
					stopExpectation
				);
				throw new WorkspaceProvisioningError(safeMessage, "sandbox_unavailable");
			}
			const cleanupClearedSandbox = { ...claimed.sandbox };
			delete cleanupClearedSandbox.cleanup;
			await this.updateRuntime(
				claimed._id,
				{ sandbox: cleanupClearedSandbox, updatedAt: stoppedAt },
				stopExpectation
			);
			claimed = { ...claimed, sandbox: cleanupClearedSandbox, updatedAt: stoppedAt };
			stopExpectation = stopClaimExpectation(claimed, ["provisioning"]);
		}
		if (stopRuntime) {
			return this.retireSandbox(
				claimed,
				stopRuntime.sandbox,
				stopRuntime.jobsToken,
				"HF Sandbox stopped",
				stopExpectation,
				options.onProgress
			);
		}
		const now = this.now();
		const stoppedSandbox = stoppedSandboxMetadata(claimed.sandbox, now);
		await emit(options.onProgress, "cleanup", "Stopping HF Sandbox");
		await this.updateRuntime(
			claimed._id,
			{
				sandbox: stoppedSandbox,
				piSession: claimed.piSession ? { ...claimed.piSession, bridgePid: undefined } : null,
				llmProxy: null,
				updatedAt: now,
			},
			stopExpectation
		);
		return this.requireExpectedWorkspace(
			claimed._id,
			stopClaimExpectation({ ...claimed, sandbox: stoppedSandbox })
		);
	}

	async restart(options: RestartWorkspaceOptions): Promise<WorkspaceRuntimeConnection> {
		const dead = await this.kill(options);
		return this.provision({ ...options, workspace: dead });
	}

	private resolveJobsToken(requestHfToken?: string): string {
		const token =
			this.configValue("AGENT_HF_TOKEN") || this.configValue("HF_TOKEN") || requestHfToken;
		if (!token) {
			throw new WorkspaceProvisioningError(
				"Agent Sandboxes require AGENT_HF_TOKEN, HF_TOKEN, or an authenticated HF session",
				"invalid_configuration"
			);
		}
		return token;
	}

	private resolvePublicOrigin(override?: string): string {
		const configured = override || this.configValue("PUBLIC_ORIGIN");
		if (!configured) {
			throw new WorkspaceProvisioningError(
				"PUBLIC_ORIGIN is required so the sandbox can reach the inference proxy",
				"invalid_configuration"
			);
		}
		let origin: URL;
		try {
			origin = new URL(configured);
		} catch {
			throw new WorkspaceProvisioningError(
				"PUBLIC_ORIGIN must be an absolute HTTP(S) URL",
				"invalid_configuration"
			);
		}
		if (!["http:", "https:"].includes(origin.protocol) || origin.username || origin.password) {
			throw new WorkspaceProvisioningError(
				"PUBLIC_ORIGIN must be a credential-free HTTP(S) URL",
				"invalid_configuration"
			);
		}
		if (origin.search || origin.hash) {
			throw new WorkspaceProvisioningError(
				"PUBLIC_ORIGIN must not contain a query or fragment",
				"invalid_configuration"
			);
		}
		return trimTrailingSlash(origin.toString());
	}

	private async resolveNamespace(token: string, hubUrl: string): Promise<string> {
		const response = await this.requestFetch(`${hubUrl}/api/whoami-v2`, {
			headers: { Authorization: `Bearer ${token}` },
			redirect: "manual",
			signal: AbortSignal.timeout(10_000),
		});
		if (!response.ok) {
			throw new WorkspaceProvisioningError(
				"Unable to resolve the HF Jobs namespace",
				"invalid_configuration"
			);
		}
		const whoami = (await response.json()) as { name?: unknown };
		if (typeof whoami.name !== "string" || !whoami.name.trim()) {
			throw new WorkspaceProvisioningError(
				"HF authentication returned no Jobs namespace",
				"invalid_configuration"
			);
		}
		return whoami.name;
	}

	private async prepareRepository(
		sandbox: SandboxRuntime,
		repository: ValidatedRepository,
		branch: string | undefined,
		publicOrigin: string,
		agentSessionToken: string
	): Promise<void> {
		const selectedBranch = validateBranch(branch || DEFAULT_BRANCH);
		if (repository.kind === "git" || repository.kind === "hf") {
			const cloneUrl =
				repository.kind === "hf" ? `${publicOrigin}/api/agent/git/repository.git` : repository.url;
			const args = ["git", "clone", "--origin", "origin"];
			if (branch || repository.kind === "hf") {
				args.push("--branch", selectedBranch, "--single-branch");
			}
			args.push("--", cloneUrl, AGENT_WORKSPACE_DIR);
			await assertExec(
				sandbox,
				args,
				repository.kind === "hf" ? "clone Hugging Face repository" : "clone public Git repository",
				300,
				gitCloneEnvironment(
					repository.kind === "hf" ? cloneUrl : undefined,
					repository.kind === "hf" ? agentSessionToken : undefined
				),
				repository.kind === "hf" ? mapHuggingFaceCloneFailure : undefined
			);
		} else {
			await assertExec(sandbox, ["mkdir", "-p", AGENT_WORKSPACE_DIR], "create workspace directory");
			await assertExec(
				sandbox,
				["git", "init", "-b", selectedBranch, "--", AGENT_WORKSPACE_DIR],
				"initialize Git repository"
			);
		}

		await assertExec(
			sandbox,
			["git", "config", "--global", "--add", "safe.directory", AGENT_WORKSPACE_DIR],
			"trust dedicated Git workspace"
		);
		await assertExec(
			sandbox,
			["git", "-C", AGENT_WORKSPACE_DIR, "config", "user.name", "Chat UI Agent"],
			"configure Git author"
		);
		await assertExec(
			sandbox,
			["git", "-C", AGENT_WORKSPACE_DIR, "config", "user.email", "agent@chat-ui.local"],
			"configure Git email"
		);

		if (repository.kind === "blank") {
			await sandbox.writeFile(
				`${AGENT_WORKSPACE_DIR}/.gitignore`,
				BLANK_WORKSPACE_GITIGNORE,
				"0644"
			);
			await assertExec(
				sandbox,
				["git", "-C", AGENT_WORKSPACE_DIR, "add", "-A", "--"],
				"stage workspace baseline",
				300
			);
			await assertExec(
				sandbox,
				[
					"git",
					"-C",
					AGENT_WORKSPACE_DIR,
					"commit",
					"--allow-empty",
					"--no-gpg-sign",
					"-m",
					"Initialize agent workspace",
				],
				"commit workspace baseline",
				300
			);
		}
	}

	private async configurePi(
		sandbox: SandboxRuntime,
		workspace: Workspace,
		credential: IssuedAgentProxyCredential,
		publicOrigin: string,
		runtimeModels: AgentRuntimeModelDefinition[]
	): Promise<void> {
		await assertExec(
			sandbox,
			["mkdir", "-p", AGENT_CONFIG_DIR, AGENT_SESSION_DIR, "/run/chat-ui-agent"],
			"create pi configuration directory"
		);
		await sandbox.writeFile(AGENT_POLICY_PATH, `${workspace.policy}\n`, "0600");
		// Shape verified against @earendil-works/pi-coding-agent 0.84.2 docs/models.md.
		const models = {
			providers: {
				[AGENT_PROVIDER_ID]: {
					baseUrl: `${publicOrigin}/api/agent/llm`,
					api: "openai-completions",
					// A non-secret placeholder makes the custom model available to pi.
					apiKey: "agent-proxy",
					authHeader: false,
					headers: { "X-Agent-Session": credential.token },
					models: runtimeModels.map((model) => ({
						id: model.id,
						name: model.name || model.id,
						reasoning:
							model.id === workspace.model
								? workspace.thinkingLevel !== "off"
								: Boolean(model.reasoning),
						input: model.multimodal ? ["text", "image"] : ["text"],
						contextWindow: model.contextWindow ?? 128_000,
						maxTokens: model.maxTokens ?? 16_384,
						cost: {
							input: model.cost?.input ?? 0,
							output: model.cost?.output ?? 0,
							cacheRead: model.cost?.cacheRead ?? 0,
							cacheWrite: model.cost?.cacheWrite ?? 0,
						},
					})),
				},
			},
		};
		await sandbox.writeFile(AGENT_MODELS_PATH, `${JSON.stringify(models, null, 2)}\n`, "0600");
	}

	private async waitForBridge(sandbox: SandboxRuntime): Promise<void> {
		const deadline = Date.now() + BRIDGE_START_TIMEOUT_MS;
		let lastStatus = 0;
		while (Date.now() < deadline) {
			try {
				const response = await sandbox.proxy(AGENT_BRIDGE_PORT, "/health", {
					signal: AbortSignal.timeout(2_000),
				});
				lastStatus = response.status;
				const available = response.ok;
				await response.body?.cancel().catch(() => undefined);
				if (available) return;
			} catch {
				// The bridge port is expected to refuse connections during its short startup.
			}
			await this.wait(100);
		}
		throw new WorkspaceProvisioningError(
			`Coding-agent bridge did not become ready${lastStatus ? ` (HTTP ${lastStatus})` : ""}`,
			"sandbox_unavailable"
		);
	}

	private async requireWorkspace(workspaceId: ObjectId): Promise<Workspace> {
		const workspace = await this.store.get(workspaceId);
		if (!workspace) {
			throw new WorkspaceProvisioningError("Agent workspace no longer exists", "workspace_missing");
		}
		return workspace;
	}

	private async requireExpectedWorkspace(
		workspaceId: ObjectId,
		expected: WorkspaceSandboxExpectation
	): Promise<Workspace> {
		const workspace = await this.requireWorkspace(workspaceId);
		if (!matchesWorkspaceExpectation(workspace, expected)) {
			await this.throwWorkspaceChanged(workspaceId);
		}
		return workspace;
	}

	private async updateRuntime(
		workspaceId: ObjectId,
		patch: RuntimePatch,
		expected: WorkspaceSandboxExpectation
	): Promise<void> {
		if (await this.store.update(workspaceId, patch, expected)) return;
		await this.throwWorkspaceChanged(workspaceId);
	}

	private async throwWorkspaceChanged(workspaceId: ObjectId): Promise<never> {
		if (!(await this.store.get(workspaceId))) {
			throw new WorkspaceProvisioningError("Agent workspace no longer exists", "workspace_missing");
		}
		throw new WorkspaceProvisioningError(
			"Agent workspace changed during the sandbox operation",
			"sandbox_unavailable"
		);
	}

	private async retireSandbox(
		workspace: Workspace,
		sandbox: SandboxRuntime,
		jobsToken: string,
		reason: string,
		expected: WorkspaceSandboxExpectation,
		onProgress?: WorkspaceProvisioningCallback
	): Promise<Workspace> {
		await emit(onProgress, "cleanup", "Stopping HF Sandbox");
		const stoppedAt = this.now();
		try {
			await sandbox.kill();
		} catch (error) {
			const safeMessage = redactError(error, [jobsToken]);
			await this.updateRuntime(
				workspace._id,
				{
					sandbox: {
						...workspace.sandbox,
						status: "error",
						lastSeenAt: stoppedAt,
						error: `${reason}. Automatic sandbox cleanup failed: ${safeMessage}`.slice(
							0,
							MAX_STORED_ERROR_CHARS
						),
					},
					piSession: workspace.piSession ? { ...workspace.piSession, bridgePid: undefined } : null,
					llmProxy: null,
					updatedAt: stoppedAt,
				},
				expected
			);
			throw new WorkspaceProvisioningError(safeMessage, "sandbox_unavailable");
		}

		const stoppedSandbox = stoppedSandboxMetadata(workspace.sandbox, stoppedAt);
		await this.updateRuntime(
			workspace._id,
			{
				sandbox: stoppedSandbox,
				piSession: workspace.piSession ? { ...workspace.piSession, bridgePid: undefined } : null,
				llmProxy: null,
				updatedAt: stoppedAt,
			},
			expected
		);
		return this.requireExpectedWorkspace(
			workspace._id,
			stopClaimExpectation({ ...workspace, sandbox: stoppedSandbox })
		);
	}
}

type ValidatedRepository = { kind: "blank" } | { kind: "git"; url: string } | { kind: "hf" };

function normalizeRuntimeModels(
	workspace: Workspace,
	available: AgentRuntimeModelDefinition[] | undefined
): AgentRuntimeModelDefinition[] {
	const candidates = available?.length
		? available
		: [
				{
					id: workspace.model,
					name: workspace.model,
					reasoning: workspace.thinkingLevel !== "off",
				},
			];
	const byId = new Map<string, AgentRuntimeModelDefinition>();
	for (const candidate of candidates.slice(0, 500)) {
		const id = candidate.id.trim();
		if (!id || id.length > 512 || hasAsciiControl(id, true)) continue;
		byId.set(id, {
			id,
			...(candidate.name?.trim() ? { name: candidate.name.trim().slice(0, 512) } : {}),
			reasoning: Boolean(candidate.reasoning),
			multimodal: Boolean(candidate.multimodal),
			...(candidate.contextWindow && candidate.contextWindow > 0
				? { contextWindow: Math.min(Math.floor(candidate.contextWindow), 10_000_000) }
				: {}),
			...(candidate.maxTokens && candidate.maxTokens > 0
				? { maxTokens: Math.min(Math.floor(candidate.maxTokens), 1_000_000) }
				: {}),
			...(candidate.cost &&
			Number.isFinite(candidate.cost.input) &&
			candidate.cost.input >= 0 &&
			Number.isFinite(candidate.cost.output) &&
			candidate.cost.output >= 0
				? {
						cost: {
							input: candidate.cost.input,
							output: candidate.cost.output,
							cacheRead: Math.max(0, candidate.cost.cacheRead ?? 0),
							cacheWrite: Math.max(0, candidate.cost.cacheWrite ?? 0),
						},
					}
				: {}),
		});
	}
	if (!byId.has(workspace.model)) {
		byId.set(workspace.model, {
			id: workspace.model,
			name: workspace.model,
			reasoning: workspace.thinkingLevel !== "off",
		});
	}
	return [...byId.values()];
}

function validateRepository(workspace: Workspace): ValidatedRepository {
	if (
		!workspace.model.trim() ||
		workspace.model.length > 512 ||
		hasAsciiControl(workspace.model, false)
	) {
		throw new WorkspaceProvisioningError("Invalid agent model id", "invalid_configuration");
	}
	if (workspace.branch) validateBranch(workspace.branch);
	if (workspace.repo.kind === "blank") return { kind: "blank" };

	if (workspace.repo.kind === "git") {
		let url: URL;
		try {
			url = new URL(workspace.repo.ref);
		} catch {
			throw new WorkspaceProvisioningError(
				"Git repository must be a public HTTPS URL",
				"invalid_repository"
			);
		}
		if (
			url.protocol !== "https:" ||
			!url.hostname ||
			url.username ||
			url.password ||
			url.search ||
			url.hash
		) {
			throw new WorkspaceProvisioningError(
				"Git repository must be a credential-free public HTTPS URL",
				"invalid_repository"
			);
		}
		return { kind: "git", url: url.toString() };
	}

	if (
		workspace.repo.ref.length > 256 ||
		!/^[A-Za-z0-9][A-Za-z0-9._-]*\/[A-Za-z0-9][A-Za-z0-9._-]*$/.test(workspace.repo.ref) ||
		!(["model", "dataset", "space"] as const).includes(workspace.repo.repoType)
	) {
		throw new WorkspaceProvisioningError(
			"HF repository must use namespace/name",
			"invalid_repository"
		);
	}
	return { kind: "hf" };
}

/**
 * Disable every ambient credential helper. For an HF source, scope the existing
 * short-lived agent-session header to the exact server-side proxy URL. The raw
 * owner/Jobs Hub credentials never enter this environment or the resulting Git
 * configuration. LFS smudging is skipped because authenticated LFS action URLs
 * are intentionally outside this read-only smart-HTTP proxy's narrow surface.
 */
function gitCloneEnvironment(repoUrl?: string, agentSessionToken?: string): Record<string, string> {
	const environment: Record<string, string> = {
		GIT_TERMINAL_PROMPT: "0",
		GIT_ASKPASS: "/bin/false",
		SSH_ASKPASS: "/bin/false",
		GIT_CONFIG_COUNT: agentSessionToken ? "2" : "1",
		GIT_CONFIG_KEY_0: "credential.helper",
		GIT_CONFIG_VALUE_0: "",
	};
	if (repoUrl && agentSessionToken) {
		environment.GIT_LFS_SKIP_SMUDGE = "1";
		environment.GIT_CONFIG_KEY_1 = `http.${repoUrl}.extraHeader`;
		environment.GIT_CONFIG_VALUE_1 = `X-Agent-Session: ${agentSessionToken}`;
	}
	return environment;
}

function validateBranch(branch: string): string {
	const components = branch.split("/");
	if (
		branch.length === 0 ||
		branch.length > 200 ||
		branch === "HEAD" ||
		branch.startsWith("-") ||
		branch.endsWith(".") ||
		branch.endsWith("/") ||
		branch.includes("..") ||
		branch.includes("//") ||
		branch.includes("@{") ||
		components.some(
			(component) => !component || component.startsWith(".") || component.endsWith(".lock")
		) ||
		hasInvalidBranchCharacter(branch)
	) {
		throw new WorkspaceProvisioningError("Invalid Git branch name", "invalid_repository");
	}
	return branch;
}

function bridgeEnvironment(
	workspace: Workspace,
	session: WorkspacePiSession
): Record<string, string> {
	if (!session.id || !session.dir) {
		throw new WorkspaceProvisioningError(
			"Agent session metadata is incomplete",
			"invalid_configuration"
		);
	}
	const prefixArgs = [
		"--extension",
		AGENT_APPROVAL_EXTENSION_PATH,
		"--thinking",
		workspace.thinkingLevel,
		"--session-id",
		session.id,
	];
	return {
		BRIDGE_HOST: "127.0.0.1",
		BRIDGE_PORT: String(AGENT_BRIDGE_PORT),
		WORKSPACE_DIR: AGENT_WORKSPACE_DIR,
		PI_PROVIDER: AGENT_PROVIDER_ID,
		PI_MODEL: workspace.model,
		AGENT_POLICY: workspace.policy,
		AGENT_POLICY_PATH,
		PI_SESSION_DIR: session.dir,
		PI_SESSION_NAME: `chat-ui-${workspace.conversationId.toString()}`,
		PI_BIN_PREFIX_ARGS_JSON: JSON.stringify(prefixArgs),
		PI_TELEMETRY: "0",
	};
}

function hasAsciiControl(value: string, includeDelete: boolean): boolean {
	for (const character of value) {
		const code = character.charCodeAt(0);
		if (code <= 31 || (includeDelete && code === 127)) return true;
	}
	return false;
}

function hasInvalidBranchCharacter(value: string): boolean {
	for (const character of value) {
		const code = character.charCodeAt(0);
		if (code <= 32 || code === 127 || "~^:?*\\[".includes(character)) return true;
	}
	return false;
}

async function assertExec(
	sandbox: SandboxRuntime,
	command: string[],
	action: string,
	timeoutSeconds = 60,
	environment?: Record<string, string>,
	mapFailure?: (result: SandboxExecResult) => WorkspaceProvisioningError | undefined
): Promise<void> {
	const result = await sandbox.exec(command, {
		timeoutSeconds,
		...(environment ? { environment } : {}),
	});
	if (result.exitCode === 0 && !result.timedOut) return;
	const mapped = mapFailure?.(result);
	if (mapped) throw mapped;
	const detail = (result.stderr || result.stdout).trim().slice(-500);
	throw new WorkspaceProvisioningError(
		`Unable to ${action}${detail ? `: ${detail}` : ""}`,
		"sandbox_unavailable"
	);
}

function mapHuggingFaceCloneFailure(
	result: SandboxExecResult
): WorkspaceProvisioningError | undefined {
	if (result.timedOut) return undefined;
	const detail = `${result.stderr}\n${result.stdout}`;
	if (
		/askpass|could not read Username|authentication|HTTP[^\n]*\b(?:401|403)\b|error:\s*403/i.test(
			detail
		)
	) {
		return new WorkspaceProvisioningError(
			"Sign in with Hugging Face to access this repository, then retry",
			"invalid_repository"
		);
	}
	if (
		/remote branch .+ not found|repository not found|not found in upstream origin/i.test(detail)
	) {
		return new WorkspaceProvisioningError(
			"The selected Hugging Face repository or branch is unavailable",
			"invalid_repository"
		);
	}
	return undefined;
}

function toPersistedCredential(credential: IssuedAgentProxyCredential): WorkspaceLlmProxy {
	return {
		tokenHash: credential.tokenHash,
		createdAt: credential.createdAt,
		expiresAt: credential.expiresAt,
		allowedModels: credential.allowedModels,
	};
}

function workspaceSandboxFromConnection(
	connection: SandboxConnection,
	status: "provisioning" | "ready",
	seenAt: Date,
	generation: string,
	claimedAt?: Date
): WorkspaceSandbox {
	return {
		status,
		generation,
		...(claimedAt ? { claimedAt } : {}),
		id: connection.id,
		namespace: connection.namespace,
		baseUrl: connection.baseUrl,
		nonce: connection.nonce,
		image: connection.image,
		flavor: connection.flavor,
		createdAt: connection.createdAt,
		lastSeenAt: seenAt,
	};
}

function sandboxOwnershipLabels(workspaceId: ObjectId, generation: string): Record<string, string> {
	return {
		"chat-ui-agent": "1",
		"chat-ui-workspace": workspaceId.toString(),
		"chat-ui-generation": generation,
	};
}

function abandonedUnpersistedGeneration(
	workspace: Workspace,
	now: Date,
	leaseMs: number
): string | undefined {
	const { sandbox } = workspace;
	if (
		!sandbox.generation ||
		hasCompleteConnection(sandbox) ||
		(sandbox.status !== "error" &&
			(sandbox.status !== "provisioning" ||
				(sandbox.claimedAt && sandbox.claimedAt.getTime() > now.getTime() - leaseMs)))
	) {
		return undefined;
	}
	return sandbox.generation;
}

function sameSandboxConnection(left: SandboxConnection, right: SandboxConnection): boolean {
	return (
		left.id === right.id &&
		left.namespace === right.namespace &&
		left.baseUrl === right.baseUrl &&
		left.nonce === right.nonce &&
		left.image === right.image &&
		left.flavor === right.flavor &&
		left.createdAt.getTime() === right.createdAt.getTime()
	);
}

function sandboxExpectation(
	sandbox: WorkspaceSandbox,
	statuses: readonly WorkspaceSandboxStatus[] = [sandbox.status]
): WorkspaceSandboxExpectation {
	return {
		statuses,
		...(sandbox.id ? { id: sandbox.id } : {}),
		...(sandbox.nonce ? { nonce: sandbox.nonce } : {}),
		...(sandbox.generation ? { generation: sandbox.generation } : {}),
		cleanup: sandbox.cleanup ? { ...sandbox.cleanup } : null,
	};
}

function provisionClaimExpectation(workspace: Workspace): WorkspaceSandboxExpectation {
	return {
		...sandboxExpectation(workspace.sandbox),
		updatedAt: workspace.updatedAt,
		retirement: null,
	};
}

function stopClaimExpectation(
	workspace: Workspace,
	statuses: readonly WorkspaceSandboxStatus[] = [workspace.sandbox.status]
): WorkspaceSandboxExpectation {
	return {
		...sandboxExpectation(workspace.sandbox, statuses),
		retirement: workspace.retirement ? { id: workspace.retirement.id } : null,
	};
}

function liveSandboxExpectation(sandbox: WorkspaceSandbox): WorkspaceSandboxExpectation {
	if (!sandbox.id || !sandbox.nonce) {
		throw new WorkspaceProvisioningError(
			"Workspace has no reconnectable HF Sandbox metadata",
			"sandbox_unavailable"
		);
	}
	return sandboxExpectation(sandbox, ["ready", "idle"]);
}

function liveWorkspaceExpectation(workspace: Workspace): WorkspaceSandboxExpectation {
	return {
		...liveSandboxExpectation(workspace.sandbox),
		retirement: workspace.retirement ? { id: workspace.retirement.id } : null,
	};
}

function matchesSandboxExpectation(
	sandbox: WorkspaceSandbox,
	expected: WorkspaceSandboxExpectation
): boolean {
	return (
		expected.statuses.includes(sandbox.status) &&
		(!expected.id || sandbox.id === expected.id) &&
		(!expected.nonce || sandbox.nonce === expected.nonce) &&
		(!expected.generation || sandbox.generation === expected.generation) &&
		(expected.cleanup === undefined ||
			(expected.cleanup === null
				? sandbox.cleanup === undefined
				: sameSandboxCleanup(sandbox.cleanup, expected.cleanup)))
	);
}

function sameSandboxCleanup(
	actual: WorkspaceSandboxCleanup | undefined,
	expected: WorkspaceSandboxCleanup
): boolean {
	return (
		actual?.generation === expected.generation &&
		actual.namespace === expected.namespace &&
		actual.hubUrl === expected.hubUrl
	);
}

function matchesWorkspaceExpectation(
	workspace: Workspace,
	expected: WorkspaceSandboxExpectation
): boolean {
	return (
		matchesSandboxExpectation(workspace.sandbox, expected) &&
		(!expected.updatedAt || workspace.updatedAt.getTime() === expected.updatedAt.getTime()) &&
		(expected.retirement === undefined ||
			(expected.retirement === null
				? workspace.retirement === undefined
				: workspace.retirement?.id === expected.retirement.id))
	);
}

function stoppedSandboxMetadata(sandbox: WorkspaceSandbox, stoppedAt: Date): WorkspaceSandbox {
	return {
		status: "dead",
		flavor: sandbox.flavor,
		...(sandbox.generation ? { generation: sandbox.generation } : {}),
		...(sandbox.id ? { id: sandbox.id } : {}),
		...(sandbox.namespace ? { namespace: sandbox.namespace } : {}),
		...(sandbox.image ? { image: sandbox.image } : {}),
		...(sandbox.createdAt ? { createdAt: sandbox.createdAt } : {}),
		lastSeenAt: stoppedAt,
	};
}

function connectionFromWorkspace(workspace: Workspace): SandboxConnection {
	const sandbox = workspace.sandbox;
	if (!hasCompleteConnection(sandbox)) {
		throw new WorkspaceProvisioningError(
			"Workspace has no reconnectable HF Sandbox metadata",
			"sandbox_unavailable"
		);
	}
	return {
		id: sandbox.id,
		namespace: sandbox.namespace,
		baseUrl: sandbox.baseUrl,
		nonce: sandbox.nonce,
		image: sandbox.image,
		flavor: sandbox.flavor,
		createdAt: sandbox.createdAt || workspace.createdAt,
	};
}

function hasCompleteConnection(
	sandbox: WorkspaceSandbox
): sandbox is WorkspaceSandbox &
	Required<Pick<WorkspaceSandbox, "id" | "namespace" | "baseUrl" | "nonce" | "image">> {
	return Boolean(
		sandbox.id && sandbox.namespace && sandbox.baseUrl && sandbox.nonce && sandbox.image
	);
}

function runtimeConnection(
	workspace: Workspace,
	sandbox: SandboxRuntime
): WorkspaceRuntimeConnection {
	return {
		workspace,
		sandbox: sandbox as SandboxClient,
		bridgeUrl: sandbox.proxyUrlFor(AGENT_BRIDGE_PORT, "/ws", "wss"),
	};
}

function redactError(error: unknown, secrets: string[]): string {
	let message = error instanceof Error ? error.message : "Sandbox provisioning failed";
	for (const secret of secrets) {
		if (secret) message = message.split(secret).join("[redacted]");
	}
	return message.slice(0, MAX_STORED_ERROR_CHARS);
}

async function emit(
	callback: WorkspaceProvisioningCallback | undefined,
	phase: WorkspaceProvisioningPhase,
	message: string
): Promise<void> {
	try {
		await callback?.({ phase, message });
	} catch {
		// Progress delivery is best-effort and must never strand a billable Job.
	}
}

function readConfigValue(name: string): string {
	const value = (config as unknown as Record<string, unknown>)[name];
	return typeof value === "string" ? value.trim() : "";
}

function trimTrailingSlash(value: string): string {
	return value.replace(/\/+$/, "");
}

function delay(milliseconds: number): Promise<void> {
	return new Promise((resolve) => setTimeout(resolve, milliseconds));
}

export const mongoWorkspaceStore: WorkspaceProvisionerStore = {
	async claim(workspaceId, expected, llmProxy, generation, cleanup, now, staleBefore) {
		return collections.workspaces.findOneAndUpdate(
			{
				...workspaceExpectationFilter(workspaceId, expected),
				$or: [
					{ "sandbox.status": { $in: ["unprovisioned", "dead", "error"] } },
					{
						"sandbox.status": "provisioning",
						"sandbox.claimedAt": { $lt: staleBefore },
					},
					{ "sandbox.status": "provisioning", "sandbox.claimedAt": { $exists: false } },
				],
			},
			{
				$set: {
					"sandbox.status": "provisioning",
					"sandbox.generation": generation,
					"sandbox.claimedAt": now,
					...(cleanup ? { "sandbox.cleanup": cleanup } : {}),
					llmProxy,
					updatedAt: now,
				},
				$unset: {
					"sandbox.error": "",
					"piSession.bridgePid": "",
					...(!cleanup ? { "sandbox.cleanup": "" } : {}),
				},
			},
			{ returnDocument: "after", includeResultMetadata: false }
		);
	},
	async claimStop(workspaceId, expected, generation, cleanup, now) {
		return collections.workspaces.findOneAndUpdate(
			workspaceExpectationFilter(workspaceId, expected),
			{
				$set: {
					"sandbox.status": "provisioning",
					"sandbox.generation": generation,
					"sandbox.claimedAt": now,
					...(cleanup ? { "sandbox.cleanup": cleanup } : {}),
					updatedAt: now,
				},
				$unset: {
					llmProxy: "",
					...(!cleanup ? { "sandbox.cleanup": "" } : {}),
				},
			},
			{ returnDocument: "after", includeResultMetadata: false }
		);
	},
	get(workspaceId) {
		return collections.workspaces.findOne({ _id: workspaceId });
	},
	async update(workspaceId, patch, expected) {
		const set: Record<string, unknown> = { updatedAt: patch.updatedAt };
		const unset: Record<string, ""> = {};
		if (patch.sandbox !== undefined) set.sandbox = patch.sandbox;
		if (patch.piSession === null) unset.piSession = "";
		else if (patch.piSession !== undefined) set.piSession = patch.piSession;
		if (patch.llmProxy === null) unset.llmProxy = "";
		else if (patch.llmProxy !== undefined) set.llmProxy = patch.llmProxy;
		const result = await collections.workspaces.updateOne(
			workspaceExpectationFilter(workspaceId, expected),
			{ $set: set, ...(Object.keys(unset).length ? { $unset: unset } : {}) }
		);
		return result.matchedCount === 1;
	},
};

function workspaceExpectationFilter(
	workspaceId: ObjectId,
	expected: WorkspaceSandboxExpectation
): Record<string, unknown> {
	return {
		_id: workspaceId,
		"sandbox.status": { $in: [...expected.statuses] },
		...(expected.id ? { "sandbox.id": expected.id } : {}),
		...(expected.nonce ? { "sandbox.nonce": expected.nonce } : {}),
		...(expected.generation ? { "sandbox.generation": expected.generation } : {}),
		...(expected.cleanup === null
			? { "sandbox.cleanup": { $exists: false } }
			: expected.cleanup
				? {
						"sandbox.cleanup.generation": expected.cleanup.generation,
						"sandbox.cleanup.namespace": expected.cleanup.namespace,
						"sandbox.cleanup.hubUrl": expected.cleanup.hubUrl,
					}
				: {}),
		...(expected.updatedAt ? { updatedAt: expected.updatedAt } : {}),
		...(expected.retirement === null
			? { retirement: { $exists: false } }
			: expected.retirement
				? { "retirement.id": expected.retirement.id }
				: {}),
	};
}

export const workspaceProvisioner = new WorkspaceProvisioner();