From 0ccbcd5f5fc028b4b57ce4fd202c1bc153890c1c Mon Sep 17 00:00:00 2001
From: Hyu <chenhyudada@gmail.com>
Date: Sun, 2 Aug 2026 00:56:20 +0800
Subject: [PATCH] fix(migrations): preserve published Cloud revision head
 (#2375)

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
---
 .../versions/0016_space_launch_replay.py      | 57 +++++++++++++++++++
 .../versions/0018_merge_launch_replay.py      | 21 +++++++
 .../persistence/test_migrations.py            | 12 ++++
 3 files changed, 90 insertions(+)
 create mode 100644 src/langbot/pkg/persistence/alembic/versions/0016_space_launch_replay.py
 create mode 100644 src/langbot/pkg/persistence/alembic/versions/0018_merge_launch_replay.py

diff --git a/src/langbot/pkg/persistence/alembic/versions/0016_space_launch_replay.py b/src/langbot/pkg/persistence/alembic/versions/0016_space_launch_replay.py
new file mode 100644
index 000000000..33f44c19c
--- /dev/null
+++ b/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)
diff --git a/src/langbot/pkg/persistence/alembic/versions/0018_merge_launch_replay.py b/src/langbot/pkg/persistence/alembic/versions/0018_merge_launch_replay.py
new file mode 100644
index 000000000..68f4e1007
--- /dev/null
+++ b/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
diff --git a/tests/integration/persistence/test_migrations.py b/tests/integration/persistence/test_migrations.py
index c1a455fb1..3447c9320 100644
--- a/tests/integration/persistence/test_migrations.py
+++ b/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):
         """