File size: 5,992 Bytes
2a87ede | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | """HTTP client layer for the dense_evolution_mcp adapter: the shared
httpx.AsyncClient, the low-level _request helper every tool calls, and the
@catch_errors decorator that replaces the try/except-return-_handle_error(e)
boilerplate every tool used to repeat individually.
KERNEL_URL and _TEST_TRANSPORT live here (not config.py) specifically
because tests/integration/test_mcp_server.py mutates them directly
(`from mcp_server import client as mcp_client; mcp_client.KERNEL_URL = ...`)
to exercise the unreachable-kernel and in-process-ASGI-transport paths --
_get_client/_request must see those mutations as their OWN module globals,
which only works if the value's single source of truth lives in the same
module that reads it. Moving _TEST_TRANSPORT into a proper pytest fixture
instead of a raw mutable global is planned for Phase 4 (see prog.txt
Sezione 3.3), not done here.
"""
import functools
import os
from typing import Optional
import httpx
from .config import DEFAULT_TIMEOUT
KERNEL_URL = os.environ.get("DENSE_EVOLUTION_KERNEL_URL", "http://127.0.0.1:8800").rstrip("/")
# Set by tests to route through httpx.ASGITransport straight into the real,
# in-process local_site.app.server.app (see tests/integration/test_mcp_server.py) --
# None here means "use a real TCP connection to KERNEL_URL", unchanged from
# before this existed. This is the one seam in the whole adapter: it lets
# tests exercise every tool function against the real FastAPI kernel (real
# DenseSVSimulator, real PennyLane Hamiltonians -- no mocked physics) without
# a live subprocess bound to a real port in CI.
_TEST_TRANSPORT = None
# Lazily-created, reused across calls -- _request used to open a fresh
# httpx.AsyncClient (TCP handshake to the kernel) per tool call, real
# overhead for an MCP session that calls many tools in a row. Keyed on
# (_TEST_TRANSPORT identity, KERNEL_URL): production code never changes
# KERNEL_URL after import (it's a plain module-level constant read once
# from os.environ), but the test suite mutates both KERNEL_URL and
# _TEST_TRANSPORT directly to exercise the unreachable-kernel path, and its
# autouse fixture swaps _TEST_TRANSPORT to a fresh ASGITransport before
# every test and back to None after -- a client cached against a stale
# transport or base_url must not be reused across either kind of swap.
_shared_client: Optional[httpx.AsyncClient] = None
_shared_client_key = None # (transport identity, base_url) the cached client was built with
def _get_client() -> httpx.AsyncClient:
global _shared_client, _shared_client_key
current_key = (id(_TEST_TRANSPORT), KERNEL_URL)
if _shared_client is None or _shared_client_key != current_key:
# Client-level timeout is just the fallback ceiling -- every real
# call site passes its own timeout= to _request, sized to that
# endpoint's actual expected cost (a health check and a 10-minute
# VQE run have nothing in common).
_shared_client = httpx.AsyncClient(base_url=KERNEL_URL, transport=_TEST_TRANSPORT, timeout=DEFAULT_TIMEOUT)
_shared_client_key = current_key
return _shared_client
async def _request(method: str, path: str, timeout: float = DEFAULT_TIMEOUT, **kwargs) -> dict:
"""Reusable request helper for every tool. `timeout` (seconds) should be
sized to the specific endpoint's real expected cost -- see each tool's
own call site."""
client = _get_client()
try:
resp = await client.request(method, path, timeout=timeout, **kwargs)
except httpx.ConnectError:
raise RuntimeError(
f"Dense Evolution kernel not reachable at {KERNEL_URL}. Start it with "
"`dense-evolution serve` (or `python -m local_site.app.server` from the "
"repo root), then retry. Use dense_evolution_health to check connectivity."
)
except httpx.TimeoutException:
raise RuntimeError(
f"Request to the Dense Evolution kernel timed out after {timeout:g}s -- the "
"simulation may be too large or slow for this request (e.g. high qubit "
"count, many VQE iterations, or a long MD trajectory). Try reducing its size."
)
if resp.status_code >= 400:
try:
detail = resp.json().get("detail", resp.text)
except Exception:
detail = resp.text
raise RuntimeError(f"Dense Evolution kernel returned HTTP {resp.status_code}: {detail}")
return resp.json()
def _handle_error(e: Exception) -> str:
"""Consistent error formatting across all tools."""
return f"Error: {e}"
def catch_errors(func):
"""Wraps an async tool function so any exception raised anywhere in its
body becomes the same "Error: ..." string every tool used to return
from its own try/except -- replaces ~20 repeated
`try: ... except Exception as e: return _handle_error(e)` blocks with
one decorator.
A strict superset of the old behavior for the handful of tools whose
try block only wrapped part of the function -- e.g.
dense_evolution_run_circuit's post-request image saving/truncation used
to sit OUTSIDE its try block, so an exception there would have
propagated uncaught; now it's caught too.
Tools that do their own internal per-item exception handling
(dense_evolution_energy_scan, dense_evolution_wormhole_scan -- a single
failed point shouldn't abort the whole batch) are unaffected by this:
it only adds an outer safety net around exceptions those inner
try/excepts don't already catch (e.g. a bug in the post-loop
aggregation code), it does not change their partial-failure behavior.
"""
@functools.wraps(func)
async def wrapper(*args, **kwargs):
try:
return await func(*args, **kwargs)
except Exception as e:
return _handle_error(e)
return wrapper
|