From ff5a74dac683d0fd70000eb63eb37cb56886f01c Mon Sep 17 00:00:00 2001 From: pedrohsdb Date: Fri, 13 Feb 2026 12:16:45 -0800 Subject: [PATCH] Fix PII leakage in cached run_task() prompts (#SKY-7989) (#4744) --- .../script_generations/generate_script.py | 85 +++++-- .../test_run_task_prompt_parameterization.py | 226 ++++++++++++++++++ 2 files changed, 290 insertions(+), 21 deletions(-) create mode 100644 tests/unit/test_run_task_prompt_parameterization.py diff --git a/skyvern/core/script_generations/generate_script.py b/skyvern/core/script_generations/generate_script.py index 0586255a..1604f62a 100644 --- a/skyvern/core/script_generations/generate_script.py +++ b/skyvern/core/script_generations/generate_script.py @@ -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: diff --git a/tests/unit/test_run_task_prompt_parameterization.py b/tests/unit/test_run_task_prompt_parameterization.py new file mode 100644 index 00000000..ba0fdb6f --- /dev/null +++ b/tests/unit/test_run_task_prompt_parameterization.py @@ -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, "", "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, "", "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, "", "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}", "", "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