Implement progress reporting for in-flight commands

inorbit_edge/robot.py

@@ -99,6 +99,7 @@
 # CustomCommand execution status
 CUSTOM_COMMAND_STATUS_FINISHED = "finished"
 CUSTOM_COMMAND_STATUS_ABORTED = "aborted"
+CUSTOM_COMMAND_STATUS_RUNNING = "running"
 
 ROBOT_PATH_POINTS_LIMIT = 1000
 
@@ -1023,9 +1024,23 @@ def result_function(
                         stderr,
                     )
 
-            # TODO: Implement progress reporting function
-            def progress_function(output, error):
-                return 1
+            def progress_function(output=None, error=None):
+                """Report that a command is still running.
+
+                A command that takes a while is otherwise silent until it
+                finishes, leaving no way to tell slow progress apart from a
+                stalled command.
+
+                Call it only when a handler wants progress observed: some
+                consumers act on the first status update they see, so a
+                handler that does not opt in keeps reporting a single final
+                status.
+                """
+                if execution_id is None:
+                    return
+                return self.report_command_progress(
+                    command_name, args, execution_id, output, error
+                )
 
             options = {
                 "result_function": result_function,
@@ -1066,6 +1081,28 @@ def report_command_result(
         msg.ts = int(time.time() * 1000)
         self.publish_protobuf(MQTT_SCRIPT_OUTPUT_TOPIC, msg)
 
+    def report_command_progress(
+        self, command_name, args, execution_id, stdout=None, stderr=None
+    ):
+        """Send to server a "still running" update for an in-flight command.
+
+        Mirrors `report_command_result`, but carries no return code: the
+        command has not finished, so there is no outcome to report yet.
+        """
+
+        msg = CustomScriptStatusMessage()
+        msg.file_name = (
+            args[0] if args and isinstance(args[0], (str, bytes)) else command_name
+        )
+        msg.execution_id = execution_id
+        msg.execution_status = CUSTOM_COMMAND_STATUS_RUNNING
+        if stdout:
+            msg.stdout = stdout
+        if stderr:
+            msg.stderr = stderr
+        msg.ts = int(time.time() * 1000)
+        self.publish_protobuf(MQTT_SCRIPT_OUTPUT_TOPIC, msg)
+
     def register_commands_path(self, path="./user_scripts", exec_name_regex=r".*"):
         """Registers executable commands that handle InOrbit custom command actions.
         Use `exec_name_regex` and `path` to customize which executables can be

inorbit_edge/tests/demo/README.md

@@ -14,7 +14,7 @@ cd /path/to/edge-sdk-python
 pip install -e '.[video,telemetry]'
 cd inorbit_edge/tests/demo
 
-export INORBIT_CONNECTION_CONFIG_URL="https://control.inorbit.ai"
+export INORBIT_CONNECTION_CONFIG_URL="https://control.inorbit.ai/cloud_sdk_robot_config"
 export INORBIT_API_URL="https://api.inorbit.ai"
 export INORBIT_API_KEY="foobar123"
 export INORBIT_ROBOT_ID_PREFIX="$(hostname)"
@@ -63,3 +63,32 @@ docker run --rm -p 9464:9464 \
 
 The image sets `INORBIT_METRICS_PORT=9464` and `INORBIT_METRICS_ADDR=0.0.0.0` by default; override or unset
 `INORBIT_METRICS_PORT` to disable the metrics HTTP server.
+
+## Long-running commands
+
+The demo handles two custom commands that take a while and report progress
+while they work, for exercising commands that outlive however long a caller is
+willing to wait for a result:
+
+| Filename | Outcome |
+| --- | --- |
+| `slow_success` | reports progress, then succeeds |
+| `slow_failure` | reports progress, then fails with details |
+
+`cac.yaml` in this directory defines the `slow_success` action; set its scope
+and apply it to try the commands from the UI.
+
+Both take an optional `seconds` argument (default 20):
+
+```
+slow_success seconds 30
+```
+
+Each reports once before starting, then every 5 seconds, then the real result.
+A command that says nothing until it finishes cannot be told apart from one
+that has stalled, so `progress_function` is what makes a slow command
+distinguishable from a stuck one.
+
+They run on their own thread, so the session keeps publishing while they work —
+worth copying if a handler of yours does anything slow, since blocking the
+callback holds up everything else on that session.

inorbit_edge/tests/demo/cac.yaml

@@ -0,0 +1,22 @@
+apiVersion: v0.1
+kind: ActionDefinition
+metadata:
+  id: slow-success
+  scope: account/[your_account_id]
+spec:
+  label: Slow Success
+  group: Demo
+  type: RunScript
+  lock: false
+  confirmation:
+    required: false
+  description: Reports progress for N seconds, then succeeds
+  arguments:
+    - name: filename
+      type: string
+      value: slow_success
+    - name: seconds
+      type: string
+      value: "30"
+      input:
+        control: text

inorbit_edge/tests/demo/example.py

@@ -5,6 +5,7 @@
 import os
 import socket
 import sys
+import threading
 from time import sleep
 from random import randint, uniform, random
 from math import pi, inf
@@ -40,6 +41,16 @@
 NUM_ROBOTS = 2
 NUM_LASERS = 3
 
+# Custom command script names handled by `long_running_command_handler`, for
+# exercising commands that take longer than a caller may be willing to wait.
+CMD_SLOW_SUCCESS = "slow_success"
+CMD_SLOW_FAILURE = "slow_failure"
+LONG_RUNNING_COMMANDS = (CMD_SLOW_SUCCESS, CMD_SLOW_FAILURE)
+# Long enough that the command is still going after a caller would have given
+# up on a silent one.
+SLOW_COMMAND_DEFAULT_SECS = 20
+SLOW_COMMAND_PROGRESS_INTERVAL_SECS = 5
+
 
 def _mqtt_use_ssl():
     """Use TLS for MQTT unless INORBIT_USE_SSL is false, 0, no, or off."""
@@ -148,11 +159,85 @@ def my_command_handler(robot_id, command_name, args, options):
             information about the received command request.
     """
     if command_name == "customCommand":
+        # The long-running commands answer for themselves, on their own thread.
+        if args and args[0] in LONG_RUNNING_COMMANDS:
+            return
         print(f"Received '{command_name}' for robot '{robot_id}'!. {args}")
         # Return '0' for success
         options["result_function"]("0")
 
 
+def _slow_command_seconds(script_args):
+    """Read an optional `seconds` argument from a custom command's arguments."""
+    args = dict(zip(script_args[::2], script_args[1::2])) if script_args else {}
+    try:
+        return max(1, int(args.get("seconds", SLOW_COMMAND_DEFAULT_SECS)))
+    except (TypeError, ValueError):
+        return SLOW_COMMAND_DEFAULT_SECS
+
+
+def long_running_command_handler(robot_id, command_name, args, options):
+    """Handler for commands that take a while, reporting progress as they go.
+
+    A command that says nothing until it finishes cannot be told apart from one
+    that has stalled, and a caller waiting on it may give up before the real
+    result arrives. `progress_function` reports that the command is still
+    running; `result_function` still reports the outcome once it is known.
+
+    Try it from a custom command with a filename of `slow_success` or
+    `slow_failure`, optionally with a `seconds` argument:
+
+        slow_success seconds 30
+
+    Runs on its own thread so the session keeps publishing while it works.
+
+    Args:
+        robot_id (str): InOrbit robot ID
+        command_name (str): InOrbit command e.g. 'customCommand'
+        args (list): Command arguments
+        options (dict): see `my_command_handler`
+    """
+    if command_name != "customCommand" or not args:
+        return
+    script_name = args[0]
+    if script_name not in LONG_RUNNING_COMMANDS:
+        return
+
+    seconds = _slow_command_seconds(args[1] if len(args) > 1 else [])
+
+    def run():
+        logging.info(
+            "%s: starting %r, reporting progress for %ss",
+            robot_id,
+            script_name,
+            seconds,
+        )
+        elapsed = 0
+        # Report before doing any work, so the command is known to be running
+        # rather than merely silent.
+        options["progress_function"](f"{script_name} started, {seconds}s to go")
+        while elapsed < seconds:
+            sleep(min(SLOW_COMMAND_PROGRESS_INTERVAL_SECS, seconds - elapsed))
+            elapsed += SLOW_COMMAND_PROGRESS_INTERVAL_SECS
+            remaining = max(0, seconds - elapsed)
+            options["progress_function"](
+                f"{script_name} still running, {remaining}s to go"
+            )
+
+        if script_name == CMD_SLOW_SUCCESS:
+            logging.info("%s: %r finished", robot_id, script_name)
+            options["result_function"]("0", stdout=f"{script_name} finished")
+        else:
+            logging.info("%s: %r failed", robot_id, script_name)
+            options["result_function"](
+                "1",
+                execution_status_details="Demo failure after a long run",
+                stderr=f"{script_name} was asked to fail",
+            )
+
+    threading.Thread(target=run, name=f"slow_command_{robot_id}", daemon=True).start()
+
+
 def _init_prometheus_metrics():
     """Serve /metrics when INORBIT_METRICS_PORT is set (pip extra telemetry)."""
     port_s = os.environ.get("INORBIT_METRICS_PORT", "").strip()
@@ -226,6 +311,7 @@ def main():
     )
     robot_session_factory.register_command_callback(log_command)
     robot_session_factory.register_command_callback(my_command_handler)
+    robot_session_factory.register_command_callback(long_running_command_handler)
     robot_session_factory.register_commands_path("./user_scripts", r".*\.sh")
 
     robot_session_pool = RobotSessionPool(robot_session_factory)

inorbit_edge/tests/test_robot_session_callbacks.py

@@ -18,6 +18,7 @@
 )
 from inorbit_edge.robot import (
     CUSTOM_COMMAND_STATUS_FINISHED,
+    CUSTOM_COMMAND_STATUS_RUNNING,
     MQTT_SCRIPT_OUTPUT_TOPIC,
     RobotSession,
 )
@@ -474,3 +475,110 @@ def test_report_command_result_file_name(
     assert msg.file_name == expected_file_name
     assert msg.execution_id == "exec_123"
     assert msg.execution_status == CUSTOM_COMMAND_STATUS_FINISHED
+
+
+def _script_status_messages(robot_session):
+    """Parse every CustomScriptStatusMessage published by a session."""
+    messages = []
+    for call in robot_session.client.publish.call_args_list:
+        if MQTT_SCRIPT_OUTPUT_TOPIC not in str(call):
+            continue
+        msg = CustomScriptStatusMessage()
+        msg.ParseFromString(bytes(call[1]["payload"]))
+        messages.append(msg)
+    return messages
+
+
+def test_report_command_progress_publishes_running(
+    mock_mqtt_client, mock_inorbit_api, mock_sleep
+):
+    """A command that takes a while is otherwise silent until it finishes."""
+    robot_session = RobotSession(
+        robot_id="id_123", robot_name="name_123", api_key="apikey_123"
+    )
+    robot_session.connect()
+    robot_session._on_connect(None, None, None, 0, None)
+
+    robot_session.report_command_progress(
+        command_name="customCommand",
+        args=["my_script.sh", ["arg1"]],
+        execution_id="exec_123",
+        stdout="still working",
+    )
+
+    (msg,) = _script_status_messages(robot_session)
+    assert msg.execution_status == CUSTOM_COMMAND_STATUS_RUNNING
+    assert msg.file_name == "my_script.sh"
+    assert msg.execution_id == "exec_123"
+    assert msg.stdout == "still working"
+    # No outcome yet, so no return code to report.
+    assert msg.return_code == ""
+
+
+def test_progress_function_reports_running_then_result(
+    mock_mqtt_client, mock_inorbit_api, mock_sleep
+):
+    """The pair a slow handler emits: still running, then the real outcome."""
+    robot_session = RobotSession(
+        robot_id="id_123", robot_name="name_123", api_key="apikey_123"
+    )
+    robot_session.connect()
+    robot_session._on_connect(None, None, None, 0, None)
+
+    def handler(command_name, args, options):
+        options["progress_function"]()
+        options["result_function"]("0")
+
+    robot_session.register_command_callback(handler)
+    robot_session.dispatch_command(
+        command_name="customCommand",
+        args=["my_script.sh", []],
+        execution_id="exec_123",
+    )
+
+    running, finished = _script_status_messages(robot_session)
+    assert running.execution_status == CUSTOM_COMMAND_STATUS_RUNNING
+    assert finished.execution_status == CUSTOM_COMMAND_STATUS_FINISHED
+    assert finished.return_code == "0"
+    assert running.execution_id == finished.execution_id == "exec_123"
+
+
+def test_progress_is_opt_in(mock_mqtt_client, mock_inorbit_api, mock_sleep):
+    """A handler that does not opt in keeps reporting a single final status.
+
+    Some consumers act on the first status update they see, so progress must
+    never be emitted on a handler's behalf.
+    """
+    robot_session = RobotSession(
+        robot_id="id_123", robot_name="name_123", api_key="apikey_123"
+    )
+    robot_session.connect()
+    robot_session._on_connect(None, None, None, 0, None)
+
+    robot_session.register_command_callback(
+        lambda command_name, args, options: options["result_function"]("0")
+    )
+    robot_session.dispatch_command(
+        command_name="customCommand", args=["my_script.sh", []], execution_id="exec_123"
+    )
+
+    (only,) = _script_status_messages(robot_session)
+    assert only.execution_status == CUSTOM_COMMAND_STATUS_FINISHED
+
+
+def test_progress_function_without_execution_id_publishes_nothing(
+    mock_mqtt_client, mock_inorbit_api, mock_sleep
+):
+    """Nothing can correlate the update, so there is nothing worth sending."""
+    robot_session = RobotSession(
+        robot_id="id_123", robot_name="name_123", api_key="apikey_123"
+    )
+    robot_session.connect()
+    robot_session._on_connect(None, None, None, 0, None)
+
+    robot_session.register_command_callback(
+        lambda command_name, args, options: options["progress_function"]()
+    )
+    robot_session.dispatch_command(command_name="customCommand", args=["s.sh", []])
+
+    assert _script_status_messages(robot_session) == []