From 3dab193c9f237c59c4b536507bc37ff3703fa4db Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Thu, 25 Jun 2026 17:32:55 +0000
Subject: [PATCH 1/4] experimental/air: authenticate `air get` up front and
 fail fast

The get command validated authentication only lazily: Config.Authenticate
(in PreRunE) merely attaches credentials, so a bad PAT or missing profile
surfaced as a confusing error partway through rendering, after a run's
config had already been printed.

Validate the workspace client up front instead:

- PreRunE maps MustWorkspaceClient failures to an actionable auth error: a
  missing-profile hint when no default profile is set and --profile (-p)
  was not passed (ErrCannotConfigureDefault), otherwise "authentication was
  not successful".
- RunE calls CurrentUser.Me before fetching or rendering anything, so a
  credential that resolves locally but is rejected by the workspace fails
  fast with the same clear message and no partial status/config output.

Co-authored-by: Isaac
---
 experimental/air/cmd/get.go      | 36 +++++++++++++++++++++++++---
 experimental/air/cmd/get_test.go | 41 ++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go
index 0da1f082e12..037a3cada7f 100644
--- a/experimental/air/cmd/get.go
+++ b/experimental/air/cmd/get.go
@@ -1,6 +1,7 @@
 package aircmd
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strconv"
@@ -9,6 +10,8 @@ import (
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/config"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/spf13/cobra"
 )
@@ -63,6 +66,22 @@ Sweep Tasks:
 {{- end}}
 `
 
+// errNoProfile is the actionable message shown when no credentials are
+// configured: no default profile, no --profile (-p), and no auth environment.
+var errNoProfile = errors.New("no default profile is set: pass --profile (-p) or configure a default profile in your .databrickscfg")
+
+// authError converts a workspace-client or authentication failure into an air
+// error envelope with an actionable message: a missing-profile hint when no
+// credentials are configured, otherwise "authentication was not successful".
+// Both are permanent (not retryable).
+func authError(ctx context.Context, cmd *cobra.Command, err error) error {
+	if errors.Is(err, config.ErrCannotConfigureDefault) {
+		return renderError(ctx, cmd, "UNAUTHENTICATED", "PERMANENT", false, errNoProfile)
+	}
+	return renderError(ctx, cmd, "UNAUTHENTICATED", "PERMANENT", false,
+		fmt.Errorf("authentication was not successful: %w", err))
+}
+
 // newGetCommand returns the `air get JOB_RUN_ID` command, which shows status,
 // configuration, and timing details for a specific run.
 func newGetCommand() *cobra.Command {
@@ -75,14 +94,16 @@ func newGetCommand() *cobra.Command {
 		},
 	}
 
-	// Match Python: a client/auth failure is a JSON error envelope in -o json mode,
-	// not a bare error. ErrAlreadyPrinted passes through (it was handled upstream).
+	// Resolve and authenticate the workspace client up front so an auth failure
+	// fails fast here, before any run status or config is fetched or printed.
+	// ErrAlreadyPrinted passes through (it was handled upstream); other failures
+	// become an actionable auth error (JSON envelope in -o json mode).
 	cmd.PreRunE = func(cmd *cobra.Command, args []string) error {
 		err := root.MustWorkspaceClient(cmd, args)
 		if err == nil || errors.Is(err, root.ErrAlreadyPrinted) {
 			return err
 		}
-		return renderError(cmd.Context(), cmd, "INTERNAL_ERROR", "TRANSIENT", true, err)
+		return authError(cmd.Context(), cmd, err)
 	}
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
@@ -95,6 +116,15 @@ func newGetCommand() *cobra.Command {
 				fmt.Errorf("invalid JOB_RUN_ID %q: must be a positive integer", args[0]))
 		}
 
+		// Validate authentication against the workspace before fetching or
+		// rendering anything. MustWorkspaceClient's Config.Authenticate only
+		// attaches credentials (e.g. it does not check a PAT server-side), so
+		// without this a bad credential would surface as a confusing failure
+		// mid-render instead of a clear "not authenticated" error here.
+		if _, err := w.CurrentUser.Me(ctx, iam.MeRequest{}); err != nil {
+			return authError(ctx, cmd, err)
+		}
+
 		run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID})
 		if err != nil {
 			// The backend returns this when the run ID is unknown to the user.
diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go
index 8ff864e387c..e77abeb2196 100644
--- a/experimental/air/cmd/get_test.go
+++ b/experimental/air/cmd/get_test.go
@@ -3,6 +3,7 @@ package aircmd
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"testing"
 	"text/template"
 
@@ -11,13 +12,21 @@ import (
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
+// expectAuthSuccess stubs CurrentUser.Me so the command's up-front auth check
+// passes and the test can exercise the run-fetch path.
+func expectAuthSuccess(m *mocks.MockWorkspaceClient) {
+	m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(&iam.User{UserName: "u@example.com"}, nil)
+}
+
 // renderGet renders the get template against the JSON envelope, exactly as the
 // command does for a sweep run, so the test covers the real template branches.
 func renderGet(t *testing.T, data getData) string {
@@ -55,6 +64,7 @@ func TestGetRunInvalidID(t *testing.T) {
 
 func TestGetRunNotFound(t *testing.T) {
 	m := mocks.NewMockWorkspaceClient(t)
+	expectAuthSuccess(m)
 	m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 5}).Return(
 		nil, apierr.ErrResourceDoesNotExist)
 	ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient)
@@ -66,9 +76,40 @@ func TestGetRunNotFound(t *testing.T) {
 	assert.Contains(t, err.Error(), "run 5 not found")
 }
 
+func TestGetRunAuthFailed(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+	// Auth is validated before the run is fetched, so GetRun is never reached
+	// (no GetRun expectation is set) and nothing is rendered.
+	m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, errors.New("403 Forbidden"))
+	ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient)
+	cmd := withOutput(newGetCommand(), flags.OutputText)
+	cmd.SetContext(ctx)
+
+	err := cmd.RunE(cmd, []string{"5"})
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "authentication was not successful")
+}
+
+func TestAuthError(t *testing.T) {
+	ctx := cmdio.MockDiscard(t.Context())
+	cmd := withOutput(newGetCommand(), flags.OutputText)
+
+	// No configurable credentials maps to the missing-profile hint.
+	noProfile := authError(ctx, cmd, config.ErrCannotConfigureDefault)
+	require.Error(t, noProfile)
+	assert.Contains(t, noProfile.Error(), "no default profile is set")
+
+	// Any other failure maps to the generic auth message, preserving the cause.
+	other := authError(ctx, cmd, errors.New("token expired"))
+	require.Error(t, other)
+	assert.Contains(t, other.Error(), "authentication was not successful")
+	assert.Contains(t, other.Error(), "token expired")
+}
+
 func TestGetRunNotFoundJSON(t *testing.T) {
 	var buf bytes.Buffer
 	m := mocks.NewMockWorkspaceClient(t)
+	expectAuthSuccess(m)
 	m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 5}).Return(
 		nil, apierr.ErrResourceDoesNotExist)
 	ctx := cmdctx.SetWorkspaceClient(t.Context(), m.WorkspaceClient)

From 5d0feb4165c00ed91974189de6aeda508a27fa6d Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Mon, 6 Jul 2026 15:23:47 +0000
Subject: [PATCH 2/4] experimental/air: render `air get` config and metadata
 for ai_runtime_task runs

The typed SDK's jobs.RunTask has no ai_runtime_task field, so w.Jobs.GetRun
drops the task for runs submitted by the current `air run` (which emits
ai_runtime_task). This left the Configuration box empty and Experiment /
Accelerators as N/A for those runs.

Mirror the Python CLI: for a single run with no gen_ai_compute_task, read the
raw runs/get payload, take the ai_runtime_task command_path, and fill the config
box from the sibling training_config.yaml plus the experiment and accelerator
cells. Legacy gen_ai_compute_task runs keep their existing path.

Co-authored-by: Isaac
---
 .../air/get-ai-runtime/out.test.toml          |  3 +
 .../air/get-ai-runtime/output.txt             | 52 ++++++++++++++++++
 .../experimental/air/get-ai-runtime/script    |  9 +++
 .../experimental/air/get-ai-runtime/test.toml | 55 +++++++++++++++++++
 .../air/get-ai-runtime/training_config.yaml   |  8 +++
 experimental/air/cmd/get.go                   | 29 ++++++++++
 experimental/air/cmd/get_test.go              | 40 ++++++++++++++
 experimental/air/cmd/joblist.go               | 16 +++++-
 experimental/air/cmd/joblist_test.go          | 50 +++++++++++++++++
 experimental/air/cmd/render.go                | 20 +++++--
 experimental/air/cmd/render_test.go           | 27 +++++++++
 11 files changed, 303 insertions(+), 6 deletions(-)
 create mode 100644 acceptance/experimental/air/get-ai-runtime/out.test.toml
 create mode 100644 acceptance/experimental/air/get-ai-runtime/output.txt
 create mode 100644 acceptance/experimental/air/get-ai-runtime/script
 create mode 100644 acceptance/experimental/air/get-ai-runtime/test.toml
 create mode 100644 acceptance/experimental/air/get-ai-runtime/training_config.yaml
 create mode 100644 experimental/air/cmd/joblist_test.go

diff --git a/acceptance/experimental/air/get-ai-runtime/out.test.toml b/acceptance/experimental/air/get-ai-runtime/out.test.toml
new file mode 100644
index 00000000000..d6187dcb046
--- /dev/null
+++ b/acceptance/experimental/air/get-ai-runtime/out.test.toml
@@ -0,0 +1,3 @@
+Local = true
+Cloud = false
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = []
diff --git a/acceptance/experimental/air/get-ai-runtime/output.txt b/acceptance/experimental/air/get-ai-runtime/output.txt
new file mode 100644
index 00000000000..150f394801f
--- /dev/null
+++ b/acceptance/experimental/air/get-ai-runtime/output.txt
@@ -0,0 +1,52 @@
+
+=== get (text)
+>>> [CLI] experimental air get 123
+
+╭─ Configuration ────────────────────────────────────────────────╮
+│                                                                │
+│  experiment_name: my-exp                                       │
+│  compute:                                                      │
+│    accelerator_type: a10                                       │
+│    num_accelerators: 1                                         │
+│  command: |                                                    │
+│    for i in $(seq 1 10); do                                    │
+│      echo "step $i"                                            │
+│    done                                                        │
+│                                                                │
+╰────────────────────────────────────────────────────────────────╯
+
+╭─ Metadata ─────────────────────────────────────────────────────╮
+│                                                                │
+│  Run ID        123                                             │
+│  Status        ● SUCCESS                                       │
+│  Submitted     2023-11-14 22:13 UTC                            │
+│  Retries       0                                               │
+│  Max Retries   3                                               │
+│  Duration      12s                                             │
+│  Experiment    my-exp                                          │
+│  MLflow Run    my-run                                          │
+│  User          user@example.com                                │
+│  Accelerators  1x A10                                          │
+│  Environment   N/A                                             │
+│                                                                │
+╰────────────────────────────────────────────────────────────────╯
+
+Run URL:    [DATABRICKS_URL]/jobs/runs/123?o=[NUMID]
+MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1
+
+=== get (json)
+>>> [CLI] experimental air get 123 -o json
+{
+  "v": 1,
+  "ts": "[TIMESTAMP]",
+  "data": {
+    "run_id": "123",
+    "status": "SUCCESS",
+    "started_at": "[TIMESTAMP]",
+    "duration_seconds": 12,
+    "attempt_number": 0,
+    "experiment_name": "my-exp",
+    "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[NUMID]",
+    "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/artifacts/logs/node_0"
+  }
+}
diff --git a/acceptance/experimental/air/get-ai-runtime/script b/acceptance/experimental/air/get-ai-runtime/script
new file mode 100644
index 00000000000..616e0e3087d
--- /dev/null
+++ b/acceptance/experimental/air/get-ai-runtime/script
@@ -0,0 +1,9 @@
+# Seed the run's training_config.yaml next to command.sh so `air get` can
+# download and render it in the Configuration box.
+$CLI workspace import "/Workspace/Users/user@example.com/.air/cli_launch/my-exp/my-exp_abc/training_config.yaml" --file training_config.yaml --format AUTO &> LOG.import
+
+title "get (text)"
+trace $CLI experimental air get 123
+
+title "get (json)"
+trace $CLI experimental air get 123 -o json
diff --git a/acceptance/experimental/air/get-ai-runtime/test.toml b/acceptance/experimental/air/get-ai-runtime/test.toml
new file mode 100644
index 00000000000..c344dbb7af6
--- /dev/null
+++ b/acceptance/experimental/air/get-ai-runtime/test.toml
@@ -0,0 +1,55 @@
+# This command does not deploy a bundle, so no engine matrix is needed.
+[EnvMatrix]
+DATABRICKS_BUNDLE_ENGINE = []
+
+# The SDK occasionally probes host reachability with a HEAD request; stub it so
+# the test is deterministic.
+[[Server]]
+Pattern = "HEAD /"
+Response.Body = ''
+
+# The typed SDK GetRun response: an ai_runtime_task run has no gen_ai_compute_task,
+# so the task comes back empty (the SDK has no field for ai_runtime_task).
+[[Server]]
+Pattern = "GET /api/2.2/jobs/runs/get"
+Response.Body = '''
+{
+  "run_id": 123,
+  "run_page_url": "https://my-workspace.cloud.databricks.test/jobs/runs/123",
+  "creator_user_name": "user@example.com",
+  "start_time": 1700000000000,
+  "end_time": 1700000012000,
+  "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"},
+  "tasks": [
+    {
+      "task_key": "train",
+      "run_id": 456,
+      "attempt_number": 0,
+      "max_retries": 3,
+      "ai_runtime_task": {
+        "experiment": "my-exp",
+        "deployments": [
+          {
+            "command_path": "/Workspace/Users/user@example.com/.air/cli_launch/my-exp/my-exp_abc/command.sh",
+            "compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1}
+          }
+        ]
+      }
+    }
+  ]
+}
+'''
+
+# MLflow identifiers for the deep-link (runs/get-output is not modeled by the typed SDK).
+[[Server]]
+Pattern = "GET /api/2.2/jobs/runs/get-output"
+Response.Body = '''
+{"gen_ai_compute_output": {"run_info": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}}
+'''
+
+# The MLflow Run cell shows the run's name, fetched from the MLflow REST API.
+[[Server]]
+Pattern = "GET /api/2.0/mlflow/runs/get"
+Response.Body = '''
+{"run": {"info": {"run_name": "my-run"}}}
+'''
diff --git a/acceptance/experimental/air/get-ai-runtime/training_config.yaml b/acceptance/experimental/air/get-ai-runtime/training_config.yaml
new file mode 100644
index 00000000000..82102bef3c0
--- /dev/null
+++ b/acceptance/experimental/air/get-ai-runtime/training_config.yaml
@@ -0,0 +1,8 @@
+experiment_name: my-exp
+compute:
+  accelerator_type: a10
+  num_accelerators: 1
+command: |
+  for i in $(seq 1 10); do
+    echo "step $i"
+  done
diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go
index 037a3cada7f..e1044a7ada0 100644
--- a/experimental/air/cmd/get.go
+++ b/experimental/air/cmd/get.go
@@ -4,11 +4,14 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"path"
 	"strconv"
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/flags"
+	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/service/iam"
@@ -41,6 +44,8 @@ type getData struct {
 	AcceleratorsDisplay string `json:"-"`
 	EnvironmentDisplay  string `json:"-"`
 	MaxRetriesDisplay   string `json:"-"`
+	// TrainingConfigPath is the run's config file, downloaded for the config box.
+	TrainingConfigPath string `json:"-"`
 	// Sweep replaces the single-run view for foreach runs.
 	Sweep *sweepInfo `json:"-"`
 }
@@ -151,6 +156,9 @@ func newGetCommand() *cobra.Command {
 		}
 		if task := findForEachTask(run); task != nil {
 			data.Sweep = buildSweepInfo(ctx, w, task)
+		} else if genAIComputeTask(run) == nil {
+			// The typed SDK drops ai_runtime_task, so read it from the raw run.
+			enrichFromRawRun(ctx, w, runID, &data)
 		}
 
 		if root.OutputType(cmd) != flags.OutputText {
@@ -172,6 +180,27 @@ func newGetCommand() *cobra.Command {
 	return cmd
 }
 
+// enrichFromRawRun fills the config path, experiment, and accelerators from the
+// raw run. Best-effort: on failure the existing "N/A" fallbacks stand.
+func enrichFromRawRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64, data *getData) {
+	raw, err := fetchJobRun(ctx, w, runID)
+	if err != nil {
+		log.Debugf(ctx, "air get: raw run lookup failed for run %d: %v", runID, err)
+		return
+	}
+
+	if cmdPath := raw.commandPath(); cmdPath != "" {
+		data.TrainingConfigPath = path.Join(path.Dir(cmdPath), trainingConfigName)
+	}
+	if exp := jobExperiment(raw); exp != "" {
+		data.ExperimentName = &exp
+		data.ExperimentDisplay = exp
+	}
+	if a := acceleratorLabel(jobCompute(raw)); a != "" {
+		data.AcceleratorsDisplay = a
+	}
+}
+
 // buildGetData extracts the fields we display from a run. The text-view cells
 // are pre-rendered here with their "N/A" fallbacks; the styled renderer adds the
 // hyperlinks and colors once the dashboard and MLflow identifiers are known.
diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go
index e77abeb2196..1bf663c019b 100644
--- a/experimental/air/cmd/get_test.go
+++ b/experimental/air/cmd/get_test.go
@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"encoding/json"
 	"errors"
+	"net/http"
+	"net/http/httptest"
 	"testing"
 	"text/template"
 
@@ -190,6 +192,44 @@ func TestBuildGetData(t *testing.T) {
 	assert.Equal(t, int64(12), *d.DurationSeconds)
 }
 
+func TestEnrichFromRawRun(t *testing.T) {
+	t.Run("fills config path, experiment, and accelerators", func(t *testing.T) {
+		body := `{"run_id":5,"tasks":[{"ai_runtime_task":{
+			"experiment":"/Users/me@example.com/my-exp",
+			"deployments":[{"command_path":"/Workspace/run/command.sh","compute":{"accelerator_type":"GPU_1xA10","accelerator_count":1}}]
+		}}]}`
+		srv := runGetServer(t, body)
+		data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na}
+		enrichFromRawRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5, data)
+		assert.Equal(t, "/Workspace/run/training_config.yaml", data.TrainingConfigPath)
+		assert.Equal(t, "my-exp", data.ExperimentDisplay)
+		require.NotNil(t, data.ExperimentName)
+		assert.Equal(t, "my-exp", *data.ExperimentName)
+		assert.Equal(t, "1x A10", data.AcceleratorsDisplay)
+	})
+
+	t.Run("leaves fallbacks when the run has no ai_runtime_task", func(t *testing.T) {
+		srv := runGetServer(t, `{"run_id":5,"tasks":[{}]}`)
+		data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na}
+		enrichFromRawRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5, data)
+		assert.Empty(t, data.TrainingConfigPath)
+		assert.Equal(t, na, data.ExperimentDisplay)
+		assert.Equal(t, na, data.AcceleratorsDisplay)
+	})
+
+	t.Run("leaves fallbacks when the raw lookup fails", func(t *testing.T) {
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusInternalServerError)
+		}))
+		t.Cleanup(srv.Close)
+		data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na}
+		enrichFromRawRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5, data)
+		assert.Empty(t, data.TrainingConfigPath)
+		assert.Equal(t, na, data.ExperimentDisplay)
+		assert.Equal(t, na, data.AcceleratorsDisplay)
+	})
+}
+
 func TestBuildGetDataEmpty(t *testing.T) {
 	// A run with no tasks, creator, or timing renders every text cell as "N/A".
 	d := buildGetData(&jobs.Run{RunId: 7})
diff --git a/experimental/air/cmd/joblist.go b/experimental/air/cmd/joblist.go
index c84f87c9df0..4c9e05e9b8c 100644
--- a/experimental/air/cmd/joblist.go
+++ b/experimental/air/cmd/joblist.go
@@ -66,6 +66,8 @@ type jobAiRuntimeTask struct {
 
 type aiRuntimeDeploy struct {
 	Compute airCompute `json:"compute"`
+	// CommandPath is command.sh; training_config.yaml sits beside it.
+	CommandPath string `json:"command_path"`
 }
 
 type airCompute struct {
@@ -192,7 +194,8 @@ func fetchJobRunsPage(ctx context.Context, w *databricks.WorkspaceClient, query
 }
 
 // fetchJobRun fetches a single run via runs/get. The response mirrors one
-// runs/list element, so it deserializes into jobRun directly.
+// runs/list element, so it deserializes into jobRun directly; expand_tasks
+// pulls the ai_runtime_task the typed SDK omits.
 func fetchJobRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (*jobRun, error) {
 	apiClient, err := client.New(w.Config)
 	if err != nil {
@@ -203,7 +206,7 @@ func fetchJobRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64
 	query := map[string]any{"run_id": runID, "expand_tasks": true}
 	err = apiClient.Do(ctx, http.MethodGet, jobsRunsGetPath, nil, nil, query, &run)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to get run %d: %w", runID, err)
 	}
 	return &run, nil
 }
@@ -242,3 +245,12 @@ func hydrateJobRuns(ctx context.Context, w *databricks.WorkspaceClient, ids []in
 	}
 	return hydrated, nil
 }
+
+// commandPath returns the run's command.sh path, or "" when it has no deployment.
+func (r *jobRun) commandPath() string {
+	ai, _ := r.airTask()
+	if ai == nil || len(ai.Deployments) == 0 {
+		return ""
+	}
+	return ai.Deployments[0].CommandPath
+}
diff --git a/experimental/air/cmd/joblist_test.go b/experimental/air/cmd/joblist_test.go
new file mode 100644
index 00000000000..d8f03fc60ba
--- /dev/null
+++ b/experimental/air/cmd/joblist_test.go
@@ -0,0 +1,50 @@
+package aircmd
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// runGetServer serves one runs/get body and a stub for the SDK config probe.
+func runGetServer(t *testing.T, body string) *httptest.Server {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == jobsRunsGetPath {
+			_, _ = w.Write([]byte(body))
+			return
+		}
+		_, _ = w.Write([]byte(`{}`))
+	}))
+	t.Cleanup(srv.Close)
+	return srv
+}
+
+func TestFetchJobRunParsesAiRuntimeTask(t *testing.T) {
+	body := `{
+		"run_id": 5,
+		"tasks": [{
+			"ai_runtime_task": {
+				"experiment": "my-exp",
+				"deployments": [{
+					"command_path": "/Workspace/Users/me/.air/cli_launch/my-exp/my-exp_abc/command.sh",
+					"compute": {"accelerator_count": 1, "accelerator_type": "GPU_1xA10"}
+				}]
+			}
+		}]
+	}`
+	srv := runGetServer(t, body)
+
+	run, err := fetchJobRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5)
+	require.NoError(t, err)
+	assert.Equal(t, "/Workspace/Users/me/.air/cli_launch/my-exp/my-exp_abc/command.sh", run.commandPath())
+}
+
+func TestCommandPathEmpty(t *testing.T) {
+	// No ai_runtime_task deployment: no command path.
+	assert.Empty(t, (&jobRun{Tasks: []jobTask{{}}}).commandPath())
+	assert.Empty(t, (&jobRun{}).commandPath())
+}
diff --git a/experimental/air/cmd/render.go b/experimental/air/cmd/render.go
index ff7c0298671..3d1fe9637c5 100644
--- a/experimental/air/cmd/render.go
+++ b/experimental/air/cmd/render.go
@@ -115,10 +115,8 @@ func renderRunText(ctx context.Context, out io.Writer, w *databricks.WorkspaceCl
 	}
 
 	var sections []string
-	if task := genAIComputeTask(run); task != nil {
-		if body := colorizeConfig(p, configYAML(ctx, w, task)); body != "" {
-			sections = append(sections, renderBox(p, configBoxTitle, body))
-		}
+	if body := colorizeConfig(p, resolveConfigYAML(ctx, w, run, data)); body != "" {
+		sections = append(sections, renderBox(p, configBoxTitle, body))
 	}
 	sections = append(sections, renderBox(p, metadataBoxTitle, renderFields(p, colorOn, view)))
 
@@ -148,6 +146,20 @@ func genAIComputeTask(run *jobs.Run) *jobs.GenAiComputeTask {
 	return run.Tasks[0].GenAiComputeTask
 }
 
+// resolveConfigYAML returns the config box body: from the downloaded config file
+// when we have its path, else from the legacy task.
+func resolveConfigYAML(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run, data *getData) string {
+	if data.TrainingConfigPath != "" {
+		if raw := downloadConfig(ctx, w, data.TrainingConfigPath); len(raw) > 0 {
+			return strings.TrimRight(reformatYAMLForDisplay(raw), "\n")
+		}
+	}
+	if task := genAIComputeTask(run); task != nil {
+		return configYAML(ctx, w, task)
+	}
+	return ""
+}
+
 // configYAML returns the run's resolved training config as YAML for the box. The
 // full config (including the command/script) lives in the run's parameters, not
 // the structured task fields, so we prefer the inline parameters, then the
diff --git a/experimental/air/cmd/render_test.go b/experimental/air/cmd/render_test.go
index 71e2732fe0f..e950bb22331 100644
--- a/experimental/air/cmd/render_test.go
+++ b/experimental/air/cmd/render_test.go
@@ -55,6 +55,33 @@ func TestConfigYAML(t *testing.T) {
 	})
 }
 
+func TestResolveConfigYAML(t *testing.T) {
+	ctx := t.Context()
+
+	t.Run("downloads the training config path for new runs", func(t *testing.T) {
+		m := mocks.NewMockWorkspaceClient(t)
+		m.GetMockWorkspaceAPI().EXPECT().
+			Download(mock.Anything, "/Workspace/run/training_config.yaml").
+			Return(io.NopCloser(strings.NewReader("experiment_name: from-air\n")), nil)
+		run := &jobs.Run{Tasks: []jobs.RunTask{{}}}
+		data := &getData{TrainingConfigPath: "/Workspace/run/training_config.yaml"}
+		assert.Equal(t, "experiment_name: from-air", resolveConfigYAML(ctx, m.WorkspaceClient, run, data))
+	})
+
+	t.Run("falls back to the legacy task when no path is set", func(t *testing.T) {
+		run := &jobs.Run{Tasks: []jobs.RunTask{{GenAiComputeTask: &jobs.GenAiComputeTask{
+			MlflowExperimentName: "/Users/me@example.com/exp",
+		}}}}
+		got := resolveConfigYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, run, &getData{})
+		assert.Contains(t, got, "experiment_name: exp")
+	})
+
+	t.Run("empty when neither source is present", func(t *testing.T) {
+		run := &jobs.Run{Tasks: []jobs.RunTask{{}}}
+		assert.Empty(t, resolveConfigYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, run, &getData{}))
+	})
+}
+
 func TestSynthConfigYAML(t *testing.T) {
 	task := &jobs.GenAiComputeTask{
 		MlflowExperimentName: "/Users/me@example.com/stream-latency-test",

From 3eabc54ba148b6de90128f76c17fa9a673980bb8 Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Wed, 8 Jul 2026 04:30:23 +0000
Subject: [PATCH 3/4] experimental/air: address `air get` review feedback
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two fixes from review of the fail-fast/render commits:

- authError only maps genuinely auth-shaped failures (missing profile, or
  an apierr 401/403 via ErrUnauthenticated/ErrPermissionDenied) to
  UNAUTHENTICATED/PERMANENT. A transient failure at the Me() probe (429,
  5xx, network blip) is now reported as INTERNAL_ERROR/TRANSIENT so it
  stays retryable instead of a misleading non-retryable auth error.

- `air get` fetched the run twice on the ai_runtime_task path: w.Jobs.GetRun
  (typed) then a second runs/get in enrichFromRawRun. Add fetchRun, which
  makes one runs/get call and unmarshals the body into both the typed
  jobs.Run (display path) and the raw jobRun (preserves ai_runtime_task).
  enrichFromRawRun now takes the already-fetched raw run — no second
  roundtrip.

Tests: command-level not-found tests move to a real runs/get HTTP server
(the run fetch is no longer the typed, mockable API); add fetchRun coverage
and a transient-auth case.

Co-authored-by: Isaac
---
 experimental/air/cmd/get.go      |  39 +++++-----
 experimental/air/cmd/get_test.go | 121 +++++++++++++++++++++----------
 experimental/air/cmd/joblist.go  |  29 ++++++++
 3 files changed, 133 insertions(+), 56 deletions(-)

diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go
index e1044a7ada0..a6b47c977f7 100644
--- a/experimental/air/cmd/get.go
+++ b/experimental/air/cmd/get.go
@@ -10,8 +10,6 @@ import (
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/flags"
-	"github.com/databricks/cli/libs/log"
-	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/service/iam"
@@ -75,16 +73,20 @@ Sweep Tasks:
 // configured: no default profile, no --profile (-p), and no auth environment.
 var errNoProfile = errors.New("no default profile is set: pass --profile (-p) or configure a default profile in your .databrickscfg")
 
-// authError converts a workspace-client or authentication failure into an air
-// error envelope with an actionable message: a missing-profile hint when no
-// credentials are configured, otherwise "authentication was not successful".
-// Both are permanent (not retryable).
+// authError classifies a workspace-client or Me() probe failure. Only genuinely
+// auth-shaped errors surface as UNAUTHENTICATED/PERMANENT: missing profile,
+// SDK auth wrappers, or an API 401/403. Anything else (network blip, 429, 5xx)
+// is transient and reported as INTERNAL_ERROR/TRANSIENT so the caller can retry.
 func authError(ctx context.Context, cmd *cobra.Command, err error) error {
 	if errors.Is(err, config.ErrCannotConfigureDefault) {
 		return renderError(ctx, cmd, "UNAUTHENTICATED", "PERMANENT", false, errNoProfile)
 	}
-	return renderError(ctx, cmd, "UNAUTHENTICATED", "PERMANENT", false,
-		fmt.Errorf("authentication was not successful: %w", err))
+	if errors.Is(err, apierr.ErrUnauthenticated) || errors.Is(err, apierr.ErrPermissionDenied) {
+		return renderError(ctx, cmd, "UNAUTHENTICATED", "PERMANENT", false,
+			fmt.Errorf("authentication was not successful: %w", err))
+	}
+	return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true,
+		fmt.Errorf("failed to verify authentication: %w", err))
 }
 
 // newGetCommand returns the `air get JOB_RUN_ID` command, which shows status,
@@ -130,7 +132,10 @@ func newGetCommand() *cobra.Command {
 			return authError(ctx, cmd, err)
 		}
 
-		run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID})
+		// Fetch the run once, in both the typed and raw shapes: the typed jobs.Run
+		// drives the display path, and the raw jobRun preserves the ai_runtime_task
+		// the typed model drops (used below without a second roundtrip).
+		run, rawRun, err := fetchRun(ctx, w, runID)
 		if err != nil {
 			// The backend returns this when the run ID is unknown to the user.
 			if errors.Is(err, apierr.ErrResourceDoesNotExist) {
@@ -157,8 +162,9 @@ func newGetCommand() *cobra.Command {
 		if task := findForEachTask(run); task != nil {
 			data.Sweep = buildSweepInfo(ctx, w, task)
 		} else if genAIComputeTask(run) == nil {
-			// The typed SDK drops ai_runtime_task, so read it from the raw run.
-			enrichFromRawRun(ctx, w, runID, &data)
+			// The typed SDK drops ai_runtime_task, so read it from the raw run we
+			// already fetched above.
+			enrichFromRawRun(rawRun, &data)
 		}
 
 		if root.OutputType(cmd) != flags.OutputText {
@@ -181,14 +187,9 @@ func newGetCommand() *cobra.Command {
 }
 
 // enrichFromRawRun fills the config path, experiment, and accelerators from the
-// raw run. Best-effort: on failure the existing "N/A" fallbacks stand.
-func enrichFromRawRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64, data *getData) {
-	raw, err := fetchJobRun(ctx, w, runID)
-	if err != nil {
-		log.Debugf(ctx, "air get: raw run lookup failed for run %d: %v", runID, err)
-		return
-	}
-
+// raw run (the ai_runtime_task the typed model drops). Best-effort: empty fields
+// leave the existing "N/A" fallbacks in place.
+func enrichFromRawRun(raw *jobRun, data *getData) {
 	if cmdPath := raw.commandPath(); cmdPath != "" {
 		data.TrainingConfigPath = path.Join(path.Dir(cmdPath), trainingConfigName)
 	}
diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go
index 1bf663c019b..88a7bdfc4bd 100644
--- a/experimental/air/cmd/get_test.go
+++ b/experimental/air/cmd/get_test.go
@@ -16,19 +16,12 @@ import (
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/experimental/mocks"
-	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
-// expectAuthSuccess stubs CurrentUser.Me so the command's up-front auth check
-// passes and the test can exercise the run-fetch path.
-func expectAuthSuccess(m *mocks.MockWorkspaceClient) {
-	m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(&iam.User{UserName: "u@example.com"}, nil)
-}
-
 // renderGet renders the get template against the JSON envelope, exactly as the
 // command does for a sweep run, so the test covers the real template branches.
 func renderGet(t *testing.T, data getData) string {
@@ -64,12 +57,27 @@ func TestGetRunInvalidID(t *testing.T) {
 	assert.Contains(t, err.Error(), "invalid JOB_RUN_ID")
 }
 
+// notFoundGetServer serves the auth probe plus a runs/get that reports the run
+// as missing (400 INVALID_PARAMETER_VALUE, which the SDK maps to
+// ErrResourceDoesNotExist for this path).
+func notFoundGetServer(t *testing.T) *httptest.Server {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == jobsRunsGetPath {
+			w.WriteHeader(http.StatusBadRequest)
+			_, _ = w.Write([]byte(`{"error_code":"INVALID_PARAMETER_VALUE","message":"Run 5 does not exist."}`))
+			return
+		}
+		// Me() probe and any other config discovery.
+		_, _ = w.Write([]byte(`{"userName":"u@example.com"}`))
+	}))
+	t.Cleanup(srv.Close)
+	return srv
+}
+
 func TestGetRunNotFound(t *testing.T) {
-	m := mocks.NewMockWorkspaceClient(t)
-	expectAuthSuccess(m)
-	m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 5}).Return(
-		nil, apierr.ErrResourceDoesNotExist)
-	ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient)
+	srv := notFoundGetServer(t)
+	ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL))
 	cmd := withOutput(newGetCommand(), flags.OutputText)
 	cmd.SetContext(ctx)
 
@@ -80,9 +88,9 @@ func TestGetRunNotFound(t *testing.T) {
 
 func TestGetRunAuthFailed(t *testing.T) {
 	m := mocks.NewMockWorkspaceClient(t)
-	// Auth is validated before the run is fetched, so GetRun is never reached
-	// (no GetRun expectation is set) and nothing is rendered.
-	m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, errors.New("403 Forbidden"))
+	// A genuine auth failure (permission denied) is validated before the run is
+	// fetched, so GetRun is never reached and nothing is rendered.
+	m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, apierr.ErrPermissionDenied)
 	ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient)
 	cmd := withOutput(newGetCommand(), flags.OutputText)
 	cmd.SetContext(ctx)
@@ -92,6 +100,21 @@ func TestGetRunAuthFailed(t *testing.T) {
 	assert.Contains(t, err.Error(), "authentication was not successful")
 }
 
+func TestGetRunAuthTransient(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+	// A transient failure at the auth probe must not be misreported as an auth
+	// error; it surfaces as a retryable internal error instead.
+	m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, errors.New("connection reset"))
+	ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient)
+	cmd := withOutput(newGetCommand(), flags.OutputText)
+	cmd.SetContext(ctx)
+
+	err := cmd.RunE(cmd, []string{"5"})
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to verify authentication")
+	assert.NotContains(t, err.Error(), "authentication was not successful")
+}
+
 func TestAuthError(t *testing.T) {
 	ctx := cmdio.MockDiscard(t.Context())
 	cmd := withOutput(newGetCommand(), flags.OutputText)
@@ -101,20 +124,25 @@ func TestAuthError(t *testing.T) {
 	require.Error(t, noProfile)
 	assert.Contains(t, noProfile.Error(), "no default profile is set")
 
-	// Any other failure maps to the generic auth message, preserving the cause.
-	other := authError(ctx, cmd, errors.New("token expired"))
-	require.Error(t, other)
-	assert.Contains(t, other.Error(), "authentication was not successful")
-	assert.Contains(t, other.Error(), "token expired")
+	// A 401 / 403 (via the SDK sentinels) is a real auth failure.
+	unauth := authError(ctx, cmd, apierr.ErrUnauthenticated)
+	require.Error(t, unauth)
+	assert.Contains(t, unauth.Error(), "authentication was not successful")
+
+	denied := authError(ctx, cmd, apierr.ErrPermissionDenied)
+	require.Error(t, denied)
+	assert.Contains(t, denied.Error(), "authentication was not successful")
+
+	transient := authError(ctx, cmd, errors.New("connection reset"))
+	require.Error(t, transient)
+	assert.Contains(t, transient.Error(), "failed to verify authentication")
+	assert.Contains(t, transient.Error(), "connection reset")
 }
 
 func TestGetRunNotFoundJSON(t *testing.T) {
 	var buf bytes.Buffer
-	m := mocks.NewMockWorkspaceClient(t)
-	expectAuthSuccess(m)
-	m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 5}).Return(
-		nil, apierr.ErrResourceDoesNotExist)
-	ctx := cmdctx.SetWorkspaceClient(t.Context(), m.WorkspaceClient)
+	srv := notFoundGetServer(t)
+	ctx := cmdctx.SetWorkspaceClient(t.Context(), newTestWorkspaceClient(t, srv.URL))
 	ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputJSON, nil, &buf, &buf, "", ""))
 	cmd := withOutput(newGetCommand(), flags.OutputJSON)
 	cmd.SetContext(ctx)
@@ -193,14 +221,20 @@ func TestBuildGetData(t *testing.T) {
 }
 
 func TestEnrichFromRawRun(t *testing.T) {
+	rawRun := func(t *testing.T, body string) *jobRun {
+		t.Helper()
+		var r jobRun
+		require.NoError(t, json.Unmarshal([]byte(body), &r))
+		return &r
+	}
+
 	t.Run("fills config path, experiment, and accelerators", func(t *testing.T) {
-		body := `{"run_id":5,"tasks":[{"ai_runtime_task":{
+		raw := rawRun(t, `{"run_id":5,"tasks":[{"ai_runtime_task":{
 			"experiment":"/Users/me@example.com/my-exp",
 			"deployments":[{"command_path":"/Workspace/run/command.sh","compute":{"accelerator_type":"GPU_1xA10","accelerator_count":1}}]
-		}}]}`
-		srv := runGetServer(t, body)
+		}}]}`)
 		data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na}
-		enrichFromRawRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5, data)
+		enrichFromRawRun(raw, data)
 		assert.Equal(t, "/Workspace/run/training_config.yaml", data.TrainingConfigPath)
 		assert.Equal(t, "my-exp", data.ExperimentDisplay)
 		require.NotNil(t, data.ExperimentName)
@@ -209,24 +243,37 @@ func TestEnrichFromRawRun(t *testing.T) {
 	})
 
 	t.Run("leaves fallbacks when the run has no ai_runtime_task", func(t *testing.T) {
-		srv := runGetServer(t, `{"run_id":5,"tasks":[{}]}`)
+		raw := rawRun(t, `{"run_id":5,"tasks":[{}]}`)
 		data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na}
-		enrichFromRawRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5, data)
+		enrichFromRawRun(raw, data)
 		assert.Empty(t, data.TrainingConfigPath)
 		assert.Equal(t, na, data.ExperimentDisplay)
 		assert.Equal(t, na, data.AcceleratorsDisplay)
 	})
+}
 
-	t.Run("leaves fallbacks when the raw lookup fails", func(t *testing.T) {
+func TestFetchRun(t *testing.T) {
+	t.Run("parses the run into both the typed and raw shapes from one call", func(t *testing.T) {
+		body := `{"run_id":5,"creator_user_name":"me@example.com","tasks":[{"ai_runtime_task":{
+			"experiment":"/Users/me@example.com/my-exp",
+			"deployments":[{"command_path":"/Workspace/run/command.sh","compute":{"accelerator_type":"GPU_1xA10","accelerator_count":1}}]
+		}}]}`
+		srv := runGetServer(t, body)
+		typed, raw, err := fetchRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5)
+		require.NoError(t, err)
+		// Typed jobs.Run carries the common fields; raw jobRun keeps ai_runtime_task.
+		assert.Equal(t, int64(5), typed.RunId)
+		assert.Equal(t, "me@example.com", typed.CreatorUserName)
+		assert.Equal(t, "/Workspace/run/command.sh", raw.commandPath())
+	})
+
+	t.Run("propagates a fetch failure", func(t *testing.T) {
 		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(http.StatusInternalServerError)
 		}))
 		t.Cleanup(srv.Close)
-		data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na}
-		enrichFromRawRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5, data)
-		assert.Empty(t, data.TrainingConfigPath)
-		assert.Equal(t, na, data.ExperimentDisplay)
-		assert.Equal(t, na, data.AcceleratorsDisplay)
+		_, _, err := fetchRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5)
+		require.Error(t, err)
 	})
 }
 
diff --git a/experimental/air/cmd/joblist.go b/experimental/air/cmd/joblist.go
index 4c9e05e9b8c..c6f94e1a69f 100644
--- a/experimental/air/cmd/joblist.go
+++ b/experimental/air/cmd/joblist.go
@@ -2,6 +2,7 @@ package aircmd
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
@@ -9,6 +10,7 @@ import (
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/client"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -211,6 +213,33 @@ func fetchJobRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64
 	return &run, nil
 }
 
+// fetchRun fetches a run once and returns it in both shapes: the typed SDK
+// jobs.Run the display path consumes, and the raw jobRun that preserves the
+// ai_runtime_task the typed model drops. The single runs/get body is unmarshalled
+// into both, avoiding a second identical roundtrip.
+func fetchRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (*jobs.Run, *jobRun, error) {
+	apiClient, err := client.New(w.Config)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create API client: %w", err)
+	}
+
+	var body json.RawMessage
+	query := map[string]any{"run_id": runID, "expand_tasks": true}
+	if err := apiClient.Do(ctx, http.MethodGet, jobsRunsGetPath, nil, nil, query, &body); err != nil {
+		return nil, nil, err
+	}
+
+	var typed jobs.Run
+	if err := json.Unmarshal(body, &typed); err != nil {
+		return nil, nil, fmt.Errorf("failed to parse run %d: %w", runID, err)
+	}
+	var raw jobRun
+	if err := json.Unmarshal(body, &raw); err != nil {
+		return nil, nil, fmt.Errorf("failed to parse run %d: %w", runID, err)
+	}
+	return &typed, &raw, nil
+}
+
 // hydrateJobRuns fetches the given run ids concurrently via runs/get, preserving
 // input order. runs/get enforces per-run view ACLs, so an id the caller can't
 // view (403) or that has been purged (404) is dropped; any other error is

From 00be4b25913f317d6e9a0048ec0a58526ab22c29 Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Wed, 8 Jul 2026 04:54:33 +0000
Subject: [PATCH 4/4] experimental/air: yamlfmt the get-ai-runtime training
 config

yamlfmt normalizes the block scalar to `command: |-` (strip trailing
newline); the CLI echoes the config back, so regenerate the acceptance
golden to match. Fixes the `task fmt` CI gate.

Co-authored-by: Isaac
---
 acceptance/experimental/air/get-ai-runtime/output.txt           | 2 +-
 acceptance/experimental/air/get-ai-runtime/training_config.yaml | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/acceptance/experimental/air/get-ai-runtime/output.txt b/acceptance/experimental/air/get-ai-runtime/output.txt
index 150f394801f..c47790eaa31 100644
--- a/acceptance/experimental/air/get-ai-runtime/output.txt
+++ b/acceptance/experimental/air/get-ai-runtime/output.txt
@@ -8,7 +8,7 @@
 │  compute:                                                      │
 │    accelerator_type: a10                                       │
 │    num_accelerators: 1                                         │
-│  command: |                                                    │
+│  command: |-                                                   │
 │    for i in $(seq 1 10); do                                    │
 │      echo "step $i"                                            │
 │    done                                                        │
diff --git a/acceptance/experimental/air/get-ai-runtime/training_config.yaml b/acceptance/experimental/air/get-ai-runtime/training_config.yaml
index 82102bef3c0..5f6060ecbbc 100644
--- a/acceptance/experimental/air/get-ai-runtime/training_config.yaml
+++ b/acceptance/experimental/air/get-ai-runtime/training_config.yaml
@@ -2,7 +2,7 @@ experiment_name: my-exp
 compute:
   accelerator_type: a10
   num_accelerators: 1
-command: |
+command: |-
   for i in $(seq 1 10); do
     echo "step $i"
   done