fix(migrations): preserve published Cloud revision head (#2375)

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>

src/langbot/pkg/persistence/alembic/versions/0016_space_launch_replay.py

@@ -0,0 +1,57 @@
+"""add durable replay protection for signed Space launch assertions
+
+Revision ID: 0016_space_launch_replay
+Revises: 0015_cloud_core_collab
+Create Date: 2026-07-31
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+revision = '0016_space_launch_replay'
+down_revision = '0015_cloud_core_collab'
+branch_labels = None
+depends_on = None
+
+_TABLE = 'space_launch_assertion_consumptions'
+_POLICY = 'langbot_directory_projection'
+_SETTING = "NULLIF(current_setting('langbot.directory_instance_uuid', true), '')"
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    if _TABLE not in set(sa.inspect(conn).get_table_names()):
+        op.create_table(
+            _TABLE,
+            sa.Column('instance_uuid', sa.String(255), nullable=False),
+            sa.Column('jti', sa.String(255), nullable=False),
+            sa.Column('expires_at', sa.DateTime(timezone=True), nullable=False),
+            sa.Column('consumed_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
+            sa.PrimaryKeyConstraint('instance_uuid', 'jti'),
+        )
+        op.create_index(
+            'ix_space_launch_assertion_consumptions_expiry',
+            _TABLE,
+            ['instance_uuid', 'expires_at'],
+            unique=False,
+        )
+    if conn.dialect.name == 'postgresql':
+        table = conn.dialect.identifier_preparer.quote(_TABLE)
+        policy = conn.dialect.identifier_preparer.quote(_POLICY)
+        expression = f'instance_uuid::text = {_SETTING}'
+        op.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY'))
+        op.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY'))
+        op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}'))
+        op.execute(
+            sa.text(
+                f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC '
+                f'USING ({expression}) WITH CHECK ({expression})'
+            )
+        )
+
+
+def downgrade() -> None:
+    if _TABLE in set(sa.inspect(op.get_bind()).get_table_names()):
+        op.drop_table(_TABLE)

src/langbot/pkg/persistence/alembic/versions/0018_merge_launch_replay.py

@@ -0,0 +1,21 @@
+"""merge the published Space launch replay and main migration branches
+
+Revision ID: 0018_merge_launch_replay
+Revises: 0016_space_launch_replay, 0017_oss_workspace_identity
+Create Date: 2026-08-01
+"""
+
+from __future__ import annotations
+
+revision = '0018_merge_launch_replay'
+down_revision = ('0016_space_launch_replay', '0017_oss_workspace_identity')
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    pass
+
+
+def downgrade() -> None:
+    pass

src/langbot/pkg/telemetry/heartbeat.py

@@ -37,6 +37,7 @@ class WorkspaceResourceSnapshot(typing.TypedDict):
     extension_count: int
     skill_count: int
     adapters: list[str]
+    execution_generation: int
 
 
 async def _count(
@@ -81,6 +82,7 @@ async def _cloud_workspace_resource_counts(ap: core_app.Application, bindings) -
             'extension_count': 0,
             'skill_count': 0,
             'adapters': [],
+            'execution_generation': binding.placement_generation,
         }
         for binding in bindings
     }

tests/integration/persistence/test_migrations.py

@@ -95,6 +95,18 @@ class TestSQLiteMigrationBaseline:
 class TestSQLiteMigrationUpgrade:
     """Tests for upgrade to head workflow."""
 
+    @pytest.mark.asyncio
+    async def test_upgrade_from_published_space_launch_head_to_merged_head(self, sqlite_engine):
+        """A database released at the production-only 0016 head must remain upgradable."""
+        async with sqlite_engine.begin() as conn:
+            await conn.run_sync(Base.metadata.create_all)
+
+        await run_alembic_stamp(sqlite_engine, '0016_space_launch_replay')
+        await run_alembic_upgrade(sqlite_engine, 'head')
+
+        assert await get_alembic_current(sqlite_engine) == _get_script_head()
+        assert _get_script_head() == '0018_merge_launch_replay'
+
     @pytest.mark.asyncio
     async def test_upgrade_from_baseline_to_head(self, sqlite_engine):
         """

tests/unit_tests/telemetry/test_heartbeat.py

@@ -139,8 +139,8 @@ class TestBuildHeartbeatPayload:
         }
         ap.workspace_service.list_active_execution_bindings = AsyncMock(
             return_value=[
-                SimpleNamespace(workspace_uuid='workspace-a'),
-                SimpleNamespace(workspace_uuid='workspace-b'),
+                SimpleNamespace(workspace_uuid='workspace-a', placement_generation=7),
+                SimpleNamespace(workspace_uuid='workspace-b', placement_generation=9),
             ],
         )
         ap.platform_mgr._bots_by_key[('instance-a', 'workspace-b', 'bot-b')] = SimpleNamespace(
@@ -159,10 +159,12 @@ class TestBuildHeartbeatPayload:
         assert by_workspace['workspace-a']['plugin_count'] == 2
         assert by_workspace['workspace-a']['extension_count'] == 5
         assert by_workspace['workspace-a']['skill_count'] == 2
+        assert by_workspace['workspace-a']['execution_generation'] == 7
         assert by_workspace['workspace-a']['adapters'] == ['WorkspaceAAdapter']
         assert by_workspace['workspace-b']['bot_count'] == 1
         assert by_workspace['workspace-b']['pipeline_count'] == 0
         assert by_workspace['workspace-b']['skill_count'] == 1
+        assert by_workspace['workspace-b']['execution_generation'] == 9
         assert by_workspace['workspace-b']['adapters'] == ['WorkspaceBAdapter']
         assert 'workspace_resources' not in by_workspace['workspace-a']
         ap.persistence_mgr.execute_async.assert_not_awaited()