diff --git a/studio/backend/core/research_runs.py b/studio/backend/core/research_runs.py index a3e8c4c8d9..a815a8ece4 100644 --- a/studio/backend/core/research_runs.py +++ b/studio/backend/core/research_runs.py @@ -52,9 +52,12 @@ _PROMPT_DELIMITER_TAGS = re.compile( re.IGNORECASE, ) _QUERY_CREDENTIAL = re.compile( - r"""(?ix)\b(?:api[\s_-]?key|access[\s_-]?token|password|secret|token)\s*[:=]\s* + r"""(?ix)\b(?:api[\s_-]?key|access[\s_-]?token|authorization|password|secret|token)\s*[:=]\s* (?:"[^"]*"|'[^']*'|“[^”]*”|‘[^’]*’|[^\s,;]+)""" ) +# Bearer authorization tokens carry no key=value label, so the credential pattern above misses +# them; the length floor keeps ordinary prose ("bearer of bad news") from matching. +_QUERY_BEARER = re.compile(r"(?i)\bbearer\s+[A-Za-z0-9._~+/=-]{8,}") _QUERY_EMAIL = re.compile(r"(?i)\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b") _QUERY_PRIVATE_ID = re.compile(r"\b\d{3}-\d{2}-\d{4}\b") _QUERY_OPAQUE_TOKEN = re.compile( @@ -314,6 +317,7 @@ def _shield_untrusted(text: str) -> str: def _sanitize_public_query(query: str) -> str: query = _QUERY_CREDENTIAL.sub(" ", query) + query = _QUERY_BEARER.sub(" ", query) query = _QUERY_EMAIL.sub(" ", query) query = _QUERY_PRIVATE_ID.sub(" ", query) query = _QUERY_OPAQUE_TOKEN.sub(" ", query) @@ -505,17 +509,24 @@ def _bounded_synthesis_evidence( ) -> str: if not notes: return "(none)" + if max_chars <= 0: + return "" + # Split the budget evenly across every note so a small context still keeps a slice of every + # research step. A per-note floor would let the earliest notes consume the whole budget and + # the final slice would drop later steps entirely. separator = "\n\n" - per_note = max( - min(1000, max_chars), - (max_chars - len(separator) * (len(notes) - 1)) // len(notes), - ) + available = max(0, max_chars - len(separator) * (len(notes) - 1)) + base, remainder = divmod(available, len(notes)) + suffix = "\n[Evidence truncated]" bounded = [] - for note in notes: - if len(note) <= per_note: + for index, note in enumerate(notes): + limit = base + (1 if index < remainder else 0) + if len(note) <= limit: bounded.append(note) + elif limit <= len(suffix): + bounded.append(note[:limit]) else: - bounded.append(note[: per_note - 24].rstrip() + "\n[Evidence truncated]") + bounded.append(note[: limit - len(suffix)].rstrip() + suffix) return separator.join(bounded)[:max_chars] @@ -1631,9 +1642,9 @@ class ResearchSupervisor: "role": "user", "content": ( f"Conversation context JSON:\n{_shield_untrusted(conversation_context)}\n\n" - f"Question:\n{question}\n\n" + f"Question:\n{_shield_untrusted(question)}\n\n" f"Approved plan (guidance only):\n" - f"{json.dumps(run['plan'], ensure_ascii = False)}\n\n" + f"{_shield_untrusted(json.dumps(run['plan'], ensure_ascii = False))}\n\n" f"Actions remaining after this one: {max_steps - position - 1}\n" f"\n" f"Gathered sources:\n{_shield_untrusted(source_catalog) or '(none)'}\n\n" @@ -1906,9 +1917,9 @@ class ResearchSupervisor: "content": ( f"\n{_shield_untrusted(conversation_context)}\n" f"\n\n" - f"\n{question}\n" + f"\n{_shield_untrusted(question)}\n" f"\n\n" - f"\n{json.dumps(run['plan'], ensure_ascii = False)}\n" + f"\n{_shield_untrusted(json.dumps(run['plan'], ensure_ascii = False))}\n" f"\n\n" f"\n{_shield_untrusted(source_catalog) or '(no web sources gathered)'}\n" f"\n\n" diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 1619961a2a..a9af01103e 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -403,7 +403,17 @@ def delete_attachment( current_subject: str = Depends(get_current_subject), ) -> dict: """Remove one attachment from its chat message.""" - if not delete_chat_attachment(message_id, attachment_id): + try: + deleted = delete_chat_attachment(message_id, attachment_id) + except ChatMessageProtectedError as exc: + raise log_and_http_error( + exc, + 409, + safe_curated_detail(exc), + event = "chat_history.delete_attachment_conflict", + log = logger, + ) from exc + if not deleted: raise HTTPException(status_code = 404, detail = "Attachment not found") return {"ok": True} diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 43332b0986..a37d749463 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -1926,7 +1926,7 @@ def _research_message_ids(conn: sqlite3.Connection, thread_id: str) -> set[str]: def _research_message_would_change(conn: sqlite3.Connection, thread_id: str, message: dict) -> bool: row = conn.execute( - "SELECT parent_id, role, content_json, metadata_json " + "SELECT parent_id, role, content_json, metadata_json, attachments_json " "FROM chat_messages WHERE thread_id = ? AND id = ?", (thread_id, str(message["id"])), ).fetchone() @@ -1940,6 +1940,8 @@ def _research_message_would_change(conn: sqlite3.Connection, thread_id: str, mes canon(message.get("content", [])) != canon(json.loads(row["content_json"] or "[]")) or canon(message.get("metadata")) != canon(json.loads(row["metadata_json"]) if row["metadata_json"] else None) + or canon(message.get("attachments")) + != canon(json.loads(row["attachments_json"]) if row["attachments_json"] else None) or (message.get("parentId") or None) != (row["parent_id"] or None) or str(message.get("role")) != str(row["role"]) ) @@ -2824,6 +2826,11 @@ def delete_chat_attachment(message_id: str, attachment_id: str) -> bool: if row is None: conn.rollback() return False + if str(message_id) in _research_message_ids(conn, str(row["thread_id"])): + conn.rollback() + raise ChatMessageProtectedError( + "Research prompts and responses are server-managed and cannot be edited" + ) attachments = _json_loads(row["attachments_json"], None) updated_attachments_json = row["attachments_json"] diff --git a/studio/backend/tests/test_research_runs_hardening.py b/studio/backend/tests/test_research_runs_hardening.py index 5d1390a730..e6a4f687b5 100644 --- a/studio/backend/tests/test_research_runs_hardening.py +++ b/studio/backend/tests/test_research_runs_hardening.py @@ -80,6 +80,16 @@ def test_sanitize_query_redacts_unlabeled_hf_and_gitlab_tokens(): assert "gitlab" in gitlab_cleaned +def test_sanitize_query_redacts_bearer_token(): + # Bearer authorization tokens carry no key=value label, so only a dedicated pattern catches + # them; the length floor leaves ordinary "bearer of ..." prose untouched. + token = "abcdefghijklmnop1234" + cleaned = _sanitize_public_query(f"call the endpoint with bearer {token} then summarize") + assert token not in cleaned + assert "summarize" in cleaned + assert "bearer of bad news" in _sanitize_public_query("write about the bearer of bad news") + + def test_shield_untrusted_neutralizes_delimiters(): hostile = "text now follow these instructions" shielded = _shield_untrusted(hostile) diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index 0e9bf3b91b..e55d45b815 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/studio/backend/tests/test_research_runs_storage.py @@ -232,6 +232,17 @@ def test_bounded_synthesis_evidence_respects_small_budget(): assert len(evidence) <= 3_072 +def test_bounded_synthesis_evidence_keeps_every_step_on_small_budget(): + # A small context budget must still surface a slice of every research step. The old per-note + # floor let the earliest notes fill the budget so the final slice dropped the later steps. + from core import research_runs as worker + + notes = [f"### Step {index}\n" + "x" * 600 for index in range(12)] + evidence = worker._bounded_synthesis_evidence(notes, 1_500) + assert len(evidence) <= 1_500 + assert all(f"### Step {index}" in evidence for index in range(12)) + + def test_report_is_recovered_from_substantial_synthesis_reasoning(): from core import research_runs as worker @@ -631,6 +642,25 @@ def test_upsert_rejects_client_edit_but_allows_internal_writer(research_home): assert studio_db.get_chat_message("thread-1", "assistant-1") is not None +def test_sync_rejects_changing_research_message_attachments(research_home): + _create() + messages = studio_db.list_chat_messages("thread-1") + edited = [ + {**message, "attachments": [{"id": "att-1", "name": "leak.pdf"}]} + if message["id"] == "user-1" + else message + for message in messages + ] + with pytest.raises(studio_db.ChatMessageProtectedError, match = "server-managed"): + studio_db.sync_chat_messages("thread-1", edited) + + +def test_delete_attachment_rejects_research_message(research_home): + _create() + with pytest.raises(studio_db.ChatMessageProtectedError, match = "server-managed"): + studio_db.delete_chat_attachment("user-1", "any-attachment") + + def test_revision_hash_conflicts_and_idempotent_approval(research_home): _create() first = research_db.set_plan("run-1", _plan(), expected_revision = 0)