From e5cc1e6e546534bc07421a41b5d933829e889bd8 Mon Sep 17 00:00:00 2001 From: George Ng Date: Fri, 18 Sep 2026 16:52:11 -0700 Subject: [PATCH] Share dispatcher request cancellation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../dispatcher/dispatcher/src/dispatcher.ts | 13 +-- .../dispatcher/src/queue/requestQueue.ts | 26 ++++++ .../src/structuredAction/execution.ts | 7 +- .../test/queue/requestQueue.spec.ts | 82 +++++++++++++++++++ 4 files changed, 112 insertions(+), 16 deletions(-) diff --git a/ts/packages/dispatcher/dispatcher/src/dispatcher.ts b/ts/packages/dispatcher/dispatcher/src/dispatcher.ts index 5b980e801c..f3d8eef09c 100644 --- a/ts/packages/dispatcher/dispatcher/src/dispatcher.ts +++ b/ts/packages/dispatcher/dispatcher/src/dispatcher.ts @@ -33,6 +33,7 @@ import { DispatcherOptions, initializeCommandHandlerContext, } from "./context/commandHandlerContext.js"; +import { cancelQueuedRequest } from "./queue/requestQueue.js"; import { randomUUID } from "node:crypto"; import { context as otelContext } from "@opentelemetry/api"; import { getAgentSchemas } from "./context/system/describe/agentSchemaInfo.js"; @@ -413,17 +414,7 @@ export function createDispatcherFromContext( structuredExecution.continueAction(request), cancelAction: (request) => structuredExecution.cancelAction(request), async cancelCommand(requestId: string): Promise { - const kind = context.requestQueue.classifyCancel(requestId, "user"); - if (kind === "queued") { - return { kind: "cancelled_queued", requestId }; - } - if (kind === "running") { - context.requestQueue.cancelRunning(requestId, "user"); - const controller = context.activeRequests.get(requestId); - controller?.abort(); - return { kind: "cancelled_running", requestId }; - } - return { kind: "not_found", requestId }; + return cancelQueuedRequest(context, requestId, "user"); }, async promoteCommand(requestId: string): Promise { return context.requestQueue.promote(requestId); diff --git a/ts/packages/dispatcher/dispatcher/src/queue/requestQueue.ts b/ts/packages/dispatcher/dispatcher/src/queue/requestQueue.ts index b373d1d0e9..938c6dcd23 100644 --- a/ts/packages/dispatcher/dispatcher/src/queue/requestQueue.ts +++ b/ts/packages/dispatcher/dispatcher/src/queue/requestQueue.ts @@ -750,5 +750,31 @@ export function entryCompletion( return entry.completion; } +type RequestCancellationContext = { + requestQueue: RequestQueue; + activeRequests: ReadonlyMap; +}; + +/** + * Cancel a queue-managed request by its server request id. + * Running cancellation is published before its controller is aborted. + */ +export function cancelQueuedRequest( + context: RequestCancellationContext, + requestId: string, + reason: QueueCancelReason, +): CancelResult { + const kind = context.requestQueue.classifyCancel(requestId, reason); + if (kind === "queued") { + return { kind: "cancelled_queued", requestId }; + } + if (kind === "running") { + context.requestQueue.cancelRunning(requestId, reason); + context.activeRequests.get(requestId)?.abort(); + return { kind: "cancelled_running", requestId }; + } + return { kind: "not_found", requestId }; +} + /** Re-export for consumers that build CancelResult without going via the queue. */ export type { CancelResult }; diff --git a/ts/packages/dispatcher/dispatcher/src/structuredAction/execution.ts b/ts/packages/dispatcher/dispatcher/src/structuredAction/execution.ts index e0a395f8d6..7a055deb71 100644 --- a/ts/packages/dispatcher/dispatcher/src/structuredAction/execution.ts +++ b/ts/packages/dispatcher/dispatcher/src/structuredAction/execution.ts @@ -29,6 +29,7 @@ import { import type { CommandHandlerContext } from "../context/commandHandlerContext.js"; import type { executeActions } from "../execute/actionHandlers.js"; import type { getActionContext } from "../execute/actionContext.js"; +import { cancelQueuedRequest } from "../queue/requestQueue.js"; import { getAppAgentName } from "../translation/agentTranslators.js"; import type { StructuredActionDiscovery } from "./discovery.js"; import { @@ -675,11 +676,7 @@ class Operation implements StructuredExecutionHooks { this.possibleEffects ? "execution_uncertain" : "cancelled", ), ); - const queue = this.context.requestQueue; - if (!queue.cancelQueued(this.id, "user")) { - queue.cancelRunning(this.id, "user"); - this.context.activeRequests.get(this.id)?.abort(); - } + cancelQueuedRequest(this.context, this.id, "user"); } retire(): void { diff --git a/ts/packages/dispatcher/dispatcher/test/queue/requestQueue.spec.ts b/ts/packages/dispatcher/dispatcher/test/queue/requestQueue.spec.ts index 02ad247117..4f0f6d48f8 100644 --- a/ts/packages/dispatcher/dispatcher/test/queue/requestQueue.spec.ts +++ b/ts/packages/dispatcher/dispatcher/test/queue/requestQueue.spec.ts @@ -14,6 +14,7 @@ import { import { context, createContextKey, type Context } from "@opentelemetry/api"; import { + cancelQueuedRequest, RequestQueue, MAX_QUEUE_DEPTH, QueueBroadcaster, @@ -244,6 +245,87 @@ describe("RequestQueue", () => { expect(queue.cancelQueued("does-not-exist", "user")).toBe(false); }); + it("cancelQueuedRequest maps queued and unknown request ids", async () => { + const dispatcher = new ControllableDispatcher(); + const { queue } = makeQueue(dispatcher); + const activeRequests = new Map(); + const running = queue.submit({ + text: "running", + originatorConnectionId: "c1", + }); + const queued = queue.submit({ + text: "queued", + originatorConnectionId: "c1", + }); + await flush(); + + expect( + cancelQueuedRequest( + { requestQueue: queue, activeRequests }, + queued.requestId, + "user", + ), + ).toEqual({ + kind: "cancelled_queued", + requestId: queued.requestId, + }); + expect( + cancelQueuedRequest( + { requestQueue: queue, activeRequests }, + queued.requestId, + "user", + ), + ).toEqual({ kind: "not_found", requestId: queued.requestId }); + + dispatcher.calls[0].resolve({}); + await running.completion; + }); + + it("cancelQueuedRequest broadcasts before abort and is idempotent", async () => { + const dispatcher = new ControllableDispatcher(); + const { queue, events } = makeQueue(dispatcher); + const running = queue.submit({ + text: "running", + originatorConnectionId: "c1", + }); + await flush(); + const controller = new AbortController(); + const activeRequests = new Map([[running.requestId, controller]]); + let cancelEventsAtAbort = 0; + controller.signal.addEventListener("abort", () => { + cancelEventsAtAbort = events.filter( + (event) => + event.type === "cancelled" && + event.requestId === running.requestId, + ).length; + }); + + const context = { requestQueue: queue, activeRequests }; + expect(cancelQueuedRequest(context, running.requestId, "user")).toEqual( + { + kind: "cancelled_running", + requestId: running.requestId, + }, + ); + expect(cancelEventsAtAbort).toBe(1); + expect(cancelQueuedRequest(context, running.requestId, "user")).toEqual( + { + kind: "cancelled_running", + requestId: running.requestId, + }, + ); + expect( + events.filter( + (event) => + event.type === "cancelled" && + event.requestId === running.requestId, + ), + ).toHaveLength(1); + + dispatcher.calls[0].resolve({}); + await running.completion; + }); + it("promote moves a queued entry to the front so it runs next", async () => { const dispatcher = new ControllableDispatcher(); const { queue, events } = makeQueue(dispatcher);