Expose per-task SLayer MCP query engine for teardown (DEV-1656)

[ZmeiGorynych]

Why

Under --query-mode slayer on a Postgres benchmark, the SLayer MCP query path opens asyncpg connection pools that are never disposed for the life of a long-running task. On the cloud runner one Ray actor process is reused across many tasks, so per-task engines/pools accumulate across tasks until Postgres max_connections is exhausted. This change makes robust cloud cross-task disposal possible (contract for DEV-1654, which is coupled to this issue).

Root of the un-disposability: the MCP query/query_nested tools run the async engine.execute(...) path (which, unlike execute_sync, does not auto-dispose), and the closure engine that owns the pools was not reachable from the returned server.

Changes

1. slayer/mcp/server.py — attach the closure engine to the returned FastMCP server as mcp._slayer_engine (the cross-repo teardown contract). Callers dispose via:

   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; a :memory: SQLite StaticPool survives).

2. slayer/mcp/server.py — validate_models and recommend_root_model reuse the closure engine instead of each constructing a fresh per-call SlayerQueryEngine. One engine now holds every cached SQL client for the server's lifetime, shared across query/query_nested/validate_models/recommend_root_model (all read-only introspection paths).

3. slayer/engine/schema_drift.py — _collect_sql_diffs 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 that client was constructed standalone and leaked even after the closure-engine reuse. When no engine dict is provided (direct/test callers), behaviour is unchanged.

Tests

New tests/test_mcp_engine_teardown.py (10 tests): the _slayer_engine attribute is the same engine query/query_nested use; validate_models/recommend_root_model construct no fresh engine; aclose() disposes cached clients' async engines while leaving the sync StaticPool intact (guarded); aclose() is no-op-safe and leaves the engine reusable; the sql-mode drift client is cached on the shared engine and that specific client is disposed by aclose(); query-then-validate_models reuses the same client object.

Full non-integration suite: 6611 passed, 46 skipped, 5 xfailed. Lint clean.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features
  • Improved cleanup of shared engine resources so background pools and cached clients are released more reliably.
  • Reused a shared engine and client cache across multiple validation and query actions for better consistency.
• Bug Fixes
  • Fixed cases where repeated validation or query actions could create extra engine instances or leave cached SQL resources undisposed.
  • Improved teardown behavior so queries continue working after cleanup and no-op cleanup is safe when no async resources exist.

[linear]

DEV-1656

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: d3bfe208-60d2-40eb-8bf1-9840ece681e2

📥 Commits

Reviewing files that changed from the base of the PR and between 3ca6c68 and e09f084.

📒 Files selected for processing (3)

• slayer/engine/schema_drift.py
• slayer/mcp/server.py
• tests/test_mcp_engine_teardown.py

---

📝 Walkthrough

Walkthrough

This PR updates the SQL trial-execution diffing path to cache SlayerSQLClient instances by key for lifecycle management, modifies the MCP server to store and reuse a shared SlayerQueryEngine instance across validate_models and recommend_root_model tool calls instead of creating new engines per call, and adds a comprehensive test module verifying engine reuse identity and teardown disposal behavior.

Changes

Shared MCP engine reuse and cached SQL client disposal

Layer / File(s)	Summary
SQL client caching in schema drift slayer/engine/schema_drift.py	_collect_sql_diffs computes a shared client cache key and writes newly created SlayerSQLClient instances back into the sql_clients dict when provided.
Shared engine wiring in MCP server slayer/mcp/server.py	create_mcp_server stores the engine as mcp._slayer_engine; validate_models and recommend_root_model reuse this closure engine instead of instantiating new engines.
Teardown and engine-reuse test suite tests/test_mcp_engine_teardown.py	New test module with fixtures and test classes verifying shared engine identity, tool reuse of the closure engine, aclose() disposal of cached async clients (with sqlite sync engine preserved), no-op safety, and drift client caching/reuse.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant MCPServer
  participant SlayerQueryEngine
  participant SlayerSQLClient

  Client->>MCPServer: call validate_models / recommend_root_model
  MCPServer->>SlayerQueryEngine: reuse mcp._slayer_engine
  SlayerQueryEngine->>SlayerSQLClient: get or create cached client via cache key
  SlayerQueryEngine-->>MCPServer: result
  MCPServer-->>Client: response

  Client->>MCPServer: teardown (aclose)
  MCPServer->>SlayerQueryEngine: aclose()
  SlayerQueryEngine->>SlayerSQLClient: dispose cached async engines

Loading

Possibly related PRs

• MotleyAI/slayer#194: Builds on the same SlayerQueryEngine.aclose() disposal and cached SlayerSQLClient lifecycle behavior extended in this PR.

Suggested reviewers: AivanF, whimo

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: exposing the MCP query engine so it can be torn down after each task.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch egor/dev-1656-slayer-mcp-expose-the-per-task-query-engine-for-teardown

---

_{Comment @coderabbitai help to get the list of available commands.}

[sonarqubecloud]

Quality Gate passed

Issues
0 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud