diff --git a/sdk/packages/core/src/hub/server/boundary.test.ts b/sdk/packages/core/src/hub/server/boundary.test.ts index c1dcdde657..4d62e137b4 100644 --- a/sdk/packages/core/src/hub/server/boundary.test.ts +++ b/sdk/packages/core/src/hub/server/boundary.test.ts @@ -1749,18 +1749,18 @@ describe("HubServerTransport boundaries", () => { ); }); - it("fails capability requests to a client that unregistered and did not come back", async () => { + it("holds capability requests for a disconnected owner through a short reconnect", async () => { vi.useFakeTimers(); try { const transport = createTransport(); const ctx = getContext(transport); - const register = (clientId: string) => + const register = () => transport.handleCommand({ version: "v1", - requestId: `req-register-${clientId}`, + requestId: "req-register", command: "client.register", - clientId, - payload: { clientId, clientType: "desktop" }, + clientId: "owner-client", + payload: { clientId: "owner-client", clientType: "desktop" }, }); const unregister = () => transport.handleCommand({ @@ -1769,12 +1769,27 @@ describe("HubServerTransport boundaries", () => { command: "client.unregister", clientId: "owner-client", }); - await register("owner-client"); + const requested: string[] = []; + const resubscribe = async () => { + await register(); + const unsubscribe = transport.subscribe( + "owner-client", + (event) => { + if (event.event === "capability.requested") { + requested.push(String(event.payload?.capabilityName)); + } + }, + { sessionId: "session-1" }, + ); + await vi.advanceTimersByTimeAsync(0); + return unsubscribe; + }; + await register(); await unregister(); - // A brief socket drop: the same client reconnects within the grace - // period and gets the request. - void requestCapability( + // A brief socket drop: the request waits, and reaches the same + // client once it registers and subscribes again. + const answered = requestCapability( ctx, "session-1", "hook.afterTool", @@ -1782,31 +1797,57 @@ describe("HubServerTransport boundaries", () => { "owner-client", ); await vi.advanceTimersByTimeAsync(1_000); - expect(ctx.pendingCapabilityRequests.size).toBe(0); - await register("owner-client"); - await vi.advanceTimersByTimeAsync(200); - expect( - [...ctx.pendingCapabilityRequests.values()].map( - (r) => r.targetClientId, - ), - ).toEqual(["owner-client"]); - ctx.pendingCapabilityRequests.clear(); + expect(requested).toEqual([]); + const unsubscribe = await resubscribe(); + expect(requested).toEqual(["hook.afterTool"]); + const [requestId] = [...ctx.pendingCapabilityRequests.keys()]; + await transport.handleCommand({ + version: "v1", + requestId: "req-respond", + command: "capability.respond", + clientId: "owner-client", + payload: { requestId, ok: true, payload: { done: true } }, + }); + await expect(answered).resolves.toEqual({ done: true }); + unsubscribe(); await unregister(); - // A run still in flight after its client quit raises its next hook: - // it fails once the grace period passes instead of waiting forever. - const request = requestCapability( + // Stop while waiting for the owner cancels the request; it is not + // delivered when the owner comes back. + const stopped = requestCapability( ctx, "session-1", "hook.beforeTool", {}, "owner-client", ); - const failed = expect(request).rejects.toThrow( + const stoppedFails = expect(stopped).rejects.toThrow("user stop"); + await transport.handleCommand({ + version: "v1", + requestId: "req-abort", + command: "run.abort", + sessionId: "session-1", + payload: { reason: "user stop" }, + }); + await stoppedFails; + (await resubscribe())(); + expect(requested).toEqual(["hook.afterTool"]); + await unregister(); + + // The owner never comes back: the request fails once the grace + // period passes instead of waiting forever. + const abandoned = requestCapability( + ctx, + "session-1", + "hook.beforeTool", + {}, + "owner-client", + ); + const abandonedFails = expect(abandoned).rejects.toThrow( /disconnected before it could respond/, ); await vi.advanceTimersByTimeAsync(CAPABILITY_OWNER_RECONNECT_GRACE_MS); - await failed; + await abandonedFails; expect(ctx.pendingCapabilityRequests.size).toBe(0); } finally { vi.useRealTimers(); diff --git a/sdk/packages/core/src/hub/server/handlers/capability-handlers.ts b/sdk/packages/core/src/hub/server/handlers/capability-handlers.ts index 9467b3fd25..5e16e54b59 100644 --- a/sdk/packages/core/src/hub/server/handlers/capability-handlers.ts +++ b/sdk/packages/core/src/hub/server/handlers/capability-handlers.ts @@ -1,4 +1,8 @@ -import type { HubCommandEnvelope, HubReplyEnvelope } from "@cline/shared"; +import type { + HubCommandEnvelope, + HubEventEnvelope, + HubReplyEnvelope, +} from "@cline/shared"; import { createSessionId } from "@cline/shared"; import { captureCapabilityOwnerGone } from "../../../services/telemetry/core-events"; import { logHubMessage } from "../hub-server-logging"; @@ -19,16 +23,28 @@ export const CAPABILITY_OWNER_GONE_MESSAGE = */ export const CAPABILITY_OWNER_RECONNECT_GRACE_MS = 7_000; -async function waitForClient( +/** + * Requests that waited for their owner to reconnect, to deliver to that owner + * now that it is subscribed again. Each is returned once. + */ +export function takeUndeliveredCapabilityEvents( ctx: HubTransportContext, clientId: string, -): Promise { - // ponytail: polls every 100ms; switch to a register listener if many requests wait at once. - const deadline = Date.now() + CAPABILITY_OWNER_RECONNECT_GRACE_MS; - while (!ctx.clients.has(clientId) && Date.now() < deadline) { - await new Promise((resolve) => setTimeout(resolve, 100)); + sessionId?: string, +): HubEventEnvelope[] { + const events: HubEventEnvelope[] = []; + for (const pending of ctx.pendingCapabilityRequests.values()) { + if ( + !pending.undeliveredEvent || + pending.targetClientId !== clientId || + (sessionId && pending.sessionId !== sessionId) + ) { + continue; + } + events.push(pending.undeliveredEvent); + pending.undeliveredEvent = undefined; } - return ctx.clients.has(clientId); + return events; } export async function requestCapability( @@ -47,43 +63,51 @@ export async function requestCapability( capabilityName, targetClientId, }); - // Unregistering or detaching cancels the owner's pending requests, but a - // request raised after that (a run still in flight when its client - // restarted) would wait forever: nobody answers it and nothing times it - // out. A disconnected owner gets a short grace period to reconnect under - // the same id. Requests are never handed to another client: a client that - // joined later may enforce different hooks or tools under the same name. - const goneReason = ctx.detachedClientsBySession - ?.get(sessionId) - ?.has(targetClientId) - ? "detached" - : !ctx.clients.has(targetClientId) && - !(await waitForClient(ctx, targetClientId)) - ? "disconnected" - : undefined; - if (goneReason) { + const fail = (reason: "detached" | "disconnected") => { logHubMessage("warn", "capability.request.end", { requestId, sessionId, capabilityName, targetClientId, ok: false, - error: `target_client_${goneReason}`, + error: `target_client_${reason}`, durationMs: Math.round(performance.now() - startedAt), }); captureCapabilityOwnerGone(ctx.telemetry, { sessionId, capabilityName, - reason: goneReason, + reason, }); - throw new Error(CAPABILITY_OWNER_GONE_MESSAGE); + return new Error(CAPABILITY_OWNER_GONE_MESSAGE); + }; + // Unregistering or detaching cancels the owner's pending requests, but a + // request raised after that (a run still in flight when its client + // restarted) would wait forever: nobody answers it and nothing times it + // out. Requests are never handed to another client: a client that joined + // later may enforce different hooks or tools under the same name. + if (ctx.detachedClientsBySession?.get(sessionId)?.has(targetClientId)) { + throw fail("detached"); } return await new Promise((resolve, reject) => { + const event = ctx.buildEvent( + "capability.requested", + { + requestId, + targetClientId, + capabilityName, + payload, + }, + sessionId, + ); + const connected = ctx.clients.has(targetClientId); + // Registered before any waiting, so Stop (or a detach) cancels it + // whether it is waiting for the owner to reconnect or to answer. ctx.pendingCapabilityRequests.set(requestId, { sessionId, targetClientId, capabilityName, onProgress, + undeliveredEvent: connected ? undefined : event, resolve: (result) => { logHubMessage(result.ok ? "info" : "warn", "capability.request.end", { requestId, @@ -113,18 +137,19 @@ export async function requestCapability( resolve(result.payload); }, }); - ctx.publish( - ctx.buildEvent( - "capability.requested", - { - requestId, - targetClientId, - capabilityName, - payload, - }, - sessionId, - ), - ); + if (!connected) { + // A disconnected owner gets a short grace period to come back under the + // same id (a brief socket drop); the request reaches it when it + // subscribes again (see takeUndeliveredCapabilityEvents). + setTimeout(() => { + const pending = ctx.pendingCapabilityRequests.get(requestId); + if (!pending?.undeliveredEvent) return; + ctx.pendingCapabilityRequests.delete(requestId); + reject(fail("disconnected")); + }, CAPABILITY_OWNER_RECONNECT_GRACE_MS); + return; + } + ctx.publish(event); logHubMessage("info", "capability.request.published", { requestId, sessionId, diff --git a/sdk/packages/core/src/hub/server/handlers/context.ts b/sdk/packages/core/src/hub/server/handlers/context.ts index a00a69fdd4..c6afe8351d 100644 --- a/sdk/packages/core/src/hub/server/handlers/context.ts +++ b/sdk/packages/core/src/hub/server/handlers/context.ts @@ -44,6 +44,11 @@ export type PendingCapabilityRequest = { targetClientId: string; capabilityName: string; onProgress?: (payload: Record) => void; + /** + * Set while the owner is disconnected and the request has not reached it + * yet; delivered once when the owner subscribes again. + */ + undeliveredEvent?: HubEventEnvelope; resolve: (result: { ok: boolean; payload?: Record; diff --git a/sdk/packages/core/src/hub/server/hub-server-transport.ts b/sdk/packages/core/src/hub/server/hub-server-transport.ts index 0cf276fe6b..57c2ee89ae 100644 --- a/sdk/packages/core/src/hub/server/hub-server-transport.ts +++ b/sdk/packages/core/src/hub/server/hub-server-transport.ts @@ -64,6 +64,7 @@ import { handleCapabilityRequest, handleCapabilityRespond, requestCapability as requestCapabilityHandler, + takeUndeliveredCapabilityEvents, } from "./handlers/capability-handlers"; import { handleClientList, @@ -1264,26 +1265,36 @@ export class HubServerTransport implements NativeHubTransport { current.add(entry); this.listeners.set(clientId, current); // Re-issue pending approvals so a (re)connecting client can answer a - // request raised while it was away instead of leaving the turn parked. + // request raised while it was away instead of leaving the turn parked, + // and deliver capability requests that waited for this client to + // reconnect. const pending = pendingApprovalEvents(this.ctx, options?.sessionId); - if (pending.length > 0) { - queueMicrotask(() => { - const listeners = this.listeners.get(clientId); - if (!listeners?.has(entry)) { - return; + queueMicrotask(() => { + const listeners = this.listeners.get(clientId); + if (!listeners?.has(entry)) { + return; + } + // Taken here, not at subscribe time, so a listener removed before + // this runs leaves them for the next subscription. + const events = [ + ...pending, + ...takeUndeliveredCapabilityEvents( + this.ctx, + clientId, + options?.sessionId, + ), + ]; + for (const event of events) { + try { + entry.listener(event); + } catch (error) { + logHubBoundaryError( + "listener threw while re-issuing pending request", + error, + ); } - for (const event of pending) { - try { - entry.listener(event); - } catch (error) { - logHubBoundaryError( - "listener threw while re-issuing pending approval", - error, - ); - } - } - }); - } + } + }); return () => { const listeners = this.listeners.get(clientId); if (!listeners) {