AIR CLI Integration: show environment for AI Runtime runs in air get

acceptance/experimental/air/get-ai-runtime/output.txt

@@ -27,7 +27,7 @@
 │  MLflow Run    my-run                                          │
 │  User          user@example.com                                │
 │  Accelerators  1x A10                                          │
-│  Environment   N/A                                             │
+│  Environment   4                                               │
 │                                                                │
 ╰────────────────────────────────────────────────────────────────╯
 

acceptance/experimental/air/get-ai-runtime/test.toml

@@ -9,8 +9,10 @@ MSYS_NO_PATHCONV = "1"
 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).
+# The GetRun response for an ai_runtime_task run. It has no gen_ai_compute_task,
+# so the Environment cell instead comes from the run's environments[].spec (keyed
+# by the task's environment_key), which the typed SDK Run drops — `air get` reads
+# it with a raw request, so this route is hit twice (typed GetRun + that request).
 [[Server]]
 Pattern = "GET /api/2.2/jobs/runs/get"
 Response.Body = '''
@@ -21,12 +23,16 @@ Response.Body = '''
   "start_time": 1700000000000,
   "end_time": 1700000012000,
   "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"},
+  "environments": [
+    {"environment_key": "default", "spec": {"environment_version": "4"}}
+  ],
   "tasks": [
     {
       "task_key": "train",
       "run_id": 456,
       "attempt_number": 0,
       "max_retries": 3,
+      "environment_key": "default",
       "ai_runtime_task": {
         "experiment": "my-exp",
         "deployments": [

acceptance/experimental/air/get/output.txt

@@ -27,7 +27,7 @@
 │  MLflow Run    my-run                                          │
 │  User          user@example.com                                │
 │  Accelerators  1x A10                                          │
-│  Environment   ml-runtime-gpu:1.0                              │
+│  Environment   N/A                                             │
 │                                                                │
 ╰────────────────────────────────────────────────────────────────╯
 

experimental/air/cmd/format.go

@@ -285,19 +285,6 @@ func gpuDisplayName(gpuType string) string {
 	return gpuType
 }
 
-// environment returns the run's runtime image (the training environment), or an
-// empty string if the run has no GenAI-compute task.
-func environment(run *jobs.Run) string {
-	if len(run.Tasks) == 0 {
-		return ""
-	}
-	task := run.Tasks[0].GenAiComputeTask
-	if task == nil {
-		return ""
-	}
-	return task.DlRuntimeImage
-}
-
 // maxRetries returns the configured retry limit for the run's latest task as a
 // display string: "unlimited" for the backend's -1, otherwise the count.
 func maxRetries(run *jobs.Run) string {

experimental/air/cmd/get.go

@@ -236,7 +236,10 @@ func buildGetData(run *jobs.Run) getData {
 	}
 	data.UserDisplay = orNA(run.CreatorUserName)
 	data.AcceleratorsDisplay = orNA(accelerators(run))
-	data.EnvironmentDisplay = orNA(environment(run))
+	// EnvironmentDisplay is resolved at render time: the serverless environment
+	// version needs a raw GetRun read (the typed SDK Run drops it), so it is not
+	// filled here alongside the fields read straight off the run.
+	data.EnvironmentDisplay = na
 	data.MaxRetriesDisplay = maxRetries(run)
 	return data
 }

experimental/air/cmd/render.go

@@ -4,13 +4,15 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"net/http"
 	"strconv"
 	"strings"
 
 	"github.com/charmbracelet/lipgloss"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/client"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/muesli/termenv"
 	"go.yaml.in/yaml/v3"
@@ -22,6 +24,10 @@ const (
 	metadataBoxTitle = "Metadata"
 )
 
+// jobsRunsGetPath is the Jobs GetRun endpoint, called with a raw request to read
+// the run's environments[] block, which the typed SDK Run does not expose.
+const jobsRunsGetPath = "/api/2.2/jobs/runs/get"
+
 // minBoxInnerWidth keeps all boxes a uniform, comfortable width; boxHPad and
 // boxVPad are the horizontal and vertical padding inside each box.
 const (
@@ -91,6 +97,11 @@ func renderRunText(ctx context.Context, out io.Writer, w *databricks.WorkspaceCl
 	renderer, colorOn := cmdio.NewRenderer(ctx, out)
 	p := newPalette(renderer)
 
+	// The serverless environment version lives on the run's environments[].spec,
+	// which the typed SDK Run drops, so read it with a raw request. Empty (a run
+	// with no serverless environment) stays "N/A".
+	data.EnvironmentDisplay = orNA(aiRuntimeEnvironmentVersion(ctx, w, run.RunId))
+
 	view := runView{
 		runID:        data.RunID,
 		dashboardURL: data.DashboardURL,
@@ -143,6 +154,40 @@ func genAIComputeTask(run *jobs.Run) *jobs.GenAiComputeTask {
 	return run.Tasks[0].GenAiComputeTask
 }
 
+// aiRuntimeEnvironmentVersion returns the serverless environment version an
+// ai_runtime run used (e.g. "4"), read from the Jobs GetRun response's
+// environments[] entry keyed by aiRuntimeEnvironmentKey. The typed SDK Run has
+// no environments field, so the value is fetched with a raw request. Best-effort:
+// returns "" (logged) on any error or when the run declares no environment.
+func aiRuntimeEnvironmentVersion(ctx context.Context, w *databricks.WorkspaceClient, runID int64) string {
+	apiClient, err := client.New(w.Config)
+	if err != nil {
+		log.Warnf(ctx, "air get: could not create client to read environment: %v", err)
+		return ""
+	}
+	var resp struct {
+		Environments []struct {
+			EnvironmentKey string `json:"environment_key"`
+			Spec           struct {
+				EnvironmentVersion string `json:"environment_version"`
+			} `json:"spec"`
+		} `json:"environments"`
+	}
+	// For a GET the SDK serializes the request value into query parameters, so
+	// run_id is passed as the request, mirroring the other raw calls in this package.
+	query := map[string]any{"run_id": runID}
+	if err := apiClient.Do(ctx, http.MethodGet, jobsRunsGetPath, nil, nil, query, &resp); err != nil {
+		log.Warnf(ctx, "air get: could not read environment for run %d: %v", runID, err)
+		return ""
+	}
+	for _, e := range resp.Environments {
+		if e.EnvironmentKey == aiRuntimeEnvironmentKey {
+			return e.Spec.EnvironmentVersion
+		}
+	}
+	return ""
+}
+
 // 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 {

experimental/air/cmd/render_test.go

@@ -153,15 +153,15 @@ func TestRenderFields(t *testing.T) {
 		mlflowURL:    "https://h.test/ml/experiments/E1/runs/R1",
 		user:         "user@example.com",
 		accelerators: "1x A10",
-		environment:  "ml-runtime-gpu:1.0",
+		environment:  "4",
 	})
 
 	// Labels are padded to the longest ("Accelerators"), so values align.
 	assert.Contains(t, out, "Run ID        ")
 	assert.Contains(t, out, "Accelerators  1x A10")
 	// Max retries and environment show alongside the other fields.
 	assert.Contains(t, out, "Max Retries   3")
-	assert.Contains(t, out, "Environment   ml-runtime-gpu:1.0")
+	assert.Contains(t, out, "Environment   4")
 	// The status carries its dot prefix.
 	assert.Contains(t, out, "● SUCCESS")
 	// Off a terminal, links render as the bare label (URLs live in JSON output).