Exit air run --watch when the streamed node succeeds

.nextchanges/air/watch-exit-on-completion.md

@@ -0,0 +1 @@
+* `air run --watch` now exits as soon as the command succeeds on the node it streams, instead of waiting for the job run to be marked terminal. ([#7031](https://github.com/databricks/cli/pull/7031))

cmd/air/logcompletion.go

@@ -0,0 +1,53 @@
+package aircmd
+
+import "strings"
+
+// The remote entry script prints this block from its exit handler, as the last
+// thing it does on each node:
+//
+//	================================
+//	Execution Summary
+//	================================
+//	Total time: 32s
+//	User code execution time: 32s
+//	Completed at Fri Oct  9 09:46:18 PM UTC 2026
+//	================================
+//
+// "Completed at" is printed only for exit code 0; a failure prints "ERROR: Script
+// failed ..." instead, which is not treated as completion because the run may be
+// retried. Jobs reports the run terminal tens of seconds after the block, so
+// `air run --watch` uses it to stop as soon as the streamed node has succeeded.
+const (
+	summaryHeader         = "Execution Summary"
+	summaryRulePrefix     = "===="
+	summaryTotalPrefix    = "Total time: "
+	summaryUserCodePrefix = "User code execution time: "
+	summarySuccessPrefix  = "Completed at "
+)
+
+// completionMatcher recognizes the entry script's success summary in one node's
+// log lines, fed oldest-first across polls. Only a consecutive block matches, so
+// user output that happens to contain one of the lines does not.
+type completionMatcher struct {
+	step      int
+	completed bool
+}
+
+func (m *completionMatcher) observe(line string) {
+	line = strings.TrimSpace(line)
+	switch {
+	case line == summaryHeader:
+		m.step = 1
+	case m.step == 1 && strings.HasPrefix(line, summaryRulePrefix):
+		m.step = 2
+	case m.step == 2 && strings.HasPrefix(line, summaryTotalPrefix):
+		m.step = 3
+	case m.step == 3 && strings.HasPrefix(line, summaryUserCodePrefix):
+		// Optional line between the total time and the outcome.
+	case m.step == 3 && strings.HasPrefix(line, summarySuccessPrefix):
+		m.completed = true
+		m.step = 0
+	default:
+		m.step = 0
+	}
+}

cmd/air/logcompletion_test.go

@@ -0,0 +1,219 @@
+package aircmd
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"slices"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+var successSummary = []string{
+	"================================",
+	"Execution Summary",
+	"================================",
+	"Total time: 32s",
+	"User code execution time: 32s",
+	"Completed at Fri Oct  9 09:46:18 PM UTC 2026",
+	"================================",
+}
+
+var failureSummary = []string{
+	"================================",
+	"Execution Summary",
+	"================================",
+	"Total time: 32s",
+	"User code execution time: 32s",
+	"ERROR: Script failed with exit code 1 after 32s",
+	"================================",
+}
+
+func TestCompletionMatcher(t *testing.T) {
+	tests := []struct {
+		name  string
+		lines []string
+		want  bool
+	}{
+		{name: "success summary", lines: successSummary, want: true},
+		{
+			name:  "success summary without user code time",
+			lines: slices.Delete(slices.Clone(successSummary), 4, 5),
+			want:  true,
+		},
+		{name: "failure summary", lines: failureSummary, want: false},
+		{
+			name:  "trailing whitespace",
+			lines: []string{"Execution Summary\r", "====\r", "Total time: 1s", "Completed at now\r"},
+			want:  true,
+		},
+		{
+			name:  "interrupted block",
+			lines: []string{"Execution Summary", "====", "step 5", "Total time: 1s", "Completed at now"},
+			want:  false,
+		},
+		{name: "completed line alone", lines: []string{"Completed at epoch 3"}, want: false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var m completionMatcher
+			for _, line := range tt.lines {
+				m.observe(line)
+			}
+			assert.Equal(t, tt.want, m.completed)
+		})
+	}
+}
+
+// logsBody renders a log endpoint response for the streamed node. Records are
+// timestamped one second apart from 1700000000s and keyed by position and
+// content, so a poll that re-serves earlier lines dedups them while new lines,
+// including the summary's repeated rules, stay distinct.
+func logsBody(t *testing.T, lines []string) string {
+	t.Helper()
+	type record struct {
+		RecordID     string `json:"record_id"`
+		TimeUnixNano int64  `json:"time_unix_nano"`
+		Body         string `json:"body"`
+	}
+	records := make([]record, len(lines))
+	for i, line := range lines {
+		records[i] = record{
+			RecordID:     strconv.Itoa(i) + ":" + line,
+			TimeUnixNano: (1_700_000_000 + int64(i)) * 1_000_000_000,
+			Body:         line,
+		}
+	}
+	b, err := json.Marshal(map[string]any{"log_records": records})
+	assert.NoError(t, err)
+	return string(b)
+}
+
+const (
+	runningRunBody   = `{"run_id":1,"start_time":1700000000000,"state":{"life_cycle_state":"RUNNING"},"tasks":[{"run_id":2,"attempt_number":0}]}`
+	succeededRunBody = `{"run_id":1,"start_time":1700000000000,"end_time":1700000100000,"state":{"life_cycle_state":"TERMINATED","result_state":"SUCCESS"},"tasks":[{"run_id":2,"attempt_number":0}]}`
+	failedRunBody    = `{"run_id":1,"start_time":1700000000000,"end_time":1700000100000,"state":{"life_cycle_state":"TERMINATED","result_state":"FAILED"},"tasks":[{"run_id":2,"attempt_number":0}]}`
+	retryingRunBody  = `{"run_id":1,"start_time":1700000000000,"state":{"life_cycle_state":"RUNNING"},"tasks":[{"run_id":2,"attempt_number":0},{"run_id":3,"attempt_number":1}]}`
+)
+
+// completionServer counts the requests the streamer made to the fake server.
+type completionServer struct {
+	runGets     int
+	logRequests int
+}
+
+// newCompletionServer serves an active run. runBodies are returned by successive
+// runs/get calls (the last repeats); lines returns node 0's log lines for the
+// n-th log request (1-based).
+func newCompletionServer(t *testing.T, runBodies []string, lines func(n int) []string) (*completionServer, *httptest.Server) {
+	t.Helper()
+	original := bricklensPollInterval
+	bricklensPollInterval = time.Millisecond
+	t.Cleanup(func() { bricklensPollInterval = original })
+
+	cs := &completionServer{}
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch {
+		case strings.HasSuffix(r.URL.Path, "/logs"):
+			assert.Equal(t, "0", r.URL.Query().Get("filter.node_index"))
+			cs.logRequests++
+			_, _ = w.Write([]byte(logsBody(t, lines(cs.logRequests))))
+		case r.URL.Path == "/api/2.2/jobs/runs/get":
+			_, _ = w.Write([]byte(runBodies[min(cs.runGets, len(runBodies)-1)]))
+			cs.runGets++
+		default:
+			_, _ = w.Write([]byte(`{}`))
+		}
+	}))
+	t.Cleanup(srv.Close)
+	return cs, srv
+}
+
+func streamForCompletion(t *testing.T, srvURL string, exitOnCompletion bool) (bool, string, error) {
+	t.Helper()
+	// Every scenario ends within a few polls; fail rather than hang if one does not.
+	ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
+	defer cancel()
+	var buf bytes.Buffer
+	ok, err := streamBricklensLogs(ctx, newTestWorkspaceClient(t, srvURL), &buf,
+		logRequest{runID: 1, attempt: -1, tailLines: -1, jsonOutput: true, exitOnCompletion: exitOnCompletion},
+		logRunStatus{lifeCycleState: "RUNNING", startTimeMs: 1_700_000_000_000})
+	return ok, buf.String(), err
+}
+
+func TestStreamBricklensExitsOnSuccessSummary(t *testing.T) {
+	cs, srv := newCompletionServer(t, []string{runningRunBody}, func(int) []string {
+		return append([]string{"training"}, successSummary...)
+	})
+
+	ok, out, err := streamForCompletion(t, srv.URL, true)
+	require.NoError(t, err)
+	assert.True(t, ok)
+	// Exited on the first poll, without Jobs ever reporting the run terminal.
+	assert.Equal(t, 0, cs.runGets)
+	assert.Equal(t, 1, cs.logRequests)
+	assert.Contains(t, out, `"line":"training"`)
+	assert.Contains(t, out, `"line":"Completed at `)
+}
+
+func TestStreamBricklensMatchesSummarySplitAcrossPolls(t *testing.T) {
+	cs, srv := newCompletionServer(t, []string{runningRunBody}, func(n int) []string {
+		lines := append([]string{"training"}, successSummary...)
+		if n == 1 {
+			return lines[:4]
+		}
+		return lines
+	})
+
+	ok, _, err := streamForCompletion(t, srv.URL, true)
+	require.NoError(t, err)
+	assert.True(t, ok)
+	assert.Equal(t, 2, cs.logRequests)
+	assert.Equal(t, 1, cs.runGets)
+}
+
+func TestStreamBricklensIgnoresFailureSummary(t *testing.T) {
+	cs, srv := newCompletionServer(t, []string{runningRunBody, failedRunBody}, func(int) []string {
+		return failureSummary
+	})
+
+	ok, _, err := streamForCompletion(t, srv.URL, true)
+	require.NoError(t, err)
+	// Waited for Jobs: a failed attempt may still be retried.
+	assert.False(t, ok)
+	assert.GreaterOrEqual(t, cs.runGets, 2)
+}
+
+func TestStreamBricklensExitsOnRetrySuccess(t *testing.T) {
+	// Attempt 0 fails, so the stream keeps following the run; the retry then
+	// succeeds and the stream exits without waiting for Jobs.
+	cs, srv := newCompletionServer(t, []string{retryingRunBody}, func(n int) []string {
+		if n == 1 {
+			return failureSummary
+		}
+		return append([]string{"attempt 1"}, successSummary...)
+	})
+
+	ok, _, err := streamForCompletion(t, srv.URL, true)
+	require.NoError(t, err)
+	assert.True(t, ok)
+	assert.Equal(t, 2, cs.logRequests)
+}
+
+func TestStreamBricklensCompletionRequiresOptIn(t *testing.T) {
+	cs, srv := newCompletionServer(t, []string{runningRunBody, succeededRunBody}, func(int) []string {
+		return successSummary
+	})
+
+	ok, _, err := streamForCompletion(t, srv.URL, false)
+	require.NoError(t, err)
+	assert.True(t, ok)
+	assert.GreaterOrEqual(t, cs.runGets, 2)
+}

cmd/air/logstream.go

@@ -114,6 +114,10 @@ type logRequest struct {
 	// `air run --watch -o json` to emit STATUS events.
 	onStatusChange func(current, previous string)
 	retryTracker   *retryTracker
+	// exitOnCompletion ends the stream once the streamed node prints the entry
+	// script's success summary, without waiting for Jobs to report the run
+	// terminal. Used by `air run --watch`.
+	exitOnCompletion bool
 }
 
 // logRunStatus is the subset of a run's state the log path needs, resolved once
@@ -355,6 +359,8 @@ type bricklensStreamer struct {
 	onFirstLog func()
 	// updateSpinner, when set, refreshes the waiting-spinner text each poll.
 	updateSpinner func(string)
+	// completion watches the streamed node's new lines for the success summary.
+	completion completionMatcher
 }
 
 // waitingSpinnerText returns the waiting-spinner text: the server-set STATUS
@@ -518,6 +524,10 @@ func (st *bricklensStreamer) run() (bool, error) {
 			log.Infof(st.ctx, "air logs: run %d finished in state %s", st.req.runID, st.status.displayState())
 			return st.status.succeeded(), nil
 		}
+		if st.req.exitOnCompletion && st.completion.completed {
+			log.Infof(st.ctx, "air logs: run %d succeeded on node %d; not waiting for Jobs to report it terminal", st.req.runID, st.req.node)
+			return true, nil
+		}
 		if !st.firstLogSeen {
 			emptyStreamPolls++
 			if emptyStreamPolls%bricklensEmptyMLflowProbeEveryNPolls == 0 && mlflowLogsExist(st.ctx, st.w, st.req) {
@@ -634,6 +644,7 @@ func (st *bricklensStreamer) drainPages(toSec int64, suppressInitialHistory bool
 			if st.seen.has(rec) {
 				continue
 			}
+			st.completion.observe(rec.Body)
 			if !suppressInitialHistory || (nano != 0 && nano > initialTailNano) {
 				st.emit(rec.Body)
 				emitted++

cmd/air/run.go

@@ -202,15 +202,18 @@ The path must be a separate argument: cobra reserves -h as a boolean, so
 
 		// --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).
+		// `air logs <run>` (Bricklens with MLflow fallback). It exits early once
+		// the streamed node's entry script reports success, since Jobs lags the
+		// user command's exit by tens of seconds.
 		maxRetries := cfg.maxRetries()
 		req := logRequest{
-			runID:        runID,
-			attempt:      -1,
-			tailLines:    -1,
-			jsonOutput:   jsonOut,
-			maxRetries:   &maxRetries,
-			retryTracker: newRetryTracker(logRunStatus{}),
+			runID:            runID,
+			attempt:          -1,
+			tailLines:        -1,
+			jsonOutput:       jsonOut,
+			maxRetries:       &maxRetries,
+			retryTracker:     newRetryTracker(logRunStatus{}),
+			exitOnCompletion: true,
 		}
 
 		watchCtx, stop := notifyInterrupt(ctx)
@@ -258,7 +261,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so
 		// 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)
+		printTerminalEvent(out, runIDStr, watchTerminalStatus(ctx, w, runID, err), dashboardURL)
 		return err
 	}
 
@@ -360,12 +363,17 @@ func printLogsDivider(ctx context.Context, out io.Writer) {
 }
 
 // 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 {
+// closing --watch envelope. If the status can't be re-fetched, "UNKNOWN" is
+// reported rather than guessing.
+func watchTerminalStatus(ctx context.Context, w *databricks.WorkspaceClient, runID int64, streamErr error) string {
 	status, err := resolveRunStatus(ctx, w, runID)
 	if err != nil {
 		return "UNKNOWN"
 	}
+	// A stream that returned success on a still-active run stopped early on the
+	// streamed node's success summary; Jobs reports the run terminal only later.
+	if streamErr == nil && !status.terminal() {
+		return "SUCCESS"
+	}
 	return status.displayState()
 }