air: flatten code_source configuration (#6927)

## Changes
- Allow snapshot fields directly under `code_source` without requiring
`type` or a nested `snapshot` block.
- Preserve legacy configurations and overrides while normalizing
supported forms to one internal path.
- Hide legacy fields from help and reject mixed forms with actionable
errors.

## Why
The snapshot discriminator and wrapper are redundant now that snapshot
is the only supported code source type. Flattening the configuration
simplifies workload YAML and keeps future snapshot fields on one
implementation path.

## Tests
- `./task checks`
- `go test ./cmd/air`
- `go test ./acceptance -run TestAccept/air/config-help -tail`
- `go vet ./cmd/air`
- Built the PR CLI and successfully ran the AIR hello-world and a
one-GPU Ray Core cookbook example with flattened `code_source`,
including snapshot upload and extraction.
- `./task test` ran 11,580 tests; affected AIR tests passed, but an
unrelated `libs/dyn/jsonloader` assertion fails on changed Go JSON error
text.
- Full lint is blocked by a golangci-lint/Staticcheck `buildir` panic in
dependency package `poll`.

_This PR was written by Codex._

---------

Co-authored-by: Caroline Chen <324939130+caroline-db@users.noreply.github.com>

Author: caroline-db

.nextchanges/air/code-source-flattening.md

@@ -0,0 +1 @@
+* Support snapshot fields directly under `code_source` without requiring `type` or a nested `snapshot` block. ([#6927](https://github.com/databricks/cli/pull/6927))

acceptance/air/config-help/output.txt

@@ -115,29 +115,34 @@ config.compute.accelerator_type
   Type:     string
   Required: no
 
-=== conditionally required leaf
->>> [CLI] air run -h config.code_source.type
-config.code_source.type
-  Kind of code source. Must be "snapshot", the only supported type.
+=== code source fields
+>>> [CLI] air run -h config.code_source
+config.code_source
+  Local code to upload and make available to the run.
 
-  Type:     string
-  Required: when code_source is set
+  Fields:
+    root_path      Root of the code source to archive.
+    remote_volume  Volume to upload the archive to.
+    git            Pin the snapshot to a specific git revision.
+    include_paths  Restrict the archive to these paths, relative to root_path and without "..".
 
-=== union field reports both accepted shapes
->>> [CLI] air run -h config.code_source.snapshot.git.remote
-config.code_source.snapshot.git.remote
+Use "-h config.code_source.<field>" for details on a field.
+
+=== union field
+>>> [CLI] air run -h config.code_source.git.remote
+config.code_source.git.remote
   No longer supported: the snapshot archives your local copy. Only false is accepted; use commit to pin a revision.
 
   Type:     bool or string
   Required: no
 
-=== git-pinned subdirectory snapshots document root_path scoping
->>> [CLI] air run -h config.code_source.snapshot.root_path
-config.code_source.snapshot.root_path
+=== git-pinned subdirectory documents root_path scoping
+>>> [CLI] air run -h config.code_source.root_path
+config.code_source.root_path
   Root of the code source to archive. A git-pinned subdirectory packages only that subtree.
 
   Type:     string
-  Required: when code_source.snapshot is set
+  Required: when code_source is set
 
 === the config. prefix is optional
 >>> [CLI] air run -h compute.num_accelerators

acceptance/air/config-help/script

@@ -22,14 +22,14 @@ trace $CLI air run -h config.environment.version
 title "leaf field"
 trace $CLI air run -h config.compute.accelerator_type
 
-title "conditionally required leaf"
-trace $CLI air run -h config.code_source.type
+title "code source fields"
+trace $CLI air run -h config.code_source
 
-title "union field reports both accepted shapes"
-trace $CLI air run -h config.code_source.snapshot.git.remote
+title "union field"
+trace $CLI air run -h config.code_source.git.remote
 
-title "git-pinned subdirectory snapshots document root_path scoping"
-trace $CLI air run -h config.code_source.snapshot.root_path
+title "git-pinned subdirectory documents root_path scoping"
+trace $CLI air run -h config.code_source.root_path
 
 title "the config. prefix is optional"
 trace $CLI air run -h compute.num_accelerators

cmd/air/convert_to_dabs.go

@@ -130,7 +130,7 @@ func convertToDabs(ctx context.Context, cfg *runConfig, configPath, bundleDir st
 		// artifact location is set bundle-wide via workspace.artifact_path, not
 		// per-code-source, so a per-source Volume isn't representable here.
 		if snap.RemoteVolume != nil {
-			return nil, nil, errors.New("code_source.snapshot.remote_volume is not supported by convert-to-dabs; set workspace.artifact_path in the bundle instead")
+			return nil, nil, errors.New("remote_volume in code_source is not supported by convert-to-dabs; set workspace.artifact_path in the bundle instead")
 		}
 	}
 

cmd/air/convert_to_dabs_test.go

@@ -374,7 +374,7 @@ code_source:
 	loaded, err := loadRunConfig(path)
 	require.NoError(t, err)
 	_, _, err = convertToDabs(t.Context(), loaded, path, filepath.Dir(path))
-	require.ErrorContains(t, err, "remote_volume is not supported")
+	require.ErrorContains(t, err, "remote_volume in code_source is not supported")
 }
 
 // env_variables / secrets / parameters ride as sidecar files (the ai_runtime_task

cmd/air/runconfig.go

@@ -421,52 +421,72 @@ func (s *stringOrInt) UnmarshalYAML(node *yaml.Node) error {
 
 // codeSourceConfig is the `code_source` block. Only the "snapshot" type exists.
 type codeSourceConfig struct {
-	Type     string                `yaml:"type" help:"Kind of code source. Must be \"snapshot\", the only supported type." required:"when code_source is set"`
-	Snapshot *snapshotSourceConfig `yaml:"snapshot" help:"Which local directory to archive and upload." required:"when code_source.type is \"snapshot\""`
+	Type           *string               `yaml:"type,omitempty" hidden:"yes"`
+	Snapshot       *snapshotSourceConfig `yaml:",inline"`
+	LegacySnapshot *snapshotSourceConfig `yaml:"snapshot,omitempty" hidden:"yes"`
 }
 
 func (c *codeSourceConfig) validate() error {
-	if c.Type != "snapshot" {
-		return fmt.Errorf("code_source.type must be 'snapshot', got %q", c.Type)
+	if c.Type != nil && strings.TrimSpace(*c.Type) == "" {
+		return errors.New("code_source.type cannot be empty; omit it to use snapshot")
 	}
-	if c.Snapshot == nil {
-		return errors.New("code_source.type='snapshot' requires a snapshot configuration")
+	if c.Type != nil && *c.Type != "snapshot" {
+		return fmt.Errorf("code_source.type must be 'snapshot', got %q", *c.Type)
 	}
-	return c.Snapshot.validate()
+	if c.Snapshot != nil && c.LegacySnapshot != nil {
+		return errors.New("code_source cannot combine the nested 'snapshot' block with direct fields; remove the 'snapshot' wrapper and place all fields directly under code_source")
+	}
+
+	snapshot := c.Snapshot
+	prefix := "code_source"
+	if snapshot == nil {
+		snapshot = c.LegacySnapshot
+		prefix = "code_source.snapshot"
+	}
+	if snapshot == nil {
+		return errors.New("code_source requires snapshot configuration with root_path")
+	}
+	if err := snapshot.validate(prefix); err != nil {
+		return err
+	}
+
+	c.Snapshot = snapshot
+	c.LegacySnapshot = nil
+	return nil
 }
 
 // snapshotSourceConfig describes a local directory to tar and upload.
 type snapshotSourceConfig struct {
-	RootPath     string   `yaml:"root_path" help:"Root of the code source to archive. A git-pinned subdirectory packages only that subtree." required:"when code_source.snapshot is set"`
+	RootPath     string   `yaml:"root_path" help:"Root of the code source to archive. A git-pinned subdirectory packages only that subtree." required:"when code_source is set"`
 	RemoteVolume *string  `yaml:"remote_volume" help:"Volume to upload the archive to. Must start with /Volumes/."`
 	Git          *gitRef  `yaml:"git" help:"Pin the snapshot to a specific git revision."`
 	IncludePaths []string `yaml:"include_paths" help:"Restrict the archive to these paths, relative to root_path and without \"..\". Omit to include everything."`
 }
 
-func (s *snapshotSourceConfig) validate() error {
+func (s *snapshotSourceConfig) validate(prefix string) error {
 	if strings.TrimSpace(s.RootPath) == "" {
-		return errors.New("code_source.snapshot.root_path cannot be empty")
+		return fmt.Errorf("%s.root_path cannot be empty", prefix)
 	}
 
 	if s.RemoteVolume != nil && !strings.HasPrefix(*s.RemoteVolume, "/Volumes/") {
-		return errors.New("code_source.snapshot.remote_volume must start with '/Volumes/'")
+		return fmt.Errorf("%s.remote_volume must start with '/Volumes/'", prefix)
 	}
 
 	// A non-nil but empty include_paths is an explicit mistake (omit it instead).
 	if s.IncludePaths != nil && len(s.IncludePaths) == 0 {
-		return errors.New("code_source.snapshot.include_paths cannot be an empty list; either omit it or provide paths")
+		return fmt.Errorf("%s.include_paths cannot be an empty list; either omit it or provide paths", prefix)
 	}
 	for _, p := range s.IncludePaths {
 		p = strings.TrimSpace(p)
 		if p == "" {
-			return errors.New("code_source.snapshot.include_paths entry cannot be empty")
+			return fmt.Errorf("%s.include_paths entry cannot be empty", prefix)
 		}
 		if strings.HasPrefix(p, "/") {
-			return fmt.Errorf("code_source.snapshot.include_paths must be relative paths, got: %s", p)
+			return fmt.Errorf("%s.include_paths must be relative paths, got: %s", prefix, p)
 		}
 		// No parent traversal: snapshots must stay within root_path.
 		if slices.Contains(strings.Split(p, "/"), "..") {
-			return fmt.Errorf("code_source.snapshot.include_paths cannot contain '..' traversal, got: %s", p)
+			return fmt.Errorf("%s.include_paths cannot contain '..' traversal, got: %s", prefix, p)
 		}
 	}
 
@@ -616,6 +636,7 @@ type configField struct {
 	typeName string
 	required string
 	help     string
+	hidden   bool
 	// freeForm marks a user-keyed map (parameters/secrets/env_variables): no
 	// children, yet any sub-path into it is valid.
 	freeForm bool
@@ -651,7 +672,7 @@ func resolveConfigField(path string) (configField, error) {
 			return configField{}, fmt.Errorf("%q is not an object, so it has no field %q", current.path, part)
 		}
 		child, ok := findConfigChild(current.children, part)
-		if !ok {
+		if !ok || child.hidden {
 			return configField{}, unknownConfigFieldError(current, part, strings.Split(trimmed, ".")[:i+1])
 		}
 		current = child
@@ -692,6 +713,9 @@ func cutLast(s, sep string) (before, after string, found bool) {
 func unknownConfigFieldError(parent configField, part string, matched []string) error {
 	names := make([]string, 0, len(parent.children))
 	for _, c := range parent.children {
+		if c.hidden {
+			continue
+		}
 		names = append(names, configLeafName(c.path))
 	}
 	slices.Sort(names)
@@ -749,7 +773,13 @@ func describeStruct(t reflect.Type, prefix string) []configField {
 		if tag == "" || tag == "-" {
 			continue
 		}
-		name, _, _ := strings.Cut(tag, ",")
+		name, options, _ := strings.Cut(tag, ",")
+		if slices.Contains(strings.Split(options, ","), "inline") {
+			if nested := underlyingConfigStruct(f.Type); nested != nil {
+				out = append(out, describeStruct(nested, prefix)...)
+			}
+			continue
+		}
 		if name == "" || name == "-" {
 			continue
 		}
@@ -759,6 +789,7 @@ func describeStruct(t reflect.Type, prefix string) []configField {
 			typeName: configTypeName(f.Type),
 			required: f.Tag.Get("required"),
 			help:     f.Tag.Get("help"),
+			hidden:   f.Tag.Get("hidden") == "yes",
 			freeForm: freeFormConfigFields[name],
 		}
 		if nested := underlyingConfigStruct(f.Type); nested != nil && !field.freeForm {
@@ -833,10 +864,16 @@ func renderConfigField(w io.Writer, f configField) {
 
 	width := 0
 	for _, c := range f.children {
+		if c.hidden {
+			continue
+		}
 		width = max(width, len(configLeafName(c.path)))
 	}
 	fmt.Fprintf(w, "\n  Fields:\n")
 	for _, c := range f.children {
+		if c.hidden {
+			continue
+		}
 		fmt.Fprintf(w, "    %-*s  %s\n", width, configLeafName(c.path), configFieldSummary(c))
 	}
 	fmt.Fprintf(w, "\nUse \"-h %s.<field>\" for details on a field.\n", f.path)

cmd/air/runconfig_load.go

@@ -102,6 +102,15 @@ func finishRunConfigLoad(path string, document *yaml.Node, validationYAML []byte
 	if err != nil {
 		return nil, err
 	}
+	// YAML decodes omitted and null pointers alike. Preserve explicit nulls so validation rejects them.
+	if cfg.CodeSource != nil {
+		if cfg.CodeSource.Type == nil && codeSourceFieldPresent(document, "type") {
+			cfg.CodeSource.Type = new(string)
+		}
+		if cfg.CodeSource.LegacySnapshot == nil && codeSourceFieldPresent(document, "snapshot") {
+			cfg.CodeSource.LegacySnapshot = &snapshotSourceConfig{}
+		}
+	}
 	if err := validateRunConfig(cfg); err != nil {
 		return nil, err
 	}
@@ -115,6 +124,15 @@ func finishRunConfigLoad(path string, document *yaml.Node, validationYAML []byte
 	return cfg, nil
 }
 
+func codeSourceFieldPresent(document *yaml.Node, field string) bool {
+	codeSource, _ := mappingValue(document.Content[0], "code_source")
+	if codeSource == nil || codeSource.Kind != yaml.MappingNode {
+		return false
+	}
+	fieldNode, _ := mappingValue(codeSource, field)
+	return fieldNode != nil
+}
+
 func stripYAMLComments(node *yaml.Node) {
 	node.HeadComment = ""
 	node.LineComment = ""

cmd/air/runconfig_override.go

@@ -67,6 +67,9 @@ func checkOverridePath(parts []string, node configField, fullPath string) error
 	if !ok {
 		names := make([]string, 0, len(node.children))
 		for _, c := range node.children {
+			if c.hidden {
+				continue
+			}
 			names = append(names, configLeafName(c.path))
 		}
 		slices.Sort(names)

cmd/air/runconfig_override_test.go

@@ -138,13 +138,15 @@ func TestLoadRunConfigWithOverrides(t *testing.T) {
 		assert.Equal(t, "0", cfg.EnvVariables["RANK"])
 	})
 
-	t.Run("intermediate maps are auto-created", func(t *testing.T) {
-		cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"code_source.type=snapshot", "code_source.snapshot.root_path=/Workspace/x"})
-		require.NoError(t, err)
-		require.NotNil(t, cfg.CodeSource)
-		require.NotNil(t, cfg.CodeSource.Snapshot)
-		assert.Equal(t, "/Workspace/x", cfg.CodeSource.Snapshot.RootPath)
-	})
+	for _, path := range []string{"code_source.root_path", "code_source.snapshot.root_path"} {
+		t.Run(path+" override", func(t *testing.T) {
+			cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{path + "=/Workspace/x"})
+			require.NoError(t, err)
+			require.NotNil(t, cfg.CodeSource)
+			require.NotNil(t, cfg.CodeSource.Snapshot)
+			assert.Equal(t, "/Workspace/x", cfg.CodeSource.Snapshot.RootPath)
+		})
+	}
 
 	t.Run("unity catalog image override applies", func(t *testing.T) {
 		cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"environment.unity_catalog_image=main.air.training:prod"})