Add command (API-dispatch) backend to oz-agent-worker

.gitignore

@@ -31,3 +31,6 @@ vendor/
 
 # OS
 .DS_Store
+
+# Python bytecode
+__pycache__/

README.md

@@ -16,6 +16,7 @@ Self-hosted worker for Oz cloud agents.
   - Docker daemon access for the Docker backend
   - Local `oz` CLI access plus a writable workspace root for the Direct backend
   - Kubernetes API access plus cluster credentials for the Kubernetes backend
+  - An operator-provided dispatch command for the Command backend (dispatch to any runtime over any transport)
 
 ## Usage
 
@@ -64,6 +65,40 @@ backend:
     oz_path: "/usr/local/bin/oz"
 ```
 
+### Command
+
+The command backend hands task execution to an operator-owned runtime over **any transport**, without recompiling the worker. Instead of running the agent itself, the worker invokes an operator-configured `dispatch_command` and lets that command dispatch the task however it likes (HTTP, gRPC, a cloud SDK, a message queue, SSH, etc.). Execution is **fire-and-forget**: once the dispatch command reports success, the remote `oz` agent — launched by your runtime with the worker-provided arguments — reports task progress and terminal state directly to Warp. The worker does not finalize a dispatched task itself; it only reports a failure if the dispatch command itself fails.
+
+Example config:
+
+```yaml
+worker_id: "my-worker"
+backend:
+  command:
+    dispatch_command: "/opt/oz/dispatch.sh"
+    cancel_command: "/opt/oz/cancel.sh"
+    dispatch_timeout: "60s"
+    environment:
+      - name: MY_RUNTIME_TOKEN
+```
+
+Config keys:
+
+- `dispatch_command` (required): shell command (run via `/bin/sh -c`) invoked once per task to dispatch it.
+- `cancel_command` (optional): shell command invoked best-effort when a dispatched task is cancelled. If unset, the worker relies on agent-side cancellation.
+- `dispatch_timeout` (optional): how long the dispatch command may run before it is considered failed (humantime format, e.g. `60s`). Defaults to `60s`.
+- `environment`: extra environment variables exposed to the dispatch/cancel commands (same `name`/`value` semantics as the other backends; omit `value` to inherit from the host).
+
+The dispatch contract:
+
+- The dispatch command receives the task payload as JSON on **stdin**. This is the only place task environment variables and secrets appear — they are deliberately kept out of the subprocess environment and argv.
+- The following non-secret variables are also set in the command's environment for convenience: `OZ_TASK_ID`, `OZ_EXECUTION_ID`, `OZ_WORKER_BACKEND=command`, `OZ_SERVER_ROOT_URL`, `OZ_DOCKER_IMAGE`.
+- The JSON payload contains: `version`, `task_id`, `execution_id`, `server_root_url`, `worker_id`, `docker_image`, `base_args`, `env`, `sidecars`, and `task`. `base_args` is the `oz agent run …` argument vector your runtime should launch the agent with, inside an environment built from `docker_image` and `sidecars`.
+- Exit code `0` means the task was dispatched successfully; the worker will not finalize it (the remote agent reports terminal state to Warp itself). A non-zero exit or a dispatch that exceeds `dispatch_timeout` marks the task failed.
+- The cancel command (when configured) receives `OZ_TASK_ID`, `OZ_EXECUTION_ID`, and `OZ_WORKER_BACKEND=command` in its environment.
+
+Because dispatched tasks run independently of the worker process, the command backend does not consume a local concurrency slot for the lifetime of the remote task, and worker shutdown does not cancel already-dispatched tasks.
+
 ### Kubernetes
 
 The Kubernetes backend creates one Job per task. Cluster selection is controlled by the Kubernetes client config:

examples/command-backend/README.md

@@ -0,0 +1,111 @@
+# Command backend — HTTP REST reference
+
+A templatable reference for the `oz-agent-worker` [`command` backend](../../README.md#command). The worker invokes a dispatch command per task and hands it the task payload as JSON on stdin; this reference **transforms** that payload into a runtime's API shape and forwards it to an HTTP REST endpoint so a self-hosted runtime can launch the agent on demand.
+
+It is written in Python (standard library only — no dependencies to install) so the transformation logic is easy to read and extend.
+
+## Files
+
+- `dispatch.py` — reads the JSON payload on stdin, transforms it (see `transform()`), and `POST`s it to `OZ_DISPATCH_URL`. Exit `0` means dispatched (fire-and-forget); non-zero means the worker fails the task. This is the template for delegating to a remote HTTP runtime.
+- `cancel.py` — `POST`s `{task_id, execution_id}` to `OZ_CANCEL_URL` when a dispatched task is cancelled (best-effort).
+- `dispatch-oz-local.py` — a **local end-to-end** variant that, instead of forwarding to a remote API, launches the real `oz` agent on the host using the payload's `base_args` (fire-and-forget). Use it to exercise the command backend locally and confirm the worker forwards the right payload. See [Local end-to-end testing](#local-end-to-end-testing).
+
+Requires `python3` on the worker host.
+
+## The transformation
+
+Real runtimes rarely accept the worker's payload verbatim. `dispatch.py` keeps the not-hard transformation in one place — the `transform()` function — that you replace to match your API. The example renames/reshapes the worker payload into a `run` object:
+
+```python
+def transform(payload):
+    task = payload.get("task") or {}
+    definition = task.get("task_definition") or {}
+    return {
+        "run": {
+            "task_id": payload["task_id"],
+            "execution_id": payload.get("execution_id", ""),
+            "image": payload.get("docker_image", ""),
+            "command": payload.get("base_args", []),   # the `oz agent run ...` argv
+            "env": payload.get("env", {}),
+            "mounts": [                                  # mount_path -> path
+                {"image": s.get("image", ""), "path": s.get("mount_path", ""),
+                 "read_write": s.get("read_write", False)}
+                for s in (payload.get("sidecars") or [])
+            ],
+            "callback_url": payload.get("server_root_url", ""),
+            "metadata": {
+                "worker_id": payload.get("worker_id", ""),
+                "payload_version": payload.get("version"),
+                "title": task.get("title", ""),
+                "prompt": definition.get("prompt", ""),
+            },
+        }
+    }
+```
+
+## Wiring it into the worker
+
+Point the command backend at the scripts and template the endpoints via `environment` (or host env):
+
+```yaml
+worker_id: "my-worker"
+backend:
+  command:
+    dispatch_command: "python3 /opt/oz/dispatch.py"
+    cancel_command: "python3 /opt/oz/cancel.py"
+    dispatch_timeout: "60s"
+    environment:
+      - name: OZ_DISPATCH_URL
+        value: "https://my-runtime.internal/oz/dispatch"
+      - name: OZ_CANCEL_URL
+        value: "https://my-runtime.internal/oz/cancel"
+      # Omit `value` to inherit the secret from the worker's host environment.
+      - name: OZ_DISPATCH_AUTH_HEADER
+```
+
+(The scripts are executable, so `dispatch_command: "/opt/oz/dispatch.py"` also works.)
+
+## What the worker hands the script (stdin)
+
+```json
+{
+  "version": 1,
+  "task_id": "...",
+  "execution_id": "...",
+  "server_root_url": "https://app.warp.dev",
+  "worker_id": "my-worker",
+  "docker_image": "ubuntu:22.04",
+  "base_args": ["agent", "run", "--task-id", "...", "--server-root-url", "..."],
+  "env": { "GITHUB_ACCESS_TOKEN": "...", "...": "..." },
+  "sidecars": [ { "image": "...", "mount_path": "/agent", "read_write": false } ],
+  "task": { "id": "...", "title": "...", "task_definition": { "prompt": "..." } }
+}
+```
+
+The non-secret identifiers `OZ_TASK_ID`, `OZ_EXECUTION_ID`, `OZ_WORKER_BACKEND`, `OZ_SERVER_ROOT_URL`, and `OZ_DOCKER_IMAGE` are also set in the script's environment. Secrets appear only in the stdin payload.
+
+Your runtime should launch the agent with `base_args` inside an environment built from `docker_image` + `sidecars`, injecting `env`. Because `base_args` already includes `--task-id` and `--server-root-url`, the agent reports its own progress and terminal state to Warp — the worker does not. Keep the exit-code contract: exit `0` only when the task is durably accepted for execution.
+
+## Local end-to-end testing
+
+`dispatch-oz-local.py` lets you exercise the whole command-backend path against a local stack — local warp-server, local session-sharing-server, and a running `oz-agent-worker` — using a real agent run. It reads the payload, logs a summary of what the worker forwarded (so you can verify the contract), writes the full payload to `OZ_LOCAL_RUN_LOG_DIR/payload-<task_id>.json`, then launches `$OZ_BIN <base_args...>` detached with the payload's `env` applied.
+
+Required/optional environment for this script:
+
+- `OZ_BIN` (required): the local `oz`/Warp binary to exec (the same kind of binary the `direct` backend uses).
+- `OZ_LOCAL_RUN_LOG_DIR` (optional): where to write per-task payloads and run logs (defaults to a temp dir).
+
+The easiest way to run the full stack is `warp-server`'s `script/oz-local`, which boots the servers and the worker for you. Once it supports the command backend, run:
+
+```bash
+# from warp-server, with WARP_API_KEY exported and a local oz bundle built
+./script/oz-local --worker-backend command --oz-path <path-to-oz-binary>
+```
+
+Then trigger a run routed to this worker (e.g. from `warp-internal`):
+
+```bash
+WITH_LOCAL_SERVER=1 WITH_LOCAL_SESSION_SHARING_SERVER=1 ./script/run --host-id local-dev
+```
+
+In the `oz-agent-worker` log you should see the worker claim the task and the `[dispatch-oz-local]` lines showing the forwarded `task_id`, `base_args`, and `env` keys; the agent then runs against your local server and reports its own status. Because runs are launched detached (fire-and-forget), they keep running after `script/oz-local` is stopped — find and stop them with `pgrep -fl 'agent run'` / `pkill -f 'agent run'` if needed.

examples/command-backend/cancel.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+"""Reference cancel command for the oz-agent-worker "command" backend.
+
+Invoked best-effort when a dispatched task is cancelled. The worker sets
+``OZ_TASK_ID`` and ``OZ_EXECUTION_ID`` in the environment; this script POSTs them
+to ``OZ_CANCEL_URL`` so your runtime can stop the corresponding agent. Adapt the
+request body to your runtime's API if needed.
+
+Required environment:
+  OZ_CANCEL_URL             REST endpoint to POST the cancellation to.
+
+Optional environment:
+  OZ_DISPATCH_AUTH_HEADER   Authorization header value (e.g. "Bearer ...").
+  OZ_DISPATCH_TIMEOUT_SECS  Request timeout in seconds (default 30).
+"""
+
+import json
+import os
+import sys
+import urllib.error
+import urllib.request
+
+
+def main():
+    url = os.environ.get("OZ_CANCEL_URL")
+    if not url:
+        sys.stderr.write("OZ_CANCEL_URL must be set\n")
+        return 2
+    timeout = float(os.environ.get("OZ_DISPATCH_TIMEOUT_SECS", "30"))
+
+    body = json.dumps(
+        {
+            "task_id": os.environ.get("OZ_TASK_ID", ""),
+            "execution_id": os.environ.get("OZ_EXECUTION_ID", ""),
+        }
+    ).encode("utf-8")
+
+    request = urllib.request.Request(url, data=body, method="POST")
+    request.add_header("Content-Type", "application/json")
+    auth = os.environ.get("OZ_DISPATCH_AUTH_HEADER")
+    if auth:
+        request.add_header("Authorization", auth)
+
+    try:
+        with urllib.request.urlopen(request, timeout=timeout) as response:
+            response.read()
+    except urllib.error.HTTPError as exc:
+        sys.stderr.write(f"cancel endpoint returned HTTP {exc.code}\n")
+        return 1
+    except urllib.error.URLError as exc:
+        sys.stderr.write(f"failed to reach cancel endpoint: {exc}\n")
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

examples/command-backend/dispatch-oz-local.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+"""Local real-run dispatch command for the oz-agent-worker "command" backend.
+
+Unlike dispatch.py (which forwards the payload to an HTTP runtime), this
+reference runs the agent *for real* on the host using the worker-provided
+``base_args``. It exists for local end-to-end testing of the command backend —
+for example via warp-server's ``script/oz-local --worker-backend command`` — so
+you can both (a) verify the worker forwards the right payload to the command and
+(b) trigger an actual run against your local server.
+
+It reads the ``DispatchPayload`` (JSON) on stdin and:
+  1. logs a summary of what it received and writes the full payload to
+     ``OZ_LOCAL_RUN_LOG_DIR/payload-<task_id>.json`` for inspection,
+  2. launches ``$OZ_BIN <base_args...>`` detached (fire-and-forget) with the
+     payload's ``env`` applied, writing the run's output to a per-task log file,
+  3. exits 0 once the run is launched. The agent reports its own terminal state
+     to the server (base_args already carries --task-id / --server-root-url), so
+     the worker does not finalize the task.
+
+This mirrors how the `direct` backend executes the host oz binary, but routed
+through the command backend, so it ignores docker_image / sidecars.
+
+Required environment:
+  OZ_BIN                The oz/Warp binary to exec; base_args is appended to it.
+
+Optional environment:
+  OZ_LOCAL_RUN_LOG_DIR  Directory for per-task payload + run logs (default: tmp).
+
+Exit 0  => run launched (the worker treats the task as dispatched).
+Exit !=0 => failed to launch; the worker marks the task failed.
+"""
+
+import json
+import os
+import shlex
+import subprocess
+import sys
+import tempfile
+
+
+def main():
+    oz_bin = os.environ.get("OZ_BIN")
+    if not oz_bin:
+        sys.stderr.write("OZ_BIN must be set to the oz/Warp binary path\n")
+        return 2
+
+    try:
+        payload = json.load(sys.stdin)
+    except json.JSONDecodeError as exc:
+        sys.stderr.write(f"invalid dispatch payload on stdin: {exc}\n")
+        return 1
+
+    task_id = payload.get("task_id", "unknown")
+    base_args = payload.get("base_args") or []
+    env_overlay = payload.get("env") or {}
+
+    log_dir = os.environ.get("OZ_LOCAL_RUN_LOG_DIR") or tempfile.gettempdir()
+    os.makedirs(log_dir, exist_ok=True)
+
+    # Persist the full payload so the forwarded contents are easy to inspect.
+    payload_path = os.path.join(log_dir, f"payload-{task_id}.json")
+    with open(payload_path, "w", encoding="utf-8") as payload_file:
+        json.dump(payload, payload_file, indent=2, sort_keys=True)
+
+    # Log a summary so the worker log shows exactly what was forwarded.
+    sys.stderr.write(
+        "[dispatch-oz-local] task_id=%s execution_id=%s image=%s\n"
+        % (task_id, payload.get("execution_id", ""), payload.get("docker_image", ""))
+    )
+    sys.stderr.write(
+        "[dispatch-oz-local] base_args: %s\n"
+        % " ".join(shlex.quote(arg) for arg in base_args)
+    )
+    sys.stderr.write(
+        "[dispatch-oz-local] env keys: %s\n" % ", ".join(sorted(env_overlay))
+    )
+    sys.stderr.write("[dispatch-oz-local] full payload: %s\n" % payload_path)
+
+    if not base_args:
+        sys.stderr.write("payload has no base_args; nothing to run\n")
+        return 1
+
+    # Inherit the current environment and overlay the task env (incl. secrets).
+    child_env = dict(os.environ)
+    child_env.update({str(key): str(value) for key, value in env_overlay.items()})
+
+    argv = [oz_bin, *base_args]
+    run_log_path = os.path.join(log_dir, f"oz-run-{task_id}.log")
+
+    # Launch detached so the run outlives this short-lived dispatch command,
+    # matching the command backend's fire-and-forget contract.
+    try:
+        run_log = open(run_log_path, "ab")
+    except OSError as exc:
+        sys.stderr.write(f"failed to open run log {run_log_path}: {exc}\n")
+        return 1
+
+    try:
+        proc = subprocess.Popen(
+            argv,
+            env=child_env,
+            stdin=subprocess.DEVNULL,
+            stdout=run_log,
+            stderr=run_log,
+            start_new_session=True,
+        )
+    except OSError as exc:
+        sys.stderr.write(f"failed to launch oz run: {exc}\n")
+        return 1
+    finally:
+        run_log.close()
+
+    sys.stderr.write(
+        f"[dispatch-oz-local] launched pid={proc.pid}; run output -> {run_log_path}\n"
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())