diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 0ed8920285..6995df4f41 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -80,21 +80,60 @@ _MAX_REPROMPTS = 3 # require ALL of (intent signal, length < _REPROMPT_MAX_CHARS, no # answer artifact) to fire. # -# `\r?\n` is used everywhere a newline is required so Windows-authored or -# CRLF-converted content still matches. The numbered-list indent uses -# `[ \t]*` (spaces / tabs only) rather than `\s*` so the regex stays -# linear on long whitespace runs -- greedy `\s*` + failing `\d+` caused -# O(n^2) backtracking through embedded `\r\n` characters on adversarial -# inputs. +# Notes on the patterns: +# * `\r?\n` everywhere a newline is required so Windows-authored or +# CRLF-converted content still matches. +# * Code-fence info string is `[^\r\n]{0,200}` so common languages with +# digits / symbols (python3, c++, c#, objective-c, ts-node, ...) are +# all recognised; closing fence may be indented (` ``` ` inside a +# list or blockquote). +# * HTML branches require a closing `` so plan-only mentions of +# `` or `` do not bypass the re-prompt. +# * All `[\s\S]{...}?` runs are length-bounded so the search stays +# linear on adversarial input (CRLF spam, repeated `` etc.). _HAS_ANSWER_ARTIFACT = re.compile( - r"```[a-zA-Z]*\r?\n[\s\S]+?\r?\n```" # closed code fence - r"|" # complete SVG - r"|(?:^|\r?\n)[ \t]*\d+\.[ \t]+\S.*?\r?\n[ \t]*\d+\.", # 2+ numbered list items + # Closed code fence (any markdown info string, optional indent on close). + r"```[^\r\n]{0,200}\r?\n[\s\S]{1,4000}?\r?\n[ \t]*```" + # Complete HTML page; doctype prefix is optional. + r"|(?:" + # Complete SVG document. + r"|", re.IGNORECASE, ) +# Two or more numbered list items at column 0. Indent is spaces / tabs +# only so the regex stays linear on long whitespace runs. +_NUMBERED_LIST_ARTIFACT = re.compile( + r"(?:^|\r?\n)[ \t]*\d+\.[ \t]+\S.*?\r?\n[ \t]*\d+\.", +) + +# Markers that a numbered list is a plan (still re-promptable), not a +# final answer. Explicit "Here's my plan" / "plan:" / "approach:", OR +# intent phrasing followed shortly by a tool-action verb. +_PLAN_LIST_FRAMING = re.compile( + r"\b(?:here['’]?s (?:my |the |a )?(?:plan|approach)|step \d+|" + r"i['’]?ll|i will|i am going to|let me|now i|next i)\b" + r"[\s\S]{0,80}" + r"\b(?:search|look up|call|use|fetch|browse|run|execute|" + r"check|find|open|verify|compare|summari[sz]e)\b" + r"|\b(?:plan|approach):", + re.IGNORECASE, +) + + +def _has_answer_artifact(text: str) -> bool: + """True if ``text`` looks like a completed answer artifact. + + Code fences, complete HTML, and complete SVG count directly. A + numbered list counts only when there is no plan framing, so stalls + like ``Here's my plan:\\n1. search\\n2. summarise`` still re-prompt. + """ + if _HAS_ANSWER_ARTIFACT.search(text): + return True + if _NUMBERED_LIST_ARTIFACT.search(text): + return _PLAN_LIST_FRAMING.search(text) is None + return False + # Without max_tokens, llama-server defaults to n_predict = n_ctx (up to # 262144 for Qwen3.5), producing many-minute zombie decodes when cancel # fails. t_max_predict_ms is a wall-clock backstop applied unconditionally, @@ -4842,7 +4881,7 @@ class LlamaCppBackend: and _reprompt_count < _MAX_REPROMPTS and 0 < len(_stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(_stripped) - and not _HAS_ANSWER_ARTIFACT.search(_stripped) + and not _has_answer_artifact(_stripped) ): _reprompt_count += 1 logger.info( diff --git a/studio/backend/tests/test_llama_cpp_reprompt_guard.py b/studio/backend/tests/test_llama_cpp_reprompt_guard.py index 53eb8b06bd..49954918f0 100644 --- a/studio/backend/tests/test_llama_cpp_reprompt_guard.py +++ b/studio/backend/tests/test_llama_cpp_reprompt_guard.py @@ -20,9 +20,12 @@ would still match (length < 2000, intent signal present) and the next synthetic user turn ("STOP. Do NOT write code or explain.") wiped the visible code from the conversation. -The new ``_HAS_ANSWER_ARTIFACT`` regex blocks the re-prompt whenever -the response already contains a real answer artifact: a closed code -fence, an HTML page, a complete SVG, or a numbered list of items. +The guard recognises completed code fences (any markdown info string, +indented closing fence allowed), complete HTML documents, and complete +SVGs as answer artifacts. A numbered list is an artifact only when the +response does NOT also contain plan framing ("Here's my plan", a tool- +action verb following intent phrasing, etc.), so plan-only stalls of +the form ``Here's my plan:\\n1. search\\n2. summarise`` still re-prompt. """ from __future__ import annotations @@ -46,6 +49,9 @@ sys.modules.setdefault("structlog", _structlog_stub) from core.inference.llama_cpp import ( # noqa: E402 _HAS_ANSWER_ARTIFACT, _INTENT_SIGNAL, + _NUMBERED_LIST_ARTIFACT, + _PLAN_LIST_FRAMING, + _has_answer_artifact, ) @@ -81,23 +87,78 @@ def test_intent_signal_ignores_direct_answers(): assert not _INTENT_SIGNAL.search(s), f"_INTENT_SIGNAL must not match {s!r}" -# ── _HAS_ANSWER_ARTIFACT recognises substantive content ──────────── +# ── Code fence artifact detection ────────────────────────────────── def test_artifact_regex_detects_closed_code_fence(): """Closed Python code fence is an answer artifact.""" text = "First, let me set up pygame.\n```python\nimport pygame\npygame.init()\n```" - assert _HAS_ANSWER_ARTIFACT.search( - text - ), "Closed code fence must be detected as an answer artifact" + assert _has_answer_artifact(text) + + +def test_artifact_regex_detects_non_alpha_info_strings(): + """Common languages with digits / symbols in the fence info string + (python3, c++, c#, objective-c, ts-node, bash-session) must all be + recognised as complete code answers.""" + samples = [ + "First, let me write it.\n```python3\nprint('hi')\n```", + "First, let me write it.\n```c++\nint main() { return 0; }\n```", + "First, let me write it.\n```c#\nConsole.WriteLine(\"hi\");\n```", + "First, let me write it.\n```objective-c\nNSLog(@\"hi\");\n```", + "First, let me write it.\n```ts-node\nconsole.log('hi')\n```", + "First, let me script it.\n```bash-session\n$ echo hi\n```", + "First, let me show it.\n```python linenums=\"1\"\nprint('hi')\n```", + ] + for text in samples: + assert _has_answer_artifact(text), text + assert not _would_reprompt(text), text + + +def test_artifact_regex_detects_indented_close_fence(): + """A closing fence indented under a list / blockquote still counts. + Common when the model nests code in markdown structure.""" + text = "First, let me show:\n```python\nx = 1\n ```" + assert _has_answer_artifact(text) + + +def test_artifact_regex_ignores_open_code_fence(): + """An UNCLOSED code fence is not yet a complete artifact.""" + text = "Let me set up pygame.\n```python\nimport pygame" + assert not _has_answer_artifact(text) + + +def test_artifact_regex_ignores_plain_text(): + """Plain conversational text contains no artifact.""" + text = "First, I will search for the songs that charted #3 in 2015." + assert not _has_answer_artifact(text) + + +# ── HTML artifact detection ──────────────────────────────────────── def test_artifact_regex_detects_html_page(): - """HTML pages (doctype or root) are answer artifacts.""" + """Complete HTML pages (doctype optional, required) match.""" text_a = "" text_b = "Sure, here is the dashboard:\n..." - assert _HAS_ANSWER_ARTIFACT.search(text_a) - assert _HAS_ANSWER_ARTIFACT.search(text_b) + assert _has_answer_artifact(text_a) + assert _has_answer_artifact(text_b) + + +def test_artifact_regex_ignores_incomplete_html_mention(): + """A plan-only mention of / without close + must NOT be treated as a completed answer. Pre-fix the guard matched + bare `` skeleton, then add CSS and JavaScript.", + "First, I'll write a complete page with a button.", + "Let me design a structure for the dashboard.", + ] + for s in samples: + assert not _has_answer_artifact(s), s + + +# ── SVG artifact detection ───────────────────────────────────────── def test_artifact_regex_detects_complete_svg(): @@ -109,32 +170,57 @@ def test_artifact_regex_detects_complete_svg(): "" "" ) - assert _HAS_ANSWER_ARTIFACT.search(text) + assert _has_answer_artifact(text) -def test_artifact_regex_detects_numbered_list(): - """A list of 2+ numbered items is an answer artifact.""" +def test_artifact_regex_ignores_incomplete_svg(): + text = "Let me draw a sloth: bool: return bool( 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) - and not _HAS_ANSWER_ARTIFACT.search(stripped) + and not _has_answer_artifact(stripped) ) @@ -165,9 +251,7 @@ def test_no_reprompt_on_complete_python_game(): " if e.type == pygame.QUIT: break\n" "```" ) - assert not _would_reprompt( - content - ), "Re-prompt must not fire after a complete code block was produced" + assert not _would_reprompt(content) def test_no_reprompt_on_complete_svg(): @@ -185,12 +269,12 @@ def test_no_reprompt_on_complete_svg(): def test_no_reprompt_on_numbered_list_answer(): - """Response with intent + numbered list (Billboard-style) does NOT re-prompt.""" + """A list answer without plan framing does NOT re-prompt.""" content = ( "Here's my list of #3 hits:\n" - "1. Animals — Maroon 5\n" - "2. Take Me to Church — Hozier\n" - "3. Drag Me Down — One Direction\n" + "1. Animals - Maroon 5\n" + "2. Take Me to Church - Hozier\n" + "3. Drag Me Down - One Direction\n" ) assert not _would_reprompt(content) @@ -198,7 +282,7 @@ def test_no_reprompt_on_numbered_list_answer(): def test_reprompts_on_plan_only_stall(): """Response that is purely a plan and no artifact STILL re-prompts.""" content = "I'll search the web for the answer." - assert _would_reprompt(content), "Plan-only stalls must still trigger the re-prompt" + assert _would_reprompt(content) def test_reprompts_on_intent_with_open_fence(): @@ -207,44 +291,49 @@ def test_reprompts_on_intent_with_open_fence(): assert _would_reprompt(content) +def test_reprompts_on_numbered_plan_only_stall(): + """Numbered plan ("Here's my plan: 1. search 2. summarise") STILL + re-prompts. Pre-fix the numbered-list artifact branch suppressed + the tool-call nudge, which contradicted the PR's stated invariant.""" + content = ( + "Here's my plan:\n" + "1. Search the web for the current Billboard Hot 100 2015 data.\n" + "2. Use python to categorise the matching songs." + ) + assert _would_reprompt(content) + + +def test_reprompts_on_intent_with_numbered_action_plan(): + """Numbered list where each item is an action (search, fetch, ...) + paired with intent phrasing is treated as a plan, not an answer.""" + content = ( + "First, I'll do these:\n" + "1. Search the web\n" + "2. Compare the sources\n" + "3. Answer concisely" + ) + assert _would_reprompt(content) + + +def test_reprompts_on_incomplete_html_intent(): + """A plan-only mention of without close STILL re-prompts.""" + content = "First, I'll create an skeleton, then add CSS." + assert _would_reprompt(content) + + # ── Cross-platform line endings ──────────────────────────────────── def test_artifact_regex_handles_crlf_code_fence(): """Windows / CRLF-converted content still detects a closed fence.""" content = "First, let me code.\r\n```python\r\nimport sys\r\nprint('hi')\r\n```" - assert _HAS_ANSWER_ARTIFACT.search( - content - ), "CRLF (\\r\\n) line endings inside a code fence must still match" - - -def test_artifact_regex_handles_crlf_numbered_list(): - """CRLF numbered list also matches.""" - content = "Here's the plan:\r\n1. one\r\n2. two\r\n" - assert _HAS_ANSWER_ARTIFACT.search(content) + assert _has_answer_artifact(content) def test_artifact_regex_handles_mixed_lf_crlf(): """Mixed line endings (real-world: paste-and-edit on Windows).""" content = "Here's the code:\r\n```python\nimport sys\r\n```" - assert _HAS_ANSWER_ARTIFACT.search(content) - - -def test_no_backtrack_on_crlf_spam(): - """10K of `\\r\\n` repeats must complete fast. - - Pre-fix the numbered-list alternative `(?:^|\\r?\\n)\\s*\\d+\\.` would - O(n^2)-backtrack on this kind of input (measured at ~630ms for 10KB - of `\\r\\n` repeats). The post-fix `[ \\t]*` indent restriction - keeps it linear. - """ - import time - - payload = "\r\n" * 5000 - t0 = time.time() - _HAS_ANSWER_ARTIFACT.search(payload) - elapsed_ms = (time.time() - t0) * 1000 - assert elapsed_ms < 50, f"regex took {elapsed_ms:.1f}ms on 10KB CRLF spam" + assert _has_answer_artifact(content) def test_no_reprompt_on_crlf_complete_python_game(): @@ -259,6 +348,36 @@ def test_no_reprompt_on_crlf_complete_python_game(): " if e.type == pygame.QUIT: break\r\n" "```" ) - assert not _would_reprompt( - content - ), "CRLF-encoded complete fence must also suppress the re-prompt" + assert not _would_reprompt(content) + + +# ── ReDoS guards ─────────────────────────────────────────────────── + + +def test_no_backtrack_on_crlf_spam(): + """10K of `\\r\\n` repeats must complete fast. + + The numbered-list alternative previously used greedy `\\s*` which + O(n^2)-backtracked through embedded `\\r\\n` characters (~630 ms on + 10 KB). The current `[ \\t]*` indent restriction plus length-bounded + `[\\s\\S]{...}?` runs keep every alternative linear.""" + import time + + payload = "\r\n" * 5000 + t0 = time.time() + _has_answer_artifact(payload) + elapsed_ms = (time.time() - t0) * 1000 + assert elapsed_ms < 50, f"guard took {elapsed_ms:.1f}ms on 10KB CRLF spam" + + +def test_no_backtrack_on_open_html_spam(): + """Many `` close must still complete + quickly. Bounded `[\\s\\S]{0,4000}?` between the open and close caps + the scan per occurrence.""" + import time + + payload = "