From c689b10c0da4b002ad976b554f7e7a9a6fd10446 Mon Sep 17 00:00:00 2001
From: RockChinQ <rockchinq@gmail.com>
Date: Sun, 21 Jun 2026 12:04:37 -0400
Subject: [PATCH] fix(mcp): ruff format remote-mode files; make migration head
 test revision-agnostic

CI follow-up to the local/remote MCP work:

- Apply ruff format to provider/tools/loaders/mcp.py and the 0006
  normalize-remote-mode migration (Lint job failed on formatting).
- test_migrations.py hardcoded the head revision as 0005_*, which broke
  once 0006 landed. Resolve the actual head from the Alembic
  ScriptDirectory so future migrations don't require editing the test.
---
 .../0006_normalize_mcp_remote_mode.py         |  8 ++------
 src/langbot/pkg/provider/tools/loaders/mcp.py |  4 +---
 .../persistence/test_migrations.py            | 20 +++++++++++++++++--
 3 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/src/langbot/pkg/persistence/alembic/versions/0006_normalize_mcp_remote_mode.py b/src/langbot/pkg/persistence/alembic/versions/0006_normalize_mcp_remote_mode.py
index 1b5410d95..61a2669e8 100644
--- a/src/langbot/pkg/persistence/alembic/versions/0006_normalize_mcp_remote_mode.py
+++ b/src/langbot/pkg/persistence/alembic/versions/0006_normalize_mcp_remote_mode.py
@@ -31,9 +31,7 @@ def upgrade() -> None:
     inspector = sa.inspect(conn)
     if 'mcp_servers' not in inspector.get_table_names():
         return
-    conn.execute(
-        sa.text("UPDATE mcp_servers SET mode = 'remote' WHERE mode IN ('sse', 'http')")
-    )
+    conn.execute(sa.text("UPDATE mcp_servers SET mode = 'remote' WHERE mode IN ('sse', 'http')"))
 
 
 def downgrade() -> None:
@@ -46,6 +44,4 @@ def downgrade() -> None:
     inspector = sa.inspect(conn)
     if 'mcp_servers' not in inspector.get_table_names():
         return
-    conn.execute(
-        sa.text("UPDATE mcp_servers SET mode = 'http' WHERE mode = 'remote'")
-    )
+    conn.execute(sa.text("UPDATE mcp_servers SET mode = 'http' WHERE mode = 'remote'"))
diff --git a/src/langbot/pkg/provider/tools/loaders/mcp.py b/src/langbot/pkg/provider/tools/loaders/mcp.py
index 8ef7d85e7..2cc83b1c7 100644
--- a/src/langbot/pkg/provider/tools/loaders/mcp.py
+++ b/src/langbot/pkg/provider/tools/loaders/mcp.py
@@ -191,9 +191,7 @@ class RuntimeMCPSession:
         try:
             await self.exit_stack.aclose()
         except Exception as cleanup_err:
-            self.ap.logger.debug(
-                f'MCP server {self.server_name}: error cleaning up before SSE fallback: {cleanup_err}'
-            )
+            self.ap.logger.debug(f'MCP server {self.server_name}: error cleaning up before SSE fallback: {cleanup_err}')
         self.exit_stack = AsyncExitStack()
         self.session = None
 
diff --git a/tests/integration/persistence/test_migrations.py b/tests/integration/persistence/test_migrations.py
index f9872f829..25d3e5c82 100644
--- a/tests/integration/persistence/test_migrations.py
+++ b/tests/integration/persistence/test_migrations.py
@@ -17,7 +17,21 @@ from langbot.pkg.persistence.alembic_runner import (
     run_alembic_upgrade,
     run_alembic_stamp,
     get_alembic_current,
+    _ALEMBIC_DIR,
 )
+from alembic.config import Config
+from alembic.script import ScriptDirectory
+
+
+def _get_script_head() -> str:
+    """Resolve the current Alembic head revision from the script directory.
+
+    Avoids hardcoding a revision number in assertions so adding a new
+    migration doesn't require editing the migration tests.
+    """
+    cfg = Config()
+    cfg.set_main_option('script_location', _ALEMBIC_DIR)
+    return ScriptDirectory.from_config(cfg).get_current_head()
 
 
 pytestmark = pytest.mark.integration
@@ -103,8 +117,10 @@ class TestSQLiteMigrationUpgrade:
         # Verify revision
         rev = await get_alembic_current(sqlite_engine)
         assert rev is not None, 'Expected a revision after upgrade'
-        # Head should be the latest migration
-        assert rev.startswith('0005'), f'Expected head to be 0005_*, got {rev}'
+        # Head should be the latest migration. Resolve the actual head from the
+        # Alembic script directory instead of hardcoding a revision number, so
+        # adding a new migration doesn't require editing this assertion.
+        assert rev == _get_script_head(), f'Expected head {_get_script_head()}, got {rev}'
 
     @pytest.mark.asyncio
     async def test_upgrade_idempotent(self, sqlite_engine):