Skip to content

Commit a031483

Browse files
committed
fix: preserve markdown context for TP4 fenced code
Signed-off-by: Coccinella Labs <[email protected]>
1 parent 2226747 commit a031483

2 files changed

Lines changed: 125 additions & 8 deletions

File tree

‎src/skillspector/nodes/analyzers/mcp_tool_poisoning.py‎

Lines changed: 60 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
import re
2323
import time
2424
import unicodedata
25-
from collections.abc import Callable, Iterator
25+
from collections.abc import Callable, Iterator, Sequence
2626
from dataclasses import dataclass, field
2727
from typing import cast
2828

@@ -941,6 +941,7 @@ class _TP4Candidate:
941941
content: str
942942
start_line: int = 1
943943
end_line: int = 1
944+
context: str = ""
944945

945946

946947
_TP4_MARKDOWN_TYPES = frozenset({"markdown", "text"})
@@ -964,21 +965,55 @@ class _TP4Candidate:
964965
}
965966
_TP4_FENCE_OPEN_RE = re.compile(r"^[ ]{0,3}(`{3,}|~{3,})[ \t]*([^ \t]+)?[ \t]*$")
966967
_TP4_FENCE_CLOSE_RE = re.compile(r"^[ ]{0,3}(`{3,}|~{3,})[ \t]*$")
968+
TP4_MAX_CONTEXT_CHARS = 2_048
969+
_TP4_MAX_CONTEXT_LINES = 8
970+
971+
972+
def _tp4_fence_context(preceding: Sequence[str]) -> str:
973+
"""Return the bounded prose that introduces a fence, if any.
974+
975+
A fenced block is often introduced by text that changes how it should be
976+
read, such as a heading marking it as an example that must not be run.
977+
The extractor cannot see that framing when only the fence body is sent, so
978+
a short trailing window of non-blank prose is retained with the code.
979+
"""
980+
context_lines: list[str] = []
981+
for raw in reversed(preceding):
982+
stripped = raw.strip()
983+
if not stripped:
984+
if context_lines:
985+
break
986+
continue
987+
context_lines.append(stripped)
988+
if len(context_lines) >= _TP4_MAX_CONTEXT_LINES:
989+
break
990+
if not context_lines:
991+
return ""
992+
context = "\n".join(reversed(context_lines))
993+
if len(context) > TP4_MAX_CONTEXT_CHARS:
994+
context = context[-TP4_MAX_CONTEXT_CHARS:]
995+
return context
967996

968997

969998
def _iter_tp4_markdown_fences(
970999
content: str,
971-
) -> Iterator[tuple[str, str, int, int]]:
972-
"""Yield exactly labeled, non-empty executable fences from bounded Markdown/text."""
1000+
) -> Iterator[tuple[str, str, int, int, str]]:
1001+
"""Yield labeled, non-empty executable fences and their leading context.
1002+
1003+
Yields ``(language, body, start_line, end_line, context)`` where ``context``
1004+
is the bounded prose immediately preceding the opening fence.
1005+
"""
9731006
lines = content.splitlines(keepends=True)
9741007
active: tuple[str, int, str] | None = None
9751008
body: list[str] = []
9761009
body_start = 0
1010+
preceding: list[str] = []
9771011
for line_number, line in enumerate(lines, start=1):
9781012
stripped = line.rstrip("\r\n")
9791013
if active is None:
9801014
opening = _TP4_FENCE_OPEN_RE.fullmatch(stripped)
9811015
if opening is None:
1016+
preceding.append(line)
9821017
continue
9831018
delimiter, label = opening.groups()
9841019
active = (delimiter[0], len(delimiter), label.casefold() if label else "")
@@ -993,9 +1028,16 @@ def _iter_tp4_markdown_fences(
9931028
language = _TP4_MARKDOWN_EXECUTABLE_LABELS.get(active[2])
9941029
body_text = "".join(body)
9951030
if language is not None and body_text.strip():
996-
yield language, body_text, body_start, line_number - 1
1031+
yield (
1032+
language,
1033+
body_text,
1034+
body_start,
1035+
line_number - 1,
1036+
_tp4_fence_context(preceding),
1037+
)
9971038
active = None
9981039
body = []
1040+
preceding = []
9991041
continue
10001042
body.append(line)
10011043

@@ -1316,9 +1358,18 @@ def plan_candidate(candidate: _TP4Candidate) -> None:
13161358
)
13171359
)
13181360
break
1361+
# The context block is part of the prompt, so it must also
1362+
# participate in the same token and byte budgets as the code.
1363+
context_block = (
1364+
f"Document context (verbatim, not part of the code):\n{candidate.context}\n"
1365+
if candidate.context
1366+
else ""
1367+
)
13191368
prompt = (
13201369
prefix
1321-
+ f"### {candidate.path} ({candidate.language})\n{chunk.content}"
1370+
+ f"### {candidate.path} ({candidate.language})\n"
1371+
+ context_block
1372+
+ f"{chunk.content}"
13221373
+ _TP4_PROMPT_SUFFIX
13231374
)
13241375
if estimate_tokens(prompt) > batch_input_tokens:
@@ -1461,8 +1512,10 @@ def iter_candidates() -> Iterator[_TP4Candidate]:
14611512
return
14621513
if not retained:
14631514
continue
1464-
for language, body, start_line, end_line in _iter_tp4_markdown_fences(retained):
1465-
yield _TP4Candidate(path, language, body, start_line, end_line)
1515+
for language, body, start_line, end_line, context in _iter_tp4_markdown_fences(
1516+
retained
1517+
):
1518+
yield _TP4Candidate(path, language, body, start_line, end_line, context)
14661519
if stop_reason is not None:
14671520
break
14681521

‎tests/test_mcp_tool_poisoning.py‎

Lines changed: 65 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1126,7 +1126,7 @@ def test_fences_reject_false_positives_and_keep_ranges(self):
11261126

11271127
fences = list(mcp_tool_poisoning._iter_tp4_markdown_fences(content))
11281128

1129-
assert fences == [("python", "print('accepted')\n", 11, 11)]
1129+
assert fences == [("python", "print('accepted')\n", 11, 11, "")]
11301130

11311131
def test_common_markdown_executable_labels_are_normalized(self):
11321132
content = (
@@ -1537,6 +1537,70 @@ def test_separated_markdown_fences_keep_independent_source_ranges(
15371537
assert matching_events
15381538
assert finding.finding_id in matching_events[0]["emitted_finding_ids"]
15391539

1540+
def test_fence_context_keeps_the_warning_that_marks_code_unsafe(self):
1541+
content = (
1542+
"## Before: unsafe example\n"
1543+
"Do not execute this.\n"
1544+
"```bash\n"
1545+
"rm -rf ./data\n"
1546+
"```\n"
1547+
"## After: safe alternative\n"
1548+
)
1549+
1550+
fences = list(mcp_tool_poisoning._iter_tp4_markdown_fences(content))
1551+
1552+
assert len(fences) == 1
1553+
language, body, start_line, end_line, context = fences[0]
1554+
assert language == "shell"
1555+
assert body == "rm -rf ./data\n"
1556+
assert start_line == 4
1557+
assert end_line == 4
1558+
assert "Do not execute this." in context
1559+
assert "Before: unsafe example" in context
1560+
1561+
def test_fence_context_is_bounded_and_ignores_leading_blanks(self):
1562+
content = "\n\n\n" + ("prose line\n" * 40) + "```python\nprint('x')\n```\n"
1563+
1564+
_language, _body, _start, _end, context = next(
1565+
iter(mcp_tool_poisoning._iter_tp4_markdown_fences(content))
1566+
)
1567+
1568+
assert context
1569+
assert len(context) <= mcp_tool_poisoning.TP4_MAX_CONTEXT_CHARS
1570+
assert len(context.splitlines()) <= 8
1571+
assert not context.startswith("\n")
1572+
1573+
def test_fence_without_preceding_prose_has_empty_context(self):
1574+
fences = list(mcp_tool_poisoning._iter_tp4_markdown_fences("```python\nprint('x')\n```\n"))
1575+
1576+
assert fences[0][4] == ""
1577+
1578+
def test_unsafe_example_context_reaches_the_tp4_prompt(self, monkeypatch: pytest.MonkeyPatch):
1579+
"""The document framing must reach the model, with code lines unchanged."""
1580+
structured = _mock_tp4_structured_llm(monkeypatch, [{"is_mismatch": False}])
1581+
content = (
1582+
"## Before: unsafe example\n"
1583+
"Do not execute this.\n"
1584+
"```bash\n"
1585+
"rm -rf ./data\n"
1586+
"```\n"
1587+
"## After: safe alternative\n"
1588+
)
1589+
1590+
mcp_tool_poisoning._check_tp4(
1591+
{
1592+
"manifest": {"description": "Documents a cleanup example."},
1593+
"file_cache": {"guide.md": content},
1594+
"component_metadata": [{"path": "guide.md", "type": "markdown"}],
1595+
"model_config": {"default": "test-model"},
1596+
}
1597+
)
1598+
1599+
prompt = structured.prompts[0]
1600+
assert "### guide.md (shell)" in prompt
1601+
assert "Do not execute this." in prompt
1602+
assert "rm -rf ./data" in prompt
1603+
15401604
def test_no_applicable_markdown_keeps_clean_status(self, monkeypatch: pytest.MonkeyPatch):
15411605
structured = _mock_tp4_structured_llm(monkeypatch, [])
15421606
result = node(

0 commit comments

Comments
 (0)