AIR CLI: implement air run --watch (stream logs until completion)

acceptance/experimental/air/run-submit/output.txt

@@ -4,6 +4,8 @@
 Submitted run 555
 View at: [DATABRICKS_URL]/jobs/runs/555
 
+Tip: use --watch to stream logs until the run completes.
+
 === the ai_runtime_task carries the code_source_path
 >>> print_requests.py //api/2.2/jobs/runs/submit
 {

acceptance/experimental/air/run/output.txt

@@ -33,11 +33,9 @@ Error: compute.num_accelerators must be positive, got 0
 
 Exit code: 1
 
-=== watch not yet supported
+=== watch is ignored with dry-run (nothing is submitted)
 >>> [CLI] experimental air run -f valid.yaml --dry-run --watch
-Error: --watch is not yet supported
-
-Exit code: 1
+Dry run: configuration for "smoke-test" is valid; not submitting.
 
 === code_source config passes validation
 >>> [CLI] experimental air run -f with-code-source.yaml --dry-run

acceptance/experimental/air/run/script

@@ -13,8 +13,8 @@ errcode trace $CLI experimental air run -f valid.yaml --dry-run --override bogus
 title "override still runs schema validation"
 errcode trace $CLI experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=0
 
-title "watch not yet supported"
-errcode trace $CLI experimental air run -f valid.yaml --dry-run --watch
+title "watch is ignored with dry-run (nothing is submitted)"
+trace $CLI experimental air run -f valid.yaml --dry-run --watch
 
 title "code_source config passes validation"
 trace $CLI experimental air run -f with-code-source.yaml --dry-run

experimental/air/cmd/logstream.go

@@ -55,6 +55,10 @@ type logRequest struct {
 	// would poll forever waiting for the run (not the attempt) to finish.
 	staticView bool
 	jsonOutput bool
+	// onStatusChange, when set, is called on each lifecycle transition while
+	// following the run (current, previous display states). Used by
+	// `air run --watch -o json` to emit STATUS events.
+	onStatusChange func(current, previous string)
 }
 
 // logRunStatus is the subset of a run's state the log path needs, resolved once
@@ -192,6 +196,8 @@ type bricklensStreamer struct {
 	lastNano     int64
 	firstLogSeen bool
 	seen         *seenSet
+	// previousState is the last display state reported to onStatusChange.
+	previousState string
 	// onFirstLog, when set, is called once just before the first log line is
 	// emitted — used to stop the "waiting for run to start" spinner before any
 	// log byte reaches stdout.
@@ -200,6 +206,20 @@ type bricklensStreamer struct {
 	updateSpinner func(string)
 }
 
+// reportStatusChange fires onStatusChange when the run's display state differs
+// from the last reported one.
+func (st *bricklensStreamer) reportStatusChange() {
+	if st.req.onStatusChange == nil {
+		return
+	}
+	current := st.status.displayState()
+	if current == st.previousState {
+		return
+	}
+	st.req.onStatusChange(current, st.previousState)
+	st.previousState = current
+}
+
 func (st *bricklensStreamer) run() (bool, error) {
 	now := time.Now()
 	st.fromSec = st.req.fromSeconds(st.status, now)
@@ -244,6 +264,8 @@ func (st *bricklensStreamer) run() (bool, error) {
 			st.status = status
 		}
 
+		st.reportStatusChange()
+
 		terminal := st.status.terminal()
 		toSec := st.req.toSeconds(st.status)
 

experimental/air/cmd/logstream_support.go

@@ -71,6 +71,78 @@ func printLogEvent(out io.Writer, eventType string, node int, line string) {
 	fmt.Fprintln(out, string(b))
 }
 
+// submittedEvent is the JSONL event `air run --watch -o json` emits before the
+// streamed log events, so a consumer sees the run id immediately.
+type submittedEvent struct {
+	Type         string `json:"type"`
+	TS           string `json:"ts"`
+	RunID        string `json:"run_id"`
+	DashboardURL string `json:"dashboard_url"`
+}
+
+// printSubmittedEvent writes the SUBMITTED JSONL event.
+func printSubmittedEvent(out io.Writer, runID, dashboardURL string) {
+	b, err := json.Marshal(submittedEvent{
+		Type:         "SUBMITTED",
+		TS:           time.Now().UTC().Format(time.RFC3339),
+		RunID:        runID,
+		DashboardURL: dashboardURL,
+	})
+	if err != nil {
+		return
+	}
+	fmt.Fprintln(out, string(b))
+}
+
+// statusEvent is a JSONL event emitted on each lifecycle transition while
+// following a run with --watch.
+type statusEvent struct {
+	Type     string `json:"type"`
+	TS       string `json:"ts"`
+	Status   string `json:"status"`
+	Previous string `json:"previous_status,omitempty"`
+}
+
+// printStatusEvent writes a STATUS JSONL event for a lifecycle transition.
+func printStatusEvent(out io.Writer, current, previous string) {
+	b, err := json.Marshal(statusEvent{
+		Type:     "STATUS",
+		TS:       time.Now().UTC().Format(time.RFC3339),
+		Status:   current,
+		Previous: previous,
+	})
+	if err != nil {
+		return
+	}
+	fmt.Fprintln(out, string(b))
+}
+
+// terminalEvent is the closing envelope `air run --watch -o json` emits after
+// streaming, carrying the run's terminal status.
+type terminalEvent struct {
+	V    int       `json:"v"`
+	TS   string    `json:"ts"`
+	Data runResult `json:"data"`
+}
+
+// printTerminalEvent writes the closing terminal-status envelope, matching the
+// shape of renderEnvelope(runResult).
+func printTerminalEvent(out io.Writer, runID, status, dashboardURL string) {
+	b, err := json.Marshal(terminalEvent{
+		V:  envelopeVersion,
+		TS: time.Now().UTC().Format(time.RFC3339),
+		Data: runResult{
+			Status:       status,
+			RunID:        runID,
+			DashboardURL: dashboardURL,
+		},
+	})
+	if err != nil {
+		return
+	}
+	fmt.Fprintln(out, string(b))
+}
+
 // emitLogLine writes one log line: raw in text mode, or a JSONL LOG event under
 // --json. In --json mode a line matching a fatal-failure pattern also emits an
 // ALERT event first, giving an agent an immediate actionable signal.

experimental/air/cmd/run.go

@@ -1,14 +1,15 @@
 package aircmd
 
 import (
-	"errors"
+	"context"
 	"fmt"
 	"strconv"
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/flags"
+	"github.com/databricks/databricks-sdk-go"
 	"github.com/spf13/cobra"
 )
 
@@ -57,11 +58,6 @@ The workload is described by a YAML config file (see --file).`,
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 
-		// --watch's pipeline is not ported yet; reject rather than silently ignore it.
-		if watch {
-			return errors.New("--watch is not yet supported")
-		}
-
 		cfg, err := loadRunConfigWithOverrides(ctx, file, overrides)
 		if err != nil {
 			return err
@@ -82,13 +78,63 @@ The workload is described by a YAML config file (see --file).`,
 		}
 
 		runIDStr := strconv.FormatInt(runID, 10)
-		if root.OutputType(cmd) == flags.OutputText {
+		jsonOut := root.OutputType(cmd) == flags.OutputJSON
+
+		if !watch {
+			if !jsonOut {
+				cmdio.LogString(ctx, "Submitted run "+runIDStr)
+				cmdio.LogString(ctx, "View at: "+dashboardURL)
+				cmdio.LogString(ctx, "\nTip: use --watch to stream logs until the run completes.")
+				return nil
+			}
+			return renderEnvelope(ctx, runResult{Status: "SUBMITTED", RunID: runIDStr, DashboardURL: dashboardURL})
+		}
+
+		// --watch: stream the submitted run's logs until it reaches a terminal
+		// state, then exit with the run's outcome. This is the same pipeline as
+		// `air logs <run>` (Bricklens with MLflow fallback).
+		req := logRequest{
+			runID:      runID,
+			attempt:    -1,
+			tailLines:  -1,
+			jsonOutput: jsonOut,
+		}
+
+		if !jsonOut {
 			cmdio.LogString(ctx, "Submitted run "+runIDStr)
 			cmdio.LogString(ctx, "View at: "+dashboardURL)
-			return nil
+			cmdio.LogString(ctx, "Monitoring run and streaming logs...")
+			return runLogs(ctx, cmd, req)
+		}
+
+		// --json: emit SUBMITTED first (so a consumer sees the run id immediately),
+		// STATUS events on each lifecycle transition, and a closing terminal-status
+		// envelope after streaming. Mirrors the Python CLI's --watch JSONL contract.
+		out := cmd.OutOrStdout()
+		printSubmittedEvent(out, runIDStr, dashboardURL)
+		req.onStatusChange = func(current, previous string) {
+			printStatusEvent(out, current, previous)
 		}
-		return renderEnvelope(ctx, runResult{Status: "SUBMITTED", RunID: runIDStr, DashboardURL: dashboardURL})
+		err = runLogs(ctx, cmd, req)
+
+		// Re-resolve the run for the closing envelope. STATUS events only fire on
+		// the Bricklens path, so the terminal status must come from the run's
+		// actual state — correct whether Bricklens or the MLflow fallback served
+		// the logs.
+		printTerminalEvent(out, runIDStr, watchTerminalStatus(ctx, w, runID), dashboardURL)
+		return err
 	}
 
 	return cmd
 }
+
+// watchTerminalStatus resolves a watched run's final display state for the
+// closing --watch envelope. The run is terminal once streaming returns; if the
+// status can't be re-fetched, "UNKNOWN" is reported rather than guessing.
+func watchTerminalStatus(ctx context.Context, w *databricks.WorkspaceClient, runID int64) string {
+	status, err := resolveRunStatus(ctx, w, runID)
+	if err != nil {
+		return "UNKNOWN"
+	}
+	return status.displayState()
+}