diff --git a/src/specify_cli/integrations/junie/__init__.py b/src/specify_cli/integrations/junie/__init__.py index 2d4a6b32d9..9e243f42b8 100644 --- a/src/specify_cli/integrations/junie/__init__.py +++ b/src/specify_cli/integrations/junie/__init__.py @@ -132,6 +132,61 @@ def _rewrite_handoff_references(content: str) -> str: lambda m: f"{m.group(1)}{format_junie_command_name(m.group(2))}", content, ) + + @staticmethod + def _inject_allow_prompt_argument(content: str, allow_prompt: bool = True) -> str: + """Inject allowPromptArgument: true/false into the YAML frontmatter. + + If frontmatter exists, it ensures the key is set to the desired value (overwriting if needed). + If not, it creates a minimal frontmatter. + """ + value = "true" if allow_prompt else "false" + if not content.startswith("---"): + # No frontmatter at all? Create one. + return f"---\nallowPromptArgument: {value}\n---\n\n" + content + + parts = re.split(r"(?m)^---\s*$", content, maxsplit=2) + if len(parts) < 3: + # Malformed frontmatter (e.g. missing closing dashes)? + return content + + frontmatter = parts[1] + body = parts[2] + + if "allowPromptArgument:" in frontmatter: + # Overwrite existing key + frontmatter = re.sub( + r"(?m)^(allowPromptArgument:\s*).*", + fr"\1{value}", + frontmatter + ) + else: + # Append to frontmatter. Ensure it ends with newline. + # Check if the last line of frontmatter is a key-value pair. + lines = frontmatter.splitlines() + if lines and not lines[-1].strip(): + # Remove trailing empty lines in frontmatter + while lines and not lines[-1].strip(): + lines.pop() + frontmatter = "\n".join(lines) + "\n" + elif not frontmatter.endswith("\n"): + frontmatter += "\n" + + frontmatter += f"allowPromptArgument: {value}\n" + + return f"---{frontmatter}---{body}" + + @staticmethod + def _transform_body_variables(content: str) -> str: + """Transform $ARGUMENTS to $prompt and escape other $word by doubling $.""" + def replacer(match: re.Match[str]) -> str: + word = match.group(1) + if word == "ARGUMENTS": + return "$prompt" + return "$$" + word + + return re.sub(r"(? str: """Apply Junie-specific transformations to command content. @@ -140,8 +195,36 @@ def post_process_command_content(self, content: str) -> str: ``post_process_command_content``) applies these transforms to extension/preset command files too, not just core commands. """ + # If it has frontmatter, we must isolate it to avoid transforming variables + # inside the YAML header (e.g. key: $VAL should not become key: $$VAL). + if content.startswith("---"): + parts = re.split(r"(?m)^---\s*$", content, maxsplit=2) + if len(parts) < 3: + # Malformed frontmatter (missing closing dashes) - return as is. + return content + + frontmatter_block = f"---{parts[1]}---" + body = parts[2] + + has_arguments = bool(re.search(r"(? true + content = "---\nallowPromptArgument: false\ndescription: Test\n---\n$ARGUMENTS" + updated = junie.post_process_command_content(content) + assert updated.count("allowPromptArgument:") == 1 + assert "allowPromptArgument: true" in updated + + # Without $ARGUMENTS -> false + content = "---\nallowPromptArgument: true\ndescription: Test\n---\nBody" + updated = junie.post_process_command_content(content) + assert updated.count("allowPromptArgument:") == 1 + assert "allowPromptArgument: false" in updated + + def test_junie_variable_escaping_regex(self): + """Verify generic variable escaping using regex. (US2)""" + junie = get_integration("junie") + content = "Hook: $speckit-git-branch, User: $user_name, Prompt: $prompt" + # Note: $prompt should also be escaped to $$prompt to isolate Junie's reserved token + updated = junie.post_process_command_content(content) + assert "$$speckit-git-branch" in updated + assert "$$user_name" in updated + assert "$$prompt" in updated + assert "$prompt" in updated # Still matches because $$prompt contains $prompt + # More specific check: + assert "Prompt: $$prompt" in updated + + def test_junie_token_isolation(self): + """Verify $ARGUMENTS is isolated from substrings like $ARGUMENTS_SUFFIX.""" + junie = get_integration("junie") + content = "Use $ARGUMENTS but not $ARGUMENTS_SUFFIX" + updated = junie.post_process_command_content(content) + assert "allowPromptArgument: true" in updated + assert "$prompt" in updated + assert "$$ARGUMENTS_SUFFIX" in updated + assert "$prompt_SUFFIX" not in updated + + content = "Only $ARGUMENTS_SUFFIX" + updated = junie.post_process_command_content(content) + assert "allowPromptArgument: false" in updated + assert "$$ARGUMENTS_SUFFIX" in updated + assert "$prompt" not in updated + + def test_junie_missing_frontmatter(self): + """Verify frontmatter is created if missing.""" + junie = get_integration("junie") + # With $ARGUMENTS + content = "$ARGUMENTS" + updated = junie.post_process_command_content(content) + assert updated.startswith("---\nallowPromptArgument: true\n---\n\n") + assert "$prompt" in updated + + # Without $ARGUMENTS + content = "Body" + updated = junie.post_process_command_content(content) + assert updated.startswith("---\nallowPromptArgument: false\n---\n\n") + + def test_junie_malformed_frontmatter(self): + """Verify malformed frontmatter is handled gracefully (returned as is).""" + junie = get_integration("junie") + content = "---\nMalformed frontmatter (missing closing dashes)\nBody" + updated = junie.post_process_command_content(content) + assert updated == content + + def test_junie_frontmatter_non_standalone_delimiter(self): + """Verify --- not on standalone line is NOT identified as delimiter.""" + junie = get_integration("junie") + content = "---\ntitle: Foo\n--- not a delimiter\nbody" + updated = junie.post_process_command_content(content) + assert updated == content + + def test_junie_frontmatter_with_horizontal_rule_in_body(self): + """Verify valid frontmatter is preserved and horizontal rule in body is untouched.""" + junie = get_integration("junie") + content = "---\ntitle: Foo\n---\nBody with\n---\nhorizontal rule" + updated = junie.post_process_command_content(content) + assert "---" in updated + assert "horizontal rule" in updated + assert updated.count("---") == 3 + + def test_junie_nested_yaml_protection(self): + """Verify regex does not match nested allowPromptArgument.""" + junie = get_integration("junie") + content = "---\nconfig:\n allowPromptArgument: nested\n---\nBody" + updated = junie.post_process_command_content(content) + assert "allowPromptArgument: nested" in updated + assert "allowPromptArgument: false" not in updated + # -- Overrides for MarkdownIntegrationTests --------------------------- def test_setup_creates_files(self, tmp_path): diff --git a/tests/integrations/test_junie_token_mangling.py b/tests/integrations/test_junie_token_mangling.py new file mode 100644 index 0000000000..8cd067764b --- /dev/null +++ b/tests/integrations/test_junie_token_mangling.py @@ -0,0 +1,105 @@ +"""Reproduction tests for Junie token mangling bug.""" +from specify_cli.integrations import get_integration + +def test_junie_token_mangling_reproduction(): + """Reproduce the token mangling and false positive detection issues.""" + junie = get_integration("junie") + + # Case 1: $ARGUMENTS should be exactly matched and replaced by $prompt + content = "Run with $ARGUMENTS" + updated = junie.post_process_command_content(content) + assert "$prompt" in updated + assert "allowPromptArgument: true" in updated + + # Case 2: $ARGUMENTS_SUFFIX should NOT be replaced by $prompt_SUFFIX + # Currently it is mangled. + content = "Variable: $ARGUMENTS_SUFFIX" + updated = junie.post_process_command_content(content) + # EXPECTED: $$ARGUMENTS_SUFFIX (escaped) + assert "$$ARGUMENTS_SUFFIX" in updated + assert "$prompt" not in updated + + # Case 3: $prompt (original) should be escaped to $$prompt + # Currently it is NOT escaped. + content = "Reserved: $prompt" + updated = junie.post_process_command_content(content) + # EXPECTED: $$prompt + assert "$$prompt" in updated + + # Case 4: has_arguments false positive + # Currently '$ARGUMENTS' in content returns True for '$ARGUMENTS_SUFFIX' + content = "Only $ARGUMENTS_SUFFIX here" + updated = junie.post_process_command_content(content) + # EXPECTED: allowPromptArgument: false + assert "allowPromptArgument: false" in updated + +def test_junie_escaped_literals_preservation(): + """Verify that already escaped literals are preserved and not double-escaped.""" + junie = get_integration("junie") + + # $$HOME should stay $$HOME + content = "Literal: $$HOME" + updated = junie.post_process_command_content(content) + assert "$$HOME" in updated + assert "$$$HOME" not in updated + + # $$ARGUMENTS should stay $$ARGUMENTS and NOT trigger allowPromptArgument + content = "Help: $$ARGUMENTS literal" + updated = junie.post_process_command_content(content) + assert "$$ARGUMENTS" in updated + assert "$$$ARGUMENTS" not in updated + assert "allowPromptArgument: false" in updated + +def test_junie_frontmatter_preservation(): + """Verify that variables in YAML frontmatter are not escaped.""" + junie = get_integration("junie") + + content = "---\nkey: $VAL\n---\nbody with $HOME" + updated = junie.post_process_command_content(content) + + # Frontmatter should be preserved + assert "key: $VAL" in updated + assert "key: $$VAL" not in updated + + # Body should still be transformed + assert "body with $$HOME" in updated + +def test_junie_arguments_scope(): + """Verify that $ARGUMENTS in frontmatter does not trigger allowPromptArgument.""" + junie = get_integration("junie") + + # Case 1: $ARGUMENTS in frontmatter only + content = "---\ndescription: Use $ARGUMENTS\n---\nbody" + updated = junie.post_process_command_content(content) + assert "allowPromptArgument: false" in updated + + # Case 2: $ARGUMENTS in body only + content = "---\ndescription: none\n---\nRun $ARGUMENTS" + updated = junie.post_process_command_content(content) + assert "allowPromptArgument: true" in updated + +def test_transform_body_variables_direct(): + """Directly test _transform_body_variables with various tokens.""" + junie = get_integration("junie") + + # Desired behavior: + assert junie._transform_body_variables("$ARGUMENTS") == "$prompt" + assert junie._transform_body_variables("$ARGUMENTS_SUFFIX") == "$$ARGUMENTS_SUFFIX" + assert junie._transform_body_variables("$prompt") == "$$prompt" + assert junie._transform_body_variables("$foo") == "$$foo" + +def test_malformed_frontmatter_short_circuit(): + """Verify that malformed frontmatter short-circuits the pipeline.""" + junie = get_integration("junie") + + # Input starts with --- but has no closing --- + content = "---\n$ARGUMENTS" + updated = junie.post_process_command_content(content) + + # It should be returned unchanged + assert updated == content + + # Case with spaces after dashes + content = "--- \n$ARGUMENTS" + updated = junie.post_process_command_content(content) + assert updated == content