benchmarks/collection_ingest/mcp_session.py
Rama main · réplica de lectura del repositorio original.
"""One long-lived stdio MCP session against the real server, with a call ledger.
Shared by ``drive.py`` and ``score.py`` so both observe the product the same way:
only public tool results, every call timed, every refusal kept verbatim. All
server state (vault, state root, logs, bytecode cache) points at the caller's
scratch directory, never at the repository or a real vault.
"""
from __future__ import annotations
import asyncio
import json
import os
import time
from pathlib import Path
from typing import Any
# A result that names `retry_after_ms` asks the caller to repeat the identical call.
RETRY_FIELD = "retry_after_ms"
DEFAULT_RETRY_DEADLINE_S = 180.0
ARG_LIST_CAP = 12
def isolated_env(
*, state: Path, vault: Path, server_root: Path, extra: dict[str, str] | None = None
) -> dict[str, str]:
"""Environment for a hermetic server: no EXOMEM_* inherited, every path in scratch."""
env = {
key: value
for key, value in os.environ.items()
if not key.startswith("EXOMEM_") and key not in {"PYTHONPATH", "XDG_STATE_HOME"}
}
env.update(
{
"EXOMEM_VAULT_PATH": str(vault),
"EXOMEM_STATE_ROOT": str(state / "state"),
"EXOMEM_CONFIG_PATH": str(state / "config.json"),
"EXOMEM_WRITER_LEASE_STATE_DIR": str(state / "leases"),
"EXOMEM_CALL_LEDGER_DIR": str(state / "ledger"),
"EXOMEM_LOG_DIR": str(state / "logs"),
# No model weights or network in this instrument: lexical retrieval only.
"EXOMEM_DISABLE_EMBEDDINGS": "1",
"EXOMEM_DISABLE_CLIP": "1",
"EXOMEM_DISABLE_RELEVANCE_CHECK": "1",
"EXOMEM_DISABLE_QUERY_LOG": "1",
"EXOMEM_DISABLE_RANKING_CONFIG": "1",
"EXOMEM_EAGER_BOOT": "1",
"HF_HUB_OFFLINE": "1",
"XDG_STATE_HOME": str(state / "xdg"),
"PYTHONPATH": str(server_root / "src"),
"PYTHONPYCACHEPREFIX": str(state / "pycache"),
"FASTMCP_CHECK_FOR_UPDATES": "off",
"FASTMCP_SHOW_SERVER_BANNER": "false",
}
)
env.update(extra or {})
return env
def summarise_args(args: dict[str, Any]) -> dict[str, Any]:
out: dict[str, Any] = {}
for key, value in args.items():
if isinstance(value, list) and len(value) > ARG_LIST_CAP:
out[key] = {"len": len(value), "head": value[:5]}
else:
out[key] = value
return out
def split_error_text(text: str) -> tuple[str | None, str]:
"""``CODE: reason`` is the product's error text format; anything else has no code."""
head, sep, tail = text.partition(": ")
if sep and head.replace("_", "").isalnum() and head == head.upper():
return head, tail
return None, text
def _texts(result: Any) -> list[str]:
return [t for t in (getattr(item, "text", None) for item in result.content or ()) if t]
def decode(result: Any) -> tuple[str, Any, str | None, str | None]:
"""Return ``(outcome, data, error_code, error_message)`` for a raw tool result."""
if result.is_error:
text = "\n".join(_texts(result))
code, _ = split_error_text(text.removeprefix("Error calling tool").split("': ", 1)[-1])
return "tool_error", None, code, text[:800]
data = result.structured_content
if data is None:
for text in _texts(result):
try:
data = json.loads(text)
except json.JSONDecodeError:
continue
break
return classify(data)
def classify(data: Any) -> tuple[str, Any, str | None, str | None]:
"""Return ``(outcome, data, error_code, error_message)`` for a tool's decoded reply object."""
if isinstance(data, dict) and set(data) == {"result"}:
data = data["result"]
if isinstance(data, dict):
error = data.get("error")
if data.get("success") is False or isinstance(error, dict):
err = error if isinstance(error, dict) else {}
return "refused", data, str(err.get("code") or "UNKNOWN"), json.dumps(err)[:800]
if data is None:
return "transport_error", None, None, "no structured or JSON content"
return "ok", data, None, None
def retry_hint_ms(data: Any) -> int | None:
"""The wait a refusal asks for, read from its ``error`` object; None if it asks nothing.
A minimal refusal names the wait as ``retry.after_s``; an older one names ``retry_after_ms``.
"""
error = data.get("error") if isinstance(data, dict) else None
if not isinstance(error, dict):
return None
value = error.get(RETRY_FIELD)
if isinstance(value, int | float):
return int(value)
retry = error.get("retry")
seconds = retry.get("after_s") if isinstance(retry, dict) else None
if isinstance(seconds, int | float):
return int(seconds * 1000)
return None
class LedgerSession:
"""The call ledger over any transport: a subclass supplies ``invoke``, one raw tool call.
``call`` repeats a call while the product asks for a retry and records every logical call
in ``self.calls``, so a drive log reads the same whichever transport carried it.
"""
def __init__(self) -> None:
self.calls: list[dict[str, Any]] = []
self.tools: list[str] = []
self.tool_schemas: dict[str, Any] = {}
async def invoke(self, tool: str, args: dict[str, Any]) -> tuple[str, Any, str | None, str | None]:
raise NotImplementedError
async def call(
self,
tool: str,
args: dict[str, Any],
*,
journey: str,
step: str,
retry_deadline_s: float = DEFAULT_RETRY_DEADLINE_S,
) -> tuple[str, Any, dict[str, Any]]:
"""Run one logical call, repeating it while the product asks for a retry.
Returns ``(outcome, data, record)``; the record is also appended to ``self.calls``.
"""
started = time.perf_counter()
attempts = 0
waited_ms = 0
warming_codes: list[str] = []
outcome, data, code, message = "transport_error", None, None, None
while True:
attempts += 1
try:
outcome, data, code, message = await self.invoke(tool, args)
except Exception as exc: # noqa: BLE001 - a measurement harness keeps every failure verbatim
outcome, data, code, message = "transport_error", None, type(exc).__name__, str(exc)[:800]
break
hint = retry_hint_ms(data) if outcome == "refused" else None
if hint is None or (time.perf_counter() - started) > retry_deadline_s:
break
warming_codes.append(code or "")
waited_ms += hint
await asyncio.sleep(hint / 1000.0)
record = {
"n": len(self.calls) + 1,
"journey": journey,
"step": step,
"tool": tool,
"args": summarise_args(args),
"outcome": outcome,
"error_code": code,
"error_message": message,
"attempts": attempts,
"retried_codes": sorted(set(warming_codes)),
"retry_wait_ms": waited_ms,
"wall_ms": round((time.perf_counter() - started) * 1000.0, 1),
}
self.calls.append(record)
return outcome, data, record
class McpSession(LedgerSession):
"""Async context manager over one server process; ``call`` records into ``self.calls``."""
def __init__(self, *, python: Path, state: Path, env: dict[str, str], timeout: float = 600.0):
super().__init__()
self.python = python
self.state = state
self.env = env
self.timeout = timeout
self._client: Any = None
async def __aenter__(self) -> McpSession:
from fastmcp import Client
from fastmcp.client.transports import StdioTransport
self.state.mkdir(parents=True, exist_ok=True)
transport = StdioTransport(
command=str(self.python),
args=["-m", "exomem", "--transport", "stdio"],
env=self.env,
cwd=str(self.state),
keep_alive=False,
log_file=self.state / "stdio.log",
)
self._client = Client(transport, timeout=self.timeout, init_timeout=self.timeout)
await self._client.__aenter__()
listed = await self._client.list_tools()
self.tools = sorted(tool.name for tool in listed)
self.tool_schemas = {tool.name: tool.input_schema for tool in listed}
# Deliver a public refusal envelope to the harness instead of a client-side schema error.
for owner in (self._client, getattr(self._client, "session", None)):
schemas = getattr(owner, "_tool_output_schemas", None)
if isinstance(schemas, dict):
for name in tuple(schemas):
schemas[name] = None
return self
async def __aexit__(self, *exc_info: object) -> None:
await self._client.__aexit__(*exc_info)
async def invoke(self, tool: str, args: dict[str, Any]) -> tuple[str, Any, str | None, str | None]:
return decode(await self._client.call_tool_mcp(tool, args))