diff --git a/bin/context-kit b/bin/context-kit index 22bd41b..535802f 100755 --- a/bin/context-kit +++ b/bin/context-kit @@ -382,19 +382,28 @@ cleanup_owned_container() { docker rm -f "${container_id}" >/dev/null 2>&1 || true } +cleanup_owned_stdio_container() { + local container_id="$1" owner="$2" attach_pid="$3" + cleanup_owned_container "${container_id}" "${owner}" + if [[ -n "${attach_pid}" ]]; then + kill "${attach_pid}" >/dev/null 2>&1 || true + wait "${attach_pid}" 2>/dev/null || true + fi +} + run_owned_stdio_container() { local role="$1" shift - local uid owner name container_id='' status=0 + local uid owner name container_id='' attach_pid='' status=0 uid="$(id -u)" owner="${PROJECT}:${role}:${uid}:$$" name="${PROJECT}-${role}-${uid}-$$" - trap 'cleanup_owned_container "${container_id}" "${owner}"' EXIT + trap 'cleanup_owned_stdio_container "${container_id}" "${owner}" "${attach_pid}"' EXIT trap 'exit 129' HUP trap 'exit 130' INT trap 'exit 143' TERM - container_id="$(docker create -i --rm \ + container_id="$(docker create -i --rm --init \ --name "${name}" \ --label dev.context-kit=true \ --label dev.context-kit.lifecycle=client \ @@ -411,7 +420,10 @@ run_owned_stdio_container() { printf '%s\n' "${container_id}" > "${CONTEXT_KIT_DOCKER_CIDFILE}" fi - docker start -ai "${container_id}" <&0 || status=$? + docker start -ai "${container_id}" <&0 & + attach_pid=$! + wait "${attach_pid}" || status=$? + attach_pid='' cleanup_owned_container "${container_id}" "${owner}" trap - EXIT HUP INT TERM return "${status}" diff --git a/docker/web-search/.dockerignore b/docker/web-search/.dockerignore index b912d3c..e0f74cd 100644 --- a/docker/web-search/.dockerignore +++ b/docker/web-search/.dockerignore @@ -2,6 +2,7 @@ !Dockerfile !http-entrypoint.mjs !mcp-probe.mjs +!patch-mcp-proxy.py !patch-mcp-web-search.mjs !overrides/ !overrides/bing.js diff --git a/docker/web-search/Dockerfile b/docker/web-search/Dockerfile index 5e65bc2..d177ce5 100644 --- a/docker/web-search/Dockerfile +++ b/docker/web-search/Dockerfile @@ -3,6 +3,7 @@ FROM node:22-bookworm-slim@sha256:813a7480f28fdadac1f7f5c824bcdad435b5bc1322a596 ARG MCP_WEB_SEARCH_VERSION=1.3.0 ARG MCP_WEB_SEARCH_MAX_BYTES=52428800 ARG MCP_PROXY_VERSION=0.12.0 +ARG MCP_PYTHON_SDK_VERSION=1.28.1 # Chromium intentionally tracks Debian security updates inside the pinned base # image family; Bing's browser path is more likely to break with stale Chromium @@ -15,8 +16,13 @@ RUN apt-get update \ python3-venv \ && rm -rf /var/lib/apt/lists/* +COPY patch-mcp-proxy.py /tmp/patch-mcp-proxy.py + RUN python3 -m venv /opt/mcp-proxy \ - && /opt/mcp-proxy/bin/pip install --no-cache-dir "mcp-proxy==${MCP_PROXY_VERSION}" \ + && /opt/mcp-proxy/bin/pip install --no-cache-dir \ + "mcp-proxy==${MCP_PROXY_VERSION}" \ + "mcp==${MCP_PYTHON_SDK_VERSION}" \ + && /opt/mcp-proxy/bin/python /tmp/patch-mcp-proxy.py \ && /opt/mcp-proxy/bin/mcp-proxy --version COPY patch-mcp-web-search.mjs /tmp/patch-mcp-web-search.mjs @@ -35,7 +41,7 @@ RUN npm install -g "@zhafron/mcp-web-search@${MCP_WEB_SEARCH_VERSION}" \ && cp /tmp/context-kit-fetch/browser-fetch.js /usr/local/lib/node_modules/@zhafron/mcp-web-search/dist/src/fetch/browser.js \ && cp /tmp/context-kit-fetch/bounds.mjs /usr/local/lib/node_modules/@zhafron/mcp-web-search/dist/src/fetch/bounds.js \ && node /tmp/patch-mcp-web-search.mjs \ - && rm -rf /tmp/patch-mcp-web-search.mjs /tmp/context-kit-bing-provider.js /tmp/context-kit-providers /tmp/context-kit-fetch \ + && rm -rf /tmp/patch-mcp-proxy.py /tmp/patch-mcp-web-search.mjs /tmp/context-kit-bing-provider.js /tmp/context-kit-providers /tmp/context-kit-fetch \ && npm cache clean --force RUN chmod -R a+rX /usr/local/lib/context-kit \ diff --git a/docker/web-search/http-entrypoint.mjs b/docker/web-search/http-entrypoint.mjs index 9f34c6b..ec7f3ac 100644 --- a/docker/web-search/http-entrypoint.mjs +++ b/docker/web-search/http-entrypoint.mjs @@ -88,13 +88,22 @@ export function createSecureMcpServer({ upstream = defaultUpstream, probe = prob path: request.url, headers: copyRequestHeaders(request.headers, target.host) }, upstreamResponse => { + upstreamResponse.on("error", () => response.destroy()); response.writeHead( upstreamResponse.statusCode || 502, copyResponseHeaders(upstreamResponse.headers) ); upstreamResponse.pipe(response); }); + const abortUpstream = () => { + if (!upstreamRequest.destroyed) upstreamRequest.destroy(new Error("downstream disconnected")); + }; + request.once("aborted", abortUpstream); + response.once("close", () => { + if (!response.writableEnded) abortUpstream(); + }); upstreamRequest.on("error", error => { + if (response.destroyed) return; if (!response.headersSent) response.writeHead(502, { "Content-Type": "text/plain" }); response.end(`backend unavailable: ${error.message}`); }); @@ -122,6 +131,23 @@ export function superviseBackend({ probe, intervalMs = 10000, onFailure }) { }; } +export async function terminateChild(child, { graceMs = 3000 } = {}) { + if (child.exitCode !== null || child.signalCode !== null) return; + const gracefulExit = once(child, "exit").then(() => true); + child.kill("SIGTERM"); + const exited = await Promise.race([ + gracefulExit, + delay(graceMs).then(() => false) + ]); + if (exited || child.exitCode !== null || child.signalCode !== null) return; + + const forcedExit = once(child, "exit"); + if (!child.kill("SIGKILL") && child.exitCode === null && child.signalCode === null) { + throw new Error("failed to terminate mcp-proxy child"); + } + await forcedExit; +} + async function waitForBackend(child, url) { let lastError; for (let attempt = 0; attempt < 60; attempt += 1) { @@ -152,11 +178,7 @@ async function main() { stopSupervisor(); server?.close(); server?.closeAllConnections(); - if (child.exitCode === null) { - child.kill("SIGTERM"); - await Promise.race([once(child, "exit"), delay(3000)]).catch(() => {}); - if (child.exitCode === null) child.kill("SIGKILL"); - } + await terminateChild(child); process.exitCode = code; }; diff --git a/docker/web-search/mcp-probe.mjs b/docker/web-search/mcp-probe.mjs index 022460e..318eeee 100644 --- a/docker/web-search/mcp-probe.mjs +++ b/docker/web-search/mcp-probe.mjs @@ -3,7 +3,8 @@ import { fileURLToPath } from "node:url"; const protocolVersion = "2024-11-05"; const expectedTools = ["fetch_url", "search_web"]; -export async function rpc(url, id, method, params = {}, timeoutMs = 5000) { +export async function rpc(url, id, method, params = {}, timeoutMs = 5000, signal) { + const timeoutSignal = AbortSignal.timeout(timeoutMs); const response = await fetch(url, { method: "POST", headers: { @@ -12,7 +13,7 @@ export async function rpc(url, id, method, params = {}, timeoutMs = 5000) { "MCP-Protocol-Version": protocolVersion }, body: JSON.stringify({ jsonrpc: "2.0", id, method, params }), - signal: AbortSignal.timeout(timeoutMs) + signal: signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal }); const text = await response.text(); if (!response.ok) throw new Error(`${method} returned HTTP ${response.status}: ${text.slice(0, 300)}`); diff --git a/docker/web-search/overrides/browser-fetch.js b/docker/web-search/overrides/browser-fetch.js index c08e329..673a49f 100644 --- a/docker/web-search/overrides/browser-fetch.js +++ b/docker/web-search/overrides/browser-fetch.js @@ -12,61 +12,69 @@ function responseHeaders(headers) { return record; } -export async function fetchBrowserResource(url, timeoutMs = HTTP_TIMEOUT) { +export async function fetchBrowserResource(url, timeoutMs = HTTP_TIMEOUT, signal) { + signal?.throwIfAborted(); await assertSafeUrl(url); return browserPool.withBrowser(async browser => { const page = await browser.newPage(); - const devtools = await page.target().createCDPSession(); - await devtools.send("Network.enable"); - await devtools.send("Network.setBlockedURLs", { - urls: ["ws://*", "wss://*", "file://*", "ftp://*"] - }); - await page.evaluateOnNewDocument(() => { - const blockedTransport = name => class { - constructor() { - throw new DOMException(`${name} is disabled by the safe browser fetcher`, "SecurityError"); - } - }; - for (const name of ["WebSocket", "WebTransport", "RTCPeerConnection", "webkitRTCPeerConnection"]) { - if (name in globalThis) { - Object.defineProperty(globalThis, name, { - configurable: false, - writable: false, - value: blockedTransport(name) - }); - } - } - }); - await page.setBypassServiceWorker(true); - let requests = 0; - let totalBytes = 0; - let blockedError; - await page.setRequestInterception(true); - page.on("request", request => { - void (async () => { - try { - const requestUrl = new URL(request.url()); - if (!["http:", "https:"].includes(requestUrl.protocol)) throw new Error("unsupported browser request scheme"); - if (request.method() !== "GET") throw new Error("browser fetch blocks non-GET requests"); - requests += 1; - if (requests > MAX_BROWSER_REQUESTS) throw new Error("browser request limit exceeded"); - await assertSafeUrl(requestUrl); - const upstream = await fetchViaVettedAddress(requestUrl, timeoutMs); - const body = Buffer.from(await upstream.arrayBuffer()); - totalBytes += body.byteLength; - if (totalBytes > MAX_BROWSER_TOTAL_BYTES) throw new Error("browser byte limit exceeded"); - await request.respond({ - status: upstream.status, - headers: responseHeaders(upstream.headers), - body - }); - } catch (error) { - blockedError ||= error; - await request.abort("blockedbyclient").catch(() => undefined); - } - })(); - }); + const pendingRequests = new Set(); + const abort = () => void page.close().catch(() => undefined); + signal?.addEventListener("abort", abort, { once: true }); try { + signal?.throwIfAborted(); + const devtools = await page.target().createCDPSession(); + await devtools.send("Network.enable"); + await devtools.send("Network.setBlockedURLs", { + urls: ["ws://*", "wss://*", "file://*", "ftp://*"] + }); + await page.evaluateOnNewDocument(() => { + const blockedTransport = name => class { + constructor() { + throw new DOMException(`${name} is disabled by the safe browser fetcher`, "SecurityError"); + } + }; + for (const name of ["WebSocket", "WebTransport", "RTCPeerConnection", "webkitRTCPeerConnection"]) { + if (name in globalThis) { + Object.defineProperty(globalThis, name, { + configurable: false, + writable: false, + value: blockedTransport(name) + }); + } + } + }); + await page.setBypassServiceWorker(true); + let requests = 0; + let totalBytes = 0; + let blockedError; + await page.setRequestInterception(true); + page.on("request", request => { + const pending = (async () => { + try { + const requestUrl = new URL(request.url()); + if (!["http:", "https:"].includes(requestUrl.protocol)) throw new Error("unsupported browser request scheme"); + if (request.method() !== "GET") throw new Error("browser fetch blocks non-GET requests"); + requests += 1; + if (requests > MAX_BROWSER_REQUESTS) throw new Error("browser request limit exceeded"); + await assertSafeUrl(requestUrl); + const upstream = await fetchViaVettedAddress(requestUrl, timeoutMs, signal); + const body = Buffer.from(await upstream.arrayBuffer()); + totalBytes += body.byteLength; + if (totalBytes > MAX_BROWSER_TOTAL_BYTES) throw new Error("browser byte limit exceeded"); + await request.respond({ + status: upstream.status, + headers: responseHeaders(upstream.headers), + body + }); + } catch (error) { + blockedError ||= error; + await request.abort("blockedbyclient").catch(() => undefined); + } + })(); + pendingRequests.add(pending); + void pending.finally(() => pendingRequests.delete(pending)); + }); + signal?.throwIfAborted(); const navigation = await page.goto(url.toString(), { waitUntil: "networkidle2", timeout: timeoutMs @@ -88,8 +96,13 @@ export async function fetchBrowserResource(url, timeoutMs = HTTP_TIMEOUT) { buffer, byteLength: buffer.byteLength }; + } catch (error) { + signal?.throwIfAborted(); + throw error; } finally { - await page.close(); + signal?.removeEventListener("abort", abort); + if (!page.isClosed()) await page.close(); + await Promise.allSettled(pendingRequests); } }); } diff --git a/docker/web-search/overrides/diagnostics.mjs b/docker/web-search/overrides/diagnostics.mjs index 02e3acf..5e6c7ba 100644 --- a/docker/web-search/overrides/diagnostics.mjs +++ b/docker/web-search/overrides/diagnostics.mjs @@ -17,24 +17,22 @@ export async function attemptProvider(provider, query, limit, lang, options = {} const timeoutMs = Math.max(10, Math.min(options.timeoutMs || DEFAULT_TIMEOUT_MS, 60_000)); const now = options.now || (() => performance.now()); const started = now(); + options.signal?.throwIfAborted(); if (provider.configured === false) { return { items: [], diagnostic: { provider: provider.name, status: "unavailable", duration_ms: 0, result_count: 0 } }; } - let timer; const controller = new AbortController(); + const signal = options.signal + ? AbortSignal.any([options.signal, controller.signal]) + : controller.signal; + const timer = setTimeout(() => { + controller.abort(new Error(`provider timed out after ${timeoutMs}ms`)); + }, timeoutMs); try { - const items = await Promise.race([ - provider.search(query, limit, lang, controller.signal), - new Promise((_, reject) => { - timer = setTimeout(() => { - controller.abort(new Error(`provider timed out after ${timeoutMs}ms`)); - reject(new Error(`provider timed out after ${timeoutMs}ms`)); - }, timeoutMs); - }) - ]); + const items = await provider.search(query, limit, lang, signal); const bounded = Array.isArray(items) ? items.slice(0, limit) : []; return { items: bounded, @@ -46,6 +44,7 @@ export async function attemptProvider(provider, query, limit, lang, options = {} } }; } catch (error) { + options.signal?.throwIfAborted(); return { items: [], diagnostic: { @@ -57,7 +56,6 @@ export async function attemptProvider(provider, query, limit, lang, options = {} } }; } finally { - controller.abort(); clearTimeout(timer); } } diff --git a/docker/web-search/overrides/registry.js b/docker/web-search/overrides/registry.js index b3c72da..7ca17cc 100644 --- a/docker/web-search/overrides/registry.js +++ b/docker/web-search/overrides/registry.js @@ -23,15 +23,19 @@ export class ProviderRegistry { return this.providers.get(name); } - async searchWithFallback(q, limit, lang, preferredProvider) { + async searchWithFallback(q, limit, lang, preferredProvider, signal) { const defaultProvider = preferredProvider || DEFAULT_SEARCH_PROVIDER; const order = [defaultProvider, ...PROVIDERS.filter(name => name !== defaultProvider)].slice(0, MAX_PROVIDER_ATTEMPTS); const attempts = []; const started = performance.now(); for (const providerName of order) { + signal?.throwIfAborted(); const provider = this.providers.get(providerName); if (!provider) continue; - const attempt = await attemptProvider(provider, q, limit, lang, { timeoutMs: PROVIDER_TIMEOUT_MS }); + const attempt = await attemptProvider(provider, q, limit, lang, { + timeoutMs: PROVIDER_TIMEOUT_MS, + signal + }); attempts.push(attempt.diagnostic); if (attempt.items.length) { return { diff --git a/docker/web-search/patch-mcp-proxy.py b/docker/web-search/patch-mcp-proxy.py new file mode 100644 index 0000000..fb1d3d6 --- /dev/null +++ b/docker/web-search/patch-mcp-proxy.py @@ -0,0 +1,168 @@ +#!/usr/bin/env python3 +from importlib.metadata import version +from importlib.util import find_spec +from pathlib import Path + + +EXPECTED_VERSIONS = { + "mcp-proxy": "0.12.0", + "mcp": "1.28.1", +} + + +def module_path(name: str) -> Path: + spec = find_spec(name) + if spec is None or spec.origin is None: + raise RuntimeError(f"module not found: {name}") + return Path(spec.origin) + + +def replace_once(path: Path, before: str, after: str) -> None: + source = path.read_text() + count = source.count(before) + if count != 1: + raise RuntimeError(f"expected one patch target in {path}, found {count}: {before}") + path.write_text(source.replace(before, after)) + + +for package, expected in EXPECTED_VERSIONS.items(): + actual = version(package) + if actual != expected: + raise RuntimeError(f"expected {package} {expected}, found {actual}") + + +proxy_path = module_path("mcp_proxy.proxy_server") +replace_once( + proxy_path, + "import logging\nimport typing as t\n\nfrom mcp import server, types", + "import logging\nimport typing as t\n\nimport anyio\n\nfrom mcp import server, types", +) +replace_once( + proxy_path, + """ result = await remote_app.call_tool( + req.params.name, + (req.params.arguments or {}), + meta=meta_dict, + progress_callback=progress_callback, + ) +""", + """ completed = anyio.Event() + disconnected = False + downstream_request = request_ctx.get().request + + async def watch_downstream_disconnect() -> None: + nonlocal disconnected + while not completed.is_set(): + if await downstream_request.is_disconnected(): + disconnected = True + task_group.cancel_scope.cancel() + return + await anyio.sleep(0.05) + + async with anyio.create_task_group() as task_group: + if downstream_request is not None and hasattr(downstream_request, "is_disconnected"): + task_group.start_soon(watch_downstream_disconnect) + try: + result = await remote_app.call_tool( + req.params.name, + (req.params.arguments or {}), + meta=meta_dict, + progress_callback=progress_callback, + ) + finally: + completed.set() + task_group.cancel_scope.cancel() + + if disconnected: + raise ConnectionError("downstream client disconnected") +""", +) + + +session_path = module_path("mcp.shared.session") +replace_once( + session_path, + """ finally: + self._response_streams.pop(request_id, None) + self._progress_callbacks.pop(request_id, None) +""", + """ except anyio.get_cancelled_exc_class(): + # Context Kit: forward cancellation before abandoning the remote request. + with anyio.move_on_after(1, shield=True): + try: + await self.send_notification( + CancelledNotification( + params={"requestId": request_id, "reason": "upstream request cancelled"} + ) + ) + except Exception: + pass + raise + finally: + self._response_streams.pop(request_id, None) + self._progress_callbacks.pop(request_id, None) +""", +) + + +streamable_http_path = module_path("mcp.client.streamable_http") +replace_once( + streamable_http_path, + """ self.url = url + self.session_id = None + self.protocol_version = None +""", + """ self.url = url + self.session_id = None + self.protocol_version = None + self._request_cancel_scopes: dict[RequestId, anyio.CancelScope] = {} +""", +) +replace_once( + streamable_http_path, + """ async def handle_request_async(): + if is_resumption: + await self._handle_resumption_request(ctx) + else: + await self._handle_post_request(ctx) + + # If this is a request, start a new task to handle it + if isinstance(message.root, JSONRPCRequest): + tg.start_soon(handle_request_async) + else: + await handle_request_async() +""", + """ async def handle_request_async( + request_context: RequestContext = ctx, + resume: bool = is_resumption, + ) -> None: + root = request_context.session_message.message.root + request_id = root.id if isinstance(root, JSONRPCRequest) else None + with anyio.CancelScope() as request_scope: + if request_id is not None: + self._request_cancel_scopes[request_id] = request_scope + try: + if resume: + await self._handle_resumption_request(request_context) + else: + await self._handle_post_request(request_context) + finally: + if self._request_cancel_scopes.get(request_id) is request_scope: + self._request_cancel_scopes.pop(request_id, None) + + # If this is a request, start a new task to handle it + if isinstance(message.root, JSONRPCRequest): + tg.start_soon(handle_request_async) + else: + if ( + isinstance(message.root, JSONRPCNotification) + and message.root.method == "notifications/cancelled" + ): + request_scope = self._request_cancel_scopes.get( + (message.root.params or {}).get("requestId") + ) + if request_scope is not None: + request_scope.cancel() + await handle_request_async() +""", +) diff --git a/docker/web-search/patch-mcp-web-search.mjs b/docker/web-search/patch-mcp-web-search.mjs index 2cdb035..4847130 100644 --- a/docker/web-search/patch-mcp-web-search.mjs +++ b/docker/web-search/patch-mcp-web-search.mjs @@ -23,6 +23,22 @@ const replacements = [ [ "Search the web using multiple providers (DuckDuckGo, Bing, SearXNG). Automatically falls back to other providers if the default fails. No API keys required for DuckDuckGo and SearXNG.", "Search the web with bounded provider fallback and per-attempt diagnostics. SearXNG is local; Brave is available when BRAVE_SEARCH_API_KEY is configured." + ], + [ + '}, async ({ q, limit = DEFAULT_LIMIT, lang = "en", provider }) => {', + '}, async ({ q, limit = DEFAULT_LIMIT, lang = "en", provider }, { signal }) => {' + ], + [ + "providerRegistry.searchWithFallback(q, Math.min(Math.max(1, limit), 50), lang, provider)", + "providerRegistry.searchWithFallback(q, Math.min(Math.max(1, limit), 50), lang, provider, signal)" + ], + [ + '}, async ({ url, format, max_length, start_index, engine, include_links, include_media, include_metadata, include_comments, comment_limit, comment_sort, max_depth, timeout_ms, fresh, download, download_dir, download_ttl_seconds, max_download_bytes }) => {', + '}, async ({ url, format, max_length, start_index, engine, include_links, include_media, include_metadata, include_comments, comment_limit, comment_sort, max_depth, timeout_ms, fresh, download, download_dir, download_ttl_seconds, max_download_bytes }, { signal }) => {' + ], + [ + " max_download_bytes\n });", + " max_download_bytes,\n signal\n });" ] ]; @@ -39,7 +55,13 @@ const httpPath = "/usr/local/lib/node_modules/@zhafron/mcp-web-search/dist/src/f let httpSource = fs.readFileSync(httpPath, "utf8"); const privateTransport = "async function fetchViaVettedAddress(url, timeoutMs)"; if (!httpSource.includes(privateTransport)) throw new Error(`mcp-web-search patch target not found: ${privateTransport}`); -httpSource = httpSource.replace(privateTransport, "export async function fetchViaVettedAddress(url, timeoutMs)"); +httpSource = httpSource.replace(privateTransport, "export async function fetchViaVettedAddress(url, timeoutMs, signal)"); +const requestTimeout = " timeout: timeoutMs\n }, response => {"; +if (!httpSource.includes(requestTimeout)) throw new Error(`mcp-web-search patch target not found: ${requestTimeout}`); +httpSource = httpSource.replace(requestTimeout, " timeout: timeoutMs,\n signal\n }, response => {"); +const transportCall = "response = await transport(currentUrl, timeoutMs);"; +if (!httpSource.includes(transportCall)) throw new Error(`mcp-web-search patch target not found: ${transportCall}`); +httpSource = httpSource.replace(transportCall, "response = await transport(currentUrl, timeoutMs, options?.signal);"); fs.writeFileSync(httpPath, httpSource); const utilityHttpPath = "/usr/local/lib/node_modules/@zhafron/mcp-web-search/dist/src/utils/http.js"; @@ -59,17 +81,20 @@ const extractReplacements = [ ], [ "fetchCache.set(cacheKey, siteResult);\n return siteResult;", - "const boundedSiteResult = boundFetchCollections(siteResult);\n fetchCache.set(cacheKey, boundedSiteResult);\n return boundedSiteResult;" + "const boundedSiteResult = boundFetchCollections(siteResult);\n options?.signal?.throwIfAborted();\n fetchCache.set(cacheKey, boundedSiteResult);\n return boundedSiteResult;" ], [ "const resource = await fetchResource(parsedUrl, options?.timeout_ms, transport, options);", - 'const resource = options?.engine === "browser"\n ? await fetchBrowserResource(parsedUrl, options?.timeout_ms)\n : await fetchResource(parsedUrl, options?.timeout_ms, transport, options);' + 'const resource = options?.engine === "browser"\n ? await fetchBrowserResource(parsedUrl, options?.timeout_ms, options?.signal)\n : await fetchResource(parsedUrl, options?.timeout_ms, transport, options);\n options?.signal?.throwIfAborted();' ], [ "fetchCache.set(cacheKey, result);\n return result;", - "result = boundFetchCollections(result);\n fetchCache.set(cacheKey, result);\n return result;" + "result = boundFetchCollections(result);\n options?.signal?.throwIfAborted();\n fetchCache.set(cacheKey, result);\n return result;" ] ]; +const fetchStart = "export async function fetchAndExtract(url, options, transport) {\n const parsedUrl = new URL(url);"; +if (!extractSource.includes(fetchStart)) throw new Error(`mcp-web-search extract patch target not found: ${fetchStart}`); +extractSource = extractSource.replace(fetchStart, "export async function fetchAndExtract(url, options, transport) {\n options?.signal?.throwIfAborted();\n const parsedUrl = new URL(url);"); for (const [before, after] of extractReplacements) { if (!extractSource.includes(before)) throw new Error(`mcp-web-search extract patch target not found: ${before}`); extractSource = extractSource.replace(before, after); diff --git a/scripts/fixtures/web/mock-server.mjs b/scripts/fixtures/web/mock-server.mjs index 617d2a7..87f8121 100644 --- 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) => { `); 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(); diff --git a/scripts/mcp-smoke-client.mjs b/scripts/mcp-smoke-client.mjs index 1b6189f..9edd0c0 100644 --- 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"); diff --git a/scripts/release-check b/scripts/release-check index 99ad07e..f6d6d93 100755 --- 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" diff --git a/scripts/test-lifecycle.sh b/scripts/test-lifecycle.sh index 1ab43bc..8108e7e 100644 --- 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 || 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 "${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 diff --git a/scripts/test-web-search-candidate.mjs b/scripts/test-web-search-candidate.mjs index c836a90..0d7e188 100644 --- 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"); diff --git a/scripts/test-web-search-candidate.sh b/scripts/test-web-search-candidate.sh index ac71e75..4088809 100644 --- 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 diff --git a/scripts/test-web-search-http.mjs b/scripts/test-web-search-http.mjs index aa73c24..a1942a9 100644 --- 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(); diff --git a/scripts/test-web-search-quality.mjs b/scripts/test-web-search-quality.mjs index e7a737b..57e4d6d 100644 --- 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}` })), diff --git a/scripts/test-web-search-stdio-cancellation.mjs b/scripts/test-web-search-stdio-cancellation.mjs new file mode 100644 index 0000000..08fa07f --- /dev/null +++ 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 [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(); +}