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

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 {