From e4264499e3f99770a7920580fe51fef37209dcc1 Mon Sep 17 00:00:00 2001 From: alkinun Date: Sat, 18 Jul 2026 19:12:50 +0300 Subject: [PATCH] Studio: protect durable research turns --- studio/backend/routes/chat_history.py | 3 +- studio/backend/storage/studio_db.py | 23 +++++- .../backend/tests/test_chat_history_routes.py | 25 ++++++- .../tests/test_research_runs_storage.py | 22 ++++++ .../src/components/assistant-ui/thread.tsx | 73 ++++++++----------- .../test_deep_research_frontend_contract.py | 9 ++- 6 files changed, 104 insertions(+), 51 deletions(-) diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 7a27a58a52..6113f48c3b 100644 --- a/studio/backend/routes/chat_history.py +++ b/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, diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index f870052b72..e9acd448de 100644 --- a/studio/backend/storage/studio_db.py +++ b/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: diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py index a60ac700bf..7eb26d0e3c 100644 --- a/studio/backend/tests/test_chat_history_routes.py +++ b/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}" ) diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index 9613400116..061767408c 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/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) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 612cf1c1c1..15ffed253a 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/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 ( {/* Recovery path for interrupted/failed turns: regenerate in place. */} - {!research.runId && ( + {!researchRunId && (