vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fix PII leakage in cached run_task() prompts (#SKY-7989) (#4744)

Browse Source
This commit is contained in:
pedrohsdb
2026-02-13 12:16:45 -08:00
committed by GitHub
parent 53a89c3e44
commit ff5a74dac6
2 changed files with 290 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

85
skyvern/core/script_generations/generate_script.py
View File

@@ -229,6 +229,16 @@ def _build_value_to_param_lookup(
return dict(sorted(raw.items(), key=lambda kv: len(kv[0]), reverse=True))
def _escape_for_fstring_text(text: str) -> str:
"""Escape ``{`` and ``}`` so they survive inside a ``FormattedStringText`` node.
Jinja2 templates (e.g. ``{{param}}``) would otherwise be interpreted as
f-string expressions. Doubling the braces turns them into literal braces
in the rendered f-string.
"""
return text.replace("{", "{{").replace("}", "}}")
def _build_parameterized_prompt_cst(
intention: str,
value_to_param: dict[str, str],
@@ -268,7 +278,7 @@ def _build_parameterized_prompt_cst(
for start, end, field_name in matches:
# Text segment before this match.
if start > cursor:
parts.append(cst.FormattedStringText(intention[cursor:start]))
parts.append(cst.FormattedStringText(_escape_for_fstring_text(intention[cursor:start])))
# The {context.parameters['field_name']} expression.
parts.append(
cst.FormattedStringExpression(
@@ -289,9 +299,17 @@ def _build_parameterized_prompt_cst(
# Trailing text after last match.
if cursor < len(intention):
parts.append(cst.FormattedStringText(intention[cursor:]))
parts.append(cst.FormattedStringText(_escape_for_fstring_text(intention[cursor:])))
return cst.FormattedString(parts=parts)
# Use triple-quote f-string when the content contains newlines or quotes
# (run_task prompts always have newlines from the appended navigation_payload).
raw_text = intention
if "\n" in raw_text or '"' in raw_text or "'" in raw_text:
quote = '"""'
else:
quote = '"'
return cst.FormattedString(parts=parts, start=f"f{quote}", end=quote)
def _requires_mini_agent(act: dict[str, Any]) -> bool:
@@ -1085,10 +1103,13 @@ def _build_generated_model_from_schema(schema_code: str) -> cst.ClassDef | None:
def _build_run_task_statement(
block_title: str, block: dict[str, Any], data_variable_name: str | None = None
block_title: str,
block: dict[str, Any],
data_variable_name: str | None = None,
value_to_param: dict[str, str] | None = None,
) -> cst.SimpleStatementLine:
"""Build a skyvern.run_task statement."""
args = __build_base_task_statement(block_title, block, data_variable_name)
args = __build_base_task_statement(block_title, block, data_variable_name, value_to_param=value_to_param)
call = cst.Call(
func=cst.Attribute(value=cst.Name("skyvern"), attr=cst.Name("run_task")),
args=args,
@@ -1102,10 +1123,13 @@ def _build_run_task_statement(
def _build_download_statement(
block_title: str, block: dict[str, Any], data_variable_name: str | None = None
block_title: str,
block: dict[str, Any],
data_variable_name: str | None = None,
value_to_param: dict[str, str] | None = None,
) -> cst.SimpleStatementLine:
"""Build a skyvern.download statement."""
args = __build_base_task_statement(block_title, block, data_variable_name)
args = __build_base_task_statement(block_title, block, data_variable_name, value_to_param=value_to_param)
call = cst.Call(
func=cst.Attribute(value=cst.Name("skyvern"), attr=cst.Name("download")),
args=args,
@@ -1168,10 +1192,13 @@ def _build_action_statement(
def _build_login_statement(
block_title: str, block: dict[str, Any], data_variable_name: str | None = None
block_title: str,
block: dict[str, Any],
data_variable_name: str | None = None,
value_to_param: dict[str, str] | None = None,
) -> cst.SimpleStatementLine:
"""Build a skyvern.login statement."""
args = __build_base_task_statement(block_title, block, data_variable_name)
args = __build_base_task_statement(block_title, block, data_variable_name, value_to_param=value_to_param)
call = cst.Call(
func=cst.Attribute(value=cst.Name("skyvern"), attr=cst.Name("login")),
args=args,
@@ -1256,10 +1283,13 @@ def _build_extract_statement(
def _build_navigate_statement(
block_title: str, block: dict[str, Any], data_variable_name: str | None = None
block_title: str,
block: dict[str, Any],
data_variable_name: str | None = None,
value_to_param: dict[str, str] | None = None,
) -> cst.SimpleStatementLine:
"""Build a skyvern.navigate statement."""
args = __build_base_task_statement(block_title, block, data_variable_name)
args = __build_base_task_statement(block_title, block, data_variable_name, value_to_param=value_to_param)
call = cst.Call(
func=cst.Attribute(value=cst.Name("skyvern"), attr=cst.Name("run_task")),
args=args,
@@ -1995,7 +2025,10 @@ def _mark_last_arg_as_comma(args: list[cst.Arg]) -> None:
def __build_base_task_statement(
block_title: str, block: dict[str, Any], data_variable_name: str | None = None
block_title: str,
block: dict[str, Any],
data_variable_name: str | None = None,
value_to_param: dict[str, str] | None = None,
) -> list[cst.Arg]:
block_type = block.get("block_type")
prompt = block.get("prompt") if block_type == "task_v2" else block.get("navigation_goal")
@@ -2011,10 +2044,17 @@ def __build_base_task_statement(
prompt = prompt or ""
prompt = f"{prompt}\n{navigation_payload}"
# Try to parameterize PII in the prompt with f-string context.parameters refs.
prompt_value: cst.BaseExpression | None = None
if value_to_param and prompt:
prompt_value = _build_parameterized_prompt_cst(prompt, value_to_param)
if prompt_value is None:
prompt_value = _value(prompt)
args = [
cst.Arg(
keyword=cst.Name("prompt"),
value=_value(prompt),
value=prompt_value,
whitespace_after_arg=cst.ParenthesizedWhitespace(
indent=True,
last_line=cst.SimpleWhitespace(INDENT),
@@ -2132,7 +2172,10 @@ def __build_base_task_statement(
def _build_block_statement(
block: dict[str, Any], data_variable_name: str | None = None, assign_output: bool = False
block: dict[str, Any],
data_variable_name: str | None = None,
assign_output: bool = False,
value_to_param: dict[str, str] | None = None,
) -> cst.SimpleStatementLine:
"""Build a block statement."""
block_type = block.get("block_type")
@@ -2141,23 +2184,23 @@ def _build_block_statement(
if block_type in SCRIPT_TASK_BLOCKS:
# For task blocks, call the custom function with cache_key
if block_type == "task":
stmt = _build_run_task_statement(block_title, block, data_variable_name)
stmt = _build_run_task_statement(block_title, block, data_variable_name, value_to_param=value_to_param)
elif block_type == "file_download":
stmt = _build_download_statement(block_title, block, data_variable_name)
stmt = _build_download_statement(block_title, block, data_variable_name, value_to_param=value_to_param)
elif block_type == "action":
stmt = _build_action_statement(block_title, block, data_variable_name)
elif block_type == "login":
stmt = _build_login_statement(block_title, block, data_variable_name)
stmt = _build_login_statement(block_title, block, data_variable_name, value_to_param=value_to_param)
elif block_type == "extraction":
stmt = _build_extract_statement(block_title, block, data_variable_name, assign_output)
elif block_type == "navigation":
stmt = _build_navigate_statement(block_title, block, data_variable_name)
stmt = _build_navigate_statement(block_title, block, data_variable_name, value_to_param=value_to_param)
elif block_type == "validation":
stmt = _build_validate_statement(block_title, block, data_variable_name)
elif block_type == "human_interaction":
stmt = _build_human_interaction_statement(block)
elif block_type == "task_v2":
stmt = _build_run_task_statement(block_title, block, data_variable_name)
stmt = _build_run_task_statement(block_title, block, data_variable_name, value_to_param=value_to_param)
elif block_type == "send_email":
stmt = _build_send_email_statement(block)
elif block_type == "text_prompt":
@@ -2398,7 +2441,7 @@ async def generate_workflow_script_python_code(
temp_module = cst.Module(body=[block_fn_def])
block_code = temp_module.code
block_stmt = _build_block_statement(task)
block_stmt = _build_block_statement(task, value_to_param=value_to_param)
run_signature_module = cst.Module(body=[block_stmt])
run_signature = run_signature_module.code.strip()
@@ -2464,7 +2507,7 @@ async def generate_workflow_script_python_code(
temp_module = cst.Module(body=task_v2_block_body)
block_code = temp_module.code
task_v2_stmt = _build_block_statement(task_v2)
task_v2_stmt = _build_block_statement(task_v2, value_to_param=value_to_param)
run_signature = cst.Module(body=[task_v2_stmt]).code.strip()
if script_id and script_revision_id and organization_id:

226
tests/unit/test_run_task_prompt_parameterization.py Normal file
View File

@@ -0,0 +1,226 @@
"""
Tests for run_task() prompt parameterization in cached script generation.
When generating cached scripts, block-level prompts (run_task, navigate, etc.) should
replace literal parameter values with f-string references to context.parameters[...],
so that re-runs with different input values don't leak PII from prior runs.
"""
from typing import Any
import libcst as cst
from skyvern.core.script_generations.generate_script import (
_build_parameterized_prompt_cst,
_build_run_task_statement,
_escape_for_fstring_text,
)
# ---------------------------------------------------------------------------
# _escape_for_fstring_text
# ---------------------------------------------------------------------------
class TestEscapeForFstringText:
def test_escapes_single_braces(self) -> None:
assert _escape_for_fstring_text("{hello}") == "{{hello}}"
def test_escapes_jinja2_double_braces(self) -> None:
# Jinja2 {{param}} → {{{{param}}}} in an f-string text node
assert _escape_for_fstring_text("{{param}}") == "{{{{param}}}}"
def test_leaves_plain_text_unchanged(self) -> None:
assert _escape_for_fstring_text("no braces here") == "no braces here"
def test_mixed_content(self) -> None:
text = "Name: {first} and {{last}}"
assert _escape_for_fstring_text(text) == "Name: {{first}} and {{{{last}}}}"
def test_empty_string(self) -> None:
assert _escape_for_fstring_text("") == ""
# ---------------------------------------------------------------------------
# _build_parameterized_prompt_cst — brace escaping
# ---------------------------------------------------------------------------
class TestBuildParameterizedPromptCstBraces:
"""Tests that Jinja2 templates / braces in the prompt text survive parameterization."""
def test_jinja2_templates_preserved(self) -> None:
"""Prompt with {{param}} Jinja2 template + PII should produce a valid f-string."""
prompt = "Fill in First Name: MASOOD for {{firstName}}"
value_to_param = {"MASOOD": "firstName"}
result = _build_parameterized_prompt_cst(prompt, value_to_param)
assert result is not None
code = cst.Module(body=[]).code_for_node(result)
# PII should be replaced
assert "MASOOD" not in code
assert "context.parameters" in code
# The Jinja2 template should survive as literal braces
assert "firstName" in code
# Should be compilable Python
compile(code, "<test>", "eval")
def test_braces_in_non_match_segments(self) -> None:
"""Braces in text segments that aren't PII should be escaped."""
prompt = "Look for {item} with ID 542-641-668"
value_to_param = {"542-641-668": "patient_id"}
result = _build_parameterized_prompt_cst(prompt, value_to_param)
assert result is not None
code = cst.Module(body=[]).code_for_node(result)
assert "542-641-668" not in code
assert "patient_id" in code
# Should be compilable (braces are escaped)
compile(code, "<test>", "eval")
# ---------------------------------------------------------------------------
# _build_parameterized_prompt_cst — triple-quote support
# ---------------------------------------------------------------------------
class TestBuildParameterizedPromptCstTripleQuote:
"""Tests that multiline prompts use triple-quote f-strings."""
def test_multiline_prompt_uses_triple_quote(self) -> None:
prompt = "Fill in the form:\nFirst Name: MASOOD\nLast Name: SABIR"
value_to_param = {"MASOOD": "firstName", "SABIR": "lastName"}
result = _build_parameterized_prompt_cst(prompt, value_to_param)
assert result is not None
code = cst.Module(body=[]).code_for_node(result)
assert code.startswith('f"""')
assert code.endswith('"""')
assert "MASOOD" not in code
assert "SABIR" not in code
compile(code, "<test>", "eval")
def test_single_line_prompt_uses_regular_quote(self) -> None:
prompt = "Find patient MASOOD SABIR"
value_to_param = {"MASOOD SABIR": "fullName"}
result = _build_parameterized_prompt_cst(prompt, value_to_param)
assert result is not None
code = cst.Module(body=[]).code_for_node(result)
assert code.startswith('f"')
assert not code.startswith('f"""')
def test_prompt_with_quotes_uses_triple_quote(self) -> None:
prompt = """Click the "Submit" button for MASOOD"""
value_to_param = {"MASOOD": "firstName"}
result = _build_parameterized_prompt_cst(prompt, value_to_param)
assert result is not None
code = cst.Module(body=[]).code_for_node(result)
assert code.startswith('f"""')
# ---------------------------------------------------------------------------
# End-to-end: _build_run_task_statement with value_to_param
# ---------------------------------------------------------------------------
def _make_block(
navigation_goal: str,
block_type: str = "task",
parameters: list[dict[str, str]] | None = None,
label: str = "test_block",
url: str = "",
) -> dict[str, Any]:
block: dict[str, Any] = {
"block_type": block_type,
"navigation_goal": navigation_goal,
"label": label,
}
if parameters:
block["parameters"] = parameters
if url:
block["url"] = url
return block
def _render_stmt(stmt: cst.SimpleStatementLine) -> str:
return cst.Module(body=[stmt]).code
class TestBuildRunTaskStatementParameterization:
"""End-to-end tests: _build_run_task_statement with value_to_param."""
def test_pii_in_navigation_goal_is_parameterized(self) -> None:
"""PII in Steps section should be replaced with context.parameters refs."""
block = _make_block(
navigation_goal=("Navigate to the form.\nSteps:\nFill in (First Name: MASOOD, Last Name: SABIR)"),
)
value_to_param = {"MASOOD": "firstName", "SABIR": "lastName"}
stmt = _build_run_task_statement("test_block", block, value_to_param=value_to_param)
code = _render_stmt(stmt)
assert "MASOOD" not in code
assert "SABIR" not in code
assert "context.parameters" in code
assert "firstName" in code
assert "lastName" in code
def test_jinja2_templates_in_navigation_payload_preserved(self) -> None:
"""Jinja2 {{param}} templates appended by navigation_payload should survive."""
block = _make_block(
navigation_goal="Fill in First Name: MASOOD",
parameters=[{"key": "firstName"}, {"key": "lastName"}],
)
value_to_param = {"MASOOD": "firstName"}
stmt = _build_run_task_statement("test_block", block, value_to_param=value_to_param)
code = _render_stmt(stmt)
# PII replaced
assert "MASOOD" not in code
assert "context.parameters" in code
# The code should be syntactically valid Python (wrap in async def since it has await)
compile(f"async def _test():\n {code}", "<test>", "exec")
def test_no_parameterization_without_value_to_param(self) -> None:
"""Without value_to_param, prompt should be a plain string literal."""
block = _make_block(navigation_goal="Fill in First Name: MASOOD")
stmt = _build_run_task_statement("test_block", block, value_to_param=None)
code = _render_stmt(stmt)
assert "MASOOD" in code
assert "context.parameters" not in code
def test_no_parameterization_when_no_match(self) -> None:
"""When value_to_param doesn't match anything, prompt is a plain literal."""
block = _make_block(navigation_goal="Navigate to the dashboard")
value_to_param = {"MASOOD": "firstName"}
stmt = _build_run_task_statement("test_block", block, value_to_param=value_to_param)
code = _render_stmt(stmt)
assert "Navigate to the dashboard" in code
assert "context.parameters" not in code
def test_task_v2_block_parameterized(self) -> None:
"""task_v2 blocks use 'prompt' key instead of 'navigation_goal'."""
block: dict[str, Any] = {
"block_type": "task_v2",
"prompt": "Fill in the form for MASOOD SABIR",
"label": "test_v2",
}
value_to_param = {"MASOOD SABIR": "fullName"}
stmt = _build_run_task_statement("test_v2", block, value_to_param=value_to_param)
code = _render_stmt(stmt)
assert "MASOOD SABIR" not in code
assert "context.parameters" in code
assert "fullName" in code