Propagate web-search cancellation

2026-07-25 21:18:08 -07:00
parent 802fc5339e
commit b4efe82ce2
19 changed files with 634 additions and 91 deletions
--- a/scripts/fixtures/web/mock-server.mjs
+++ b/scripts/fixtures/web/mock-server.mjs
@@ -1,6 +1,7 @@
 import http from "node:http";

 let websocketUpgrades = 0;
+let slowRequests = 0;

 const server = http.createServer((request, response) => {
  const url = new URL(request.url, "http://mock-search.test");
@@ -23,6 +24,23 @@ const server = http.createServer((request, response) => {
      </body></html>`);
    return;
  }
+  if (url.pathname === "/slow") {
+    slowRequests += 1;
+    let closed = false;
+    response.once("close", () => {
+      if (closed) return;
+      closed = true;
+      slowRequests -= 1;
+    });
+    response.writeHead(200, { "Content-Type": "text/plain" });
+    response.write("pending");
+    return;
+  }
+  if (url.pathname === "/slow-count") {
+    response.writeHead(200, { "Content-Type": "text/plain" });
+    response.end(String(slowRequests));
+    return;
+  }
  if (url.pathname === "/redirect-private") {
    response.writeHead(302, { Location: "http://127.0.0.1:8765/private" });
    response.end();
--- a/scripts/mcp-smoke-client.mjs
+++ b/scripts/mcp-smoke-client.mjs
@@ -49,7 +49,7 @@ export async function runSmoke({ usage, tmpPrefix, timeoutMs, clientInfo, scenar
  }
 }

-class McpSmokeClient {
+export class McpSmokeClient {
  constructor({ command, args, tmpPrefix }) {
    this.tmpDir = mkdtempSync(join(tmpdir(), tmpPrefix));
    this.cidFile = join(this.tmpDir, "container.cid");
--- a/scripts/release-check
+++ b/scripts/release-check
@@ -170,12 +170,28 @@ if (actual !== expected) process.exit(1);
 const serverPath = "/usr/local/lib/node_modules/@zhafron/mcp-web-search/dist/src/server.js";
 const server = fs.readFileSync(serverPath, "utf8");
 if (!server.includes("max_download_bytes: z.number().int().min(1).max(MAX_BYTES).optional()")) process.exit(1);
+if (!server.includes("provider, signal)")) process.exit(1);
+if (!server.includes("max_download_bytes,\n        signal")) process.exit(1);

 const bingPath = "/usr/local/lib/node_modules/@zhafron/mcp-web-search/dist/src/providers/bing.js";
 const bing = fs.readFileSync(bingPath, "utf8");
 if (!bing.includes("Context Kit override for @zhafron/mcp-web-search 1.3.0")) process.exit(1);
 if (!bing.includes("waitForSelector")) process.exit(1);
 if (!bing.includes("decodeBingRedirect")) process.exit(1);
+' >/dev/null
+
+  docker run --rm --entrypoint /opt/mcp-proxy/bin/python \
+    "${CONTEXT_KIT_WEB_SEARCH_IMAGE}" \
+    -c '
+from pathlib import Path
+
+root = Path("/opt/mcp-proxy/lib/python3.11/site-packages")
+session = (root / "mcp/shared/session.py").read_text()
+proxy = (root / "mcp_proxy/proxy_server.py").read_text()
+transport = (root / "mcp/client/streamable_http.py").read_text()
+assert "Context Kit: forward cancellation" in session
+assert "downstream client disconnected" in proxy
+assert "_request_cancel_scopes" in transport
 ' >/dev/null

  docker run --rm --entrypoint /usr/bin/test \
@@ -249,6 +265,7 @@ assert_web_search_backend_supervision() {
 git diff --check HEAD
 git show --check --format= HEAD >/dev/null
 git ls-files --cached --error-unmatch \
+  docker/web-search/patch-mcp-proxy.py \
  docker/web-search/patch-mcp-web-search.mjs \
  docker/web-search/overrides/bing.js \
  docker/docs/constraints.txt \
@@ -258,7 +275,10 @@ git ls-files --cached --error-unmatch \
  scripts/smoke-repomix.mjs \
  scripts/test-compose-upgrade.sh \
  scripts/test-lifecycle.sh \
+  scripts/test-web-search-candidate.sh \
  scripts/test-web-search-http.mjs \
+  scripts/test-web-search-quality.mjs \
+  scripts/test-web-search-stdio-cancellation.mjs \
  docker/web-search/mcp-probe.mjs \
  docker/web-search/http-entrypoint.mjs \
  scripts/release-check >/dev/null
@@ -267,7 +287,8 @@ bash -n scripts/release-check
 bash -n scripts/test-compose-upgrade.sh
 bash -n scripts/test-lifecycle.sh
 sh -n docker/docs/entrypoint.sh
-check_node docker/web-search/patch-mcp-web-search.mjs docker/web-search/overrides/bing.js docker/web-search/overrides/brave.js docker/web-search/overrides/browser-fetch.js docker/web-search/overrides/registry.js docker/web-search/mcp-probe.mjs docker/web-search/http-entrypoint.mjs scripts/docs-rebuild.mjs scripts/mcp-smoke-client.mjs scripts/smoke-web-search.mjs scripts/smoke-docs.mjs scripts/smoke-repomix.mjs scripts/test-docs-candidate.mjs scripts/test-web-search-candidate.mjs scripts/test-web-search-http.mjs scripts/test-web-search-quality.mjs
+check_node docker/web-search/patch-mcp-web-search.mjs docker/web-search/overrides/bing.js docker/web-search/overrides/brave.js docker/web-search/overrides/browser-fetch.js docker/web-search/overrides/registry.js docker/web-search/mcp-probe.mjs docker/web-search/http-entrypoint.mjs scripts/docs-rebuild.mjs scripts/mcp-smoke-client.mjs scripts/smoke-web-search.mjs scripts/smoke-docs.mjs scripts/smoke-repomix.mjs scripts/test-docs-candidate.mjs scripts/test-web-search-candidate.mjs scripts/test-web-search-http.mjs scripts/test-web-search-quality.mjs scripts/test-web-search-stdio-cancellation.mjs scripts/fixtures/web/mock-server.mjs
+python3 -c 'import ast, pathlib; ast.parse(pathlib.Path("docker/web-search/patch-mcp-proxy.py").read_text())'

 node -e 'const fs=require("node:fs"); JSON.parse(fs.readFileSync("snippets/opencode.json", "utf8")); JSON.parse(fs.readFileSync("snippets/claude.mcp.json", "utf8"));'
 CONTEXT_KIT_WEB_SEARCH_HTTP_URL="http://127.0.0.1:8777/mcp" CONTEXT_KIT_DOCS_HTTP_URL="http://127.0.0.1:8776/mcp" bin/context-kit install opencode > "${tmp_dir}/opencode-default.json"
--- a/scripts/test-lifecycle.sh
+++ b/scripts/test-lifecycle.sh
@@ -189,6 +189,14 @@ docker() {
      if service="$(fake_service_for_container "${container_id}" 2>/dev/null)"; then
        assert_docs_sources_restored_before_state_change
        touch "${FAKE_DOCKER_STATE}/service.${service}.running"
+      elif [[ -f "${FAKE_DOCKER_STATE}/owner.${container_id}" && "${FAKE_CLIENT_START_BLOCK:-0}" -eq 1 ]]; then
+        /bin/sh -c '
+          touch "$1"
+          while [ -f "$2" ] && [ ! -f "$3" ]; do /bin/sleep 0.02; done
+        ' sh \
+          "${FAKE_DOCKER_STATE}/client-attach.started" \
+          "${FAKE_DOCKER_STATE}/owner.${container_id}" \
+          "${FAKE_DOCKER_STATE}/client-attach.release"
      fi
      ;;
    rm)
@@ -259,7 +267,8 @@ new_case() {
  unset CONTEXT_KIT_DOCKER_CIDFILE CONTEXT_KIT_RUNTIME_DIR FAKE_DOCS_UID FAKE_WEB_UID \
    FAKE_REPLACEMENT_REQUIRED FAKE_RESTART_FAIL FAKE_DROP_RUNNING FAKE_SEARXNG_FAIL \
    FAKE_WEB_SEARCH_FAIL FAKE_DOCS_FAIL FAKE_LEGACY_CONTAINER FAKE_CLIENT_OWNER_MISMATCH \
-    FAKE_EXPECT_DOCS_SOURCES FAKE_EXPECT_DOCS_SOURCES_ABSENT CONTEXT_KIT_DOCS_LOCAL_SOURCES_DIR
+    FAKE_CLIENT_START_BLOCK FAKE_EXPECT_DOCS_SOURCES FAKE_EXPECT_DOCS_SOURCES_ABSENT \
+    CONTEXT_KIT_DOCS_LOCAL_SOURCES_DIR
  mkdir -p "${FAKE_DOCKER_STATE}" "${HOME}"
  : > "${FAKE_DOCKER_LOG}"
 }
@@ -339,11 +348,40 @@ touch "${FAKE_DOCKER_STATE}/network"
 seed_service web-search-mcp
 "${CONTEXT_KIT}" web-search </dev/null
 grep -E 'docker create .*dev.context-kit.lifecycle=client .*--entrypoint mcp-proxy .*http://web-search-mcp:8000/mcp' "${FAKE_DOCKER_LOG}" >/dev/null || fail_test "stdio bridge does not reuse the shared service"
+grep -F 'docker create -i --rm --init' "${FAKE_DOCKER_LOG}" >/dev/null || fail_test "stdio bridge container does not use Docker init"
 [[ -f "${FAKE_DOCKER_STATE}/service.web-search-mcp.running" ]] || fail_test "stdio bridge stopped the shared service"
 if compgen -G "${FAKE_DOCKER_STATE}/owner.*" >/dev/null; then
  fail_test "stdio bridge did not clean up its own container"
 fi

+new_case client-signal-cleanup
+touch "${FAKE_DOCKER_STATE}/network"
+seed_service web-search-mcp
+export FAKE_CLIENT_START_BLOCK=1
+"${CONTEXT_KIT}" web-search </dev/null >"${CASE_ROOT}/client.out" 2>&1 &
+client_pid=$!
+for _ in {1..100}; do
+  [[ -f "${FAKE_DOCKER_STATE}/client-attach.started" ]] && break
+  /bin/sleep 0.01
+done
+[[ -f "${FAKE_DOCKER_STATE}/client-attach.started" ]] || fail_test "blocking stdio attach did not start"
+kill -TERM "${client_pid}"
+owner_removed=0
+for _ in {1..50}; do
+  if ! compgen -G "${FAKE_DOCKER_STATE}/owner.*" >/dev/null; then
+    owner_removed=1
+    break
+  fi
+  /bin/sleep 0.01
+done
+touch "${FAKE_DOCKER_STATE}/client-attach.release"
+set +e
+wait "${client_pid}"
+client_status=$?
+set -e
+[[ "${owner_removed}" -eq 1 ]] || fail_test "SIGTERM did not promptly remove the owned stdio container"
+[[ "${client_status}" -eq 143 ]] || fail_test "SIGTERM returned ${client_status} instead of 143"
+
 new_case client-owner-isolation
 touch "${FAKE_DOCKER_STATE}/network"
 seed_service web-search-mcp
--- a/scripts/test-web-search-candidate.mjs
+++ b/scripts/test-web-search-candidate.mjs
@@ -1,4 +1,5 @@
 import assert from "node:assert/strict";
+import { setTimeout as delay } from "node:timers/promises";

 import { probeMcp, rpc } from "../docker/web-search/mcp-probe.mjs";

@@ -56,4 +57,40 @@ const websocketCount = payload(await rpc(url, 9, "tools/call", {
 }, 30_000));
 assert.equal(websocketCount.content.trim(), "0");

-console.log("pass web-search candidate diagnostics, browser rendering, and SSRF rejection");
+let requestId = 10;
+async function slowCount() {
+  const result = payload(await rpc(url, requestId++, "tools/call", {
+    name: "fetch_url",
+    arguments: { url: "http://mock-search.test:8080/slow-count", engine: "http", format: "text", fresh: true }
+  }, 5_000));
+  return Number(result.content.trim());
+}
+
+async function waitForSlowCount(expected, timeoutMs) {
+  const deadline = Date.now() + timeoutMs;
+  let actual;
+  while (Date.now() < deadline) {
+    actual = await slowCount();
+    if (actual === expected) return;
+    await delay(50);
+  }
+  assert.equal(actual, expected, `slow request count did not reach ${expected}`);
+}
+
+const cancellation = new AbortController();
+const pendingFetch = rpc(url, requestId++, "tools/call", {
+  name: "fetch_url",
+  arguments: {
+    url: "http://mock-search.test:8080/slow",
+    engine: "browser",
+    format: "text",
+    fresh: true,
+    timeout_ms: 120_000
+  }
+}, 120_000, cancellation.signal);
+await waitForSlowCount(1, 10_000);
+cancellation.abort(new Error("candidate client disconnected"));
+await assert.rejects(pendingFetch, /candidate client disconnected/);
+await waitForSlowCount(0, 3_000);
+
+console.log("pass web-search candidate diagnostics, browser rendering, SSRF rejection, and cancellation");
--- a/scripts/test-web-search-candidate.sh
+++ b/scripts/test-web-search-candidate.sh
@@ -6,9 +6,10 @@ IMAGE="${CONTEXT_KIT_WEB_SEARCH_CANDIDATE_IMAGE:-context-kit/web-search-mcp:qual
 NETWORK="context-kit-web-quality-$RANDOM-$$"
 MOCK="${NETWORK}-mock"
 SERVER="${NETWORK}-server"
+BRIDGE="${NETWORK}-bridge"

 cleanup() {
-  docker rm -f "${SERVER}" "${MOCK}" >/dev/null 2>&1 || true
+  docker rm -f "${BRIDGE}" "${SERVER}" "${MOCK}" >/dev/null 2>&1 || true
  docker network rm "${NETWORK}" >/dev/null 2>&1 || true
 }
 trap cleanup EXIT
@@ -18,17 +19,38 @@ docker run -d --name "${MOCK}" --network "${NETWORK}" --ip 203.0.113.10 \
  --network-alias mock-search.test \
  -v "${ROOT}/scripts/fixtures/web/mock-server.mjs:/fixture/mock-server.mjs:ro" \
  node:22-bookworm-slim node /fixture/mock-server.mjs >/dev/null
-docker run -d --name "${SERVER}" --network "${NETWORK}" --ip 203.0.113.11 \
+docker run -d --init --name "${SERVER}" --network "${NETWORK}" --ip 203.0.113.11 \
+  --network-alias web-search-mcp \
  -p 127.0.0.1::8000 \
  -e SEARXNG_URL=http://mock-search.test:8080 \
  -e DEFAULT_SEARCH_PROVIDER=searxng \
  "${IMAGE}" >/dev/null

+[[ "$(docker inspect -f '{{.HostConfig.Init}}' "${SERVER}")" == true ]] || {
+  printf 'candidate web-search container does not use Docker init\n' >&2
+  exit 1
+}
+
 binding="$(docker port "${SERVER}" 8000/tcp)"
 port="${binding##*:}"
 for _ in {1..120}; do
  if curl -fsS "http://127.0.0.1:${port}/healthz" >/dev/null 2>&1; then
    node "${ROOT}/scripts/test-web-search-candidate.mjs" "http://127.0.0.1:${port}/mcp"
+    node "${ROOT}/scripts/test-web-search-stdio-cancellation.mjs" \
+      docker run --rm --init -i --name "${BRIDGE}" --network "${NETWORK}" \
+      --entrypoint mcp-proxy "${IMAGE}" --transport streamablehttp "http://web-search-mcp:8000/mcp"
+    docker exec "${SERVER}" sh -eu -c '
+      for status in /proc/[0-9]*/status; do
+        while IFS=: read -r key value; do
+          if [ "$key" = State ]; then
+            case "$value" in
+              *Z*) printf "zombie process found in %s: %s\n" "$status" "$value" >&2; exit 1 ;;
+            esac
+            break
+          fi
+        done < "$status"
+      done
+    '
    exit 0
  fi
  sleep 0.25
--- a/scripts/test-web-search-http.mjs
+++ b/scripts/test-web-search-http.mjs
@@ -1,15 +1,22 @@
 import assert from "node:assert/strict";
 import http from "node:http";
-import { once } from "node:events";
+import { EventEmitter, once } from "node:events";
+import { setTimeout as delay } from "node:timers/promises";

 import {
  createSecureMcpServer,
  mcpProxyArguments,
-  superviseBackend
+  superviseBackend,
+  terminateChild
 } from "../docker/web-search/http-entrypoint.mjs";
 import { probeMcp } from "../docker/web-search/mcp-probe.mjs";

 let backendAlive = true;
+let hangingBackendResponse;
+let resolveHangingBackendStarted;
+let resolveHangingBackendClosed;
+const hangingBackendStarted = new Promise(resolve => { resolveHangingBackendStarted = resolve; });
+const hangingBackendClosed = new Promise(resolve => { resolveHangingBackendClosed = resolve; });
 assert(mcpProxyArguments.includes("--stateless"));
 const backend = http.createServer(async (request, response) => {
  if (request.url === "/status") {
@@ -47,6 +54,14 @@ const backend = http.createServer(async (request, response) => {
    }));
    return;
  }
+  if (message.method === "tools/call" && message.params?.name === "hang") {
+    hangingBackendResponse = response;
+    response.once("close", resolveHangingBackendClosed);
+    response.writeHead(200, { "Content-Type": "application/json" });
+    response.write("pending");
+    resolveHangingBackendStarted();
+    return;
+  }
  response.writeHead(500, { "Content-Type": "text/plain" });
  response.end("backend dead");
 });
@@ -89,6 +104,35 @@ assert.equal(await rawRequest({ host: "attacker.example" }), 421);
 assert.equal(await rawRequest({ host: `127.0.0.1:${frontPort}`, origin: "https://attacker.example" }), 403);
 assert.equal(await rawRequest({ host: `127.0.0.1:${frontPort}`, origin: `http://127.0.0.1:${frontPort}` }), 403);

+let downstreamResponse;
+try {
+  const responseReceived = new Promise((resolve, reject) => {
+    const request = http.request({
+      hostname: "127.0.0.1",
+      port: frontPort,
+      path: "/mcp",
+      method: "POST",
+      headers: {
+        Accept: "application/json, text/event-stream",
+        "Content-Type": "application/json"
+      }
+    }, resolve);
+    request.once("error", reject);
+    request.end('{"jsonrpc":"2.0","id":9,"method":"tools/call","params":{"name":"hang","arguments":{}}}');
+  });
+  await hangingBackendStarted;
+  downstreamResponse = await responseReceived;
+  downstreamResponse.on("error", () => {});
+  downstreamResponse.destroy();
+  await Promise.race([
+    hangingBackendClosed,
+    delay(250).then(() => { throw new Error("upstream request remained open after downstream disconnect"); })
+  ]);
+} finally {
+  downstreamResponse?.destroy();
+  hangingBackendResponse?.destroy();
+}
+
 backendAlive = false;
 assert.equal((await fetch(`${upstream}/status`)).status, 200);
 assert.equal((await fetch(`http://127.0.0.1:${frontPort}/healthz`)).status, 503);
@@ -106,6 +150,28 @@ await new Promise((resolve, reject) => {
  });
 });

+class StubbornChild extends EventEmitter {
+  exitCode = null;
+  signalCode = null;
+  signals = [];
+
+  kill(signal) {
+    this.signals.push(signal);
+    if (signal === "SIGKILL") {
+      setTimeout(() => {
+        this.signalCode = signal;
+        this.emit("exit", null, signal);
+      }, 10);
+    }
+    return true;
+  }
+}
+
+const stubbornChild = new StubbornChild();
+await terminateChild(stubbornChild, { graceMs: 1 });
+assert.deepEqual(stubbornChild.signals, ["SIGTERM", "SIGKILL"]);
+assert.equal(stubbornChild.signalCode, "SIGKILL");
+
 front.close();
 front.closeAllConnections();
 backend.close();
--- a/scripts/test-web-search-quality.mjs
+++ b/scripts/test-web-search-quality.mjs
@@ -30,19 +30,44 @@ assert.equal(failed.diagnostic.status, "error");
 assert.equal(failed.diagnostic.error.category, "network");

 let underlyingAborted = false;
+let underlyingCleanupFinished = false;
 const timedOut = await attemptProvider({
  name: "slow",
  async search(_query, _limit, _lang, signal) {
    await new Promise((resolve, reject) => {
      signal.addEventListener("abort", () => {
        underlyingAborted = true;
-        reject(signal.reason);
+        setTimeout(() => {
+          underlyingCleanupFinished = true;
+          reject(signal.reason);
+        }, 25);
      }, { once: true });
    });
  }
 }, "q", 3, "en", { timeoutMs: 20 });
 assert.equal(timedOut.diagnostic.error.category, "timeout");
 assert.equal(underlyingAborted, true);
+assert.equal(underlyingCleanupFinished, true);
+
+const cancellation = new AbortController();
+const cancellationReason = new Error("search request cancelled");
+let cancellationCleanupFinished = false;
+const cancelled = attemptProvider({
+  name: "cancelled",
+  async search(_query, _limit, _lang, signal) {
+    await new Promise((resolve, reject) => {
+      signal.addEventListener("abort", () => {
+        setTimeout(() => {
+          cancellationCleanupFinished = true;
+          reject(signal.reason);
+        }, 10);
+      }, { once: true });
+    });
+  }
+}, "q", 3, "en", { timeoutMs: 1000, signal: cancellation.signal });
+cancellation.abort(cancellationReason);
+await assert.rejects(cancelled, error => error === cancellationReason);
+assert.equal(cancellationCleanupFinished, true);

 const result = boundFetchCollections({
  links: Array.from({ length: 550 }, (_, index) => ({ url: `https://example.test/${index}` })),
--- a/scripts/test-web-search-stdio-cancellation.mjs
+++ b/scripts/test-web-search-stdio-cancellation.mjs
@@ -0,0 +1,66 @@
+import assert from "node:assert/strict";
+import { setTimeout as delay } from "node:timers/promises";
+
+import {
+  McpSmokeClient,
+  requireToolSuccess,
+  textFrom
+} from "./mcp-smoke-client.mjs";
+
+const command = process.argv[2];
+const args = process.argv.slice(3);
+if (!command) throw new Error("usage: node scripts/test-web-search-stdio-cancellation.mjs <command> [args...]");
+
+const client = new McpSmokeClient({
+  command,
+  args,
+  tmpPrefix: "context-kit-stdio-cancellation-"
+});
+
+function payload(result) {
+  const text = textFrom(requireToolSuccess("fetch_url", result));
+  return JSON.parse(text);
+}
+
+async function slowCount() {
+  const result = await client.callTool("fetch_url", {
+    url: "http://mock-search.test:8080/slow-count",
+    engine: "http",
+    format: "text",
+    fresh: true
+  });
+  return Number(payload(result).content.trim());
+}
+
+async function waitForSlowCount(expected, timeoutMs) {
+  const deadline = Date.now() + timeoutMs;
+  let actual;
+  while (Date.now() < deadline) {
+    actual = await slowCount();
+    if (actual === expected) return;
+    await delay(50);
+  }
+  assert.equal(actual, expected, `stdio slow request count did not reach ${expected}`);
+}
+
+try {
+  await client.initialize({ name: "context-kit-stdio-cancellation", version: "1" });
+  const requestId = client.nextId;
+  const pendingFetch = client.callTool("fetch_url", {
+    url: "http://mock-search.test:8080/slow",
+    engine: "browser",
+    format: "text",
+    fresh: true,
+    timeout_ms: 120_000
+  });
+  await waitForSlowCount(1, 10_000);
+  client.notify("notifications/cancelled", {
+    requestId,
+    reason: "stdio client disconnected"
+  });
+  await assert.rejects(pendingFetch, /cancel/i);
+  await waitForSlowCount(0, 3_000);
+  console.log("pass web-search stdio bridge cancellation");
+} finally {
+  await client.stop();
+}