diff --git a/README.md b/README.md index 0908257..a9cf962 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,30 @@ go build -o oz-agent-worker ./oz-agent-worker --api-key "wk-abc123" --worker-id "my-worker" ``` +## Environment Variables for Task Containers + +Use `-e` / `--env` to pass environment variables into task containers: + +```bash +# Explicit key=value +oz-agent-worker --api-key "wk-abc123" --worker-id "my-worker" -e MY_SECRET=hunter2 + +# Pass through from host environment +export MY_SECRET=hunter2 +oz-agent-worker --api-key "wk-abc123" --worker-id "my-worker" -e MY_SECRET + +# Multiple variables +oz-agent-worker --api-key "wk-abc123" --worker-id "my-worker" -e FOO=bar -e BAZ=qux +``` + +When using Docker to run the worker, note that `-e` flags for the worker itself (task containers) are passed as arguments, while `-e` flags for the worker container use Docker's syntax: + +```bash +docker run -v /var/run/docker.sock:/var/run/docker.sock \ + -e WARP_API_KEY="wk-abc123" \ + warpdotdev/oz-agent-worker --worker-id "my-worker" -e MY_SECRET=hunter2 +``` + ## Docker Connectivity The worker automatically discovers the Docker daemon using standard Docker client mechanisms, in this order: diff --git a/internal/worker/worker.go b/internal/worker/worker.go index d60e49d..93902b4 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -43,6 +43,7 @@ type Config struct { LogLevel string NoCleanup bool Volumes []string + Env map[string]string } type Worker struct { @@ -502,6 +503,11 @@ func (w *Worker) executeTaskInDocker(ctx context.Context, assignment *types.Task envVars = append(envVars, fmt.Sprintf("%s=%s", key, value)) } + // Append user-specified CLI env vars last so they take precedence. + for key, value := range w.config.Env { + envVars = append(envVars, fmt.Sprintf("%s=%s", key, value)) + } + cmd := []string{ "/bin/sh", "/agent/entrypoint.sh", diff --git a/main.go b/main.go index 9fcef05..d43ae16 100644 --- a/main.go +++ b/main.go @@ -2,6 +2,7 @@ package main import ( "context" + "fmt" "os" "os/signal" "strings" @@ -20,6 +21,7 @@ var CLI struct { LogLevel string `help:"Log level (debug, info, warn, error)" default:"info" enum:"debug,info,warn,error"` NoCleanup bool `help:"Do not remove containers after execution (for debugging)"` Volumes []string `help:"Volume mounts for task containers (format: HOST_PATH:CONTAINER_PATH or HOST_PATH:CONTAINER_PATH:MODE)" short:"v"` + Env []string `help:"Environment variables for task containers (format: KEY=VALUE or KEY to pass through from host)" short:"e"` } func main() { @@ -38,6 +40,11 @@ func main() { log.SetLevel(CLI.LogLevel) + envMap, err := parseEnvFlags(CLI.Env) + if err != nil { + log.Fatalf(ctx, "%v", err) + } + config := worker.Config{ APIKey: CLI.APIKey, WorkerID: CLI.WorkerID, @@ -46,6 +53,7 @@ func main() { LogLevel: CLI.LogLevel, NoCleanup: CLI.NoCleanup, Volumes: CLI.Volumes, + Env: envMap, } w, err := worker.New(ctx, config) @@ -72,3 +80,30 @@ func main() { log.Infof(ctx, "Worker shutdown complete") } + +// parseEnvFlags parses -e/--env flag values into a map. +// "KEY=VALUE" is used as-is; bare "KEY" inherits from the host environment. +// Empty keys and keys containing whitespace are rejected. +func parseEnvFlags(raw []string) (map[string]string, error) { + result := make(map[string]string, len(raw)) + for _, entry := range raw { + if entry == "" { + return nil, fmt.Errorf("invalid --env flag: empty value") + } + + key, value, hasEquals := strings.Cut(entry, "=") + if key == "" { + return nil, fmt.Errorf("invalid --env flag: missing key in %q", entry) + } + if strings.ContainsAny(key, " \t") { + return nil, fmt.Errorf("invalid --env flag: key contains whitespace in %q", entry) + } + + if hasEquals { + result[key] = value + } else { + result[key] = os.Getenv(key) + } + } + return result, nil +}