feat: web search MCP v1 (oh-my-pi pattern, Python stdlib)
web_search auto chain Perplexity/Exa/Tavily/Brave/SearXNG/DuckDuckGo with pinned-provider no-fallback, web_fetch with head+tail truncation, codex-home seeding helper (idempotent, never writes keys). Tests: cloud/tests/test_websearch_mcp.py 12/12 (mocked HTTP). Not wired into worker, not deployed, paid backends not live-tested.
This commit is contained in:
host committed
1 parent
5026e75b2c
commit
f8acdffb4a
5 files changed
+1194
No files matched your search
@@ -0,0 +1,253 @@
|
||||
"""Web search MCP v1 tests: chain fallback, fetch, MCP framing, config.
|
||||
|
||||
No network, no keys: every HTTP call goes through a FakeHttp, and
|
||||
the fake key values are deliberately distinctive so a leak into an
|
||||
error string or a tool result would be caught by assertion, not by
|
||||
eyeballing.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from cloud.websearch import (
|
||||
WebSearchError, ensure_codex_home_websearch, fetch_url, search)
|
||||
from cloud.websearch_mcp import handle_request
|
||||
except ImportError:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.websearch import (
|
||||
WebSearchError, ensure_codex_home_websearch, fetch_url, 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",
|
||||
}
|
||||
|
||||
|
||||
class FakeHttp:
|
||||
"""post_json/get keyed by URL substring; failures raise like HTTP."""
|
||||
|
||||
def __init__(self, routes):
|
||||
# routes: list of (url_substring, body_dict_or_text, status)
|
||||
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:
|
||||
# Mimic HttpClient._read: raise with status only.
|
||||
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 _exa_body():
|
||||
return {"results": [
|
||||
{"title": "Exa Hit", "url": "https://example.com/exa",
|
||||
"text": "exa snippet", "score": 0.9,
|
||||
"highlights": ["exa highlight"]}]}
|
||||
|
||||
|
||||
class ChainTestCase(unittest.TestCase):
|
||||
def test_fallback_perplexity_to_exa_and_key_not_leaked(self):
|
||||
http = FakeHttp([
|
||||
("perplexity", {}, 500),
|
||||
("api.exa.ai", _exa_body(), 200),
|
||||
])
|
||||
out = search("test query", provider="auto",
|
||||
env=dict(_FAKE_KEYS), http=http)
|
||||
self.assertEqual(out["provider"], "exa")
|
||||
self.assertEqual(out["results"][0]["url"],
|
||||
"https://example.com/exa")
|
||||
self.assertEqual(out["attempted"][0]["provider"], "perplexity")
|
||||
self.assertFalse(out["attempted"][0]["ok"])
|
||||
blob = json.dumps(out)
|
||||
for secret in _FAKE_KEYS.values():
|
||||
self.assertNotIn(secret, blob)
|
||||
|
||||
def test_all_fail_error_lists_attempts_without_keys(self):
|
||||
http = FakeHttp([("perplexity", {}, 401),
|
||||
("api.exa.ai", {}, 429),
|
||||
("api.tavily.com", {}, 500),
|
||||
("api.search.brave.com", {}, 403)])
|
||||
with self.assertRaises(WebSearchError) as ctx:
|
||||
search("q", provider="perplexity,exa,tavily,brave",
|
||||
env=dict(_FAKE_KEYS), http=http)
|
||||
for secret in _FAKE_KEYS.values():
|
||||
self.assertNotIn(secret, str(ctx.exception))
|
||||
self.assertEqual(
|
||||
[a["provider"] for a in ctx.exception.attempted],
|
||||
["perplexity", "exa", "tavily", "brave"])
|
||||
|
||||
def test_pinned_provider_does_not_fall_back(self):
|
||||
http = FakeHttp([("perplexity", {}, 500),
|
||||
("api.exa.ai", _exa_body(), 200)])
|
||||
with self.assertRaises(WebSearchError):
|
||||
search("q", provider="perplexity",
|
||||
env=dict(_FAKE_KEYS), http=http)
|
||||
# Exa was never even attempted.
|
||||
self.assertFalse(any("exa" in url for _m, url, _p in http.seen))
|
||||
|
||||
def test_brave_mapping_and_tavily_mapping(self):
|
||||
http = FakeHttp([
|
||||
("api.search.brave.com",
|
||||
{"web": {"results": [
|
||||
{"title": "Brave Hit", "url": "https://example.com/b",
|
||||
"description": "desc",
|
||||
"extra_snippets": ["extra"]}]}}, 200),
|
||||
])
|
||||
out = search("q", provider="brave",
|
||||
env=dict(_FAKE_KEYS), http=http)
|
||||
self.assertIn("desc", out["results"][0]["snippet"])
|
||||
self.assertIn("extra", out["results"][0]["snippet"])
|
||||
|
||||
http2 = FakeHttp([
|
||||
("api.tavily.com",
|
||||
{"results": [{"title": "T", "url": "https://example.com/t",
|
||||
"content": "tavily content", "score": 0.5}]},
|
||||
200),
|
||||
])
|
||||
out2 = search("q", provider="tavily",
|
||||
env=dict(_FAKE_KEYS), http=http2)
|
||||
self.assertEqual(out2["results"][0]["snippet"], "tavily content")
|
||||
# Tavily request must not ask for a synthesised answer: the
|
||||
# tool returns results, the model synthesises.
|
||||
self.assertFalse(http2.seen[0][2]["include_answer"])
|
||||
|
||||
def test_perplexity_answer_and_citations(self):
|
||||
http = FakeHttp([
|
||||
("api.perplexity.ai",
|
||||
{"choices": [{"message": {"content": "synthesised answer"}}],
|
||||
"search_results": [
|
||||
{"title": "P", "url": "https://example.com/p",
|
||||
"snippet": "p snippet"}]}, 200),
|
||||
])
|
||||
out = search("q", provider="perplexity",
|
||||
env=dict(_FAKE_KEYS), http=http)
|
||||
self.assertEqual(out["answer"], "synthesised answer")
|
||||
self.assertEqual(out["results"][0]["url"], "https://example.com/p")
|
||||
|
||||
def test_empty_results_continue_chain(self):
|
||||
http = FakeHttp([
|
||||
("api.exa.ai", {"results": []}, 200),
|
||||
("api.tavily.com",
|
||||
{"results": [{"title": "T", "url": "https://example.com/t",
|
||||
"content": "c"}]}, 200),
|
||||
])
|
||||
out = search("q", provider="exa,tavily",
|
||||
env=dict(_FAKE_KEYS), http=http)
|
||||
self.assertEqual(out["provider"], "tavily")
|
||||
|
||||
def test_duckduckgo_html_parser(self):
|
||||
html = (
|
||||
'<a class="result__a" href="//duckduckgo.com/l/?uddg='
|
||||
'https%3A%2F%2Fexample.com%2Fddg">DDG Title</a>'
|
||||
'<a class="result__snippet">ddg snippet</a>')
|
||||
http = FakeHttp([("duckduckgo", html, 200)])
|
||||
out = search("q", provider="duckduckgo", env={}, http=http)
|
||||
self.assertEqual(out["results"][0]["url"], "https://example.com/ddg")
|
||||
self.assertEqual(out["results"][0]["title"], "DDG Title")
|
||||
|
||||
|
||||
class FetchTestCase(unittest.TestCase):
|
||||
def test_strips_script_and_truncates_head_tail(self):
|
||||
body = ("<html><head><title>My Page</title></head><body>"
|
||||
"<script>var secret = 1;</script>"
|
||||
"<p>" + "A" * 10000 + "</p>"
|
||||
"<p>" + "B" * 10000 + "</p></body></html>")
|
||||
http = FakeHttp([("example.com", body, 200)])
|
||||
out = fetch_url("https://example.com/page", char_limit=2000,
|
||||
http=http)
|
||||
self.assertEqual(out["title"], "My Page")
|
||||
self.assertNotIn("secret", out["text"])
|
||||
self.assertTrue(out["truncated"])
|
||||
self.assertIn("[TRUNCATED]", out["text"])
|
||||
self.assertIn("B" * 10, out["text"])
|
||||
|
||||
def test_rejects_non_http_and_localhost(self):
|
||||
http = FakeHttp([])
|
||||
with self.assertRaises(WebSearchError):
|
||||
fetch_url("file:///etc/passwd", http=http)
|
||||
with self.assertRaises(WebSearchError):
|
||||
fetch_url("http://127.0.0.1/admin", http=http)
|
||||
with self.assertRaises(WebSearchError):
|
||||
fetch_url("http://10.0.0.5/admin", http=http)
|
||||
|
||||
|
||||
class McpFramingTestCase(unittest.TestCase):
|
||||
def test_initialize_tools_list_and_call(self):
|
||||
resp = handle_request(
|
||||
{"jsonrpc": "2.0", "id": 1, "method": "initialize",
|
||||
"params": {"protocolVersion": "2024-11-05"}})
|
||||
self.assertEqual(resp["result"]["serverInfo"]["name"], "websearch")
|
||||
resp = handle_request(
|
||||
{"jsonrpc": "2.0", "id": 2, "method": "tools/list"})
|
||||
names = [t["name"] for t in resp["result"]["tools"]]
|
||||
self.assertEqual(names, ["web_search", "web_fetch"])
|
||||
|
||||
def fake_search(query, limit=5, provider="auto"):
|
||||
return {"provider": "fake", "results": [], "answer": None,
|
||||
"attempted": []}
|
||||
|
||||
resp = handle_request(
|
||||
{"jsonrpc": "2.0", "id": 3, "method": "tools/call",
|
||||
"params": {"name": "web_search",
|
||||
"arguments": {"query": "hi"}}},
|
||||
search_fn=fake_search)
|
||||
self.assertFalse(resp["result"]["isError"])
|
||||
self.assertIn("fake", resp["result"]["content"][0]["text"])
|
||||
|
||||
def test_notification_gets_no_response_and_unknown_tool_errors(self):
|
||||
self.assertIsNone(handle_request(
|
||||
{"jsonrpc": "2.0", "method": "notifications/initialized"}))
|
||||
resp = handle_request(
|
||||
{"jsonrpc": "2.0", "id": 9, "method": "tools/call",
|
||||
"params": {"name": "nope", "arguments": {}}})
|
||||
self.assertEqual(resp["error"]["code"], -32602)
|
||||
|
||||
|
||||
class CodexHomeTestCase(unittest.TestCase):
|
||||
def test_merge_is_idempotent_and_preserves_existing(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
home = os.path.join(tmp, "codex-home")
|
||||
os.makedirs(home)
|
||||
existing = ("# tenant config\n"
|
||||
"[model_providers.relay]\n"
|
||||
"api_key = \"keep-me\"\n")
|
||||
with open(os.path.join(home, "config.toml"), "w") as fh:
|
||||
fh.write(existing)
|
||||
out = ensure_codex_home_websearch(
|
||||
home, command_args=["python3", "/tmp/ws/websearch_mcp.py"])
|
||||
self.assertTrue(out["config_changed"])
|
||||
with open(out["config_path"]) as fh:
|
||||
merged = fh.read()
|
||||
self.assertIn("keep-me", merged)
|
||||
self.assertIn("[mcp_servers.websearch]", merged)
|
||||
out2 = ensure_codex_home_websearch(
|
||||
home, command_args=["python3", "/tmp/ws/websearch_mcp.py"])
|
||||
self.assertFalse(out2["config_changed"])
|
||||
with open(out2["config_path"]) as fh:
|
||||
self.assertEqual(fh.read().count("[mcp_servers.websearch]"),
|
||||
1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,635 @@
|
||||
"""Web search + fetch for codex tasks, following the oh-my-pi pattern.
|
||||
|
||||
Zero third-party dependencies; Python standard library only.
|
||||
|
||||
Design notes (what is copied from oh-my-pi, and what is not):
|
||||
- Copied: one unified ``web_search`` tool, a ``SearchProvider`` abstraction,
|
||||
and an ``auto`` chain that walks providers in order and falls back to
|
||||
the next one on failure. Pinning a single provider by name disables
|
||||
fallback, exactly like oh-my-pi's pinned provider.
|
||||
- Copied: a two-step workflow -- ``web_search`` finds candidate URLs,
|
||||
``web_fetch`` reads a specific page -- and site-agnostic extraction
|
||||
that returns clean text plus a link list, not raw HTML.
|
||||
- NOT copied: oh-my-pi's 23 providers. v1 ships the five that matter
|
||||
for a codex task: Perplexity (synthesised answer + citations),
|
||||
Exa (semantic search + highlights), Tavily (research + extract),
|
||||
Brave (fast, independent index) and two keyless fallbacks
|
||||
(SearXNG, DuckDuckGo). Model-native OAuth search providers
|
||||
(Gemini / Anthropic / Codex / xAI) are deliberately excluded: inside
|
||||
a codex task they add no provider the task does not already have.
|
||||
- NOT copied: oh-my-pi's TypeScript / Rust extraction stack. Fetch
|
||||
here is a small stdlib HTML parser: script/style stripped, text
|
||||
collapsed, long pages truncated head+tail like Hermes' deterministic
|
||||
budget (default 15000 chars, 75% head / 25% tail, ``[TRUNCATED]``).
|
||||
|
||||
Security notes:
|
||||
- API keys come only from the process environment, never from
|
||||
arguments, config files, or this module's return values. Error
|
||||
strings contain provider names and HTTP statuses only; they never
|
||||
format headers or key values, and the chain test asserts this.
|
||||
- ``web_fetch`` is not a security boundary: a codex task running
|
||||
danger-full-access can already curl anything, so fetch does not
|
||||
pretend otherwise. It does refuse non-http(s) URLs and literal
|
||||
localhost / private-IP URLs, which stops the most trivial SSRF
|
||||
accidents, but it is not a sandbox.
|
||||
"""
|
||||
|
||||
import html as _html
|
||||
import ipaddress
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from html.parser import HTMLParser
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_CHAIN",
|
||||
"WebSearchError",
|
||||
"SearchResult",
|
||||
"available_providers",
|
||||
"resolve_provider_chain",
|
||||
"search",
|
||||
"fetch_url",
|
||||
"ensure_codex_home_websearch",
|
||||
"DEFAULT_FETCH_CHAR_LIMIT",
|
||||
]
|
||||
|
||||
DEFAULT_CHAIN = ("perplexity", "exa", "tavily", "brave", "searxng",
|
||||
"duckduckgo")
|
||||
DEFAULT_FETCH_CHAR_LIMIT = 15000
|
||||
_USER_AGENT = "codex-websearch-mcp/0.1 (+codex-app-server)"
|
||||
|
||||
_KEY_ENV = {
|
||||
"perplexity": "PERPLEXITY_API_KEY",
|
||||
"exa": "EXA_API_KEY",
|
||||
"tavily": "TAVILY_API_KEY",
|
||||
"brave": "BRAVE_API_KEY",
|
||||
}
|
||||
_FREE_PROVIDERS = ("searxng", "duckduckgo")
|
||||
|
||||
|
||||
class WebSearchError(Exception):
|
||||
"""All providers in the chain failed (or the pinned one did)."""
|
||||
|
||||
def __init__(self, message, attempted=None):
|
||||
super().__init__(message)
|
||||
self.attempted = list(attempted or [])
|
||||
|
||||
|
||||
class SearchResult(dict):
|
||||
"""A single search hit: title, url, snippet, optional score."""
|
||||
|
||||
def __init__(self, title, url, snippet="", score=None):
|
||||
super().__init__(title=title or "", url=url or "",
|
||||
snippet=snippet or "")
|
||||
if score is not None:
|
||||
self["score"] = score
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- HTTP
|
||||
|
||||
class HttpClient:
|
||||
"""Thin urllib wrapper, injectable so tests never touch the net."""
|
||||
|
||||
def __init__(self, timeout=20):
|
||||
self.timeout = timeout
|
||||
|
||||
def post_json(self, url, headers, payload):
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
req = urllib.request.Request(
|
||||
url, data=data,
|
||||
headers=dict(headers, **{"Content-Type": "application/json"}),
|
||||
method="POST")
|
||||
return self._read(req)
|
||||
|
||||
def get(self, url, headers=None):
|
||||
req = urllib.request.Request(
|
||||
url, headers=headers or {}, method="GET")
|
||||
return self._read(req)
|
||||
|
||||
def _read(self, req):
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=self.timeout) as resp:
|
||||
body = resp.read()
|
||||
return resp.status, body, resp.headers.get_content_type()
|
||||
except urllib.error.HTTPError as exc:
|
||||
# Never include headers (keys live there) in the error.
|
||||
body = exc.read()[:500]
|
||||
raise WebSearchError(
|
||||
"HTTP %s" % exc.code) from _ProviderHttpError(
|
||||
exc.code, body)
|
||||
|
||||
def get_text(self, url, headers=None):
|
||||
status, body, _ctype = self.get(url, headers)
|
||||
charset = "utf-8"
|
||||
try:
|
||||
text = body.decode(charset, errors="replace")
|
||||
except Exception:
|
||||
text = body.decode("utf-8", errors="replace")
|
||||
return status, text
|
||||
|
||||
|
||||
class _ProviderHttpError(Exception):
|
||||
def __init__(self, status, body):
|
||||
super().__init__("HTTP %s" % status)
|
||||
self.status = status
|
||||
self.body = body
|
||||
|
||||
|
||||
def _http_error_message(exc):
|
||||
"""Render a provider failure without ever leaking credentials."""
|
||||
if isinstance(exc, _ProviderHttpError):
|
||||
return "HTTP %s" % exc.status
|
||||
if isinstance(exc, WebSearchError):
|
||||
return str(exc)
|
||||
return "%s: %s" % (type(exc).__name__, exc)
|
||||
|
||||
|
||||
def _default_headers(env, provider):
|
||||
key_env = _KEY_ENV.get(provider)
|
||||
key = env.get(key_env) if key_env else None
|
||||
if provider == "exa":
|
||||
return {"x-api-key": key, "User-Agent": _USER_AGENT}
|
||||
if provider == "tavily":
|
||||
return {"Authorization": "Bearer %s" % key,
|
||||
"User-Agent": _USER_AGENT}
|
||||
if provider == "brave":
|
||||
return {"X-Subscription-Token": key, "Accept": "application/json",
|
||||
"User-Agent": _USER_AGENT}
|
||||
if provider == "perplexity":
|
||||
return {"Authorization": "Bearer %s" % key,
|
||||
"User-Agent": _USER_AGENT}
|
||||
return {"User-Agent": _USER_AGENT}
|
||||
|
||||
|
||||
def _clamp_limit(limit, default=5, lo=1, hi=10):
|
||||
try:
|
||||
value = int(limit)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
return max(lo, min(hi, value))
|
||||
|
||||
|
||||
# ------------------------------------------------------------ providers
|
||||
|
||||
def _provider_perplexity(query, limit, env, http):
|
||||
# Sonar is a synthesis API: one OpenAI-shaped chat call returns a
|
||||
# written, cited answer. We surface both the answer and the
|
||||
# underlying sources, so web_search keeps its "answer + citations"
|
||||
# shape instead of degrading to a link list.
|
||||
_status, body, _ctype = http.post_json(
|
||||
"https://api.perplexity.ai/chat/completions",
|
||||
_default_headers(env, "perplexity"),
|
||||
{"model": "sonar",
|
||||
"messages": [{"role": "user", "content": query}]})
|
||||
data = json.loads(body.decode("utf-8", errors="replace"))
|
||||
choices = data.get("choices") or []
|
||||
message = (choices[0].get("message") or {}) if choices else {}
|
||||
answer = message.get("content") or None
|
||||
results = []
|
||||
for item in data.get("search_results") or []:
|
||||
if isinstance(item, dict) and item.get("url"):
|
||||
results.append(SearchResult(item.get("title") or item["url"],
|
||||
item["url"],
|
||||
item.get("snippet") or ""))
|
||||
if not results:
|
||||
for url in data.get("citations") or []:
|
||||
if isinstance(url, str) and url.startswith("http"):
|
||||
results.append(SearchResult(url, url, ""))
|
||||
return results, answer
|
||||
|
||||
|
||||
def _provider_exa(query, limit, env, http):
|
||||
_status, body, _ctype = http.post_json(
|
||||
"https://api.exa.ai/search",
|
||||
_default_headers(env, "exa"),
|
||||
{"query": query, "numResults": limit, "type": "auto",
|
||||
"contents": {"text": {"maxCharacters": 2000},
|
||||
"highlights": True}})
|
||||
data = json.loads(body.decode("utf-8", errors="replace"))
|
||||
results = []
|
||||
for item in data.get("results") or []:
|
||||
if not isinstance(item, dict) or not item.get("url"):
|
||||
continue
|
||||
highlights = item.get("highlights") or []
|
||||
snippet = item.get("text") or ""
|
||||
if highlights and not snippet:
|
||||
snippet = " ".join(str(h) for h in highlights[:3])
|
||||
elif highlights:
|
||||
snippet = snippet[:2000]
|
||||
results.append(SearchResult(item.get("title") or item["url"],
|
||||
item["url"], snippet,
|
||||
item.get("score")))
|
||||
return results, None
|
||||
|
||||
|
||||
def _provider_tavily(query, limit, env, http):
|
||||
_status, body, _ctype = http.post_json(
|
||||
"https://api.tavily.com/search",
|
||||
_default_headers(env, "tavily"),
|
||||
{"query": query, "search_depth": "basic",
|
||||
"max_results": limit, "include_answer": False,
|
||||
"include_raw_content": False})
|
||||
data = json.loads(body.decode("utf-8", errors="replace"))
|
||||
results = []
|
||||
for item in data.get("results") or []:
|
||||
if not isinstance(item, dict) or not item.get("url"):
|
||||
continue
|
||||
results.append(SearchResult(item.get("title") or item["url"],
|
||||
item["url"],
|
||||
item.get("content") or "",
|
||||
item.get("score")))
|
||||
return results, data.get("answer") or None
|
||||
|
||||
|
||||
def _provider_brave(query, limit, env, http):
|
||||
url = ("https://api.search.brave.com/res/v1/web/search?"
|
||||
+ urllib.parse.urlencode(
|
||||
{"q": query, "count": limit, "extra_snippets": "true"}))
|
||||
_status, body, _ctype = http.get(url, _default_headers(env, "brave"))
|
||||
data = json.loads(body.decode("utf-8", errors="replace"))
|
||||
results = []
|
||||
for item in (data.get("web") or {}).get("results") or []:
|
||||
if not isinstance(item, dict) or not item.get("url"):
|
||||
continue
|
||||
parts = [item.get("description") or ""]
|
||||
parts.extend(str(s) for s in (item.get("extra_snippets") or []))
|
||||
snippet = " ".join(p for p in parts if p)
|
||||
results.append(SearchResult(item.get("title") or item["url"],
|
||||
item["url"], snippet))
|
||||
return results, None
|
||||
|
||||
|
||||
def _provider_searxng(query, limit, env, http):
|
||||
base = (env.get("SEARXNG_URL") or "").rstrip("/")
|
||||
if not base:
|
||||
raise WebSearchError("searxng not configured (SEARXNG_URL unset)")
|
||||
url = (base + "/search?"
|
||||
+ urllib.parse.urlencode({"q": query, "format": "json"}))
|
||||
_status, body, _ctype = http.get(url, _default_headers(env, "searxng"))
|
||||
data = json.loads(body.decode("utf-8", errors="replace"))
|
||||
results = []
|
||||
for item in data.get("results") or []:
|
||||
if not isinstance(item, dict) or not item.get("url"):
|
||||
continue
|
||||
results.append(SearchResult(item.get("title") or item["url"],
|
||||
item["url"],
|
||||
item.get("content") or ""))
|
||||
return results[:limit], None
|
||||
|
||||
|
||||
class _DuckDuckGoParser(HTMLParser):
|
||||
"""Best-effort parser for DuckDuckGo's HTML results page.
|
||||
|
||||
Only the last-resort free path uses this; it is deliberately
|
||||
tolerant of layout drift (missing classes just yield no results,
|
||||
which moves the chain on, it never crashes the tool).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.results = []
|
||||
self._current = None
|
||||
self._in_title = False
|
||||
self._in_snippet = False
|
||||
|
||||
def handle_starttag(self, tag, attrs):
|
||||
attrs = dict(attrs)
|
||||
cls = attrs.get("class") or ""
|
||||
if tag == "a" and "result__a" in cls:
|
||||
href = attrs.get("href") or ""
|
||||
self._current = {"title": "", "url": _unwrap_ddg(href),
|
||||
"snippet": ""}
|
||||
self._in_title = True
|
||||
elif tag == "a" and "result__snippet" in cls:
|
||||
self._in_snippet = True
|
||||
|
||||
def handle_endtag(self, tag):
|
||||
if tag == "a" and self._in_title:
|
||||
self._in_title = False
|
||||
if self._current and self._current["url"]:
|
||||
self.results.append(self._current)
|
||||
self._current = None
|
||||
if tag == "a" and self._in_snippet:
|
||||
self._in_snippet = False
|
||||
|
||||
def handle_data(self, data):
|
||||
if self._in_title and self._current is not None:
|
||||
self._current["title"] += data
|
||||
elif self._in_snippet and self.results:
|
||||
self.results[-1]["snippet"] += data
|
||||
|
||||
|
||||
def _unwrap_ddg(href):
|
||||
"""DuckDuckGo wraps result links as //duckduckgo.com/l/?uddg=<url>."""
|
||||
if not href:
|
||||
return ""
|
||||
parsed = urllib.parse.urlparse(href)
|
||||
if "duckduckgo.com" in (parsed.netloc or ""):
|
||||
target = urllib.parse.parse_qs(parsed.query).get("uddg")
|
||||
if target:
|
||||
return target[0]
|
||||
return ""
|
||||
if href.startswith("//"):
|
||||
return "https:" + href
|
||||
return href
|
||||
|
||||
|
||||
def _provider_duckduckgo(query, limit, env, http):
|
||||
url = ("https://html.duckduckgo.com/html/?"
|
||||
+ urllib.parse.urlencode({"q": query}))
|
||||
_status, body, _ctype = http.get(
|
||||
url, {"User-Agent": _USER_AGENT})
|
||||
parser = _DuckDuckGoParser()
|
||||
parser.feed(body.decode("utf-8", errors="replace"))
|
||||
results = []
|
||||
for item in parser.results[:limit]:
|
||||
results.append(SearchResult(item["title"].strip(),
|
||||
item["url"],
|
||||
item["snippet"].strip()))
|
||||
return results, None
|
||||
|
||||
|
||||
_PROVIDERS = {
|
||||
"perplexity": _provider_perplexity,
|
||||
"exa": _provider_exa,
|
||||
"tavily": _provider_tavily,
|
||||
"brave": _provider_brave,
|
||||
"searxng": _provider_searxng,
|
||||
"duckduckgo": _provider_duckduckgo,
|
||||
}
|
||||
|
||||
|
||||
def available_providers(env=None):
|
||||
"""Providers usable right now: keyed ones with a key, plus free ones."""
|
||||
env = env if env is not None else os.environ
|
||||
out = []
|
||||
for name in DEFAULT_CHAIN:
|
||||
if name in _FREE_PROVIDERS:
|
||||
if name == "searxng" and not env.get("SEARXNG_URL"):
|
||||
continue
|
||||
out.append(name)
|
||||
elif env.get(_KEY_ENV[name]):
|
||||
out.append(name)
|
||||
return out
|
||||
|
||||
|
||||
def resolve_provider_chain(provider="auto", env=None):
|
||||
"""Resolve the chain oh-my-pi style: auto walks, pinned does not.
|
||||
|
||||
``provider`` may be "auto" (default chain, filtered to what is
|
||||
configured), a single name (that provider only, even if it fails),
|
||||
or a comma-separated list (that order only). Unknown names raise
|
||||
ValueError before any HTTP happens.
|
||||
"""
|
||||
env = env if env is not None else os.environ
|
||||
if provider is None or str(provider).strip() in ("", "auto"):
|
||||
chain = list(available_providers(env))
|
||||
if not chain:
|
||||
chain = list(_FREE_PROVIDERS)
|
||||
# duckduckgo needs no configuration; searxng without a URL
|
||||
# will fail fast with a clear message, which is better
|
||||
# than pretending no provider exists at all.
|
||||
chain = ["duckduckgo"]
|
||||
return chain
|
||||
names = [p.strip() for p in str(provider).split(",") if p.strip()]
|
||||
for name in names:
|
||||
if name not in _PROVIDERS:
|
||||
raise ValueError("unknown search provider: %r" % name)
|
||||
return names
|
||||
|
||||
|
||||
def search(query, limit=5, provider="auto", env=None, http=None):
|
||||
"""Search the web. Returns {provider, results, answer, attempted}.
|
||||
|
||||
Raises WebSearchError (with .attempted) when every provider in
|
||||
the chain fails or returns no results. Error strings never contain
|
||||
API keys; see the module docstring.
|
||||
"""
|
||||
if not query or not str(query).strip():
|
||||
raise ValueError("query must be a non-empty string")
|
||||
env = env if env is not None else os.environ
|
||||
http = http or HttpClient(
|
||||
timeout=int(env.get("WEBSEARCH_TIMEOUT", "20") or 20))
|
||||
limit = _clamp_limit(limit)
|
||||
chain = resolve_provider_chain(provider, env)
|
||||
attempted = []
|
||||
for name in chain:
|
||||
if name in _KEY_ENV and not env.get(_KEY_ENV[name]):
|
||||
attempted.append({"provider": name, "ok": False,
|
||||
"error": "missing API key (%s unset)"
|
||||
% _KEY_ENV[name]})
|
||||
continue
|
||||
try:
|
||||
results, answer = _PROVIDERS[name](
|
||||
str(query), limit, env, http)
|
||||
except Exception as exc: # noqa: BLE001 - chain must continue
|
||||
attempted.append({"provider": name, "ok": False,
|
||||
"error": _http_error_message(exc)})
|
||||
continue
|
||||
if not results:
|
||||
attempted.append({"provider": name, "ok": False,
|
||||
"error": "no results"})
|
||||
continue
|
||||
attempted.append({"provider": name, "ok": True, "error": None})
|
||||
return {"provider": name,
|
||||
"results": [dict(r) for r in results[:limit]],
|
||||
"answer": answer,
|
||||
"attempted": attempted}
|
||||
raise WebSearchError(
|
||||
"all search providers failed: %s"
|
||||
% "; ".join("%s (%s)" % (a["provider"], a["error"])
|
||||
for a in attempted),
|
||||
attempted=attempted)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- fetch
|
||||
|
||||
class _TextExtractor(HTMLParser):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.title = ""
|
||||
self._in_title = False
|
||||
self._skip = 0
|
||||
self.blocks = []
|
||||
self.links = []
|
||||
self._current_block = []
|
||||
|
||||
def handle_starttag(self, tag, attrs):
|
||||
if tag in ("script", "style", "noscript"):
|
||||
self._skip += 1
|
||||
return
|
||||
if tag == "title":
|
||||
self._in_title = True
|
||||
if tag == "a":
|
||||
href = dict(attrs).get("href") or ""
|
||||
if href.startswith("http"):
|
||||
self.links.append(href)
|
||||
if tag in ("p", "h1", "h2", "h3", "h4", "li", "pre", "blockquote"):
|
||||
self._flush()
|
||||
|
||||
def handle_endtag(self, tag):
|
||||
if tag in ("script", "style", "noscript"):
|
||||
self._skip = max(0, self._skip - 1)
|
||||
if tag == "title":
|
||||
self._in_title = False
|
||||
if tag in ("p", "h1", "h2", "h3", "h4", "li", "pre", "blockquote"):
|
||||
self._flush()
|
||||
|
||||
def handle_data(self, data):
|
||||
if self._skip:
|
||||
return
|
||||
if self._in_title:
|
||||
self.title += data
|
||||
return
|
||||
text = data.strip()
|
||||
if text:
|
||||
self._current_block.append(text)
|
||||
|
||||
def _flush(self):
|
||||
if self._current_block:
|
||||
self.blocks.append(" ".join(self._current_block))
|
||||
self._current_block = []
|
||||
|
||||
def text(self):
|
||||
self._flush()
|
||||
return "\n\n".join(self.blocks)
|
||||
|
||||
|
||||
def _truncate(text, char_limit):
|
||||
if len(text) <= char_limit:
|
||||
return text, False
|
||||
head = int(char_limit * 0.75)
|
||||
tail = char_limit - head
|
||||
return (text[:head] + "\n\n[TRUNCATED]\n\n" + text[-tail:]
|
||||
if tail > 0 else text[:head] + "\n\n[TRUNCATED]"), True
|
||||
|
||||
|
||||
def _is_blocked_fetch_url(url):
|
||||
parsed = urllib.parse.urlparse(url)
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
return "only http(s) URLs can be fetched"
|
||||
host = (parsed.hostname or "").lower()
|
||||
if host in ("localhost",):
|
||||
return "localhost URLs are blocked"
|
||||
try:
|
||||
addr = ipaddress.ip_address(host)
|
||||
except ValueError:
|
||||
return None
|
||||
if addr.is_private or addr.is_loopback or addr.is_link_local:
|
||||
return "private / loopback URLs are blocked"
|
||||
return None
|
||||
|
||||
|
||||
def fetch_url(url, char_limit=DEFAULT_FETCH_CHAR_LIMIT, env=None,
|
||||
http=None):
|
||||
"""Fetch one URL and return clean text + links (truncated)."""
|
||||
if not url or not str(url).strip():
|
||||
raise ValueError("url must be a non-empty string")
|
||||
env = env if env is not None else os.environ
|
||||
if env.get("WEBSEARCH_ALLOW_PRIVATE_URLS") != "1":
|
||||
blocked = _is_blocked_fetch_url(str(url))
|
||||
if blocked:
|
||||
raise WebSearchError(blocked)
|
||||
http = http or HttpClient(
|
||||
timeout=int(env.get("WEBSEARCH_TIMEOUT", "20") or 20))
|
||||
try:
|
||||
char_limit = int(char_limit)
|
||||
except (TypeError, ValueError):
|
||||
char_limit = DEFAULT_FETCH_CHAR_LIMIT
|
||||
char_limit = max(2000, min(500000, char_limit))
|
||||
_status, body, ctype = http.get(str(url), {"User-Agent": _USER_AGENT})
|
||||
if ctype and not any(
|
||||
token in ctype for token in
|
||||
("text/html", "text/plain", "application/json",
|
||||
"application/xml", "text/")):
|
||||
raise WebSearchError("unsupported content type: %s" % ctype)
|
||||
text_in = body.decode("utf-8", errors="replace")
|
||||
if "text/html" in (ctype or "") or "<html" in text_in[:2000].lower():
|
||||
parser = _TextExtractor()
|
||||
parser.feed(text_in)
|
||||
title = _html.unescape(parser.title.strip())
|
||||
text = _html.unescape(parser.text())
|
||||
links = parser.links
|
||||
else:
|
||||
title = ""
|
||||
text = text_in
|
||||
links = []
|
||||
text, truncated = _truncate(text, char_limit)
|
||||
return {"url": str(url), "title": title, "text": text,
|
||||
"links": links[:100], "truncated": truncated,
|
||||
"char_limit": char_limit}
|
||||
|
||||
|
||||
# ------------------------------------------------------- codex-home glue
|
||||
|
||||
def ensure_codex_home_websearch(codex_home_dir, command_args=None,
|
||||
script_dir=None):
|
||||
"""Seed websearch MCP files + config into a tenant codex-home.
|
||||
|
||||
Idempotent and non-destructive by design:
|
||||
- copies ``websearch.py`` and ``websearch_mcp.py`` only when the
|
||||
destination is missing or differs (so a tenant's own edits to
|
||||
the script are never silently overwritten -- differing content
|
||||
is left alone and reported);
|
||||
- appends the ``[mcp_servers.websearch]`` section only when that
|
||||
section header is absent; an existing section is never rewritten
|
||||
(keys are never written into config.toml by this helper).
|
||||
Returns {config_path, script_path, config_changed, script_changed}.
|
||||
``command_args`` defaults to ``python3 <codex-home>/websearch_mcp.py``
|
||||
on the host; Docker tenants must pass the container path
|
||||
(``/var/lib/codex-app-server/codex-home/websearch_mcp.py``)
|
||||
instead, because the same config.toml cannot hold both paths.
|
||||
"""
|
||||
import shutil
|
||||
|
||||
home = os.path.abspath(codex_home_dir)
|
||||
os.makedirs(home, exist_ok=True)
|
||||
here = os.path.abspath(script_dir or os.path.dirname(__file__))
|
||||
script_changed = False
|
||||
script_differs = []
|
||||
for name in ("websearch.py", "websearch_mcp.py"):
|
||||
src = os.path.join(here, name)
|
||||
dst = os.path.join(home, name)
|
||||
if not os.path.exists(src):
|
||||
continue
|
||||
if not os.path.exists(dst):
|
||||
shutil.copyfile(src, dst)
|
||||
script_changed = True
|
||||
else:
|
||||
with open(src, "rb") as fh:
|
||||
src_bytes = fh.read()
|
||||
with open(dst, "rb") as fh:
|
||||
dst_bytes = fh.read()
|
||||
if src_bytes != dst_bytes:
|
||||
# Do not clobber a tenant-modified script; report it.
|
||||
script_differs.append(name)
|
||||
continue
|
||||
config_path = os.path.join(home, "config.toml")
|
||||
if command_args is None:
|
||||
command_args = ["python3",
|
||||
os.path.join(home, "websearch_mcp.py")]
|
||||
section = ("\n[mcp_servers.websearch]\n"
|
||||
"command = %s\n"
|
||||
"args = [%s]\n"
|
||||
% (json.dumps(command_args[0]),
|
||||
", ".join(json.dumps(a) for a in command_args[1:])))
|
||||
config_changed = False
|
||||
existing = ""
|
||||
if os.path.exists(config_path):
|
||||
with open(config_path, encoding="utf-8") as fh:
|
||||
existing = fh.read()
|
||||
if "[mcp_servers.websearch]" not in existing:
|
||||
with open(config_path, "a", encoding="utf-8") as fh:
|
||||
if existing and not existing.endswith("\n"):
|
||||
fh.write("\n")
|
||||
fh.write(section)
|
||||
config_changed = True
|
||||
return {"config_path": config_path,
|
||||
"script_path": os.path.join(home, "websearch_mcp.py"),
|
||||
"config_changed": config_changed,
|
||||
"script_changed": script_changed,
|
||||
"script_differs": script_differs}
|
||||
@@ -0,0 +1,181 @@
|
||||
"""Stdio MCP server exposing web_search / web_fetch to codex.
|
||||
|
||||
Protocol: JSON-RPC 2.0 over stdio, one JSON object per line.
|
||||
Logs go to stderr only -- anything on stdout would corrupt the
|
||||
JSON-RPC stream.
|
||||
|
||||
This file is deliberately thin: all provider logic lives in
|
||||
``websearch.py``. It can run two ways:
|
||||
- inside the repo: ``python3 -m cloud.websearch_mcp`` (imports
|
||||
``cloud.websearch``);
|
||||
- copied flat into a tenant codex-home next to ``websearch.py``
|
||||
(imports plain ``websearch`` from its own directory).
|
||||
The ``ensure_codex_home_websearch`` helper in websearch.py does the
|
||||
copying; it never writes API keys into config.toml -- keys are read
|
||||
from the MCP server process environment by websearch.py.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
try:
|
||||
from cloud.websearch import WebSearchError, fetch_url, search
|
||||
except ImportError: # flat copy in a tenant codex-home
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from websearch import WebSearchError, fetch_url, search # noqa: E402
|
||||
|
||||
SERVER_NAME = "websearch"
|
||||
SERVER_VERSION = "0.1.0"
|
||||
|
||||
_TOOLS = [
|
||||
{
|
||||
"name": "web_search",
|
||||
"description": (
|
||||
"Search the web and return ranked results with title, URL "
|
||||
"and snippet, plus a synthesised answer when the provider "
|
||||
"supplies one. Use web_fetch to read specific pages. "
|
||||
"provider may be 'auto' (walk Perplexity, Exa, Tavily, "
|
||||
"Brave, then keyless SearXNG / DuckDuckGo, falling back "
|
||||
"on failure) or a single provider name to pin it."
|
||||
),
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {"type": "string",
|
||||
"description": "Search query"},
|
||||
"limit": {"type": "integer", "default": 5,
|
||||
"description": "Max results, 1-10"},
|
||||
"provider": {"type": "string", "default": "auto",
|
||||
"description": "auto or a provider name"},
|
||||
},
|
||||
"required": ["query"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "web_fetch",
|
||||
"description": (
|
||||
"Fetch one URL and return clean text plus its links, "
|
||||
"truncated head+tail when over char_limit. "
|
||||
"Only http(s) URLs; localhost / private IPs are blocked."
|
||||
),
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"url": {"type": "string",
|
||||
"description": "URL to fetch"},
|
||||
"char_limit": {"type": "integer", "default": 15000,
|
||||
"description": "Max characters, "
|
||||
"clamped 2000-500000"},
|
||||
},
|
||||
"required": ["url"],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _err(code, message):
|
||||
return {"code": code, "message": message}
|
||||
|
||||
|
||||
def handle_request(request, search_fn=None, fetch_fn=None):
|
||||
"""Handle one JSON-RPC request dict. Returns a response dict or
|
||||
None for notifications (no id). Pure function for testability:
|
||||
the stdio loop passes the real search/fetch, tests pass fakes."""
|
||||
search_fn = search_fn or search
|
||||
fetch_fn = fetch_fn or fetch_url
|
||||
if not isinstance(request, dict):
|
||||
return {"jsonrpc": "2.0", "id": None,
|
||||
"error": _err(-32700, "request is not an object")}
|
||||
method = request.get("method")
|
||||
req_id = request.get("id")
|
||||
is_notification = "id" not in request
|
||||
|
||||
if method == "initialize":
|
||||
if is_notification:
|
||||
return None
|
||||
return {"jsonrpc": "2.0", "id": req_id,
|
||||
"result": {
|
||||
"protocolVersion": request.get("params", {}).get(
|
||||
"protocolVersion", "2024-11-05"),
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": SERVER_NAME,
|
||||
"version": SERVER_VERSION}}}
|
||||
if method in ("notifications/initialized", "notifications/cancelled"):
|
||||
return None
|
||||
if method == "ping":
|
||||
if is_notification:
|
||||
return None
|
||||
return {"jsonrpc": "2.0", "id": req_id, "result": {}}
|
||||
if method == "tools/list":
|
||||
if is_notification:
|
||||
return None
|
||||
return {"jsonrpc": "2.0", "id": req_id,
|
||||
"result": {"tools": _TOOLS}}
|
||||
if method == "tools/call":
|
||||
if is_notification:
|
||||
return None
|
||||
params = request.get("params") or {}
|
||||
name = params.get("name")
|
||||
args = params.get("arguments") or {}
|
||||
try:
|
||||
if name == "web_search":
|
||||
result = search_fn(
|
||||
args.get("query"),
|
||||
limit=args.get("limit", 5),
|
||||
provider=args.get("provider", "auto"))
|
||||
elif name == "web_fetch":
|
||||
result = fetch_fn(
|
||||
args.get("url"),
|
||||
char_limit=args.get("char_limit", 15000))
|
||||
else:
|
||||
return {"jsonrpc": "2.0", "id": req_id,
|
||||
"error": _err(-32602, "unknown tool: %r" % name)}
|
||||
except Exception as exc: # noqa: BLE001 - tool errors are data
|
||||
# WebSearchError messages are already key-free by
|
||||
# construction in websearch.py; other exceptions are
|
||||
# rendered as type + message, never with env or headers.
|
||||
return {"jsonrpc": "2.0", "id": req_id,
|
||||
"result": {"content": [{"type": "text",
|
||||
"text": "error: %s" % exc}],
|
||||
"isError": True}}
|
||||
return {"jsonrpc": "2.0", "id": req_id,
|
||||
"result": {"content": [{"type": "text",
|
||||
"text": json.dumps(
|
||||
result, ensure_ascii=False)}],
|
||||
"isError": False}}
|
||||
if is_notification:
|
||||
return None
|
||||
return {"jsonrpc": "2.0", "id": req_id,
|
||||
"error": _err(-32601, "method not found: %r" % method)}
|
||||
|
||||
|
||||
def main():
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
request = json.loads(line)
|
||||
except json.JSONDecodeError as exc:
|
||||
sys.stdout.write(json.dumps(
|
||||
{"jsonrpc": "2.0", "id": None,
|
||||
"error": _err(-32700, "parse error: %s" % exc)}) + "\n")
|
||||
sys.stdout.flush()
|
||||
continue
|
||||
try:
|
||||
response = handle_request(request)
|
||||
except Exception as exc: # noqa: BLE001 - server must not die
|
||||
print("websearch_mcp: handler error: %s" % exc, file=sys.stderr)
|
||||
if "id" in request:
|
||||
response = {"jsonrpc": "2.0", "id": request["id"],
|
||||
"error": _err(-32603, "internal error")}
|
||||
else:
|
||||
response = None
|
||||
if response is not None:
|
||||
sys.stdout.write(json.dumps(response) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user