Merge pull request #225 from AllenNeuralDynamics/fix-xml-rpc-server

Fix edge cases with `xml-rpc-server` remote execution

Author: bruno-f-cruz

examples/_mocks.py

@@ -97,4 +97,4 @@ def create_fake_rig():
     computer_name = os.getenv("COMPUTERNAME")
     os.makedirs(_dir := f"{LIB_CONFIG}/Rig/{computer_name}", exist_ok=True)
     with open(f"{_dir}/rig1.json", "w", encoding="utf-8") as f:
-        f.write(RigModel().model_dump_json(indent=2))
+        f.write(RigModel(data_directory=r"./local/data").model_dump_json(indent=2))

examples/client_behavior_launcher.py

@@ -0,0 +1,86 @@
+import logging
+from pathlib import Path
+
+from _mocks import (
+    LIB_CONFIG,
+    AindBehaviorSessionModel,
+    RigModel,
+    TaskLogicModel,
+    create_fake_rig,
+    create_fake_subjects,
+)
+from pydantic_settings import CliApp
+
+from clabe import resource_monitor
+from clabe.apps import BonsaiApp
+from clabe.launcher import Launcher, LauncherCliArgs, experiment
+from clabe.pickers import DefaultBehaviorPicker, DefaultBehaviorPickerSettings
+from clabe.xml_rpc import XmlRpcClient, XmlRpcClientSettings
+
+logger = logging.getLogger(__name__)
+
+
+@experiment()
+async def client_experiment(launcher: Launcher) -> None:
+    """Demo experiment showcasing CLABE functionality."""
+    picker = DefaultBehaviorPicker(
+        launcher=launcher,
+        settings=DefaultBehaviorPickerSettings(config_library_dir=LIB_CONFIG),
+        experimenter_validator=lambda _: True,
+    )
+
+    session = picker.pick_session(AindBehaviorSessionModel)
+    rig = picker.pick_rig(RigModel)
+    launcher.register_session(session, rig.data_directory)
+    trainer_state, task_logic = picker.pick_trainer_state(TaskLogicModel)
+
+    resource_monitor.ResourceMonitor(
+        constrains=[
+            resource_monitor.available_storage_constraint_factory_from_rig(rig, 2e11),
+        ]
+    ).run()
+
+    xml_rpc_client = XmlRpcClient(settings=XmlRpcClientSettings(server_url="http://localhost:8000", token="42"))
+
+    bonsai_root = Path(r"C:\git\AllenNeuralDynamics\Aind.Behavior.VrForaging")
+    session_response = xml_rpc_client.upload_model(session, "session.json")
+    rig_response = xml_rpc_client.upload_model(rig, "rig.json")
+    task_logic_response = xml_rpc_client.upload_model(task_logic, "task_logic.json")
+    assert rig_response.path is not None
+    assert session_response.path is not None
+    assert task_logic_response.path is not None
+
+    bonsai_app_result = await xml_rpc_client.run_async(
+        BonsaiApp(
+            workflow=bonsai_root / "src/test_deserialization.bonsai",
+            executable=bonsai_root / "bonsai/bonsai.exe",
+            additional_externalized_properties={
+                "RigPath": rig_response.path,
+                "SessionPath": session_response.path,
+                "TaskLogicPath": task_logic_response.path,
+            },
+        ).command
+    )
+    print(bonsai_app_result)
+    return
+
+
+def main():
+    create_fake_subjects()
+    create_fake_rig()
+    behavior_cli_args = CliApp.run(
+        LauncherCliArgs,
+        cli_args=[
+            "--debug-mode",
+            "--allow-dirty",
+            "--skip-hardware-validation",
+        ],
+    )
+
+    launcher = Launcher(settings=behavior_cli_args)
+    launcher.run_experiment(client_experiment)
+    return None
+
+
+if __name__ == "__main__":
+    main()

src/clabe/apps/__init__.py

@@ -5,8 +5,8 @@
     CommandResult,
     ExecutableApp,
     Executor,
-    OutputParser,
     StdCommand,
+    _OutputParser,
     identity_parser,
 )
 from ._bonsai import AindBehaviorServicesBonsaiApp, BonsaiApp
@@ -26,7 +26,7 @@
     "AsyncExecutor",
     "Executor",
     "identity_parser",
-    "OutputParser",
+    "_OutputParser",
     "PythonScriptApp",
     "ExecutableApp",
     "StdCommand",

src/clabe/apps/_base.py

@@ -111,7 +111,7 @@ class ExecutableApp(Protocol):
         class MyApp(ExecutableApp):
             @property
             def command(self) -> Command:
-                return Command(cmd="echo hello", output_parser=identity_parser)
+                return Command(cmd=["echo", "hello"], output_parser=identity_parser)
         ```
     """
 
@@ -170,7 +170,7 @@ async def run_async(self, command: "Command") -> CommandResult:
 
 TOutput = TypeVar("TOutput")
 
-OutputParser: TypeAlias = Callable[[CommandResult], TOutput]
+_OutputParser: TypeAlias = Callable[[CommandResult], TOutput]
 
 
 class Command(Generic[TOutput]):
@@ -181,14 +181,20 @@ class Command(Generic[TOutput]):
     Supports both synchronous and asynchronous execution patterns with type-safe
     output parsing.
 
+    Commands are provided as a list of strings, which is consistent with subprocess
+    and executed directly without shell interpretation. This approach:
+    - Avoids shell injection vulnerabilities
+    - Handles arguments with spaces correctly without manual quoting
+    - Is more portable across platforms
+
     Attributes:
-        cmd: The command string to execute
+        cmd: The command to execute as a list of strings
         result: The result of command execution (available after execution)
 
     Example:
         ```python
-        # Create a simple command
-        cmd = Command(cmd="echo hello", output_parser=identity_parser)
+        # Create a command
+        cmd = Command(cmd=["python", "-c", "print('hello')"], output_parser=identity_parser)
 
         # Execute with a synchronous executor
         executor = LocalExecutor()
@@ -198,24 +204,25 @@ class Command(Generic[TOutput]):
         def parse_json(result: CommandResult) -> dict:
             return json.loads(result.stdout)
 
-        cmd = Command(cmd="get-data --json", output_parser=parse_json)
+        cmd = Command(cmd=["get-data", "--json"], output_parser=parse_json)
         data = cmd.execute(executor)
         ```
     """
 
-    def __init__(self, cmd: str, output_parser: OutputParser[TOutput]) -> None:
+    def __init__(self, cmd: list[str], output_parser: _OutputParser[TOutput]) -> None:
         """Initialize the Command instance.
+
         Args:
-            cmd: The command string to execute
+            cmd: The command to execute as a list of strings. The first element
+                is the program to run, followed by its arguments.
             output_parser: Function to parse the command result into desired output type
 
         Example:
             ```python
-            # Create a simple command
-            cmd = Command(cmd="echo hello", output_parser=identity_parser)
+            cmd = Command(cmd=["echo", "hello"], output_parser=identity_parser)
             ```
         """
-        self._cmd = cmd
+        self._cmd: list[str] = cmd
         self._output_parser = output_parser
         self._result: Optional[CommandResult] = None
 
@@ -227,16 +234,30 @@ def result(self) -> CommandResult:
         return self._result
 
     @property
-    def cmd(self) -> str:
-        """Get the command string."""
+    def cmd(self) -> list[str]:
+        """Get the command as a list of strings."""
         return self._cmd
 
     def append_arg(self, args: str | list[str]) -> Self:
-        """Append an argument to the command."""
+        """Append arguments to the command.
+
+        Args:
+            args: Argument(s) to append. Can be a single string or list of strings.
+                Empty strings are filtered out.
+
+        Returns:
+            Self for method chaining.
+
+        Example:
+            ```python
+            cmd = Command(cmd=["python"], output_parser=identity_parser)
+            cmd.append_arg(["-m", "pytest"])  # Results in ["python", "-m", "pytest"]
+            ```
+        """
         if isinstance(args, str):
             args = [args]
         args = [arg for arg in args if arg]
-        self._cmd = (self.cmd + f" {' '.join(args)}").strip()
+        self._cmd = self._cmd + args
         return self
 
     def execute(self, executor: Executor) -> TOutput:
@@ -267,9 +288,18 @@ def _parse_output(self, result: CommandResult) -> TOutput:
 
 
 class StdCommand(Command[CommandResult]):
-    """Standard command that returns the raw CommandResult."""
+    """Standard command that returns the raw CommandResult.
+
+    A convenience class that creates a Command with the identity_parser,
+    returning the raw CommandResult without transformation.
+
+    Example:
+        ```python
+        cmd = StdCommand(["echo", "hello"])
+        ```
+    """
 
-    def __init__(self, cmd: str) -> None:
+    def __init__(self, cmd: list[str]) -> None:
         super().__init__(cmd, identity_parser)
 
 

src/clabe/apps/_bonsai.py

@@ -4,7 +4,7 @@
 import random
 from os import PathLike
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 import pydantic
 from aind_behavior_services import AindBehaviorRigModel, AindBehaviorSessionModel, AindBehaviorTaskLogicModel
@@ -117,13 +117,18 @@ def _build_bonsai_process_command(
         is_editor_mode: bool = True,
         is_start_flag: bool = True,
         additional_properties: Optional[Dict[str, str]] = None,
-    ) -> str:
+    ) -> List[str]:
         """
-        Builds a shell command for running a Bonsai workflow via subprocess.
+        Builds a command list for running a Bonsai workflow via subprocess.
 
-        Constructs the complete command string with all necessary flags and properties
-        for executing a Bonsai workflow. Handles editor mode, start flag, and
-        externalized properties.
+        Constructs the complete command as a list of arguments with all necessary
+        flags and properties for executing a Bonsai workflow. Handles editor mode,
+        start flag, and externalized properties.
+
+        Using list format is preferred over string format as it:
+        - Avoids shell injection vulnerabilities
+        - Handles paths with spaces correctly without manual quoting
+        - Is more portable across platforms
 
         Args:
             workflow_file: Path to the Bonsai workflow file
@@ -133,7 +138,7 @@ def _build_bonsai_process_command(
             additional_properties: Dictionary of externalized properties to pass. Defaults to None
 
         Returns:
-            str: The complete command string
+            List[str]: The complete command as a list of arguments
 
         Example:
             ```python
@@ -142,19 +147,20 @@ def _build_bonsai_process_command(
                 is_editor_mode=False,
                 additional_properties={"SubjectName": "Mouse123"}
             )
-            # Returns: '"bonsai.exe" "workflow.bonsai" --no-editor -p:"SubjectName"="Mouse123"'
+            # Returns: ["bonsai.exe", "workflow.bonsai", "--no-editor", "-p:SubjectName=Mouse123"]
             ```
         """
-        output_cmd: str = f'"{bonsai_exe}" "{workflow_file}"'
+        output_cmd: List[str] = [str(bonsai_exe), str(workflow_file)]
+
         if is_editor_mode:
             if is_start_flag:
-                output_cmd += " --start"
+                output_cmd.append("--start")
         else:
-            output_cmd += " --no-editor"
+            output_cmd.append("--no-editor")
 
         if additional_properties:
             for param, value in additional_properties.items():
-                output_cmd += f' -p:"{param}"="{value}"'
+                output_cmd.append(f"-p:{param}={value}")
 
         return output_cmd
 

src/clabe/apps/_curriculum.py

@@ -69,7 +69,7 @@ class CurriculumSettings(ServiceSettings):
 
     __yml_section__: t.ClassVar[t.Optional[str]] = "curriculum"
 
-    script: str = "curriculum run"
+    script: list[str] = ["curriculum", "run"]
     project_directory: os.PathLike = Path(".")
     input_trainer_state: t.Optional[os.PathLike] = None
     data_directory: t.Optional[os.PathLike] = None
@@ -145,18 +145,18 @@ def __init__(
             raise ValueError("Data directory is not set.")
 
         kwargs: dict[str, t.Any] = {  # Must use kebab casing
-            "data-directory": f'"{self._settings.data_directory}"',
-            "input-trainer-state": f'"{self._settings.input_trainer_state}"',
+            "data-directory": str(self._settings.data_directory),
+            "input-trainer-state": str(self._settings.input_trainer_state),
         }
         if self._settings.curriculum is not None:
-            kwargs["curriculum"] = f'"{self._settings.curriculum}"'
+            kwargs["curriculum"] = str(self._settings.curriculum)
 
         python_script_app_kwargs = python_script_app_kwargs or {}
         self._python_script_app = PythonScriptApp(
             script=settings.script,
             project_directory=settings.project_directory,
             extra_uv_arguments="-q",
-            additional_arguments=" ".join(f"--{key} {value}" for key, value in kwargs.items()),
+            additional_arguments=[arg for kv in kwargs.items() for arg in ("--" + kv[0], str(kv[1]))],
             **python_script_app_kwargs,
         )