fix(hub): settle failed relationship connections

This commit is contained in:
Mohamed Boudra
2026-07-13 12:28:20 +02:00
parent 7c5bb7c4ae
commit ca393ab436
9 changed files with 440 additions and 12 deletions
+1
View File
@@ -1047,6 +1047,7 @@ export async function createPaseoDaemon(
agentManager,
agentStorage,
createAgent,
logger,
registerAutoArchive: ({ agentId, createdWorktree }) =>
hubAgentLifecycle.registerAutoArchiveIfRequested({
autoArchive: true,
@@ -15,12 +15,14 @@ import { HubSession } from "./hub-session.js";
class InMemoryHubExecutions implements HubExecutions {
private readonly listeners = new Set<(event: OwnedAgentEvent) => void>();
private reconcileError: Error | null = null;
async create(_input: HubAgentCreateInput): Promise<OwnedAgentSnapshot> {
return { executionId: "execution-1", agent: createAgentSnapshot() };
}
async reconcile(_executionId: string): Promise<OwnedAgentSnapshot | null> {
if (this.reconcileError) throw this.reconcileError;
return { executionId: "execution-1", agent: createAgentSnapshot() };
}
@@ -34,6 +36,10 @@ class InMemoryHubExecutions implements HubExecutions {
listener({ type: "update", executionId: "execution-1", agent: createAgentSnapshot() });
}
}
failReconcile(error: Error): void {
this.reconcileError = error;
}
}
class HubSessionBoundary {
@@ -52,6 +58,19 @@ class HubSessionBoundary {
this.session.cleanup();
}
async failReconcile(): Promise<void> {
this.executions.failReconcile(new Error("storage unavailable"));
await this.session.handleMessage({
type: "hub.execution.reconcile.request",
requestId: "reconcile-1",
executionId: "execution-1",
});
}
deliveredMessages(): SessionOutboundMessage[] {
return this.messages.slice();
}
deliveredTypes(): string[] {
return this.messages.map((message) => message.type);
}
@@ -67,6 +86,24 @@ test("disconnect unsubscribes Hub from later owned-agent updates", () => {
expect(hub.deliveredTypes()).toEqual(["hub.agent.update"]);
});
test("reconcile storage failures receive a terminal empty response", async () => {
const hub = new HubSessionBoundary();
await hub.failReconcile();
expect(hub.deliveredMessages()).toEqual([
{
type: "hub.execution.reconcile.response",
payload: {
requestId: "reconcile-1",
executionId: "execution-1",
agentId: null,
agent: null,
},
},
]);
});
function createAgentSnapshot(): AgentSnapshotPayload {
return AgentSnapshotPayloadSchema.parse({
id: "agent-1",
@@ -90,7 +90,13 @@ export class HubSession {
}
private async reconcile(message: HubExecutionReconcileRequest): Promise<void> {
const result = await this.executions.reconcile(message.executionId);
let result: Awaited<ReturnType<HubExecutions["reconcile"]>> = null;
try {
result = await this.executions.reconcile(message.executionId);
} catch {
// Reconcile has no error variant on the wire. A terminal empty response
// lets Hub stop waiting and decide whether to replay the execution.
}
this.send({
type: "hub.execution.reconcile.response",
payload: {
@@ -192,10 +192,11 @@ describe("Hub relationship", () => {
expect(reconciled).toMatchObject({
executionId: "execution-running",
agentId: created.payload.agentId,
agent: { id: created.payload.agentId, status: "running" },
agent: { id: created.payload.agentId, status: "idle" },
});
expect(durableAgentIds).toEqual([created.payload.agentId]);
expect(relationship.providerCreations()).toBe(1);
expect(relationship.providerResumes()).toBe(1);
});
test("stale socket generations cannot replace or unregister the current socket", async () => {
@@ -301,9 +301,7 @@ export class HubRelationshipController implements HubRelationshipManagement {
connected: (socket) => this.socketConnected(generation, record, socket),
rejected: (statusCode) => this.socketRejected(generation, statusCode),
closed: (code) => this.socketClosed(generation, record, code),
failed: (error) => {
if (generation === this.generation) this.lastError = error.message;
},
failed: (error) => this.socketFailed(generation, record, error),
};
this.socket = this.options.remote.openSocket(
{
@@ -351,6 +349,12 @@ export class HubRelationshipController implements HubRelationshipManagement {
if (this.record?.state === "active") this.scheduleSocket(record);
}
private socketFailed(generation: number, record: ActiveRecord, error: Error): void {
if (generation !== this.generation) return;
this.lastError = error.message;
if (this.record?.state === "active") this.scheduleSocket(record);
}
private scheduleSocket(record: ActiveRecord): void {
this.state = "reconnecting";
this.schedule(() => this.openSocket(record, true));
@@ -3,9 +3,11 @@ import type {
AgentStreamEventPayload,
CreateAgentWorktreeTarget,
} from "@getpaseo/protocol/messages";
import type { Logger } from "pino";
import type { AgentManager, AgentManagerEvent, ManagedAgent } from "../agent/agent-manager.js";
import type { AgentStorage, StoredAgentRecord } from "../agent/agent-storage.js";
import { ensureAgentLoaded } from "../agent/agent-loading.js";
import type { BoundCreateAgentCommand } from "../agent/create-agent/create.js";
import type { CreatePaseoWorktreeWorkflowResult } from "../worktree-session.js";
import { buildStoredAgentPayload } from "../agent/agent-projections.js";
@@ -46,6 +48,7 @@ interface RelationshipOwnedExecutionsOptions {
agentManager: AgentManager;
agentStorage: AgentStorage;
createAgent: BoundCreateAgentCommand;
logger: Logger;
registerAutoArchive?: (input: {
agentId: string;
createdWorktree: CreatePaseoWorktreeWorkflowResult | null;
@@ -67,6 +70,7 @@ export class RelationshipOwnedExecutions implements HubExecutions {
private readonly agentManager: AgentManager;
private readonly agentStorage: AgentStorage;
private readonly createAgentCommand: BoundCreateAgentCommand;
private readonly logger: Logger;
private readonly pendingCreates = new Map<string, Promise<OwnedAgentSnapshot>>();
private readonly registerAutoArchive: (input: {
agentId: string;
@@ -81,6 +85,7 @@ export class RelationshipOwnedExecutions implements HubExecutions {
this.agentManager = options.agentManager;
this.agentStorage = options.agentStorage;
this.createAgentCommand = options.createAgent;
this.logger = options.logger;
this.registerAutoArchive = options.registerAutoArchive ?? (() => undefined);
this.cleanupFailedCreate = options.cleanupFailedCreate ?? (async () => undefined);
}
@@ -104,7 +109,7 @@ export class RelationshipOwnedExecutions implements HubExecutions {
async reconcile(executionId: string): Promise<OwnedAgentSnapshot | null> {
const record = await this.agentStorage.findByHubExecution(this.owner(executionId));
return record ? this.projectRecord(record) : null;
return record ? this.resolveRecord(record) : null;
}
subscribe(listener: (event: OwnedAgentEvent) => void): () => void {
@@ -125,7 +130,7 @@ export class RelationshipOwnedExecutions implements HubExecutions {
): Promise<OwnedAgentSnapshot> {
const existing = await this.agentStorage.findByHubExecution(owner);
if (existing) {
return this.projectRecord(existing);
return this.resolveRecord(existing);
}
let createdWorktree: CreatePaseoWorktreeWorkflowResult | null = null;
@@ -166,6 +171,21 @@ export class RelationshipOwnedExecutions implements HubExecutions {
};
}
private async resolveRecord(record: StoredAgentRecord): Promise<OwnedAgentSnapshot> {
if (
!this.agentManager.getAgent(record.id) &&
!record.archivedAt &&
record.lastStatus !== "closed"
) {
await ensureAgentLoaded(record.id, {
agentManager: this.agentManager,
agentStorage: this.agentStorage,
logger: this.logger,
});
}
return this.projectRecord(record);
}
private projectRecord(record: StoredAgentRecord): OwnedAgentSnapshot {
const owner = record.owner;
if (owner?.kind !== "hub" || owner.relationshipId !== this.relationshipId) {
@@ -1,12 +1,32 @@
import { createServer } from "node:http";
import type { AddressInfo } from "node:net";
import { mkdtemp, rm } from "node:fs/promises";
import type { AddressInfo, Socket } from "node:net";
import { tmpdir } from "node:os";
import path from "node:path";
import pino from "pino";
import { afterEach, expect, test } from "vitest";
import { DirectHubRelationshipRemote } from "./relationship-remote.js";
import { WebSocket } from "ws";
import type { HubExecutions } from "./relationship-owned-executions.js";
import {
HubRelationshipController,
type HubRelationshipClock,
type HubRelationshipRetryPolicy,
type ScheduledRelationshipTask,
} from "./relationship-controller.js";
import {
DirectHubRelationshipRemote,
type HubSocketConnection,
type HubSocketEvents,
} from "./relationship-remote.js";
const openServers: ReturnType<typeof createServer>[] = [];
const openUpgradeHubs: UpgradeRejectingHub[] = [];
const openPaseoHomes: string[] = [];
afterEach(async () => {
for (const hub of openUpgradeHubs.splice(0)) hub.destroyConnections();
await Promise.all(openServers.splice(0).map((server) => closeServer(server)));
await Promise.all(openPaseoHomes.splice(0).map((home) => rm(home, { recursive: true })));
});
test.each([401, 403, 404])(
@@ -30,6 +50,107 @@ test("transient revocation failures remain retryable", async () => {
).rejects.toThrow("Hub revocation failed (503)");
});
test.each([401, 403] as const)(
"socket authentication status %s is terminal and releases the failed upgrade",
async (status) => {
const hub = await UpgradeRejectingHub.start([status]);
const outcomes = new SocketOutcomes();
const remote = new DirectHubRelationshipRemote();
const socket = remote.openSocket(
{
relationshipId: "relationship-1",
webSocketUrl: hub.webSocketUrl,
credential: "invalid",
},
outcomes.events,
) as HubSocketConnection & WebSocket;
expect(await outcomes.next()).toEqual({ type: "rejected", status });
await hub.expectAttemptReleased(1);
expect(await outcomes.afterCleanup()).toEqual([{ type: "rejected", status }]);
expect(socket.readyState).toBe(WebSocket.CLOSED);
expect(hub.openConnections()).toBe(0);
},
);
test("a transient socket response releases the failed upgrade and reports one connection loss", async () => {
const hub = await UpgradeRejectingHub.start([503]);
const outcomes = new SocketOutcomes();
const remote = new DirectHubRelationshipRemote();
const socket = remote.openSocket(
{
relationshipId: "relationship-1",
webSocketUrl: hub.webSocketUrl,
credential: "credential",
},
outcomes.events,
) as HubSocketConnection & WebSocket;
expect(await outcomes.next()).toEqual({ type: "closed", code: 1006 });
await hub.expectAttemptReleased(1);
expect(await outcomes.afterCleanup()).toEqual([{ type: "closed", code: 1006 }]);
expect(socket.readyState).toBe(WebSocket.CLOSED);
expect(hub.openConnections()).toBe(0);
});
test.each([401, 403] as const)(
"controller treats socket authentication status %s as terminal without redialing",
async (status) => {
const hub = await UpgradeRejectingHub.start([status]);
const clock = new ManualRelationshipClock();
const controller = await connectController(hub, clock);
await hub.expectAttemptReleased(1);
expect(controller.status()).toMatchObject({
state: "revoked",
lastError: `Hub rejected socket authentication (${status})`,
});
expect(hub.attemptCount()).toBe(1);
expect(clock.pendingTasks()).toBe(0);
expect(hub.openConnections()).toBe(0);
},
);
test("controller redials after a transient socket response and releases both failed upgrades", async () => {
const hub = await UpgradeRejectingHub.start([503, 401]);
const clock = new ManualRelationshipClock();
const controller = await connectController(hub, clock);
await hub.expectAttemptReleased(1);
expect(controller.status().state).toBe("reconnecting");
expect(clock.pendingTasks()).toBe(1);
clock.runNext();
await hub.expectAttemptReleased(2);
expect(hub.attemptCount()).toBe(2);
expect(controller.status().state).toBe("revoked");
expect(clock.pendingTasks()).toBe(0);
expect(hub.openConnections()).toBe(0);
});
test("controller redials once after a failed upgrade without also handling its close", async () => {
const hub = await UpgradeRejectingHub.start([0, 401]);
const clock = new ManualRelationshipClock();
const controller = await connectController(hub, clock);
await hub.expectAttemptReleased(1);
expect(controller.status().state).toBe("reconnecting");
expect(clock.scheduledAttempts()).toEqual([0]);
expect(clock.pendingTasks()).toBe(1);
clock.runNext();
await hub.expectAttemptReleased(2);
expect(hub.attemptCount()).toBe(2);
expect(controller.status().state).toBe("revoked");
expect(clock.pendingTasks()).toBe(0);
});
async function startHubReturning(status: number): Promise<string> {
const server = createServer((_request, response) => {
response.writeHead(status).end();
@@ -54,3 +175,214 @@ function closeServer(server: ReturnType<typeof createServer>): Promise<void> {
});
});
}
type SocketOutcome =
| { type: "connected" }
| { type: "rejected"; status: 401 | 403 }
| { type: "closed"; code: number }
| { type: "failed"; message: string };
interface Deferred<T> {
promise: Promise<T>;
resolve(value: T): void;
}
function deferred<T>(): Deferred<T> {
let resolve!: (value: T) => void;
const promise = new Promise<T>((accept) => {
resolve = accept;
});
return { promise, resolve };
}
class SocketOutcomes {
private readonly outcomes: SocketOutcome[] = [];
private nextOutcome = deferred<SocketOutcome>();
readonly events: HubSocketEvents = {
connected: () => this.record({ type: "connected" }),
rejected: (status) => this.record({ type: "rejected", status }),
closed: (code) => this.record({ type: "closed", code }),
failed: (error) => this.record({ type: "failed", message: error.message }),
};
next(): Promise<SocketOutcome> {
return withDeadline(this.nextOutcome.promise, "Hub socket produced no outcome");
}
async afterCleanup(): Promise<SocketOutcome[]> {
await new Promise<void>((resolve) => setImmediate(resolve));
return this.outcomes.slice();
}
private record(outcome: SocketOutcome): void {
this.outcomes.push(outcome);
this.nextOutcome.resolve(outcome);
this.nextOutcome = deferred<SocketOutcome>();
}
}
class UpgradeRejectingHub {
private readonly server = createServer(async (request, response) => {
let body = "";
for await (const chunk of request) body += chunk;
const enrollment = JSON.parse(body) as { relationshipId: string };
response.writeHead(200, { "content-type": "application/json" }).end(
JSON.stringify({
relationshipId: enrollment.relationshipId,
scopes: ["hub.*"],
webSocketUrl: this.webSocketUrl,
}),
);
});
private readonly sockets = new Set<Socket>();
private readonly releasedAttempts = new Map<number, Deferred<void>>();
private attemptObserved = deferred<void>();
private attempts = 0;
private constructor(private readonly statuses: number[]) {
this.server.on("upgrade", (_request, socket) => {
const attempt = ++this.attempts;
this.attemptObserved.resolve();
this.attemptObserved = deferred<void>();
const released = deferred<void>();
this.releasedAttempts.set(attempt, released);
this.sockets.add(socket);
socket.once("close", () => {
this.sockets.delete(socket);
released.resolve();
});
const status = this.statuses[attempt - 1] ?? this.statuses.at(-1) ?? 503;
if (status === 0) {
socket.destroy();
return;
}
socket.end(
`HTTP/1.1 ${status} Rejected\r\nContent-Length: 4\r\nConnection: close\r\n\r\ndeny`,
);
});
}
static async start(statuses: number[]): Promise<UpgradeRejectingHub> {
const hub = new UpgradeRejectingHub(statuses);
openUpgradeHubs.push(hub);
openServers.push(hub.server);
await new Promise<void>((resolve, reject) => {
hub.server.once("error", reject);
hub.server.listen(0, "127.0.0.1", resolve);
});
return hub;
}
get webSocketUrl(): string {
const address = this.server.address() as AddressInfo;
return `ws://127.0.0.1:${address.port}/daemon`;
}
get origin(): string {
const address = this.server.address() as AddressInfo;
return `http://127.0.0.1:${address.port}`;
}
attemptCount(): number {
return this.attempts;
}
openConnections(): number {
return this.sockets.size;
}
destroyConnections(): void {
for (const socket of this.sockets) socket.destroy();
}
async expectAttemptReleased(attempt: number): Promise<void> {
while (this.attempts < attempt) {
await withDeadline(
this.attemptObserved.promise,
`Hub did not receive socket attempt ${attempt}`,
);
}
const released = this.releasedAttempts.get(attempt);
if (!released) throw new Error(`Hub did not receive socket attempt ${attempt}`);
await withDeadline(released.promise, `Hub socket attempt ${attempt} remained open`);
await new Promise<void>((resolve) => setImmediate(resolve));
}
}
class ManualRelationshipClock implements HubRelationshipClock, HubRelationshipRetryPolicy {
private readonly tasks: Array<{ cancelled: boolean; task: () => void }> = [];
private readonly attempts: number[] = [];
now(): Date {
return new Date("2026-07-13T00:00:00.000Z");
}
delay(attempt: number): number {
this.attempts.push(attempt);
return 0;
}
schedule(_delayMs: number, task: () => void): ScheduledRelationshipTask {
const scheduled = { cancelled: false, task };
this.tasks.push(scheduled);
return { cancel: () => (scheduled.cancelled = true) };
}
pendingTasks(): number {
return this.tasks.filter((task) => !task.cancelled).length;
}
scheduledAttempts(): number[] {
return this.attempts.slice();
}
runNext(): void {
const scheduled = this.tasks.shift();
if (!scheduled || scheduled.cancelled) throw new Error("No Hub reconnect is scheduled");
scheduled.task();
}
}
const unusedExecutions: HubExecutions = {
create: async () => {
throw new Error("Unexpected Hub agent create");
},
reconcile: async () => null,
subscribe: () => () => undefined,
};
async function connectController(
hub: UpgradeRejectingHub,
clock: ManualRelationshipClock,
): Promise<HubRelationshipController> {
const paseoHome = await mkdtemp(path.join(tmpdir(), "paseo-hub-socket-"));
openPaseoHomes.push(paseoHome);
const controller = new HubRelationshipController({
paseoHome,
serverId: "server-1",
daemonPublicKey: "daemon-public-key",
logger: pino({ level: "silent" }),
remote: new DirectHubRelationshipRemote(),
clock,
retryPolicy: clock,
attachSocket: async () => undefined,
createExecutions: () => unusedExecutions,
});
await controller.connect({ hubUrl: hub.origin, token: "enrollment-token" });
return controller;
}
async function withDeadline<T>(promise: Promise<T>, message: string): Promise<T> {
let timer: ReturnType<typeof setTimeout> | undefined;
try {
return await Promise.race([
promise,
new Promise<T>((_resolve, reject) => {
timer = setTimeout(() => reject(new Error(message)), 2_000);
}),
]);
} finally {
if (timer) clearTimeout(timer);
}
}
@@ -104,14 +104,34 @@ export class DirectHubRelationshipRemote implements HubRelationshipRemote {
"x-paseo-relationship-id": input.relationshipId,
},
});
socket.once("open", () => events.connected(socket as WebSocketLike));
let settled = false;
socket.once("open", () => {
if (!settled) events.connected(socket as WebSocketLike);
});
socket.once("unexpected-response", (_request, response) => {
if (settled) {
response.destroy();
return;
}
settled = true;
response.destroy();
socket.terminate();
if (response.statusCode === 401 || response.statusCode === 403) {
events.rejected(response.statusCode);
return;
}
events.closed(1006);
});
socket.once("close", (code) => {
if (settled) return;
settled = true;
events.closed(code);
});
socket.once("error", (error) => {
if (settled) return;
settled = true;
events.failed(error);
});
socket.once("close", (code) => events.closed(code));
socket.once("error", (error) => events.failed(error));
return socket;
}
}
@@ -241,6 +241,7 @@ class ControlledAgentClient implements AgentClient {
private gate: Deferred<void> | null = null;
private creationObserved = deferred<void>();
creations = 0;
resumes = 0;
createdConfigs: AgentSessionConfig[] = [];
constructor(private readonly client: AgentClient) {
@@ -282,6 +283,7 @@ class ControlledAgentClient implements AgentClient {
overrides?: Partial<AgentSessionConfig>,
launchContext?: AgentLaunchContext,
): Promise<AgentSession> {
this.resumes++;
return this.client.resumeSession(handle, overrides, launchContext);
}
@@ -631,6 +633,10 @@ export class HubRelationshipHarness {
return this.codex.creations;
}
providerResumes(): number {
return this.codex.resumes;
}
latestCreatedCwd(): string | null {
return this.codex.createdConfigs.at(-1)?.cwd ?? null;
}
@@ -1133,6 +1139,7 @@ export class HubRelationshipHarness {
relationshipId: this.relationshipFile()!.relationship.id,
agentManager: manager,
agentStorage: storage,
logger: pino({ level: "silent" }),
createAgent: (input) =>
createAgentCommand(
{