219 lines
8.9 KiB
Python
219 lines
8.9 KiB
Python
"""Task-level code-scope extension (search(..., allowed_hosts=[...])).
|
|
|
|
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. Extension values are the caller's own input echoed in
|
|
its own result only -- no test here writes them to a log, and the
|
|
code under test never formats them into error strings.
|
|
"""
|
|
|
|
import json
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from cloud.websearch import (
|
|
CODE_SEARCH_DOMAINS, WebSearchError, _is_code_domain,
|
|
_normalize_allowed_hosts, search)
|
|
from cloud.websearch_mcp import handle_request
|
|
except ImportError:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud.websearch import (
|
|
CODE_SEARCH_DOMAINS, WebSearchError, _is_code_domain,
|
|
_normalize_allowed_hosts, search)
|
|
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",
|
|
}
|
|
|
|
_EXTRA = "docs.example-project.dev"
|
|
|
|
|
|
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 ExtensionTestCase(unittest.TestCase):
|
|
def test_specified_extra_domain_returned_and_native_param_sent(self):
|
|
http = FakeHttp([
|
|
("api.tavily.com", _tavily_body([
|
|
"https://docs.example-project.dev/guide",
|
|
"https://github.com/example/project",
|
|
"https://random-blog.example.com/post",
|
|
]), 200),
|
|
])
|
|
out = search("project guide", provider="tavily", scope="code",
|
|
allowed_hosts=[_EXTRA],
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
payload = http.seen[0][2]
|
|
self.assertIn(_EXTRA, payload["include_domains"])
|
|
self.assertEqual(out["code_domains"],
|
|
list(CODE_SEARCH_DOMAINS) + [_EXTRA])
|
|
self.assertEqual(
|
|
[r["url"] for r in out["results"]],
|
|
["https://docs.example-project.dev/guide",
|
|
"https://github.com/example/project"])
|
|
self.assertEqual(out["attempted"][0]["domain_filtered_out"], 1)
|
|
for secret in _FAKE_KEYS.values():
|
|
self.assertNotIn(secret, json.dumps(out))
|
|
|
|
def test_unspecified_extra_domain_still_filtered_out(self):
|
|
http = FakeHttp([
|
|
("api.tavily.com", _tavily_body([
|
|
"https://docs.example-project.dev/guide",
|
|
"https://github.com/example/project",
|
|
]), 200),
|
|
])
|
|
out = search("project guide", provider="tavily", scope="code",
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
self.assertEqual(out["code_domains"], list(CODE_SEARCH_DOMAINS))
|
|
self.assertEqual([r["url"] for r in out["results"]],
|
|
["https://github.com/example/project"])
|
|
|
|
def test_dot_boundary_positive_and_negative(self):
|
|
domains = list(CODE_SEARCH_DOMAINS) + [_EXTRA]
|
|
self.assertTrue(_is_code_domain(
|
|
"https://api.docs.example-project.dev/v1", domains))
|
|
self.assertTrue(_is_code_domain(
|
|
"https://docs.example-project.dev/", domains))
|
|
self.assertFalse(_is_code_domain(
|
|
"https://evildocs.example-project.dev/", domains))
|
|
self.assertFalse(_is_code_domain(
|
|
"https://docs.example-project.dev.evil.com/", domains))
|
|
# Same boundary rule for a builtin domain, unchanged.
|
|
self.assertFalse(_is_code_domain(
|
|
"https://evilgithub.com/x", domains))
|
|
self.assertTrue(_is_code_domain(
|
|
"https://gist.github.com/x", domains))
|
|
|
|
def test_builtin_entries_are_deduped_not_errors(self):
|
|
self.assertEqual(
|
|
_normalize_allowed_hosts(["github.com", _EXTRA]), [_EXTRA])
|
|
self.assertEqual(
|
|
_normalize_allowed_hosts(["Docs.Example-Project.Dev."]),
|
|
["docs.example-project.dev"])
|
|
|
|
def test_empty_and_illegal_values_raise_before_http(self):
|
|
bad_values = [
|
|
[], [""], [" "], ["*.example.com"],
|
|
["http://example.com"], ["https://example.com/x"],
|
|
["example.com:8080"], ["example.com/path"],
|
|
[".example.com"], ["example.com..dev"], ["exa mple.com"],
|
|
["user@example.com"], [123], [None], ["github.com"],
|
|
]
|
|
for bad in bad_values:
|
|
http = FakeHttp([("api.tavily.com", _tavily_body([]), 200)])
|
|
with self.assertRaises(ValueError, msg=repr(bad)):
|
|
search("q", provider="tavily", scope="code",
|
|
allowed_hosts=bad, env=dict(_FAKE_KEYS),
|
|
http=http)
|
|
self.assertEqual(http.seen, [],
|
|
"invalid %r must fail before any HTTP"
|
|
% (bad,))
|
|
# A bare string is iterable but is not a list of hostnames.
|
|
with self.assertRaises(ValueError):
|
|
search("q", provider="tavily", scope="code",
|
|
allowed_hosts="example.com",
|
|
env=dict(_FAKE_KEYS), http=FakeHttp([]))
|
|
|
|
def test_allowed_hosts_with_general_scope_raises_not_ignored(self):
|
|
http = FakeHttp([("api.tavily.com", _tavily_body(
|
|
["https://docs.example-project.dev/guide"]), 200)])
|
|
with self.assertRaises(ValueError):
|
|
search("q", provider="tavily", scope="general",
|
|
allowed_hosts=[_EXTRA],
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
self.assertEqual(http.seen, [])
|
|
|
|
def test_perplexity_over_cap_skips_native_filter_keeps_postfilter(self):
|
|
extras = ["extra%d.example.com" % i for i in range(6)]
|
|
self.assertGreater(len(CODE_SEARCH_DOMAINS) + len(extras), 20)
|
|
http = FakeHttp([
|
|
("api.perplexity.ai",
|
|
{"choices": [{"message": {"content": "answer"}}],
|
|
"search_results": [
|
|
{"title": "P",
|
|
"url": "https://extra0.example.com/doc",
|
|
"snippet": "extra"},
|
|
{"title": "Ad", "url": "https://shop.example.com/x",
|
|
"snippet": "ad"}]}, 200),
|
|
])
|
|
out = search("q", provider="perplexity", scope="code",
|
|
allowed_hosts=extras,
|
|
env=dict(_FAKE_KEYS), http=http)
|
|
payload = http.seen[0][2]
|
|
self.assertNotIn("search_domain_filter", payload)
|
|
self.assertEqual([r["url"] for r in out["results"]],
|
|
["https://extra0.example.com/doc"])
|
|
self.assertEqual(len(out["code_domains"]), 21)
|
|
|
|
|
|
class McpPassthroughTestCase(unittest.TestCase):
|
|
def test_mcp_forwards_allowed_hosts_when_supplied(self):
|
|
captured = {}
|
|
|
|
def fake_search(query, limit=5, provider="auto", scope="general",
|
|
allowed_hosts=None):
|
|
captured["allowed_hosts"] = allowed_hosts
|
|
return {"provider": "fake", "results": [], "answer": None,
|
|
"attempted": [], "scope": scope}
|
|
|
|
resp = handle_request(
|
|
{"jsonrpc": "2.0", "id": 1, "method": "tools/call",
|
|
"params": {"name": "web_search",
|
|
"arguments": {"query": "hi", "scope": "code",
|
|
"allowed_hosts": [_EXTRA]}}},
|
|
search_fn=fake_search)
|
|
self.assertFalse(resp["result"]["isError"])
|
|
self.assertEqual(captured["allowed_hosts"], [_EXTRA])
|
|
|
|
def test_mcp_without_allowed_hosts_keeps_old_fake_working(self):
|
|
def fake_search(query, limit=5, provider="auto",
|
|
scope="general"):
|
|
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": "hi"}}},
|
|
search_fn=fake_search)
|
|
self.assertFalse(resp["result"]["isError"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|