ajaynomics/context-kit
1
0
Fork 0
Code Issues Pull Requests Releases Activity

Propagate web-search cancellation

Browse Source
This commit is contained in:
Ajay Krishnan
2026-07-25 21:18:08 -07:00
parent 802fc5339e
commit b4efe82ce2
19 changed files with 634 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
docker/web-search/.dockerignore
View File

@@ -2,6 +2,7 @@
!Dockerfile
!http-entrypoint.mjs
!mcp-probe.mjs
!patch-mcp-proxy.py
!patch-mcp-web-search.mjs
!overrides/
!overrides/bing.js

10
docker/web-search/Dockerfile
View File

@@ -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 \

32
docker/web-search/http-entrypoint.mjs
View File

@@ -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;
};

5
docker/web-search/mcp-probe.mjs
View File

@@ -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)}`);

117
docker/web-search/overrides/browser-fetch.js
View File

@@ -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);
}
});
}

20
docker/web-search/overrides/diagnostics.mjs
View File

@@ -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);
}
}

8
docker/web-search/overrides/registry.js
View File

@@ -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 {

168
docker/web-search/patch-mcp-proxy.py Normal file
View File

@@ -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()
""",
)

33
docker/web-search/patch-mcp-web-search.mjs
View File

@@ -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);