diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8cabfd7ed..b7824c235 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,7 +19,9 @@ jobs: with: python-version: ${{ matrix.python-version }} - name: Install - run: pip install -e ".[dev]" + # `parallel` is an opt-in extra, but the OfficeQA Parallel search tests + # skip without it, so CI installs it to keep that provider covered. + run: pip install -e ".[dev,parallel]" - name: Pytest run: python -m pytest -q env: diff --git a/CHANGELOG.md b/CHANGELOG.md index c61a74dd6..47d0e99a7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,15 @@ All notable changes to SkillOpt are documented here. This project adheres to ## [Unreleased] ### Added +- **OfficeQA Parallel Search MCP** as an opt-in `env.search_mode=parallel` + provider for `openai_chat`/`qwen_chat`. The default stays `offline` and no + config selects Parallel automatically. Selecting it sends the generated + search queries and the full task question (as the MCP `objective`) to + Parallel's hosted endpoint, over an anonymous, lower-rate free tier, so + its scores are not directly comparable to the paper's Google-backed search + setting. The MCP client ships in a new `parallel` extra + (`pip install 'skillopt[parallel]'`) and is imported lazily, so OfficeQA + runs and offline installs do not require `mcp`. - **SkillOpt-Sleep paired A/B evalkit** (`python -m skillopt_sleep.evalkit`): McNemar plus percentile-bootstrap CIs on a fixed task manifest, with task-cluster multi-seed inference and seeded null calibration for exact-test diff --git a/docs/reference/config.md b/docs/reference/config.md index 6371f69dd..0645c0a67 100644 --- a/docs/reference/config.md +++ b/docs/reference/config.md @@ -210,6 +210,50 @@ forwarded consistently to both SDK and CLI execution paths. Benchmark-specific `env` keys are passed through to the adapter. +### OfficeQA web search + +`env.search_mode` selects how an OfficeQA rollout gathers evidence. The +default is unchanged: `offline` searches only the configured local document +corpus and makes no network search call. + +| Setting | Behavior | +|---|---| +| `env.search_mode: offline` (default) | Keep search local to the configured document corpus. `env.use_local_tools: false` also disables local function tools in this mode. | +| `env.search_mode: custom_search` | Use the token-authenticated HTTP service. Set `search_api_url`, `search_auth_env` and `search_provider` for your service. | +| `env.search_mode: azure_search` | Use the model-provided Azure search path. | +| `env.search_mode: parallel` | Opt in to Parallel's hosted Search MCP. Requires `openai_chat` or `qwen_chat` and the `parallel` extra. | +| `env.search_timeout_seconds` | Whole MCP discovery/search deadline in seconds (default `20`). | +| `env.search_max_num_results` | Maximum number of Parallel evidence results per round (default `4`). | +| `env.max_queries_per_turn` / `env.max_tool_turns` | Bound model-requested queries and model rounds. | + +#### Parallel Search MCP (opt-in) + +[Parallel Search MCP](https://docs.parallel.ai/integrations/mcp/search-mcp) +is never selected automatically. It is used only when you set +`env.search_mode=parallel` explicitly, for example with +`--cfg-options env.search_mode=parallel`. It also needs the optional MCP +client, installed with `pip install 'skillopt[parallel]'`; selecting the mode +without that extra raises an error naming the install command. Every other +mode, and every OfficeQA import, works without `mcp` installed. + +**Data boundary.** When this mode is selected, each search round sends data +to Parallel's hosted endpoint at `https://search.parallel.ai/mcp`: the +model's generated ``, and the **full text of the task +question** as the MCP `objective` argument. Do not select this mode for +tasks whose question text cannot leave your environment. Requests are +anonymous — they carry no `PARALLEL_API_KEY` and no saved MCP credentials — +so they run on the anonymous, lower-rate free tier described in Parallel's +documentation. That tier is rate-limited, not unmetered: expect throttling on +sustained runs, and use a keyed provider for large sweeps. Rate-limit and +other service errors fail +the rollout with a search diagnostic; an explicit provider never silently +falls back to another one. + +Parallel search does not require a local OfficeQA document corpus; parsed +source-page context is included when that corpus is available. This mode +changes the evidence source, so its scores are not directly comparable to, +and are not a reproduction of, the paper's Google-backed search setting. + ## Credential Environment Variables ### Azure-family backend diff --git a/pyproject.toml b/pyproject.toml index 5d50b8fef..3808b69e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,8 @@ claude = ["claude-agent-sdk>=0.1.0", "json_repair>=0.61.0"] qwen = ["vllm>=0.4.0", "json_repair>=0.61.0"] # SearchQA data materialization searchqa = ["datasets>=2.18.0"] +# OfficeQA `search_mode: parallel` web search over Parallel's hosted Search MCP +parallel = ["mcp>=1.27.0,<2"] # Documentation site docs = ["mkdocs-material>=9.5.0", "mkdocstrings[python]>=0.24.0"] # WebUI dashboard diff --git a/skillopt/envs/officeqa/parallel_search.py b/skillopt/envs/officeqa/parallel_search.py new file mode 100644 index 000000000..0a9d4d42e --- /dev/null +++ b/skillopt/envs/officeqa/parallel_search.py @@ -0,0 +1,107 @@ +"""Parallel Search MCP client for the OfficeQA controller's search rounds. + +Reached only through the explicit `search_mode=parallel` selection. Every call +sends the model's generated queries and the full task question (as the MCP +`objective`) to Parallel's hosted endpoint, anonymously, on the documented +lower-rate free tier. +""" +from __future__ import annotations + +import asyncio +import json +from datetime import timedelta +from urllib.parse import urlsplit + +import httpx +from mcp import ClientSession +from mcp.client.streamable_http import streamable_http_client +from mcp.types import Implementation + +from skillopt import __version__ + +MCP_URL = "https://search.parallel.ai/mcp" +USER_AGENT = f"SkillOpt/{__version__} (+https://github.com/microsoft/SkillOpt)" + + +async def _search(queries: list[str], objective: str, session_id: str, timeout: float) -> dict: + # This anonymous connection deliberately does not read API keys or saved MCP + # credentials. Keep the same project identity on every transport request. + async with httpx.AsyncClient( + headers={"User-Agent": USER_AGENT}, + timeout=httpx.Timeout(timeout), + follow_redirects=False, + ) as client: + async with streamable_http_client(MCP_URL, http_client=client) as (read, write, _): + async with ClientSession( + read, write, + read_timeout_seconds=timedelta(seconds=timeout), + client_info=Implementation(name="SkillOpt", version=__version__), + ) as session: + await session.initialize() + tools = await session.list_tools() + if not any(tool.name == "web_search" for tool in tools.tools): + raise RuntimeError("Parallel MCP did not advertise web_search") + result = await session.call_tool("web_search", { + "objective": objective, + "search_queries": queries, + "session_id": session_id, + }) + text = "\n".join(block.text for block in result.content if block.type == "text") + if result.isError: + raise RuntimeError(f"Parallel search error: {text[:1000]}") + payload = result.structuredContent + if payload is None: + try: + payload = json.loads(text) + except (TypeError, ValueError) as exc: + raise RuntimeError("Parallel search returned invalid JSON evidence") from exc + if not isinstance(payload, dict) or not isinstance(payload.get("results"), list): + raise RuntimeError("Parallel search returned invalid results") + return payload + + +def parallel_search( + queries: list[str], *, objective: str, session_id: str, + max_num_results: int = 4, timeout: float = 20, +) -> str: + """Search related queries with a bounded request and render URL evidence. + + The controller supplies one stable session ID per OfficeQA task, including + subsequent rounds. Timeouts and cancellation propagate to the controller; + an exceeded deadline always raises the built-in `TimeoutError`. + """ + if not queries or any(not isinstance(query, str) or not query.strip() for query in queries): + raise ValueError("Parallel search requires non-empty queries") + if timeout <= 0 or max_num_results <= 0: + raise ValueError("Parallel search timeout and result limit must be positive") + + async def bounded() -> dict: + return await asyncio.wait_for(_search(queries, objective, session_id, timeout), timeout=timeout) + + try: + payload = asyncio.run(bounded()) + except Exception as exc: + # SDK task groups may wrap a useful transport/tool error. Preserve the + # underlying diagnostic, including explicit anonymous quota responses. + error = exc + while len(getattr(error, "exceptions", ())) == 1: + error = error.exceptions[0] + if isinstance(error, asyncio.TimeoutError): + # Before 3.11 `asyncio.TimeoutError` is its own class rather than the + # built-in one, so normalize the deadline for callers on every version. + raise TimeoutError(f"Parallel search exceeded its {timeout}s deadline") from exc + if error is not exc: + raise error from exc + raise + blocks = [] + for item in payload["results"][:max_num_results]: + if not isinstance(item, dict): + raise RuntimeError("Parallel search returned an invalid result item") + url = item.get("url") + excerpts = item.get("excerpts") + if not isinstance(url, str) or urlsplit(url).scheme not in {"http", "https"} or not urlsplit(url).netloc: + raise RuntimeError("Parallel search returned an invalid result URL") + if not isinstance(excerpts, list) or any(not isinstance(part, str) for part in excerpts): + raise RuntimeError("Parallel search returned invalid excerpts") + blocks.append(f"[{len(blocks) + 1}] {item.get('title') or url}\nURL: {url}\n" + "\n".join(excerpts)[:6000]) + return "\n\n".join(blocks) if blocks else "[no search results]" diff --git a/skillopt/envs/officeqa/rollout.py b/skillopt/envs/officeqa/rollout.py index 7d4c8e045..e45629cd5 100644 --- a/skillopt/envs/officeqa/rollout.py +++ b/skillopt/envs/officeqa/rollout.py @@ -3,6 +3,7 @@ import json import os import re +import uuid from concurrent.futures import ThreadPoolExecutor, as_completed from skillopt.envs.officeqa.evaluator import evaluate @@ -68,13 +69,30 @@ _DEFAULT_SEARCH_MODE = "offline" _CUSTOM_SEARCH_MODE = "custom_search" _AZURE_SEARCH_MODE = "azure_search" +_PARALLEL_SEARCH_MODE = "parallel" def _normalize_search_mode(search_mode: str | None) -> str: normalized = str(search_mode or _DEFAULT_SEARCH_MODE).strip().lower() if normalized in {"custom", _CUSTOM_SEARCH_MODE}: return _CUSTOM_SEARCH_MODE if normalized in {"azure", _AZURE_SEARCH_MODE}: return _AZURE_SEARCH_MODE + if normalized == _PARALLEL_SEARCH_MODE: + return _PARALLEL_SEARCH_MODE return _DEFAULT_SEARCH_MODE +def _load_parallel_search(): + """Import the Parallel client only when `parallel` mode is actually selected. + + The `mcp` package is an optional extra, so importing it at module scope + would make it a hard requirement of every OfficeQA rollout. + """ + try: + from skillopt.envs.officeqa.parallel_search import parallel_search + except ImportError as exc: + raise RuntimeError( + "search_mode='parallel' requires the optional MCP client. " + "Install it with: pip install 'skillopt[parallel]'" + ) from exc + return parallel_search def _build_system( skill_content: str, *, @@ -95,7 +113,7 @@ def _build_system( "answer inside ....\n\n" + skill_section ).rstrip() - if normalized_search_mode == _CUSTOM_SEARCH_MODE: + if normalized_search_mode in {_CUSTOM_SEARCH_MODE, _PARALLEL_SEARCH_MODE}: protocol = ( "You are an expert OfficeQA research assistant. Solve the question using the provided oracle parsed " "OfficeQA page(s) and evidence returned by the controller-managed custom search loop.\n\n" @@ -169,7 +187,7 @@ def _build_user( parts.append("## File Hints\n" + "\n".join(f"- {hint}" for hint in item["source_files"])) if diagnostic_mode and diagnostic_instruction.strip(): parts.append(f"## Training Readout\n{diagnostic_instruction.strip()}") - if normalized_search_mode == _CUSTOM_SEARCH_MODE: + if normalized_search_mode in {_CUSTOM_SEARCH_MODE, _PARALLEL_SEARCH_MODE}: parts.append( _build_round_instruction( turn=turn, @@ -352,16 +370,17 @@ def _run_custom_search_process( search_max_num_results: int, search_timeout_seconds: int, oracle_context: str = "", + search_mode: str = _CUSTOM_SEARCH_MODE, ) -> tuple[str, str, str, str, list[dict], str, dict]: - if not str(search_api_url or "").strip(): + if search_mode == _CUSTOM_SEARCH_MODE and not str(search_api_url or "").strip(): raise ValueError("custom_search mode requires a non-empty search_api_url") - if not os.environ.get(search_auth_env, "").strip(): + if search_mode == _CUSTOM_SEARCH_MODE and not os.environ.get(search_auth_env, "").strip(): raise ValueError(f"custom_search mode requires auth token env var {search_auth_env}") if get_target_backend() not in {"openai_chat", "qwen_chat"}: - raise ValueError("custom_search mode is only supported with target_backend='openai_chat' or 'qwen_chat'") + raise ValueError(f"{search_mode} mode is only supported with target_backend='openai_chat' or 'qwen_chat'") system = _build_system( skill_content, - search_mode=_CUSTOM_SEARCH_MODE, + search_mode=search_mode, max_tool_turns=max_tool_turns, max_queries_per_turn=max_queries_per_turn, ) @@ -369,7 +388,7 @@ def _run_custom_search_process( item, diagnostic_mode=diagnostic_mode, diagnostic_instruction=diagnostic_instruction, - search_mode=_CUSTOM_SEARCH_MODE, + search_mode=search_mode, turn=1, max_tool_turns=max_tool_turns, max_queries_per_turn=max_queries_per_turn, @@ -385,6 +404,7 @@ def _run_custom_search_process( final_answer = "" fail_reason = "" last_response_metadata: dict = {} + session_id = uuid.uuid4().hex for turn in range(1, max_tool_turns + 1): message, _ = chat_target_messages( messages=messages, @@ -411,15 +431,22 @@ def _run_custom_search_process( if not queries: fail_reason = "Model neither produced search queries nor a final answer" break - results_text = _execute_custom_search_round( - queries, - api_url=search_api_url, - auth_env=search_auth_env, - provider=search_provider, - max_num_results=search_max_num_results, - timeout=search_timeout_seconds, - ) - conversation.append({"type": "tool_call", "turn": turn, "cmd": f"custom_search({queries!r})", "obs": results_text}) + if search_mode == _PARALLEL_SEARCH_MODE: + # A failed explicit provider stays an error; never retry via another service. + results_text = _load_parallel_search()( + queries, objective=item["question"], session_id=session_id, + max_num_results=search_max_num_results, timeout=search_timeout_seconds, + ) + else: + results_text = _execute_custom_search_round( + queries, + api_url=search_api_url, + auth_env=search_auth_env, + provider=search_provider, + max_num_results=search_max_num_results, + timeout=search_timeout_seconds, + ) + conversation.append({"type": "tool_call", "turn": turn, "cmd": f"{search_mode}({queries!r})", "obs": results_text}) latest_user = ( f"## Search Results Round {turn}\n{results_text}\n\n" + _build_round_instruction( @@ -550,8 +577,14 @@ def process_one( else "Treat it as primary document evidence for answering the question." ), ) - elif normalized_search_mode == _CUSTOM_SEARCH_MODE: - docs_roots = resolve_docs_roots(data_dirs) + elif normalized_search_mode in {_CUSTOM_SEARCH_MODE, _PARALLEL_SEARCH_MODE}: + try: + docs_roots = resolve_docs_roots(data_dirs) + except FileNotFoundError: + if normalized_search_mode == _CUSTOM_SEARCH_MODE: + raise + # Web evidence works without downloading a local OfficeQA corpus. + docs_roots = [] if item.get("source_files"): candidate_files = resolve_candidate_files(item.get("source_files", []), docs_roots) oracle_context = build_oracle_parsed_pages_context( @@ -583,7 +616,7 @@ def process_one( last_response_metadata: dict = {} allowed_files = [os.path.basename(path) for path in candidate_files] try: - if normalized_search_mode == _CUSTOM_SEARCH_MODE: + if normalized_search_mode in {_CUSTOM_SEARCH_MODE, _PARALLEL_SEARCH_MODE}: system, user, final_response, final_answer, conversation, fail_reason, last_response_metadata = _run_custom_search_process( item, skill_content, @@ -598,6 +631,7 @@ def process_one( search_max_num_results=search_max_num_results, search_timeout_seconds=search_timeout_seconds, oracle_context=oracle_context, + search_mode=normalized_search_mode, ) elif normalized_search_mode == _AZURE_SEARCH_MODE: system, user, final_response, final_answer, conversation, fail_reason, last_response_metadata = _run_azure_search_process( diff --git a/tests/test_officeqa_parallel_search.py b/tests/test_officeqa_parallel_search.py new file mode 100644 index 000000000..22ad76f97 --- /dev/null +++ b/tests/test_officeqa_parallel_search.py @@ -0,0 +1,289 @@ +"""OfficeQA search-mode selection, optional-dependency handling and MCP caller coverage.""" +from __future__ import annotations + +import asyncio +import builtins +import importlib +import importlib.util +import json +import sys +from types import SimpleNamespace +from unittest import mock + +import pytest + +from skillopt.envs.officeqa import rollout +from skillopt.envs.officeqa.adapter import OfficeQAAdapter + +HAS_MCP = importlib.util.find_spec("mcp") is not None +requires_mcp = pytest.mark.skipif(not HAS_MCP, reason="requires the optional 'parallel' extra") +mcp_search = importlib.import_module("skillopt.envs.officeqa.parallel_search") if HAS_MCP else None + +ITEM = {"id": "treasury", "question": "When was the declaration signed?", "ground_truth": "1776"} +PAYLOAD = {"results": [{"url": "https://example.org/treasury?year=1776", "title": "Treasury history", "excerpts": ["The declaration was signed in 1776."]}]} + + +@pytest.fixture(autouse=True) +def isolated_search_config(monkeypatch): + for name in ("OFFICEQA_CUSTOM_SEARCH_AUTH", "OFFICEQA_SEARCH_API_URL", "PARALLEL_API_KEY"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(rollout, "get_target_backend", lambda: "openai_chat") + + +@pytest.fixture +def transport(monkeypatch): + """Observe actual SDK HTTP requests without substituting its MCP protocol.""" + import httpx + + requests = [] + behavior = {"payload": PAYLOAD, "error": None, "delay": False} + original_client = httpx.AsyncClient + + async def handler(request): + body = json.loads(request.content) if request.content else {} + requests.append((request, body)) + method = body.get("method") + if method == "notifications/initialized": + return httpx.Response(202) + if method == "initialize": + result = {"protocolVersion": "2025-03-26", "capabilities": {"tools": {}}, "serverInfo": {"name": "fixture", "version": "1"}} + elif method == "tools/list": + result = {"tools": [{"name": "web_search", "inputSchema": {"type": "object"}}]} + elif method == "tools/call": + if behavior["delay"]: + await asyncio.Event().wait() + if behavior["error"]: + result = {"content": [{"type": "text", "text": behavior["error"]}], "isError": True} + else: + result = {"content": [{"type": "text", "text": json.dumps(behavior["payload"])}], "structuredContent": behavior["payload"]} + else: + raise AssertionError(f"Unexpected MCP request: {request.method} {body}") + return httpx.Response(200, json={"jsonrpc": "2.0", "id": body["id"], "result": result}) + + def client(**kwargs): + return original_client(transport=httpx.MockTransport(handler), **kwargs) + + monkeypatch.setattr(mcp_search.httpx, "AsyncClient", client) + return requests, behavior + + +def scripted_model(monkeypatch, *, rounds=1): + calls = [] + + def chat(**kwargs): + # Copy now because the caller appends further messages after the call. + messages = [dict(message) for message in kwargs["messages"]] + calls.append(messages) + if len(calls) <= rounds: + text = '["declaration treasury history"]' + else: + assert "https://example.org/treasury?year=1776" in messages[-1]["content"] + assert "signed in 1776" in messages[-1]["content"] + text = "1776" + return SimpleNamespace(content=text), {} + + monkeypatch.setattr(rollout, "chat_target_messages", chat) + return calls + + +def block_mcp_import(monkeypatch): + """Make `import mcp` fail the way an install without the extra would.""" + real_import = builtins.__import__ + + def guarded(name, *args, **kwargs): + if name == "mcp" or name.startswith("mcp."): + raise ImportError("No module named 'mcp'") + return real_import(name, *args, **kwargs) + + for name in [key for key in sys.modules if key == "mcp" or key.startswith("mcp.")]: + monkeypatch.delitem(sys.modules, name) + monkeypatch.delitem(sys.modules, "skillopt.envs.officeqa.parallel_search", raising=False) + monkeypatch.setattr(builtins, "__import__", guarded) + + +# ── Default selection: Parallel is never automatic ─────────────────────────── + +def test_shipped_config_keeps_the_offline_default(): + from scripts.train import get_adapter + from skillopt.config import flatten_config, load_config + + cfg = flatten_config(load_config("configs/officeqa/default.yaml")) + assert cfg["search_mode"] == "offline" + adapter = get_adapter(cfg) + assert adapter.search_mode == "offline" + assert adapter.search_provider == "duckduckgo" + assert rollout._normalize_search_mode(adapter.search_mode) == rollout._DEFAULT_SEARCH_MODE + + +@pytest.mark.parametrize("mode", ["auto", "", None, "fresh", "web", "parallel_search"]) +def test_auto_and_unrecognized_modes_never_select_parallel(mode): + assert rollout._normalize_search_mode(mode) == rollout._DEFAULT_SEARCH_MODE + + +@pytest.mark.parametrize("mode", ["parallel", "PARALLEL", " Parallel "]) +def test_parallel_requires_explicit_selection(mode): + assert rollout._normalize_search_mode(mode) == rollout._PARALLEL_SEARCH_MODE + + +@pytest.mark.parametrize( + ("mode", "expected"), + [("offline", "offline"), ("custom", "custom_search"), ("custom_search", "custom_search"), ("azure", "azure_search"), ("azure_search", "azure_search")], +) +def test_existing_modes_resolve_unchanged(mode, expected): + assert rollout._normalize_search_mode(mode) == expected + + +@pytest.mark.parametrize("mode", ["auto", "offline"]) +def test_fresh_auto_config_runs_offline_without_search(tmp_path, monkeypatch, mode): + docs = tmp_path / "docs" + docs.mkdir() + monkeypatch.setattr(rollout, "chat_target_messages", lambda **kwargs: (SimpleNamespace(content="1776"), {})) + with ( + mock.patch.object(rollout, "_load_parallel_search", side_effect=AssertionError("parallel not selected")), + mock.patch.object(rollout, "custom_search", side_effect=AssertionError("custom search not selected")), + mock.patch.object(rollout, "run_tool", side_effect=AssertionError("local tools disabled")), + ): + row = rollout.process_one(ITEM, str(tmp_path / "out"), "", search_mode=mode, use_local_tools=False, data_dirs=[str(docs)]) + assert row["predicted_answer"] == "1776" + + +# ── Optional `mcp` dependency ──────────────────────────────────────────────── + +def test_officeqa_env_imports_without_mcp(monkeypatch): + block_mcp_import(monkeypatch) + for name in [key for key in sys.modules if key.startswith("skillopt.envs.officeqa")]: + monkeypatch.delitem(sys.modules, name) + reloaded = importlib.import_module("skillopt.envs.officeqa.adapter") + assert reloaded.OfficeQAAdapter().search_mode == "offline" + assert importlib.import_module("skillopt.envs.officeqa.rollout")._normalize_search_mode("offline") == "offline" + + +def test_parallel_without_the_extra_names_the_install_command(monkeypatch): + block_mcp_import(monkeypatch) + with pytest.raises(RuntimeError, match=r"skillopt\[parallel\]"): + rollout._load_parallel_search() + + +def test_parallel_rollout_without_the_extra_fails_with_the_install_command(tmp_path, monkeypatch): + block_mcp_import(monkeypatch) + scripted_model(monkeypatch) + row = rollout.process_one(ITEM, str(tmp_path), "", search_mode="parallel", max_tool_turns=2) + assert "skillopt[parallel]" in row["fail_reason"] + assert row["hard"] == 0 + + +# ── Parallel MCP caller ────────────────────────────────────────────────────── + +@requires_mcp +@pytest.mark.parametrize("backend", ["openai_chat", "qwen_chat"]) +def test_explicit_parallel_runs_mcp_and_consumes_evidence(tmp_path, monkeypatch, transport, backend): + monkeypatch.setattr(rollout, "get_target_backend", lambda: backend) + calls = scripted_model(monkeypatch, rounds=2) + adapter = OfficeQAAdapter(search_mode="parallel", workers=1, max_tool_turns=3) + rows = adapter.rollout([ITEM], "Read evidence before answering.", str(tmp_path)) + assert rows[0]["predicted_answer"] == "1776" + assert rows[0]["hard"] == 1 + assert len(calls) == 3 + requests, _ = transport + methods = [body.get("method") for _, body in requests] + assert methods.count("initialize") == methods.count("tools/list") == methods.count("tools/call") == 2 + searches = [body["params"] for _, body in requests if body.get("method") == "tools/call"] + assert all(call["name"] == "web_search" for call in searches) + assert searches[0]["arguments"]["session_id"] == searches[1]["arguments"]["session_id"] + assert len(searches[0]["arguments"]["session_id"]) == 32 + assert searches[0]["arguments"]["search_queries"] == ["declaration treasury history"] + # Documented data boundary: the whole task question leaves as the objective. + assert searches[0]["arguments"]["objective"] == ITEM["question"] + assert all(str(request.url) == "https://search.parallel.ai/mcp" for request, _ in requests) + assert all(request.headers["user-agent"] == mcp_search.USER_AGENT for request, _ in requests) + assert all("authorization" not in request.headers for request, _ in requests) + conversation = json.loads((tmp_path / "predictions" / ITEM["id"] / "conversation.json").read_text()) + assert any(event.get("type") == "tool_call" and "signed in 1776" in event["obs"] for event in conversation) + + +@requires_mcp +def test_environment_key_is_not_sent_to_anonymous_mcp(monkeypatch, transport): + monkeypatch.setenv("PARALLEL_API_KEY", "must-not-be-sent") + text = mcp_search.parallel_search(["treasury"], objective="history", session_id="a" * 32) + assert "https://example.org/treasury?year=1776" in text + assert all("authorization" not in request.headers for request, _ in transport[0]) + + +@requires_mcp +def test_explicit_parallel_ignores_legacy_credentials(tmp_path, monkeypatch, transport): + monkeypatch.setenv("OFFICEQA_CUSTOM_SEARCH_AUTH", "configured-token") + monkeypatch.setenv("OFFICEQA_SEARCH_API_URL", "https://configured.example") + scripted_model(monkeypatch) + row = rollout.process_one(ITEM, str(tmp_path), "", search_mode="parallel", search_provider="google", max_tool_turns=2) + assert row["hard"] == 1 + assert transport[0] + assert all("authorization" not in request.headers for request, _ in transport[0]) + + +@requires_mcp +def test_saved_custom_provider_failure_does_not_use_parallel(tmp_path, monkeypatch, transport): + docs = tmp_path / "docs" + docs.mkdir() + monkeypatch.setenv("OFFICEQA_CUSTOM_SEARCH_AUTH", "configured-token") + calls = [] + + def chat(**kwargs): + calls.append(kwargs) + if len(calls) == 1: + return SimpleNamespace(content='["treasury"]'), {} + assert "configured service failed" in kwargs["messages"][-1]["content"] + return SimpleNamespace(content="unknown"), {} + + monkeypatch.setattr(rollout, "chat_target_messages", chat) + with mock.patch.object(rollout, "custom_search", side_effect=RuntimeError("configured service failed")) as custom: + row = rollout.process_one(ITEM, str(tmp_path / "out"), "", search_mode="custom_search", search_api_url="https://configured.example", search_provider="google", data_dirs=[str(docs)], max_tool_turns=2) + assert custom.call_args.kwargs["provider"] == "google" + assert custom.call_args.kwargs["api_url"] == "https://configured.example" + assert not transport[0] + assert row["hard"] == 0 + + +@requires_mcp +@pytest.mark.parametrize("error", ["anonymous rate limit exceeded", "upstream unavailable"]) +def test_mcp_errors_fail_task_without_fallback(tmp_path, monkeypatch, transport, error): + transport[1]["error"] = error + calls = scripted_model(monkeypatch) + with mock.patch.object(rollout, "custom_search", side_effect=AssertionError("no fallback")): + row = rollout.process_one(ITEM, str(tmp_path), "", search_mode="parallel", max_tool_turns=2) + assert error in row["fail_reason"] + assert row["hard"] == 0 + assert len(calls) == 1 + + +@requires_mcp +def test_whole_mcp_request_has_a_deadline(transport): + transport[1]["delay"] = True + with pytest.raises(TimeoutError): + mcp_search.parallel_search(["treasury"], objective="history", session_id="a" * 32, timeout=0.05) + + +@requires_mcp +def test_cancellation_closes_actual_mcp_session(transport): + transport[1]["delay"] = True + + async def cancel(): + task = asyncio.create_task(mcp_search._search(["treasury"], "history", "a" * 32, 20)) + for _ in range(500): + if any(body.get("method") == "tools/call" for _, body in transport[0]): + break + await asyncio.sleep(0.001) + else: + pytest.fail("MCP caller never reached search") + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(cancel()) + + +@requires_mcp +@pytest.mark.parametrize("payload", [{"results": [{"url": "file:///etc/passwd", "excerpts": ["bad"]}]}, {"results": "malformed"}, {"results": [{"url": "https://example.org", "excerpts": "not a list"}]}]) +def test_bad_mcp_evidence_is_rejected(transport, payload): + transport[1]["payload"] = payload + with pytest.raises(RuntimeError, match="invalid"): + mcp_search.parallel_search(["treasury"], objective="history", session_id="a" * 32)