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.
This commit is contained in:
danielhanchen
parent
0308f63391
commit
5c129f0380
7 changed files with 207 additions and 27 deletions
|
|
@ -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})"
|
||||
|
|
|
|||
|
|
@ -65,6 +65,10 @@ _QUERY_PHONE = re.compile(
|
|||
r"(?<!\w)\+\d[\d\s().-]{7,17}\d(?!\w)|(?<!\w)\(?\d{3}\)?[\s.-]\d{3}[\s.-]\d{4}(?!\w)"
|
||||
)
|
||||
_QUERY_IPV4 = re.compile(r"(?<![\w.])(?:\d{1,3}\.){3}\d{1,3}(?![\w.])")
|
||||
_QUERY_IPV6 = re.compile(
|
||||
r"(?<![0-9A-Fa-f:])\[?(?:[0-9A-Fa-f]{0,4}:){2,}[0-9A-Fa-f.]*(?:%[A-Za-z0-9_.-]+)?\]?"
|
||||
r"(?![0-9A-Fa-f:])"
|
||||
)
|
||||
_QUERY_LABELED_PRIVATE_ID = re.compile(
|
||||
r"(?ix)\b(?:passport|driver(?:'s)?[\s_-]?licen[cs]e|national[\s_-]?id"
|
||||
r"|tax[\s_-]?id|account[\s_-]?(?:number|no))\s*[:=#-]?\s*[A-Za-z0-9][A-Za-z0-9_-]{4,24}\b"
|
||||
|
|
@ -190,6 +194,34 @@ def _redact_nonpublic_ip(match: "re.Match[str]") -> 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"])
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)):
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue