"""Jinja2-based system prompt renderer with SSTI hardening.
Loads a ``.j2`` template file once at startup and renders it on each call
with room-specific and tool-specific context variables. Uses a
:class:`~jinja2.sandbox.SandboxedEnvironment` and recursively sanitises
user-controllable values to prevent server-side template injection.
"""
from __future__ import annotations
import logging
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from jinja2.sandbox import SandboxedEnvironment
from jinja2 import FileSystemLoader
from observability import observability
logger = logging.getLogger(__name__)
[docs]
class LoggingSandboxedEnvironment(SandboxedEnvironment):
"""Sandboxed Jinja2 environment that audits its creation and blocked attribute access.
Thin subclass of :class:`jinja2.sandbox.SandboxedEnvironment` that adds
observability to prompt compilation: it announces every environment it
constructs (tagged with the template name) and turns each silent sandbox
refusal into a logged warning, so that a server-side template-injection
probe leaves an auditable trail instead of failing quietly. The security
policy itself is unchanged -- only the logging is added on top.
Instantiated by :class:`PromptRenderer.__init__` while loading the
system-prompt template, by ``kg_agentic_extraction.py`` for the KG
extraction prompt, and by ``knowledge_anchoring/template_loader.py`` for
its shared sandbox; the SSTI test suite (``tests/test_jinja_sandbox_ssti``)
also exercises it directly.
"""
[docs]
def __init__(self, *args, **kwargs) -> None:
"""Build a logging sandboxed Jinja2 environment and announce it.
Pops the bespoke ``template_name`` keyword (used only for log
attribution) out of *kwargs* before delegating the rest of the
positional and keyword arguments to
:meth:`jinja2.sandbox.SandboxedEnvironment.__init__`, so the parent
never sees an unexpected argument. After the base environment is
constructed it emits an ``info`` log line through the module
``logger`` recording which template this environment was created
for, giving operators a breadcrumb whenever a new prompt-compilation
sandbox is spun up.
This constructor is invoked indirectly by
:class:`PromptRenderer.__init__`, which instantiates a
``LoggingSandboxedEnvironment`` while loading the system-prompt
template; no other internal caller constructs it directly.
Args:
*args: Positional arguments forwarded verbatim to the
:class:`~jinja2.sandbox.SandboxedEnvironment` base class.
**kwargs: Keyword arguments forwarded to the base class, except
for the consumed ``template_name`` entry which defaults to
``"unknown"`` and is used purely for the log message.
"""
template_name = kwargs.pop("template_name", "unknown")
super().__init__(*args, **kwargs)
logger.info(
"Instantiated sandboxed rendering environment for prompt compilation: template_name=%s",
template_name,
)
[docs]
def is_safe_attribute(self, obj, attr, value) -> bool:
"""Decide whether template code may read an attribute, logging denials.
Delegates the actual security judgement to
:meth:`jinja2.sandbox.SandboxedEnvironment.is_safe_attribute`, which
applies Jinja2's sandbox policy (blocking dunder access and other
attributes flagged unsafe). When the base class refuses access, this
override emits a ``warning`` through the module ``logger`` naming the
attribute and the object type involved, turning otherwise-silent
sandbox refusals into an auditable signal of a possible server-side
template-injection attempt. The boolean verdict is then returned
unchanged so the sandbox still enforces the same restriction.
Jinja2's sandbox machinery calls this automatically during template
rendering for every attribute access; there is no direct internal
caller, though the sibling sandbox in
``knowledge_anchoring/template_loader.py`` mirrors this pattern.
Args:
obj: The object whose attribute is being accessed inside a
template.
attr: The attribute name the template is attempting to read.
value: The resolved attribute value supplied by Jinja2.
Returns:
bool: ``True`` if the attribute access is permitted, ``False`` if
the sandbox blocks it.
"""
safe = super().is_safe_attribute(obj, attr, value)
if not safe:
logger.warning(
"Blocked access to unsafe attribute or code injection block in template: attribute_name=%s, error=%s",
attr,
f"Access to attribute {attr!r} of {type(obj).__name__!r} object is blocked",
)
return safe
_JINJA_META_RE = re.compile(r"\{\{|\}\}|\{%|%\}|\{#|#\}")
[docs]
def sanitize_context(value: Any) -> Any:
"""Recursively strip Jinja2 metacharacters from user-controllable strings.
Replaces ``{{``, ``}}``, ``{%``, ``%}``, ``{#``, ``#}`` with
Unicode full-width lookalikes so they cannot be interpreted as
template syntax if an ``| tojson`` filter is ever omitted.
Non-string leaves (ints, floats, bools, ``None``) pass through
unchanged. Dicts and lists are walked recursively.
"""
if isinstance(value, str):
return _JINJA_META_RE.sub(_replace_meta, value)
if isinstance(value, dict):
return {k: sanitize_context(v) for k, v in value.items()}
if isinstance(value, (list, tuple)):
sanitized = [sanitize_context(v) for v in value]
return type(value)(sanitized)
return value
def _replace_meta(match: re.Match) -> str:
"""Map each Jinja2 metacharacter pair to a safe full-width substitute.
Translates one matched delimiter run -- an opening or closing brace pair
such as the expression, statement, or comment markers -- into a visually
similar Unicode full-width lookalike so the substring can no longer be
parsed as template syntax. Pure in-memory lookup with no I/O or side
effects; it is the replacement callback handed to ``_JINJA_META_RE.sub``
inside :func:`sanitize_context` and has no other callers.
Args:
match: The regex match for one Jinja2 metacharacter pair captured by
``_JINJA_META_RE``.
Returns:
The full-width Unicode replacement string for the matched delimiter.
"""
return {
"{{": "\uff5b\uff5b",
"}}": "\uff5d\uff5d",
"{%": "\uff5b\uff05",
"%}": "\uff05\uff5d",
"{#": "\uff5b\uff03",
"#}": "\uff03\uff5d",
}[match.group()]
[docs]
class PromptRenderer:
"""Render a Jinja2 system-prompt template with per-request context.
Uses :class:`~jinja2.sandbox.SandboxedEnvironment` to prevent
template injection even if a caller accidentally passes unsanitised
user data.
Parameters
----------
template_path:
Path to the ``.j2`` template file (e.g. ``"system_prompt.j2"``).
default_extras:
Optional dict of variables injected into **every** render call
(e.g. the list of registered tools). Per-call context takes
precedence over these defaults.
"""
[docs]
def __init__(
self,
template_path: str | Path,
default_extras: dict[str, Any] | None = None,
) -> None:
"""Load and compile the system-prompt template into a sandboxed renderer.
Resolves and validates the template path, then builds a
:class:`LoggingSandboxedEnvironment` (with trailing-newline preservation
and block trimming/stripping enabled) backed by a
:class:`~jinja2.FileSystemLoader` rooted at the template's parent
directory, compiles the named template once, and stashes the
``default_extras`` injected into every later render. The compilation
happens here so per-request renders stay cheap. Reads the filesystem to
confirm the template exists and emits an ``info`` log line through the
module ``logger`` once it is loaded.
Constructed by callers across the codebase -- the inference worker
(``inference_main.py``), message-context injection
(``message_processor/context_injections.py``), subagent tooling
(``tools/subagent_tools.py``), the SWORD overlay test suites, and the
SSTI tests -- typically once at startup with ``system_prompt.j2``.
Args:
template_path: Path to the ``.j2`` template file (e.g.
``"system_prompt.j2"``); accepts a string or
:class:`~pathlib.Path`.
default_extras: Optional variables injected into every
:meth:`render` call (e.g. the registered-tool list);
per-call context overrides these. Defaults to an empty dict.
Raises:
FileNotFoundError: If no file exists at ``template_path``.
"""
path = Path(template_path)
if not path.exists():
raise FileNotFoundError(
f"System prompt template not found: {path}",
)
self._env = LoggingSandboxedEnvironment(
loader=FileSystemLoader(str(path.parent)),
keep_trailing_newline=True,
trim_blocks=True,
lstrip_blocks=True,
template_name=path.name,
)
self._template = self._env.get_template(path.name)
self.default_extras: dict[str, Any] = default_extras or {}
logger.info("Loaded system prompt template from %s", path)
[docs]
def render(self, context: dict[str, Any] | None = None) -> str:
"""Render the template with the supplied *context*.
All values in *context* are recursively sanitised to strip
Jinja2 metacharacters before rendering.
The following keys are automatically injected if not already present:
* ``current_date`` -- today's date in ``YYYY-MM-DD`` format (UTC).
Keys from *default_extras* (set at init or later) are included but
can be overridden by *context*.
"""
ctx: dict[str, Any] = {}
ctx.update(sanitize_context(dict(self.default_extras)))
if context:
ctx.update(sanitize_context(context))
ctx.setdefault(
"current_date",
datetime.now(timezone.utc).strftime("%Y-%m-%d"),
)
# SWORD System-Scale Overlay
pseudo_origins = ctx.pop("_sword_pseudo_origins", None)
overlay_active = ctx.pop("_sword_overlay_active", False)
system_prompt = self._template.render(ctx)
# NOTE: Whitespace inconsistency exists here. The rendered template output
# has varying indentation. When pseudo_origins are present, json.dumps()
# reformats the entire prompt with strictly indent=2. We accept this diff
# to avoid the CPU cost of roundtripping JSON on every single message when
# no overlay is active.
if overlay_active or pseudo_origins is not None:
try:
import json
from sword.overlay import SpiritGraphOverlayBuilder
# Parse the rendered output string into a JSON dictionary
skeleton = json.loads(system_prompt)
if isinstance(skeleton, dict) and "ROOT" in skeleton:
builder = SpiritGraphOverlayBuilder()
# Apply the overlay sync using either pre-fetched origins or Postgres cache
overlaid = builder.build_overlay_sync(skeleton, pseudo_origins)
# Serialize the result back to string
system_prompt = json.dumps(overlaid, indent=2, ensure_ascii=False)
logger.info("[sword] System-Scale SpiritGraph prompt overlay applied successfully.")
else:
logger.debug("[sword] Rendered skeleton is not a SpiritGraph (missing 'ROOT' key); skipping overlay.")
except json.JSONDecodeError as e:
logger.warning("[sword] Failed to parse rendered system_prompt.j2 as JSON; overlay aborted: %s", e)
observability.increment("sword_overlay_parse_failure")
except Exception as e:
logger.warning("[sword] Failed to apply system-scale SpiritGraph prompt overlay: %s", e)
observability.increment("sword_overlay_error")
return system_prompt