feat(agent): strip raw tool-scaffolding from streamed answers — deterministic filter with one bounded recovery

This commit is contained in:
2026-09-03 13:39:15 -04:00
parent 801639efcc
commit 575d6c88d0
38 changed files with 2793 additions and 50 deletions
+302 -6
View File
@@ -32,7 +32,7 @@ import logging
import uuid
from collections.abc import AsyncGenerator, AsyncIterator
from copy import deepcopy
from typing import Any, cast
from typing import TYPE_CHECKING, Any, cast
import pytest
from sqlalchemy.orm import Session
@@ -43,11 +43,15 @@ from app.rag import agent
from app.rag.agent import (
AGENT_TOOLS,
AgentHolder,
MalformedReplyError,
run_agent,
)
from app.rag.llm import LLMClient, LLMError, RetryPiece, StreamPiece, ToolCallPiece
from app.rag.prompts import TOOLS_SECTION, _base, build_deflect_prompt, build_high_prompt
if TYPE_CHECKING:
from app.rag.scaffolding import ScaffoldingFilter
def _settings(**kwargs: Any) -> Settings:
kwargs.setdefault("_env_file", None)
@@ -68,7 +72,13 @@ def _doc(source: str, path: str, title: str = "Title", content: str = "CONTENT")
class ScriptedLLM:
"""Canned stream sequences; records every ``chat_stream`` request so
the tests can assert on the messages and the ``tools`` passthrough."""
the tests can assert on the messages and the ``tools`` passthrough.
Phase 71: when the caller passes a ``ScaffoldingFilter``, the canned
content pieces are fed through it exactly like
``LLMClient.chat_stream`` (an empty clean result yields nothing; the
held tail is flushed on normal completion) — so a scaffolding-only
canned round streams no content pieces and leaves ``stripped_chars``
behind for the recovery policy to key on."""
def __init__(self, *streams: list[StreamPiece | ToolCallPiece]) -> None:
self.streams: list[list[StreamPiece | ToolCallPiece]] = list(streams)
@@ -78,12 +88,26 @@ class ScriptedLLM:
self,
messages: list[dict[str, str]],
tools: list[dict[str, Any]] | None = None,
scaffolding: ScaffoldingFilter | None = None,
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
self.requests.append((deepcopy(messages), tools))
if not self.streams:
raise AssertionError("ScriptedLLM ran out of canned streams")
for piece in self.streams.pop(0):
yield piece
pieces = self.streams.pop(0)
if scaffolding is None:
for piece in pieces:
yield piece
return
for piece in pieces:
if isinstance(piece, StreamPiece) and piece.kind == "content":
cleaned = scaffolding.feed(piece.text)
if cleaned:
yield StreamPiece("content", cleaned)
else:
yield piece
tail = scaffolding.flush()
if tail:
yield StreamPiece("content", tail)
async def _run(
@@ -1102,6 +1126,7 @@ class FailingLLM:
self,
messages: list[dict[str, str]],
tools: list[dict[str, Any]] | None = None,
scaffolding: ScaffoldingFilter | None = None, # phase 71: fed like the real client
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
index = len(self.requests)
pieces, error = (
@@ -1112,19 +1137,35 @@ class FailingLLM:
self.requests.append(
(deepcopy(messages), deepcopy(tools) if tools is not None else None)
)
return self._attempt(index, pieces, error)
return self._attempt(index, pieces, error, scaffolding)
async def _attempt(
self,
index: int,
pieces: list[StreamPiece | ToolCallPiece],
error: Exception | None,
scaffolding: ScaffoldingFilter | None,
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
try:
for piece in pieces:
yield piece
if (
scaffolding is not None
and isinstance(piece, StreamPiece)
and piece.kind == "content"
):
cleaned = scaffolding.feed(piece.text)
if cleaned:
yield StreamPiece("content", cleaned)
else:
yield piece
if error is not None:
# The tail is NOT flushed on a failed attempt — the real
# client only flushes a cleanly completed stream.
raise error
if scaffolding is not None:
tail = scaffolding.flush()
if tail:
yield StreamPiece("content", tail)
finally:
self.closed.append(index)
@@ -1402,11 +1443,16 @@ def test_high_prompt_tools_section_with_notes_and_kb() -> None:
def test_low_prompt_is_byte_identical_and_tool_free() -> None:
# Phase 71: the LOW prompt carries the owner-permitted plain-text
# line after the DEFLECT_MODE sentence (the marker-keying contract
# is unchanged; the line must not leak into the HIGH prompt —
# pinned in tests/unit/test_prompts.py).
expected = (
_base("LOW")
+ "\nDEFLECT_MODE: retrieval was weak — the titles below are the closest "
"your notes come to the question. They are titles only; do not pretend "
"they answer it. Use them to propose 2-3 alternative questions.\n"
"Reply in plain text only — you have no tools in this mode.\n"
+ "- T1\n- T2"
)
assert build_deflect_prompt(["T1", "T2"]) == expected
@@ -1418,3 +1464,253 @@ def test_low_prompt_is_byte_identical_and_tool_free() -> None:
):
assert "<tools>" not in prompt
assert TOOLS_SECTION not in prompt
# ---------- phase 71: the scaffolding recovery policy (deterministic only) ----------
#: The raw span from the 2026-09-03 incident (the E2E mock's trigger,
#: task 05) — a complete span the filter strips in full.
_INCIDENT_SPAN = (
"<|tool_call_start|>[read(path='/homelab/backup-notes.md')]<|tool_call_end|>"
)
def test_correction_instruction_is_the_harness_constant() -> None:
"""Verbatim constant: the E2E mock (task 05) keys on a stable
substring of it, so it must not drift."""
assert agent.CORRECTION_INSTRUCTION == (
"Your previous reply contained raw tool-call markup, which is not "
"interpreted here. Answer the user's question directly in plain "
"text — no tool syntax."
)
def test_malformed_reply_error_subclasses_llm_error() -> None:
assert issubclass(MalformedReplyError, LLMError)
assert not issubclass(LLMError, MalformedReplyError)
def test_scaffolding_only_round_gets_exactly_one_recovery(
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None:
"""A round whose visible content is pure scaffolding → exactly TWO
model requests: the normal round, then the ONE recovery —
``tools=None`` with :data:`CORRECTION_INSTRUCTION` folded into the
original single system message (the user message stays last). The
clean recovery answer ends the turn, the holder is untouched by the
recovery, and the strip was captured in the warning log."""
holder = AgentHolder()
llm = ScriptedLLM(
[StreamPiece("content", _INCIDENT_SPAN)],
[StreamPiece("content", "The clean recovery answer.")],
)
with caplog.at_level(logging.WARNING, logger="app.agent"):
pieces = asyncio.run(_run(llm, holder, _settings()))
# Nothing of the round's scaffolding was yielded — only the recovery.
assert pieces == [StreamPiece("content", "The clean recovery answer.")]
assert len(llm.requests) == 2 # the round + the one recovery
# The round: the original system prompt, the tools offered.
assert llm.requests[0][1] == AGENT_TOOLS
assert llm.requests[0][0] == [
{"role": "system", "content": "SYSTEM_PROMPT"},
{"role": "user", "content": "QUESTION"},
]
# The recovery: no tools, a SINGLE system message carrying the folded
# correction, the user message last.
assert llm.requests[1][1] is None
assert llm.requests[1][0] == [
{
"role": "system",
"content": "SYSTEM_PROMPT\n" + agent.CORRECTION_INSTRUCTION,
},
{"role": "user", "content": "QUESTION"},
]
# The recovery is a fixed policy, not a conversation: the holder is
# untouched by it.
assert holder.read_docs == [] and holder.tool_calls == 0
# The turn total feeds the API layer's ``scaffold_stripped=N`` field.
assert holder.scaffold_stripped == len(_INCIDENT_SPAN)
# One strip warning per stripped span, the span truncated to 200 chars.
strip_logs = [
r
for r in caplog.records
if r.levelno == logging.WARNING and r.getMessage().startswith("agent: stripped")
]
assert len(strip_logs) == 1
message = strip_logs[0].getMessage()
assert message.startswith(
f"agent: stripped {len(_INCIDENT_SPAN)} chars of tool-scaffolding in round 1:"
)
assert _INCIDENT_SPAN[:200] in message
def test_scaffolding_twice_settles_with_malformed_reply(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""The recovery answer is scaffolding again (a second empty reply) —
terminal: :class:`MalformedReplyError` (an :class:`LLMError` subclass)
after exactly two requests — no third request, no recovery of a
recovery."""
holder = AgentHolder()
llm = ScriptedLLM(
[StreamPiece("content", _INCIDENT_SPAN)],
[StreamPiece("content", _INCIDENT_SPAN)],
)
with pytest.raises(MalformedReplyError) as excinfo:
asyncio.run(_run(llm, holder, _settings()))
assert isinstance(excinfo.value, LLMError)
assert len(llm.requests) == 2 # exactly one recovery per turn
assert llm.requests[1][1] is None
assert agent.CORRECTION_INSTRUCTION in llm.requests[1][0][0]["content"]
assert holder.scaffold_stripped == 2 * len(_INCIDENT_SPAN)
def test_scaffolding_with_real_content_needs_no_recovery(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A round with real visible content PLUS scaffolding: the clean
content stands — one request only, the clean remainder yielded (no
raw tokens), no correction in any system prompt."""
holder = AgentHolder()
llm = ScriptedLLM(
[
StreamPiece("content", "Here it is: "),
StreamPiece("content", _INCIDENT_SPAN),
StreamPiece("content", " hope that helps."),
],
)
pieces = asyncio.run(_run(llm, holder, _settings()))
assert [p for p in pieces if isinstance(p, StreamPiece)] == [
StreamPiece("content", "Here it is: "),
StreamPiece("content", " hope that helps."),
]
assert len(llm.requests) == 1 # no recovery
for messages, _tools in llm.requests:
assert all(agent.CORRECTION_INSTRUCTION not in m["content"] for m in messages)
assert holder.scaffold_stripped == len(_INCIDENT_SPAN)
def test_clean_turn_carries_no_correction_and_no_strip(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A clean turn: one request, no correction in any system prompt,
zero stripped (the log field stays 0 — uniform)."""
holder = AgentHolder()
llm = ScriptedLLM(
[StreamPiece("thinking", "hmm "), StreamPiece("content", "a clean answer")]
)
pieces = asyncio.run(_run(llm, holder, _settings()))
assert pieces == [
StreamPiece("thinking", "hmm "),
StreamPiece("content", "a clean answer"),
]
assert len(llm.requests) == 1
for messages, _tools in llm.requests:
assert all(agent.CORRECTION_INSTRUCTION not in m["content"] for m in messages)
assert holder.scaffold_stripped == 0
def test_empty_round_without_a_strip_keeps_today_behavior(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Round content 0 with NOTHING stripped (an empty/thinking-only
answer) → return as today — no recovery, no error."""
holder = AgentHolder()
llm = ScriptedLLM([StreamPiece("thinking", "nothing to say, honestly")])
pieces = asyncio.run(_run(llm, holder, _settings()))
assert pieces == [StreamPiece("thinking", "nothing to say, honestly")]
assert len(llm.requests) == 1
assert holder.scaffold_stripped == 0
def test_scaffolding_round_with_tool_calls_needs_no_recovery(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A scaffolding-only round that ALSO carried tool calls: the tool
ran, and the policy keys on the no-calls exit only — no recovery (the
next round is a normal tools-offered round carrying the tool
history)."""
monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")])
holder = AgentHolder()
llm = ScriptedLLM(
[
StreamPiece("content", _INCIDENT_SPAN),
ToolCallPiece(id="call_1", name="ls", arguments={}),
],
[StreamPiece("content", "the answer after the tool")],
)
pieces = asyncio.run(_run(llm, holder, _settings()))
# The round's scaffolding was stripped (no raw delta), the tool frame
# and the next round's answer flowed on.
assert pieces == [
ToolCallPiece(id="call_1", name="ls", arguments={}),
StreamPiece("content", "the answer after the tool"),
]
assert len(llm.requests) == 2
assert llm.requests[1][1] == AGENT_TOOLS # a normal round, not a recovery
for messages, _tools in llm.requests:
assert all(
m["content"] is None or agent.CORRECTION_INSTRUCTION not in m["content"]
for m in messages
)
assert holder.tool_calls == 1
assert holder.scaffold_stripped == len(_INCIDENT_SPAN)
def test_recovery_after_tool_rounds_keeps_the_history(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A scaffolding-only answer round after a tool round: the recovery
keeps the SINGLE (folded) system message at the front and the tool
history intact behind it — no second system message, no duplicated
correction."""
monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")])
holder = AgentHolder()
llm = ScriptedLLM(
[ToolCallPiece(id="call_1", name="ls", arguments={})],
[StreamPiece("content", _INCIDENT_SPAN)],
[StreamPiece("content", "recovered after a tool round")],
)
pieces = asyncio.run(_run(llm, holder, _settings()))
assert pieces == [
ToolCallPiece(id="call_1", name="ls", arguments={}),
StreamPiece("content", "recovered after a tool round"),
]
assert len(llm.requests) == 3 # tool round + stripped round + recovery
assert llm.requests[2][1] is None
recovered = llm.requests[2][0]
assert recovered[0] == {
"role": "system",
"content": "SYSTEM_PROMPT\n" + agent.CORRECTION_INSTRUCTION,
}
assert recovered[1] == {"role": "user", "content": "QUESTION"}
assert len(recovered) == 4
assert recovered[2]["role"] == "assistant"
assert recovered[3] == {
"role": "tool",
"tool_call_id": "call_1",
"content": "1 documents:\nsource: S | path: a.md | title: A",
}
assert sum(1 for m in recovered if m["role"] == "system") == 1
assert holder.scaffold_stripped == len(_INCIDENT_SPAN)
def test_forced_final_scaffolding_only_settles_malformed(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""The round-cap forced final (``tools=None``) is filtered too: a
scaffolding-only forced answer never reaches the user raw — the turn
settles with :class:`MalformedReplyError` (the same terminal
semantics; this turn used no recovery, so nothing is doubled up)."""
monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")])
holder = AgentHolder()
llm = ScriptedLLM(
[ToolCallPiece(id="call_1", name="ls", arguments={})],
[StreamPiece("content", _INCIDENT_SPAN)],
)
with pytest.raises(MalformedReplyError):
asyncio.run(_run(llm, holder, _settings(agent_max_rounds=1)))
assert len(llm.requests) == 2
assert llm.requests[1][1] is None # the forced final — no recovery after it
assert holder.scaffold_stripped == len(_INCIDENT_SPAN)
+12 -1
View File
@@ -11,7 +11,7 @@ from __future__ import annotations
import json
import uuid
from collections.abc import Iterator
from typing import Any
from typing import TYPE_CHECKING, Any
import pytest
from fastapi.testclient import TestClient
@@ -25,6 +25,9 @@ from app.rag.llm import StreamPiece
from app.rag.retriever import RetrievedChunk, weak_hit_titles
from app.rag.suggestions import MAX_SUGGESTIONS, derive_suggestions
if TYPE_CHECKING:
from app.rag.scaffolding import ScaffoldingFilter
ANSWER = "I haven't done anything like that — try one of these instead!"
#: A small KB outline standing in for the lite-generated one (phase 31).
@@ -332,6 +335,11 @@ def test_low_prompt_has_titles_only_no_content() -> None:
assert "<relevance>LOW</relevance>" in prompt
assert "DEFLECT_MODE" in prompt
assert "HONESTY GATE" in prompt # the LOW rule is what the model follows
# Phase 71 (owner-permitted 2026-09-03): the deflection plain-text
# line — the LOW turn offers no tools, so any tool markup there is
# always wrong (prevention at the prompt; the filter + recovery is
# the backstop).
assert "Reply in plain text only — you have no tools in this mode." in prompt
assert "- Kubernetes Homelab Cluster" in prompt
assert "- Backup Strategy" in prompt
assert "ALPHA_DOC_CONTENT" not in prompt
@@ -347,6 +355,8 @@ def test_high_path_unaffected() -> None:
assert plan.suggestions == []
assert "<relevance>HIGH</relevance>" in plan.system_prompt
assert "DEFLECT_MODE" not in plan.system_prompt
# Phase 71: the deflection plain-text line never leaks into HIGH.
assert "Reply in plain text only" not in plan.system_prompt
assert "ALPHA_DOC_CONTENT" in plan.system_prompt
assert "BETA_DOC_CONTENT" in plan.system_prompt
assert [d.title for d in plan.docs] == ["Kubernetes Homelab Cluster", "Backup Strategy"]
@@ -438,6 +448,7 @@ class _CannedLLM:
self,
messages: list[dict[str, str]],
tools: list[dict[str, Any]] | None = None,
scaffolding: ScaffoldingFilter | None = None, # phase 71 pass-through
):
self.seen.append(messages)
self.seen_tools.append(tools)
+211 -5
View File
@@ -27,6 +27,7 @@ from app.rag.llm import (
ToolCallPiece,
chat_stream_retried,
)
from app.rag.scaffolding import ScaffoldingFilter
def _settings(**kwargs: Any) -> Settings:
@@ -337,13 +338,19 @@ class _FakeCompletions:
self.completion = completion
self.kwargs: dict | None = None
self.chat_kwargs: dict | None = None
#: Every SDK-shaped stream handed out — teardown tests assert the
#: phase-48 ``close()`` on them (phase 71 task 02: with/without
#: a filter, the teardown path is the same object).
self.streams: list[_FakeChatStream] = []
async def create(self, **kwargs) -> _FakeChatStream | _FakeCompletion:
self.kwargs = kwargs
if self.fail is not None:
raise self.fail
if kwargs.get("stream"):
return _FakeChatStream(self.chunks)
stream = _FakeChatStream(self.chunks)
self.streams.append(stream)
return stream
self.chat_kwargs = kwargs
assert self.completion is not None
return self.completion
@@ -721,6 +728,146 @@ def test_chat_stream_non_object_tool_arguments_raise_llm_error() -> None:
asyncio.run(drain())
# ---------- scaffolding filter integration (phase 71, task 02) ----------
#: The incident's raw span (2026-09-03) — the filter's reason to exist.
_INCIDENT_SPAN = "<|tool_call_start|>[read(path='/homelab/backup-notes.md')]<|tool_call_end|>"
def _collect_filtered(
llm: LLMClient,
messages: list[dict[str, str]],
scaffolding: ScaffoldingFilter | None,
) -> list[StreamPiece]:
"""Collect pieces from a tools-less filtered stream (phase 71): without
tools, no ToolCallPiece can appear (the phase-37 contract)."""
async def run() -> list[StreamPiece]:
pieces = [p async for p in llm.chat_stream(messages, scaffolding=scaffolding)]
assert all(isinstance(p, StreamPiece) for p in pieces)
return cast("list[StreamPiece]", pieces)
return asyncio.run(run())
def test_chat_stream_with_filter_strips_span_mid_stream() -> None:
"""A span mid-stream never reaches the pieces: the clean remainder
flows, ``stripped_chars`` is exact, and no piece carries scaffolding."""
f = ScaffoldingFilter()
llm, _ = _make_stream_client(
[_chunk("Hello "), _chunk(_INCIDENT_SPAN + " world"), _chunk("!")]
)
pieces = _collect_filtered(llm, _RETRY_MSGS, f)
assert [(p.kind, p.text) for p in pieces] == [
("content", "Hello "),
("content", " world"),
("content", "!"),
]
assert f.stripped_chars == len(_INCIDENT_SPAN)
assert all("<|" not in p.text for p in pieces if p.kind == "content")
def test_chat_stream_with_filter_span_split_across_chunks() -> None:
"""A span split across two chunks never emits a partial marker —
nothing leaks until the span completes, then the clean text on both
sides flows and the whole span counts as stripped."""
span = _INCIDENT_SPAN
cut = len("<|tool_call_start|>") # split right after the start token
f = ScaffoldingFilter()
llm, _ = _make_stream_client([_chunk("A " + span[:cut]), _chunk(span[cut:] + " B")])
pieces = _collect_filtered(llm, _RETRY_MSGS, f)
assert [(p.kind, p.text) for p in pieces] == [
("content", "A "),
("content", " B"),
]
assert f.stripped_chars == len(span)
def test_chat_stream_with_filter_yields_nothing_for_pure_scaffolding() -> None:
"""A content stream of pure scaffolding yields ZERO content pieces —
an empty clean result yields nothing (no empty ``delta`` frames)."""
f = ScaffoldingFilter()
llm, _ = _make_stream_client(
[_chunk(_INCIDENT_SPAN), _chunk("<|tool_calls|>")]
)
pieces = _collect_filtered(llm, _RETRY_MSGS, f)
assert pieces == []
assert f.stripped_chars == len(_INCIDENT_SPAN) + len("<|tool_calls|>")
def test_chat_stream_with_filter_leaves_thinking_raw() -> None:
"""Locked (phase 71): the scratchpad stays raw — a span inside
``reasoning_content`` is yielded verbatim and never counts as
stripped."""
f = ScaffoldingFilter()
llm, _ = _make_stream_client(
[_chunk("", reasoning=_INCIDENT_SPAN), _chunk("ok")]
)
pieces = _collect_filtered(llm, _RETRY_MSGS, f)
assert [(p.kind, p.text) for p in pieces] == [
("thinking", _INCIDENT_SPAN),
("content", "ok"),
]
assert f.stripped_chars == 0
def test_chat_stream_without_filter_is_the_raw_path() -> None:
"""``scaffolding=None`` (the default, and the explicit opt-out): a span
in content is yielded verbatim — byte-identical to the pre-phase-71
raw path for callers that do not pass a filter."""
llm, _ = _make_stream_client([_chunk(_INCIDENT_SPAN)])
explicit_none = _collect_filtered(llm, _RETRY_MSGS, None)
default = asyncio.run(_collect(llm, _RETRY_MSGS))
assert explicit_none == default == [StreamPiece("content", _INCIDENT_SPAN)]
def test_chat_stream_flushed_tail_precedes_tool_call_pieces() -> None:
"""Content-before-tools wire convention (phase 71): a stream that ends
with a held filter tail (a partial marker at EOF — flushed as-is) +
tool_calls deltas yields the flushed tail content piece BEFORE the
materialized ``ToolCallPiece``."""
f = ScaffoldingFilter()
llm, _ = _make_stream_client(
[
_chunk("done <|tool_call_st"), # held: a live prefix of the start token
_chunk(
None,
tool_calls=[_tool_call(0, id="call_t", name="ls")],
finish_reason="tool_calls",
),
]
)
async def run() -> list[StreamPiece | ToolCallPiece]:
return [
p
async for p in llm.chat_stream(_RETRY_MSGS, tools=_AGENT_TOOLS, scaffolding=f)
]
pieces = asyncio.run(run())
assert pieces == [
StreamPiece("content", "done "),
StreamPiece("content", "<|tool_call_st"),
ToolCallPiece(id="call_t", name="ls", arguments={}),
]
def test_chat_stream_abandon_with_filter_closes_stream() -> None:
"""Phase-48 teardown is independent of the phase-71 filter: a consumer
abandon mid-filter still closes the endpoint stream exactly once."""
f = ScaffoldingFilter()
llm, completions = _make_stream_client([_chunk(f"piece {i} ") for i in range(1, 6)])
async def run() -> None:
gen = llm.chat_stream(_RETRY_MSGS, scaffolding=f)
first = await gen.__anext__()
assert isinstance(first, StreamPiece)
assert first.text == "piece 1 "
await gen.aclose() # the consumer stops after the first piece
asyncio.run(run())
assert completions.streams[0].close_calls == 1
# ---------- one-shot chat: LLMClient.chat (phase 30, task 01) ----------
@@ -826,11 +973,15 @@ class _ScriptedClient(LLMClient):
] = []
#: Indices of attempts whose stream teardown has run.
self.closed: list[int] = []
#: The caller-owned filter each attempt's ``chat_stream`` received
#: (phase 71 task 02 — the retry primitive forwards it).
self.scaffoldings: list[ScaffoldingFilter | None] = []
def chat_stream(
self,
messages: list[dict[str, str]],
tools: list[dict[str, Any]] | None = None,
scaffolding: ScaffoldingFilter | None = None,
) -> AsyncGenerator[StreamPiece | ToolCallPiece, None]:
index = len(self.request_args)
pieces, error = (
@@ -841,15 +992,30 @@ class _ScriptedClient(LLMClient):
self.request_args.append(
(list(messages), list(tools) if tools is not None else None)
)
return self._attempt(index, pieces, error)
self.scaffoldings.append(scaffolding)
return self._attempt(index, pieces, error, scaffolding)
async def _attempt(
self, index: int, pieces: list[StreamPiece | ToolCallPiece],
self,
index: int,
pieces: list[StreamPiece | ToolCallPiece],
error: Exception | None,
scaffolding: ScaffoldingFilter | None,
) -> AsyncGenerator[StreamPiece | ToolCallPiece, None]:
try:
for piece in pieces:
yield piece
if (
scaffolding is not None
and isinstance(piece, StreamPiece)
and piece.kind == "content"
):
# Emulate the real contract (phase 71): content feeds
# the filter, an empty clean result yields nothing.
cleaned = scaffolding.feed(piece.text)
if cleaned:
yield StreamPiece("content", cleaned)
else:
yield piece
if error is not None:
raise error
finally:
@@ -875,12 +1041,18 @@ def _collect_retried(
tools: list[dict[str, Any]] | None = None,
retries: int,
delay: float,
scaffolding: ScaffoldingFilter | None = None,
) -> list[StreamPiece | ToolCallPiece | RetryPiece]:
async def run() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
return [
p
async for p in chat_stream_retried(
client, messages, tools=tools, retries=retries, delay=delay
client,
messages,
tools=tools,
retries=retries,
delay=delay,
scaffolding=scaffolding,
)
]
@@ -1016,6 +1188,40 @@ def test_retried_zero_delay_still_notifies(monkeypatch: pytest.MonkeyPatch) -> N
assert sleeps == [0]
def test_retried_forwards_the_filter_to_every_attempt(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Phase 71: the caller-owned filter reaches EVERY attempt's
``chat_stream`` — the same object on the dead attempt and the
surviving one."""
client = _ScriptedClient(
[([], LLMError("down")), ([StreamPiece("content", "ok")], None)]
)
f = ScaffoldingFilter()
_record_sleeps(monkeypatch)
pieces = _collect_retried(client, _RETRY_MSGS, retries=1, delay=0, scaffolding=f)
assert pieces == [RetryPiece(2, 2), StreamPiece("content", "ok")]
assert client.scaffoldings == [f, f]
def test_retried_reuses_the_unfed_filter_after_a_dead_attempt(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Phase 71: reusing the SAME filter across attempts is safe by
construction — the dead attempt emitted no piece, so the filter was
never fed; the surviving attempt filters through it as if fresh."""
span = _INCIDENT_SPAN
client = _ScriptedClient(
[([], LLMError("down")), ([StreamPiece("content", span + " clean")], None)]
)
f = ScaffoldingFilter()
_record_sleeps(monkeypatch)
pieces = _collect_retried(client, _RETRY_MSGS, retries=1, delay=0, scaffolding=f)
assert pieces == [RetryPiece(2, 2), StreamPiece("content", " clean")]
assert f.stripped_chars == len(span)
assert client.scaffoldings == [f, f]
def test_retried_abandon_mid_attempt_closes_the_attempt_stream() -> None:
"""Consumer abandon at a mid-attempt piece (the stop-generation path,
phase 48): no exception leaks and the attempt's stream is torn down
+77
View File
@@ -5,6 +5,11 @@ lite-generated KB outline): its own builder contract (empty → ``""``,
char budget + ``[…truncated…]`` marker, pathological budgets), its
placement between ``<relevance>`` and ``<tuning>`` in both modes, and
the byte-identical-when-absent convention (phase 15 precedent).
And the phase-71 deflection plain-text line (owner-permitted
2026-09-03): the LOW prompt = pre-phase text + exactly the one new
line; the ``DEFLECT_MODE`` marker-keying contract is unchanged and
the line never leaks into the HIGH prompt.
"""
from __future__ import annotations
@@ -129,11 +134,15 @@ def test_zero_note_prompt_is_byte_identical_to_pre_steering() -> None:
assert build_high_prompt([doc]) == (
_base("HIGH") + "\n<documents>\n" + block + "\n</documents>" + "\n" + TOOLS_SECTION
)
# Phase 71: the LOW prompt carries the owner-permitted plain-text
# line after the DEFLECT_MODE sentence (the marker-keying contract
# is unchanged — the E2E mock keys on the marker's presence).
assert build_deflect_prompt(["T1", "T2"]) == (
_base("LOW")
+ "\nDEFLECT_MODE: retrieval was weak — the titles below are the closest "
"your notes come to the question. They are titles only; do not pretend "
"they answer it. Use them to propose 2-3 alternative questions.\n"
"Reply in plain text only — you have no tools in this mode.\n"
+ "- T1\n- T2"
)
assert "<tuning>" not in build_high_prompt([doc])
@@ -299,11 +308,14 @@ def test_no_overview_prompt_is_byte_identical_to_pre_phase() -> None:
docs_block = "\n<documents>\n" + block + "\n</documents>" + "\n" + TOOLS_SECTION
high_plain = _base("HIGH") + docs_block
high_steered = _base("HIGH") + "\n" + build_steering_section(["be concise"]) + docs_block
# Phase 71: the owner-permitted plain-text line is part of the
# DEFLECT_MODE body in every LOW build (with or without steering).
low_plain = (
_base("LOW")
+ "\nDEFLECT_MODE: retrieval was weak — the titles below are the closest "
"your notes come to the question. They are titles only; do not pretend "
"they answer it. Use them to propose 2-3 alternative questions.\n"
"Reply in plain text only — you have no tools in this mode.\n"
+ "- T1\n- T2"
)
low_steered = (
@@ -313,6 +325,7 @@ def test_no_overview_prompt_is_byte_identical_to_pre_phase() -> None:
+ "\nDEFLECT_MODE: retrieval was weak — the titles below are the closest "
"your notes come to the question. They are titles only; do not pretend "
"they answer it. Use them to propose 2-3 alternative questions.\n"
"Reply in plain text only — you have no tools in this mode.\n"
+ "- T1\n- T2"
)
for kb in (None, "", " \n\t "):
@@ -399,3 +412,67 @@ def test_prompt_kb_section_over_settings_budget_capped_with_marker(
close = section.index("</knowledge_base>")
section = section[: close + len("</knowledge_base>")]
assert len(section) <= 200
# ---------- phase 71: the deflection plain-text line (prevention) ----------
#: The owner-permitted (2026-09-03) line appended to the ``DEFLECT_MODE``
#: body — the LOW prompt's only phase-71 change. The E2E mock keys on
#: the ``DEFLECT_MODE`` marker's *presence*, not the wording, so the
#: marker-keying contract is unchanged by the appended line.
PLAIN_TEXT_ONLY_LINE = "Reply in plain text only — you have no tools in this mode."
def _pre_phase71_low_body() -> str:
"""The ``DEFLECT_MODE`` body exactly as it was before phase 71."""
return (
"DEFLECT_MODE: retrieval was weak — the titles below are the closest "
"your notes come to the question. They are titles only; do not pretend "
"they answer it. Use them to propose 2-3 alternative questions.\n"
)
def test_low_prompt_is_pre_phase_plus_exactly_the_plain_text_line() -> None:
"""Diff pin: the LOW prompt = pre-phase text + exactly the one new
line, appended to the ``DEFLECT_MODE`` body; the weak-hit title list
follows exactly as before (and the line occurs exactly once)."""
prompt = build_deflect_prompt(["T1", "T2"])
assert prompt == (
_base("LOW")
+ "\n"
+ _pre_phase71_low_body()
+ PLAIN_TEXT_ONLY_LINE
+ "\n"
+ "- T1\n- T2"
)
assert prompt.count(PLAIN_TEXT_ONLY_LINE) == 1
assert prompt.endswith("- T1\n- T2") # the title list is untouched
def test_low_prompt_carries_the_line_and_keeps_the_mock_marker() -> None:
"""The new line is present in the LOW prompt (inside the
``DEFLECT_MODE`` body, after the marker) and the ``DEFLECT_MODE``
marker the E2E mock keys on stays put."""
for titles, tail in ((["T1"], "- T1"), ([], "(nothing close at all)")):
prompt = build_deflect_prompt(titles)
assert "DEFLECT_MODE" in prompt
assert PLAIN_TEXT_ONLY_LINE in prompt
assert prompt.index("DEFLECT_MODE") < prompt.index(PLAIN_TEXT_ONLY_LINE)
# The title list (or the no-titles fallback) follows the line
# exactly as before.
assert prompt.endswith(tail)
def test_plain_text_line_never_leaks_into_high_prompt() -> None:
"""The line is the LOW prompt's: every HIGH build (with/without
steering/overview) is unchanged and carries none of it."""
doc = _doc("kubernetes.md", "Talos Linux on three nodes.", "Kubernetes Homelab Cluster")
for notes, kb in (
(None, None),
(["be concise"], None),
(None, OVERVIEW),
(["be concise"], OVERVIEW),
):
high = build_high_prompt([doc], notes=notes, kb_overview=kb)
assert PLAIN_TEXT_ONLY_LINE not in high
assert "you have no tools" not in high
+280
View File
@@ -0,0 +1,280 @@
"""Unit tests: the deterministic tool-scaffolding filter (phase 71, task 01).
The matrix from the phase overview: a span in one chunk (exact
``stripped_chars``), the observed incident span split across chunks at
**every** boundary offset of the start token (0..len) plus a mid-span
and an end-token split, multiple spans in one chunk, the standalone
sibling tokens (alone, embedded, char-by-char), look-alikes that must
NOT be stripped, partial markers at EOF (``flush`` emits as-is), and
exact preservation of surrounding text.
"""
from __future__ import annotations
import pytest
from app.rag.scaffolding import SCAFFOLD_PATTERNS, ScaffoldingFilter
# The observed incident text (2026-09-03 — the deflected round that
# streamed raw model tool-scaffolding into the UI).
SPAN_START = "<|tool_call_start|>"
SPAN_END = "<|tool_call_end|>"
SPAN_CONTENT = "[read(path='/homelab/backup-notes.md')]"
SPAN = f"{SPAN_START}{SPAN_CONTENT}{SPAN_END}"
TOOL_CALLS = "<|tool_calls|>"
TOOL_CALL = "<|tool_call|>"
def _run(text: str, chunks: list[str] | None = None) -> tuple[str, ScaffoldingFilter]:
"""Feed *chunks* (or *text* as one chunk) + ``flush``; return the
(total output, filter) pair the assertions below work on."""
f = ScaffoldingFilter()
parts = [f.feed(chunk) for chunk in (chunks if chunks is not None else [text])]
parts.append(f.flush())
return "".join(parts), f
# ---------------------------------------------------------------- the registry
def test_registry_is_the_three_observed_forms() -> None:
# Every entry traces to the 2026-09-03 incident — no speculative
# entries (the registry is the extension point; new forms need an
# observed capture + a fixture here).
assert len(SCAFFOLD_PATTERNS) == 3
assert SCAFFOLD_PATTERNS[0].pattern == r"<\|tool_call_start\|>[\s\S]*?<\|tool_call_end\|>"
assert SCAFFOLD_PATTERNS[1].pattern == r"<\|tool_calls\|>"
assert SCAFFOLD_PATTERNS[2].pattern == r"<\|tool_call\|>"
# The span form is non-greedy: two spans each strip to their own end.
assert SCAFFOLD_PATTERNS[0].sub("", SPAN + "x" + SPAN) == "x"
# ---------------------------------------------- span in one chunk, split spans
def test_span_in_one_chunk_is_stripped_with_exact_count() -> None:
out, f = _run(SPAN)
assert out == ""
assert f.stripped_chars == len(SPAN)
def _split_offsets() -> list[int]:
"""Every split offset of the start token (0..len) + a mid-span split
and an end-token split (the task-01 matrix)."""
offsets = list(range(len(SPAN_START) + 1)) # 0 .. 19, every start-token offset
offsets += [
len(SPAN_START) + 1, # one char into the span body
len(SPAN_START) + len(SPAN_CONTENT) // 2, # mid-span
len(SPAN) - len(SPAN_END), # exactly at the end token
len(SPAN) - 5, # inside the end token
len(SPAN) - 1, # the very last char
]
return offsets
@pytest.mark.parametrize("offset", _split_offsets())
def test_span_split_at_boundary_offsets_emits_nothing_until_complete(offset: int) -> None:
f = ScaffoldingFilter()
first = f.feed(SPAN[:offset])
second = f.feed(SPAN[offset:])
third = f.flush()
assert first == "" # nothing ever emits until the span completes
assert second == ""
assert third == ""
assert f.stripped_chars == len(SPAN)
def test_span_completed_by_a_later_chunk_then_content_follows() -> None:
f = ScaffoldingFilter()
assert f.feed(f"lead {SPAN_START}part") == "lead " # prose before the span flows
assert f.feed(f"{SPAN_END} tail") == " tail"
assert f.flush() == ""
assert f.stripped_chars == len(SPAN_START) + 4 + len(SPAN_END) # body: "part"
def test_span_with_nested_start_token_strips_to_first_end() -> None:
# Non-greedy: the span runs from the FIRST start to the FIRST end.
text = f"{SPAN_START}a{SPAN_START}b{SPAN_END}"
out, f = _run(text)
assert out == ""
assert f.stripped_chars == len(text)
def test_spans_across_chunks_with_a_standalone_token_after() -> None:
f = ScaffoldingFilter()
assert f.feed(SPAN[:25]) == "" # inside the span body
assert f.feed(SPAN[25:] + " mid " + TOOL_CALLS) == " mid "
assert f.flush() == ""
assert f.stripped_chars == len(SPAN) + len(TOOL_CALLS)
# ------------------------------------------------------- multiple spans / tokens
def test_two_spans_in_one_chunk_strip_each_to_their_own_end() -> None:
out, f = _run(f"{SPAN} middle {SPAN}")
assert out == " middle "
assert f.stripped_chars == 2 * len(SPAN)
def test_two_spans_with_different_content() -> None:
other = f"{SPAN_START}[grep(pattern='vault')]{SPAN_END}"
out, f = _run(f"before{SPAN}mid{other}after")
assert out == "beforemidafter"
assert f.stripped_chars == len(SPAN) + len(other)
def test_repeated_strip_until_no_complete_match_remains() -> None:
# A standalone token embedded in a would-be token: step (a) repeats
# leftmost-complete removals until the buffer is clean.
out, f = _run(f"x<|tool_call{TOOL_CALLS}|>y")
assert out == "xy"
assert f.stripped_chars == len(TOOL_CALLS) + len(TOOL_CALL)
# ------------------------------------------------------------ standalone tokens
@pytest.mark.parametrize(
("token", "expected"),
[(TOOL_CALLS, 14), (TOOL_CALL, 13)],
)
def test_standalone_token_alone_is_stripped(token: str, expected: int) -> None:
out, f = _run(token)
assert out == ""
assert f.stripped_chars == expected
@pytest.mark.parametrize("token", [TOOL_CALLS, TOOL_CALL])
def test_standalone_token_embedded_in_a_line_is_stripped(token: str) -> None:
out, f = _run(f"line {token} tail\n")
assert out == "line tail\n"
assert f.stripped_chars == len(token)
def test_standalone_tokens_adjacent_to_text() -> None:
out, f = _run(f"a{TOOL_CALL}b{TOOL_CALLS}c")
assert out == "abc"
assert f.stripped_chars == len(TOOL_CALL) + len(TOOL_CALLS)
@pytest.mark.parametrize("token", [TOOL_CALLS, TOOL_CALL])
def test_standalone_token_char_by_char_never_emits(token: str) -> None:
out, f = _run("", list(token))
assert out == ""
assert f.stripped_chars == len(token)
def test_partial_standalone_token_completing_in_a_later_chunk() -> None:
f = ScaffoldingFilter()
assert f.feed("<|tool_ca") == ""
assert f.feed("lls|>") == ""
assert f.flush() == ""
assert f.stripped_chars == len(TOOL_CALLS)
def test_partial_token_resolving_to_prose_is_emitted() -> None:
# ``<|tool_cat|>`` is no known form — the held partial prefix must be
# released as content once the input can no longer complete a token.
f = ScaffoldingFilter()
assert f.feed("x" * 50 + "<|tool_ca") == "x" * 50 # live prefix held
assert f.feed("t|> done") == "<|tool_cat|> done"
assert f.flush() == ""
assert f.stripped_chars == 0
# ------------------------------------------------------------------- look-alikes
@pytest.mark.parametrize(
"text",
[
"tool_call_start", # the word as prose, no delimiters
"the tool_call and tool_calls words", # prose words
"<|tool_call_start|", # missing delimiter close — prose at EOF
"<|some_other_token|>", # unknown token
"<|tool_call_end|>", # a lone end token without a start
"a < b", # a lone angle bracket
],
)
def test_lookalikes_are_not_stripped_and_emit_verbatim(text: str) -> None:
out, f = _run(text)
assert out == text
assert f.stripped_chars == 0
def test_lone_end_token_before_an_open_start_stays_content_until_closed() -> None:
# The lone end token is content (emitted as-is); the trailing start
# opens a span and is HELD — if the span later closes, only the
# span is stripped, the lone end token stays.
f = ScaffoldingFilter()
assert f.feed(f"{SPAN_END}{SPAN_START}") == SPAN_END
assert f.feed(f"body{SPAN_END}") == ""
assert f.flush() == ""
assert f.stripped_chars == len(SPAN_START) + 4 + len(SPAN_END)
def test_lone_end_token_at_eof_flushes_as_is() -> None:
f = ScaffoldingFilter()
out = f.feed(SPAN_END)
assert out + f.flush() == SPAN_END # verbatim total (the tail may be split)
assert f.stripped_chars == 0
# ----------------------------------------------------------- EOF / empty / clean
def test_partial_start_token_at_eof_flushes_as_is() -> None:
# A partial marker at EOF is content, not scaffolding (pinned choice).
f = ScaffoldingFilter()
assert f.feed("<|tool_call_st") == ""
assert f.flush() == "<|tool_call_st"
assert f.stripped_chars == 0
@pytest.mark.parametrize(
"partial",
[
"<",
"<|",
"<|tool_ca",
"<|tool_call|", # one char short of the standalone token
"<|tool_calls|", # one char short of the standalone token
"<|tool_call_start|", # one char short of the span opening
],
)
def test_partial_tokens_at_eof_flush_as_is(partial: str) -> None:
out, f = _run(partial)
assert out == partial
assert f.stripped_chars == 0
def test_empty_chunk_is_a_noop() -> None:
f = ScaffoldingFilter()
assert f.feed("") == ""
assert f.feed("") == ""
assert f.flush() == ""
assert f.stripped_chars == 0
def test_empty_chunk_amid_pending_is_a_noop() -> None:
f = ScaffoldingFilter()
assert f.feed("<|tool_ca") == ""
assert f.feed("") == ""
assert f.feed("lls|>") == "" # completes the standalone token
assert f.flush() == ""
assert f.stripped_chars == len(TOOL_CALLS)
def test_flush_on_a_clean_stream_emits_nothing_new() -> None:
f = ScaffoldingFilter()
assert f.feed("hello ") == "hello "
assert f.feed("world") == "world"
assert f.flush() == ""
assert f.stripped_chars == 0
def test_surrounding_text_preserved_exactly() -> None:
# No reflow beyond the removal — both spaces around the span stay.
out, f = _run(f"hello {SPAN} world")
assert out == "hello world"
assert f.stripped_chars == len(SPAN)