Adding customizable idle-on-complete for self hosted worker (#35)

Co-authored-by: Ben Navetta <ben@warp.dev>

Author: ianhodge

internal/common/task_utils.go

@@ -7,9 +7,17 @@ import (
 	"github.com/warpdotdev/oz-agent-worker/internal/types"
 )
 
+// TaskAugmentOptions contains worker-level settings that are translated into oz CLI flags
+// for every task. Add new per-worker CLI overrides here rather than as extra parameters.
+type TaskAugmentOptions struct {
+	// IdleOnComplete is passed to --idle-on-complete. Empty string uses the oz CLI default
+	// (45m). Use "0s" to exit immediately after the conversation finishes.
+	IdleOnComplete string
+}
+
 // AugmentArgsForTask allows different task sources to add CLI args in a centralized place.
 // Uses task.AgentConfigSnapshot as the source of truth when available.
-func AugmentArgsForTask(task *types.Task, args []string) []string {
+func AugmentArgsForTask(task *types.Task, args []string, opts TaskAugmentOptions) []string {
 	if task == nil {
 		return args
 	}
@@ -57,8 +65,14 @@ func AugmentArgsForTask(task *types.Task, args []string) []string {
 		}
 	}
 
-	// Keep the agent alive after task completion to allow followups.
-	args = append(args, "--idle-on-complete")
+	// Keep the agent alive after task completion to allow follow-ups.
+	// If no duration is configured, pass the flag without a value so the oz CLI
+	// uses its default of 45 minutes.
+	if opts.IdleOnComplete == "" {
+		args = append(args, "--idle-on-complete")
+	} else {
+		args = append(args, "--idle-on-complete", opts.IdleOnComplete)
+	}
 
 	return args
 }

internal/config/config.go

@@ -14,10 +14,15 @@ import (
 
 // FileConfig represents the top-level YAML configuration file.
 type FileConfig struct {
-	WorkerID           string        `yaml:"worker_id"`
-	Cleanup            *bool         `yaml:"cleanup"`
-	MaxConcurrentTasks *int          `yaml:"max_concurrent_tasks"`
-	Backend            BackendConfig `yaml:"backend"`
+	WorkerID           string `yaml:"worker_id"`
+	Cleanup            *bool  `yaml:"cleanup"`
+	MaxConcurrentTasks *int   `yaml:"max_concurrent_tasks"`
+	// IdleOnComplete controls how long the oz CLI process stays alive after a task's
+	// conversation finishes, to allow follow-up interactions via the shared session.
+	// Uses humantime format (e.g. "45m", "10m", "0s"). When nil, the oz CLI default
+	// of 45 minutes is used.
+	IdleOnComplete *string       `yaml:"idle_on_complete"`
+	Backend        BackendConfig `yaml:"backend"`
 }
 
 // BackendConfig contains the backend selection.

internal/config/config_test.go

@@ -257,6 +257,38 @@ func TestLoadFileNotFound(t *testing.T) {
 	}
 }
 
+func TestLoadIdleOnComplete(t *testing.T) {
+	t.Run("parses idle_on_complete when set", func(t *testing.T) {
+		path := writeTestConfig(t, `
+worker_id: "test"
+idle_on_complete: "10m"
+`)
+		cfg, err := Load(path)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if cfg.IdleOnComplete == nil {
+			t.Fatal("expected idle_on_complete to be set")
+		}
+		if *cfg.IdleOnComplete != "10m" {
+			t.Errorf("idle_on_complete = %q, want %q", *cfg.IdleOnComplete, "10m")
+		}
+	})
+
+	t.Run("idle_on_complete is nil when not set", func(t *testing.T) {
+		path := writeTestConfig(t, `
+worker_id: "test"
+`)
+		cfg, err := Load(path)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if cfg.IdleOnComplete != nil {
+			t.Errorf("expected idle_on_complete to be nil, got %q", *cfg.IdleOnComplete)
+		}
+	})
+}
+
 func TestLoadMaxConcurrentTasks(t *testing.T) {
 	t.Run("parses max_concurrent_tasks when set", func(t *testing.T) {
 		path := writeTestConfig(t, `

internal/worker/worker.go

@@ -33,6 +33,9 @@ type Config struct {
 	LogLevel           string
 	BackendType        string // "docker" or "direct"
 	MaxConcurrentTasks int    // 0 means unlimited
+	// IdleOnComplete is passed to the oz CLI's --idle-on-complete flag for every task.
+	// Empty string means use the oz CLI default (45m). Use "0s" to disable idle.
+	IdleOnComplete string
 
 	// Backend-specific configs. Only the one matching BackendType should be set.
 	Docker *DockerBackendConfig
@@ -371,7 +374,9 @@ func (w *Worker) prepareTaskParams(assignment *types.TaskAssignmentMessage) *Tas
 		"--server-root-url",
 		w.config.ServerRootURL,
 	}
-	baseArgs = common.AugmentArgsForTask(task, baseArgs)
+	baseArgs = common.AugmentArgsForTask(task, baseArgs, common.TaskAugmentOptions{
+		IdleOnComplete: w.config.IdleOnComplete,
+	})
 
 	// Build a unified sidecar list: the Warp agent sidecar (mounted at /agent, where
 	// entrypoint.sh lives) comes first, followed by any additional sidecars.

main.go

@@ -26,6 +26,7 @@ var CLI struct {
 	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"`
 	MaxConcurrentTasks int      `help:"Maximum number of tasks to run concurrently (0 for unlimited)" default:"0"`
+	IdleOnComplete     string   `help:"How long to keep the oz agent alive after a task completes, for follow-ups (e.g. 45m, 10m, 0s). Defaults to 45m when not set."`
 }
 
 func main() {
@@ -123,6 +124,12 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) {
 		maxConcurrentTasks = *fileConfig.MaxConcurrentTasks
 	}
 
+	// Resolve idle_on_complete: CLI (non-empty) > config file > "" (oz CLI default = 45m).
+	idleOnComplete := CLI.IdleOnComplete
+	if idleOnComplete == "" && fileConfig != nil && fileConfig.IdleOnComplete != nil {
+		idleOnComplete = *fileConfig.IdleOnComplete
+	}
+
 	wc := worker.Config{
 		APIKey:             CLI.APIKey,
 		WorkerID:           workerID,
@@ -131,6 +138,7 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) {
 		LogLevel:           CLI.LogLevel,
 		BackendType:        backendType,
 		MaxConcurrentTasks: maxConcurrentTasks,
+		IdleOnComplete:     idleOnComplete,
 	}
 
 	switch backendType {