From 94c7ec288fd34221e62fef50f3212f8a899eda95 Mon Sep 17 00:00:00 2001 From: Quratulain-bilal Date: Wed, 8 Jul 2026 17:42:34 +0500 Subject: [PATCH] fix(toml): escape control characters so generated command files parse (#3341) * fix(toml): escape control characters so generated command files parse both toml renderers (TomlIntegration._render_toml_string for gemini/tabnine and CommandRegistrar.render_toml_command for extension/preset commands) wrote control characters raw into a multiline or basic string. toml forbids literal control chars (U+0000-U+001F except tab/newline, and U+007F) in every string form, and a bare CR that is not part of a CRLF pair, so a description or body containing one produced a .toml file that fails to parse. route any value with such a character to a fully-escaped basic string that emits the leftover control chars as \uXXXX. added regression tests that round-trip through tomllib. * refactor(toml): centralize control-char escaping in one shared helper the control-char detection and basic-string escaping added for both toml renderers were copy-pasted into agents.py and integrations/base.py. move the two functions into specify_cli/_toml_string.py and have both renderers delegate to it, so the escaping rules can't drift apart later. no behavior change; both renderers now reference the same implementation. --- src/specify_cli/_toml_string.py | 60 +++++++++++++++++++ src/specify_cli/agents.py | 25 ++++---- src/specify_cli/integrations/base.py | 27 +++++---- .../test_integration_base_toml.py | 12 ++++ tests/test_extensions.py | 19 ++++++ 5 files changed, 120 insertions(+), 23 deletions(-) create mode 100644 src/specify_cli/_toml_string.py diff --git a/src/specify_cli/_toml_string.py b/src/specify_cli/_toml_string.py new file mode 100644 index 000000000..c9b024245 --- /dev/null +++ b/src/specify_cli/_toml_string.py @@ -0,0 +1,60 @@ +"""Shared TOML string-escaping helpers. + +Both TOML command renderers — ``TomlIntegration`` (gemini, tabnine) in +``specify_cli.integrations.base`` and ``CommandRegistrar.render_toml_command`` +(extension/preset commands) in ``specify_cli.agents`` — need the same rules for +detecting characters TOML forbids literally and for emitting a fully-escaped +basic string. Keeping one implementation here avoids the two drifting apart if +the escaping rules change again. +""" +from __future__ import annotations + + +def has_illegal_toml_control(value: str) -> bool: + """True when *value* contains a character TOML forbids literally. + + TOML basic/literal strings (single- or multi-line) allow tab and, in the + multiline forms, newlines — but every other control character + (``U+0000``–``U+001F`` and ``U+007F``) must be ``\\u``-escaped, which only a + basic string can do. A bare carriage return counts too: a multiline basic + string treats ``\\r`` as a newline only when paired into ``\\r\\n``; a lone + ``\\r`` is an illegal control character. + """ + length = len(value) + for i, ch in enumerate(value): + code = ord(ch) + if ch == "\r": + # Only a CR that is part of a CRLF newline is allowed literally. + if i + 1 < length and value[i + 1] == "\n": + continue + return True + if (code < 0x20 and ch not in ("\t", "\n")) or code == 0x7F: + return True + return False + + +def escape_toml_basic(value: str) -> str: + """Render *value* as a single-line basic string, escaping everything. + + Always valid TOML: backslash/quote are escaped, the common control chars + use their short escapes, and any remaining control character is emitted as + a ``\\uXXXX`` sequence. + """ + out: list[str] = [] + for ch in value: + code = ord(ch) + if ch == "\\": + out.append("\\\\") + elif ch == '"': + out.append('\\"') + elif ch == "\n": + out.append("\\n") + elif ch == "\r": + out.append("\\r") + elif ch == "\t": + out.append("\\t") + elif code < 0x20 or code == 0x7F: + out.append(f"\\u{code:04x}") + else: + out.append(ch) + return '"' + "".join(out) + '"' diff --git a/src/specify_cli/agents.py b/src/specify_cli/agents.py index c127407d4..391015394 100644 --- a/src/specify_cli/agents.py +++ b/src/specify_cli/agents.py @@ -16,6 +16,8 @@ from typing import Any, Dict, List, Optional import yaml from ._init_options import is_ai_skills_enabled, load_init_options +from ._toml_string import escape_toml_basic as _escape_toml_basic +from ._toml_string import has_illegal_toml_control as _has_illegal_toml_control from ._utils import relative_extension_path_violation @@ -258,7 +260,12 @@ class CommandRegistrar: # ``C:\\Users\\...`` whose ``\\U`` reads as an invalid unicode escape) would # produce unparseable TOML — route those to the *literal* form ('''...'''), # which does not process escapes, or to the escaped basic string. - if '"""' not in body and "\\" not in body: + # Control characters (U+0000–U+001F except tab/newline, U+007F) and a bare + # CR are illegal in every TOML string form, so a body containing them must + # go to the escaped basic string regardless of which delimiters it uses. + if self._has_illegal_toml_control(body): + toml_lines.append(f"prompt = {self._render_basic_toml_string(body)}") + elif '"""' not in body and "\\" not in body: toml_lines.append('prompt = """') toml_lines.append(body) toml_lines.append('"""') @@ -271,17 +278,11 @@ class CommandRegistrar: return "\n".join(toml_lines) - @staticmethod - def _render_basic_toml_string(value: str) -> str: - """Render *value* as a TOML basic string literal.""" - escaped = ( - value.replace("\\", "\\\\") - .replace('"', '\\"') - .replace("\n", "\\n") - .replace("\r", "\\r") - .replace("\t", "\\t") - ) - return f'"{escaped}"' + # Control-char detection and basic-string escaping are shared with the + # gemini/tabnine renderer in ``specify_cli.integrations.base`` via + # ``specify_cli._toml_string`` so the two never drift apart. + _has_illegal_toml_control = staticmethod(_has_illegal_toml_control) + _render_basic_toml_string = staticmethod(_escape_toml_basic) def render_yaml_command( self, diff --git a/src/specify_cli/integrations/base.py b/src/specify_cli/integrations/base.py index 706a3cb5d..4eb9bb19a 100644 --- a/src/specify_cli/integrations/base.py +++ b/src/specify_cli/integrations/base.py @@ -25,6 +25,9 @@ from typing import TYPE_CHECKING, Any import yaml +from .._toml_string import escape_toml_basic as _escape_toml_basic +from .._toml_string import has_illegal_toml_control as _has_illegal_toml_control + if TYPE_CHECKING: from .manifest import IntegrationManifest @@ -953,6 +956,12 @@ class TomlIntegration(IntegrationBase): body = "".join(lines[frontmatter_end + 1 :]) return frontmatter, body + # Control-char detection and basic-string escaping are shared with the + # extension/preset renderer in ``specify_cli.agents`` via + # ``specify_cli._toml_string`` so the two never drift apart. + _has_illegal_toml_control = staticmethod(_has_illegal_toml_control) + _escape_toml_basic = staticmethod(_escape_toml_basic) + @staticmethod def _render_toml_string(value: str) -> str: """Render *value* as a TOML string literal. @@ -962,6 +971,12 @@ class TomlIntegration(IntegrationBase): literal string or escaped basic string when delimiters appear in the content. """ + # Control characters other than tab/newline (and a bare CR) cannot + # appear literally in any TOML string; route them to a fully-escaped + # basic string so the generated file stays parseable. + if TomlIntegration._has_illegal_toml_control(value): + return TomlIntegration._escape_toml_basic(value) + if "\n" not in value and "\r" not in value: escaped = value.replace("\\", "\\\\").replace('"', '\\"') return f'"{escaped}"' @@ -974,17 +989,7 @@ class TomlIntegration(IntegrationBase): if "'''" not in value and not value.endswith("'"): return "'''\n" + value + "'''" - return ( - '"' - + ( - value.replace("\\", "\\\\") - .replace('"', '\\"') - .replace("\n", "\\n") - .replace("\r", "\\r") - .replace("\t", "\\t") - ) - + '"' - ) + return TomlIntegration._escape_toml_basic(value) @staticmethod def _render_toml(description: str, body: str) -> str: diff --git a/tests/integrations/test_integration_base_toml.py b/tests/integrations/test_integration_base_toml.py index 68f5fd075..37bad8a60 100644 --- a/tests/integrations/test_integration_base_toml.py +++ b/tests/integrations/test_integration_base_toml.py @@ -291,6 +291,18 @@ class TomlIntegrationTests: "closing delimiter should be inline when body does not end with a quote" ) + def test_toml_string_escapes_control_characters(self): + """A value with control chars / a bare CR must render as parseable TOML. + + TOML forbids literal control characters (U+0000–U+001F except tab and + newline, plus U+007F) in every string form, and a bare CR that is not + part of a CRLF pair. The renderer used to emit these raw into a basic or + ``\"\"\"`` multiline string, producing a config file that fails to parse.""" + value = "start\x00null\x01ctrl\x1besc\x7fdel\rlone-cr end" + rendered = TomlIntegration._render_toml_string(value) + parsed = tomllib.loads(f"prompt = {rendered}") + assert parsed["prompt"] == value + def test_toml_is_valid(self, tmp_path): """Every generated TOML file must parse without errors.""" i = get_integration(self.KEY) diff --git a/tests/test_extensions.py b/tests/test_extensions.py index 3118ccdef..288518036 100644 --- a/tests/test_extensions.py +++ b/tests/test_extensions.py @@ -1795,6 +1795,25 @@ $ARGUMENTS assert parsed["description"] == "first line\nsecond line\n" + def test_render_toml_command_escapes_control_characters(self): + """Control characters and a lone CR must be escaped so the TOML parses. + + TOML forbids literal control characters (U+0000–U+001F except tab and + newline, plus U+007F) in any string, and treats a bare CR outside a + CRLF pair as illegal. The renderer used to emit these raw — into a + basic string (single-line) or a ``\"\"\"`` multiline string (for a lone + CR) — producing a command file that fails to parse.""" + from specify_cli.agents import CommandRegistrar as AgentCommandRegistrar + + registrar = AgentCommandRegistrar() + body = "start\x00null\x01ctrl\x1besc\x7fdel\rlone-cr end" + output = registrar.render_toml_command( + {"description": "d"}, body, "extension:test-ext" + ) + + parsed = tomllib.loads(output) + assert parsed["prompt"] == body + def test_render_toml_command_preserves_backslashes_in_body(self): """A backslash in the body (e.g. a Windows path) must not break TOML.