From 545822c47ea1f267a463242ade514a9683ce5f35 Mon Sep 17 00:00:00 2001 From: zachbai Date: Wed, 24 Jun 2026 16:06:44 -0700 Subject: [PATCH 1/2] Add command (API-dispatch) backend to oz-agent-worker Introduce a transport-agnostic `command` backend that delegates task execution to an operator-configured dispatch command (fire-and-forget). The worker renders a versioned JSON dispatch payload to the command's stdin; on success it suppresses its own terminal completion so the remote oz agent reports terminal state to warp-server itself, and routes cancellations to an optional cancel command. Includes config/CLI wiring, unit + integration tests, README docs, and an HTTP REST reference dispatch script. Co-Authored-By: Oz --- .gitignore | 3 + README.md | 35 ++++ examples/command-backend/README.md | 86 +++++++++ examples/command-backend/cancel.py | 57 ++++++ examples/command-backend/dispatch.py | 113 ++++++++++++ examples/command-backend/dispatch_test.go | 146 +++++++++++++++ examples/command-backend/doc.go | 6 + internal/config/config.go | 11 ++ internal/config/config_test.go | 70 +++++++ internal/metrics/metrics.go | 11 ++ internal/worker/backend.go | 24 +++ internal/worker/command.go | 158 ++++++++++++++++ internal/worker/command_integration_test.go | 153 ++++++++++++++++ internal/worker/command_test.go | 192 ++++++++++++++++++++ internal/worker/dispatch_payload.go | 55 ++++++ internal/worker/dispatch_payload_test.go | 83 +++++++++ internal/worker/dispatch_test.go | 146 +++++++++++++++ internal/worker/worker.go | 89 ++++++++- main.go | 37 +++- main_test.go | 68 +++++++ 20 files changed, 1532 insertions(+), 11 deletions(-) create mode 100644 examples/command-backend/README.md create mode 100755 examples/command-backend/cancel.py create mode 100755 examples/command-backend/dispatch.py create mode 100644 examples/command-backend/dispatch_test.go create mode 100644 examples/command-backend/doc.go create mode 100644 internal/worker/command.go create mode 100644 internal/worker/command_integration_test.go create mode 100644 internal/worker/command_test.go create mode 100644 internal/worker/dispatch_payload.go create mode 100644 internal/worker/dispatch_payload_test.go create mode 100644 internal/worker/dispatch_test.go diff --git a/.gitignore b/.gitignore index d2ee7b6..9b7e191 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,6 @@ vendor/ # OS .DS_Store + +# Python bytecode +__pycache__/ diff --git a/README.md b/README.md index ab42058..1698439 100644 --- a/README.md +++ b/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: diff --git a/examples/command-backend/README.md b/examples/command-backend/README.md new file mode 100644 index 0000000..7b5d561 --- /dev/null +++ b/examples/command-backend/README.md @@ -0,0 +1,86 @@ +# 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. +- `cancel.py` — `POST`s `{task_id, execution_id}` to `OZ_CANCEL_URL` when a dispatched task is cancelled (best-effort). + +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. diff --git a/examples/command-backend/cancel.py b/examples/command-backend/cancel.py new file mode 100755 index 0000000..8367e9d --- /dev/null +++ b/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()) diff --git a/examples/command-backend/dispatch.py b/examples/command-backend/dispatch.py new file mode 100755 index 0000000..c79e051 --- /dev/null +++ b/examples/command-backend/dispatch.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +"""Reference dispatch command for the oz-agent-worker "command" backend. + +It reads the task ``DispatchPayload`` (JSON) on stdin, transforms it into a +hypothetical runtime's REST API shape, and POSTs it to ``OZ_DISPATCH_URL``. + +Use it as a template for delegating task execution to a self-hosted runtime that +is already configured to run oz agents on demand. The part you customize is +``transform()`` — adapt it to your runtime's request schema. Everything else +(reading stdin, auth, timeouts, exit-code contract) can stay as-is. + +Only the Python standard library is used, so there are no dependencies to +install on the worker host. + +Required environment: + OZ_DISPATCH_URL REST endpoint to POST the transformed body to. + +Optional environment: + OZ_DISPATCH_AUTH_HEADER Authorization header value (e.g. "Bearer ..."). + OZ_DISPATCH_TIMEOUT_SECS Request timeout in seconds (default 30). + +Also provided by the worker (no need to set these yourself): + OZ_TASK_ID, OZ_EXECUTION_ID, OZ_WORKER_BACKEND, OZ_SERVER_ROOT_URL, OZ_DOCKER_IMAGE + +Exit semantics (the contract the command backend relies on): + exit 0 => task accepted for dispatch; the remote runtime now owns it and + the agent reports terminal state to Warp itself. + exit != 0 => dispatch failed; the worker marks the task failed. +""" + +import json +import os +import sys +import urllib.error +import urllib.request + + +def transform(payload): + """Map the worker's DispatchPayload onto our runtime's API shape. + + This is an example of an arbitrary, not-hard transformation: fields are + renamed and nested under a ``run`` object, sidecar mounts are restructured + (``mount_path`` -> ``path``), and a few values are lifted into a metadata + block. Replace the body of this function to match your own API. + """ + 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", ""), + # base_args is the `oz agent run ...` argv the runtime should exec. + "command": payload.get("base_args", []), + "env": payload.get("env", {}), + "mounts": [ + { + "image": sidecar.get("image", ""), + "path": sidecar.get("mount_path", ""), + "read_write": sidecar.get("read_write", False), + } + for sidecar 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", ""), + }, + } + } + + +def main(): + url = os.environ.get("OZ_DISPATCH_URL") + if not url: + sys.stderr.write("OZ_DISPATCH_URL must be set\n") + return 2 + timeout = float(os.environ.get("OZ_DISPATCH_TIMEOUT_SECS", "30")) + + try: + payload = json.load(sys.stdin) + except json.JSONDecodeError as exc: + sys.stderr.write(f"invalid dispatch payload on stdin: {exc}\n") + return 1 + + body = json.dumps(transform(payload)).encode("utf-8") + + request = urllib.request.Request(url, data=body, method="POST") + request.add_header("Content-Type", "application/json") + request.add_header("X-Oz-Task-Id", os.environ.get("OZ_TASK_ID", "")) + auth = os.environ.get("OZ_DISPATCH_AUTH_HEADER") + if auth: + request.add_header("Authorization", auth) + + try: + # Any 2xx is success; urllib raises HTTPError for status >= 400. + with urllib.request.urlopen(request, timeout=timeout) as response: + response.read() + except urllib.error.HTTPError as exc: + detail = exc.read().decode("utf-8", errors="replace") + sys.stderr.write(f"dispatch endpoint returned HTTP {exc.code}: {detail}\n") + return 1 + except urllib.error.URLError as exc: + sys.stderr.write(f"failed to reach dispatch endpoint: {exc}\n") + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/command-backend/dispatch_test.go b/examples/command-backend/dispatch_test.go new file mode 100644 index 0000000..652224c --- /dev/null +++ b/examples/command-backend/dispatch_test.go @@ -0,0 +1,146 @@ +package commandbackendexample + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "strings" + "testing" +) + +// samplePayload mirrors the worker's DispatchPayload JSON for a single task. +const samplePayload = `{ + "version": 1, + "task_id": "task-1", + "execution_id": "exec-1", + "server_root_url": "https://app.warp.dev", + "worker_id": "my-worker", + "docker_image": "ubuntu:22.04", + "base_args": ["agent", "run", "--task-id", "task-1"], + "env": {"GITHUB_ACCESS_TOKEN": "secret-token"}, + "sidecars": [{"image": "warpdotdev/warp-agent:latest", "mount_path": "/agent", "read_write": false}], + "task": {"id": "task-1", "title": "do the thing", "task_definition": {"prompt": "go"}} +}` + +// transformedBody is the shape dispatch.py's transform() produces. +type transformedBody struct { + Run struct { + TaskID string `json:"task_id"` + ExecutionID string `json:"execution_id"` + Image string `json:"image"` + Command []string `json:"command"` + Env map[string]string `json:"env"` + Mounts []struct { + Image string `json:"image"` + Path string `json:"path"` + ReadWrite bool `json:"read_write"` + } `json:"mounts"` + CallbackURL string `json:"callback_url"` + Metadata struct { + WorkerID string `json:"worker_id"` + PayloadVersion int `json:"payload_version"` + Title string `json:"title"` + Prompt string `json:"prompt"` + } `json:"metadata"` + } `json:"run"` +} + +func requirePython(t *testing.T) { + t.Helper() + if _, err := exec.LookPath("python3"); err != nil { + t.Skip("python3 not available; skipping reference dispatch script test") + } +} + +func runDispatch(t *testing.T, env []string, payload string) error { + t.Helper() + requirePython(t) + cmd := exec.Command("python3", "dispatch.py") + cmd.Env = append(os.Environ(), env...) + cmd.Stdin = strings.NewReader(payload) + out, err := cmd.CombinedOutput() + if err != nil { + t.Logf("dispatch.py output:\n%s", out) + } + return err +} + +func TestDispatchScriptTransformsAndPostsPayload(t *testing.T) { + var rawBody []byte + var gotTaskID, gotContentType string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rawBody, _ = io.ReadAll(r.Body) + gotTaskID = r.Header.Get("X-Oz-Task-Id") + gotContentType = r.Header.Get("Content-Type") + w.WriteHeader(http.StatusAccepted) + })) + defer srv.Close() + + if err := runDispatch(t, []string{ + "OZ_DISPATCH_URL=" + srv.URL, + "OZ_TASK_ID=task-1", + }, samplePayload); err != nil { + t.Fatalf("dispatch script returned error on 2xx: %v", err) + } + + if gotContentType != "application/json" { + t.Errorf("Content-Type = %q, want application/json", gotContentType) + } + if gotTaskID != "task-1" { + t.Errorf("X-Oz-Task-Id = %q, want task-1", gotTaskID) + } + + var got transformedBody + if err := json.Unmarshal(rawBody, &got); err != nil { + t.Fatalf("posted body is not the expected transformed JSON: %v\nbody: %s", err, rawBody) + } + if got.Run.TaskID != "task-1" { + t.Errorf("run.task_id = %q, want task-1", got.Run.TaskID) + } + if got.Run.Image != "ubuntu:22.04" { + t.Errorf("run.image = %q, want ubuntu:22.04", got.Run.Image) + } + if len(got.Run.Command) != 4 || got.Run.Command[0] != "agent" || got.Run.Command[1] != "run" { + t.Errorf("run.command = %v, want the agent run argv", got.Run.Command) + } + if got.Run.Env["GITHUB_ACCESS_TOKEN"] != "secret-token" { + t.Errorf("run.env[GITHUB_ACCESS_TOKEN] = %q, want secret-token", got.Run.Env["GITHUB_ACCESS_TOKEN"]) + } + if len(got.Run.Mounts) != 1 || got.Run.Mounts[0].Path != "/agent" { + t.Errorf("run.mounts = %+v, want a single mount at /agent (mount_path -> path)", got.Run.Mounts) + } + if got.Run.Metadata.PayloadVersion != 1 { + t.Errorf("run.metadata.payload_version = %d, want 1", got.Run.Metadata.PayloadVersion) + } + if got.Run.Metadata.Title != "do the thing" { + t.Errorf("run.metadata.title = %q, want 'do the thing'", got.Run.Metadata.Title) + } +} + +func TestDispatchScriptFailsOnServerError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer srv.Close() + + if err := runDispatch(t, []string{ + "OZ_DISPATCH_URL=" + srv.URL, + "OZ_TASK_ID=task-1", + }, samplePayload); err == nil { + t.Fatal("expected non-zero exit when the endpoint returns 500") + } +} + +func TestDispatchScriptRequiresURL(t *testing.T) { + requirePython(t) + cmd := exec.Command("python3", "dispatch.py") + // Start from a clean env so OZ_DISPATCH_URL is definitely unset. + cmd.Env = []string{"PATH=" + os.Getenv("PATH")} + cmd.Stdin = strings.NewReader(samplePayload) + if err := cmd.Run(); err == nil { + t.Fatal("expected non-zero exit when OZ_DISPATCH_URL is unset") + } +} diff --git a/examples/command-backend/doc.go b/examples/command-backend/doc.go new file mode 100644 index 0000000..9c32b41 --- /dev/null +++ b/examples/command-backend/doc.go @@ -0,0 +1,6 @@ +// Package commandbackendexample holds a reference dispatch command for the +// oz-agent-worker "command" backend that transforms the task payload and +// delegates execution to an HTTP REST endpoint. The Python scripts (dispatch.py, +// cancel.py) are the actual artifacts; the Go test in this package verifies the +// reference dispatch script's transformation and exit-code contract. +package commandbackendexample diff --git a/internal/config/config.go b/internal/config/config.go index 864629c..7bcdcb0 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -34,6 +34,17 @@ type BackendConfig struct { Docker *DockerConfig `yaml:"docker"` Direct *DirectConfig `yaml:"direct"` Kubernetes *KubernetesConfig `yaml:"kubernetes"` + Command *CommandConfig `yaml:"command"` +} + +// CommandConfig holds command-backend-specific configuration. The command +// backend dispatches tasks to an operator-owned runtime over any transport by +// invoking dispatch_command. +type CommandConfig struct { + DispatchCommand string `yaml:"dispatch_command" validate:"required"` + CancelCommand string `yaml:"cancel_command"` + DispatchTimeout string `yaml:"dispatch_timeout"` + Environment []EnvEntry `yaml:"environment" validate:"dive"` } // DockerConfig holds Docker-backend-specific configuration. diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 8c43506..e869db9 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -376,6 +376,76 @@ func TestLoadFileNotFound(t *testing.T) { } } +func TestLoadValidCommandConfig(t *testing.T) { + path := writeTestConfig(t, ` +worker_id: "command-worker" +backend: + command: + dispatch_command: "/opt/oz/dispatch.sh" + cancel_command: "/opt/oz/cancel.sh" + dispatch_timeout: "30s" + environment: + - name: MY_VAR + value: "hello" +`) + + cfg, err := Load(path) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if cfg.Backend.Command == nil { + t.Fatal("expected command backend to be set") + } + if cfg.Backend.Docker != nil { + t.Error("docker backend should be nil") + } + if cfg.Backend.Command.DispatchCommand != "/opt/oz/dispatch.sh" { + t.Errorf("dispatch_command = %q, want %q", cfg.Backend.Command.DispatchCommand, "/opt/oz/dispatch.sh") + } + if cfg.Backend.Command.CancelCommand != "/opt/oz/cancel.sh" { + t.Errorf("cancel_command = %q, want %q", cfg.Backend.Command.CancelCommand, "/opt/oz/cancel.sh") + } + if cfg.Backend.Command.DispatchTimeout != "30s" { + t.Errorf("dispatch_timeout = %q, want %q", cfg.Backend.Command.DispatchTimeout, "30s") + } + if len(cfg.Backend.Command.Environment) != 1 { + t.Errorf("environment count = %d, want 1", len(cfg.Backend.Command.Environment)) + } +} + +func TestLoadCommandConfigRequiresDispatchCommand(t *testing.T) { + path := writeTestConfig(t, ` +worker_id: "command-worker" +backend: + command: + cancel_command: "/opt/oz/cancel.sh" +`) + + _, err := Load(path) + if err == nil { + t.Fatal("expected error for missing dispatch_command") + } + if !strings.Contains(err.Error(), "DispatchCommand") && !strings.Contains(err.Error(), "required") { + t.Errorf("error = %v, want it to mention the required dispatch command", err) + } +} + +func TestLoadCommandAndDockerBackendsError(t *testing.T) { + path := writeTestConfig(t, ` +backend: + command: + dispatch_command: "/opt/oz/dispatch.sh" + docker: + volumes: [] +`) + + _, err := Load(path) + if err == nil { + t.Fatal("expected error when command and docker backends are both set") + } +} + func TestLoadIdleOnComplete(t *testing.T) { t.Run("parses idle_on_complete when set", func(t *testing.T) { path := writeTestConfig(t, ` diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go index 8a347d4..e516134 100644 --- a/internal/metrics/metrics.go +++ b/internal/metrics/metrics.go @@ -110,6 +110,9 @@ const ( TaskFailureReasonSetupCommand = "setup_command" TaskFailureReasonAgentInvocation = "agent_invocation" TaskFailureReasonTeardownCommand = "teardown_command" + TaskFailureReasonDispatchCommand = "dispatch_command" + TaskFailureReasonDispatchTimeout = "dispatch_timeout" + TaskFailureReasonCancelCommand = "cancel_command" TaskFailureReasonJobCreate = "job_create" TaskFailureReasonJobWatch = "job_watch" TaskFailureReasonJobFailed = "job_failed" @@ -127,6 +130,7 @@ var ( TaskResultSucceeded, TaskResultFailed, TaskResultCancelled, + TaskResultDispatched, } taskFailurePhases = []string{ TaskFailurePhaseAssignment, @@ -148,6 +152,9 @@ var ( TaskFailureReasonSetupCommand, TaskFailureReasonAgentInvocation, TaskFailureReasonTeardownCommand, + TaskFailureReasonDispatchCommand, + TaskFailureReasonDispatchTimeout, + TaskFailureReasonCancelCommand, TaskFailureReasonJobCreate, TaskFailureReasonJobWatch, TaskFailureReasonJobFailed, @@ -460,6 +467,10 @@ const ( TaskResultSucceeded TaskResult = "succeeded" TaskResultFailed TaskResult = "failed" TaskResultCancelled TaskResult = "cancelled" + // TaskResultDispatched is recorded when a backend hands a task off to a + // remote runtime (fire-and-forget) rather than running it to completion + // on the worker. + TaskResultDispatched TaskResult = "dispatched" ) // RecordTaskCompleted records a completed task, regardless of outcome, along diff --git a/internal/worker/backend.go b/internal/worker/backend.go index 1807ba3..b66d781 100644 --- a/internal/worker/backend.go +++ b/internal/worker/backend.go @@ -2,10 +2,17 @@ package worker import ( "context" + "errors" "github.com/warpdotdev/oz-agent-worker/internal/types" ) +// ErrTaskDispatched is returned by Backend.ExecuteTask to signal a successful +// fire-and-forget dispatch to a remote runtime. The worker treats the task as +// accepted-but-not-finalized: it must NOT send a terminal completion message, +// because the remote runtime/agent owns terminal reporting back to warp-server. +var ErrTaskDispatched = errors.New("task dispatched to remote runtime") + // TaskParams contains pre-processed task parameters common to all backends. // This provides a layer of abstraction between the wire-format TaskAssignmentMessage // and the backend interface, so backends don't need to handle common concerns like @@ -45,3 +52,20 @@ type Backend interface { // Shutdown cleans up backend resources. Shutdown(ctx context.Context) } + +// CancelParams carries the minimal, non-secret identifiers needed to attempt +// cancellation of a task that has already been handed off. It deliberately +// excludes env/secrets so the worker need not retain secrets for the lifetime +// of a dispatched task. +type CancelParams struct { + TaskID string + ExecutionID string +} + +// CancelableBackend is implemented by backends that can attempt to cancel a +// task that has already been handed off (e.g. fire-and-forget dispatch). The +// worker invokes CancelTask when it receives a cancellation for a task that is +// no longer executing locally. +type CancelableBackend interface { + CancelTask(ctx context.Context, params *CancelParams) error +} diff --git a/internal/worker/command.go b/internal/worker/command.go new file mode 100644 index 0000000..46bf17b --- /dev/null +++ b/internal/worker/command.go @@ -0,0 +1,158 @@ +package worker + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/log" + "github.com/warpdotdev/oz-agent-worker/internal/metrics" +) + +// defaultDispatchTimeout bounds how long the operator dispatch command may run. +// The command is expected to hand the task off to a remote runtime and return +// promptly; it is not meant to stay alive for the task's lifetime. +const defaultDispatchTimeout = 60 * time.Second + +// CommandBackendConfig configures the command backend, which dispatches tasks to +// an operator-owned runtime over any transport by invoking a shell command. +type CommandBackendConfig struct { + // DispatchCommand is the shell command invoked (via /bin/sh -c) to dispatch + // a task. It receives the DispatchPayload as JSON on stdin. Required. + DispatchCommand string + // CancelCommand, when set, is invoked (via /bin/sh -c) to best-effort cancel + // a previously dispatched task. + CancelCommand string + // DispatchTimeout bounds the dispatch command's runtime. Zero uses defaultDispatchTimeout. + DispatchTimeout time.Duration + // Env contains extra environment variables exposed to the dispatch and + // cancel commands. Task env/secrets are NOT placed here; they travel only in + // the JSON payload on stdin. + Env map[string]string + // ServerRootURL and WorkerID are copied into the dispatch payload. + ServerRootURL string + WorkerID string +} + +// CommandBackend hands task execution to an operator-configured command, which +// dispatches the task to a remote runtime over any transport. Execution is +// fire-and-forget: a successful dispatch returns ErrTaskDispatched, and the +// remote oz agent reports terminal state to warp-server itself. +type CommandBackend struct { + config CommandBackendConfig +} + +var ( + _ Backend = (*CommandBackend)(nil) + _ CancelableBackend = (*CommandBackend)(nil) +) + +// NewCommandBackend constructs a command backend, requiring a dispatch command. +func NewCommandBackend(ctx context.Context, config CommandBackendConfig) (*CommandBackend, error) { + if config.DispatchCommand == "" { + return nil, fmt.Errorf("command backend requires a dispatch command") + } + if config.DispatchTimeout <= 0 { + config.DispatchTimeout = defaultDispatchTimeout + } + + hasCancel := config.CancelCommand != "" + log.Infof(ctx, "Using command backend (cancel command configured: %t, dispatch timeout: %s)", hasCancel, config.DispatchTimeout) + + return &CommandBackend{config: config}, nil +} + +// ExecuteTask dispatches the task by invoking the configured dispatch command +// with the JSON payload on stdin. It returns ErrTaskDispatched on success. +func (b *CommandBackend) ExecuteTask(ctx context.Context, params *TaskParams) error { + payload := NewDispatchPayload(params, b.config.ServerRootURL, b.config.WorkerID) + payloadJSON, err := json.Marshal(payload) + if err != nil { + return newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonDispatchCommand, fmt.Errorf("failed to marshal dispatch payload: %w", err)) + } + + dctx, cancel := context.WithTimeout(ctx, b.config.DispatchTimeout) + defer cancel() + + env := b.commandEnv([]string{ + fmt.Sprintf("OZ_TASK_ID=%s", params.TaskID), + fmt.Sprintf("OZ_EXECUTION_ID=%s", params.ExecutionID), + "OZ_WORKER_BACKEND=command", + fmt.Sprintf("OZ_SERVER_ROOT_URL=%s", b.config.ServerRootURL), + fmt.Sprintf("OZ_DOCKER_IMAGE=%s", params.DockerImage), + }) + + cmd := exec.CommandContext(dctx, "/bin/sh", "-c", b.config.DispatchCommand) // #nosec G204 -- dispatch command is explicit operator configuration. + cmd.Stdin = bytes.NewReader(payloadJSON) + cmd.Env = env + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + + log.Infof(ctx, "Dispatching task %s via command backend", params.TaskID) + if err := cmd.Run(); err != nil { + // The parent context being cancelled means the worker is cancelling the + // task (user/shutdown), not a dispatch failure. + if ctx.Err() != nil { + return newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonTaskCancelled, ctx.Err()) + } + // The dispatch-scoped context expiring means the command overran its budget. + if dctx.Err() == context.DeadlineExceeded { + return newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonDispatchTimeout, fmt.Errorf("dispatch command timed out after %s: %w", b.config.DispatchTimeout, err)) + } + return newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonDispatchCommand, fmt.Errorf("dispatch command failed: %w", err)) + } + + log.Infof(ctx, "Task %s dispatched successfully", params.TaskID) + return ErrTaskDispatched +} + +// CancelTask best-effort cancels a dispatched task via the configured cancel +// command. It is a no-op when no cancel command is configured. +func (b *CommandBackend) CancelTask(ctx context.Context, params *CancelParams) error { + if b.config.CancelCommand == "" { + return nil + } + + env := b.commandEnv([]string{ + fmt.Sprintf("OZ_TASK_ID=%s", params.TaskID), + fmt.Sprintf("OZ_EXECUTION_ID=%s", params.ExecutionID), + "OZ_WORKER_BACKEND=command", + }) + + cmd := exec.CommandContext(ctx, "/bin/sh", "-c", b.config.CancelCommand) // #nosec G204 -- cancel command is explicit operator configuration. + cmd.Env = env + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + + log.Infof(ctx, "Cancelling dispatched task %s via command backend", params.TaskID) + if err := cmd.Run(); err != nil { + return newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonCancelCommand, fmt.Errorf("cancel command failed: %w", err)) + } + return nil +} + +// PreservesTasksOnShutdown reports true: dispatched tasks run on the operator's +// remote runtime, independent of this worker, so worker shutdown must not cancel them. +func (b *CommandBackend) PreservesTasksOnShutdown() bool { return true } + +// Shutdown has nothing to clean up; the backend owns no local resources. +func (b *CommandBackend) Shutdown(ctx context.Context) { + log.Debugf(ctx, "Command backend shutdown (no-op)") +} + +// commandEnv builds the subprocess environment: the host environment overlaid +// with the well-known OZ_* vars and the operator-configured Env. Task env vars +// (which may include secrets) are deliberately excluded — they travel only in +// the JSON payload on stdin. +func (b *CommandBackend) commandEnv(wellKnown []string) []string { + overlay := make([]string, 0, len(wellKnown)+len(b.config.Env)) + overlay = append(overlay, wellKnown...) + for key, value := range b.config.Env { + overlay = append(overlay, fmt.Sprintf("%s=%s", key, value)) + } + return mergeEnvVars(os.Environ(), overlay) +} diff --git a/internal/worker/command_integration_test.go b/internal/worker/command_integration_test.go new file mode 100644 index 0000000..50027fa --- /dev/null +++ b/internal/worker/command_integration_test.go @@ -0,0 +1,153 @@ +package worker + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "testing" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +// waitFor polls cond until it returns true or the timeout elapses. +func waitFor(t *testing.T, timeout time.Duration, cond func() bool) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if cond() { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatal("condition not met within timeout") +} + +func (w *Worker) activeTaskCount() int { + w.tasksMutex.Lock() + defer w.tasksMutex.Unlock() + return len(w.activeTasks) +} + +func drainMessages(t *testing.T, ch <-chan []byte) []types.WebSocketMessage { + t.Helper() + var msgs []types.WebSocketMessage + for { + select { + case b := <-ch: + var m types.WebSocketMessage + if err := json.Unmarshal(b, &m); err != nil { + t.Fatalf("failed to unmarshal websocket message: %v", err) + } + msgs = append(msgs, m) + default: + return msgs + } + } +} + +func newCommandWorker(t *testing.T, dispatchCommand string) *Worker { + t.Helper() + w, err := New(context.Background(), Config{ + WorkerID: "it-worker", + ServerRootURL: "https://app.warp.dev", + BackendType: "command", + Command: &CommandBackendConfig{ + DispatchCommand: dispatchCommand, + ServerRootURL: "https://app.warp.dev", + WorkerID: "it-worker", + }, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + return w +} + +func commandTaskAssignment() *types.TaskAssignmentMessage { + return &types.TaskAssignmentMessage{ + TaskID: "task-1", + ExecutionID: "exec-1", + Task: &types.Task{ + ID: "task-1", + Title: "integration task", + Owner: &types.TaskOwner{Type: "TEAM", Id: 1}, + }, + DockerImage: "ubuntu:22.04", + EnvVars: map[string]string{"GITHUB_ACCESS_TOKEN": "secret-token"}, + } +} + +func TestIntegrationCommandBackendDispatchSuppressesTerminalMessage(t *testing.T) { + outFile := filepath.Join(t.TempDir(), "payload.json") + w := newCommandWorker(t, "cat > "+outFile) + + w.handleTaskAssignment(commandTaskAssignment()) + + // The dispatch script writes the captured payload, then executeTask returns + // and the deferred cleanup empties activeTasks. + waitFor(t, 5*time.Second, func() bool { + _, err := os.Stat(outFile) + return err == nil + }) + waitFor(t, 5*time.Second, func() bool { return w.activeTaskCount() == 0 }) + + // Verify the captured dispatch payload. + data, err := os.ReadFile(outFile) // #nosec G304 -- test-controlled temp path. + if err != nil { + t.Fatalf("failed to read captured payload: %v", err) + } + var payload DispatchPayload + if err := json.Unmarshal(data, &payload); err != nil { + t.Fatalf("captured stdin is not valid DispatchPayload JSON: %v", err) + } + if payload.Version != DispatchPayloadVersion { + t.Errorf("Version = %d, want %d", payload.Version, DispatchPayloadVersion) + } + if payload.TaskID != "task-1" { + t.Errorf("TaskID = %q, want task-1", payload.TaskID) + } + if len(payload.BaseArgs) < 2 || payload.BaseArgs[0] != "agent" || payload.BaseArgs[1] != "run" { + t.Errorf("BaseArgs = %v, want to start with [agent run]", payload.BaseArgs) + } + if payload.Env["GITHUB_ACCESS_TOKEN"] != "secret-token" { + t.Errorf("Env[GITHUB_ACCESS_TOKEN] = %q, want secret-token", payload.Env["GITHUB_ACCESS_TOKEN"]) + } + + // The worker must have sent task_claimed but no terminal message. + msgs := drainMessages(t, w.sendChan) + var sawClaimed bool + for _, m := range msgs { + switch m.Type { + case types.MessageTypeTaskClaimed: + sawClaimed = true + case types.MessageTypeTaskCompleted, types.MessageTypeTaskFailed: + t.Fatalf("worker must not send a terminal message for a dispatched task, got %q", m.Type) + } + } + if !sawClaimed { + t.Error("expected the worker to send task_claimed") + } +} + +func TestIntegrationCommandBackendDispatchFailureReportsTaskFailed(t *testing.T) { + w := newCommandWorker(t, "exit 1") + + w.handleTaskAssignment(commandTaskAssignment()) + + waitFor(t, 5*time.Second, func() bool { return w.activeTaskCount() == 0 }) + + var failedCount int + waitFor(t, 5*time.Second, func() bool { + for _, m := range drainMessages(t, w.sendChan) { + if m.Type == types.MessageTypeTaskFailed { + failedCount++ + } + } + return failedCount > 0 + }) + if failedCount != 1 { + t.Fatalf("task_failed message count = %d, want 1", failedCount) + } +} diff --git a/internal/worker/command_test.go b/internal/worker/command_test.go new file mode 100644 index 0000000..25624f7 --- /dev/null +++ b/internal/worker/command_test.go @@ -0,0 +1,192 @@ +package worker + +import ( + "context" + "encoding/json" + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/metrics" + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +func testTaskParams() *TaskParams { + return &TaskParams{ + TaskID: "task-1", + ExecutionID: "exec-1", + Task: &types.Task{ID: "task-1", Title: "test"}, + DockerImage: "ubuntu:22.04", + BaseArgs: []string{"agent", "run", "--task-id", "task-1"}, + EnvVars: []string{"SUPER_SECRET_TOKEN=hunter2"}, + Sidecars: []types.SidecarMount{{Image: "warpdotdev/warp-agent:latest", MountPath: "/agent"}}, + } +} + +func newTestCommandBackend(t *testing.T, cfg CommandBackendConfig) *CommandBackend { + t.Helper() + if cfg.DispatchCommand == "" { + cfg.DispatchCommand = "true" + } + b, err := NewCommandBackend(context.Background(), cfg) + if err != nil { + t.Fatalf("NewCommandBackend() error = %v", err) + } + return b +} + +func assertBackendFailureReason(t *testing.T, err error, wantReason string) { + t.Helper() + var failure *backendFailureError + if !errors.As(err, &failure) { + t.Fatalf("error %v is not a *backendFailureError", err) + } + if failure.reason != wantReason { + t.Fatalf("failure reason = %q, want %q", failure.reason, wantReason) + } +} + +func TestCommandBackendDispatchesPayloadOnStdin(t *testing.T) { + outFile := filepath.Join(t.TempDir(), "payload.json") + b := newTestCommandBackend(t, CommandBackendConfig{ + DispatchCommand: "cat > " + outFile, + ServerRootURL: "https://app.warp.dev", + WorkerID: "my-worker", + }) + + err := b.ExecuteTask(context.Background(), testTaskParams()) + if !errors.Is(err, ErrTaskDispatched) { + t.Fatalf("ExecuteTask() = %v, want ErrTaskDispatched", err) + } + + data, readErr := os.ReadFile(outFile) // #nosec G304 -- test-controlled temp path. + if readErr != nil { + t.Fatalf("failed to read captured payload: %v", readErr) + } + var got DispatchPayload + if err := json.Unmarshal(data, &got); err != nil { + t.Fatalf("captured stdin is not valid DispatchPayload JSON: %v", err) + } + if got.Version != DispatchPayloadVersion { + t.Errorf("Version = %d, want %d", got.Version, DispatchPayloadVersion) + } + if got.TaskID != "task-1" { + t.Errorf("TaskID = %q, want task-1", got.TaskID) + } + if got.DockerImage != "ubuntu:22.04" { + t.Errorf("DockerImage = %q", got.DockerImage) + } + if len(got.BaseArgs) < 2 || got.BaseArgs[0] != "agent" || got.BaseArgs[1] != "run" { + t.Errorf("BaseArgs = %v, want to start with [agent run]", got.BaseArgs) + } + if got.Env["SUPER_SECRET_TOKEN"] != "hunter2" { + t.Errorf("Env[SUPER_SECRET_TOKEN] = %q, want hunter2", got.Env["SUPER_SECRET_TOKEN"]) + } + if len(got.Sidecars) != 1 || got.Sidecars[0].MountPath != "/agent" { + t.Errorf("Sidecars = %v", got.Sidecars) + } +} + +func TestCommandBackendSuccessReturnsErrTaskDispatched(t *testing.T) { + b := newTestCommandBackend(t, CommandBackendConfig{DispatchCommand: "exit 0"}) + err := b.ExecuteTask(context.Background(), testTaskParams()) + if !errors.Is(err, ErrTaskDispatched) { + t.Fatalf("ExecuteTask() = %v, want ErrTaskDispatched", err) + } +} + +func TestCommandBackendNonZeroExitIsClassifiedFailure(t *testing.T) { + b := newTestCommandBackend(t, CommandBackendConfig{DispatchCommand: "exit 7"}) + err := b.ExecuteTask(context.Background(), testTaskParams()) + if errors.Is(err, ErrTaskDispatched) { + t.Fatal("non-zero exit must not be treated as a successful dispatch") + } + assertBackendFailureReason(t, err, metrics.TaskFailureReasonDispatchCommand) +} + +func TestCommandBackendTimeoutIsClassifiedFailure(t *testing.T) { + b := newTestCommandBackend(t, CommandBackendConfig{ + DispatchCommand: "sleep 5", + DispatchTimeout: 50 * time.Millisecond, + }) + err := b.ExecuteTask(context.Background(), testTaskParams()) + if errors.Is(err, ErrTaskDispatched) { + t.Fatal("timed-out dispatch must not be treated as success") + } + assertBackendFailureReason(t, err, metrics.TaskFailureReasonDispatchTimeout) +} + +func TestCommandBackendDoesNotLeakSecretsIntoSubprocessEnv(t *testing.T) { + outFile := filepath.Join(t.TempDir(), "env.txt") + b := newTestCommandBackend(t, CommandBackendConfig{ + DispatchCommand: "env > " + outFile, + Env: map[string]string{"OPERATOR_VAR": "ok"}, + }) + + if err := b.ExecuteTask(context.Background(), testTaskParams()); !errors.Is(err, ErrTaskDispatched) { + t.Fatalf("ExecuteTask() = %v, want ErrTaskDispatched", err) + } + + data, err := os.ReadFile(outFile) // #nosec G304 -- test-controlled temp path. + if err != nil { + t.Fatalf("failed to read captured env: %v", err) + } + env := string(data) + if !strings.Contains(env, "OZ_TASK_ID=task-1") { + t.Errorf("expected OZ_TASK_ID in subprocess env, got:\n%s", env) + } + if !strings.Contains(env, "OPERATOR_VAR=ok") { + t.Errorf("expected operator env var in subprocess env") + } + if strings.Contains(env, "SUPER_SECRET_TOKEN") { + t.Errorf("task secret must NOT appear in subprocess env; got:\n%s", env) + } +} + +func TestNewCommandBackendRequiresDispatchCommand(t *testing.T) { + if _, err := NewCommandBackend(context.Background(), CommandBackendConfig{}); err == nil { + t.Fatal("expected error for empty dispatch command") + } +} + +func TestCommandBackendCancelTask(t *testing.T) { + sentinel := filepath.Join(t.TempDir(), "cancelled") + b := newTestCommandBackend(t, CommandBackendConfig{ + DispatchCommand: "true", + CancelCommand: "touch " + sentinel, + }) + + if err := b.CancelTask(context.Background(), &CancelParams{TaskID: "task-1", ExecutionID: "exec-1"}); err != nil { + t.Fatalf("CancelTask() error = %v", err) + } + if _, err := os.Stat(sentinel); err != nil { + t.Fatalf("cancel command did not run (sentinel missing): %v", err) + } +} + +func TestCommandBackendCancelTaskNoCommandIsNoop(t *testing.T) { + b := newTestCommandBackend(t, CommandBackendConfig{DispatchCommand: "true"}) + if err := b.CancelTask(context.Background(), &CancelParams{TaskID: "task-1"}); err != nil { + t.Fatalf("CancelTask() with no cancel command should be a no-op, got %v", err) + } +} + +func TestNewWorkerSelectsCommandBackend(t *testing.T) { + w, err := New(context.Background(), Config{ + BackendType: "command", + Command: &CommandBackendConfig{DispatchCommand: "true"}, + }) + if err != nil { + t.Fatalf("New() error = %v", err) + } + if _, ok := w.backend.(*CommandBackend); !ok { + t.Fatalf("backend type = %T, want *CommandBackend", w.backend) + } + + if _, err := New(context.Background(), Config{BackendType: "command"}); err == nil { + t.Fatal("expected error when command backend selected but no config provided") + } +} diff --git a/internal/worker/dispatch_payload.go b/internal/worker/dispatch_payload.go new file mode 100644 index 0000000..bef13ac --- /dev/null +++ b/internal/worker/dispatch_payload.go @@ -0,0 +1,55 @@ +package worker + +import ( + "strings" + + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +// DispatchPayloadVersion is the schema version of DispatchPayload. It is bumped +// when the payload contract changes in a way operator dispatch commands must be +// aware of. +const DispatchPayloadVersion = 1 + +// DispatchPayload is the stable, versioned JSON contract handed to an operator's +// dispatch command (on stdin) by the command backend. It contains everything a +// remote runtime needs to launch the oz agent for a task. Secrets (e.g. GitHub +// tokens) travel only inside Env here, never via the dispatch subprocess's own +// environment or argv. +type DispatchPayload struct { + Version int `json:"version"` + TaskID string `json:"task_id"` + ExecutionID string `json:"execution_id"` + ServerRootURL string `json:"server_root_url"` + WorkerID string `json:"worker_id"` + DockerImage string `json:"docker_image"` + BaseArgs []string `json:"base_args"` + Env map[string]string `json:"env"` + Sidecars []types.SidecarMount `json:"sidecars"` + Task *types.Task `json:"task"` +} + +// NewDispatchPayload builds a DispatchPayload from backend-agnostic TaskParams +// plus the worker-level server URL and worker ID. The TaskParams.EnvVars slice +// (KEY=VALUE entries) is converted into the Env map, splitting on the first '=' +// so values may themselves contain '='. Later entries win on duplicate keys. +func NewDispatchPayload(params *TaskParams, serverRootURL, workerID string) *DispatchPayload { + env := make(map[string]string, len(params.EnvVars)) + for _, entry := range params.EnvVars { + key, value, _ := strings.Cut(entry, "=") + env[key] = value + } + + return &DispatchPayload{ + Version: DispatchPayloadVersion, + TaskID: params.TaskID, + ExecutionID: params.ExecutionID, + ServerRootURL: serverRootURL, + WorkerID: workerID, + DockerImage: params.DockerImage, + BaseArgs: params.BaseArgs, + Env: env, + Sidecars: params.Sidecars, + Task: params.Task, + } +} diff --git a/internal/worker/dispatch_payload_test.go b/internal/worker/dispatch_payload_test.go new file mode 100644 index 0000000..10779a4 --- /dev/null +++ b/internal/worker/dispatch_payload_test.go @@ -0,0 +1,83 @@ +package worker + +import ( + "encoding/json" + "testing" + + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +func TestNewDispatchPayload(t *testing.T) { + task := &types.Task{ID: "task-1", Title: "do the thing"} + sidecars := []types.SidecarMount{{Image: "warpdotdev/warp-agent:latest", MountPath: "/agent"}} + params := &TaskParams{ + TaskID: "task-1", + ExecutionID: "exec-1", + Task: task, + DockerImage: "ubuntu:22.04", + BaseArgs: []string{"agent", "run", "--task-id", "task-1"}, + EnvVars: []string{"A=1", "B=2", "A=3", "URL=a=b"}, + Sidecars: sidecars, + } + + got := NewDispatchPayload(params, "https://app.warp.dev", "my-worker") + + if got.Version != DispatchPayloadVersion { + t.Errorf("Version = %d, want %d", got.Version, DispatchPayloadVersion) + } + if got.TaskID != "task-1" || got.ExecutionID != "exec-1" { + t.Errorf("identifiers = (%q, %q), want (task-1, exec-1)", got.TaskID, got.ExecutionID) + } + if got.ServerRootURL != "https://app.warp.dev" || got.WorkerID != "my-worker" { + t.Errorf("server/worker = (%q, %q)", got.ServerRootURL, got.WorkerID) + } + if got.DockerImage != "ubuntu:22.04" { + t.Errorf("DockerImage = %q", got.DockerImage) + } + if got.Task != task { + t.Errorf("Task pointer not preserved") + } + if len(got.BaseArgs) != 4 || got.BaseArgs[0] != "agent" || got.BaseArgs[1] != "run" { + t.Errorf("BaseArgs = %v", got.BaseArgs) + } + if len(got.Sidecars) != 1 || got.Sidecars[0].MountPath != "/agent" { + t.Errorf("Sidecars = %v", got.Sidecars) + } + + // Env: last-wins on duplicate keys; values may contain '='. + wantEnv := map[string]string{"A": "3", "B": "2", "URL": "a=b"} + if len(got.Env) != len(wantEnv) { + t.Fatalf("Env = %v, want %v", got.Env, wantEnv) + } + for k, v := range wantEnv { + if got.Env[k] != v { + t.Errorf("Env[%q] = %q, want %q", k, got.Env[k], v) + } + } +} + +func TestNewDispatchPayloadMarshalsSnakeCaseKeys(t *testing.T) { + params := &TaskParams{ + TaskID: "task-1", + ExecutionID: "exec-1", + Task: &types.Task{ID: "task-1"}, + EnvVars: []string{"K=v"}, + } + data, err := json.Marshal(NewDispatchPayload(params, "https://app.warp.dev", "w")) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + + var generic map[string]json.RawMessage + if err := json.Unmarshal(data, &generic); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + for _, key := range []string{ + "version", "task_id", "execution_id", "server_root_url", "worker_id", + "docker_image", "base_args", "env", "sidecars", "task", + } { + if _, ok := generic[key]; !ok { + t.Errorf("payload JSON missing key %q", key) + } + } +} diff --git a/internal/worker/dispatch_test.go b/internal/worker/dispatch_test.go new file mode 100644 index 0000000..660a276 --- /dev/null +++ b/internal/worker/dispatch_test.go @@ -0,0 +1,146 @@ +package worker + +import ( + "context" + "testing" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/types" + "go.opentelemetry.io/otel/trace" + "golang.org/x/sync/semaphore" +) + +// dispatchBackend is a fake Backend that reports a successful fire-and-forget +// dispatch. It does not implement CancelableBackend. +type dispatchBackend struct{} + +func (b *dispatchBackend) ExecuteTask(context.Context, *TaskParams) error { return ErrTaskDispatched } +func (b *dispatchBackend) Shutdown(context.Context) {} +func (b *dispatchBackend) PreservesTasksOnShutdown() bool { return true } + +// cancelableDispatchBackend is a dispatchBackend that also records CancelTask calls. +type cancelableDispatchBackend struct { + dispatchBackend + cancelCalled chan *CancelParams +} + +func (b *cancelableDispatchBackend) CancelTask(_ context.Context, params *CancelParams) error { + b.cancelCalled <- params + return nil +} + +func newDispatchWorker(backend Backend) *Worker { + return &Worker{ + ctx: context.Background(), + config: Config{}, + sendChan: make(chan []byte, 4), + activeTasks: map[string]activeTask{"task-1": {cancel: func() {}}}, + dispatchedTasks: make(map[string]*CancelParams), + backend: backend, + } +} + +func runDispatchTask(w *Worker) { + ctx := context.Background() + w.executeTask(ctx, func() {}, trace.SpanFromContext(ctx), &types.TaskAssignmentMessage{ + TaskID: "task-1", + ExecutionID: "exec-1", + Task: &types.Task{ID: "task-1", Title: "test task"}, + }, time.Now()) +} + +func TestExecuteTaskDispatchedSuppressesTerminalMessage(t *testing.T) { + w := newDispatchWorker(&dispatchBackend{}) + + runDispatchTask(w) + + if len(w.sendChan) != 0 { + msg := readWebSocketMessage(t, w.sendChan) + t.Fatalf("expected no terminal message after dispatch, got %q", msg.Type) + } + if _, ok := w.activeTasks["task-1"]; ok { + t.Error("dispatched task should be removed from active tasks") + } + w.tasksMutex.Lock() + cp, ok := w.dispatchedTasks["task-1"] + w.tasksMutex.Unlock() + if !ok { + t.Fatal("dispatched task should be registered in dispatchedTasks") + } + if cp.TaskID != "task-1" || cp.ExecutionID != "exec-1" { + t.Errorf("dispatched cancel params = %+v, want {task-1 exec-1}", cp) + } +} + +func TestExecuteTaskDispatchedReleasesSemaphore(t *testing.T) { + w := newDispatchWorker(&dispatchBackend{}) + w.config.MaxConcurrentTasks = 1 + w.taskSemaphore = semaphore.NewWeighted(1) + if !w.taskSemaphore.TryAcquire(1) { + t.Fatal("failed to acquire the only slot before dispatch") + } + + runDispatchTask(w) + + // The deferred cleanup must have released the slot so it can be re-acquired. + if !w.taskSemaphore.TryAcquire(1) { + t.Fatal("expected the concurrency slot to be released after dispatch") + } +} + +func TestHandleTaskCancellationRoutesToCancelableBackend(t *testing.T) { + backend := &cancelableDispatchBackend{cancelCalled: make(chan *CancelParams, 1)} + w := &Worker{ + ctx: context.Background(), + sendChan: make(chan []byte, 1), + activeTasks: map[string]activeTask{}, + dispatchedTasks: map[string]*CancelParams{"task-1": {TaskID: "task-1", ExecutionID: "exec-1"}}, + backend: backend, + } + + w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "task-1"}) + + select { + case cp := <-backend.cancelCalled: + if cp.TaskID != "task-1" || cp.ExecutionID != "exec-1" { + t.Errorf("CancelTask params = %+v, want {task-1 exec-1}", cp) + } + case <-time.After(2 * time.Second): + t.Fatal("CancelTask was not invoked for a dispatched task") + } + + w.tasksMutex.Lock() + _, ok := w.dispatchedTasks["task-1"] + w.tasksMutex.Unlock() + if ok { + t.Error("dispatched task should be removed after cancellation is routed") + } +} + +func TestHandleTaskCancellationDispatchedNoCancelSupportIsNoop(t *testing.T) { + w := &Worker{ + ctx: context.Background(), + sendChan: make(chan []byte, 1), + activeTasks: map[string]activeTask{}, + dispatchedTasks: map[string]*CancelParams{"task-1": {TaskID: "task-1"}}, + backend: &dispatchBackend{}, + } + + // Must not panic and must not emit any task status message. + w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "task-1"}) + + if len(w.sendChan) != 0 { + t.Fatalf("expected no message for no-cancel-support dispatched task, got %d", len(w.sendChan)) + } +} + +func TestExecuteTaskSuccessStillReportsCompleted(t *testing.T) { + w := newDispatchWorker(&recordingBackend{err: nil}) + + runDispatchTask(w) + + msg := readWebSocketMessage(t, w.sendChan) + if msg.Type != types.MessageTypeTaskCompleted { + t.Fatalf("message type = %q, want %q", msg.Type, types.MessageTypeTaskCompleted) + } +} diff --git a/internal/worker/worker.go b/internal/worker/worker.go index 6fbdfc3..30fd697 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -3,6 +3,7 @@ package worker import ( "context" "encoding/json" + "errors" "fmt" "net/url" "sync" @@ -50,6 +51,7 @@ type Config struct { Docker *DockerBackendConfig Direct *DirectBackendConfig Kubernetes *KubernetesBackendConfig + Command *CommandBackendConfig } type Worker struct { @@ -62,9 +64,14 @@ type Worker struct { lastHeartbeat time.Time sendChan chan []byte activeTasks map[string]activeTask - tasksMutex sync.Mutex - backend Backend - taskSemaphore *semaphore.Weighted // nil when unlimited + // dispatchedTasks tracks fire-and-forget tasks the backend handed off to a + // remote runtime. They are no longer executing locally, but are retained + // (by non-secret identifiers only) so a later cancellation can be routed to + // a CancelableBackend. Guarded by tasksMutex. + dispatchedTasks map[string]*CancelParams + tasksMutex sync.Mutex + backend Backend + taskSemaphore *semaphore.Weighted // nil when unlimited // heartbeatInterval is how often the worker pings the server. It defaults // to HeartbeatInterval and is overridable in tests. heartbeatInterval time.Duration @@ -100,6 +107,12 @@ func New(ctx context.Context, config Config) (*Worker, error) { return nil, fmt.Errorf("direct backend selected but no direct config provided") } backend, err = NewDirectBackend(ctx, *config.Direct) + case "command": + if config.Command == nil { + cancel() + return nil, fmt.Errorf("command backend selected but no command config provided") + } + backend, err = NewCommandBackend(ctx, *config.Command) case "docker", "": if config.Docker == nil { config.Docker = &DockerBackendConfig{} @@ -127,6 +140,7 @@ func New(ctx context.Context, config Config) (*Worker, error) { reconnectDelay: InitialReconnectDelay, sendChan: make(chan []byte, 256), activeTasks: make(map[string]activeTask), + dispatchedTasks: make(map[string]*CancelParams), backend: backend, taskSemaphore: taskSemaphore, heartbeatInterval: HeartbeatInterval, @@ -356,18 +370,55 @@ func (w *Worker) handleTaskCancellation(cancellation *types.TaskCancellationMess task.cancellationSource = taskCancellationSourceUser w.activeTasks[cancellation.TaskID] = task } + var dispatched *CancelParams + if !ok { + if dp, dok := w.dispatchedTasks[cancellation.TaskID]; dok { + dispatched = dp + delete(w.dispatchedTasks, cancellation.TaskID) + } + } w.tasksMutex.Unlock() + + if ok { + log.Infof(w.ctx, "Cancelling task from server request: taskID=%s", cancellation.TaskID) + metrics.AddTaskEvent(task.ctx, "task.cancellation_requested", + attribute.String("source", "server"), + attribute.String("task.id", cancellation.TaskID), + ) + task.cancel() + return + } + + if dispatched != nil { + w.cancelDispatchedTask(dispatched) + return + } + + log.Warnf(w.ctx, "Received cancellation for inactive task: taskID=%s", cancellation.TaskID) +} + +// cancelDispatchedTask makes a best-effort attempt to cancel a task that was +// already handed off to a remote runtime via fire-and-forget dispatch. If the +// backend cannot cancel, the worker relies on agent-side cancellation. +func (w *Worker) cancelDispatchedTask(params *CancelParams) { + cancelable, ok := w.backend.(CancelableBackend) if !ok { - log.Warnf(w.ctx, "Received cancellation for inactive task: taskID=%s", cancellation.TaskID) + log.Infof(w.ctx, "No backend cancellation for dispatched task %s; relying on agent-side cancellation", params.TaskID) return } - log.Infof(w.ctx, "Cancelling task from server request: taskID=%s", cancellation.TaskID) - metrics.AddTaskEvent(task.ctx, "task.cancellation_requested", - attribute.String("source", "server"), - attribute.String("task.id", cancellation.TaskID), - ) - task.cancel() + log.Infof(w.ctx, "Requesting backend cancellation for dispatched task %s", params.TaskID) + go func() { + ctx, cancel := context.WithTimeout(context.WithoutCancel(w.ctx), BackendShutdownTimeout) + defer cancel() + if err := cancelable.CancelTask(ctx, params); err != nil { + log.Warnf(w.ctx, "Backend cancellation failed for dispatched task %s: %v", params.TaskID, err) + metrics.AddTaskEvent(ctx, "cancel.failed", + attribute.String("reason", metrics.TaskFailureReasonCancelCommand), + attribute.String("task.id", params.TaskID), + ) + } + }() } func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage) { @@ -555,6 +606,20 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc, ) err := w.backend.ExecuteTask(ctx, params) + if errors.Is(err, ErrTaskDispatched) { + // Fire-and-forget: the backend handed the task to a remote runtime that + // owns terminal reporting to warp-server. The worker must not finalize + // the task; it only retains the identifiers needed to route a later + // cancellation. + result = metrics.TaskResultDispatched + w.tasksMutex.Lock() + w.dispatchedTasks[taskID] = &CancelParams{TaskID: taskID, ExecutionID: assignment.ExecutionID} + w.tasksMutex.Unlock() + metrics.AddTaskEvent(ctx, "task.dispatched") + span.SetStatus(codes.Ok, "task dispatched to remote runtime") + log.Infof(ctx, "Task dispatched fire-and-forget; worker will not finalize: taskID=%s", taskID) + return + } if err != nil { if ctx.Err() == context.Canceled && w.cancellationSource(taskID) == taskCancellationSourceUser { result = metrics.TaskResultCancelled @@ -760,6 +825,10 @@ func (w *Worker) Shutdown() { task.cancel() } } + // Dispatched (fire-and-forget) tasks run on the remote runtime independent + // of this worker; drop the in-memory cancellation registry without touching + // them. + clear(w.dispatchedTasks) w.tasksMutex.Unlock() if activeTaskCount > 0 && !preserveActiveTasks { diff --git a/main.go b/main.go index 1c40eef..810f631 100644 --- a/main.go +++ b/main.go @@ -26,7 +26,7 @@ var Version = "dev" var CLI struct { ConfigFile string `help:"Path to YAML config file" type:"path"` - Backend string `help:"Backend type (docker, direct, or kubernetes)" enum:"docker,direct,kubernetes," default:""` + Backend string `help:"Backend type (docker, direct, kubernetes, or command)" enum:"docker,direct,kubernetes,command," default:""` APIKey string `help:"API key for authentication" env:"WARP_API_KEY" required:""` WorkerID string `help:"Worker host identifier (required via flag or config file)"` WebSocketURL string `default:"wss://oz.warp.dev/api/v1/selfhosted/worker/ws" hidden:""` @@ -140,6 +140,8 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) { backendType = "direct" } else if fileConfig.Backend.Kubernetes != nil { backendType = "kubernetes" + } else if fileConfig.Backend.Command != nil { + backendType = "command" } else if fileConfig.Backend.Docker != nil { backendType = "docker" } @@ -299,6 +301,39 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) { Env: mergedEnv, } + case "command": + // Merge env: config file first, then CLI overlay (CLI wins on key conflict). + mergedEnv := make(map[string]string) + var dispatchCmd, cancelCmd, dispatchTimeoutStr string + if fileConfig != nil && fileConfig.Backend.Command != nil { + cc := fileConfig.Backend.Command + mergedEnv = config.ResolveEnv(cc.Environment) + dispatchCmd = cc.DispatchCommand + cancelCmd = cc.CancelCommand + dispatchTimeoutStr = cc.DispatchTimeout + } + for k, v := range cliEnv { + mergedEnv[k] = v + } + + var dispatchTimeout time.Duration + if dispatchTimeoutStr != "" { + d, err := time.ParseDuration(dispatchTimeoutStr) + if err != nil { + return worker.Config{}, fmt.Errorf("invalid backend.command.dispatch_timeout %q: %w", dispatchTimeoutStr, err) + } + dispatchTimeout = d + } + + wc.Command = &worker.CommandBackendConfig{ + DispatchCommand: dispatchCmd, + CancelCommand: cancelCmd, + DispatchTimeout: dispatchTimeout, + Env: mergedEnv, + ServerRootURL: CLI.ServerRootURL, + WorkerID: workerID, + } + default: // docker // Merge env: config file first, then CLI overlay (CLI wins on key conflict). mergedEnv := make(map[string]string) diff --git a/main_test.go b/main_test.go index 0980af2..1173153 100644 --- a/main_test.go +++ b/main_test.go @@ -172,6 +172,74 @@ func TestMergeConfigKubernetesCLIOverridesCleanupAndWorkerID(t *testing.T) { } } +func TestMergeConfigCommandFromFile(t *testing.T) { + resetCLIForTest() + t.Cleanup(resetCLIForTest) + + CLI.ServerRootURL = "https://app.warp.dev" + CLI.Env = []string{"CLI_ONLY=1"} + + fileConfig := &config.FileConfig{ + WorkerID: "command-worker", + Backend: config.BackendConfig{ + Command: &config.CommandConfig{ + DispatchCommand: "/opt/oz/dispatch.sh", + CancelCommand: "/opt/oz/cancel.sh", + DispatchTimeout: "30s", + }, + }, + } + + wc, err := mergeConfig(fileConfig) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if wc.BackendType != "command" { + t.Fatalf("BackendType = %q, want %q", wc.BackendType, "command") + } + if wc.Command == nil { + t.Fatal("expected command backend config") + } + if wc.Command.DispatchCommand != "/opt/oz/dispatch.sh" { + t.Errorf("DispatchCommand = %q, want %q", wc.Command.DispatchCommand, "/opt/oz/dispatch.sh") + } + if wc.Command.CancelCommand != "/opt/oz/cancel.sh" { + t.Errorf("CancelCommand = %q, want %q", wc.Command.CancelCommand, "/opt/oz/cancel.sh") + } + if wc.Command.DispatchTimeout != 30*time.Second { + t.Errorf("DispatchTimeout = %v, want 30s", wc.Command.DispatchTimeout) + } + if wc.Command.ServerRootURL != "https://app.warp.dev" { + t.Errorf("ServerRootURL = %q, want %q", wc.Command.ServerRootURL, "https://app.warp.dev") + } + if wc.Command.WorkerID != "command-worker" { + t.Errorf("WorkerID = %q, want %q", wc.Command.WorkerID, "command-worker") + } + if wc.Command.Env["CLI_ONLY"] != "1" { + t.Errorf("Env[CLI_ONLY] = %q, want %q", wc.Command.Env["CLI_ONLY"], "1") + } +} + +func TestMergeConfigCommandInvalidDispatchTimeout(t *testing.T) { + resetCLIForTest() + t.Cleanup(resetCLIForTest) + + fileConfig := &config.FileConfig{ + WorkerID: "command-worker", + Backend: config.BackendConfig{ + Command: &config.CommandConfig{ + DispatchCommand: "/opt/oz/dispatch.sh", + DispatchTimeout: "nope", + }, + }, + } + + if _, err := mergeConfig(fileConfig); err == nil { + t.Fatal("expected error for invalid dispatch_timeout") + } +} + func TestMergeConfigKubernetesAllowsZeroUnschedulableTimeout(t *testing.T) { resetCLIForTest() t.Cleanup(resetCLIForTest) From 0c6aec153fc99e885476ee480473f6994f4c8ac0 Mon Sep 17 00:00:00 2001 From: zachbai Date: Wed, 24 Jun 2026 21:14:56 -0700 Subject: [PATCH 2/2] Add local e2e command-backend reference (dispatch-oz-local.py) Adds a real-run reference dispatch command for local end-to-end testing: it logs/persists the forwarded DispatchPayload and launches the oz agent for real via base_args (fire-and-forget), so the command backend can be exercised against a local warp-server + session-sharing-server + oz-agent-worker (e.g. via warp-server's script/oz-local --worker-backend command). Includes CI-verified tests and README docs. Co-Authored-By: Oz --- examples/command-backend/README.md | 27 +++- examples/command-backend/dispatch-oz-local.py | 120 ++++++++++++++++++ examples/command-backend/dispatch_test.go | 61 +++++++++ examples/command-backend/doc.go | 11 +- 4 files changed, 213 insertions(+), 6 deletions(-) create mode 100755 examples/command-backend/dispatch-oz-local.py diff --git a/examples/command-backend/README.md b/examples/command-backend/README.md index 7b5d561..56d74d7 100644 --- a/examples/command-backend/README.md +++ b/examples/command-backend/README.md @@ -6,8 +6,9 @@ It is written in Python (standard library only — no dependencies to install) s ## 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. +- `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. @@ -84,3 +85,27 @@ backend: 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-.json`, then launches `$OZ_BIN ` 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 +``` + +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. diff --git a/examples/command-backend/dispatch-oz-local.py b/examples/command-backend/dispatch-oz-local.py new file mode 100755 index 0000000..2e65e4f --- /dev/null +++ b/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-.json`` for inspection, + 2. launches ``$OZ_BIN `` 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()) diff --git a/examples/command-backend/dispatch_test.go b/examples/command-backend/dispatch_test.go index 652224c..af7088b 100644 --- a/examples/command-backend/dispatch_test.go +++ b/examples/command-backend/dispatch_test.go @@ -7,8 +7,10 @@ import ( "net/http/httptest" "os" "os/exec" + "path/filepath" "strings" "testing" + "time" ) // samplePayload mirrors the worker's DispatchPayload JSON for a single task. @@ -144,3 +146,62 @@ func TestDispatchScriptRequiresURL(t *testing.T) { t.Fatal("expected non-zero exit when OZ_DISPATCH_URL is unset") } } + +func waitForFile(t *testing.T, path string) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if _, err := os.Stat(path); err == nil { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("file %q did not appear within timeout", path) +} + +func TestDispatchOzLocalLaunchesBinaryWithBaseArgs(t *testing.T) { + requirePython(t) + dir := t.TempDir() + + // Stub OZ_BIN records its argv and a forwarded env var, then exits. + invocation := filepath.Join(dir, "invocation.txt") + stub := filepath.Join(dir, "oz-stub.sh") + script := "#!/bin/sh\n{ echo \"argv:$*\"; echo \"WITH_LOCAL_SERVER=$WITH_LOCAL_SERVER\"; } > \"" + invocation + "\"\n" + if err := os.WriteFile(stub, []byte(script), 0o755); err != nil { + t.Fatalf("failed to write stub: %v", err) + } + + cmd := exec.Command("python3", "dispatch-oz-local.py") + cmd.Env = append(os.Environ(), "OZ_BIN="+stub, "OZ_LOCAL_RUN_LOG_DIR="+dir) + cmd.Stdin = strings.NewReader(`{"task_id":"task-1","base_args":["agent","run","--task-id","task-1"],"env":{"WITH_LOCAL_SERVER":"1"}}`) + if out, err := cmd.CombinedOutput(); err != nil { + t.Fatalf("dispatch-oz-local.py failed: %v\n%s", err, out) + } + + // The launched (detached) stub writes asynchronously. + waitForFile(t, invocation) + data, err := os.ReadFile(invocation) + if err != nil { + t.Fatalf("failed to read stub invocation: %v", err) + } + got := string(data) + if !strings.Contains(got, "argv:agent run --task-id task-1") { + t.Errorf("stub argv = %q, want it to include the base_args", got) + } + if !strings.Contains(got, "WITH_LOCAL_SERVER=1") { + t.Errorf("stub env = %q, want WITH_LOCAL_SERVER=1 applied from payload env", got) + } + + // The full payload must be persisted for inspection. + waitForFile(t, filepath.Join(dir, "payload-task-1.json")) +} + +func TestDispatchOzLocalRequiresOzBin(t *testing.T) { + requirePython(t) + cmd := exec.Command("python3", "dispatch-oz-local.py") + cmd.Env = []string{"PATH=" + os.Getenv("PATH")} + cmd.Stdin = strings.NewReader(`{"task_id":"t","base_args":["agent","run"]}`) + if err := cmd.Run(); err == nil { + t.Fatal("expected non-zero exit when OZ_BIN is unset") + } +} diff --git a/examples/command-backend/doc.go b/examples/command-backend/doc.go index 9c32b41..d2b4ca2 100644 --- a/examples/command-backend/doc.go +++ b/examples/command-backend/doc.go @@ -1,6 +1,7 @@ -// Package commandbackendexample holds a reference dispatch command for the -// oz-agent-worker "command" backend that transforms the task payload and -// delegates execution to an HTTP REST endpoint. The Python scripts (dispatch.py, -// cancel.py) are the actual artifacts; the Go test in this package verifies the -// reference dispatch script's transformation and exit-code contract. +// Package commandbackendexample holds reference dispatch commands for the +// oz-agent-worker "command" backend. dispatch.py / cancel.py transform the task +// payload and delegate execution to an HTTP REST endpoint; dispatch-oz-local.py +// launches a real oz agent run on the host for local end-to-end testing. The +// Python scripts are the actual artifacts; the Go test in this package verifies +// their transformation, launch behavior, and exit-code contracts. package commandbackendexample