fix(SKY-11793): remove non-sandboxed second render in TextPromptBlock (SSTI) (#7137)

This commit is contained in:
Aaron Perez 2026-07-06 16:44:03 -05:00 committed by GitHub
parent 4d06850f1c
commit d723de621d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 75 additions and 28 deletions

View file

@ -4837,7 +4837,6 @@ class TextPromptBlock(Block):
async def send_prompt(
self,
prompt: str,
parameter_values: dict[str, Any],
workflow_run_id: str,
organization_id: str | None = None,
workflow_run_block_id: str | None = None,
@ -4850,7 +4849,9 @@ class TextPromptBlock(Block):
)
schema_to_use = json_schema or self.json_schema or _default_text_prompt_schema()
prompt = prompt_engine.load_prompt_from_string(prompt, **parameter_values)
# `prompt` is already fully rendered by format_potential_template_parameters().
# Keep send_prompt focused on delivery/retry formatting so parameter values
# stay literal after that render.
if schema_validation_failure:
prompt = _build_schema_validation_retry_prompt(prompt, schema_validation_failure)
prompt += (
@ -4950,8 +4951,6 @@ class TextPromptBlock(Block):
workflow_run_block_id=workflow_run_block_id,
organization_id=organization_id,
)
# get all parameters into a dictionary
parameter_values = {}
for parameter in self.parameters:
if not workflow_run_context.has_value(parameter.key):
LOG.warning(
@ -4971,12 +4970,6 @@ class TextPromptBlock(Block):
workflow_run_block_id=workflow_run_block_id,
organization_id=organization_id,
)
value = workflow_run_context.get_value(parameter.key)
secret_value = workflow_run_context.get_original_secret_value_or_none(value)
if secret_value:
continue
else:
parameter_values[parameter.key] = value
response: dict[str, Any] | list | str | None = None
schema_to_use = self.json_schema or _default_text_prompt_schema()
@ -4995,7 +4988,6 @@ class TextPromptBlock(Block):
try:
response = await self.send_prompt(
self.prompt,
parameter_values,
workflow_run_id,
organization_id,
workflow_run_block_id=workflow_run_block_id,

View file

@ -124,7 +124,7 @@ async def test_text_prompt_block_uses_selected_model(monkeypatch, model_name):
lambda template, **kwargs: template,
)
response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")
response = await block.send_prompt(block.prompt, workflow_run_id="workflow-run", organization_id="org-1")
assert captured["llm_key"] == expected_llm_key
assert captured["prompt_name"] == "text-prompt"
@ -188,7 +188,7 @@ async def test_text_prompt_block_uses_workflow_handler_when_no_override(monkeypa
lambda template, **kwargs: template,
)
response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")
response = await block.send_prompt(block.prompt, workflow_run_id="workflow-run", organization_id="org-1")
assert captured["llm_key"] == "default"
assert captured["default_handler"] == fake_secondary_handler
@ -251,7 +251,7 @@ async def test_text_prompt_block_prefers_prompt_type_config_over_secondary(monke
lambda template, **kwargs: template,
)
response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")
response = await block.send_prompt(block.prompt, workflow_run_id="workflow-run", organization_id="org-1")
assert captured["default_handler"] == prompt_config_handler
prompt_config_handler.assert_awaited_once()
@ -330,7 +330,7 @@ async def test_text_prompt_block_bad_llm_key_uses_same_runtime_path_as_no_overri
)
for block in blocks:
response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")
response = await block.send_prompt(block.prompt, workflow_run_id="workflow-run", organization_id="org-1")
assert response == {"llm_response": "secondary"}
assert captured == [
@ -395,7 +395,7 @@ async def test_text_prompt_block_uses_explicit_internal_llm_key_override(monkeyp
lambda template, **kwargs: template,
)
response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")
response = await block.send_prompt(block.prompt, workflow_run_id="workflow-run", organization_id="org-1")
assert captured["llm_key"] == "SPECIAL_INTERNAL_KEY"
assert captured["default_handler"] == fake_default_handler
@ -451,7 +451,7 @@ async def test_text_prompt_block_array_schema_does_not_force_dict(monkeypatch):
lambda template, **kwargs: template,
)
response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")
response = await block.send_prompt(block.prompt, workflow_run_id="workflow-run", organization_id="org-1")
assert captured["prompt_name"] == "text-prompt"
assert captured["force_dict"] is False
@ -500,7 +500,7 @@ async def test_text_prompt_block_object_schema_does_not_force_dict_before_valida
lambda template, **kwargs: template,
)
response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")
response = await block.send_prompt(block.prompt, workflow_run_id="workflow-run", organization_id="org-1")
assert captured["force_dict"] is False
assert response == {"invoice_search_string": "062026"}
@ -508,8 +508,10 @@ async def test_text_prompt_block_object_schema_does_not_force_dict_before_valida
@pytest.mark.asyncio
async def test_text_prompt_block_retry_feedback_is_not_rendered_as_template(monkeypatch):
# send_prompt receives the already-rendered prompt and must not render it again, so
# schema-validation retry feedback (which may contain `{{ }}`) reaches the model verbatim.
block = _make_text_prompt_block(
prompt="Hello {{ name }}",
prompt="Hello Alice",
json_schema={
"type": "object",
"properties": {"answer": {"type": "string"}},
@ -517,7 +519,6 @@ async def test_text_prompt_block_retry_feedback_is_not_rendered_as_template(monk
},
)
captured: dict[str, object] = {}
rendered_templates: list[str] = []
async def fake_handler(*, prompt: str, prompt_name: str, force_dict: bool, **kwargs):
captured["prompt"] = prompt
@ -529,11 +530,6 @@ async def test_text_prompt_block_retry_feedback_is_not_rendered_as_template(monk
def fake_get_override_handler(llm_key: str | None, *, default):
return default
def fake_load_prompt_from_string(template: str, **kwargs: str) -> str:
rendered_templates.append(template)
assert "previous response failed JSON schema validation" not in template
return template.replace("{{ name }}", kwargs["name"])
monkeypatch.setattr(
block_module.LLMAPIHandlerFactory,
"get_override_llm_api_handler",
@ -546,13 +542,12 @@ async def test_text_prompt_block_retry_feedback_is_not_rendered_as_template(monk
fake_resolve_default_llm_handler,
raising=False,
)
monkeypatch.setattr(prompt_engine, "load_prompt_from_string", fake_load_prompt_from_string)
await block.send_prompt(
block.prompt,
{"name": "Alice"},
workflow_run_id="workflow-run",
organization_id="org-1",
workflow_run_block_id=None,
schema_validation_failure="root: has 1 unexpected properties {{ dangerous_lookup }}",
)
@ -953,3 +948,64 @@ def test_format_potential_template_parameters_no_json_schema():
assert block.json_schema is None
assert block.prompt == "simple prompt"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"payload,forbidden",
[
("{{ 7*7 }}", "49"),
("{{ (1).__class__.__name__ }}", "int"),
],
)
async def test_text_prompt_block_does_not_evaluate_template_in_parameter_value(monkeypatch, payload, forbidden):
"""A template expression delivered as a parameter *value* must reach the model as a
literal and must never be evaluated by a second render.
The prompt is rendered once (sandboxed) to substitute parameter values as inert text;
it must not be rendered a second time, or the substituted value would be interpreted as
a live Jinja template (server-side template injection).
"""
block = _make_text_prompt_block(prompt="{{ payload }}")
ctx = _make_workflow_run_context({"payload": payload})
captured: dict[str, str] = {}
async def fake_handler(*, prompt: str, prompt_name: str, **kwargs):
captured["prompt"] = prompt
return {"answer": "ok"}
async def fake_resolve_default_llm_handler(*args, **kwargs):
return fake_handler
def fake_get_override_handler(llm_key: str | None, *, default):
return default
monkeypatch.setattr(
block_module.LLMAPIHandlerFactory,
"get_override_llm_api_handler",
fake_get_override_handler,
raising=False,
)
monkeypatch.setattr(
TextPromptBlock,
"_resolve_default_llm_handler",
fake_resolve_default_llm_handler,
raising=False,
)
# Real render #1 (sandboxed): the parameter value is substituted as inert text.
block.format_potential_template_parameters(ctx)
assert block.prompt == payload
# Real send path, no monkeypatching of the render sink: the literal must survive.
await block.send_prompt(
block.prompt,
workflow_run_id="wr_test",
organization_id="org-1",
workflow_run_block_id=None,
)
sent_prompt = captured["prompt"]
assert payload in sent_prompt
assert forbidden not in sent_prompt

View file

@ -281,7 +281,6 @@ def test_text_prompt_block_forwards_system_prompt(monkeypatch) -> None:
asyncio.run(
block.send_prompt(
prompt="What is the meaning of life?",
parameter_values={},
workflow_run_id="wfr_sp",
organization_id="o_sp",
workflow_run_block_id=None,