From 783a14bd145f5eb03fb7b9d10d01a9bf03475633 Mon Sep 17 00:00:00 2001 From: Tami Komssi Date: Tue, 1 Sep 2026 23:24:17 +0300 Subject: [PATCH 1/5] fix(daemon): cancel retired extension UI requests --- packages/coding-agent/CHANGELOG.md | 1 + .../coding-agent/docs/agent-connection.md | 4 +- .../daemon-agent-connection.ts | 8 ++- .../src/modes/agent-connection/types.ts | 1 + .../modes/daemon/daemon-extension-binding.ts | 8 ++- .../src/modes/daemon/daemon-mode.ts | 5 ++ .../src/modes/daemon/daemon-protocol.ts | 13 ++++- .../src/modes/interactive/interactive-mode.ts | 11 ++++ .../test/agent-connection-daemon.test.ts | 49 ++++++++++++++-- .../test/daemon-extension-binding.test.ts | 57 +++++++++++++++++++ .../coding-agent/test/daemon-mode.test.ts | 17 +++++- .../coding-agent/test/daemon-protocol.test.ts | 9 +++ .../test/interactive-mode-status.test.ts | 39 +++++++++++++ 13 files changed, 209 insertions(+), 13 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 60589450e1..2e6631509d 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,7 @@ ## [Unreleased] +- Fixed daemon-aborted extension dialogs remaining live in the interactive client and sending a stale response after another client had already decided the request. - Fixed fullscreen wheel scrolling in Ghostty while retaining application link clicks; set `terminal.fullscreenMouse` to `false` to use native Cmd-click instead. - Changed the agents view to sort idle and inactive sessions by last message time, newest first, while keeping running agents in stable creation order. - Fixed `openai-codex` models being invisible to `rlm` subagents and `find_models` because model discovery reported Prime Agent's own version as the Codex client version ([#1375](https://github.com/PrimeIntellect-ai/prime-agent/pull/1375) by [@bilelrais](https://github.com/bilelrais)). diff --git a/packages/coding-agent/docs/agent-connection.md b/packages/coding-agent/docs/agent-connection.md index 1c76f3c31b..2e05be1572 100644 --- a/packages/coding-agent/docs/agent-connection.md +++ b/packages/coding-agent/docs/agent-connection.md @@ -80,7 +80,7 @@ An initial or replacement snapshot combines: - active RLM child snapshots; and - an in-progress assistant message when one exists. -Connection events cover session events, replacement and resynchronization snapshots, extension UI requests, connection status, and terminal closure. The adapter updates its cache before notifying the UI. +Connection events cover session events, replacement and resynchronization snapshots, extension UI requests and exact-request cancellations, connection status, and terminal closure. The adapter updates its cache before notifying the UI. Some connection types still reuse internal `AgentMessage`, `AgentEvent`, and model types. Those are local TypeScript contracts, not promises of a stable public network schema. @@ -116,7 +116,7 @@ When switching to a session already owned by another resident worker, a non-owne ## Extension UI Boundary -Daemon-owned extensions can request serializable UI operations such as select, confirm, input, editor, notification, status, widget, title, and editor-text updates. The client validates the payload and returns a serializable response. +Daemon-owned extensions can request serializable UI operations such as select, confirm, input, editor, notification, status, widget, title, and editor-text updates. The client validates the payload and returns a serializable response. If the daemon aborts or times out a pending dialog, it sends a capability-gated cancellation for that exact request ID; the interactive client retires only that request without sending a response. Executable callbacks are deliberately excluded: diff --git a/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts b/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts index 016a053279..4ecc8b3676 100644 --- a/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts +++ b/packages/coding-agent/src/modes/agent-connection/daemon-agent-connection.ts @@ -305,7 +305,7 @@ export class DaemonAgentConnection implements AgentConnection { capabilities: [ "attach_snapshot", "event_sequence", - ...(supportsExtensionUi ? (["extension_ui"] as const) : []), + ...(supportsExtensionUi ? (["extension_ui", "extension_ui_cancellation"] as const) : []), "slim_attach", "chunked_snapshot", ...(this.options.ownedSession ? (["client_owned_sessions"] as const) : []), @@ -1156,7 +1156,7 @@ export class DaemonAgentConnection implements AgentConnection { capabilities: [ "attach_snapshot", "event_sequence", - ...(supportsExtensionUi ? (["extension_ui"] as const) : []), + ...(supportsExtensionUi ? (["extension_ui", "extension_ui_cancellation"] as const) : []), "slim_attach", "chunked_snapshot", ...(this.options.ownedSession ? (["client_owned_sessions"] as const) : []), @@ -1555,6 +1555,10 @@ export class DaemonAgentConnection implements AgentConnection { }); return; } + if (message.type === "extension_ui_cancelled") { + await this.emit({ type: "extension_ui_cancelled", requestId: message.id }); + return; + } if (message.type === "extension_error") { await this.emit({ type: "extension_error", diff --git a/packages/coding-agent/src/modes/agent-connection/types.ts b/packages/coding-agent/src/modes/agent-connection/types.ts index 7e98f96681..89ef555caa 100644 --- a/packages/coding-agent/src/modes/agent-connection/types.ts +++ b/packages/coding-agent/src/modes/agent-connection/types.ts @@ -624,6 +624,7 @@ export type AgentConnectionEvent = | { type: "session_resynced"; snapshot: AgentConnectionSnapshot } | { type: "session_status"; recap?: string } | { type: "extension_ui_request"; request: AgentConnectionExtensionUiRequest } + | { type: "extension_ui_cancelled"; requestId: string } | { type: "extension_error"; extensionPath: string; event: string; error: string } | { type: "connection_status"; status: "reconnecting" | "connected"; error?: string } | { type: "heartbeats_changed" } diff --git a/packages/coding-agent/src/modes/daemon/daemon-extension-binding.ts b/packages/coding-agent/src/modes/daemon/daemon-extension-binding.ts index cecd046186..0c81f349fb 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-extension-binding.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-extension-binding.ts @@ -159,7 +159,13 @@ function createExtensionUIContext( clearTimeout(timeoutId); } opts?.signal?.removeEventListener("abort", onAbort); - state.extensionUiRequests.delete(requestId); + if (state.extensionUiRequests.delete(requestId)) { + broadcast(state, { + type: "extension_ui_cancelled", + activeSessionId: state.activeSessionId, + id: requestId, + }); + } }; const finish = (value: T) => { cleanup(); diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index 68667dce9d..f888f515ae 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -6888,6 +6888,7 @@ type SequencedDaemonOutbound = Extract< | "session_resynced" | "session_closed" | "extension_ui_request" + | "extension_ui_cancelled" | "extension_error"; } >; @@ -6900,11 +6901,15 @@ function isSequencedSessionOutbound(message: DaemonOutbound): message is Sequenc message.type === "session_resynced" || message.type === "session_closed" || message.type === "extension_ui_request" || + message.type === "extension_ui_cancelled" || message.type === "extension_error" ); } export function shouldSendDaemonOutboundToClient(client: DaemonSocketClient, message: DaemonOutbound): boolean { + if (message.type === "extension_ui_cancelled") { + return daemonClientCapabilitiesForSession(client, message.activeSessionId).has("extension_ui_cancellation"); + } return ( message.type !== "extension_ui_request" || !isDaemonDialogExtensionUiRequest(message.method) || diff --git a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts index b26eb1bdc2..a0304efc3c 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-protocol.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-protocol.ts @@ -60,8 +60,9 @@ export const DAEMON_COMMAND_ENVELOPE_MIN_PROTOCOL_VERSION = 7; // Revision 14 carries the client's monotonic telemetry opt-out on attach and reattach. // Revision 15 adds the mutate_queued_message command and queue_message_mutation capability. // Revision 16 adds the "stopping" workerState and stops reporting disconnected workers as "ready". -export const DAEMON_SCHEMA_REVISION = 16; -export const DAEMON_SCHEMA_ID = "protocol-7-schema-16-1bcb9e7f1a49"; +// Revision 17 adds capability-gated exact-request extension UI cancellation. +export const DAEMON_SCHEMA_REVISION = 17; +export const DAEMON_SCHEMA_ID = "protocol-7-schema-17-28439aeb8512"; export type DaemonProtocolName = typeof DAEMON_PROTOCOL_NAME; export type DaemonProtocolVersion = number; @@ -77,6 +78,7 @@ export type DaemonClientCapability = | "attach_snapshot" | "event_sequence" | "extension_ui" + | "extension_ui_cancellation" | "slim_attach" | "chunked_snapshot" | "client_owned_sessions"; @@ -123,6 +125,7 @@ export const DAEMON_SUPPORTED_CLIENT_CAPABILITIES: readonly DaemonClientCapabili "attach_snapshot", "event_sequence", "extension_ui", + "extension_ui_cancellation", "slim_attach", "chunked_snapshot", "client_owned_sessions", @@ -938,6 +941,7 @@ export type DaemonOutbound = payload: Record; meta?: DaemonEventMeta; } + | { type: "extension_ui_cancelled"; activeSessionId: string; id: string; meta?: DaemonEventMeta } | { type: "extension_error"; activeSessionId: string; @@ -967,6 +971,11 @@ export const DAEMON_OUTBOUND_COMPATIBILITY = { session_detached: LEGACY_DAEMON_COMMAND, session_closed: LEGACY_DAEMON_COMMAND, extension_ui_request: LEGACY_DAEMON_COMMAND, + extension_ui_cancelled: { + minProtocol: 7, + minSchemaRevision: 17, + capability: "extension_ui_cancellation", + }, extension_error: LEGACY_DAEMON_COMMAND, } as const satisfies Record; diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index db3d89d006..1f5f2b7131 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -5135,6 +5135,8 @@ export class InteractiveMode { this.handleSideQuestionEvent(event.event); } else if (event.type === "extension_ui_request") { await this.handleConnectionExtensionUiRequest(event.request); + } else if (event.type === "extension_ui_cancelled") { + this.cancelConnectionExtensionUiRequest(event.requestId); } else if (event.type === "connection_status") { this.showStatus( event.status === "connected" ? "Daemon reconnected" : "Daemon connection lost; reconnecting…", @@ -5201,6 +5203,15 @@ export class InteractiveMode { ); } + private cancelConnectionExtensionUiRequest(requestId: string): void { + const request = this.activeConnectionExtensionUiRequests.get(requestId); + if (!request) { + return; + } + this.activeConnectionExtensionUiRequests.delete(requestId); + request.cancelLocal(); + } + private cancelActiveConnectionExtensionUiRequests(): void { const requestIds = [...this.activeConnectionExtensionUiRequests.keys()]; for (const requestId of requestIds) { diff --git a/packages/coding-agent/test/agent-connection-daemon.test.ts b/packages/coding-agent/test/agent-connection-daemon.test.ts index 7ff21e88c4..edc6002255 100644 --- a/packages/coding-agent/test/agent-connection-daemon.test.ts +++ b/packages/coding-agent/test/agent-connection-daemon.test.ts @@ -2006,7 +2006,14 @@ describe("DaemonAgentConnection", () => { expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-1", - capabilities: ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"], + capabilities: [ + "attach_snapshot", + "event_sequence", + "extension_ui", + "extension_ui_cancellation", + "slim_attach", + "chunked_snapshot", + ], resumeCursor: { activeSessionId: "active-1", generation: "generation-active-1", @@ -2501,7 +2508,14 @@ describe("DaemonAgentConnection", () => { type: "attach", activeSessionId: "active-1", clientId: expect.any(String), - capabilities: ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"], + capabilities: [ + "attach_snapshot", + "event_sequence", + "extension_ui", + "extension_ui_cancellation", + "slim_attach", + "chunked_snapshot", + ], resumeCursor: { activeSessionId: "active-1", generation: "generation-active-1", @@ -2513,7 +2527,14 @@ describe("DaemonAgentConnection", () => { expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-1", - capabilities: ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"], + capabilities: [ + "attach_snapshot", + "event_sequence", + "extension_ui", + "extension_ui_cancellation", + "slim_attach", + "chunked_snapshot", + ], resumeCursor: { activeSessionId: "active-1", generation: "generation-active-1", @@ -2559,7 +2580,7 @@ describe("DaemonAgentConnection", () => { }); }); - it("forwards extension UI requests and responses", async () => { + it("forwards extension UI requests, cancellations, and responses", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const events: AgentConnectionEvent[] = []; @@ -2572,7 +2593,14 @@ describe("DaemonAgentConnection", () => { activeSessionId: "active-1", supportsExtensionUi: true, clientId: expect.any(String), - capabilities: ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"], + capabilities: [ + "attach_snapshot", + "event_sequence", + "extension_ui", + "extension_ui_cancellation", + "slim_attach", + "chunked_snapshot", + ], }); fakeClient.emitMessage({ @@ -2589,6 +2617,16 @@ describe("DaemonAgentConnection", () => { method: "confirm", payload: { title: "Ignore", message: "Wrong session" }, }); + fakeClient.emitMessage({ + type: "extension_ui_cancelled", + activeSessionId: "active-1", + id: "request-1", + }); + fakeClient.emitMessage({ + type: "extension_ui_cancelled", + activeSessionId: "other", + id: "request-2", + }); expect(events).toEqual([ { @@ -2599,6 +2637,7 @@ describe("DaemonAgentConnection", () => { payload: { title: "Confirm", message: "Proceed?" }, }, }, + { type: "extension_ui_cancelled", requestId: "request-1" }, ]); await connection.respondToExtensionUiRequest("request-1", { confirmed: true }); diff --git a/packages/coding-agent/test/daemon-extension-binding.test.ts b/packages/coding-agent/test/daemon-extension-binding.test.ts index c71f4ba565..6edb19ea15 100644 --- a/packages/coding-agent/test/daemon-extension-binding.test.ts +++ b/packages/coding-agent/test/daemon-extension-binding.test.ts @@ -148,6 +148,63 @@ describe("daemon extension binding", () => { } }); + it("broadcasts exact request cancellation when a daemon extension dialog aborts", async () => { + const abortController = new AbortController(); + let markDialogStarted: (() => void) | undefined; + const dialogStarted = new Promise((resolve) => { + markDialogStarted = resolve; + }); + let dialogResult: boolean | undefined; + const runtime = await createRuntimeForTest((pi) => { + pi.registerCommand("daemon-dialog", { + description: "daemon dialog", + handler: async (_args, ctx) => { + markDialogStarted?.(); + dialogResult = await ctx.ui.confirm("Continue?", "Confirm action", { + signal: abortController.signal, + }); + }, + }); + }, []); + + const outbound: DaemonOutbound[] = []; + const state: ActiveSessionState = { + activeSessionId: "active-dialog", + runtime, + clients: new Set([{ supportsExtensionUi: true }] as never[]), + pendingAttaches: 0, + extensionUiRequests: new Map(), + eventGeneration: "generation-dialog", + lastEventSequence: 0, + }; + await bindActiveSessionState(state, { + broadcast: (_state, message) => { + outbound.push(message); + }, + shutdown: () => {}, + }); + + const prompt = runtime.session.prompt("/daemon-dialog"); + await dialogStarted; + const request = outbound.find( + (message): message is Extract => + message.type === "extension_ui_request", + ); + expect(request).toBeDefined(); + expect(state.extensionUiRequests.has(request?.id ?? "")).toBe(true); + + abortController.abort(); + await prompt; + + expect(dialogResult).toBe(false); + expect(state.extensionUiRequests.size).toBe(0); + expect(outbound).toContainEqual({ + type: "extension_ui_cancelled", + activeSessionId: state.activeSessionId, + id: request?.id, + }); + }); + it("keeps extension replacement callbacks daemon-side and rebinds before withSession", async () => { const phases: string[] = []; let oldSessionFile: string | undefined; diff --git a/packages/coding-agent/test/daemon-mode.test.ts b/packages/coding-agent/test/daemon-mode.test.ts index cb4b03a5c5..06214039c9 100644 --- a/packages/coding-agent/test/daemon-mode.test.ts +++ b/packages/coding-agent/test/daemon-mode.test.ts @@ -4056,15 +4056,30 @@ describe("daemon mode helpers", () => { }), ).toBe(true); - setDaemonClientSessionCapabilities(uiClient, "active", new Set(["extension_ui"])); + const cancellation = { + type: "extension_ui_cancelled", + activeSessionId: "active", + id: "request-1", + } as const; + expect(shouldSendDaemonOutboundToClient(lineClient, cancellation)).toBe(false); + expect(shouldSendDaemonOutboundToClient(uiClient, cancellation)).toBe(false); + + setDaemonClientSessionCapabilities(uiClient, "active", new Set(["extension_ui", "extension_ui_cancellation"])); setDaemonClientSessionCapabilities(uiClient, "other", new Set()); expect(shouldSendDaemonOutboundToClient(uiClient, dialogRequest)).toBe(true); + expect(shouldSendDaemonOutboundToClient(uiClient, cancellation)).toBe(true); expect( shouldSendDaemonOutboundToClient(uiClient, { ...dialogRequest, activeSessionId: "other", }), ).toBe(false); + expect( + shouldSendDaemonOutboundToClient(uiClient, { + ...cancellation, + activeSessionId: "other", + }), + ).toBe(false); }); it("delivers session closure while a client is snapshotting and backpressured", () => { diff --git a/packages/coding-agent/test/daemon-protocol.test.ts b/packages/coding-agent/test/daemon-protocol.test.ts index d7eb8acf4a..2bdff0754f 100644 --- a/packages/coding-agent/test/daemon-protocol.test.ts +++ b/packages/coding-agent/test/daemon-protocol.test.ts @@ -130,6 +130,15 @@ describe("daemon protocol helpers", () => { expect(DAEMON_DEFAULT_SERVER_CAPABILITIES).toContain("prompt_admission_cancellation"); }); + it("capability-gates exact extension UI cancellation", () => { + expect(DAEMON_OUTBOUND_COMPATIBILITY.extension_ui_cancelled).toEqual({ + minProtocol: 7, + minSchemaRevision: 17, + capability: "extension_ui_cancellation", + }); + expect(DAEMON_DEFAULT_SERVER_CAPABILITIES).toContain("extension_ui_cancellation"); + }); + it("gates honest worker-state reporting at its introducing schema revision", () => { // Revision 16 adds the "stopping" workerState and stops reporting // disconnected workers as "ready". The field is optional and old clients diff --git a/packages/coding-agent/test/interactive-mode-status.test.ts b/packages/coding-agent/test/interactive-mode-status.test.ts index 3d8ae677ec..96d9e76ce4 100644 --- a/packages/coding-agent/test/interactive-mode-status.test.ts +++ b/packages/coding-agent/test/interactive-mode-status.test.ts @@ -1623,6 +1623,7 @@ describe("InteractiveMode connection extension UI", () => { }; showError(message: string): void; cancelActiveConnectionExtensionUiRequests(): void; + cancelConnectionExtensionUiRequest(requestId: string): void; }; type ConnectionExtensionUiHandlerHarness = ConnectionExtensionUiCancelHarness & { @@ -1685,6 +1686,44 @@ describe("InteractiveMode connection extension UI", () => { }); expect(fakeThis.activeConnectionExtensionUiRequests.size).toBe(0); }); + + test("daemon abort cancellation retires the exact request before a late local response", async () => { + const fakeThis = Object.create(InteractiveMode.prototype) as ConnectionExtensionUiHandlerHarness; + fakeThis.activeConnectionExtensionUiRequests = new Map(); + fakeThis.agentConnection = { + respondToExtensionUiRequest: vi.fn(async () => { + throw new Error("Unknown extension UI request: request-1"); + }), + }; + fakeThis.showError = vi.fn(); + let resolveLateResponse: ((response: AgentConnectionExtensionUiResponse) => void) | undefined; + fakeThis.resolveConnectionExtensionUiRequest = vi.fn( + () => + new Promise((resolve) => { + resolveLateResponse = resolve; + }), + ); + + const request: AgentConnectionExtensionUiRequest = { + id: "request-1", + method: "confirm", + payload: {}, + }; + const handling = prototype.handleConnectionExtensionUiRequest.call(fakeThis, request); + expect(fakeThis.activeConnectionExtensionUiRequests.has(request.id)).toBe(true); + + const cancelRequest = prototype.cancelConnectionExtensionUiRequest as ((requestId: string) => void) | undefined; + cancelRequest?.call(fakeThis, "unknown-request"); + expect(fakeThis.activeConnectionExtensionUiRequests.has(request.id)).toBe(true); + + cancelRequest?.call(fakeThis, request.id); + resolveLateResponse?.({ confirmed: true }); + await handling; + + expect(fakeThis.activeConnectionExtensionUiRequests.has(request.id)).toBe(false); + expect(fakeThis.agentConnection.respondToExtensionUiRequest).not.toHaveBeenCalled(); + expect(fakeThis.showError).not.toHaveBeenCalled(); + }); }); describe("InteractiveMode key handlers", () => { From c299890acbf4fa4e205ab9293dc308decc4e978d Mon Sep 17 00:00:00 2001 From: Tami Komssi Date: Tue, 1 Sep 2026 23:38:02 +0300 Subject: [PATCH 2/5] test(daemon): exclude cancellation after valid response --- .../test/daemon-extension-binding.test.ts | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/packages/coding-agent/test/daemon-extension-binding.test.ts b/packages/coding-agent/test/daemon-extension-binding.test.ts index 6edb19ea15..dec0f258a3 100644 --- a/packages/coding-agent/test/daemon-extension-binding.test.ts +++ b/packages/coding-agent/test/daemon-extension-binding.test.ts @@ -205,6 +205,57 @@ describe("daemon extension binding", () => { }); }); + it("does not broadcast cancellation after a valid daemon extension response", async () => { + let markDialogStarted: (() => void) | undefined; + const dialogStarted = new Promise((resolve) => { + markDialogStarted = resolve; + }); + let dialogResult: boolean | undefined; + const runtime = await createRuntimeForTest((pi) => { + pi.registerCommand("answered-dialog", { + description: "answered dialog", + handler: async (_args, ctx) => { + markDialogStarted?.(); + dialogResult = await ctx.ui.confirm("Continue?", "Confirm action"); + }, + }); + }, []); + + const outbound: DaemonOutbound[] = []; + const state: ActiveSessionState = { + activeSessionId: "active-answered-dialog", + runtime, + clients: new Set([{ supportsExtensionUi: true }] as never[]), + pendingAttaches: 0, + extensionUiRequests: new Map(), + eventGeneration: "generation-answered-dialog", + lastEventSequence: 0, + }; + await bindActiveSessionState(state, { + broadcast: (_state, message) => { + outbound.push(message); + }, + shutdown: () => {}, + }); + + const prompt = runtime.session.prompt("/answered-dialog"); + await dialogStarted; + const request = outbound.find( + (message): message is Extract => + message.type === "extension_ui_request", + ); + expect(request).toBeDefined(); + const pending = state.extensionUiRequests.get(request?.id ?? ""); + expect(pending).toBeDefined(); + + state.extensionUiRequests.delete(request?.id ?? ""); + pending?.resolve({ confirmed: true }); + await prompt; + + expect(dialogResult).toBe(true); + expect(outbound.filter((message) => message.type === "extension_ui_cancelled")).toEqual([]); + }); + it("keeps extension replacement callbacks daemon-side and rebinds before withSession", async () => { const phases: string[] = []; let oldSessionFile: string | undefined; From 20d6933e9baff32ec36b7da6e5249a9ced4b975d Mon Sep 17 00:00:00 2001 From: Tami Komssi Date: Wed, 2 Sep 2026 08:38:06 +0300 Subject: [PATCH 3/5] fix(daemon): route UI cancellation through supervisor --- .../src/modes/daemon/daemon-supervisor.ts | 22 +++-- .../test/daemon-supervisor-monitor.test.ts | 85 +++++++++++++++++++ 2 files changed, 102 insertions(+), 5 deletions(-) diff --git a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts index 43d550fe98..b19a861663 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-supervisor.ts @@ -2422,15 +2422,24 @@ export class DaemonSupervisor { if (!worker.client) { throw new Error("Session worker is not connected"); } - const supportsExtensionUi = [...this.clients].some( - (client) => client.attachedActiveSessionIds.has(activeSessionId) && client.supportsExtensionUi, + const attachedClients = [...this.clients].filter((client) => + client.attachedActiveSessionIds.has(activeSessionId), + ); + const supportsExtensionUi = attachedClients.some((client) => client.supportsExtensionUi); + const supportsExtensionUiCancellation = attachedClients.some( + (client) => client.supportsExtensionUi && client.capabilities.has("extension_ui_cancellation"), ); const response = await worker.client.requestWorker({ type: "worker_subscribe", activeSessionId, - capabilities: supportsExtensionUi - ? ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"] - : ["attach_snapshot", "event_sequence", "slim_attach", "chunked_snapshot"], + capabilities: [ + "attach_snapshot", + "event_sequence", + ...(supportsExtensionUi ? (["extension_ui"] as const) : []), + ...(supportsExtensionUiCancellation ? (["extension_ui_cancellation"] as const) : []), + "slim_attach", + "chunked_snapshot", + ], supportsExtensionUi, }); if (!response.success) { @@ -4257,6 +4266,9 @@ export class DaemonSupervisor { if (outboundType === "extension_ui_request" && !client.supportsExtensionUi) { continue; } + if (outboundType === "extension_ui_cancelled" && !client.capabilities.has("extension_ui_cancellation")) { + continue; + } if (client.snapshotActiveSessionIds?.has(activeSessionId)) { this.queueCatchup(client, activeSessionId, outboundType === "session_replaced" ? "replacement" : "resync"); continue; diff --git a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts index 9eb6f0002c..8c1e7ba178 100644 --- a/packages/coding-agent/test/daemon-supervisor-monitor.test.ts +++ b/packages/coding-agent/test/daemon-supervisor-monitor.test.ts @@ -3145,6 +3145,91 @@ describe("daemon worker supervisor monitoring", () => { }); }); + it("propagates extension UI cancellation capability to a session worker", async () => { + type SubscriptionWorker = { + client: { requestWorker: (command: unknown) => Promise<{ success: boolean }> }; + }; + const activeSessionId = "active-extension-ui-cancellation"; + const requestWorker = vi.fn(async () => ({ success: true })); + const worker: SubscriptionWorker = { client: { requestWorker } }; + const client = { + attachedActiveSessionIds: new Set([activeSessionId]), + supportsExtensionUi: true, + capabilities: new Set(["extension_ui", "extension_ui_cancellation"]), + }; + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + clients: new Set([client]), + }) as { + subscribeWorker(worker: SubscriptionWorker, activeSessionId: string): Promise; + }; + + await supervisor.subscribeWorker(worker, activeSessionId); + + expect(requestWorker).toHaveBeenCalledWith({ + type: "worker_subscribe", + activeSessionId, + capabilities: [ + "attach_snapshot", + "event_sequence", + "extension_ui", + "extension_ui_cancellation", + "slim_attach", + "chunked_snapshot", + ], + supportsExtensionUi: true, + }); + }); + + it("forwards extension UI cancellation only to capable attached clients", () => { + const activeSessionId = "active-extension-ui-cancellation"; + const capableWrites: string[] = []; + const legacyWrites: string[] = []; + const client = (writes: string[], cancellation: boolean) => + ({ + id: cancellation ? "capable" : "legacy", + socket: { + destroyed: false, + write: (value: string | Buffer) => { + writes.push(value.toString()); + return true; + }, + }, + attachedActiveSessionIds: new Set([activeSessionId]), + catchupActiveSessionIds: new Set(), + backpressured: false, + supportsExtensionUi: true, + capabilities: new Set(cancellation ? ["extension_ui", "extension_ui_cancellation"] : ["extension_ui"]), + }) as unknown as DaemonSocketClient; + const worker = { + snapshotCache: new Map(), + transcriptCaches: new Map(), + incomingTranscriptActiveSessionIds: new Set(), + duplicateIncomingTranscriptChunkIndexes: new Map(), + snapshotTransferFrames: new Map(), + }; + const supervisor = Object.assign(Object.create(DaemonSupervisor.prototype), { + clients: new Set([client(capableWrites, true), client(legacyWrites, false)]), + streamReconstructor: { observe: vi.fn() }, + invalidateWorkerSnapshot: vi.fn(), + }) as { + handleWorkerFrame(residentWorker: typeof worker, frame: PrivateFrame): void; + }; + const outbound = { + type: "extension_ui_cancelled" as const, + activeSessionId, + id: "request-1", + }; + + supervisor.handleWorkerFrame(worker, { + header: { kind: "outbound", outboundType: outbound.type, activeSessionId }, + payload: Buffer.from(`${JSON.stringify(outbound)}\n`), + }); + + expect(capableWrites).toHaveLength(1); + expect(capableWrites[0]).toContain('"id":"request-1"'); + expect(legacyWrites).toEqual([]); + }); + it("does not retain an attachment when snapshot loading fails", async () => { type AttachClient = { id: string; From 13262a60c26b43103a696c1af2c0137ff762e751 Mon Sep 17 00:00:00 2001 From: Tami Komssi Date: Wed, 2 Sep 2026 08:44:44 +0300 Subject: [PATCH 4/5] fix(tui): dismiss cancelled daemon extension dialogs --- .../src/modes/interactive/interactive-mode.ts | 36 ++++++++++++++++--- .../test/interactive-mode-status.test.ts | 6 +++- 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 1f5f2b7131..8f26a5a7c9 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -3906,18 +3906,35 @@ export class InteractiveMode { /** * Show a multi-line editor for extensions (with Ctrl+G support). */ - private showExtensionEditor(title: string, prefill?: string): Promise { + private showExtensionEditor( + title: string, + prefill?: string, + opts?: ExtensionUIDialogOptions, + ): Promise { return new Promise((resolve) => { + if (opts?.signal?.aborted) { + resolve(undefined); + return; + } + + const onAbort = () => { + this.hideExtensionEditor(); + resolve(undefined); + }; + opts?.signal?.addEventListener("abort", onAbort, { once: true }); + this.extensionEditor = new ExtensionEditorComponent( this.ui, this.keybindings, title, prefill, (value) => { + opts?.signal?.removeEventListener("abort", onAbort); this.hideExtensionEditor(); resolve(value); }, () => { + opts?.signal?.removeEventListener("abort", onAbort); this.hideExtensionEditor(); resolve(undefined); }, @@ -5162,14 +5179,21 @@ export class InteractiveMode { try { if (expectsResponse) { + const dialogController = new AbortController(); let cancelLocal: (response: AgentConnectionExtensionUiResponse) => void = () => {}; const cancelled = new Promise((resolve) => { cancelLocal = resolve; }); this.activeConnectionExtensionUiRequests.set(request.id, { - cancelLocal: () => cancelLocal({ cancelled: true }), + cancelLocal: () => { + dialogController.abort(); + cancelLocal({ cancelled: true }); + }, }); - response = await Promise.race([this.resolveConnectionExtensionUiRequest(request), cancelled]); + response = await Promise.race([ + this.resolveConnectionExtensionUiRequest(request, dialogController.signal), + cancelled, + ]); } else { response = await this.resolveConnectionExtensionUiRequest(request); } @@ -5229,6 +5253,7 @@ export class InteractiveMode { private async resolveConnectionExtensionUiRequest( request: AgentConnectionExtensionUiRequest, + signal?: AbortSignal, ): Promise { const { payload } = request; switch (request.method) { @@ -5240,6 +5265,7 @@ export class InteractiveMode { } const value = await this.showExtensionSelector(title, options, { timeout: getPayloadNumber(payload, "timeout"), + signal, }); return value === undefined ? { cancelled: true } : { value }; } @@ -5251,6 +5277,7 @@ export class InteractiveMode { } const confirmed = await this.showExtensionConfirm(title, message, { timeout: getPayloadNumber(payload, "timeout"), + signal, }); return { confirmed }; } @@ -5261,6 +5288,7 @@ export class InteractiveMode { } const value = await this.showExtensionInput(title, getPayloadString(payload, "placeholder"), { timeout: getPayloadNumber(payload, "timeout"), + signal, }); return value === undefined ? { cancelled: true } : { value }; } @@ -5269,7 +5297,7 @@ export class InteractiveMode { if (!title) { return { cancelled: true }; } - const value = await this.showExtensionEditor(title, getPayloadString(payload, "prefill")); + const value = await this.showExtensionEditor(title, getPayloadString(payload, "prefill"), { signal }); return value === undefined ? { cancelled: true } : { value }; } case "notify": { diff --git a/packages/coding-agent/test/interactive-mode-status.test.ts b/packages/coding-agent/test/interactive-mode-status.test.ts index 96d9e76ce4..2806a9731c 100644 --- a/packages/coding-agent/test/interactive-mode-status.test.ts +++ b/packages/coding-agent/test/interactive-mode-status.test.ts @@ -1629,6 +1629,7 @@ describe("InteractiveMode connection extension UI", () => { type ConnectionExtensionUiHandlerHarness = ConnectionExtensionUiCancelHarness & { resolveConnectionExtensionUiRequest( request: AgentConnectionExtensionUiRequest, + signal?: AbortSignal, ): Promise; handleConnectionExtensionUiRequest(request: AgentConnectionExtensionUiRequest): Promise; }; @@ -1697,9 +1698,11 @@ describe("InteractiveMode connection extension UI", () => { }; fakeThis.showError = vi.fn(); let resolveLateResponse: ((response: AgentConnectionExtensionUiResponse) => void) | undefined; + let dialogSignal: AbortSignal | undefined; fakeThis.resolveConnectionExtensionUiRequest = vi.fn( - () => + (_request, signal) => new Promise((resolve) => { + dialogSignal = signal; resolveLateResponse = resolve; }), ); @@ -1717,6 +1720,7 @@ describe("InteractiveMode connection extension UI", () => { expect(fakeThis.activeConnectionExtensionUiRequests.has(request.id)).toBe(true); cancelRequest?.call(fakeThis, request.id); + expect(dialogSignal?.aborted).toBe(true); resolveLateResponse?.({ confirmed: true }); await handling; From 62659fe86cdf6751f684d924b7bc96e6225b3871 Mon Sep 17 00:00:00 2001 From: Tami Komssi Date: Wed, 2 Sep 2026 08:57:52 +0300 Subject: [PATCH 5/5] test(daemon): verify extension UI response ordering --- .../coding-agent/test/daemon-mode.test.ts | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/packages/coding-agent/test/daemon-mode.test.ts b/packages/coding-agent/test/daemon-mode.test.ts index 06214039c9..784785d033 100644 --- a/packages/coding-agent/test/daemon-mode.test.ts +++ b/packages/coding-agent/test/daemon-mode.test.ts @@ -40,6 +40,7 @@ import { type DaemonAttachResult, type DaemonCommand, type DaemonOutbound, + type DaemonResponse, failure, } from "../src/modes/daemon/daemon-protocol.js"; import type { SessionSummary } from "../src/modes/daemon/daemon-session-list.js"; @@ -184,6 +185,69 @@ describe("daemon mode helpers", () => { expect(resolve).toHaveBeenCalledWith({ cancelled: true }); }); + it("removes a pending extension UI request before resolving its valid response", async () => { + const daemon = new AgentDaemon("/tmp/prime-agent-test.sock", { + defaultSessionConfig: { agentDir: "/tmp/prime-agent-test-agent", cwd: "/tmp" }, + createRuntime: async () => { + throw new Error("unexpected runtime creation"); + }, + }); + const client = makeClient("client-1", "active"); + const resolve = vi.fn(); + const state = { + ...makeState("active"), + clients: new Set([client]), + extensionUiRequests: new Map([["request-1", { resolve }]]), + }; + const internals = daemon as unknown as { + sessions: Map; + handleCommand(client: DaemonSocketClient, command: DaemonCommand): Promise; + }; + internals.sessions.set(state.activeSessionId, state); + + const response = await internals.handleCommand(client, { + id: "response-1", + type: "extension_ui_response", + activeSessionId: state.activeSessionId, + requestId: "request-1", + response: { confirmed: true }, + }); + + expect(response).toMatchObject({ success: true, command: "extension_ui_response" }); + expect(state.extensionUiRequests.has("request-1")).toBe(false); + expect(resolve).toHaveBeenCalledWith({ confirmed: true }); + }); + + it("rejects a truly unknown extension UI response", async () => { + const daemon = new AgentDaemon("/tmp/prime-agent-test.sock", { + defaultSessionConfig: { agentDir: "/tmp/prime-agent-test-agent", cwd: "/tmp" }, + createRuntime: async () => { + throw new Error("unexpected runtime creation"); + }, + }); + const client = makeClient("client-1", "active"); + const state = { + ...makeState("active"), + clients: new Set([client]), + extensionUiRequests: new Map(), + }; + const internals = daemon as unknown as { + sessions: Map; + handleCommand(client: DaemonSocketClient, command: DaemonCommand): Promise; + }; + internals.sessions.set(state.activeSessionId, state); + + await expect( + internals.handleCommand(client, { + id: "response-unknown", + type: "extension_ui_response", + activeSessionId: state.activeSessionId, + requestId: "unknown-request", + response: { confirmed: true }, + }), + ).rejects.toThrow("Unknown extension UI request: unknown-request"); + }); + it("acknowledges agent messages after target prompt preflight succeeds", async () => { const daemon = new AgentDaemon("/tmp/prime-agent-test.sock", { defaultSessionConfig: { agentDir: "/tmp/prime-agent-test-agent", cwd: "/tmp" },