Merge pull request #87 from GitHubSecurityLab/copilot/add-globals-command-line-support

Add command line support for global variables

Author: kevinbackhouse

README.md

@@ -77,6 +77,18 @@ Example: deploying a Taskflow:
 hatch run main -t examples.taskflows.example
 ```
 
+Example: deploying a Taskflow with command line global variables:
+
+```sh
+hatch run main -t examples.taskflows.example_globals -g fruit=apples
+```
+
+Multiple global variables can be set:
+
+```sh
+hatch run main -t examples.taskflows.example_globals -g fruit=apples -g color=red
+```
+
 ## Deploying from Docker
 
 You can deploy the Taskflow Agent via its Docker image using `docker/run.sh`.
@@ -104,6 +116,13 @@ Example: deploying a Taskflow (example.yaml):
 ```sh
 docker/run.sh -t example
 ```
+
+Example: deploying a Taskflow with global variables:
+
+```sh
+docker/run.sh -t example_globals -g fruit=apples
+```
+
 Example: deploying a custom taskflow (custom_taskflow.yaml):
 
 ```sh

doc/GRAMMAR.md

@@ -340,6 +340,20 @@ taskflow:
         Tell me more about {{ GLOBALS_fruit }}.
 ```
 
+Global variables can also be set or overridden from the command line using the `-g` or `--global` flag:
+
+```sh
+hatch run main -t examples.taskflows.example_globals -g fruit=apples
+```
+
+Multiple global variables can be set by repeating the flag:
+
+```sh
+hatch run main -t examples.taskflows.example_globals -g fruit=apples -g color=red
+```
+
+Command line globals override any globals defined in the taskflow YAML file, allowing you to reuse taskflows with different parameter values without editing the files.
+
 ### Reusable Tasks
 
 Tasks can reuse single step taskflows and optionally override any of its configurations. This is done by setting a `uses` field with a link to the single step taskflow YAML file as its value.

src/seclab_taskflow_agent/__main__.py

@@ -67,21 +67,35 @@ def parse_prompt_args(available_tools: AvailableTools,
     group.add_argument("-p", help="The personality to use (mutex with -t)", required=False)
     group.add_argument("-t", help="The taskflow to use (mutex with -p)", required=False)
     group.add_argument("-l", help="List available tool call models and exit", action='store_true', required=False)
+    parser.add_argument("-g", "--global", dest="globals", action='append', help="Set global variable (KEY=VALUE). Can be used multiple times.", required=False)
     parser.add_argument('prompt', nargs=argparse.REMAINDER)
     #parser.add_argument('remainder', nargs=argparse.REMAINDER, help="Remaining args")
     help_msg = parser.format_help()
     help_msg += "\nExamples:\n\n"
     help_msg += "`-p assistant explain modems to me please`\n"
+    help_msg += "`-t example -g fruit=apples`\n"
+    help_msg += "`-t example -g fruit=apples -g color=red`\n"
     try:
         args = parser.parse_known_args(user_prompt.split(' ') if user_prompt else None)
     except SystemExit as e:
         if e.code == 2:
             logging.error(f"User provided incomplete prompt: {user_prompt}")
-            return None, None, None, help_msg
+            return None, None, None, None, help_msg
     p = args[0].p.strip() if args[0].p else None
     t = args[0].t.strip() if args[0].t else None
     l = args[0].l
-    return p, t, l, ' '.join(args[0].prompt), help_msg
+    
+    # Parse global variables from command line
+    cli_globals = {}
+    if args[0].globals:
+        for g in args[0].globals:
+            if '=' not in g:
+                logging.error(f"Invalid global variable format: {g}. Expected KEY=VALUE")
+                return None, None, None, None, None, help_msg
+            key, value = g.split('=', 1)
+            cli_globals[key.strip()] = value.strip()
+    
+    return p, t, l, cli_globals, ' '.join(args[0].prompt), help_msg
 
 async def deploy_task_agents(available_tools: AvailableTools,
                              agents: dict,
@@ -378,7 +392,7 @@ async def _run_streamed():
 
 
 async def main(available_tools: AvailableTools,
-               p: str | None, t: str | None, prompt: str | None):
+               p: str | None, t: str | None, cli_globals: dict, prompt: str | None):
     last_mcp_tool_results = [] # XXX: memleaky
 
     async def on_tool_end_hook(
@@ -418,7 +432,10 @@ async def on_handoff_hook(
         await render_model_output(f"** 🤖💪 Running Task Flow: {t}\n")
 
         # optional global vars available for the taskflow tasks
+        # Start with globals from taskflow file, then override with CLI globals
         global_variables = taskflow.get('globals', {})
+        if cli_globals:
+            global_variables.update(cli_globals)
         model_config = taskflow.get('model_config', {})
         model_keys = []
         if model_config:
@@ -646,7 +663,7 @@ async def _deploy_task_agents(resolved_agents, prompt):
     cwd = pathlib.Path.cwd()
     available_tools = AvailableTools()
 
-    p, t, l, user_prompt, help_msg = parse_prompt_args(available_tools)
+    p, t, l, cli_globals, user_prompt, help_msg = parse_prompt_args(available_tools)
 
     if l:
         tool_models = list_tool_call_models(os.getenv('COPILOT_TOKEN'))
@@ -658,4 +675,4 @@ async def _deploy_task_agents(resolved_agents, prompt):
         print(help_msg)
         sys.exit(1)
 
-    asyncio.run(main(available_tools, p, t, user_prompt), debug=True)
+    asyncio.run(main(available_tools, p, t, cli_globals, user_prompt), debug=True)

tests/data/test_globals_taskflow.yaml

@@ -0,0 +1,13 @@
+# SPDX-FileCopyrightText: 2025 GitHub
+# SPDX-License-Identifier: MIT
+
+seclab-taskflow-agent:
+  version: 1
+  filetype: taskflow
+
+globals:
+  test_var: default_value
+taskflow:
+  - task:
+      run: |
+        echo "test"

tests/test_yaml_parser.py

@@ -43,5 +43,66 @@ def test_parse_example_taskflows(self):
         assert len(example_task_flow['taskflow']) == 4  # 4 tasks in taskflow
         assert example_task_flow['taskflow'][0]['task']['max_steps'] == 20
 
+class TestCliGlobals:
+    """Test CLI global variable parsing."""
+    
+    def test_parse_single_global(self):
+        """Test parsing a single global variable from command line."""
+        from seclab_taskflow_agent.__main__ import parse_prompt_args
+        available_tools = AvailableTools()
+        
+        p, t, l, cli_globals, user_prompt, _ = parse_prompt_args(
+            available_tools, "-t example -g fruit=apples")
+        
+        assert t == "example"
+        assert cli_globals == {"fruit": "apples"}
+        assert p is None
+        assert l is False
+    
+    def test_parse_multiple_globals(self):
+        """Test parsing multiple global variables from command line."""
+        from seclab_taskflow_agent.__main__ import parse_prompt_args
+        available_tools = AvailableTools()
+        
+        p, t, l, cli_globals, user_prompt, _ = parse_prompt_args(
+            available_tools, "-t example -g fruit=apples -g color=red")
+        
+        assert t == "example"
+        assert cli_globals == {"fruit": "apples", "color": "red"}
+        assert p is None
+        assert l is False
+    
+    def test_parse_global_with_spaces(self):
+        """Test parsing global variables with spaces in values."""
+        from seclab_taskflow_agent.__main__ import parse_prompt_args
+        available_tools = AvailableTools()
+        
+        p, t, l, cli_globals, user_prompt, _ = parse_prompt_args(
+            available_tools, "-t example -g message=hello world")
+        
+        assert t == "example"
+        # "world" becomes part of the prompt, not the value
+        assert cli_globals == {"message": "hello"}
+        assert "world" in user_prompt
+    
+    def test_parse_global_with_equals_in_value(self):
+        """Test parsing global variables with equals sign in value."""
+        from seclab_taskflow_agent.__main__ import parse_prompt_args
+        available_tools = AvailableTools()
+        
+        p, t, l, cli_globals, user_prompt, _ = parse_prompt_args(
+            available_tools, "-t example -g equation=x=5")
+        
+        assert t == "example"
+        assert cli_globals == {"equation": "x=5"}
+    
+    def test_globals_in_taskflow_file(self):
+        """Test that globals can be read from taskflow file."""
+        available_tools = AvailableTools()
+        
+        taskflow = available_tools.get_taskflow("tests.data.test_globals_taskflow")
+        assert 'globals' in taskflow
+        assert taskflow['globals']['test_var'] == 'default_value'
+
 if __name__ == '__main__':
     pytest.main([__file__, '-v'])