Use task-level config.idle_timeout_minutes for VM timeout

internal/common/task_utils.go

@@ -2,6 +2,7 @@ package common
 
 import (
 	"encoding/json"
+	"fmt"
 	"strings"
 
 	"github.com/warpdotdev/oz-agent-worker/internal/types"
@@ -12,6 +13,7 @@ import (
 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.
+	// Task-level config.idle_timeout_minutes takes precedence when set.
 	IdleOnComplete string
 }
 
@@ -66,13 +68,28 @@ func AugmentArgsForTask(task *types.Task, args []string, opts TaskAugmentOptions
 	}
 
 	// 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 == "" {
+	// Priority: task config idle_timeout_minutes > worker IdleOnComplete > oz CLI default (45m).
+	idleOnComplete, hasIdleOnCompleteValue := resolveIdleOnComplete(task, opts)
+	if !hasIdleOnCompleteValue {
 		args = append(args, "--idle-on-complete")
 	} else {
-		args = append(args, "--idle-on-complete", opts.IdleOnComplete)
+		args = append(args, "--idle-on-complete", idleOnComplete)
 	}
 
 	return args
 }
+
+func resolveIdleOnComplete(task *types.Task, opts TaskAugmentOptions) (string, bool) {
+	if task != nil &&
+		task.AgentConfigSnapshot != nil &&
+		task.AgentConfigSnapshot.IdleTimeoutMinutes != nil &&
+		*task.AgentConfigSnapshot.IdleTimeoutMinutes > 0 {
+		return fmt.Sprintf("%dm", *task.AgentConfigSnapshot.IdleTimeoutMinutes), true
+	}
+
+	if opts.IdleOnComplete != "" {
+		return opts.IdleOnComplete, true
+	}
+
+	return "", false
+}

internal/common/task_utils_test.go

@@ -0,0 +1,79 @@
+package common
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/warpdotdev/oz-agent-worker/internal/types"
+)
+
+func strPtr(v string) *string { return &v }
+func intPtr(v int) *int       { return &v }
+
+func TestAugmentArgsForTask_IdleOnCompletePrecedence(t *testing.T) {
+	baseArgs := []string{"agent", "run"}
+
+	tests := []struct {
+		name     string
+		task     *types.Task
+		opts     TaskAugmentOptions
+		expected []string
+	}{
+		{
+			name: "uses task idle_timeout_minutes when set",
+			task: &types.Task{
+				AgentConfigSnapshot: &types.AmbientAgentConfig{
+					IdleTimeoutMinutes: intPtr(15),
+				},
+			},
+			opts:     TaskAugmentOptions{IdleOnComplete: "30m"},
+			expected: []string{"agent", "run", "--idle-on-complete", "15m"},
+		},
+		{
+			name: "falls back to worker idle_on_complete when task timeout not set",
+			task: &types.Task{
+				AgentConfigSnapshot: &types.AmbientAgentConfig{},
+			},
+			opts:     TaskAugmentOptions{IdleOnComplete: "30m"},
+			expected: []string{"agent", "run", "--idle-on-complete", "30m"},
+		},
+		{
+			name: "uses oz cli default when neither task nor worker timeout is set",
+			task: &types.Task{
+				AgentConfigSnapshot: &types.AmbientAgentConfig{},
+			},
+			opts:     TaskAugmentOptions{},
+			expected: []string{"agent", "run", "--idle-on-complete"},
+		},
+		{
+			name: "ignores non-positive task idle_timeout_minutes and falls back to worker value",
+			task: &types.Task{
+				AgentConfigSnapshot: &types.AmbientAgentConfig{
+					IdleTimeoutMinutes: intPtr(0),
+				},
+			},
+			opts:     TaskAugmentOptions{IdleOnComplete: "20m"},
+			expected: []string{"agent", "run", "--idle-on-complete", "20m"},
+		},
+		{
+			name: "still appends other config-derived args before idle timeout",
+			task: &types.Task{
+				AgentConfigSnapshot: &types.AmbientAgentConfig{
+					ModelID:            strPtr("claude-sonnet-4"),
+					IdleTimeoutMinutes: intPtr(12),
+				},
+			},
+			opts:     TaskAugmentOptions{},
+			expected: []string{"agent", "run", "--model", "claude-sonnet-4", "--idle-on-complete", "12m"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := AugmentArgsForTask(tt.task, append([]string{}, baseArgs...), tt.opts)
+			if !reflect.DeepEqual(got, tt.expected) {
+				t.Fatalf("args mismatch\n got: %#v\nwant: %#v", got, tt.expected)
+			}
+		})
+	}
+}

internal/config/config.go

@@ -21,6 +21,9 @@ type FileConfig struct {
 	// 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.
+	// TODO: Remove idle_on_complete from worker config/schema after task-level
+	// config.idle_timeout_minutes is fully rolled out and legacy worker-level
+	// overrides are no longer needed.
 	IdleOnComplete *string       `yaml:"idle_on_complete"`
 	Backend        BackendConfig `yaml:"backend"`
 }

internal/types/messages.go

@@ -74,6 +74,7 @@ type AmbientAgentConfig struct {
 	SkillSpec          *string                    `json:"skill_spec,omitempty"`
 	MCPServers         map[string]json.RawMessage `json:"mcp_servers,omitempty"`
 	ComputerUseEnabled *bool                      `json:"computer_use_enabled,omitempty"`
+	IdleTimeoutMinutes *int                       `json:"idle_timeout_minutes,omitempty"`
 }
 
 // Task represents an ambient agent job.

internal/worker/direct.go

@@ -128,7 +128,11 @@ func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) err
 	if err := envFile.Close(); err != nil {
 		return fmt.Errorf("failed to close environment file: %w", err)
 	}
-	defer os.Remove(envFilePath)
+	defer func() {
+		if err := os.Remove(envFilePath); err != nil && !os.IsNotExist(err) {
+			log.Warnf(ctx, "Failed to remove environment file %s: %v", envFilePath, err)
+		}
+	}()
 
 	// 3. Build environment variables: common + config-level.
 	envVars := make([]string, len(params.EnvVars))

internal/worker/kubernetes.go

@@ -503,7 +503,7 @@ func workspaceVolume(sizeLimit *resource.Quantity) corev1.Volume {
 	}
 	if sizeLimit != nil {
 		copy := sizeLimit.DeepCopy()
-		volume.VolumeSource.EmptyDir.SizeLimit = &copy
+		volume.EmptyDir.SizeLimit = &copy
 	}
 	return volume
 }

internal/worker/kubernetes_test.go

@@ -219,8 +219,8 @@ func TestWatchJobReturnsWatchInterface(t *testing.T) {
 	// The fake client's watch should be open (channel not closed).
 	select {
 	case _, ok := <-watcher.ResultChan():
-		if ok {
-			// Got an event — fine for the fake client.
+		if !ok {
+			t.Fatal("watch channel unexpectedly closed")
 		}
 	case <-time.After(50 * time.Millisecond):
 		// No events yet — expected for an empty cluster.