"""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 import websearch as _ws 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 import websearch as _ws 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()