Propagate cancellation and order shutdown usage persistence

Author: GlitterKill

src/cli/commands/serve.ts

@@ -110,6 +110,7 @@ export function registerServeFinalCleanups(
   shutdownMgr: Pick<ShutdownManager, "addCleanup">,
   cleanup: {
     drainWork: () => Promise<void>;
+    persistUsage: () => Promise<void>;
     stopObservability: () => Promise<void>;
     closeDatabase: () => Promise<void>;
   },
@@ -120,6 +121,9 @@ export function registerServeFinalCleanups(
     workDrained = true;
   };
   shutdownMgr.addCleanup("workDrain", drainWork);
+  shutdownMgr.addCleanup("persistUsage", async () => {
+    if (workDrained) await cleanup.persistUsage();
+  });
   shutdownMgr.addCleanup("observability", async () => {
     if (workDrained) await cleanup.stopObservability();
   });
@@ -289,23 +293,6 @@ export async function serveCommand(options: ServeOptions): Promise<void> {
     dashboardHandle?.close(),
   );
   shutdownMgr.addCleanup("httpServer", () => httpHandle?.close());
-  shutdownMgr.addCleanup("persistUsage", async () => {
-    try {
-      if (
-        graphDbAvailable &&
-        startupReadiness.isWriteReady() &&
-        tokenAccumulator.hasUsage
-      ) {
-        await persistUsageSnapshot(tokenAccumulator.getSnapshot());
-      }
-    } catch (err) {
-      // Non-critical — don't block shutdown.
-      writeServeStderrLine(
-        "[sdl-mcp] Failed to persist usage snapshot: " +
-          (err instanceof Error ? err.message : String(err)),
-      );
-    }
-  });
   shutdownMgr.addCleanup("watchers", async () => {
     for (const watcher of watchers) {
       try {
@@ -320,6 +307,23 @@ export async function serveCommand(options: ServeOptions): Promise<void> {
   shutdownMgr.addCleanup("graphIntegrityVerifier", stopGraphIntegrityVerifier);
   registerServeFinalCleanups(shutdownMgr, {
     drainWork: drainLadybugWork,
+    persistUsage: async () => {
+      try {
+        if (
+          graphDbAvailable &&
+          startupReadiness.isWriteReady() &&
+          tokenAccumulator.hasUsage
+        ) {
+          await persistUsageSnapshot(tokenAccumulator.getSnapshot());
+        }
+      } catch (err) {
+        // Non-critical — don't block shutdown.
+        writeServeStderrLine(
+          "[sdl-mcp] Failed to persist usage snapshot: " +
+            (err instanceof Error ? err.message : String(err)),
+        );
+      }
+    },
     stopObservability,
     closeDatabase: closeLadybugDb,
   });

src/mcp/dispatch-limiter.ts

@@ -171,6 +171,7 @@ export async function runToolDispatch<T>(
   fn: () => Promise<T>,
   timeoutMs?: number,
   label = "tool-dispatch",
+  signal?: AbortSignal,
 ): Promise<T> {
   const deferredWork = label.startsWith("derived-refresh:")
     ? undefined
@@ -195,6 +196,7 @@ export async function runToolDispatch<T>(
           }
         }),
       queueTimeoutMs,
+      signal,
     );
   } catch (error) {
     if (error instanceof ConcurrencyQueueTimeoutError) {

src/mcp/tools/code.ts

@@ -94,6 +94,32 @@ type ResolvedGetHotPathRequest = Omit<
   "symbolId" | "symbolRef"
 > & { symbolId: string };
 
+function awaitWithSignal<T>(
+  promise: Promise<T>,
+  signal?: AbortSignal,
+): Promise<T> {
+  if (!signal) return promise;
+  if (signal.aborted) return Promise.reject(signal.reason);
+
+  return new Promise<T>((resolve, reject) => {
+    const onAbort = (): void => {
+      signal.removeEventListener("abort", onAbort);
+      reject(signal.reason);
+    };
+    signal.addEventListener("abort", onAbort, { once: true });
+    void promise.then(
+      (value) => {
+        signal.removeEventListener("abort", onAbort);
+        resolve(value);
+      },
+      (error: unknown) => {
+        signal.removeEventListener("abort", onAbort);
+        reject(error);
+      },
+    );
+  });
+}
+
 function buildSessionRef(key: string, etag?: string): { key: string; etag?: string } {
   const ref: { key: string; etag?: string } = { key };
   if (etag !== undefined) ref.etag = etag;
@@ -561,11 +587,21 @@ export async function handleCodeNeedWindow(
       request.repoId,
     );
     if (latestVersion) {
-      const { slice } = await buildSlice({
-        repoId: request.repoId,
-        versionId: latestVersion.versionId,
-        ...request.sliceContext,
-      });
+      const sliceSignal = context?.signal
+        ? AbortSignal.any([context.signal, AbortSignal.timeout(30_000)])
+        : undefined;
+      sliceSignal?.throwIfAborted();
+      // LadybugDB reads are not uniformly abort-aware; release the MCP
+      // dispatch slot on cancellation while also passing the signal down.
+      const { slice } = await awaitWithSignal(
+        buildSlice({
+          repoId: request.repoId,
+          versionId: latestVersion.versionId,
+          ...request.sliceContext,
+          signal: sliceSignal,
+        }),
+        sliceSignal,
+      );
       sliceContext = slice;
     }
   }

src/server.ts

@@ -1298,7 +1298,12 @@ export class MCPServer {
             const runDispatch = () =>
               shouldBypassToolDispatch(toolName, parsedArgs)
                 ? dispatchTool()
-                : runToolDispatch(dispatchTool, undefined, toolName);
+                : runToolDispatch(
+                    dispatchTool,
+                    undefined,
+                    toolName,
+                    toolContext.signal,
+                  );
             // Refresh admission must happen before the outer dispatch lease.
             // This also covers workflows, whose refresh step executes inside
             // the workflow's single outer lease rather than acquiring its own.

tests/unit/dispatch-limiter.test.ts

@@ -1,4 +1,6 @@
 import assert from "node:assert";
+import { readFileSync } from "node:fs";
+import { join } from "node:path";
 import { afterEach, beforeEach, describe, it } from "node:test";
 
 import {
@@ -180,6 +182,56 @@ describe("tool dispatch limiter", () => {
     await first;
   });
 
+  it("removes queued dispatch work when the request is cancelled", async () => {
+    configureToolDispatchLimiter({ maxConcurrency: 1, queueTimeoutMs: 1_000 });
+
+    let release: (() => void) | undefined;
+    const blocker = new Promise<void>((resolve) => {
+      release = resolve;
+    });
+    const first = runToolDispatch(async () => blocker, undefined, "first");
+    const controller = new AbortController();
+    const reason = new Error("client disconnected");
+    const queued = runToolDispatch(
+      async () => "second",
+      undefined,
+      "second",
+      controller.signal,
+    );
+
+    try {
+      await delay(20);
+      assert.strictEqual(getToolDispatchStats().queued, 1);
+
+      controller.abort(reason);
+      const outcome = await Promise.race([
+        queued.then(
+          () => ({ kind: "resolved" as const }),
+          (error: unknown) => ({ kind: "rejected" as const, error }),
+        ),
+        delay(100).then(() => ({ kind: "pending" as const })),
+      ]);
+
+      assert.strictEqual(outcome.kind, "rejected");
+      if (outcome.kind !== "rejected") return;
+      assert.strictEqual(outcome.error, reason);
+      assert.strictEqual(getToolDispatchStats().queued, 0);
+    } finally {
+      release?.();
+      await first;
+      await queued.catch(() => undefined);
+    }
+  });
+
+  it("forwards MCP request cancellation into the shared dispatch queue", () => {
+    const source = readFileSync(join(process.cwd(), "src", "server.ts"), "utf8");
+
+    assert.match(
+      source,
+      /runToolDispatch\(\s*dispatchTool,\s*undefined,\s*toolName,\s*toolContext\.signal,\s*\)/,
+    );
+  });
+
   it("lets derived-refresh deferred work finish before timing foreground dispatch out", async () => {
     configureToolDispatchLimiter({ maxConcurrency: 1, queueTimeoutMs: 30 });
 

tests/unit/http-shutdown.test.ts

@@ -77,36 +77,54 @@ describe("HTTP shutdown wiring", () => {
     }
   });
 
-  it("registers HTTP server cleanup before final DB cleanup", () => {
+  it("stops HTTP producers before draining work and persisting usage", () => {
     const source = readFileSync(
       join(process.cwd(), "src", "cli", "commands", "serve.ts"),
       "utf8",
     );
-    const httpCleanupIndex = source.indexOf(
+    const serveSource = source.slice(
+      source.indexOf("export async function serveCommand"),
+    );
+    const httpCleanupIndex = serveSource.indexOf(
       'shutdownMgr.addCleanup("httpServer"',
     );
-    const persistUsageIndex = source.indexOf(
+    const earlyPersistUsageIndex = serveSource.indexOf(
       'shutdownMgr.addCleanup("persistUsage"',
     );
-    const dbCleanupIndex = source.lastIndexOf(
+    const watcherCleanupIndex = serveSource.indexOf(
+      'shutdownMgr.addCleanup("watchers"',
+    );
+    const verifierCleanupIndex = serveSource.indexOf(
+      'shutdownMgr.addCleanup("graphIntegrityVerifier"',
+    );
+    const finalCleanupIndex = serveSource.indexOf(
       "registerServeFinalCleanups(shutdownMgr",
     );
-    const loggerCleanupIndex = source.indexOf(
+    const loggerCleanupIndex = serveSource.indexOf(
       'shutdownMgr.addCleanup("logger"',
     );
 
     assert.ok(httpCleanupIndex >= 0, "HTTP cleanup should be registered");
-    assert.ok(
-      persistUsageIndex >= 0,
-      "usage persistence cleanup should be registered",
+    assert.equal(
+      earlyPersistUsageIndex,
+      -1,
+      "usage persistence must not run before producer shutdown and work drain",
     );
-    assert.ok(dbCleanupIndex >= 0, "DB cleanup should be registered");
+    assert.ok(watcherCleanupIndex >= 0, "watcher cleanup should be registered");
+    assert.ok(verifierCleanupIndex >= 0, "verifier cleanup should be registered");
+    assert.ok(finalCleanupIndex >= 0, "final cleanup should be registered");
     assert.ok(loggerCleanupIndex >= 0, "logger cleanup should be registered");
     assert.ok(
-      httpCleanupIndex < persistUsageIndex &&
-        persistUsageIndex < dbCleanupIndex &&
-        dbCleanupIndex < loggerCleanupIndex,
-      "HTTP transport cleanup must run before usage persistence and final DB/logger cleanup",
+      httpCleanupIndex < watcherCleanupIndex &&
+        watcherCleanupIndex < verifierCleanupIndex &&
+        verifierCleanupIndex < finalCleanupIndex &&
+        finalCleanupIndex < loggerCleanupIndex,
+      "HTTP transport and producers must stop before final drain/usage/DB cleanup",
+    );
+    assert.match(
+      serveSource,
+      /registerServeFinalCleanups\(shutdownMgr,\s*\{[\s\S]*?persistUsage:\s*async\s*\(\)\s*=>/,
+      "serve cleanup must persist usage through the post-drain final cleanup path",
     );
   });
 

tests/unit/mcp-code-need-window-policy.test.ts

@@ -120,6 +120,15 @@ describe("code.needWindow policy remediation", () => {
       createdAt: now,
     });
 
+    await ladybugDb.createVersion(conn, {
+      versionId: "version-test",
+      repoId: "repo-test",
+      createdAt: now,
+      reason: "test",
+      prevVersionHash: null,
+      versionHash: "hash-version-test",
+    });
+
     await ladybugDb.upsertFile(conn, {
       fileId: "file-demo",
       repoId: "repo-test",
@@ -288,6 +297,46 @@ describe("code.needWindow policy remediation", () => {
     assert.equal(JSON.stringify(first), JSON.stringify(second));
   });
 
+  it("passes request cancellation into slice-backed code access", async () => {
+    const originalAny = AbortSignal.any;
+    const calls: AbortSignal[][] = [];
+    AbortSignal.any = ((signals: AbortSignal[]) => {
+      calls.push([...signals]);
+      return originalAny(signals);
+    }) as typeof AbortSignal.any;
+
+    const controller = new AbortController();
+    const reason = new Error("client disconnected");
+    try {
+      const response = handleCodeNeedWindow(
+        {
+          repoId: "repo-test",
+          symbolId: "sym-demo",
+          reason: "inspect important flag handling",
+          expectedLines: 20,
+          maxTokens: 120,
+          identifiersToFind: ["importantFlag"],
+          sliceContext: {
+            taskText: "inspect demoWindow",
+            entrySymbols: ["sym-demo"],
+          },
+        },
+        {
+          sendNotification: async () => {},
+          signal: controller.signal,
+        },
+      );
+      queueMicrotask(() => controller.abort(reason));
+
+      await assert.rejects(response, (error: unknown) => error === reason);
+      assert.equal(calls.length, 1);
+      assert.strictEqual(calls[0]?.[0], controller.signal);
+      assert.equal(calls[0]?.length, 2);
+    } finally {
+      AbortSignal.any = originalAny;
+    }
+  });
+
   it("resolves stringified symbolRef targets for raw code windows", async () => {
     const response = await handleCodeNeedWindow({
       repoId: "repo-test",