Fix AIR live log polling correctness

acceptance/experimental/air/help/output.txt

@@ -61,7 +61,7 @@ Flags:
       --minutes int          Fetch only logs from the last N minutes
       --node int             Fetch logs from this node
       --retry int            View logs from a specific retry attempt; -1 means latest (default -1)
-      --tail int             For completed runs, print the last N log lines (default 10000)
+      --tail int             Print the last N existing lines before following, or from a completed run (default 10000)
 
 Global Flags:
       --debug            enable debug logging

acceptance/experimental/air/logs/output.txt

@@ -22,9 +22,11 @@ CUDA out of memory
 >>> [CLI] experimental air logs 123 --tail 1
 CUDA out of memory
 
-=== logs with --tail 0 prints no log lines
+=== logs with --tail 0 is rejected
 >>> [CLI] experimental air logs 123 --tail 0
-No logs available for run 123. Run terminated in state SUCCESS
+Error: invalid --tail 0: must be positive
+
+Exit code: 1
 
 === logs from a specific retry
 >>> [CLI] experimental air logs 123 --retry 0

acceptance/experimental/air/logs/script

@@ -10,8 +10,8 @@ trace $CLI experimental air logs 123 --minutes 30
 title "logs with --tail"
 trace $CLI experimental air logs 123 --tail 1
 
-title "logs with --tail 0 prints no log lines"
-trace $CLI experimental air logs 123 --tail 0
+title "logs with --tail 0 is rejected"
+errcode trace $CLI experimental air logs 123 --tail 0
 
 title "logs from a specific retry"
 trace $CLI experimental air logs 123 --retry 0

experimental/air/cmd/logbricklens.go

@@ -18,6 +18,7 @@ const bricklensLogsPathFmt = "/api/2.0/ai-training/workflows/by-run-id/%d/logs"
 type logRecord struct {
 	// TimeUnixNano may arrive as a JSON number or string.
 	TimeUnixNano json.Number `json:"time_unix_nano"`
+	RecordID     string      `json:"record_id"`
 	Body         string      `json:"body"`
 	NodeIndex    int         `json:"node_index"`
 }

experimental/air/cmd/logs.go

@@ -36,7 +36,7 @@ func newLogsCommand() *cobra.Command {
 	}
 
 	cmd.Flags().IntVar(&node, "node", 0, "Fetch logs from this node")
-	cmd.Flags().IntVar(&tail, "tail", 0, "For completed runs, print the last N log lines (default 10000)")
+	cmd.Flags().IntVar(&tail, "tail", 0, "Print the last N existing lines before following, or from a completed run (default 10000)")
 	cmd.Flags().IntVar(&minutes, "minutes", 0, "Fetch only logs from the last N minutes")
 	cmd.Flags().IntVar(&retry, "retry", -1, "View logs from a specific retry attempt; -1 means latest")
 	cmd.Flags().StringVar(&downloadTo, "download-to", "", "Download all logs to this directory instead of printing")
@@ -75,7 +75,7 @@ func newLogsCommand() *cobra.Command {
 			return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false,
 				errors.New("cannot combine --tail with --minutes: --tail selects by line count, --minutes by time window"))
 		}
-		if tail < 0 {
+		if cmd.Flags().Changed("tail") && tail <= 0 {
 			return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false,
 				fmt.Errorf("invalid --tail %d: must be positive", tail))
 		}
@@ -98,8 +98,7 @@ func newLogsCommand() *cobra.Command {
 				fmt.Errorf("invalid JOB_RUN_ID %q: must be a positive integer", args[0]))
 		}
 
-		// -1 signals "unset" (use the default cap); an explicit --tail 0 stays 0
-		// and prints no log lines.
+		// -1 signals "unset" (use the default cap).
 		tailLines := -1
 		if cmd.Flags().Changed("tail") {
 			tailLines = tail
@@ -116,14 +115,15 @@ func newLogsCommand() *cobra.Command {
 		}
 
 		err = runLogs(streamCtx, cmd, logRequest{
-			runID:         runID,
-			node:          node,
-			nodeSet:       cmd.Flags().Changed("node"),
-			attempt:       retry,
-			windowMinutes: minutes,
-			tailLines:     tailLines,
-			downloadTo:    downloadTo,
-			jsonOutput:    root.OutputType(cmd) == flags.OutputJSON,
+			runID:            runID,
+			node:             node,
+			nodeSet:          cmd.Flags().Changed("node"),
+			attempt:          retry,
+			windowMinutes:    minutes,
+			tailLines:        tailLines,
+			boundInitialLogs: minutes == 0,
+			downloadTo:       downloadTo,
+			jsonOutput:       root.OutputType(cmd) == flags.OutputJSON,
 		})
 		if downloadTo != "" || root.OutputType(cmd) == flags.OutputJSON {
 			return err

experimental/air/cmd/logs_test.go

@@ -68,6 +68,12 @@ func TestLogsFlagValidation(t *testing.T) {
 			flags:   map[string]string{"tail": "-1"},
 			wantMsg: "invalid --tail",
 		},
+		{
+			name:    "zero tail rejected",
+			args:    []string{"5"},
+			flags:   map[string]string{"tail": "0"},
+			wantMsg: "invalid --tail",
+		},
 		{
 			name:    "negative minutes rejected",
 			args:    []string{"5"},

experimental/air/cmd/logstream.go

@@ -33,6 +33,10 @@ const (
 	// bricklensEmptyMLflowProbeEveryNPolls throttles the active-run MLflow
 	// existence probe while Bricklens has returned no records.
 	bricklensEmptyMLflowProbeEveryNPolls = 10
+	// Event time and ingestion order can differ, so live polls overlap.
+	bricklensLogLookback = 30 * time.Second
+	// A terminal job can have records that are not searchable yet.
+	bricklensTerminalEmptyPolls = 3
 )
 
 // statusMessageType tags a client-facing message packed into
@@ -68,6 +72,9 @@ func normalizeStatusMessage(raw string) string {
 // shrink it.
 var retryCheckInterval = 3 * time.Second
 
+// bricklensPollInterval is independent of slower retry and discovery backoffs.
+var bricklensPollInterval = time.Second
+
 // errBricklensFeatureDisabled signals the caller to fall back to MLflow: Bricklens
 // is gated off (FEATURE_DISABLED), not deployed (ENDPOINT_NOT_FOUND / 404),
 // persistently failing, or served every request successfully but never returned a
@@ -97,7 +104,9 @@ type logRequest struct {
 	// past retry of an active run: that attempt's logs are immutable, so streaming
 	// would poll forever waiting for the run (not the attempt) to finish.
 	staticView bool
-	jsonOutput bool
+	// boundInitialLogs limits existing output before following an active run.
+	boundInitialLogs 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.
@@ -265,10 +274,11 @@ type bricklensStreamer struct {
 	req       logRequest
 	status    logRunStatus
 
-	fromSec      int64
-	lastNano     int64
-	firstLogSeen bool
-	seen         *seenSet
+	fromSec        int64
+	streamStartSec int64
+	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
@@ -330,6 +340,7 @@ func (st *bricklensStreamer) reportStatusChange() {
 func (st *bricklensStreamer) run() (bool, error) {
 	now := time.Now()
 	st.fromSec = st.req.fromSeconds(st.status, now)
+	st.streamStartSec = st.fromSec
 
 	// A past retry's logs are immutable: render a one-shot tail rather than
 	// following the still-active run, which would poll forever.
@@ -354,8 +365,10 @@ func (st *bricklensStreamer) run() (bool, error) {
 	// redundant spinner updates.
 	statusRefreshCounter := 0
 	emptyStreamPolls := 0
+	terminalEmptyPolls := 0
 	lastSpinnerText := ""
 	for {
+		pollStarted := time.Now()
 		if !firstIteration {
 			status, err := resolveRunStatus(st.ctx, st.w, st.req.runID)
 			if err != nil {
@@ -395,20 +408,38 @@ func (st *bricklensStreamer) run() (bool, error) {
 			statusRefreshCounter++
 		}
 
-		// A run already terminal on the first iteration renders as a tail (most
-		// recent N lines). An active run streams everything with dedup, so a run
-		// that terminates while we watch doesn't re-print the boundary second.
+		// `air logs` starts an active attachment with a bounded tail; subsequent
+		// polls follow live with dedup so the overlap does not re-print it.
+		var emitted int
 		var err error
-		if firstIteration && terminal {
-			err = st.drainTail(toSec)
+		if firstIteration && (terminal || st.req.boundInitialLogs) {
+			err = st.drainTail(toSec, !terminal)
+			if err == nil && !terminal {
+				// Baseline omitted history without suppressing newer records.
+				_, err = st.drainPages(toSec, true)
+			}
 		} else {
-			err = st.drainPages(toSec)
+			emitted, err = st.drainPages(toSec, false)
 		}
 		if err != nil {
 			return false, err
 		}
 
 		if terminal {
+			if !firstIteration {
+				if emitted == 0 {
+					terminalEmptyPolls++
+				} else {
+					terminalEmptyPolls = 0
+				}
+				if terminalEmptyPolls < bricklensTerminalEmptyPolls {
+					if err := waitForNextBricklensPoll(st.ctx, pollStarted); err != nil {
+						return false, err
+					}
+					firstIteration = false
+					continue
+				}
+			}
 			if !st.firstLogSeen {
 				// A successful but empty Bricklens stream isn't proof the run has no
 				// logs; they may be in MLflow (as --download-to reads). Fall back
@@ -427,12 +458,16 @@ func (st *bricklensStreamer) run() (bool, error) {
 		}
 
 		firstIteration = false
-		if err := sleepOrCancel(st.ctx, retryCheckInterval); err != nil {
+		if err := waitForNextBricklensPoll(st.ctx, pollStarted); err != nil {
 			return false, err
 		}
 	}
 }
 
+func waitForNextBricklensPoll(ctx context.Context, pollStarted time.Time) error {
+	return sleepOrCancel(ctx, max(time.Duration(0), bricklensPollInterval-time.Since(pollStarted)))
+}
+
 // sleepOrCancel waits for d, or returns early with the context error if the
 // context is cancelled (e.g. Ctrl-C) so the poll loop exits promptly.
 func sleepOrCancel(ctx context.Context, d time.Duration) error {
@@ -449,7 +484,7 @@ func sleepOrCancel(ctx context.Context, d time.Duration) error {
 // drainStatic renders a single tail pass without following the run. Success
 // reflects the run's current result state (empty while active).
 func (st *bricklensStreamer) drainStatic(toSec int64) (bool, error) {
-	if err := st.drainTail(toSec); err != nil {
+	if err := st.drainTail(toSec, false); err != nil {
 		return false, err
 	}
 	if !st.firstLogSeen {
@@ -474,7 +509,7 @@ func (req logRequest) tailTarget() int {
 // drainTail emits the most-recent `target` records oldest-first. Bricklens
 // returns records newest-first, so it pages until it has `target`, keeps the
 // newest `target`, and reverses to chronological order.
-func (st *bricklensStreamer) drainTail(toSec int64) error {
+func (st *bricklensStreamer) drainTail(toSec int64, remember bool) error {
 	target := st.req.tailTarget()
 	if target <= 0 {
 		return nil
@@ -500,36 +535,42 @@ func (st *bricklensStreamer) drainTail(toSec int64) error {
 	}
 	for _, c := range slices.Backward(collected) {
 		st.emit(c.Body)
+		if remember {
+			st.seen.add(c)
+			st.lastNano = max(st.lastNano, c.nano())
+		}
+	}
+	if remember && st.lastNano != 0 {
+		st.fromSec = max(st.streamStartSec, st.lastNano/1_000_000_000-int64(bricklensLogLookback/time.Second))
 	}
 	return nil
 }
 
-// drainPages exhausts all pages from the current from-second in ascending order,
-// deduping against the seen-set so a re-queried boundary second is not
-// re-printed, then advances fromSec to the newest record's floor-second.
-func (st *bricklensStreamer) drainPages(toSec int64) error {
+// drainPages exhausts all pages from the current from-second in ascending order.
+// Live polls retain a bounded overlap so late records remain visible. The
+// initial baseline suppresses omitted history through the tail's watermark.
+func (st *bricklensStreamer) drainPages(toSec int64, suppressInitialHistory bool) (int, error) {
+	emitted := 0
+	initialTailNano := st.lastNano
+	var maximumEvictedNano int64
 	var pageToken string
 	for {
 		resp, err := st.requestPage(pageToken, toSec, 0, true)
 		if err != nil {
-			return err
+			return 0, err
 		}
 
 		for _, rec := range resp.LogRecords {
 			nano := rec.nano()
-			if nano != 0 {
-				// Skip a record older than the last emitted one to keep output
-				// monotonic (out of order, or a re-queried boundary record).
-				if st.lastNano != 0 && nano < st.lastNano {
-					continue
-				}
-				if st.seen.has(nano, rec.Body) {
-					continue
-				}
+			if st.seen.has(rec) {
+				continue
+			}
+			if !suppressInitialHistory || (nano != 0 && nano > initialTailNano) {
+				st.emit(rec.Body)
+				emitted++
 			}
-			st.emit(rec.Body)
+			maximumEvictedNano = max(maximumEvictedNano, st.seen.add(rec))
 			if nano != 0 {
-				st.seen.add(nano, rec.Body)
 				st.lastNano = max(st.lastNano, nano)
 			}
 		}
@@ -541,9 +582,13 @@ func (st *bricklensStreamer) drainPages(toSec int64) error {
 	}
 
 	if st.lastNano != 0 {
-		st.fromSec = st.lastNano / 1_000_000_000
+		st.fromSec = max(st.streamStartSec, st.lastNano/1_000_000_000-int64(bricklensLogLookback/time.Second))
+		st.seen.removeBefore(st.fromSec * 1_000_000_000)
+		if maximumEvictedNano != 0 {
+			st.fromSec = max(st.fromSec, maximumEvictedNano/1_000_000_000+1)
+		}
 	}
-	return nil
+	return emitted, nil
 }
 
 // requestPage fetches one page, retrying transient failures up to