From 3c5921f14379544e82a801a49a78330660ca4fde Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Mon, 13 Jul 2026 15:53:17 +0200 Subject: [PATCH] fix(hub): close relationship lifecycle gaps --- .../hub/relationship-controller.test.ts | 47 +++++++++++++++ .../src/server/hub/relationship-controller.ts | 58 +++++++++++++------ .../hub/relationship-owned-executions.test.ts | 21 +++++++ .../hub/relationship-owned-executions.ts | 31 ++++++---- .../hub/test-utils/relationship-harness.ts | 27 +++++++++ 5 files changed, 156 insertions(+), 28 deletions(-) diff --git a/packages/server/src/server/hub/relationship-controller.test.ts b/packages/server/src/server/hub/relationship-controller.test.ts index 10b3d654bb..48ff6b1797 100644 --- a/packages/server/src/server/hub/relationship-controller.test.ts +++ b/packages/server/src/server/hub/relationship-controller.test.ts @@ -256,6 +256,32 @@ describe("Hub relationship", () => { }, ); + test("a persisted non-WebSocket transport is quarantined before daemon startup", async () => { + relationship = await HubRelationshipHarness.start(); + await relationship.corruptRelationshipFile( + JSON.stringify({ + version: 1, + state: "active", + relationship: { + id: "relationship-1", + idempotencyKey: "ceremony-1", + hubOrigin: "https://hub.test", + createdAt: "2026-07-13T00:00:00.000Z", + scopes: ["hub.*"], + }, + credential: { secret: "credential" }, + transport: { kind: "direct_websocket", webSocketUrl: "ftp://hub.test/daemon" }, + }), + ); + + await relationship.startStoppedDaemon(); + + expect(await relationship.status()).toMatchObject({ state: "not_connected" }); + expect(relationship.relationshipFile()).toBeNull(); + expect(await relationship.quarantinedRelationshipFiles()).toHaveLength(1); + expect(relationship.socketAttempts()).toBe(0); + }); + test("disconnect revokes an ambiguous pending enrollment before removing local authority", async () => { relationship = await HubRelationshipHarness.start(); relationship.holdEnrollment(); @@ -284,6 +310,27 @@ describe("Hub relationship", () => { expect(relationship.relationshipFile()).toBeNull(); }); + test("disconnect revokes only after an in-flight enrollment settles", async () => { + relationship = await HubRelationshipHarness.start(); + relationship.holdEnrollment(); + const connecting = relationship.beginConnect("one-time-token"); + const enrollment = await relationship.enrollmentBegins(); + + const disconnecting = relationship.beginDisconnect(); + await relationship.relationshipStateBecomes("disconnecting"); + expect(relationship.revocationAttempts()).toBe(0); + + relationship.completeEnrollment(); + await connecting.result; + const disconnected = await disconnecting.result; + + expect(disconnected.state).toBe("not_connected"); + expect(relationship.latestRevocation()?.relationshipId).toBe(enrollment.relationshipId); + expect(relationship.revocationAttempts()).toBe(1); + expect(relationship.socketAttempts()).toBe(0); + expect(relationship.relationshipFile()).toBeNull(); + }); + test("daemon restart reconnects the same durable relationship", async () => { relationship = await HubRelationshipHarness.start(); await relationship.beginConnect().result; diff --git a/packages/server/src/server/hub/relationship-controller.ts b/packages/server/src/server/hub/relationship-controller.ts index 45caeeaf64..55b90932a8 100644 --- a/packages/server/src/server/hub/relationship-controller.ts +++ b/packages/server/src/server/hub/relationship-controller.ts @@ -28,7 +28,10 @@ const SanitizedRelationshipSchema = RelationshipSchema.omit({ idempotencyKey: tr const CredentialSchema = z.object({ secret: z.string().min(1) }); const TransportSchema = z.object({ kind: z.literal("direct_websocket"), - webSocketUrl: z.string().url(), + webSocketUrl: z + .string() + .url() + .refine((value) => ["ws:", "wss:"].includes(new URL(value).protocol)), }); const PendingSchema = z.object({ version: z.literal(1), @@ -59,12 +62,20 @@ const RevokedSchema = z.object({ transport: TransportSchema.optional(), reason: z.string().optional(), }); -const RecordSchema = z.discriminatedUnion("state", [ - PendingSchema, - ActiveSchema, - DisconnectingSchema, - RevokedSchema, -]); +const RecordSchema = z + .discriminatedUnion("state", [PendingSchema, ActiveSchema, DisconnectingSchema, RevokedSchema]) + .superRefine((record, context) => { + if (!("transport" in record) || !record.transport) return; + const hub = new URL(record.relationship.hubOrigin); + const socket = new URL(record.transport.webSocketUrl); + const expectedProtocol = hub.protocol === "https:" ? "wss:" : "ws:"; + if (socket.protocol === expectedProtocol && socket.host === hub.host) return; + context.addIssue({ + code: "custom", + path: ["transport", "webSocketUrl"], + message: "Hub WebSocket URL must match the Hub origin", + }); + }); type PendingRecord = z.infer; type ActiveRecord = z.infer; type DisconnectingRecord = z.infer; @@ -160,6 +171,7 @@ export class HubRelationshipController implements HubRelationshipManagement { private generation = 0; private enrollmentGeneration = 0; private retryAttempt = 0; + private readonly inFlightEnrollments = new Set>(); private executions: { relationshipId: string; value: HubExecutions } | null = null; constructor(private readonly options: HubRelationshipControllerOptions) { @@ -249,6 +261,7 @@ export class HubRelationshipController implements HubRelationshipManagement { async disconnect(input: { force: boolean; }): Promise<{ status: HubRelationshipStatus; warning?: string }> { + const waitForEnrollment = this.record?.state === "pending"; this.cancelLifecycle(); this.socket?.close(); this.socket = null; @@ -273,6 +286,9 @@ export class HubRelationshipController implements HubRelationshipManagement { this.persist(disconnecting); this.record = disconnecting; this.state = "disconnecting"; + if (waitForEnrollment) { + await Promise.all(this.inFlightEnrollments); + } await this.tryRevocation(disconnecting); return { status: this.status() }; } @@ -280,17 +296,23 @@ export class HubRelationshipController implements HubRelationshipManagement { private async tryEnrollment(pending: PendingRecord, enrollmentGeneration: number): Promise { if (enrollmentGeneration !== this.enrollmentGeneration) return; const verifier = createHash("sha256").update(pending.credential.secret).digest("base64url"); + const request = this.options.remote.enroll({ + relationshipId: pending.relationship.id, + idempotencyKey: pending.relationship.idempotencyKey, + hubOrigin: pending.relationship.hubOrigin, + token: pending.enrollment.token, + serverId: pending.identity.serverId, + daemonPublicKey: pending.identity.daemonPublicKey, + credentialVerifier: verifier, + scopes: pending.relationship.scopes, + }); + const settled = request.then( + () => undefined, + () => undefined, + ); + this.inFlightEnrollments.add(settled); try { - const enrollment = await this.options.remote.enroll({ - relationshipId: pending.relationship.id, - idempotencyKey: pending.relationship.idempotencyKey, - hubOrigin: pending.relationship.hubOrigin, - token: pending.enrollment.token, - serverId: pending.identity.serverId, - daemonPublicKey: pending.identity.daemonPublicKey, - credentialVerifier: verifier, - scopes: pending.relationship.scopes, - }); + const enrollment = await request; if (enrollmentGeneration !== this.enrollmentGeneration) return; if ( enrollment.relationshipId !== pending.relationship.id || @@ -318,6 +340,8 @@ export class HubRelationshipController implements HubRelationshipManagement { } this.lastError = error instanceof Error ? error.message : String(error); this.scheduleEnrollment(pending, enrollmentGeneration); + } finally { + this.inFlightEnrollments.delete(settled); } } diff --git a/packages/server/src/server/hub/relationship-owned-executions.test.ts b/packages/server/src/server/hub/relationship-owned-executions.test.ts index dd855a0430..686d4e99c5 100644 --- a/packages/server/src/server/hub/relationship-owned-executions.test.ts +++ b/packages/server/src/server/hub/relationship-owned-executions.test.ts @@ -23,7 +23,9 @@ test("sequential replay after reconstruction keeps one durable owned agent", asy const reconstructed = await hub.reconstructAndReplay(); expect(reconstructed.replay.agent.id).toBe(created.first.agentId); + expect(reconstructed.replay.agent.status).toBe("closed"); expect(reconstructed.reconciliation?.agent.id).toBe(created.first.agentId); + expect(reconstructed.reconciliation?.agent.status).toBe("closed"); expect(reconstructed.durableAgentCount).toBe(1); }); @@ -91,6 +93,25 @@ test("failed Hub auto-archive creates release their lifecycle subscriptions", as expect(hub.agentSubscriptionCount()).toBe(subscriptionBaseline); }); +test("failed Hub create cleans durable state when provider close rejects", async () => { + const hub = await launchRelationship(); + hub.failNextProviderPromptStart(); + hub.failNextProviderSessionClose(); + hub.beginOwnedCreate("failed-close-create", "failed-close-execution", { + worktree: { mode: "branch-off", newBranch: "failed-close-worktree" }, + }); + + const response = await hub.ownedCreateResult("failed-close-create"); + + expect(response).toMatchObject({ + type: "hub.agent.create.response", + payload: { success: false, executionId: "failed-close-execution" }, + }); + expect(hub.activeOwnedAgentIds()).toEqual([]); + expect(await hub.durableOwnedAgentIds()).toEqual([]); + expect(await hub.listedWorktrees()).toHaveLength(1); +}); + test("Hub checkout uses the requested branch ref", async () => { const hub = await launchRelationship(); await hub.createBranch("existing-hub-branch"); diff --git a/packages/server/src/server/hub/relationship-owned-executions.ts b/packages/server/src/server/hub/relationship-owned-executions.ts index 50b6ed4634..e27c905eb9 100644 --- a/packages/server/src/server/hub/relationship-owned-executions.ts +++ b/packages/server/src/server/hub/relationship-owned-executions.ts @@ -165,16 +165,22 @@ export class RelationshipOwnedExecutions implements HubExecutions { }, }); } catch (error) { - await autoArchiveRegistration.cancel(); - if (createdAgentId && this.agentManager.getAgent(createdAgentId)) { - await this.agentManager.closeAgent(createdAgentId); - } - await this.cleanupFailedCreate({ - createdWorktree: ownedCreatedWorktree(createdWorktree), - createdAgentId: null, - }); - if (createdAgentId) { - await this.agentStorage.remove(createdAgentId); + try { + await autoArchiveRegistration.cancel(); + if (createdAgentId && this.agentManager.getAgent(createdAgentId)) { + await this.agentManager.closeAgent(createdAgentId); + } + } finally { + try { + await this.cleanupFailedCreate({ + createdWorktree: ownedCreatedWorktree(createdWorktree), + createdAgentId: null, + }); + } finally { + if (createdAgentId) { + await this.agentStorage.remove(createdAgentId); + } + } } throw error; } @@ -196,7 +202,10 @@ export class RelationshipOwnedExecutions implements HubExecutions { executionId: owner.executionId, agent: live ? serializeAgentSnapshot(live) - : buildStoredAgentPayload(record, this.agentManager.getRegisteredProviderIds()), + : { + ...buildStoredAgentPayload(record, this.agentManager.getRegisteredProviderIds()), + status: "closed", + }, }; } diff --git a/packages/server/src/server/hub/test-utils/relationship-harness.ts b/packages/server/src/server/hub/test-utils/relationship-harness.ts index 5da47b43fe..87c06f7ca5 100644 --- a/packages/server/src/server/hub/test-utils/relationship-harness.ts +++ b/packages/server/src/server/hub/test-utils/relationship-harness.ts @@ -447,11 +447,17 @@ export class HubRelationshipHarness { private readonly cliProcesses = new Set>(); private readonly claimedCliSockets = new Set(); private failNextPromptStart = false; + private failNextSessionClose = false; private observedEnrollments = 0; private observedSockets = 0; private reconcileRequests = 0; private readonly codex = new ControlledAgentClient( createTestAgentClients({ + closeSession: async () => { + if (!this.failNextSessionClose) return; + this.failNextSessionClose = false; + throw new Error("Requested provider session close failure"); + }, onStartTurn: (prompt) => { if (this.failNextPromptStart) { this.failNextPromptStart = false; @@ -489,6 +495,23 @@ export class HubRelationshipHarness { return this.runCli(["hub", "disconnect", ...(force ? ["--force"] : [])]); } + beginDisconnect(force = false): CliProcess { + return { result: this.runCli(["hub", "disconnect", ...(force ? ["--force"] : [])]) }; + } + + async relationshipStateBecomes(expected: string): Promise { + const observed = deferred(); + const watcher = watch(this.paseoHome, () => { + if (this.relationshipFile()?.state === expected) observed.resolve(); + }); + if (this.relationshipFile()?.state === expected) observed.resolve(); + try { + await observed.promise; + } finally { + watcher.close(); + } + } + async attemptExternalManagementAsCli(): Promise { const address = Object.values(networkInterfaces()) .flat() @@ -575,6 +598,10 @@ export class HubRelationshipHarness { this.failNextPromptStart = true; } + failNextProviderSessionClose(): void { + this.failNextSessionClose = true; + } + async socketDialed(): Promise { await this.remote.socketAt(this.observedSockets++); }