diff --git a/studio/backend/core/inference/web_access_policy.py b/studio/backend/core/inference/web_access_policy.py index ebdb96c050..4d044a07f3 100644 --- a/studio/backend/core/inference/web_access_policy.py +++ b/studio/backend/core/inference/web_access_policy.py @@ -108,7 +108,7 @@ def check_url_access(url: str, policy: dict[str, Any] | None) -> tuple[bool, str except (TypeError, ValueError): return False, "Blocked: URL has an invalid hostname or port.", "" if not hostname_allowed(hostname, policy): - return False, f"Blocked by website access policy: {hostname}.", hostname + return False, f"Blocked: website access policy disallows {hostname}.", hostname return True, "", hostname diff --git a/studio/backend/core/research_runs.py b/studio/backend/core/research_runs.py index 02caf74f9d..5eb65c9133 100644 --- a/studio/backend/core/research_runs.py +++ b/studio/backend/core/research_runs.py @@ -30,7 +30,7 @@ _URL_BLOCK = re.compile( r"Title:\s*(?P[^\n]*)\nURL:\s*(?P<url>https?://[^\s]+)\nSnippet:\s*(?P<snippet>.*?)(?=\n\n---|\Z)", re.DOTALL, ) -_MARKDOWN_LINK = re.compile(r"\[([^\]]+)\]\((https?://[^)\s]+)\)") +_MARKDOWN_LINK_START = re.compile(r"\[([^\]\n]+)\]\((https?://)") _SOURCES_HEADING = re.compile( r"^(?:#{1,6}\s+|\*\*)?" r"(?:Sources?|References?|Bibliography|Works\s+Cited|Source\s+List)" @@ -387,6 +387,10 @@ def _split_rag_result(result: str) -> tuple[str, list[dict[str, Any]]]: return text.rstrip(), sources +def _research_step_failed(web_result: str, rag_sources: list[dict]) -> bool: + return is_tool_error(web_result) and not rag_sources + + def _validate_report_sources(report: str, sources: list[dict]) -> str: """Canonicalize citations and remove model-authored source lists.""" source_by_url = { @@ -408,9 +412,69 @@ def _validate_report_sources(report: str, sources: list[dict]) -> str: placeholders[token] = f"[{title or url}]({url})" return token - def replace_link(match: re.Match) -> str: - label, url = match.group(1).strip(), match.group(2) - return citation(url) or label + def replace_markdown_links(text: str) -> str: + pieces = [] + cursor = 0 + while match := _MARKDOWN_LINK_START.search(text, cursor): + destination_start = match.start(2) + index = match.end(2) + depth = 0 + escaped = False + close = None + destination_end = None + while index < len(text): + character = text[index] + if escaped: + escaped = False + elif character == "\\": + escaped = True + elif character.isspace(): + if depth != 0: + break + destination_end = index + title_start = index + while title_start < len(text) and text[title_start].isspace(): + title_start += 1 + if title_start < len(text) and text[title_start] in {'"', "'"}: + quote = text[title_start] + title_end = title_start + 1 + title_escaped = False + while title_end < len(text): + if title_escaped: + title_escaped = False + elif text[title_end] == "\\": + title_escaped = True + elif text[title_end] == quote: + break + title_end += 1 + if title_end >= len(text): + break + title_start = title_end + 1 + while title_start < len(text) and text[title_start].isspace(): + title_start += 1 + if title_start < len(text) and text[title_start] == ")": + close = title_start + break + elif character == "(": + depth += 1 + elif character == ")": + if depth == 0: + close = index + destination_end = index + break + depth -= 1 + index += 1 + if close is None: + pieces.append(text[cursor : match.start()]) + pieces.append(match.group(1).strip()) + cursor = index + continue + url = text[destination_start:destination_end].replace(r"\(", "(").replace(r"\)", ")") + pieces.append(text[cursor : match.start()]) + pieces.append(citation(url) or match.group(1).strip()) + cursor = close + 1 + pieces.append(text[cursor:]) + return "".join(pieces) def replace_number(match: re.Match) -> str: index = int(match.group(1)) - 1 @@ -421,7 +485,7 @@ 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) - validated = _MARKDOWN_LINK.sub(replace_link, report) + 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): @@ -1400,6 +1464,7 @@ class ResearchSupervisor: f"Input: {argument}\nResult:\n{result[:12000]}" ) tool_failed = is_tool_error(result) + step_failed = _research_step_failed(result, rag_sources) clean_result = strip_result_for_model(result) step_result = { "action": action["action"], @@ -1417,7 +1482,7 @@ class ResearchSupervisor: position, action["title"], argument, - "failed" if tool_failed else "completed", + "failed" if step_failed else "completed", step_result, self.worker_id, ) @@ -1426,7 +1491,7 @@ class ResearchSupervisor: db.append_worker_event, run["id"], self.worker_id, - "step.failed" if tool_failed else "step.completed", + "step.failed" if step_failed else "step.completed", { "position": position, "stepPosition": position, @@ -1434,7 +1499,7 @@ class ResearchSupervisor: "action": action["action"], "input": argument, "sourceCount": len(step_sources) + len(rag_sources), - **({"error": clean_result[:500]} if tool_failed else {}), + **({"error": clean_result[:500]} if step_failed else {}), }, ) await self._check_worker_write(run["id"], seq is not None) diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index accb29a297..3baf17e550 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/studio/backend/tests/test_research_runs_storage.py @@ -749,6 +749,32 @@ def test_report_citations_are_limited_to_gathered_sources(): assert "https://invalid.example/guess" not in validated +def test_report_citations_preserve_balanced_parentheses_in_urls(): + from core.research_runs import _validate_report_sources + + url = "https://en.wikipedia.org/wiki/Function_(mathematics)" + validated = _validate_report_sources( + f"Supported [generic label]({url}).", + [{"url": url, "title": "Function (mathematics)"}], + ) + + assert f"[Function (mathematics)]({url})" in validated + assert ( + _validate_report_sources( + f'With title [generic label]({url} "reference page").', + [{"url": url, "title": "Function (mathematics)"}], + ) + == f"With title [Function (mathematics)]({url})." + ) + assert ( + _validate_report_sources( + f"Malformed [generic label]({url}", + [{"url": url, "title": "Function (mathematics)"}], + ) + == "Malformed generic label" + ) + + def test_report_citations_use_canonical_titles_without_model_sources_section(): from core.research_runs import _validate_report_sources @@ -873,6 +899,14 @@ def test_research_agent_actions_are_model_directed_and_url_bounded(): ) +def test_rag_evidence_makes_failed_web_search_recoverable(): + from core.research_runs import _research_step_failed + + blocked = "Blocked: website access policy disallows example.com." + assert _research_step_failed(blocked, []) is True + assert _research_step_failed(blocked, [{"chunkId": "doc-1:0"}]) is False + + def test_research_budget_defaults_support_long_runs(): from routes.research_runs import CreateResearchRun, ResearchPlan, _sanitize_config diff --git a/studio/backend/tests/test_web_access_policy.py b/studio/backend/tests/test_web_access_policy.py index 38a457c847..6f05c3700f 100644 --- a/studio/backend/tests/test_web_access_policy.py +++ b/studio/backend/tests/test_web_access_policy.py @@ -165,7 +165,7 @@ def test_direct_fetch_rejects_blocked_host_before_dns(monkeypatch): "https://example.com/article", website_policy = ARXIV_ONLY, ) - assert "Blocked by website access policy" in result + assert "Blocked: website access policy" in result assert resolved == [] @@ -188,5 +188,5 @@ def test_direct_fetch_rechecks_every_redirect_before_dns(monkeypatch): "https://arxiv.org/abs/1", website_policy = ARXIV_ONLY, ) - assert "Blocked by website access policy: example.com" in result + assert "Blocked: website access policy disallows example.com" in result assert resolved == [("arxiv.org", 443)] diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index f435c63ffb..6d4a7e2a31 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -1454,6 +1454,16 @@ const Composer: FC<{ const researchThreadClaimed = useResearchRunStore((state) => researchThreadId ? Boolean(state.claimedThreadIds[researchThreadId]) : false, ); + const activeResearchRun = useResearchRunStore((state) => { + const runId = researchThreadId + ? state.latestRunByThreadId[researchThreadId] + : undefined; + return runId ? state.sessions[runId]?.run : undefined; + }); + const isResearchActive = Boolean( + activeResearchRun && + !["completed", "failed", "cancelled"].includes(activeResearchRun.status), + ); const hasResearchMessage = useAuiState(({ thread }) => thread.messages.some((message) => { const custom = ( @@ -1778,6 +1788,10 @@ const Composer: FC<{ const handleSubmit = useCallback( (event: Parameters<NonNullable<ComponentProps<"form">["onSubmit"]>>[0]) => { + if (isResearchActive) { + event.preventDefault(); + return; + } if (disabled || shouldBlockSend()) { event.preventDefault(); return; @@ -1865,6 +1879,7 @@ const Composer: FC<{ hasAttachments, hasPendingAudio, interceptSend, + isResearchActive, overlay, promptQueueActive, referenceThreadId, diff --git a/tests/studio/test_deep_research_frontend_contract.py b/tests/studio/test_deep_research_frontend_contract.py index 767f3589a2..735247fcc6 100644 --- a/tests/studio/test_deep_research_frontend_contract.py +++ b/tests/studio/test_deep_research_frontend_contract.py @@ -37,6 +37,7 @@ def test_research_api_is_isolated_and_cursor_based() -> None: def test_research_mode_is_single_chat_and_detaches_without_cancel() -> None: adapter = source("features/chat/api/chat-adapter.ts") + thread = source("components/assistant-ui/thread.tsx") assert "runtime.deepResearchEnabled" in adapter assert "!options.pairId" in adapter assert 'options.modelType === "base"' in adapter @@ -61,6 +62,9 @@ def test_research_mode_is_single_chat_and_detaches_without_cancel() -> None: assert "signal: researchFollowController.signal" in adapter assert "beginExternalResearchFollow(" in adapter assert "ragScope" in adapter + submit = thread.split("const handleSubmit = useCallback", 1)[1].split("const stopQueue", 1)[0] + assert "if (isResearchActive)" in submit + assert "event.preventDefault()" in submit assert "runtime.ragEnabled\n ? { thread_id: resolvedThreadId }" in adapter create_block = adapter.split("createdRun = await createResearchRun({", 1)[1].split("});", 1)[0] assert "modelId:" not in create_block