diff --git a/internal/common/task_utils.go b/internal/common/task_utils.go index cc57ee3..d3f7bab 100644 --- a/internal/common/task_utils.go +++ b/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 { diff --git a/internal/common/task_utils_test.go b/internal/common/task_utils_test.go index 415d465..0e9f6e1 100644 --- a/internal/common/task_utils_test.go +++ b/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 { diff --git a/internal/types/messages.go b/internal/types/messages.go index cb98d46..7a3d818 100644 --- a/internal/types/messages.go +++ b/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.