From e09f0844270303b4545371d47bed21a714ef62eb Mon Sep 17 00:00:00 2001 From: Egor Kraev Date: Tue, 7 Jul 2026 16:54:11 +0200 Subject: [PATCH] Expose per-task MCP query engine for teardown (DEV-1656) create_mcp_server now attaches its closure SlayerQueryEngine to the returned FastMCP server as mcp._slayer_engine, the cross-repo contract bird-interact-agents disposes at async task teardown: engine = getattr(mcp, "_slayer_engine", None) if engine is not None: await engine.aclose() This lets the cloud Ray runner (one actor reused across tasks) release per-task asyncpg pools that would otherwise accumulate until Postgres max_connections is exhausted. validate_models and recommend_root_model now reuse the closure engine instead of constructing a fresh per-call SlayerQueryEngine, so a single engine holds every cached SQL client for the server's lifetime. schema_drift._collect_sql_diffs now caches the SlayerSQLClient it opens for sql-mode drift probes back into the shared sql_clients dict, so the asyncpg pool it opens is reachable by SlayerQueryEngine.aclose() and disposed at teardown (previously it was constructed standalone and leaked). When no engine dict is provided the behaviour is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- slayer/engine/schema_drift.py | 10 +- slayer/mcp/server.py | 17 +- tests/test_mcp_engine_teardown.py | 317 ++++++++++++++++++++++++++++++ 3 files changed, 341 insertions(+), 3 deletions(-) create mode 100644 tests/test_mcp_engine_teardown.py diff --git a/slayer/engine/schema_drift.py b/slayer/engine/schema_drift.py index 2600ff08..9093c16e 100644 --- a/slayer/engine/schema_drift.py +++ b/slayer/engine/schema_drift.py @@ -2035,9 +2035,17 @@ async def _collect_sql_diffs( # distinct clients. Mirror that key shape here via the shared # ``_sql_client_cache_key`` helper. from slayer.engine.query_engine import _sql_client_cache_key # noqa: PLC0415 - client = (sql_clients or {}).get(_sql_client_cache_key(datasource)) + key = _sql_client_cache_key(datasource) + client = (sql_clients or {}).get(key) if client is None: client = SlayerSQLClient(datasource=datasource) + # DEV-1656: cache the client back into the shared engine dict so the + # asyncpg pool it opens (trial-execute of sql-mode models) is + # reachable by ``SlayerQueryEngine.aclose()`` and disposed at task + # teardown. When ``sql_clients`` is None (no engine — direct/test + # callers own the lifecycle), behaviour is unchanged. + if sql_clients is not None: + sql_clients[key] = client async def _diff_one(model: SlayerModel) -> None: live_cols = await _live_columns_for_sql_model(model=model, client=client) diff --git a/slayer/mcp/server.py b/slayer/mcp/server.py index 75bcb2a1..098bae49 100644 --- a/slayer/mcp/server.py +++ b/slayer/mcp/server.py @@ -278,6 +278,17 @@ def create_mcp_server( # NOSONAR(S3776) — FastMCP tool-registration factory; ), ) engine = SlayerQueryEngine(storage=storage) + # DEV-1656: expose the closure engine so callers (bird-interact-agents on + # the cloud Ray runner, where one actor process is reused across many + # tasks) can dispose its per-task asyncpg pools at task teardown: + # engine = getattr(mcp, "_slayer_engine", None) + # if engine is not None: + # await engine.aclose() # loop-bound; run before the task loop closes + # aclose() is idempotent and leaves the engine reusable (a later execute + # lazily recreates the async engine). The read-only introspection tools + # (validate_models / recommend_root_model) reuse this same engine so a + # single engine holds every cached SQL client for the server's lifetime. + mcp._slayer_engine = engine _help_description = ( "Return conceptual help on SLayer. " @@ -1619,7 +1630,9 @@ async def validate_models(data_source: str | None = None) -> str: ds = await storage.get_datasource(data_source) if ds is None: return f"Datasource '{data_source}' not found." - engine = SlayerQueryEngine(storage=storage) + # DEV-1656: reuse the closure engine (not a fresh per-call engine) so + # the schema-drift SQL client it opens is cached on the server's + # engine and disposed by ``mcp._slayer_engine.aclose()`` at teardown. try: entries = await engine.validate_models(data_source=data_source) except (sa.exc.OperationalError, sa.exc.DatabaseError) as exc: @@ -1668,7 +1681,7 @@ async def recommend_root_model( f"recommend_root_model failed: unknown format '{format}'. " f"Use 'markdown' or 'json'." ) - engine = SlayerQueryEngine(storage=storage) + # DEV-1656: reuse the closure engine (see validate_models above). try: rec = await engine.recommend_root_model( items, data_source=data_source, root_hint=root_hint diff --git a/tests/test_mcp_engine_teardown.py b/tests/test_mcp_engine_teardown.py new file mode 100644 index 00000000..707f3101 --- /dev/null +++ b/tests/test_mcp_engine_teardown.py @@ -0,0 +1,317 @@ +"""DEV-1656: the SLayer MCP server exposes its per-task query engine for +caller-controlled teardown. + +``create_mcp_server`` attaches its closure ``SlayerQueryEngine`` to the +returned FastMCP object as ``mcp._slayer_engine`` — the cross-repo contract +bird-interact-agents disposes at task teardown: + + engine = getattr(mcp, "_slayer_engine", None) + if engine is not None: + await engine.aclose() + +The two read-only introspection tools (``validate_models`` / +``recommend_root_model``) reuse that closure engine instead of constructing +their own per-call engine, so a single engine holds every cached SQL client +for the server's lifetime and one ``aclose()`` disposes them all — including +the client the ``sql``-mode schema-drift path opens. + +These tests build a *fresh* server per test (never the session-scoped shared +fixture in ``test_mcp_server.py``) so mutating / disposing the engine can't +leak across tests. +""" + +from __future__ import annotations + +import json +import sqlite3 +import tempfile +from pathlib import Path +from typing import Any +from collections.abc import Iterator +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import slayer.mcp.server as srv +from slayer.core.enums import DataType +from slayer.core.models import Column, DatasourceConfig, SlayerModel +from slayer.engine.query_engine import SlayerQueryEngine, _sql_client_cache_key +from slayer.mcp.server import create_mcp_server +from slayer.storage.yaml_storage import YAMLStorage + + +@pytest.fixture +def workspace() -> Iterator[Path]: + tmp = tempfile.TemporaryDirectory() + try: + yield Path(tmp.name) + finally: + tmp.cleanup() + + +def _make_sqlite_db(workspace: Path) -> Path: + db = workspace / "live.db" + conn = sqlite3.connect(db) + conn.execute("CREATE TABLE t (id INTEGER PRIMARY KEY, v INTEGER)") + conn.execute("INSERT INTO t (v) VALUES (10), (20), (30)") + conn.commit() + conn.close() + return db + + +def _lite_datasource(db: Path) -> DatasourceConfig: + return DatasourceConfig(name="lite", type="sqlite", database=str(db)) + + +async def _seed_storage(workspace: Path, *, with_sql_mode: bool = False) -> YAMLStorage: + """A YAML store with a ``lite`` sqlite datasource, a ``sql_table`` model + ``t`` and (optionally) a ``sql``-mode model ``sv`` over the same table.""" + db = _make_sqlite_db(workspace) + storage = YAMLStorage(base_dir=str(workspace / "store")) + await storage.save_datasource(_lite_datasource(db)) + await storage.save_model(SlayerModel( + name="t", data_source="lite", sql_table="t", + columns=[ + Column(name="id", sql="id", type=DataType.INT, primary_key=True), + Column(name="v", sql="v", type=DataType.INT), + ], + )) + if with_sql_mode: + await storage.save_model(SlayerModel( + name="sv", data_source="lite", sql="SELECT id, v FROM t", + columns=[ + Column(name="id", sql="id", type=DataType.INT, primary_key=True), + Column(name="v", sql="v", type=DataType.INT), + ], + )) + return storage + + +async def _call(server, *, name: str, arguments: dict[str, Any] | None = None) -> str: + content_blocks, _ = await server.call_tool(name=name, arguments=arguments or {}) + return content_blocks[0].text + + +async def _query_count(server, *, source_model: str) -> int: + """Run the ``query`` tool for ``*:count`` and return the integer count.""" + text = await _call( + server, name="query", + arguments={ + "source_model": source_model, + "measures": [{"formula": "*:count"}], + "format": "json", + }, + ) + # The query tool's json format appends a human-readable "Measure + # attributes:" footer after the JSON payload — decode just the leading + # JSON value and ignore the trailing text. + rows, _ = json.JSONDecoder().raw_decode(text) + row = rows[0] if isinstance(rows, list) else rows["data"][0] + return int(row[f"{source_model}._count"]) + + +class TestContractAttribute: + async def test_slayer_engine_attribute_is_a_query_engine( + self, workspace: Path, + ) -> None: + storage = await _seed_storage(workspace) + server = create_mcp_server(storage=storage) + assert isinstance(server._slayer_engine, SlayerQueryEngine) + + async def test_slayer_engine_is_the_object_query_uses( + self, workspace: Path, monkeypatch, + ) -> None: + """The exposed attribute is identity-equal to the engine the ``query`` + tool executes on — not just some engine.""" + storage = await _seed_storage(workspace) + server = create_mcp_server(storage=storage) + + captured: dict[str, Any] = {} + original = SlayerQueryEngine.execute + + async def _spy(self, *args, **kwargs): + captured["self"] = self + return await original(self, *args, **kwargs) + + monkeypatch.setattr(SlayerQueryEngine, "execute", _spy) + await _query_count(server, source_model="t") + + assert captured["self"] is server._slayer_engine + + async def test_slayer_engine_is_the_object_query_nested_uses( + self, workspace: Path, monkeypatch, + ) -> None: + """query_nested shares the same closure engine (plan: one engine + across query/query_nested/validate_models/recommend_root_model).""" + storage = await _seed_storage(workspace) + server = create_mcp_server(storage=storage) + + captured: dict[str, Any] = {} + original = SlayerQueryEngine.execute + + async def _spy(self, *args, **kwargs): + captured["self"] = self + return await original(self, *args, **kwargs) + + monkeypatch.setattr(SlayerQueryEngine, "execute", _spy) + await _call( + server, name="query_nested", + arguments={"queries": [{"source_model": "t", "measures": [{"formula": "*:count"}]}]}, + ) + + assert captured["self"] is server._slayer_engine + + +class TestToolsReuseClosureEngine: + async def test_validate_models_constructs_no_new_engine( + self, workspace: Path, monkeypatch, + ) -> None: + storage = await _seed_storage(workspace) + server = create_mcp_server(storage=storage) + + class _Boom: + def __init__(self, *args, **kwargs): + raise AssertionError("validate_models constructed a new engine") + + monkeypatch.setattr(srv, "SlayerQueryEngine", _Boom) + out = await _call(server, name="validate_models", arguments={"data_source": "lite"}) + # Sane output: a JSON list of pending deletes (empty when no drift). + assert isinstance(json.loads(out), list) + + async def test_recommend_root_model_constructs_no_new_engine( + self, workspace: Path, monkeypatch, + ) -> None: + storage = await _seed_storage(workspace) + server = create_mcp_server(storage=storage) + + class _Boom: + def __init__(self, *args, **kwargs): + raise AssertionError("recommend_root_model constructed a new engine") + + monkeypatch.setattr(srv, "SlayerQueryEngine", _Boom) + out = await _call( + server, name="recommend_root_model", + arguments={"items": ["t.v"], "data_source": "lite"}, + ) + assert "t" in out + + +class TestAcloseDisposesCachedClients: + async def test_aclose_disposes_async_engine_sync_pool_survives( + self, workspace: Path, monkeypatch, + ) -> None: + """After a ``query`` caches a client, ``server._slayer_engine.aclose()`` + disposes the (fake) async engine while leaving the sqlite StaticPool + sync engine intact, and a follow-up query still returns the same rows. + File-backed sqlite stands in for ``:memory:`` — the client-cache logic + is identical and ``:memory:`` can't be probed past DEV-1538 introspect. + """ + storage = await _seed_storage(workspace) + server = create_mcp_server(storage=storage) + + assert await _query_count(server, source_model="t") == 3 + engine = server._slayer_engine + assert len(engine._sql_clients) == 1 + client = next(iter(engine._sql_clients.values())) + sync_engine_before = client._sync_engine + + # sqlite never opens an async engine; inject a tracked one so the real + # aclose loop (engine -> client -> dispose) is exercised. + fake = MagicMock(name="fake_async_engine") + fake.dispose = AsyncMock(name="dispose") + client._async_engine = fake + + # Guard: aclose must NOT dispose the sync StaticPool engine — that's + # what pins the (potential :memory:) connection holding the data. + sync_dispose_spy = MagicMock(name="sync_dispose") + monkeypatch.setattr(sync_engine_before, "dispose", sync_dispose_spy) + + await engine.aclose() + + fake.dispose.assert_awaited_once() + assert client._async_engine is None + assert client in engine._sql_clients.values() # client kept + assert client._sync_engine is sync_engine_before # StaticPool untouched + sync_dispose_spy.assert_not_called() + + # Reusable after disposal — the sync path still serves the same data. + assert await _query_count(server, source_model="t") == 3 + + +class TestAcloseNoOpSafeAndReusable: + async def test_aclose_no_async_engine_is_noop_safe( + self, workspace: Path, + ) -> None: + storage = await _seed_storage(workspace) + server = create_mcp_server(storage=storage) + # No query yet — nothing opened. Must not raise. + await server._slayer_engine.aclose() + # And the engine is still usable. + assert await _query_count(server, source_model="t") == 3 + + async def test_query_after_aclose_still_works( + self, workspace: Path, + ) -> None: + storage = await _seed_storage(workspace) + server = create_mcp_server(storage=storage) + assert await _query_count(server, source_model="t") == 3 + await server._slayer_engine.aclose() + # Async engine lazily recreated / sync path intact. + assert await _query_count(server, source_model="t") == 3 + + +class TestSqlModeDriftClientReachable: + """DEV-1656 / Codex finding #1: the ``sql``-mode schema-drift path opens a + ``SlayerSQLClient`` that must be cached on the shared engine so teardown + reaches it (otherwise it leaks an asyncpg pool on Postgres).""" + + async def test_validate_models_caches_sql_mode_client_and_aclose_disposes_it( + self, workspace: Path, + ) -> None: + """A ``validate_models`` call with no prior query must cache the + ``sql``-mode drift client on the shared engine, and a subsequent + ``aclose()`` must dispose *that specific* client's async engine.""" + storage = await _seed_storage(workspace, with_sql_mode=True) + server = create_mcp_server(storage=storage) + engine = server._slayer_engine + assert len(engine._sql_clients) == 0 + + await _call(server, name="validate_models", arguments={"data_source": "lite"}) + + key = _sql_client_cache_key(_lite_datasource(workspace / "live.db")) + assert key in engine._sql_clients, ( + "sql-mode drift client not cached on the engine — aclose() can't " + "reach it, so the pool leaks" + ) + + # Prove teardown actually disposes this cached client (sqlite opens no + # async engine, so inject a tracked one — the Postgres leak path). + client = engine._sql_clients[key] + fake = MagicMock(name="fake_async_engine") + fake.dispose = AsyncMock(name="dispose") + client._async_engine = fake + + await engine.aclose() + + fake.dispose.assert_awaited_once() + assert client._async_engine is None + + async def test_query_then_validate_reuses_same_client( + self, workspace: Path, + ) -> None: + storage = await _seed_storage(workspace, with_sql_mode=True) + server = create_mcp_server(storage=storage) + engine = server._slayer_engine + + # Query the sql-mode model — caches a client under the datasource key. + assert await _query_count(server, source_model="sv") == 3 + assert len(engine._sql_clients) == 1 + key = next(iter(engine._sql_clients.keys())) + client_before = engine._sql_clients[key] + + # validate_models reuses that cached client object; no second one is + # built (a new client under the same key would still leak on teardown + # of the original — assert object identity, not just key equality). + await _call(server, name="validate_models", arguments={"data_source": "lite"}) + assert set(engine._sql_clients.keys()) == {key} + assert engine._sql_clients[key] is client_before