Mario Zechner
441 lines
13 KiB
TypeScript
441 lines
13 KiB
TypeScript
import { hostname, platform } from "node:os";
|
|
import type { OAuthCredential } from "@earendil-works/pi-ai";
|
|
import { readStoredCredential } from "@earendil-works/pi-coding-agent";
|
|
import { getOrchestratorDir, getSocketPath, VERSION } from "./config.ts";
|
|
import { loadMachine, saveMachine } from "./storage.ts";
|
|
import type { InstanceRecord, MachineRecord, RadiusRegistration } from "./types.ts";
|
|
|
|
const DEFAULT_RADIUS_URL = "https://radius.pi.dev/";
|
|
const DEFAULT_ORCHESTRATOR_BASE_PATH = "/v1/";
|
|
const NOT_FOUND_RETRY_THRESHOLD = 3;
|
|
const HEARTBEAT_BACKOFF_BASE_MS = 1_000;
|
|
const HEARTBEAT_BACKOFF_MAX_MS = 30_000;
|
|
const RADIUS_PROVIDER = "radius";
|
|
|
|
interface RegisterMachineResponse extends RadiusRegistration {
|
|
id: string;
|
|
}
|
|
|
|
interface RegisterPiResponse extends RadiusRegistration {
|
|
id: string;
|
|
}
|
|
|
|
interface RadiusPresenceCoordinator {
|
|
getLiveInstance(instanceId: string): InstanceRecord | undefined;
|
|
listLiveInstances(): InstanceRecord[];
|
|
updateInstance(instance: InstanceRecord): void;
|
|
}
|
|
|
|
interface PiHeartbeatState {
|
|
timer?: NodeJS.Timeout;
|
|
intervalMs: number;
|
|
radiusPiId: string;
|
|
consecutiveNotFoundCount: number;
|
|
transientFailureCount: number;
|
|
}
|
|
|
|
class RadiusHttpError extends Error {
|
|
readonly status: number;
|
|
|
|
constructor(status: number, message: string) {
|
|
super(message);
|
|
this.name = "RadiusHttpError";
|
|
this.status = status;
|
|
}
|
|
}
|
|
|
|
async function post<T>(path: string, body: unknown): Promise<T> {
|
|
const response = await fetch(new URL(path, getRadiusOrchestratorBaseUrl()), {
|
|
method: "POST",
|
|
headers: {
|
|
Authorization: `Bearer ${getRadiusAccessToken()}`,
|
|
"Content-Type": "application/json",
|
|
},
|
|
body: JSON.stringify(body),
|
|
});
|
|
|
|
if (!response.ok) {
|
|
throw new RadiusHttpError(response.status, `Radius request failed: ${response.status} ${await response.text()}`);
|
|
}
|
|
|
|
return (await response.json()) as T;
|
|
}
|
|
|
|
async function maybePost(path: string, body: unknown): Promise<void> {
|
|
const response = await fetch(new URL(path, getRadiusOrchestratorBaseUrl()), {
|
|
method: "POST",
|
|
headers: {
|
|
Authorization: `Bearer ${getRadiusAccessToken()}`,
|
|
"Content-Type": "application/json",
|
|
},
|
|
body: JSON.stringify(body),
|
|
});
|
|
if (!response.ok) {
|
|
throw new RadiusHttpError(response.status, `Radius request failed: ${response.status} ${await response.text()}`);
|
|
}
|
|
}
|
|
|
|
function isNotFoundError(error: unknown): error is RadiusHttpError {
|
|
return error instanceof RadiusHttpError && error.status === 404;
|
|
}
|
|
|
|
function computeBackoffDelayMs(failureCount: number): number {
|
|
const exponentialDelay = Math.min(
|
|
HEARTBEAT_BACKOFF_MAX_MS,
|
|
HEARTBEAT_BACKOFF_BASE_MS * 2 ** Math.max(0, failureCount - 1),
|
|
);
|
|
const jitterMs = Math.floor(Math.random() * Math.max(250, exponentialDelay / 4));
|
|
return Math.min(HEARTBEAT_BACKOFF_MAX_MS, exponentialDelay + jitterMs);
|
|
}
|
|
|
|
function formatRadiusError(error: unknown): string {
|
|
if (error instanceof RadiusHttpError) {
|
|
return `HTTP ${error.status}: ${error.message}`;
|
|
}
|
|
if (error instanceof Error) {
|
|
return error.message;
|
|
}
|
|
return String(error);
|
|
}
|
|
|
|
function logRadiusRetry(scope: string, action: string, delayMs: number, failureCount: number, error: unknown): void {
|
|
console.error(
|
|
`${scope} ${action} failed (attempt ${failureCount}); retrying in ${delayMs}ms: ${formatRadiusError(error)}`,
|
|
);
|
|
}
|
|
|
|
export function getRadiusUrl(): string {
|
|
return process.env.PI_RADIUS_URL || DEFAULT_RADIUS_URL;
|
|
}
|
|
|
|
export function getRadiusOrchestratorBaseUrl(): string {
|
|
const explicitUrl = process.env.PI_RADIUS_ORCHESTRATOR_URL;
|
|
if (explicitUrl) {
|
|
return explicitUrl;
|
|
}
|
|
|
|
return new URL(DEFAULT_ORCHESTRATOR_BASE_PATH, getRadiusUrl()).toString();
|
|
}
|
|
|
|
function getStoredRadiusCredential(): OAuthCredential | undefined {
|
|
const credential = readStoredCredential(RADIUS_PROVIDER);
|
|
return credential?.type === "oauth" ? credential : undefined;
|
|
}
|
|
|
|
export function getRadiusAccessToken(): string {
|
|
const storedCredential = getStoredRadiusCredential();
|
|
if (typeof storedCredential?.access === "string" && storedCredential.access) {
|
|
return storedCredential.access;
|
|
}
|
|
|
|
const apiKey = process.env.RADIUS_API_KEY;
|
|
if (apiKey) {
|
|
return apiKey;
|
|
}
|
|
|
|
throw new Error("Radius credentials are required in ~/.pi/agent/auth.json or RADIUS_API_KEY");
|
|
}
|
|
|
|
export function isRadiusEnabled(): boolean {
|
|
return !!getStoredRadiusCredential()?.access || !!process.env.RADIUS_API_KEY;
|
|
}
|
|
|
|
export class RadiusPresence {
|
|
private machineHeartbeatTimer?: NodeJS.Timeout;
|
|
private machineHeartbeatIntervalMs = 0;
|
|
private machineConsecutiveNotFoundCount = 0;
|
|
private machineTransientFailureCount = 0;
|
|
private readonly piHeartbeatStates = new Map<string, PiHeartbeatState>();
|
|
private machine?: MachineRecord;
|
|
private coordinator?: RadiusPresenceCoordinator;
|
|
|
|
setCoordinator(coordinator: RadiusPresenceCoordinator): void {
|
|
this.coordinator = coordinator;
|
|
}
|
|
|
|
async start(label?: string): Promise<MachineRecord | undefined> {
|
|
if (!isRadiusEnabled()) {
|
|
return undefined;
|
|
}
|
|
|
|
const registered = await this.registerMachine(label);
|
|
this.startMachineHeartbeat(registered.heartbeatIntervalMs);
|
|
return this.machine;
|
|
}
|
|
|
|
async stop(): Promise<void> {
|
|
if (this.machineHeartbeatTimer) {
|
|
clearTimeout(this.machineHeartbeatTimer);
|
|
this.machineHeartbeatTimer = undefined;
|
|
}
|
|
for (const [instanceId, state] of this.piHeartbeatStates) {
|
|
if (state.timer) {
|
|
clearTimeout(state.timer);
|
|
}
|
|
this.piHeartbeatStates.delete(instanceId);
|
|
}
|
|
if (!this.machine || !isRadiusEnabled()) {
|
|
return;
|
|
}
|
|
try {
|
|
await maybePost(`machines/${this.machine.id}/disconnect`, {});
|
|
} catch (error) {
|
|
if (!isNotFoundError(error)) {
|
|
throw error;
|
|
}
|
|
}
|
|
}
|
|
|
|
async registerPi(instance: InstanceRecord): Promise<InstanceRecord> {
|
|
if (!isRadiusEnabled()) {
|
|
return instance;
|
|
}
|
|
const machine = this.machine ?? loadMachine();
|
|
if (!machine) {
|
|
throw new Error("No registered machine available for Pi registration");
|
|
}
|
|
const registered = await post<RegisterPiResponse>("pis/register", {
|
|
machineId: machine.id,
|
|
label: instance.label,
|
|
cwd: instance.cwd,
|
|
hostname: hostname(),
|
|
pid: process.pid,
|
|
transport: "local-rpc",
|
|
capabilities: { rpc: true, relay: false, iroh: false },
|
|
sessionId: instance.sessionId,
|
|
});
|
|
const registeredInstance = { ...instance, radiusPiId: registered.id };
|
|
this.startPiHeartbeat(instance.id, registered.heartbeatIntervalMs, registered.id);
|
|
return registeredInstance;
|
|
}
|
|
|
|
async disconnectPi(instance: InstanceRecord): Promise<void> {
|
|
const state = this.piHeartbeatStates.get(instance.id);
|
|
if (state) {
|
|
if (state.timer) {
|
|
clearTimeout(state.timer);
|
|
}
|
|
this.piHeartbeatStates.delete(instance.id);
|
|
}
|
|
if (!isRadiusEnabled() || !instance.radiusPiId) {
|
|
return;
|
|
}
|
|
try {
|
|
await maybePost(`pis/${instance.radiusPiId}/disconnect`, {});
|
|
} catch (error) {
|
|
if (!isNotFoundError(error)) {
|
|
throw error;
|
|
}
|
|
}
|
|
}
|
|
|
|
private async registerMachine(label?: string): Promise<RegisterMachineResponse> {
|
|
const existingMachine = this.machine ?? loadMachine();
|
|
const registered = await post<RegisterMachineResponse>("machines/register", {
|
|
machineId: existingMachine?.id,
|
|
label,
|
|
hostname: hostname(),
|
|
platform: platform(),
|
|
arch: process.arch,
|
|
version: VERSION,
|
|
capabilities: { spawn: true, relay: false, iroh: false },
|
|
});
|
|
|
|
const timestamp = new Date().toISOString();
|
|
this.machine = {
|
|
id: registered.id,
|
|
createdAt: existingMachine?.createdAt ?? timestamp,
|
|
lastSeenAt: timestamp,
|
|
label,
|
|
};
|
|
saveMachine(this.machine);
|
|
this.machineConsecutiveNotFoundCount = 0;
|
|
this.machineTransientFailureCount = 0;
|
|
return registered;
|
|
}
|
|
|
|
private startMachineHeartbeat(intervalMs: number): void {
|
|
this.machineHeartbeatIntervalMs = intervalMs;
|
|
this.scheduleMachineHeartbeat(intervalMs);
|
|
}
|
|
|
|
private scheduleMachineHeartbeat(delayMs: number): void {
|
|
if (this.machineHeartbeatTimer) {
|
|
clearTimeout(this.machineHeartbeatTimer);
|
|
}
|
|
this.machineHeartbeatTimer = setTimeout(() => {
|
|
void this.heartbeatMachine();
|
|
}, delayMs);
|
|
}
|
|
|
|
private startPiHeartbeat(instanceId: string, intervalMs: number, radiusPiId: string): void {
|
|
const existingState = this.piHeartbeatStates.get(instanceId);
|
|
if (existingState?.timer) {
|
|
clearTimeout(existingState.timer);
|
|
}
|
|
const state: PiHeartbeatState = existingState ?? {
|
|
intervalMs,
|
|
radiusPiId,
|
|
consecutiveNotFoundCount: 0,
|
|
transientFailureCount: 0,
|
|
};
|
|
state.intervalMs = intervalMs;
|
|
state.radiusPiId = radiusPiId;
|
|
state.consecutiveNotFoundCount = 0;
|
|
state.transientFailureCount = 0;
|
|
this.piHeartbeatStates.set(instanceId, state);
|
|
this.schedulePiHeartbeat(instanceId, intervalMs);
|
|
}
|
|
|
|
private schedulePiHeartbeat(instanceId: string, delayMs: number): void {
|
|
const state = this.piHeartbeatStates.get(instanceId);
|
|
if (!state) {
|
|
return;
|
|
}
|
|
if (state.timer) {
|
|
clearTimeout(state.timer);
|
|
}
|
|
state.timer = setTimeout(() => {
|
|
void this.heartbeatPi(instanceId);
|
|
}, delayMs);
|
|
}
|
|
|
|
private async heartbeatMachine(): Promise<void> {
|
|
if (!this.machine || !isRadiusEnabled()) {
|
|
return;
|
|
}
|
|
|
|
try {
|
|
await maybePost(`machines/${this.machine.id}/heartbeat`, {
|
|
cwd: getOrchestratorDir(),
|
|
socketPath: getSocketPath(),
|
|
});
|
|
this.machineConsecutiveNotFoundCount = 0;
|
|
this.machineTransientFailureCount = 0;
|
|
this.scheduleMachineHeartbeat(this.machineHeartbeatIntervalMs);
|
|
} catch (error) {
|
|
if (!isNotFoundError(error)) {
|
|
this.machineTransientFailureCount += 1;
|
|
const delayMs = computeBackoffDelayMs(this.machineTransientFailureCount);
|
|
logRadiusRetry("Radius machine", "heartbeat", delayMs, this.machineTransientFailureCount, error);
|
|
this.scheduleMachineHeartbeat(delayMs);
|
|
return;
|
|
}
|
|
|
|
this.machineTransientFailureCount = 0;
|
|
this.machineConsecutiveNotFoundCount += 1;
|
|
if (this.machineConsecutiveNotFoundCount < NOT_FOUND_RETRY_THRESHOLD) {
|
|
this.scheduleMachineHeartbeat(this.machineHeartbeatIntervalMs);
|
|
return;
|
|
}
|
|
|
|
try {
|
|
await this.reRegisterMachineAndPis();
|
|
} catch (recoveryError) {
|
|
this.machineTransientFailureCount += 1;
|
|
const delayMs = computeBackoffDelayMs(this.machineTransientFailureCount);
|
|
logRadiusRetry(
|
|
"Radius machine",
|
|
"re-registration",
|
|
delayMs,
|
|
this.machineTransientFailureCount,
|
|
recoveryError,
|
|
);
|
|
this.scheduleMachineHeartbeat(delayMs);
|
|
}
|
|
}
|
|
}
|
|
|
|
private async heartbeatPi(instanceId: string): Promise<void> {
|
|
if (!isRadiusEnabled()) {
|
|
return;
|
|
}
|
|
|
|
const state = this.piHeartbeatStates.get(instanceId);
|
|
if (!state) {
|
|
return;
|
|
}
|
|
|
|
try {
|
|
await maybePost(`pis/${state.radiusPiId}/heartbeat`, {});
|
|
state.consecutiveNotFoundCount = 0;
|
|
state.transientFailureCount = 0;
|
|
this.schedulePiHeartbeat(instanceId, state.intervalMs);
|
|
} catch (error) {
|
|
if (!isNotFoundError(error)) {
|
|
state.transientFailureCount += 1;
|
|
const delayMs = computeBackoffDelayMs(state.transientFailureCount);
|
|
logRadiusRetry(`Radius Pi ${instanceId}`, "heartbeat", delayMs, state.transientFailureCount, error);
|
|
this.schedulePiHeartbeat(instanceId, delayMs);
|
|
return;
|
|
}
|
|
|
|
state.transientFailureCount = 0;
|
|
state.consecutiveNotFoundCount += 1;
|
|
if (state.consecutiveNotFoundCount < NOT_FOUND_RETRY_THRESHOLD) {
|
|
this.schedulePiHeartbeat(instanceId, state.intervalMs);
|
|
return;
|
|
}
|
|
|
|
try {
|
|
const recovered = await this.reRegisterPi(instanceId);
|
|
if (!recovered) {
|
|
const delayMs = computeBackoffDelayMs(1);
|
|
console.error(`Radius Pi ${instanceId} re-registration skipped; retrying in ${delayMs}ms`);
|
|
this.schedulePiHeartbeat(instanceId, delayMs);
|
|
}
|
|
} catch (recoveryError) {
|
|
state.transientFailureCount += 1;
|
|
const delayMs = computeBackoffDelayMs(state.transientFailureCount);
|
|
logRadiusRetry(
|
|
`Radius Pi ${instanceId}`,
|
|
"re-registration",
|
|
delayMs,
|
|
state.transientFailureCount,
|
|
recoveryError,
|
|
);
|
|
this.schedulePiHeartbeat(instanceId, delayMs);
|
|
}
|
|
}
|
|
}
|
|
|
|
private async reRegisterMachineAndPis(): Promise<void> {
|
|
const registered = await this.registerMachine(this.machine?.label);
|
|
this.startMachineHeartbeat(registered.heartbeatIntervalMs);
|
|
|
|
const instances = this.coordinator?.listLiveInstances() ?? [];
|
|
for (const instance of instances) {
|
|
try {
|
|
await this.reRegisterPi(instance.id);
|
|
} catch (error) {
|
|
console.error(`Radius Pi ${instance.id} re-registration failed: ${formatRadiusError(error)}`);
|
|
}
|
|
}
|
|
}
|
|
|
|
private async reRegisterPi(instanceId: string): Promise<boolean> {
|
|
const instance = this.coordinator?.getLiveInstance(instanceId);
|
|
if (!instance) {
|
|
const state = this.piHeartbeatStates.get(instanceId);
|
|
if (state) {
|
|
if (state.timer) {
|
|
clearTimeout(state.timer);
|
|
}
|
|
this.piHeartbeatStates.delete(instanceId);
|
|
}
|
|
return false;
|
|
}
|
|
|
|
if (!this.machine) {
|
|
await this.reRegisterMachineAndPis();
|
|
return true;
|
|
}
|
|
|
|
const registeredInstance = await this.registerPi(instance);
|
|
this.coordinator?.updateInstance(registeredInstance);
|
|
return true;
|
|
}
|
|
}
|
|
|
|
export const radiusPresence = new RadiusPresence();
|