Studio: protect durable research turns

studio/backend/routes/chat_history.py

@@ -15,6 +15,7 @@ from loggers import get_logger
 from utils.utils import safe_curated_detail, log_and_http_error
 from storage.studio_db import (
     ChatMessageConflictError,
+    ChatMessageProtectedError,
     CorruptSettingsError,
     clear_chat_history,
     count_chat_threads,
@@ -459,7 +460,7 @@ async def replace_thread_messages(
                 )
             ]
         )
-    except ChatMessageConflictError as exc:
+    except (ChatMessageConflictError, ChatMessageProtectedError) as exc:
         raise log_and_http_error(
             exc,
             409,

studio/backend/storage/studio_db.py

@@ -1540,6 +1540,10 @@ class ChatMessageConflictError(RuntimeError):
     """Raised when a chat message id already belongs to another thread."""
 
 
+class ChatMessageProtectedError(RuntimeError):
+    """Raised when pruning would remove a message owned by a durable feature."""
+
+
 class CorruptSettingsError(RuntimeError):
     """Raised when a partial settings patch would overwrite corrupt settings."""
 
@@ -1744,9 +1748,24 @@ def sync_chat_messages(
                     "SELECT id FROM chat_messages WHERE thread_id = ?", (thread_id,)
                 ).fetchall()
             }
+            removed_ids = existing_ids - survivor_ids
+            research_message_ids = {
+                str(message_id)
+                for row in conn.execute(
+                    """SELECT user_message_id, assistant_message_id
+                       FROM research_runs WHERE thread_id = ?""",
+                    (thread_id,),
+                ).fetchall()
+                for message_id in row
+                if message_id is not None
+            }
+            if removed_ids & research_message_ids:
+                raise ChatMessageProtectedError(
+                    "Research prompts and responses cannot be deleted from their original thread"
+                )
             conn.executemany(
                 "DELETE FROM chat_messages WHERE thread_id = ? AND id = ?",
-                [(thread_id, message_id) for message_id in existing_ids - survivor_ids],
+                [(thread_id, message_id) for message_id in removed_ids],
             )
             _recompute_chat_thread_updated_at(conn, thread_id)
         elif messages:
@@ -1755,7 +1774,7 @@ def sync_chat_messages(
             )
         conn.commit()
         return list_chat_messages(thread_id)
-    except ChatMessageConflictError:
+    except (ChatMessageConflictError, ChatMessageProtectedError):
         conn.rollback()
         raise
     except sqlite3.Error:

studio/backend/tests/test_chat_history_routes.py

@@ -57,6 +57,29 @@ def test_replace_thread_messages_rejects_body_thread_mismatch(monkeypatch):
     assert called is False
 
 
+def test_replace_thread_messages_reports_protected_research_turn(monkeypatch):
+    monkeypatch.setattr(chat_history, "get_chat_thread", lambda _thread_id: {"id": "thread-1"})
+
+    def reject_prune(*_args, **_kwargs):
+        raise chat_history.ChatMessageProtectedError(
+            "Research prompts and responses cannot be deleted from their original thread"
+        )
+
+    monkeypatch.setattr(chat_history, "sync_chat_messages", reject_prune)
+
+    with pytest.raises(HTTPException) as exc_info:
+        asyncio.run(
+            chat_history.replace_thread_messages(
+                "thread-1",
+                chat_history.ChatMessageSyncRequest(messages = [], pruneMissing = True),
+                current_subject = "test-user",
+            )
+        )
+
+    assert exc_info.value.status_code == 409
+    assert "Research prompts and responses" in str(exc_info.value.detail)
+
+
 # ---------------------------------------------------------------------------
 # /api/chat/settings
 # ---------------------------------------------------------------------------
@@ -126,7 +149,7 @@ def test_chat_inference_settings_covers_frontend_persisted_fields():
 
     backend = set(chat_history.ChatInferenceSettings.model_fields)
     assert persisted == backend, (
-        f"schema drift: frontend-only {persisted - backend}, " f"backend-only {backend - persisted}"
+        f"schema drift: frontend-only {persisted - backend}, backend-only {backend - persisted}"
     )
 
 

studio/backend/tests/test_research_runs_storage.py

@@ -417,6 +417,28 @@ def test_pruning_messages_preserves_runs_whose_user_message_survives(research_ho
     assert studio_db.get_chat_message("thread-1", "temporary") is None
 
 
+@pytest.mark.parametrize("removed_id", ["user-1", "assistant-1"])
+def test_pruning_rejects_deleting_research_turn_messages(research_home, removed_id):
+    _create()
+    plan = research_db.set_plan("run-1", _plan(), expected_revision = 0)
+    research_db.approve("run-1", 1, plan["planHash"])
+    research_db.claim_next("worker-1")
+    research_db.finish("run-1", "worker-1", "completed")
+    survivors = [
+        message
+        for message in studio_db.list_chat_messages("thread-1")
+        if message["id"] != removed_id
+    ]
+
+    with pytest.raises(studio_db.ChatMessageProtectedError, match = "cannot be deleted"):
+        studio_db.sync_chat_messages("thread-1", survivors, prune_missing = True)
+
+    assert research_db.get_run("run-1") is not None
+    assert research_db.has_thread_claim("thread-1") is True
+    assert studio_db.get_chat_message("thread-1", "user-1") is not None
+    assert studio_db.get_chat_message("thread-1", "assistant-1") is not None
+
+
 def test_revision_hash_conflicts_and_idempotent_approval(research_home):
     _create()
     first = research_db.set_plan("run-1", _plan(), expected_revision = 0)

studio/frontend/src/components/assistant-ui/thread.tsx

@@ -3614,13 +3614,13 @@ const ComposerRightControls: FC<{
 };
 
 const MessageError: FC = () => {
-  const research = useResearchMessageState();
+  const researchRunId = useResearchMessageRunId();
   return (
     <MessagePrimitive.Error>
       <ErrorPrimitive.Root className="aui-message-error-root mt-2 flex flex-wrap items-center gap-x-3 gap-y-2 rounded-md bg-destructive/10 p-3 text-destructive text-sm dark:bg-destructive/5 dark:text-red-200">
         <ErrorPrimitive.Message className="aui-message-error-message line-clamp-2 min-w-0 flex-1" />
         {/* Recovery path for interrupted/failed turns: regenerate in place. */}
-        {!research.runId && (
+        {!researchRunId && (
           <ActionBarPrimitive.Reload asChild={true}>
             <button
               type="button"
@@ -3978,49 +3978,36 @@ const ForkMessageButton: FC = () => {
   );
 };
 
-const TERMINAL_RESEARCH_MESSAGE_STATUSES = new Set([
-  "cancelled",
-  "completed",
-  "failed",
-]);
+const getResearchRunId = (metadata: unknown): string | null => {
+  const custom = (
+    metadata as
+      | {
+          custom?: {
+            researchRunId?: unknown;
+            researchRun?: { id?: unknown };
+          };
+        }
+      | undefined
+  )?.custom;
+  const runId = custom?.researchRunId ?? custom?.researchRun?.id;
+  return typeof runId === "string" ? runId : null;
+};
 
-const useResearchMessageState = () => {
-  const metadata = useAuiState(({ message }) =>
-    (
-      message.metadata as
-        | {
-            custom?: {
-              researchRunId?: unknown;
-              researchStatus?: unknown;
-              researchRun?: { id?: unknown; status?: unknown };
-            };
-          }
-        | undefined
-    )?.custom,
-  );
-  const metadataRunId = metadata?.researchRunId ?? metadata?.researchRun?.id;
-  const runId = typeof metadataRunId === "string" ? metadataRunId : null;
-  const followedStatus = useResearchRunStore((state) =>
-    runId ? state.sessions[runId]?.run.status : undefined,
-  );
-  const metadataStatus = metadata?.researchStatus ?? metadata?.researchRun?.status;
-  return {
-    runId,
-    status:
-      followedStatus ??
-      (typeof metadataStatus === "string" ? metadataStatus : null),
-  };
+const useResearchMessageRunId = () => {
+  return useAuiState(({ message }) => getResearchRunId(message.metadata));
 };
 
 const DeleteMessageButton: FC = () => {
   const aui = useAui();
   const messageId = useAuiState(({ message }) => message.id);
   const isRunning = useAuiState(({ thread }) => thread.isRunning);
-  const research = useResearchMessageState();
-  const isActiveResearchMessage = Boolean(
-    research.runId &&
-      (!research.status ||
-        !TERMINAL_RESEARCH_MESSAGE_STATUSES.has(research.status)),
+  const researchRunId = useResearchMessageRunId();
+  const ownsResearchMessage = aui
+    .thread()
+    .export()
+    .messages.some(
+      ({ parentId, message }) =>
+        parentId === messageId && Boolean(getResearchRunId(message.metadata)),
   );
 
   const handleDelete = async () => {
@@ -4066,7 +4053,7 @@ const DeleteMessageButton: FC = () => {
     }
   };
 
-  if (isActiveResearchMessage) {
+  if (researchRunId || ownsResearchMessage) {
     return null;
   }
 
@@ -4118,11 +4105,11 @@ const CopyButton: FC = () => {
 
 const EditAssistantMessageButton: FC = () => {
   const messageId = useAuiState(({ message }) => message.id);
-  const research = useResearchMessageState();
+  const researchRunId = useResearchMessageRunId();
   const isRunning = useAuiState(({ thread }) => thread.isRunning);
   const setEditingId = useChatRuntimeStore((s) => s.setEditingMessageId);
 
-  if (research.runId) return null;
+  if (researchRunId) return null;
 
   return (
     <TooltipIconButton
@@ -4141,7 +4128,7 @@ const EditAssistantMessageButton: FC = () => {
 
 const AssistantActionBar: FC = () => {
   const { forkMessage, forkDisabled } = useForkMessageAction();
-  const research = useResearchMessageState();
+  const researchRunId = useResearchMessageRunId();
   const [detailsOpen, setDetailsOpen] = useState(false);
   const ttsEnabled = useVoiceSettingsStore((s) => s.ttsEnabled);
   // hideWhenRunning is thread-level, so a new run would hide this bar and its
@@ -4156,7 +4143,7 @@ const AssistantActionBar: FC = () => {
       >
         <CopyButton />
         <EditAssistantMessageButton />
-        {!research.runId && (
+        {!researchRunId && (
           <ActionBarPrimitive.Reload asChild={true}>
             <TooltipIconButton tooltip="Refresh">
               <RefreshCwIcon strokeWidth={1.75} className="size-icon" />

tests/studio/test_deep_research_frontend_contract.py

@@ -94,13 +94,14 @@ def test_research_presentation_is_integrated() -> None:
     research_gate = thread.split("const researchDisabled =", 1)[1].split(";", 1)[0]
     assert "!modelLoaded" not in research_gate
     assert "<ResearchMessage />" in thread
-    assert "if (research.runId) return null" in thread
-    assert "!research.runId &&" in thread
-    assert "if (isActiveResearchMessage)" in thread
+    assert "if (researchRunId) return null" in thread
+    assert "!researchRunId &&" in thread
+    assert "if (researchRunId || ownsResearchMessage)" in thread
+    assert "parentId === messageId && Boolean(getResearchRunId(message.metadata))" in thread
     message_error = thread.split("const MessageError: FC = () =>", 1)[1].split(
         "const GeneratingIndicator:", 1
     )[0]
-    assert "!research.runId &&" in message_error
+    assert "!researchRunId &&" in message_error
     assert "ResearchActivityPanel" in page
     assert "ResearchActivitySheet" in page
     assert "ResearchActivityPanel" in chat_index