host
397a46f850
- search(..., scope=code) restricts to 15 programming domains (official docs, GitHub, registries, Stack Overflow, CVE/NVD): native include_domains/search_domain_filter where supported (Tavily, Exa, Perplexity), post-filtering for all providers (Brave, SearXNG, DuckDuckGo have no native filter). - web_fetch now uses fetch_best: Tavily extract, then Exa contents, then direct fetch; both-fail errors name both failures, direct-fetch 403 under the egress whitelist is fail-closed, not bypassed. Default general scope and direct fetch_url behaviour unchanged. - MCP web_search gains scope enum (general/code); docs updated. - Tests all mock HTTP, no real keys.
359 lines
15 KiB
Python
359 lines
15 KiB
Python
"""Programming search scope (scope="code") + provider extract chain.
|
|
|
|
All HTTP is mocked; no network, no real keys. Fake key values are
|
|
distinctive so a leak into an error string or result is caught by
|
|
assertion.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from cloud import websearch as _ws
|
|
from cloud.websearch import (
|
|
CODE_SEARCH_DOMAINS, WebSearchError, _is_code_domain,
|
|
ensure_codex_home_websearch, fetch_best, fetch_via_provider,
|
|
search)
|
|
from cloud import websearch_mcp as _mcp
|
|
from cloud.websearch_mcp import handle_request
|
|
except ImportError:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud import websearch as _ws
|
|
from cloud.websearch import (
|
|
CODE_SEARCH_DOMAINS, WebSearchError, _is_code_domain,
|
|
ensure_codex_home_websearch, fetch_best, fetch_via_provider,
|
|
search)
|
|
from cloud import websearch_mcp as _mcp
|
|
from cloud.websearch_mcp import handle_request
|
|
|
|
_FAKE_KEYS = {
|
|
"PERPLEXITY_API_KEY": "pplx-SECRET-111",
|
|
"EXA_API_KEY": "exa-SECRET-222",
|
|
"TAVILY_API_KEY": "tvly-SECRET-333",
|
|
"BRAVE_API_KEY": "brave-SECRET-444",
|
|
}
|
|
|
|
|
|
class FakeHttp:
|
|
"""post_json/get keyed by URL substring, in route order."""
|
|
|
|
def __init__(self, routes):
|
|
self.routes = routes
|
|
self.seen = []
|
|
|
|
def post_json(self, url, headers, payload):
|
|
self.seen.append(("POST", url, payload))
|
|
return self._route(url)
|
|
|
|
def get(self, url, headers=None):
|
|
self.seen.append(("GET", url, headers))
|
|
return self._route(url)
|
|
|
|
def _route(self, url):
|
|
for needle, body, status in self.routes:
|
|
if needle in url:
|
|
if status != 200:
|
|
raise WebSearchError("HTTP %s" % status)
|
|
if isinstance(body, dict):
|
|
return 200, json.dumps(body).encode(), \
|
|
"application/json"
|
|
return 200, body.encode(), "text/html"
|
|
raise WebSearchError("HTTP 404")
|
|
|
|
|
|
def _tavily_body(urls):
|
|
return {"results": [{"title": "T %s" % u, "url": u,
|
|
"content": "snippet %s" % u} for u in urls]}
|
|
|
|
|
|
class CodeScopeTestCase(unittest.TestCase):
|
|
def test_tavily_code_sends_include_domains_and_filters(self):
|
|
http = FakeHttp([
|
|
("api.tavily.com", _tavily_body([
|
|
"https://github.com/python/cpython",
|
|
"https://news.example.com/python-312",
|
|
"https://docs.python.org/3/library/pathlib.html",
|
|
]), 200),
|
|
])
|
|
out = search("pathlib changes", provider="tavily",
|
|
scope="code", env=dict(_FAKE_KEYS), http=http)
|
|
payload = http.seen[0][2]
|
|
self.assertEqual(payload["include_domains"],
|
|
list(CODE_SEARCH_DOMAINS))
|
|
self.assertEqual(out["scope"], "code")
|
|
self.assertEqual(out["code_domains"], list(CODE_SEARCH_DOMAINS))
|
|
kept = [r["url"] for r in out["results"]]
|
|
self.assertEqual(kept, ["https://github.com/python/cpython",
|
|
"https://docs.python.org/3/library/"
|
|
"pathlib.html"])
|
|
self.assertEqual(out["attempted"][0]["domain_filtered_out"], 1)
|
|
for secret in _FAKE_KEYS.values():
|
|
self.assertNotIn(secret, json.dumps(out))
|
|
|
|
def test_exa_code_sends_include_domains(self):
|
|
http = FakeHttp([
|
|
("api.exa.ai", {"results": [
|
|
{"title": "E", "url": "https://crates.io/crates/serde",
|
|
"text": "serde", "score": 0.9}]}, 200),
|
|
])
|
|
out = search("serde", provider="exa", scope="code",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
self.assertEqual(http.seen[0][2]["include_domains"],
|
|
list(CODE_SEARCH_DOMAINS))
|
|
self.assertEqual(out["results"][0]["url"],
|
|
"https://crates.io/crates/serde")
|
|
|
|
def test_perplexity_code_sends_domain_filter_within_cap(self):
|
|
self.assertLessEqual(len(CODE_SEARCH_DOMAINS), 20)
|
|
http = FakeHttp([
|
|
("api.perplexity.ai",
|
|
{"choices": [{"message": {"content": "answer"}}],
|
|
"search_results": [
|
|
{"title": "P", "url": "https://go.dev/doc/",
|
|
"snippet": "go docs"},
|
|
{"title": "Ad", "url": "https://shop.example.com/go",
|
|
"snippet": "buy go"}]}, 200),
|
|
])
|
|
out = search("go docs", provider="perplexity", scope="code",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
self.assertEqual(http.seen[0][2]["search_domain_filter"],
|
|
list(CODE_SEARCH_DOMAINS))
|
|
self.assertEqual([r["url"] for r in out["results"]],
|
|
["https://go.dev/doc/"])
|
|
|
|
def test_brave_code_has_no_native_filter_but_is_post_filtered(self):
|
|
http = FakeHttp([
|
|
("api.search.brave.com",
|
|
{"web": {"results": [
|
|
{"title": "B1",
|
|
"url": "https://stackoverflow.com/questions/1",
|
|
"description": "so"},
|
|
{"title": "B2", "url": "https://blog.example.com/x",
|
|
"description": "blog"}]}}, 200),
|
|
])
|
|
out = search("q", provider="brave", scope="code",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
# No native domain parameter was invented for Brave.
|
|
self.assertNotIn("include_domains", http.seen[0][1])
|
|
self.assertNotIn("site:", http.seen[0][1])
|
|
self.assertEqual([r["url"] for r in out["results"]],
|
|
["https://stackoverflow.com/questions/1"])
|
|
|
|
def test_all_filtered_out_continues_chain_then_fails_clearly(self):
|
|
http = FakeHttp([
|
|
("api.search.brave.com",
|
|
{"web": {"results": [
|
|
{"title": "B", "url": "https://blog.example.com/x",
|
|
"description": "blog"}]}}, 200),
|
|
])
|
|
with self.assertRaises(WebSearchError) as ctx:
|
|
search("q", provider="brave", scope="code",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
self.assertIn("code-domain filter removed all 1 results",
|
|
str(ctx.exception))
|
|
|
|
def test_default_scope_unchanged_no_filter_no_post_filter(self):
|
|
http = FakeHttp([
|
|
("api.tavily.com", _tavily_body([
|
|
"https://example.com/general",
|
|
"https://github.com/python/cpython",
|
|
]), 200),
|
|
])
|
|
out = search("q", provider="tavily",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
self.assertEqual(out["scope"], "general")
|
|
self.assertNotIn("code_domains", out)
|
|
self.assertNotIn("include_domains", http.seen[0][2])
|
|
self.assertEqual(len(out["results"]), 2)
|
|
|
|
def test_unknown_scope_raises_before_http(self):
|
|
http = FakeHttp([("api.tavily.com", _tavily_body([]), 200)])
|
|
with self.assertRaises(ValueError):
|
|
search("q", provider="tavily", scope="shopping",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
self.assertEqual(http.seen, [])
|
|
|
|
def test_domain_matching_is_dot_boundary(self):
|
|
self.assertTrue(_is_code_domain("https://gist.github.com/u/1"))
|
|
self.assertTrue(_is_code_domain(
|
|
"https://www.stackoverflow.com/questions/1"))
|
|
self.assertFalse(_is_code_domain("https://evilgithub.com/x"))
|
|
self.assertFalse(_is_code_domain("https://github.com.evil.com/x"))
|
|
self.assertFalse(_is_code_domain("https://notstackoverflow.com/"))
|
|
|
|
def test_missing_key_in_code_scope_reports_unconfigured(self):
|
|
http = FakeHttp([])
|
|
with self.assertRaises(WebSearchError) as ctx:
|
|
search("q", provider="tavily", scope="code",
|
|
env={}, http=http)
|
|
self.assertIn("missing API key (TAVILY_API_KEY unset)",
|
|
str(ctx.exception))
|
|
self.assertEqual(http.seen, [])
|
|
|
|
|
|
class ExtractChainTestCase(unittest.TestCase):
|
|
_PAGE = ("<html><head><title>Direct</title></head><body>"
|
|
"<p>direct page text</p></body></html>")
|
|
|
|
def test_tavily_extract_success(self):
|
|
http = FakeHttp([
|
|
("api.tavily.com/extract",
|
|
{"results": [{"url": "https://docs.python.org/3/",
|
|
"raw_content": "extracted docs text"}],
|
|
"failed_results": []}, 200),
|
|
])
|
|
out = fetch_via_provider("https://docs.python.org/3/",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
self.assertEqual(out["source"], "tavily-extract")
|
|
self.assertEqual(out["text"], "extracted docs text")
|
|
self.assertEqual(http.seen[0][1], "https://api.tavily.com/extract")
|
|
|
|
def test_tavily_fails_then_exa_contents_used(self):
|
|
http = FakeHttp([
|
|
("api.tavily.com/extract", {}, 500),
|
|
("api.exa.ai/contents",
|
|
{"results": [{"url": "https://docs.python.org/3/",
|
|
"text": "exa page text"}]}, 200),
|
|
])
|
|
out = fetch_via_provider("https://docs.python.org/3/",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
self.assertEqual(out["source"], "exa-contents")
|
|
self.assertEqual(out["text"], "exa page text")
|
|
|
|
def test_both_extracts_fail_then_direct_fetch_fallback(self):
|
|
http = FakeHttp([
|
|
("api.tavily.com/extract", {}, 500),
|
|
("api.exa.ai/contents", {}, 500),
|
|
("docs.python.org", self._PAGE, 200),
|
|
])
|
|
out = fetch_best("https://docs.python.org/3/",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
self.assertEqual(out["source"], "direct-fetch")
|
|
self.assertIn("direct page text", out["text"])
|
|
self.assertEqual(
|
|
[a["provider"] for a in out["extract_attempted"]],
|
|
["tavily", "exa"])
|
|
|
|
def test_both_fail_and_direct_fails_reports_both_no_fabrication(self):
|
|
http = FakeHttp([
|
|
("api.tavily.com/extract", {}, 500),
|
|
("api.exa.ai/contents", {}, 500),
|
|
("docs.python.org", {}, 403),
|
|
])
|
|
with self.assertRaises(WebSearchError) as ctx:
|
|
fetch_best("https://docs.python.org/3/",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
msg = str(ctx.exception)
|
|
self.assertIn("provider extraction failed", msg)
|
|
self.assertIn("direct fetch failed", msg)
|
|
for secret in _FAKE_KEYS.values():
|
|
self.assertNotIn(secret, msg)
|
|
|
|
def test_no_keys_extract_errors_but_best_still_direct_fetches(self):
|
|
with self.assertRaises(WebSearchError) as ctx:
|
|
fetch_via_provider("https://docs.python.org/3/",
|
|
env={}, http=FakeHttp([]))
|
|
self.assertIn("no extract provider configured", str(ctx.exception))
|
|
http = FakeHttp([("docs.python.org", self._PAGE, 200)])
|
|
out = fetch_best("https://docs.python.org/3/",
|
|
env={}, http=http)
|
|
self.assertEqual(out["source"], "direct-fetch")
|
|
|
|
def test_localhost_blocked_before_any_provider_call(self):
|
|
http = FakeHttp([])
|
|
with self.assertRaises(WebSearchError):
|
|
fetch_best("http://127.0.0.1/admin",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
with self.assertRaises(WebSearchError):
|
|
fetch_via_provider("http://10.0.0.5/admin",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
self.assertEqual(http.seen, [])
|
|
|
|
|
|
class McpScopeTestCase(unittest.TestCase):
|
|
def test_tools_list_exposes_scope_enum(self):
|
|
resp = handle_request(
|
|
{"jsonrpc": "2.0", "id": 1, "method": "tools/list"})
|
|
tool = [t for t in resp["result"]["tools"]
|
|
if t["name"] == "web_search"][0]
|
|
scope = tool["inputSchema"]["properties"]["scope"]
|
|
self.assertEqual(scope["enum"], ["general", "code"])
|
|
self.assertEqual(scope["default"], "general")
|
|
self.assertIn("shopping", tool["description"])
|
|
fetch_tool = [t for t in resp["result"]["tools"]
|
|
if t["name"] == "web_fetch"][0]
|
|
self.assertIn("Tavily extract", fetch_tool["description"])
|
|
|
|
def test_tools_call_passes_scope_through(self):
|
|
seen = {}
|
|
|
|
def fake_search(query, limit=5, provider="auto", scope="general"):
|
|
seen.update(query=query, scope=scope)
|
|
return {"provider": "fake", "results": [], "answer": None,
|
|
"attempted": [], "scope": scope}
|
|
|
|
resp = handle_request(
|
|
{"jsonrpc": "2.0", "id": 2, "method": "tools/call",
|
|
"params": {"name": "web_search",
|
|
"arguments": {"query": "serde",
|
|
"scope": "code"}}},
|
|
search_fn=fake_search)
|
|
self.assertFalse(resp["result"]["isError"])
|
|
self.assertEqual(seen["scope"], "code")
|
|
|
|
def test_web_fetch_default_is_extract_first_chain(self):
|
|
# Wiring assertion: the MCP default fetch is fetch_best, not
|
|
# the direct-only fetch_url.
|
|
self.assertIs(_mcp.fetch_best, _ws.fetch_best)
|
|
|
|
seen = {}
|
|
|
|
def fake_fetch(url, char_limit=15000):
|
|
seen.update(url=url, char_limit=char_limit)
|
|
return {"url": url, "text": "t", "source": "fake"}
|
|
|
|
resp = handle_request(
|
|
{"jsonrpc": "2.0", "id": 3, "method": "tools/call",
|
|
"params": {"name": "web_fetch",
|
|
"arguments": {"url": "https://example.com/x"}}},
|
|
fetch_fn=fake_fetch)
|
|
self.assertFalse(resp["result"]["isError"])
|
|
self.assertEqual(seen["url"], "https://example.com/x")
|
|
|
|
|
|
class SeededScriptTestCase(unittest.TestCase):
|
|
def test_seeded_flat_copy_exposes_code_scope(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
home = os.path.join(tmp, "codex-home")
|
|
os.makedirs(home)
|
|
out = ensure_codex_home_websearch(
|
|
home, command_args=["python3", "/tmp/ws/websearch_mcp.py"])
|
|
self.assertTrue(out["config_changed"])
|
|
# Run the seeded flat copy standalone (no repo cloud
|
|
# package on its path, exactly like a tenant codex-home).
|
|
probe = (
|
|
"import json, websearch_mcp\n"
|
|
"resp = websearch_mcp.handle_request(\n"
|
|
" {'jsonrpc': '2.0', 'id': 1, 'method': 'tools/list'})\n"
|
|
"tool = [t for t in resp['result']['tools']\n"
|
|
" if t['name'] == 'web_search'][0]\n"
|
|
"print(json.dumps(tool['inputSchema']['properties']"
|
|
"['scope']))\n")
|
|
env = {"PATH": os.environ.get("PATH", ""),
|
|
"PYTHONPATH": home}
|
|
proc = subprocess.run(
|
|
[sys.executable, "-c", probe], cwd=home, env=env,
|
|
capture_output=True, text=True, timeout=60)
|
|
self.assertEqual(proc.returncode, 0, proc.stderr)
|
|
scope = json.loads(proc.stdout.strip().splitlines()[-1])
|
|
self.assertEqual(scope["enum"], ["general", "code"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|