@0xmaxma/claude-gateway 1.7.1 → 1.7.2

package/mcp/server.ts

@@ -118,19 +118,46 @@ mcp.setRequestHandler(CallToolRequestSchema, async (req, extra) => {
   // it by requestId) — the CLI sends that when a user Stop/Ctrl-C interrupts an
   // in-flight tool call. Threaded through so a long-running module (image
   // generation's poll loop) can react instead of running to its full timeout.
-  return mod.handleTool(toolName, args, extra.signal);
+  // Also combine with shutdownController.signal: the CLI sending that
+  // notification is only a "SHOULD" in the MCP spec, not a "MUST", and in
+  // practice a Stop that lets the CLI process exit cleanly (rather than
+  // staying alive to keep chatting) closes this server's stdin without ever
+  // sending notifications/cancelled — extra.signal would then never fire, and
+  // an in-flight image generation would poll to its full timeout instead of
+  // cancelling. stdin closing is a reliable, protocol-independent signal that
+  // the turn is over either way, so it backstops the notification.
+  const combinedSignal = AbortSignal.any([extra.signal, shutdownController.signal]);
+  return mod.handleTool(toolName, args, combinedSignal);
 });
 
 // Connect MCP transport
 await mcp.connect(new StdioServerTransport());
 
-// Graceful shutdown
+// Graceful shutdown — used by stdin-close, SIGINT, and SIGTERM paths.
+// Awaits any in-flight image cancel (E3) before exiting so the provider
+// actually stops generating when the user presses Stop, rather than the
+// process dying mid-request and the cancel call never reaching the server.
+const imageModuleRef = modules.find((m) => m.id === 'image') as { drainCancel?: () => Promise<void> } | undefined;
 let shuttingDown = false;
 function shutdown(): void {
   if (shuttingDown) return;
   shuttingDown = true;
   shutdownController.abort();
-  setTimeout(() => process.exit(0), 2000);
+  // Hard cap: never let shutdown block on the drain. cancelJob() is bounded by
+  // its own AbortSignal.timeout (up to 30s), which is far longer than a
+  // supervisor's SIGTERM→SIGKILL grace period. Force-exit after a short window
+  // so a slow/hung provider cancel can't keep the process alive and risk SIGKILL
+  // mid-drain. Whichever fires first wins; unref() so the timer itself never
+  // holds the loop open.
+  const forceExit = setTimeout(() => process.exit(0), 2000);
+  forceExit.unref();
+  // Give the event loop one tick so the poll loop's sleep() onAbort listener
+  // fires and cancelledResult() sets activeCancelPromise before we try to drain it.
+  setImmediate(async () => {
+    try { await imageModuleRef?.drainCancel?.(); } catch { /* non-fatal */ }
+    clearTimeout(forceExit);
+    process.exit(0);
+  });
 }
 process.stdin.on('end', shutdown);
 process.stdin.on('close', shutdown);

package/mcp/tools/image/module.ts

@@ -382,9 +382,12 @@ export class ImageModule implements ToolModule {
       if (signal?.aborted) return this.cancelledResult(taskId);
       await sleep(DEFAULT_POLL_INTERVAL_MS, signal);
       if (signal?.aborted) return this.cancelledResult(taskId);
-      const polled = await this.fetchJob(taskId);
+      const polled = await this.fetchJob(taskId, signal);
       if (polled.__transportError) {
-        // transient transport error — keep polling until deadline
+        // transient transport error (or abort signal fired mid-fetch) — check
+        // signal before continuing so a Stop-triggered abort isn't swallowed
+        // as a retryable error and the cancel fires on the next loop iteration.
+        if (signal?.aborted) return this.cancelledResult(taskId);
         continue;
       }
       if (polled.httpError) return this.mapHttpError(polled.httpError.status, polled.httpError.body);
@@ -539,8 +542,9 @@ export class ImageModule implements ToolModule {
   /**
    * Best-effort E3 cancel — fires and never throws, so a failed cancel call can
    * never break the caller's own (already-cancelled) tool response. The api-side
-   * cancel is itself idempotent/no-op-safe on an already-terminal or unknown task,
-   * so no need to track whether this actually reached a still-running job.
+   * cancel is itself idempotent/no-op-safe on an already-terminal or unknown task.
+   * cancelledResult() tracks the returned promise (activeCancelPromise) so a
+   * process-exiting shutdown can await it via drainCancel() instead of racing it.
    */
   private async cancelJob(taskId: string): Promise<void> {
     try {
@@ -554,9 +558,23 @@ export class ImageModule implements ToolModule {
     }
   }
 
+  // Tracks the in-flight cancel call so drainCancel() can await it before the
+  // server process exits — a Stop that lets the CLI process die cleanly closes
+  // this server's stdin, which fires process.exit() shortly after; without
+  // this, that exit can race the fire-and-forget cancelJob() fetch and kill it
+  // before the request ever reaches the provider.
+  private activeCancelPromise: Promise<void> | null = null;
+
+  /** Wait for any in-flight E3 cancel to complete. Called by the shutdown handler. */
+  async drainCancel(): Promise<void> {
+    if (this.activeCancelPromise) await this.activeCancelPromise;
+  }
+
   /** Tool result for a Stop-triggered cancellation, firing the E3 cancel first. */
   private cancelledResult(taskId: string): McpToolResult {
-    void this.cancelJob(taskId);
+    this.activeCancelPromise = this.cancelJob(taskId).finally(() => {
+      this.activeCancelPromise = null;
+    });
     return {
       content: [{
         type: 'text',
@@ -567,13 +585,19 @@ export class ImageModule implements ToolModule {
   }
 
   /** Fetch a job (E2), classifying transport vs HTTP errors so the poller can retry transient ones. */
-  private async fetchJob(taskId: string): Promise<{ job?: JobResponse; httpError?: { status: number; body: string }; __transportError?: unknown }> {
+  private async fetchJob(taskId: string, signal?: AbortSignal): Promise<{ job?: JobResponse; httpError?: { status: number; body: string }; __transportError?: unknown }> {
     let res: Response;
     try {
+      // Combine with the caller's cancel signal (not just the request timeout) so
+      // a Stop that lands mid-fetch aborts the request immediately instead of
+      // waiting out REQUEST_TIMEOUT_MS before the poll loop even gets to check.
+      const fetchSignal = signal
+        ? AbortSignal.any([AbortSignal.timeout(REQUEST_TIMEOUT_MS), signal])
+        : AbortSignal.timeout(REQUEST_TIMEOUT_MS);
       res = await fetch(`${this.baseUrl()}/v1/images/jobs/${encodeURIComponent(taskId)}`, {
         method: 'GET',
         headers: this.headers(),
-        signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
+        signal: fetchSignal,
       });
     } catch (err) {
       return { __transportError: err };
@@ -764,7 +788,7 @@ export class ImageModule implements ToolModule {
 // signal.aborted itself right after, so this only needs to shorten the wait.
 // The abort listener is removed when the timer fires normally: sleep() is called
 // once per poll iteration against the SAME long-lived signal (up to ~75 times for
-// the default 150s/2s budget), so leaving { once:true } listeners around on the
+// the default 150s/2s budget), so leaving { once: true } listeners around on the
 // non-abort path would pile them onto that one signal and trip Node's
 // MaxListenersExceededWarning.
 function sleep(ms: number, signal?: AbortSignal): Promise<void> {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@0xmaxma/claude-gateway",
-  "version": "1.7.1",
+  "version": "1.7.2",
   "description": "Multi-agent gateway for Claude",
   "repository": {
     "type": "git",