From 5c129f0380a9716e17e75eb129c7c8aad063ce7d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 11:47:15 +0000 Subject: [PATCH] Studio: harden Deep Research citations, query privacy, and message protection Address review findings in the Deep Research backend: - Escape an unbalanced ")" in citation destinations so a source URL cannot close the markdown link early and inject a second link, keeping balanced parentheses literal. - Match raw-URL citations on whole tokens so a URL sharing another URL's prefix is no longer partially rewritten. - Redact non-global IPv6 addresses in public search queries, matching the existing IPv4 handling. - Detect credential key names after normalizing case and separators so nested openaiApiKey, accessToken, and clientSecret values cannot be persisted. - Reject client edits to server-managed research prompts and reports at the storage layer; only the internal writers pass allow_research_update. - Scope research searches to the first allowed domains instead of dropping site scoping for large allow lists. - Persist the same fetch evidence bound used during live synthesis so a resumed run is not shortened. - Scope run completion so it only replaces this run's message parts. Add regression tests for the above. --- .../core/inference/web_access_policy.py | 6 +- studio/backend/core/research_runs.py | 56 +++++++++++++-- studio/backend/routes/chat_history.py | 2 +- studio/backend/routes/research_runs.py | 19 ++++- studio/backend/storage/studio_db.py | 70 +++++++++++++++---- .../tests/test_research_runs_hardening.py | 44 +++++++++++- .../tests/test_research_runs_storage.py | 37 +++++++++- 7 files changed, 207 insertions(+), 27 deletions(-) diff --git a/studio/backend/core/inference/web_access_policy.py b/studio/backend/core/inference/web_access_policy.py index 4d044a07f3..21134f05eb 100644 --- a/studio/backend/core/inference/web_access_policy.py +++ b/studio/backend/core/inference/web_access_policy.py @@ -135,7 +135,9 @@ def website_policy_prompt(policy: dict[str, Any] | None) -> str: def scope_search_query(query: str, policy: dict[str, Any] | None) -> str: allowed = normalize_website_policy(policy)["allowedDomains"] - if not allowed or len(allowed) > 8: + if not allowed: return query - site_filter = " OR ".join(f"site:{domain}" for domain in allowed) + # Cap the site: filter (search engines limit OR operators) instead of dropping scoping + # entirely for large allow lists, which returned unrelated results that all got filtered out. + site_filter = " OR ".join(f"site:{domain}" for domain in allowed[:8]) return f"{query} ({site_filter})" diff --git a/studio/backend/core/research_runs.py b/studio/backend/core/research_runs.py index 08e3385d44..918db50f04 100644 --- a/studio/backend/core/research_runs.py +++ b/studio/backend/core/research_runs.py @@ -65,6 +65,10 @@ _QUERY_PHONE = re.compile( r"(? str: return match.group(0) +def _redact_nonpublic_ipv6(match: "re.Match[str]") -> str: + # Strip brackets and any zone id before validating; redact non-global addresses. + candidate = match.group(0).strip("[]").split("%", 1)[0] + try: + return " " if not ipaddress.ip_address(candidate).is_global else match.group(0) + except ValueError: + return match.group(0) + + +def _escape_link_destination(url: str) -> str: + # Escape an unbalanced ")" so a source URL cannot close the citation and inject a link. + out: list[str] = [] + depth = 0 + for char in url: + if char == "\\": + out.append("\\\\") + elif char == "(": + depth += 1 + out.append(char) + elif char == ")" and depth == 0: + out.append("\\)") + else: + if char == ")": + depth -= 1 + out.append(char) + return "".join(out) + + def _shield_untrusted(text: str) -> str: """Escape prompt-delimiter tags embedded in untrusted evidence so gathered web or document content cannot close a wrapper block and inject model instructions.""" @@ -209,6 +241,7 @@ def _sanitize_public_query(query: str) -> str: query = _QUERY_PHONE.sub(" ", query) query = _QUERY_LABELED_PRIVATE_ID.sub(" ", query) query = _QUERY_IPV4.sub(_redact_nonpublic_ip, query) + query = _QUERY_IPV6.sub(_redact_nonpublic_ipv6, query) query = _QUERY_PAYMENT_CARD.sub( lambda match: " " if _luhn_valid(match.group(0)) else match.group(0), query, @@ -469,7 +502,7 @@ def _validate_report_sources(report: str, sources: list[dict]) -> str: return None title = str(source.get("title") or url).replace("[", "").replace("]", "").strip() token = f"\x00research-citation-{len(placeholders)}\x00" - placeholders[token] = f"[{title or url}]({url})" + placeholders[token] = f"[{title or url}]({_escape_link_destination(url)})" return token def replace_markdown_links(text: str) -> str: @@ -545,12 +578,18 @@ def _validate_report_sources(report: str, sources: list[dict]) -> str: def replace_autolink(match: re.Match) -> str: return citation(match.group(1)) or match.group(1) + def replace_raw_url(match: re.Match) -> str: + # Cite whole source URLs; drop other raw URLs. Whole-match avoids prefix collisions. + raw = match.group(0) + core = raw.rstrip(".,;:!?") + if core in source_by_url: + return (citation(core) or core) + raw[len(core):] + return "" + validated = replace_markdown_links(report) validated = _AUTOLINK.sub(replace_autolink, validated) validated = _NUMBERED_CITATION.sub(replace_number, validated) - for url in sorted(source_urls, key = len, reverse = True): - validated = validated.replace(url, citation(url) or url) - validated = _RAW_URL.sub("", validated) + validated = _RAW_URL.sub(replace_raw_url, validated) for token, link in placeholders.items(): validated = validated.replace(token, link) return validated.strip() @@ -606,7 +645,9 @@ def _update_assistant( retained = [ part for part in content - if not isinstance(part, dict) or part.get("type") not in replaced_types + if not isinstance(part, dict) + or part.get("type") not in replaced_types + or part.get("researchRunId") not in (None, run["id"]) ] if reasoning: retained.append({"type": "reasoning", "text": reasoning, "researchRunId": run["id"]}) @@ -642,7 +683,8 @@ def _update_assistant( "attachments": existing.get("attachments"), "metadata": metadata, "createdAt": existing.get("createdAt") or db.now_ms(), - } + }, + allow_research_update = True, ) @@ -1545,7 +1587,7 @@ class ResearchSupervisor: "sourceCount": len(step_sources) + len(rag_sources), "sourceUrls": [source["url"] for source in step_sources], "evidenceSources": rag_sources, - **({"excerpt": clean_result[:2000]} if action["action"] == "fetch" else {}), + **({"excerpt": clean_result[:12000]} if action["action"] == "fetch" else {}), **({"error": clean_result[:500]} if tool_failed else {}), } await self._check_active(run["id"]) diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 6113f48c3b..b95ea9f4ce 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -422,7 +422,7 @@ async def save_thread_message( raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found") try: return ChatMessage(**upsert_chat_message(payload.model_dump())) - except ChatMessageConflictError as exc: + except (ChatMessageConflictError, ChatMessageProtectedError) as exc: raise log_and_http_error( exc, 409, diff --git a/studio/backend/routes/research_runs.py b/studio/backend/routes/research_runs.py index b0e0f4d6f5..875d064980 100644 --- a/studio/backend/routes/research_runs.py +++ b/studio/backend/routes/research_runs.py @@ -22,7 +22,13 @@ from storage import research_runs_db as db from storage.studio_db import get_chat_message, get_chat_thread, upsert_chat_message router = APIRouter() -_SENSITIVE_KEY = re.compile(r"^(?:api.?key|secret|token|authorization|password)$", re.IGNORECASE) +_SENSITIVE_KEY_EXACT = { + "authorization", "password", "secret", "token", "apikey", "credential", "credentials", +} +_SENSITIVE_KEY_SUFFIXES = ( + "apikey", "accesskey", "accesstoken", "authtoken", "bearertoken", + "clientsecret", "privatekey", "refreshtoken", "sessiontoken", +) _MAX_PLAN_STEPS = 30 _DELTA_ONLY_EVENTS = {"reasoning.updated", "report.updated"} @@ -118,16 +124,23 @@ def _sync_assistant(run: dict, text: str | None = None) -> None: **message, "content": content, "metadata": metadata, - } + }, + allow_research_update = True, ) +def _is_sensitive_key(key: object) -> bool: + # Match after stripping separators/case so openaiApiKey, access_token, clientSecret all hit. + normalized = re.sub(r"[^a-z0-9]", "", str(key).casefold()) + return normalized in _SENSITIVE_KEY_EXACT or normalized.endswith(_SENSITIVE_KEY_SUFFIXES) + + def _contains_sensitive_key(value: object) -> bool: """Recursively test whether any (possibly nested) mapping key looks sensitive, so credentials cannot be smuggled into a durable run via a nested dict.""" if isinstance(value, dict): return any( - bool(_SENSITIVE_KEY.search(str(key))) or _contains_sensitive_key(item) + _is_sensitive_key(key) or _contains_sensitive_key(item) for key, item in value.items() ) if isinstance(value, (list, tuple)): diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index e758d8fca2..35e6cf6f82 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -1664,10 +1664,62 @@ def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str) ) -def upsert_chat_message(message: dict) -> dict: +def _research_message_ids(conn: sqlite3.Connection, thread_id: str) -> set[str]: + return { + 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 + } + + +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 " + "FROM chat_messages WHERE thread_id = ? AND id = ?", + (thread_id, str(message["id"])), + ).fetchone() + if row is None: + return False + + def canon(value: object) -> str | None: + return json.dumps(value, sort_keys=True) if value is not None else None + + return ( + 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 (message.get("parentId") or None) != (row["parent_id"] or None) + or str(message.get("role")) != str(row["role"]) + ) + + +def _guard_research_messages( + conn: sqlite3.Connection, thread_id: str, messages: list[dict] +) -> None: + protected = _research_message_ids(conn, thread_id) + if not protected: + return + for message in messages: + if str(message["id"]) in protected and _research_message_would_change( + conn, thread_id, message + ): + raise ChatMessageProtectedError( + "Research prompts and responses are server-managed and cannot be edited" + ) + + +def upsert_chat_message(message: dict, *, allow_research_update: bool = False) -> dict: conn = get_connection() try: conn.execute("BEGIN IMMEDIATE") + if not allow_research_update: + _guard_research_messages(conn, message["threadId"], [message]) _raise_if_chat_message_thread_conflicts( conn, message["threadId"], @@ -1716,10 +1768,14 @@ def sync_chat_messages( thread_id: str, messages: list[dict], prune_missing: bool = False, + *, + allow_research_update: bool = False, ) -> list[dict]: conn = get_connection() try: conn.execute("BEGIN IMMEDIATE") + if not allow_research_update: + _guard_research_messages(conn, thread_id, messages) _raise_if_chat_message_thread_conflicts( conn, thread_id, @@ -1762,17 +1818,7 @@ def sync_chat_messages( ).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: + if removed_ids & _research_message_ids(conn, thread_id): raise ChatMessageProtectedError( "Research prompts and responses cannot be deleted from their original thread" ) diff --git a/studio/backend/tests/test_research_runs_hardening.py b/studio/backend/tests/test_research_runs_hardening.py index 34f4938db1..ed7fdfe46c 100644 --- a/studio/backend/tests/test_research_runs_hardening.py +++ b/studio/backend/tests/test_research_runs_hardening.py @@ -6,11 +6,13 @@ import pytest from core.research_runs import ( + _escape_link_destination, _sanitize_public_query, _shield_untrusted, _validate_report_document_sources, + _validate_report_sources, ) -from routes.research_runs import CreateResearchRun, _sanitize_config +from routes.research_runs import CreateResearchRun, _is_sensitive_key, _sanitize_config def test_sanitize_query_redacts_payment_card(): @@ -84,3 +86,43 @@ def test_sanitize_config_rejects_nested_rag_scope_secret(): payload = _make_payload(ragScope = {"kb_id": {"token": "rag-secret"}}) with pytest.raises(Exception): _sanitize_config(payload, {"modelId": "m"}) + + +def test_sensitive_key_matches_prefixed_and_camelcase_variants(): + for key in ( + "apiKey", "openaiApiKey", "accessToken", "access_token", + "clientSecret", "refreshToken", "authorization", + ): + assert _is_sensitive_key(key), key + # Ordinary request fields must not be flagged, so normal runs still validate. + for key in ("model", "temperature", "maxTokens", "project_id", "top_k"): + assert not _is_sensitive_key(key), key + + +def test_sanitize_query_redacts_nonpublic_ipv6_but_keeps_public(): + assert "fd00" not in _sanitize_public_query("inspect fd00::dead:beef service health") + assert "fe80" not in _sanitize_public_query("connect to fe80::1%eth0 gateway now") + assert "2606:4700:4700::1111" in _sanitize_public_query( + "what runs on 2606:4700:4700::1111 dns" + ) + + +def test_escape_link_destination_escapes_only_unbalanced_paren(): + assert _escape_link_destination("https://x.co/a)evil") == "https://x.co/a\\)evil" + # Balanced parentheses (e.g. Wikipedia-style URLs) stay literal. + assert _escape_link_destination("https://x.co/Foo_(bar)") == "https://x.co/Foo_(bar)" + + +def test_citation_injection_cannot_open_second_link(): + url = "https://allowed.example/a)evil" + out = _validate_report_sources(f"See {url} now.", [{"url": url, "title": "Allowed"}]) + assert "a\\)evil" in out + + +def test_raw_url_citation_does_not_collide_on_prefix(): + sources = [{"url": "https://ex.com/report", "title": "Report"}] + out = _validate_report_sources( + "See https://ex.com/report and https://ex.com/report-attack now.", sources + ) + assert "[Report](https://ex.com/report)" in out + assert "/report)-attack" not in out diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index 5cebcc39b0..b7794f5e4e 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/studio/backend/tests/test_research_runs_storage.py @@ -540,6 +540,40 @@ def test_pruning_rejects_deleting_research_turn_messages(research_home, removed_ 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 + + +def test_sync_rejects_editing_research_message_but_allows_noop(research_home): + _create() + unchanged = studio_db.list_chat_messages("thread-1") + # Re-syncing identical content is a no-op and must still be allowed. + studio_db.sync_chat_messages("thread-1", unchanged) + edited = [ + {**message, "content": [{"type": "text", "text": "HIJACKED"}]} + if message["id"] == "user-1" + else message + for message in unchanged + ] + with pytest.raises(studio_db.ChatMessageProtectedError, match = "server-managed"): + studio_db.sync_chat_messages("thread-1", edited) + assert studio_db.get_chat_message("thread-1", "user-1")["content"] == [ + {"type": "text", "text": "What changed?"} + ] + + +def test_upsert_rejects_client_edit_but_allows_internal_writer(research_home): + _create() + original = studio_db.get_chat_message("thread-1", "user-1") + with pytest.raises(studio_db.ChatMessageProtectedError, match = "server-managed"): + studio_db.upsert_chat_message( + {**original, "content": [{"type": "text", "text": "client edit"}]} + ) + studio_db.upsert_chat_message( + {**original, "content": [{"type": "text", "text": "server update"}]}, + allow_research_update = True, + ) + assert studio_db.get_chat_message("thread-1", "user-1")["content"] == [ + {"type": "text", "text": "server update"} + ] assert studio_db.get_chat_message("thread-1", "assistant-1") is not None @@ -1746,7 +1780,8 @@ def test_update_assistant_replaces_report_parts_without_duplication(research_hom ], "metadata": {"researchRunId": "run-1"}, "createdAt": 3, - } + }, + allow_research_update = True, ) run = research_db.get_run("run-1") source = {"url": "https://new.example", "title": "New", "snippet": "Evidence"}