Protect chat mutations with revision compare-and-swap

backend/cortex_backend/api/routes.py

@@ -27,7 +27,7 @@
     title_from_first_message,
 )
 from cortex_backend.core.settings import CortexSettings
-from cortex_backend.repositories.chats import ChatRepositoryError
+from cortex_backend.repositories.chats import ChatRepositoryError, ChatRevisionConflict
 from cortex_backend.repositories.settings import SettingsMigrationReport
 from cortex_backend.services.models import ModelPullProgress
 from cortex_backend.services.attachments import (
@@ -347,10 +347,13 @@ def add_message(
                 thoughts=payload.thoughts,
                 attachments=[attachment.model_dump(mode="json") for attachment in attachment_refs],
                 thread_title="New Chat" if existing is None else None,
+                expected_revision=payload.base_revision,
             )
             chat = deps.chats.get_chat(thread_id)
         except ChatAttachmentError as exc:
             _raise_chat_attachment_error(exc)
+        except ChatRevisionConflict as exc:
+            raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(exc)) from exc
         except ChatDomainError as exc:
             raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(exc)) from exc
         except Exception as exc:
@@ -1133,11 +1136,19 @@ async def regenerate_chat(
                         raise ChatDomainError(
                             "A regeneration request needs user input."
                         )
+                    current_revision = chat_revision(chat)
+                    if (
+                        payload.base_revision is not None
+                        and payload.base_revision != current_revision
+                    ):
+                        raise ChatDomainError(
+                            "This chat changed. Reload it before regenerating."
+                        )
                     generation_payload = GenerationRequest(
                         request_id=payload.request_id,
                         thread_id=thread_id,
                         user_input=user_input,
-                        base_revision=chat_revision(chat),
+                        base_revision=current_revision,
                         attachments=payload.attachments,
                     )
                     snapshot, _ = await _start_generation_job(
@@ -1164,6 +1175,8 @@ async def regenerate_chat(
             ) from exc
         except JobConflict as exc:
             raise HTTPException(status_code=409, detail=str(exc)) from exc
+        except ChatRevisionConflict as exc:
+            raise HTTPException(status_code=409, detail=str(exc)) from exc
         except ChatDomainError as exc:
             raise HTTPException(status_code=409, detail=str(exc)) from exc
         return _accepted(snapshot)
@@ -1424,6 +1437,7 @@ async def _start_generation_job(
     try:
         chat = await asyncio.to_thread(deps.chats.get_chat, thread_id)
         current_revision = chat_revision(chat) if chat is not None else 0
+        admission_revision = current_revision
         if (
             payload.base_revision is not None
             and current_revision != payload.base_revision
@@ -1475,18 +1489,16 @@ async def _start_generation_job(
         )
 
         user_message_id: str | None = None
+        prepared_revision: int | None = None
         prepared_history = history_messages
 
         def prepare() -> Mapping[str, Any]:
-            nonlocal prepared_history, user_message_id
+            nonlocal prepared_history, prepared_revision, user_message_id
             current_chat = deps.chats.get_chat(thread_id)
             current_revision = (
                 chat_revision(current_chat) if current_chat is not None else 0
             )
-            if (
-                payload.base_revision is not None
-                and current_revision != payload.base_revision
-            ):
+            if current_revision != admission_revision:
                 raise ChatDomainError(
                     "This chat changed. Reload it before generating again."
                 )
@@ -1510,6 +1522,7 @@ def prepare() -> Mapping[str, Any]:
                         "The selected response has no user turn to regenerate."
                     )
                 prepared_history = current_messages[:target_position]
+                prepared_revision = current_revision
             else:
                 user_message_id = deps.chats.add_message(
                     thread_id,
@@ -1520,7 +1533,12 @@ def prepare() -> Mapping[str, Any]:
                         for attachment in attachment_refs
                     ],
                     thread_title="New Chat" if current_chat is None else None,
+                    expected_revision=admission_revision,
                 )
+                updated_chat = deps.chats.get_chat(thread_id)
+                if updated_chat is None:
+                    raise ChatRepositoryError("Chat did not persist the user message.")
+                prepared_revision = chat_revision(updated_chat)
             return {"user_message_id": user_message_id}
 
         def runner(sink, cancel_event):
@@ -1568,13 +1586,15 @@ def runner(sink, cancel_event):
                     "assistant",
                     result.response,
                     thoughts=result.thoughts,
+                    expected_revision=prepared_revision,
                 )
             else:
                 deps.chats.replace_message(
                     thread_id,
                     target_message_id,
                     result.response,
                     thoughts=result.thoughts,
+                    expected_revision=prepared_revision,
                 )
                 assistant_message_id = target_message_id
 

backend/cortex_backend/api/schemas.py

@@ -135,6 +135,7 @@ class RenameChatRequest(APIModel):
 class AddMessageRequest(APIModel):
     role: ChatRole
     content: str = Field(min_length=1, max_length=100_000)
+    base_revision: int | None = Field(default=None, ge=0)
     sources: list[Any] | None = None
     thoughts: str | None = Field(default=None, max_length=100_000)
     attachments: list[ChatAttachment] | None = Field(default=None, max_length=MAX_CHAT_ATTACHMENTS)
@@ -460,6 +461,7 @@ class RegenerationRequest(APIModel):
     request_id: str | None = Field(default=None, min_length=1, max_length=200)
     message_id: str = Field(min_length=1, max_length=200)
     user_input: str | None = Field(default=None, max_length=100_000)
+    base_revision: int | None = Field(default=None, ge=0)
     attachments: list[ChatAttachment] = Field(default_factory=list, max_length=MAX_CHAT_ATTACHMENTS)
 
 

backend/cortex_backend/repositories/__init__.py

@@ -3,6 +3,7 @@
 from .chats import (
     ChatRepository,
     ChatRepositoryError,
+    ChatRevisionConflict,
     InMemoryChatRepository,
     LegacyDatabaseChatRepository,
 )
@@ -24,6 +25,7 @@
 __all__ = [
     "ChatRepository",
     "ChatRepositoryError",
+    "ChatRevisionConflict",
     "InMemoryChatRepository",
     "LegacyDatabaseChatRepository",
     "MemoryRepository",