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):