ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix report-stalling regex, uncataloged KB evidence and bracketed titles for PR #7219

Browse source 
Catastrophic backtracking in _DOCUMENT_CITATION. The alternation
(?:[^\[\]]+|\[[^\[\]]*\])* backtracks exponentially on an unterminated
"[Document:" with no later bare "]", which is ordinary malformed model output
and exactly what this sanitizer exists to handle. Runtime quadrupled every two
characters; one realistic 76-char line did not finish in 90s. It runs
synchronously inside async _research (the line below it uses asyncio.to_thread),
so a single bad report pins the event loop and stalls all of Studio, not just
the run. Replaced with the language-equivalent unrolled form, verified identical
on well-formed inputs including bracketed filenames, and linear: a 20,000-char
tail now takes 0.4ms. Not using possessive quantifiers or atomic groups, which
need Python 3.11 while this package declares >=3.9.

Uncataloged knowledge base evidence reached synthesis. When maxSources is
already full, every returned chunk hits the continue, so accepted_rag_sources
stays empty, the "if accepted_rag_sources" rebuild no-ops and rag_result keeps
the raw KB text. That text has no document_source_catalog entry, so the
validator strips any citation to it and synthesis is left building claims on
private KB chunks it cannot attribute. Cleared, gated on rag_sources so a
text-only KB reply is still passed through. The resume branch built rag_evidence
from all restored sources with the same hole, so it now mirrors the live loop.

Bracketed source titles destroyed their own citation. The catalog gave the model
the raw title while the citation writer stripped brackets. Search titles
routinely carry one ("[PDF] Annual Report"), and the prompt tells the model to
copy the title verbatim, producing a label the validator cannot match. Both
sides now share _citation_title.

Verified: 756 passed across the research/web/sandbox/chat-history/rag backend
suites. Each fix has a regression test that fails without it.
This commit is contained in:
danielhanchen 2026-07-26 12:41:05 +00:00
commit 3a816907e6
3 changed files with 151 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

66
studio/backend/core/research_runs.py
View file

@ -42,7 +42,10 @@ _SOURCES_HEADING = re.compile(
_NUMBERED_CITATION = re.compile(r"(?<!\^)\[(\d+)]")
_AUTOLINK = re.compile(r"<(https?://[^>\s]+)>")
_RAW_URL = re.compile(r"https?://[^\s<>]+")
_DOCUMENT_CITATION = re.compile(r"\[Document:(?:[^\[\]]+|\[[^\[\]]*\])*\]")
# Unrolled rather than the equivalent (?:[^\[\]]+|\[[^\[\]]*\])* : that alternation backtracks
# catastrophically on an unterminated "[Document:" (ordinary malformed model output), and this
# runs on the event loop, so one bad report would stall all of Studio.
_DOCUMENT_CITATION = re.compile(r"\[Document:[^\[\]]*(?:\[[^\[\]]*\][^\[\]]*)*\]")
# Wrapper delimiters used in the decision/synthesis prompts. Any occurrence inside
# untrusted evidence is escaped so gathered content cannot close a block early.
_PROMPT_DELIMITER_TAGS = re.compile(
@ -645,6 +648,18 @@ def _split_rag_result(result: str) -> tuple[str, list[dict[str, Any]]]:
return text.rstrip(), sources
def _citation_title(source: dict, fallback: str) -> str:
"""Title as it may appear in a markdown link label.
Brackets are stripped because the report prompt tells the model to copy titles verbatim
from the source catalog, and search titles routinely carry one ("[PDF] Annual Report").
A bracket inside the label makes the citation unmatchable, so the catalog and the
citation writer must agree on the same stripped form.
"""
title = str(source.get("title") or fallback).replace("[", "").replace("]", "").strip()
return title or fallback
def _trim_url_tail(raw: str) -> str:
"""Strip trailing prose punctuation that ``_RAW_URL`` swallowed.
@ -688,9 +703,9 @@ def _validate_report_sources(report: str, sources: list[dict]) -> str:
source = source_by_url.get(url)
if source is None:
return None
title = str(source.get("title") or url).replace("[", "").replace("]", "").strip()
title = _citation_title(source, url)
token = f"\x00research-citation-{len(placeholders)}\x00"
placeholders[token] = f"[{title or url}]({_escape_link_destination(url)})"
placeholders[token] = f"[{title}]({_escape_link_destination(url)})"
return token
def replace_markdown_links(text: str) -> str:
@ -1627,30 +1642,33 @@ class ResearchSupervisor:
)
for source in document_sources
}
# Mirrors the live loop: evidence must hold only chunks that made it into the
# catalog, else the validator strips citations to the rest and synthesis is left
# building claims on uncataloged document text.
accepted_rag_sources = []
for source in restored_rag_sources:
source_key = str(
source.get("chunkId")
or f"{source.get('documentId') or source.get('filename')}:{source.get('page') or ''}"
)
if (
source_key in document_source_keys
or len(sources) + len(document_sources) >= max_sources
):
continue
written = await asyncio.to_thread(
db.upsert_document_source,
run["id"],
int(step["position"]),
source,
self.worker_id,
)
await self._check_worker_write(run["id"], written)
document_source_keys.add(source_key)
document_sources.append({**source, "stepPosition": step["position"]})
if source_key not in document_source_keys:
if len(sources) + len(document_sources) >= max_sources:
continue
written = await asyncio.to_thread(
db.upsert_document_source,
run["id"],
int(step["position"]),
source,
self.worker_id,
)
await self._check_worker_write(run["id"], written)
document_source_keys.add(source_key)
document_sources.append({**source, "stepPosition": step["position"]})
accepted_rag_sources.append(source)
rag_evidence = "\n".join(
f"{item.get('filename') or 'Document'}: "
f"{item.get('text') or item.get('snippet') or ''}"
for item in restored_rag_sources
for item in accepted_rag_sources
)
title = str(step.get("title") or "Recovered research step")
notes.append(
@ -1671,7 +1689,7 @@ class ResearchSupervisor:
for position in range(start_position, max_steps):
await self._check_active(run["id"])
source_catalog = "\n".join(
f"- {source.get('title') or source['url']} | {source['url']} | "
f"- {_citation_title(source, source['url'])} | {source['url']} | "
f"{source.get('snippet') or ''}"
for source in sources
)
@ -1837,6 +1855,12 @@ class ResearchSupervisor:
f"{source.get('text') or source.get('snippet') or ''}"
for source in accepted_rag_sources
)
elif rag_sources:
# Every chunk was refused by the source cap, so none has a catalog entry and the
# validator would strip any citation to it. Drop the evidence rather than let
# synthesis build claims on it. Gated on rag_sources so a text-only KB reply
# ("No documents are attached to this chat.") is still passed through.
rag_result = ""
rag_sources = accepted_rag_sources
step_sources = []
for match in _URL_BLOCK.finditer(result if action["action"] == "search" else ""):
@ -1943,7 +1967,7 @@ class ResearchSupervisor:
await self._check_worker_write(run["id"], seq is not None)
await self._check_active(run["id"])
source_catalog = "\n".join(
f"{index}. Title: {source.get('title') or source['url']}\n URL: {source['url']}"
f"{index}. Title: {_citation_title(source, source['url'])}\n URL: {source['url']}"
for index, source in enumerate(sources, 1)
)
document_source_catalog = "\n".join(

27
studio/backend/tests/test_research_runs_hardening.py
View file

@ -6,6 +6,7 @@
import pytest
from core.research_runs import (
_citation_title,
_escape_link_destination,
_sanitize_public_query,
_shield_untrusted,
@ -121,6 +122,32 @@ def test_document_citation_strips_unknown_source_with_brackets():
assert out == "Ghost cite end."
def test_document_citation_regex_does_not_backtrack_catastrophically():
# An unterminated "[Document:" with no later bare "]" is ordinary malformed model output,
# which is exactly what this sanitizer exists to handle. The old alternation took longer
# than the age of the universe on one line, and it runs on the event loop.
import time
report = "Revenue rose 12 percent [Document: q3_report.pdf, p. 12 and margins improved."
start = time.perf_counter()
_validate_report_document_sources(report, [{"filename": "q3_report.pdf", "page": 12}])
assert time.perf_counter() - start < 1.0
# And a long tail stays linear rather than exponential.
start = time.perf_counter()
_validate_report_document_sources("[Document: " + "a" * 20_000, [])
assert time.perf_counter() - start < 1.0
def test_citation_title_strips_brackets_for_catalog_and_citation():
# Search titles routinely carry a bracketed prefix ("[PDF] ..."), and the report prompt tells
# the model to copy the catalog title verbatim into the link label, where a bracket makes the
# citation unmatchable. The catalog and the citation writer share this helper so they cannot
# offer a label the validator then fails to match.
assert _citation_title({"title": "[PDF] Annual Report 2024"}, "https://x/a") == "PDF Annual Report 2024"
assert _citation_title({"title": "[]"}, "https://x/a") == "https://x/a"
assert _citation_title({}, "https://x/a") == "https://x/a"
def _make_payload(**overrides) -> CreateResearchRun:
payload = {"threadId": "t1", "userMessageId": "u1", "inferenceRequest": {"model": "m"}}
payload.update(overrides)

79
studio/backend/tests/test_research_runs_storage.py
View file

@ -1923,6 +1923,85 @@ def test_recovered_running_research_resumes_durable_progress(research_home, monk
assert completed["report"].startswith("# Resumed report")
def test_knowledge_base_evidence_beyond_the_source_cap_is_not_synthesized(
research_home, monkeypatch
):
"""A knowledge-base hit that the source cap refuses to persist must not reach synthesis:
it has no document_source_catalog entry, so any citation of it is stripped from the
finished report and the claim it supports would be left unattributed."""
from core import research_runs as worker
_create(
rag_scope = {"kb_id": "kb-1", "default_top_k": 4},
budgets = {
"maxSteps": 3,
"maxSources": 1,
"modelTimeoutSeconds": 30,
"toolTimeoutSeconds": 10,
},
)
supervisor = worker.ResearchSupervisor(SimpleNamespace(state = SimpleNamespace(server_port = 1)))
decisions = iter(
(
json.dumps({"action": "search", "title": "First", "query": "first query"}),
json.dumps({"action": "search", "title": "Second", "query": "second query"}),
json.dumps({"action": "finish", "title": "Enough evidence"}),
)
)
synthesis_prompts = []
report = "# Report\n\nA finding [Document: kept.pdf, p. 1]."
async def fake_stream_completion(run, messages, **kwargs):
system = messages[0]["content"]
if "rigorous web research plan" in system:
return json.dumps(_plan()), "planned", "stop"
if "iterative research process" in system:
return next(decisions), "decided", "stop"
synthesis_prompts.append(messages[1]["content"])
research_db.set_report_progress(run["id"], report)
return report, "synthesized", "stop"
labels = iter(("kept", "capped"))
def fake_tool(name, arguments, *args, **kwargs):
if name == "search_knowledge_base":
label = next(labels)
return (
f"UNCATALOGED_{label.upper()}_KB_TEXT"
+ worker.RAG_SOURCES_SENTINEL
+ json.dumps(
[
{
"chunkId": f"doc-{label}:0",
"documentId": f"doc-{label}",
"filename": f"{label}.pdf",
"page": 1,
"text": f"{label} chunk body",
}
]
)
)
return "Title: Alpha\nURL: https://a.example.com\nSnippet: alpha snippet."
monkeypatch.setattr(supervisor, "_stream_completion", fake_stream_completion)
monkeypatch.setattr(worker, "execute_tool", fake_tool)
asyncio.run(supervisor._process(research_db.claim_next(supervisor.worker_id)))
planned = research_db.get_run("run-1")
research_db.approve("run-1", planned["planRevision"], planned["planHash"])
asyncio.run(supervisor._process(research_db.claim_next(supervisor.worker_id)))
completed = research_db.get_run("run-1")
assert completed["status"] == "completed"
# The cap admitted the first chunk only, so only it may appear in the evidence.
assert [source["filename"] for source in completed["documentSources"]] == ["kept.pdf"]
assert synthesis_prompts, "synthesis must have run"
assert "kept chunk body" in synthesis_prompts[0]
assert "UNCATALOGED_KEPT_KB_TEXT" not in synthesis_prompts[0]
assert "capped chunk body" not in synthesis_prompts[0]
assert "UNCATALOGED_CAPPED_KB_TEXT" not in synthesis_prompts[0]
def test_create_without_assistant_id_does_not_eagerly_create_message(research_home):
from routes.research_runs import CreateResearchRun, create_research_run