Support passing environment variables in to agent containers

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:

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",

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
+}