feat(chat): add response checkpoints (#16040)

backend/alembic/versions/6b4e821c0a93_add_chat_response_checkpoints.py

@@ -0,0 +1,34 @@
+"""Store resumable response checkpoints.
+
+Revision ID: 6b4e821c0a93
+Revises: 7a03b6e90c12
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+revision = "6b4e821c0a93"
+down_revision = "7a03b6e90c12"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "chat_response_checkpoint",
+        sa.Column(
+            "chat_message_id",
+            sa.Integer(),
+            sa.ForeignKey("chat_message.id", ondelete="CASCADE"),
+            primary_key=True,
+        ),
+        sa.Column(
+            "ownership_version", sa.BigInteger(), server_default="0", nullable=False
+        ),
+        sa.Column("state", postgresql.JSONB(), nullable=False),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("chat_response_checkpoint")

backend/tests/external_dependency_unit/db/test_response_checkpoint_migration.py

@@ -0,0 +1,89 @@
+"""Checkpoint migration leaves response items and current writers intact."""
+
+from pathlib import Path
+from uuid import uuid4
+
+from alembic.migration import MigrationContext
+from alembic.operations import Operations
+from alembic.script import ScriptDirectory
+from sqlalchemy import inspect, text
+from sqlalchemy.orm import Session
+from sqlalchemy.schema import CreateSchema, DropSchema
+
+
+def test_checkpoint_upgrade_downgrade_upgrade(db_session: Session) -> None:
+    schema = f"response_checkpoint_{uuid4().hex}"
+    with db_session.connection().engine.connect() as connection:
+        try:
+            connection.execute(CreateSchema(schema))
+            connection.execute(
+                text("SELECT set_config('search_path', :schema, false)"),
+                {"schema": schema},
+            )
+            connection.execute(
+                text("""
+                CREATE TABLE chat_message (id integer PRIMARY KEY);
+                CREATE TABLE tool_call (id integer PRIMARY KEY);
+                INSERT INTO chat_message VALUES (1), (2);
+            """)
+            )
+            scripts = ScriptDirectory(
+                str(Path(__file__).resolve().parents[3] / "alembic")
+            )
+            items = scripts.get_revision("7a03b6e90c12")
+            checkpoint = scripts.get_revision("6b4e821c0a93")
+            assert items is not None and checkpoint is not None
+            with Operations.context(MigrationContext.configure(connection)):
+                items.module.upgrade()
+                connection.execute(
+                    text("""
+                    INSERT INTO chat_response_item (id, chat_message_id, step_index, generation_status, message)
+                    VALUES ('retained', 2, 0, 'complete', '{}');
+                """)
+                )
+                checkpoint.module.upgrade()
+                connection.execute(
+                    text("""
+                    INSERT INTO chat_response_checkpoint (chat_message_id, state)
+                    VALUES (1, '{"runtime_revision": 7}');
+                """)
+                )
+                assert connection.execute(
+                    text(
+                        "SELECT ownership_version, state FROM chat_response_checkpoint"
+                    )
+                ).one() == (0, {"runtime_revision": 7})
+                assert {
+                    column["name"]
+                    for column in inspect(connection).get_columns(
+                        "chat_message", schema=schema
+                    )
+                } == {"id"}
+                connection.execute(text("DELETE FROM chat_message WHERE id=1"))
+                assert (
+                    connection.scalar(
+                        text("SELECT count(*) FROM chat_response_checkpoint")
+                    )
+                    == 0
+                )
+                checkpoint.module.downgrade()
+                assert "chat_response_checkpoint" not in inspect(
+                    connection
+                ).get_table_names(schema=schema)
+                assert (
+                    connection.scalar(text("SELECT id FROM chat_response_item"))
+                    == "retained"
+                )
+                checkpoint.module.upgrade()
+                assert "chat_response_checkpoint" in inspect(
+                    connection
+                ).get_table_names(schema=schema)
+                assert (
+                    connection.scalar(text("SELECT id FROM chat_response_item"))
+                    == "retained"
+                )
+        finally:
+            connection.rollback()
+            connection.execute(DropSchema(schema, cascade=True, if_exists=True))
+            connection.execute(text("RESET search_path"))
+            connection.commit()