add snapshot flags for oz worker (#56)

Author: harryalbert

internal/common/task_utils.go

@@ -86,6 +86,16 @@ func AugmentArgsForTask(task *types.Task, args []string, opts TaskAugmentOptions
 				args = append(args, "--share", fmt.Sprintf("public:%s", level))
 			}
 		}
+
+		if task.AgentConfigSnapshot.SnapshotDisabled != nil && *task.AgentConfigSnapshot.SnapshotDisabled {
+			args = append(args, "--no-snapshot")
+		}
+		if task.AgentConfigSnapshot.SnapshotUploadTimeoutSecs != nil && *task.AgentConfigSnapshot.SnapshotUploadTimeoutSecs > 0 {
+			args = append(args, "--snapshot-upload-timeout", fmt.Sprintf("%ds", *task.AgentConfigSnapshot.SnapshotUploadTimeoutSecs))
+		}
+		if task.AgentConfigSnapshot.SnapshotScriptTimeoutSecs != nil && *task.AgentConfigSnapshot.SnapshotScriptTimeoutSecs > 0 {
+			args = append(args, "--snapshot-script-timeout", fmt.Sprintf("%ds", *task.AgentConfigSnapshot.SnapshotScriptTimeoutSecs))
+		}
 	}
 
 	if task.AgentConfigSnapshot != nil && task.AgentConfigSnapshot.EnvironmentID != nil {

internal/common/task_utils_test.go

@@ -9,6 +9,7 @@ import (
 
 func strPtr(v string) *string                          { return &v }
 func intPtr(v int) *int                                { return &v }
+func boolPtr(v bool) *bool                             { return &v }
 func accessPtr(v types.AccessLevel) *types.AccessLevel { return &v }
 
 func TestAugmentArgsForTask_IdleOnCompletePrecedence(t *testing.T) {
@@ -170,6 +171,24 @@ func TestAugmentArgsForTask_IdleOnCompletePrecedence(t *testing.T) {
 			opts:     TaskAugmentOptions{},
 			expected: []string{"agent", "run", "--idle-on-complete"},
 		},
+		{
+			name: "adds snapshot controls when configured",
+			task: &types.Task{
+				AgentConfigSnapshot: &types.AmbientAgentConfig{
+					SnapshotDisabled:          boolPtr(true),
+					SnapshotUploadTimeoutSecs: intPtr(90),
+					SnapshotScriptTimeoutSecs: intPtr(45),
+				},
+			},
+			opts: TaskAugmentOptions{},
+			expected: []string{
+				"agent", "run",
+				"--no-snapshot",
+				"--snapshot-upload-timeout", "90s",
+				"--snapshot-script-timeout", "45s",
+				"--idle-on-complete",
+			},
+		},
 	}
 
 	for _, tt := range tests {

internal/types/messages.go

@@ -107,18 +107,21 @@ type SessionSharingConfig struct {
 
 // AmbientAgentConfig represents the agent configuration.
 type AmbientAgentConfig struct {
-	EnvironmentID        *string                    `json:"environment_id,omitempty"`
-	BasePrompt           *string                    `json:"base_prompt,omitempty"`
-	ModelID              *string                    `json:"model_id,omitempty"`
-	ProfileID            *string                    `json:"profile_id,omitempty"`
-	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"`
-	Harness              *Harness                   `json:"harness,omitempty"`
-	HarnessAuthSecrets   *HarnessAuthSecrets        `json:"harness_auth_secrets,omitempty"`
-	BedrockInferenceRole *string                    `json:"bedrock_inference_role,omitempty"`
-	SessionSharing       *SessionSharingConfig      `json:"session_sharing,omitempty"`
+	EnvironmentID             *string                    `json:"environment_id,omitempty"`
+	BasePrompt                *string                    `json:"base_prompt,omitempty"`
+	ModelID                   *string                    `json:"model_id,omitempty"`
+	ProfileID                 *string                    `json:"profile_id,omitempty"`
+	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"`
+	Harness                   *Harness                   `json:"harness,omitempty"`
+	HarnessAuthSecrets        *HarnessAuthSecrets        `json:"harness_auth_secrets,omitempty"`
+	BedrockInferenceRole      *string                    `json:"bedrock_inference_role,omitempty"`
+	SessionSharing            *SessionSharingConfig      `json:"session_sharing,omitempty"`
+	SnapshotDisabled          *bool                      `json:"snapshot_disabled,omitempty"`
+	SnapshotUploadTimeoutSecs *int                       `json:"snapshot_upload_timeout_secs,omitempty"`
+	SnapshotScriptTimeoutSecs *int                       `json:"snapshot_script_timeout_secs,omitempty"`
 }
 
 // Task represents an ambient agent job.