, ) are NOT escaped. Clean input is wrapped in plain-text boundary markers as a secondary semantic defense (OWASP structured-prompt guidance). """ from __future__ import annotations import logging import re from collections.abc import Awaitable, Callable from typing import override from langchain.agents import AgentState from langchain.agents.middleware import AgentMiddleware from langchain.agents.middleware.types import ( ModelCallResult, ModelRequest, ModelResponse, ) from langchain_core.messages import HumanMessage from langgraph.errors import GraphBubbleUp from deerflow.agents.middlewares.message_utils import requires_input_sanitization from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY, message_content_to_text logger = logging.getLogger(__name__) # Finite set of blocked tag names: system-reserved + common injection patterns. # # Maintenance: when adding a new framework block tag that the system emits into # model input, you MUST also update the expected count in # test_input_sanitization_middleware.py::test_denylist_covers_framework_authority_blocks. # The test pins the exact number of blocked tags so a new framework tag cannot # be added without the corresponding regression guard. _BLOCKED_TAG_NAMES: frozenset[str] = frozenset( { # Framework-injected structured/authority blocks. The lead-agent system # prompt's "System-Context Confidentiality" section (agents/lead_agent/ # prompt.py) declares *every* such tag trusted internal data — it names a # few then says "and all other structured tags". So the denylist must # cover the framework's authority blocks as a class, not a hand-picked # subset: any one of them, forged in untrusted input, mimics trusted # framework context. Enumerated from the block tags the framework actually # emits into model input (system prompt + hidden-context/reminder # middlewares) and pinned against drift by # test_input_sanitization_middleware.py::test_denylist_covers_framework_authority_blocks. # Both spellings of the reminder block are covered: "system-reminder" # (dynamic-context) and "system_reminder" (todo/terminal middlewares). # # Subagents share this denylist: build_subagent_runtime_middlewares reuses # the same _build_runtime_middlewares base, so both sanitization paths guard # subagent model input too. The subagent system-prompt blocks # (file_editing_workflow / guidelines / output_format / working_directory) # are therefore authority blocks of the same class as the lead-agent ones. "system-reminder", "system_reminder", "memory", "current_date", "think", "analysis", "role", "soul", "self_update", "thinking_style", "clarification_system", "critical_reminders", "response_style", "citations", "current_uploads", "subagent_system", "skill_system", "skill_index", "available_skills", "disabled_skills", "memory_tool_system", "todo_list_system", "durable_context_data", "slash_skill_activation", "mcp_routing_hints", "available-deferred-tools", "goal_continuation", "background_task_event", "file_editing_workflow", "guidelines", "output_format", "working_directory", # Subagent system-prompt block (general_purpose.py): declares the task # tool off-limits. Forging this in untrusted input could trick the # model into believing it has (or lacks) tool restrictions it does not. "tool_restrictions", # Subagent report-contract blocks (subagents/report_contract.py, RFC # #4651 PR3): injected by the executor into every subagent system # prompt (the criteria pointer note carries no criterion values — # those stay in the untrusted task message). Forging them in # untrusted input could impersonate the verification contract (e.g. # pre-declaring acceptance criteria as met). "report_contract", "acceptance_criteria", # Common prompt-injection tag patterns "system", "instruction", "important", "override", "ignore", "prompt", } ) # Matches a full blocked tag: , , , , bare ]*>?", re.IGNORECASE, ) # Plain-text boundary markers (OWASP structured-prompt guidance). _USER_INPUT_BEGIN = "--- BEGIN USER INPUT ---" _USER_INPUT_END = "--- END USER INPUT ---" # Neutralized forms injected when the user's text already contains a marker. # These look visually similar but do not match the real boundary delimiters. _NEUTRALIZED_BEGIN = "[BEGIN USER INPUT]" _NEUTRALIZED_END = "[END USER INPUT]" # Matches either boundary token as a standalone line or embedded in text. _BOUNDARY_TOKEN_RE = re.compile( re.escape(_USER_INPUT_BEGIN) + r"|" + re.escape(_USER_INPUT_END), ) def _escape_tag_match(match: re.Match) -> str: """Escape < and > in a blocked-tag match so it renders as literal text.""" return match.group(0).replace("<", "<").replace(">", ">") def _neutralize_boundary_tokens(text: str) -> str: """Replace real BEGIN/END USER INPUT markers with look-alike inert forms.""" return _BOUNDARY_TOKEN_RE.sub( lambda m: _NEUTRALIZED_BEGIN if m.group(0) == _USER_INPUT_BEGIN else _NEUTRALIZED_END, text, ) def neutralize_untrusted_tags(text: str) -> str: """Neutralize framework/injection control tokens in untrusted text. Shared primitive for any content that originates outside the trust boundary and is about to enter the model context as *data* — currently the genuine user message (via :func:`frame_untrusted_text`) and remote tool results (web_fetch / web_search and friends, via :class:`ToolResultSanitizationMiddleware`). Applies exactly the two structural defenses, and nothing else: * blocked framework/injection tags (e.g. ````) are HTML-escaped to ``<system-reminder>`` so they lose their structural meaning while staying human-readable; * the plain-text ``--- BEGIN/END USER INPUT ---`` boundary markers are neutralized so untrusted content cannot forge or break out of the user-input boundary. It intentionally does **not** wrap the text in boundary markers: that framing is specific to the user message. Empty/whitespace-only text is returned unchanged so callers do not emit marker noise. """ if not text.strip(): return text text = _BLOCKED_TAG_PATTERN.sub(_escape_tag_match, text) return _neutralize_boundary_tokens(text) def frame_untrusted_text(text: str) -> str: """Sanitize untrusted text, then wrap it in user-input boundary markers. * Empty/whitespace-only → return unchanged (no marker noise). * Blocked tags → HTML-escape ``<``/``>`` (e.g. ```` → ``<system>``). * Boundary tokens in user text → neutralized so they cannot forge boundaries. * Already wrapped (strict prefix+suffix) → return text unchanged (idempotent). * Otherwise → wrap in boundary markers. """ if not text.strip(): return text text = _BLOCKED_TAG_PATTERN.sub(_escape_tag_match, text) # Idempotency: only skip if text is *exactly* wrapped (prefix+suffix), # not if the user merely typed the begin token somewhere. if text.startswith(_USER_INPUT_BEGIN) and text.endswith(_USER_INPUT_END): # Still neutralize boundary tokens in the inner content — a user # can forge the outer wrapping to bypass the neutralization below # and inject inner boundary markers (break-out attack). inner = text[len(_USER_INPUT_BEGIN) : -len(_USER_INPUT_END)] neutralized_inner = _neutralize_boundary_tokens(inner) if neutralized_inner == inner: return text return f"{_USER_INPUT_BEGIN}{neutralized_inner}{_USER_INPUT_END}" # Neutralize any boundary tokens the user may have embedded, preventing # both self-suppression (begin token skips wrapping) and break-out # (end token creates a premature boundary inside the payload). text = _neutralize_boundary_tokens(text) return f"{_USER_INPUT_BEGIN}\n{text}\n{_USER_INPUT_END}" def _check_user_content(text: str) -> str: """Backward-compatible internal alias for untrusted text framing.""" return frame_untrusted_text(text) class InputSanitizationMiddleware(AgentMiddleware[AgentState]): """Guardrail middleware that escapes prompt-injection tags in user input. Blocked tags are HTML-escaped (not rejected) so the user's intent is preserved while the tags lose their semantic significance. Clean input is wrapped in plain-text boundary markers. Transformation is temporary (wrap_model_call) — never written to state, which is why every genuine user message is re-sanitized on each call rather than only the newest. """ @staticmethod def _extract_text_from_content(content: str | list) -> tuple[str, list | None]: """Extract concatenated text from a plain-string or content-block-list. Returns ``(text, extracted_blocks)``. *extracted_blocks* is None when *content* is a string, or the list of text-content blocks when a list. A list can hold bare ``str`` items next to content-block dicts (``message_content_to_text`` treats both as text, and some IM/SDK clients send exactly that shape), so bare strings are collected too — skipping them would skip sanitization entirely for that message. """ if isinstance(content, str): return content, None if not isinstance(content, list): return "", None text_parts: list[str] = [] text_blocks: list[dict | str] = [] for block in content: if isinstance(block, str): if not block: # skip empty items — matches message_content_to_text behaviour continue text_parts.append(block) text_blocks.append(block) elif isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str): text = block["text"] if not text: # skip empty blocks — matches message_content_to_text behaviour continue text_parts.append(text) text_blocks.append(block) return "\n".join(text_parts), text_blocks @staticmethod def _rebuild_content( original_content: list, processed_text: str, text_blocks: list, ) -> list: """Replace text blocks with a single merged text block, preserving interleaved non-text blocks. For ``[text, image, text]`` the image block between the two text blocks is kept in place — only the text blocks are collapsed into one. """ text_block_ids = {id(b) for b in text_blocks} first = last = None for i, block in enumerate(original_content): if id(block) in text_block_ids: if first is None: first = i last = i if first is None: return original_content result: list = [*original_content[:first], {"type": "text", "text": processed_text}] # Re-insert any non-text blocks that sat between text blocks for i in range(first + 1, last + 1): if id(original_content[i]) not in text_block_ids: result.append(original_content[i]) result.extend(original_content[last + 1 :]) return result def _sanitize_message(self, msg: HumanMessage) -> HumanMessage | None: """Return a sanitized copy of *msg*, or None when nothing needs changing. Blocked tags are HTML-escaped (not rejected) so the user's intent is preserved while the tags lose their semantic significance. The original message is never mutated. """ content = msg.content text_content, text_blocks = self._extract_text_from_content(content) # No text at all (e.g. image-only message) — pass through if not text_content and not isinstance(content, str): logger.debug("_sanitize_message: no text content in message — passing through") return None # Sanitize only the user's original input when available (set by # UploadsMiddleware before it prepends the block), # so server-injected trusted blocks are never scanned for blocked # tags. Fall back to full-content scanning only when the marker is # absent — UploadsMiddleware sets it on upload turns, so plain text # messages without uploads won't have it. Full-content scanning is # safe for those: no server-injected block exists # to accidentally escape. preserved_kwargs = dict(msg.additional_kwargs or {}) original_user_content = preserved_kwargs.get(ORIGINAL_USER_CONTENT_KEY) if isinstance(original_user_content, str) and original_user_content: processed_user = _check_user_content(original_user_content) if processed_user != original_user_content: # Replace only the user's text suffix within the full # content — server-prepended blocks stay untouched. idx = text_content.rfind(original_user_content) if idx >= 0: processed = text_content[:idx] + processed_user else: # _extract_text_from_content and message_content_to_text # disagreed on text extraction — rfind failed (only # reachable for multimodal list content; see Decision 18). if isinstance(content, list) and len(content) >= 2: # content[0] is the server-injected # block (UploadsMiddleware # prepends it as the first element for list # content). Sanitize only user blocks (content[1:]) # and rebuild directly — _rebuild_content only # handles type:"text" blocks and would miss raw # strings or non-standard dict blocks that # message_content_to_text sees. logger.warning( "rfind failed on multimodal content; sanitizing user content blocks individually", ) new_content: list = [content[0]] for block in content[1:]: if isinstance(block, str): new_content.append(neutralize_untrusted_tags(block)) elif isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str): sanitized = neutralize_untrusted_tags(block["text"]) if sanitized != block["text"]: new_content.append({**block, "text": sanitized}) else: new_content.append(block) else: new_content.append(block) return HumanMessage( content=new_content, id=msg.id, name=msg.name, additional_kwargs=preserved_kwargs, ) # Cannot distinguish server block from user blocks # (non-list content or len(content) < 2). # Degrade to full-content sanitization — server # block may be escaped (UX degradation) but user # forgeries are still neutralized (no security # regression). logger.warning( "rfind failed with original_user_content set; cannot distinguish blocks, falling back to full-content sanitization", ) processed = _check_user_content(text_content) else: processed = text_content # no change needed elif isinstance(original_user_content, str): # Key is present but empty string (e.g. file upload with no # text input). No user text to sanitize; server-injected # blocks must survive untouched. processed = text_content else: processed = _check_user_content(text_content) # fallback if processed == text_content: # Already clean / already wrapped — no override needed return None if text_blocks: new_content = self._rebuild_content(content, processed, text_blocks) else: new_content = processed # Preserve the pre-sanitization user text so downstream consumers that # must see the genuine input (slash skill activation, regenerate) can # recover it after the BEGIN/END wrapping. Keep a valid value set by # UploadsMiddleware or an IM channel, but repair malformed metadata so # persistence never falls back to the wrapped model-facing content. if not isinstance(original_user_content, str): if ORIGINAL_USER_CONTENT_KEY in preserved_kwargs: logger.warning( "InputSanitizationMiddleware replaced non-string %s metadata: type=%s", ORIGINAL_USER_CONTENT_KEY, type(original_user_content).__name__, ) preserved_kwargs[ORIGINAL_USER_CONTENT_KEY] = message_content_to_text(content) logger.debug( "InputSanitizationMiddleware: original=%r -> processed=%r", content if isinstance(content, str) else "[content-blocks]", processed, ) return HumanMessage( content=new_content, id=msg.id, name=msg.name, additional_kwargs=preserved_kwargs, ) def _process_request(self, request: ModelRequest) -> ModelRequest: """Return a request with every genuine user message sanitized. Each genuine message is processed, not just the newest one. The transformation is request-scoped, so thread state keeps the raw text: sanitizing only the last turn would make the guardrail last exactly one turn, replaying an earlier turn's payload to the model verbatim — and outside the boundary markers, which the lead-agent prompt declares trusted internal framework data. Covering the whole history also keeps each turn's rendering stable across model calls, so the prompt prefix no longer changes shape behind the newest turn. Scope comes from ``requires_input_sanitization``: framework-injected messages stay excluded, because escaping their blocks would corrupt trusted context, while a caller-supplied message is covered even when it carries a framework marker — the Gateway marks those on the way in. """ messages = list(request.messages) changed = False for index, msg in enumerate(messages): if not requires_input_sanitization(msg): if isinstance(msg, HumanMessage): logger.debug( "_process_request: skipping non-genuine HumanMessage at pos=%d name=%s hide_from_ui=%s content_preview=%.80r", index, msg.name, msg.additional_kwargs.get("hide_from_ui"), msg.content, ) continue logger.debug("_process_request: found genuine user message at pos=%d content=%.120r", index, msg.content) # Recover per message rather than per request. Unexpected errors # fail open, and one unprocessable history row must not widen that # into "no sanitization this turn" — an attacker who can land such a # row would otherwise buy themselves an unescaped newest turn. try: sanitized = self._sanitize_message(msg) except GraphBubbleUp: raise except Exception: logger.warning( "Input guardrail failed on message at pos=%d; leaving it unchanged", index, exc_info=True, ) continue if sanitized is None: continue messages[index] = sanitized changed = True if not changed: return request return request.override(messages=messages) def _try_process(self, request: ModelRequest) -> ModelRequest: """Sanitize request; fail-open on unexpected errors. GraphBubbleUp propagates; other exceptions return the original request. """ try: return self._process_request(request) except GraphBubbleUp: raise except Exception: logger.warning( "Input guardrail processing failed; passing original request to model", exc_info=True, ) return request @override def wrap_model_call( self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse], ) -> ModelCallResult: return handler(self._try_process(request)) @override async def awrap_model_call( self, request: ModelRequest, handler: Callable[[ModelRequest], Awaitable[ModelResponse]], ) -> ModelCallResult: return await handler(self._try_process(request))