diff --git a/internal/common/task_utils.go b/internal/common/task_utils.go index 1cd20c9..b89a8d5 100644 --- a/internal/common/task_utils.go +++ b/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 +} diff --git a/internal/common/task_utils_test.go b/internal/common/task_utils_test.go new file mode 100644 index 0000000..3348a66 --- /dev/null +++ b/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) + } + }) + } +} diff --git a/internal/config/config.go b/internal/config/config.go index 54e68e4..cf4382f 100644 --- a/internal/config/config.go +++ b/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"` } diff --git a/internal/types/messages.go b/internal/types/messages.go index d8e622f..e11aa9d 100644 --- a/internal/types/messages.go +++ b/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. diff --git a/internal/worker/direct.go b/internal/worker/direct.go index 38d14c4..6ba458b 100644 --- a/internal/worker/direct.go +++ b/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)) diff --git a/internal/worker/kubernetes.go b/internal/worker/kubernetes.go index 7495e3d..e7b0ec2 100644 --- a/internal/worker/kubernetes.go +++ b/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 = © + volume.EmptyDir.SizeLimit = © } return volume } diff --git a/internal/worker/kubernetes_test.go b/internal/worker/kubernetes_test.go index a07e7e6..efd58d7 100644 --- a/internal/worker/kubernetes_test.go +++ b/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.