Escape complex config in llama-swap

Author: lapy

backend/llama_swap_config.py

@@ -3,6 +3,7 @@
 import subprocess
 import re
 import json
+import shlex
 from typing import Dict, Any, Set, Optional
 from backend.logging_config import get_logger
 
@@ -12,6 +13,35 @@
 _supported_flags_cache: Dict[str, Set[str]] = {}
 
 
+def _quote_arg_if_needed(arg: str) -> str:
+    """
+    Quote an argument if it contains spaces, special characters, or is a complex value.
+    
+    Args:
+        arg: The argument value to potentially quote
+        
+    Returns:
+        The argument, quoted if necessary
+    """
+    if not isinstance(arg, str):
+        arg = str(arg)
+    
+    # Always quote if it contains spaces, quotes, or shell special characters
+    if any(char in arg for char in [' ', '\t', '\n', '|', '&', ';', '(', ')', '<', '>', '*', '?', '[', ']', '{', '}', '$', '`', '\\']):
+        # Use shlex.quote for proper shell escaping
+        return shlex.quote(arg)
+    
+    # Quote regex patterns (common in --override-tensor)
+    if re.search(r'[.*+?^${}|()\[\]\\]', arg):
+        return shlex.quote(arg)
+    
+    # Quote if it starts with a dash (could be confused with a flag)
+    if arg.startswith('-') and not arg.startswith('--'):
+        return shlex.quote(arg)
+    
+    return arg
+
+
 def _coerce_model_config(config_value: Optional[Any]) -> Dict[str, Any]:
     if not config_value:
         return {}
@@ -242,7 +272,9 @@ def generate_llama_swap_config(models: Dict[str, Dict[str, Any]], llama_server_p
                 logger.debug(f"Model {proxy_model_name}: jinja={config.get('jinja')} (type: {type(config.get('jinja'))})")
 
             # Build llama.cpp command arguments
-            cmd_args = [llama_server_path, "--model", model_path, "--port", "${PORT}"]
+            # Quote model path if it contains spaces or special characters
+            quoted_model_path = _quote_arg_if_needed(model_path)
+            cmd_args = [llama_server_path, "--model", quoted_model_path, "--port", "${PORT}"]
 
             # Default values to skip (these cause errors if flag isn't supported)
             default_values = {
@@ -374,7 +406,11 @@ def generate_llama_swap_config(models: Dict[str, Dict[str, Any]], llama_server_p
                                 logger.debug(f"Skipping duplicate --temp flag for {proxy_model_name}")
                                 continue
                             temp_flag_added = True
-                        cmd_args.extend([flag_options[0], str(value)])
+                        # Quote complex values (grammar, json_schema, yaml, etc.)
+                        value_str = str(value)
+                        if key in ("grammar", "json_schema", "yaml") or flag_options[0] in ("--grammar", "--json-schema", "--yaml"):
+                            value_str = _quote_arg_if_needed(value_str)
+                        cmd_args.extend([flag_options[0], value_str])
 
             # Special handling: MoE offload flags
             # Check for direct cpu_moe or n_cpu_moe parameters first (these take precedence)
@@ -442,7 +478,9 @@ def generate_llama_swap_config(models: Dict[str, Dict[str, Any]], llama_server_p
         llama_cpp_config = model_data["config"]
 
         # Build llama.cpp command arguments (using full path to llama-server)
-        cmd_args = [llama_server_path, "--model", model_path, "--port", "${PORT}"]
+        # Quote model path if it contains spaces or special characters
+        quoted_model_path = _quote_arg_if_needed(model_path)
+        cmd_args = [llama_server_path, "--model", quoted_model_path, "--port", "${PORT}"]
 
         # Default values to skip (these cause errors if flag isn't supported)
         default_values = {
@@ -569,7 +607,11 @@ def generate_llama_swap_config(models: Dict[str, Dict[str, Any]], llama_server_p
                             logger.debug(f"Skipping duplicate --temp flag")
                             continue
                         temp_flag_added = True
-                    cmd_args.extend([flag_options[0], str(value)])
+                    # Quote complex values (grammar, json_schema, yaml, etc.)
+                    value_str = str(value)
+                    if key in ("grammar", "json_schema", "yaml") or flag_options[0] in ("--grammar", "--json-schema", "--yaml"):
+                        value_str = _quote_arg_if_needed(value_str)
+                    cmd_args.extend([flag_options[0], value_str])
 
         # Special handling: MoE offload flags
         # Check for direct cpu_moe or n_cpu_moe parameters first (these take precedence)