fix: only parse marked output as chat messages

Prompt.chat_messages() parsed every line of the rendered template as a
possible ChatMessage, so template data that rendered to message JSON was
returned as a privileged system, assistant or tool message.

Each Prompt now carries a secret sentinel in its render context. The chat
tag, the tool filter and the media filters mark everything they emit with
it, and the parsers accept only marked lines and markers; anything else
stays plain text. The sentinel is per instance rather than per render
because the render cache keys on the context, so a sentinel that changed
between renders would make cached text unparseable. text() strips it, so
its output is unchanged and the marker cannot leak to whoever supplies
the data.

The same round trip through rendered text also affected
CompletionExtension._body_to_messages(), where data could inject tool
definitions as well as messages, and chat_message_from_text(), where data
could inject content blocks such as remote image URLs through the
no-chat-block fallback. Both are covered.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: masci

src/banks/env.py

@@ -6,6 +6,7 @@
 
 from .config import config
 from .filters import audio, cache_control, document, image, lemmatize, tool, video, xml
+from .utils import ensure_environment_sentinel
 
 
 def _add_extensions(_env):
@@ -43,4 +44,8 @@ def _add_extensions(_env):
 env.filters["document"] = document
 env.filters["to_xml"] = xml
 
+# Fallback for templates rendered straight off `env` instead of through a `Prompt`,
+# which would otherwise carry no sentinel and produce no parseable messages.
+ensure_environment_sentinel(env)
+
 _add_extensions(env)

src/banks/extensions/chat.py

@@ -6,6 +6,7 @@
 from jinja2.ext import Extension
 
 from banks.types import chat_message_from_text
+from banks.utils import ensure_environment_sentinel, sentinel_from_context
 
 SUPPORTED_TYPES = ("system", "user", "assistant")
 
@@ -25,6 +26,10 @@ class ChatExtension(Extension):
     # a set of names that trigger the extension.
     tags = {"chat"}  # noqa
 
+    def __init__(self, environment):
+        super().__init__(environment)
+        ensure_environment_sentinel(environment)
+
     def parse(self, parser):
         # We get the line number of the first token for error reporting
         lineno = next(parser.stream).lineno
@@ -54,18 +59,22 @@ def parse(self, parser):
             msg = f"Unknown role type '{attr_value.value}', use one of ({types})"
             raise TemplateSyntaxError(msg, lineno)
 
-        # Pass the role name to the CallBlock node
-        args: list[nodes.Expr] = [nodes.Const(attr_value.value)]
+        # Pass the render context and the role name to the CallBlock node
+        args: list[nodes.Expr] = [nodes.ContextReference(), nodes.Const(attr_value.value)]
 
         # Message body
         body = parser.parse_statements(("name:endchat",), drop_needle=True)
 
         # Build messages list
         return nodes.CallBlock(self.call_method("_store_chat_messages", args), [], [], body).set_lineno(lineno)
 
-    def _store_chat_messages(self, role, caller):
+    def _store_chat_messages(self, context, role, caller):
         """
         Helper callback.
         """
-        cm = chat_message_from_text(role=role, content=caller())
-        return cm.model_dump_json(exclude_none=True) + "\n"
+        sentinel = sentinel_from_context(context)
+        cm = chat_message_from_text(role=role, content=caller(), sentinel=sentinel)
+        # The sentinel marks this line as coming from a `chat` tag, so that template data
+        # rendering to a JSON message can't pick its own role. `model_dump_json` escapes
+        # newlines, so the content can't break out onto a line of its own either.
+        return sentinel + cm.model_dump_json(exclude_none=True) + "\n"

src/banks/extensions/completion.py

@@ -10,6 +10,7 @@
 
 from banks.errors import InvalidPromptError, LLMError
 from banks.types import ChatMessage, Tool
+from banks.utils import ensure_environment_sentinel, sentinel_from_context
 
 if TYPE_CHECKING:
     from litellm.types.utils import ChatCompletionMessageToolCall
@@ -44,6 +45,10 @@ class CompletionExtension(Extension):
     # Resolution never goes through importlib — only callables registered here are invocable.
     _callable_registry: ClassVar[dict[str, Callable[..., Any]]] = {}
 
+    def __init__(self, environment):
+        super().__init__(environment)
+        ensure_environment_sentinel(environment)
+
     @classmethod
     def register_callable(cls, name: str, func: Callable[..., Any]) -> None:
         cls._callable_registry[name] = func
@@ -72,8 +77,8 @@ def parse(self, parser):
         if attr_name.value not in SUPPORTED_KWARGS or attr_assign.value != "=":
             raise TemplateSyntaxError(error_msg, lineno)
 
-        # Pass the role name to the CallBlock node
-        args: list[nodes.Expr] = [nodes.Const(attr_value.value)]
+        # Pass the render context and the model name to the CallBlock node
+        args: list[nodes.Expr] = [nodes.ContextReference(), nodes.Const(attr_value.value)]
 
         # Message body
         body = parser.parse_statements(("name:endcompletion",), drop_needle=True)
@@ -105,7 +110,7 @@ def _get_tool_callable(self, tools: list[Tool], tool_call: "ChatCompletionMessag
             raise ValueError(msg)
         return self._callable_registry[name]
 
-    def _do_completion(self, model_name, caller):
+    def _do_completion(self, context, model_name, caller):
         """
         Helper callback.
         """
@@ -115,7 +120,7 @@ def _do_completion(self, model_name, caller):
         except ImportError as e:
             raise ImportError(LITELLM_INSTALL_MSG) from e
 
-        messages, tools = self._body_to_messages(caller())
+        messages, tools = self._body_to_messages(caller(), sentinel_from_context(context))
         message_dicts = [m.model_dump() for m in messages]
         tool_dicts = [t.model_dump(exclude={"import_path"}) for t in tools] or None
 
@@ -145,7 +150,7 @@ def _do_completion(self, model_name, caller):
         choices = cast(list[Choices], response.choices)
         return choices[0].message.content
 
-    async def _do_completion_async(self, model_name, caller):
+    async def _do_completion_async(self, context, model_name, caller):
         """
         Helper callback.
         """
@@ -155,7 +160,7 @@ async def _do_completion_async(self, model_name, caller):
         except ImportError as e:
             raise ImportError(LITELLM_INSTALL_MSG) from e
 
-        messages, tools = self._body_to_messages(caller())
+        messages, tools = self._body_to_messages(caller(), sentinel_from_context(context))
         message_dicts = [m.model_dump() for m in messages]
         tool_dicts = [t.model_dump(exclude={"import_path"}) for t in tools] or None
 
@@ -186,19 +191,27 @@ async def _do_completion_async(self, model_name, caller):
         choices = cast(list[Choices], response.choices)
         return choices[0].message.content
 
-    def _body_to_messages(self, body: str) -> tuple[list[ChatMessage], list[Tool]]:
-        """Converts each line in the body of a block into a chat message."""
+    def _body_to_messages(self, body: str, sentinel: str) -> tuple[list[ChatMessage], list[Tool]]:
+        """Converts each line in the body of a block into a chat message.
+
+        Only lines marked with the render sentinel are parsed: they are the ones the `chat`
+        tag and the `tool` filter produced. Unmarked lines are template data, which must not
+        be able to declare its own message role or register a tool.
+        """
         body = body.strip()
         messages = []
         tools = []
         for line in body.split("\n"):
+            if not line.startswith(sentinel):
+                continue
+            payload = line.removeprefix(sentinel)
             try:
                 # Try to parse a chat message
-                messages.append(ChatMessage.model_validate_json(line))
+                messages.append(ChatMessage.model_validate_json(payload))
             except ValidationError:  # pylint: disable=R0801
                 try:
                     # If not a chat message, try to parse a tool
-                    tools.append(Tool.model_validate_json(line))
+                    tools.append(Tool.model_validate_json(payload))
                 except ValidationError:
                     # Give up
                     pass

src/banks/extensions/docs.py

@@ -6,6 +6,9 @@ def chat(role: str):  # pylint: disable=W0613
     Text inside `chat` tags will be rendered as JSON strings representing chat messages. Calling `Prompt.chat_messages`
     will return a list of `ChatMessage` instances.
 
+    Only `chat` tags produce messages: template data that happens to render to message JSON is returned as plain
+    text instead, so untrusted input cannot pick its own role.
+
     Example:
         ```jinja
         {% chat role="system" %}

src/banks/filters/audio.py

@@ -9,8 +9,10 @@
 from urllib.parse import urlparse
 
 import filetype  # type: ignore[import-untyped]
+from jinja2 import pass_context
 
-from banks.types import AudioFormat, ContentBlock, InputAudio, resolve_binary
+from banks.types import CONTENT_BLOCK_END, AudioFormat, ContentBlock, InputAudio, content_block_start, resolve_binary
+from banks.utils import sentinel_from_context
 
 BASE64_AUDIO_REGEX = re.compile(r"audio\/.*;base64,.*")
 
@@ -53,7 +55,8 @@ def _get_audio_format_from_bytes(data: bytes) -> AudioFormat:
     return "mp3"
 
 
-def audio(value: str | bytes) -> str:
+@pass_context
+def audio(context, value: str | bytes) -> str:
     """Wrap the filtered value into a ContentBlock of type audio.
 
     The resulting ChatMessage will have the field `content` populated with a list of ContentBlock objects.
@@ -75,4 +78,4 @@ def audio(value: str | bytes) -> str:
     else:
         input_audio = InputAudio.from_path(Path(value))
     block = ContentBlock.model_validate({"type": "audio", "input_audio": input_audio})
-    return f"<content_block>{block.model_dump_json()}</content_block>"
+    return f"{content_block_start(sentinel_from_context(context))}{block.model_dump_json()}{CONTENT_BLOCK_END}"

src/banks/filters/cache_control.py

@@ -1,10 +1,14 @@
 # SPDX-FileCopyrightText: 2023-present Massimiliano Pippi <mpippi@gmail.com>
 #
 # SPDX-License-Identifier: MIT
-from banks.types import ContentBlock
+from jinja2 import pass_context
 
+from banks.types import CONTENT_BLOCK_END, ContentBlock, content_block_start
+from banks.utils import sentinel_from_context
 
-def cache_control(value: str, cache_type: str = "ephemeral") -> str:
+
+@pass_context
+def cache_control(context, value: str, cache_type: str = "ephemeral") -> str:
     """Wrap the filtered value into a ContentBlock with the proper cache_control field set.
 
     The resulting ChatMessage will have the field `content` populated with a list of ContentBlock objects.
@@ -17,8 +21,8 @@ def cache_control(value: str, cache_type: str = "ephemeral") -> str:
         ```
 
     Important:
-        this filter marks the content to cache by surrounding it with `<content_block>` and
-        `</content_block>`, so it's only useful when used within a `{% chat %}` block.
+        this filter marks the content to cache by surrounding it with content block markers,
+        so it's only useful when used within a `{% chat %}` block.
     """
     block = ContentBlock.model_validate({"type": "text", "text": value, "cache_control": {"type": cache_type}})
-    return f"<content_block>{block.model_dump_json()}</content_block>"
+    return f"{content_block_start(sentinel_from_context(context))}{block.model_dump_json()}{CONTENT_BLOCK_END}"

src/banks/filters/document.py

@@ -10,8 +10,17 @@
 from urllib.parse import urlparse
 
 import filetype  # type: ignore[import-untyped]
-
-from banks.types import ContentBlock, DocumentFormat, InputDocument, resolve_binary
+from jinja2 import pass_context
+
+from banks.types import (
+    CONTENT_BLOCK_END,
+    ContentBlock,
+    DocumentFormat,
+    InputDocument,
+    content_block_start,
+    resolve_binary,
+)
+from banks.utils import sentinel_from_context
 
 BASE64_DOCUMENT_REGEX = re.compile(r"(text|application)\/.*;base64,.*")
 
@@ -112,7 +121,8 @@ def _get_document_format_from_bytes(data: bytes) -> DocumentFormat:
     raise ValueError("Unsupported document format: " + kind.extension)
 
 
-def document(value: str | bytes) -> str:
+@pass_context
+def document(context, value: str | bytes) -> str:
     """Wrap the filtered value into a ContentBlock of type document.
 
     The resulting ChatMessage will have the field `content` populated with a list of ContentBlock objects.
@@ -134,4 +144,4 @@ def document(value: str | bytes) -> str:
     else:
         input_document = InputDocument.from_path(Path(value))
     block = ContentBlock.model_validate({"type": "document", "input_document": input_document})
-    return f"<content_block>{block.model_dump_json()}</content_block>"
+    return f"{content_block_start(sentinel_from_context(context))}{block.model_dump_json()}{CONTENT_BLOCK_END}"

src/banks/filters/image.py

@@ -7,7 +7,10 @@
 from pathlib import Path
 from urllib.parse import urlparse
 
-from banks.types import ContentBlock, ImageUrl
+from jinja2 import pass_context
+
+from banks.types import CONTENT_BLOCK_END, ContentBlock, ImageUrl, content_block_start
+from banks.utils import sentinel_from_context
 
 BASE64_PATH_REGEX = re.compile(r"image\/.*;base64,.*")
 
@@ -24,7 +27,8 @@ def _is_url(string: str) -> bool:
     return True
 
 
-def image(value: str | bytes) -> str:
+@pass_context
+def image(context, value: str | bytes) -> str:
     """Wrap the filtered value into a ContentBlock of type image.
 
     The resulting ChatMessage will have the field `content` populated with a list of ContentBlock objects.
@@ -37,8 +41,8 @@ def image(value: str | bytes) -> str:
         ```
 
     Important:
-        this filter marks the content to cache by surrounding it with `<content_block>` and
-        `</content_block>`, so it's only useful when used within a `{% chat %}` block.
+        this filter marks the content to cache by surrounding it with content block markers,
+        so it's only useful when used within a `{% chat %}` block.
     """
     if isinstance(value, bytes):
         image_url = ImageUrl.from_bytes(bytes_str=value)
@@ -48,4 +52,4 @@ def image(value: str | bytes) -> str:
         image_url = ImageUrl.from_path(Path(value))
 
     block = ContentBlock.model_validate({"type": "image_url", "image_url": image_url})
-    return f"<content_block>{block.model_dump_json()}</content_block>"
+    return f"{content_block_start(sentinel_from_context(context))}{block.model_dump_json()}{CONTENT_BLOCK_END}"

src/banks/filters/tool.py

@@ -3,10 +3,14 @@
 # SPDX-License-Identifier: MIT
 from typing import Callable
 
+from jinja2 import pass_context
+
 from banks.types import Tool
+from banks.utils import sentinel_from_context
 
 
-def tool(function: Callable) -> str:
+@pass_context
+def tool(context, function: Callable) -> str:
     """Inspect a Python callable and generates a JSON-schema ready for LLM function calling.
 
     Important:
@@ -16,4 +20,4 @@ def tool(function: Callable) -> str:
 
     t = Tool.from_callable(function)
     CompletionExtension.register_callable(function.__name__, function)
-    return t.model_dump_json() + "\n"
+    return sentinel_from_context(context) + t.model_dump_json() + "\n"

src/banks/filters/video.py

@@ -10,8 +10,10 @@
 
 import filetype  # type: ignore[import-untyped]
 from filetype.types.video import IsoBmff  # type: ignore[import-untyped]
+from jinja2 import pass_context
 
-from banks.types import ContentBlock, InputVideo, VideoFormat, resolve_binary
+from banks.types import CONTENT_BLOCK_END, ContentBlock, InputVideo, VideoFormat, content_block_start, resolve_binary
+from banks.utils import sentinel_from_context
 
 BASE64_VIDEO_REGEX = re.compile(r"video\/.*;base64,.*")
 
@@ -85,7 +87,8 @@ def _get_video_format_from_bytes(data: bytes) -> VideoFormat:
     return "mp4"
 
 
-def video(value: str | bytes) -> str:
+@pass_context
+def video(context, value: str | bytes) -> str:
     """Wrap the filtered value into a ContentBlock of type video.
 
     The resulting ChatMessage will have the field `content` populated with a list of ContentBlock objects.
@@ -107,4 +110,4 @@ def video(value: str | bytes) -> str:
     else:
         input_video = InputVideo.from_path(Path(value))
     block = ContentBlock.model_validate({"type": "video", "input_video": input_video})
-    return f"<content_block>{block.model_dump_json()}</content_block>"
+    return f"{content_block_start(sentinel_from_context(context))}{block.model_dump_json()}{CONTENT_BLOCK_END}"