apps/canvas/back/src/project_monitor.ts

import { z } from "zod";

const LogItemSchema = z.object({
	runId: z.string(),
	timestampMilli: z.number(),
	commit: z.string().optional(),
	contents: z.preprocess((val) => {
		if (typeof val === "string") {
			return Buffer.from(val, "base64").toString("utf-8");
		}
		throw new Error("Log item contents is not a string");
	}, z.string()),
});

export type LogItem = z.infer<typeof LogItemSchema>;

const LogItemsSchema = z.array(LogItemSchema);

export const WorkerSchema = z.object({
	id: z.string(),
	service: z.string(),
	address: z.string().url(),
	status: z.optional(
		z.object({
			commit: z
				.optional(
					z.object({
						hash: z.string(),
						message: z.string(),
					}),
				)
				.nullable(),
			commands: z.optional(
				z.array(
					z.object({
						command: z.string(),
						state: z.string(),
					}),
				),
			),
		}),
	),
	logs: LogItemsSchema.optional(),
});

export type Worker = z.infer<typeof WorkerSchema>;

class ServiceMonitor {
	private workers: Map<string, string> = new Map();
	private statuses: Map<string, Worker["status"]> = new Map();
	private lastPings: Map<string, number> = new Map();

	constructor(public readonly serviceName: string) {}

	registerWorker(workerId: string, workerAddress: string, workerStatus?: Worker["status"]): void {
		this.workers.set(workerId, workerAddress);
		if (workerStatus) {
			this.statuses.set(workerId, workerStatus);
		}
		this.lastPings.set(workerId, Date.now());
	}

	getWorkerAddress(workerId: string): string | undefined {
		return this.workers.get(workerId);
	}

	getWorkerStatus(workerId: string): Worker["status"] | undefined {
		return this.statuses.get(workerId);
	}

	getAllStatuses(): Map<string, Worker["status"]> {
		return new Map(this.statuses);
	}

	getWorkerAddresses(): string[] {
		return Array.from(this.workers.values());
	}

	getWorkerIds(): string[] {
		return Array.from(this.workers.keys());
	}

	async reloadWorker(workerId: string): Promise<void> {
		const workerAddress = this.workers.get(workerId);
		if (!workerAddress) {
			throw new Error(`Worker ${workerId} not found in service ${this.serviceName}`);
		}
		try {
			const response = await fetch(`${workerAddress}/update`, { method: "POST" });
			if (!response.ok) {
				throw new Error(
					`Failed to trigger reload for worker ${workerId} at ${workerAddress}: ${response.statusText}`,
				);
			}
			console.log(`Reload triggered for worker ${workerId} in service ${this.serviceName}`);
		} catch (error) {
			console.error(`Error reloading worker ${workerId} in service ${this.serviceName}:`, error);
			throw error; // Re-throw to be caught by ProjectMonitor
		}
	}

	async terminateWorker(workerId: string): Promise<void> {
		const workerAddress = this.workers.get(workerId);
		if (!workerAddress) {
			throw new Error(`Worker ${workerId} not found in service ${this.serviceName}`);
		}
		try {
			const response = await fetch(`${workerAddress}/quitquitquit`, { method: "POST" });
			if (!response.ok) {
				throw new Error(`Failed to terminate worker ${workerId} at ${workerAddress}: ${response.statusText}`);
			}
			console.log(`Terminated worker ${workerId} in service ${this.serviceName}`);
		} catch (error) {
			console.error(`Error terminating worker ${workerId} in service ${this.serviceName}:`, error);
			throw error; // Re-throw to be caught by ProjectMonitor
		}
	}

	cleanupWorkers(now: number) {
		const toDelete: string[] = [];
		this.workers.forEach((_, workerId) => {
			if (now - (this.lastPings.get(workerId) ?? 0) > 1000 * 60) {
				toDelete.push(workerId);
			}
		});
		for (const workerId of toDelete) {
			this.workers.delete(workerId);
			this.statuses.delete(workerId);
			this.lastPings.delete(workerId);
		}
	}
}

export class ProjectMonitor {
	private serviceMonitors: Map<string, ServiceMonitor> = new Map();

	constructor() {}

	registerWorker(workerData: Worker): void {
		let serviceMonitor = this.serviceMonitors.get(workerData.service);
		if (!serviceMonitor) {
			serviceMonitor = new ServiceMonitor(workerData.service);
			this.serviceMonitors.set(workerData.service, serviceMonitor);
		}
		serviceMonitor.registerWorker(workerData.id, workerData.address, workerData.status);
	}

	getWorkerAddresses(): string[] {
		let allAddresses: string[] = [];
		for (const serviceMonitor of this.serviceMonitors.values()) {
			allAddresses = allAddresses.concat(serviceMonitor.getWorkerAddresses());
		}
		return Array.from(new Set(allAddresses));
	}

	getAllServiceNames(): string[] {
		return Array.from(this.serviceMonitors.keys());
	}

	getServiceMonitor(serviceName: string): ServiceMonitor | undefined {
		return this.serviceMonitors.get(serviceName);
	}

	getWorkerStatusesForService(serviceName: string): Map<string, Worker["status"]> {
		const serviceMonitor = this.serviceMonitors.get(serviceName);
		if (serviceMonitor) {
			return serviceMonitor.getAllStatuses();
		}
		return new Map();
	}

	async reloadWorker(serviceName: string, workerId: string): Promise<void> {
		const serviceMonitor = this.serviceMonitors.get(serviceName);
		if (!serviceMonitor) {
			throw new Error(`Service ${serviceName} not found`);
		}
		await serviceMonitor.reloadWorker(workerId);
	}

	async terminateWorker(serviceName: string, workerId: string): Promise<void> {
		const serviceMonitor = this.serviceMonitors.get(serviceName);
		if (!serviceMonitor) {
			throw new Error(`Service ${serviceName} not found`);
		}
		await serviceMonitor.terminateWorker(workerId);
	}

	cleanupWorkers(now: number) {
		this.serviceMonitors.forEach((serviceMonitor) => {
			serviceMonitor.cleanupWorkers(now);
		});
	}
}