Skip to content

Commit 7e75290

Browse files
benjibcxzrderek
andauthored
Cache airline flight DB and harden adapters (#191)
* Cache airline flight DB and harden adapters * revert langfuse / langsmith changes * try diff prots * test * hard code port --------- Co-authored-by: Derek Xu <xzrderek@gmail.com>
1 parent 057e132 commit 7e75290

File tree

2 files changed

+31
-12
lines changed

2 files changed

+31
-12
lines changed

eval_protocol/mcp_servers/tau2/airplane_environment/airline_environment.py

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,7 @@
99
import json
1010
import logging
1111
import os
12-
import time
13-
from copy import deepcopy
12+
from functools import lru_cache
1413
from dataclasses import dataclass, field
1514
from enum import Enum
1615
from pathlib import Path
@@ -24,6 +23,16 @@
2423
from vendor.tau2.domains.airline.utils import AIRLINE_DB_PATH
2524

2625

26+
@lru_cache(maxsize=1)
27+
def _load_flight_db(path: str) -> FlightDB:
28+
"""Load and cache the flight database for reuse across resets."""
29+
30+
logger.info("🗂️ Loading airline database from disk (cached)")
31+
db_loaded = FlightDB.load(path)
32+
assert isinstance(db_loaded, FlightDB)
33+
return db_loaded
34+
35+
2736
class AirlineEnvironment:
2837
"""
2938
Airline environment that integrates τ²-Bench simulation pattern
@@ -37,13 +46,10 @@ def __init__(self, config: Optional[Dict[str, Any]] = None):
3746

3847
def reset(self, seed: Optional[int] = None) -> Tuple[Dict[str, Any], Dict[str, Any]]:
3948
"""Reset the environment to initial state"""
40-
logger.info("🔄 Resetting airline environment - reloading database from disk")
41-
# FlightDB.load expects a str path
42-
# Ensure type matches expected FlightDB
43-
# FlightDB.load returns vendor.tau2.domains.airline.data_model.FlightDB which is compatible
44-
db_loaded = FlightDB.load(str(AIRLINE_DB_PATH))
45-
assert isinstance(db_loaded, FlightDB)
46-
self.db = db_loaded
49+
logger.info("🔄 Resetting airline environment - using cached airline database")
50+
cached_db = _load_flight_db(str(AIRLINE_DB_PATH))
51+
# Provide a fresh copy for each environment reset without re-reading from disk.
52+
self.db = cached_db.model_copy(deep=True)
4753
self.airline_tools = AirlineTools(self.db)
4854

4955
return {}, {}

tests/pytest/test_mcp_session_autocreate.py

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616
def _run_airline_server():
1717
import os
1818

19-
os.environ["PORT"] = "9780"
19+
python_version = os.environ.get("PYTHON_VERSION", "3.10").replace(".", "")
20+
port = str(9780 + int(python_version[-1:]))
21+
os.environ["PORT"] = port
2022
from eval_protocol.mcp_servers.tau2.tau2_mcp import AirlineDomainMcp
2123

2224
server = AirlineDomainMcp(seed=None)
@@ -25,24 +27,35 @@ def _run_airline_server():
2527

2628
@pytest.mark.asyncio
2729
async def test_tool_call_returns_json_without_prior_initial_state():
30+
import os
31+
2832
proc = Process(target=_run_airline_server, daemon=True)
2933
proc.start()
3034

3135
try:
32-
base_url = "http://127.0.0.1:9780/mcp"
36+
python_version = os.environ.get("PYTHON_VERSION", "3.10").replace(".", "")
37+
port = str(9780 + int(python_version[-1:]))
38+
39+
base_url = f"http://127.0.0.1:{port}/mcp"
3340
client = httpx.Client(timeout=1.0)
34-
deadline = time.time() + 20
41+
start_time = time.time()
42+
deadline = start_time + 20
43+
ready_time = None
3544
while time.time() < deadline:
3645
try:
3746
r = client.get(base_url)
3847
if r.status_code in (200, 307, 406):
48+
ready_time = time.time()
3949
break
4050
except Exception:
4151
pass
4252
time.sleep(0.2)
4353
else:
4454
pytest.fail("Server did not start on port 9780 in time")
4555

56+
assert ready_time is not None, "Server did not return a successful status before exiting loop"
57+
assert ready_time - start_time < 20, f"Server took too long to respond: {ready_time - start_time:.2f}s"
58+
4659
session = MCPSession(base_url=base_url, session_id="test-autocreate", seed=None, model_id="test-model")
4760

4861
mgr = MCPConnectionManager()

0 commit comments

Comments
 (0)