From dcf4b3c68c9cc7500502d9427a4897c7f376d4ed Mon Sep 17 00:00:00 2001
From: Sankalp-Mittal <sankalp.mittal@databricks.com>
Date: Wed, 4 Nov 2026 12:27:42 +0000
Subject: [PATCH 1/7] Send heartbeats for open DMS deployment versions

DMS marks a version in progress at CreateVersion. If the CLI dies before
CompleteVersion, nothing refreshes the version and the UI shows it as
deploying forever. The service will expose a per-version last update time,
refreshed by UpdateOperation and Heartbeat, and the UI will show a version
as interrupted once that time is over two minutes old.

The CLI now keeps that time fresh while a version is open. A heartbeat is
sent only after 30s with no successful UpdateOperation, so a fast deploy
sends none. Heartbeats are best effort: a failed one is retried quickly,
then reported as a warning, then tried again at the next interval. It
never fails or stops the deploy.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 bundle/direct/dstate/state.go |  11 +++-
 libs/dms/heartbeat.go         | 121 ++++++++++++++++++++++++++++++++++
 libs/dms/heartbeat_test.go    | 115 ++++++++++++++++++++++++++++++++
 libs/dms/operation_buffer.go  |   7 +-
 4 files changed, 252 insertions(+), 2 deletions(-)
 create mode 100644 libs/dms/heartbeat.go
 create mode 100644 libs/dms/heartbeat_test.go

diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go
index ac8f8226c4b..8a74b072b2d 100644
--- a/bundle/direct/dstate/state.go
+++ b/bundle/direct/dstate/state.go
@@ -98,6 +98,9 @@ type DeploymentState struct {
 	// history, in which case InitializeOperationBuffer installs it once the version exists.
 	operationBuffer *dms.OperationBuffer
 
+	// heartbeat keeps the open version's last update time fresh. Set together with operationBuffer.
+	heartbeat *dms.Heartbeat
+
 	// versionCompleted makes CompleteVersion a no-op after the first call, so a deferred safety-net
 	// completion after an explicit one does nothing.
 	versionCompleted bool
@@ -177,10 +180,12 @@ type WALEntry struct {
 // (after approval) - which is why it is not an Open option. It also records the id a first deploy
 // just created, which Open could not know, so CompleteVersion later has it.
 func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploymentID string, versionID int) {
-	buf := dms.StartOperationBuffer(ctx, deploymentID, versionID)
+	heartbeat := dms.StartHeartbeat(ctx, deploymentID, versionID)
+	buf := dms.StartOperationBuffer(ctx, deploymentID, versionID, heartbeat)
 
 	db.mu.Lock()
 	defer db.mu.Unlock()
+	db.heartbeat = heartbeat
 	db.operationBuffer = buf
 	db.DeploymentID = deploymentID
 	db.VersionID = versionID
@@ -222,8 +227,12 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b
 	}
 	db.versionCompleted = true
 	deploymentID, versionID := db.DeploymentID, db.VersionID
+	heartbeat := db.heartbeat
 	db.mu.Unlock()
 
+	// Stopped first: a heartbeat racing the completion would be rejected and warn for nothing.
+	heartbeat.Stop()
+
 	// A recording failure fails the version even when the caller counted the deploy a success: the
 	// service does not then hold everything the WAL does. Finalize already drained and surfaced it;
 	// this reads the drained buffer's error so a destroy whose uploads failed keeps its record.
diff --git a/libs/dms/heartbeat.go b/libs/dms/heartbeat.go
new file mode 100644
index 00000000000..73a57b28a95
--- /dev/null
+++ b/libs/dms/heartbeat.go
@@ -0,0 +1,121 @@
+package dms
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/diag"
+	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/cli/libs/logdiag"
+	"github.com/databricks/databricks-sdk-go/service/bundledeployments"
+)
+
+// heartbeatInterval is how long a version may go without an update before a heartbeat is sent.
+// The UI shows a version as interrupted once its last update is over 2 minutes old.
+const heartbeatInterval = 30 * time.Second
+
+// heartbeatRetryDelays are the quick retries after a failed heartbeat, before waiting a full interval.
+var heartbeatRetryDelays = []time.Duration{time.Second, 2 * time.Second, 4 * time.Second}
+
+// Heartbeat keeps an open deployment version's last update time fresh, which is how the UI tells
+// a running deploy from an interrupted one. CreateVersion and every successful write update it
+// too, so a heartbeat is sent only after an interval without one. Best effort: a failure can only
+// make the deploy show as interrupted, so it never fails the deploy.
+type Heartbeat struct {
+	send func(ctx context.Context) error
+
+	renewed  chan struct{}
+	stop     chan struct{}
+	done     chan struct{}
+	stopOnce func()
+}
+
+// StartHeartbeat starts heartbeating the version the caller just created. CreateVersion counts
+// as the first update, so the first heartbeat goes out an interval later.
+func StartHeartbeat(ctx context.Context, deploymentID string, versionNum int) *Heartbeat {
+	name := VersionName(deploymentID, versionNum)
+	send := func(ctx context.Context) error {
+		_, err := cmdctx.WorkspaceClient(ctx).BundleDeployments.Heartbeat(ctx, bundledeployments.HeartbeatRequest{Name: name})
+		return err
+	}
+	return startHeartbeat(ctx, send)
+}
+
+func startHeartbeat(ctx context.Context, send func(ctx context.Context) error) *Heartbeat {
+	h := &Heartbeat{
+		send:    send,
+		renewed: make(chan struct{}, 1),
+		stop:    make(chan struct{}),
+		done:    make(chan struct{}),
+	}
+	h.stopOnce = sync.OnceFunc(func() {
+		close(h.stop)
+		<-h.done
+	})
+	go h.run(ctx)
+	return h
+}
+
+// renew records that a successful write updated the version, which the service counts as a
+// heartbeat, so the next heartbeat is pushed back a full interval.
+func (h *Heartbeat) renew() {
+	select {
+	case h.renewed <- struct{}{}:
+	default:
+	}
+}
+
+func (h *Heartbeat) run(ctx context.Context) {
+	defer close(h.done)
+
+	next := time.NewTimer(heartbeatInterval)
+	defer next.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-h.stop:
+			return
+		case <-h.renewed:
+			next.Reset(heartbeatInterval)
+		case <-next.C:
+			h.beat(ctx)
+			next.Reset(heartbeatInterval)
+		}
+	}
+}
+
+// beat sends a heartbeat, retrying quickly a few times before leaving it to the next interval.
+func (h *Heartbeat) beat(ctx context.Context) {
+	err := h.send(ctx)
+	for _, delay := range heartbeatRetryDelays {
+		if err == nil || ctx.Err() != nil {
+			break
+		}
+		log.Warnf(ctx, "Failed to send deployment heartbeat, retrying in %s: %s", delay, err)
+		select {
+		case <-time.After(delay):
+		case <-ctx.Done():
+			return
+		case <-h.stop:
+			return
+		}
+		err = h.send(ctx)
+	}
+
+	// A cancelled run needs no warning.
+	if err != nil && ctx.Err() == nil {
+		logdiag.LogDiag(ctx, diag.Diagnostic{
+			Severity: diag.Warning,
+			Summary:  fmt.Sprintf("Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err),
+		})
+	}
+}
+
+// Stop ends the heartbeat and waits for it. Safe to call twice.
+func (h *Heartbeat) Stop() {
+	h.stopOnce()
+}
diff --git a/libs/dms/heartbeat_test.go b/libs/dms/heartbeat_test.go
new file mode 100644
index 00000000000..b7d71c4b2f9
--- /dev/null
+++ b/libs/dms/heartbeat_test.go
@@ -0,0 +1,115 @@
+package dms
+
+import (
+	"context"
+	"errors"
+	"sync/atomic"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"github.com/databricks/cli/libs/logdiag"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// Timing is not reachable from acceptance tests, so the heartbeat is tested here on synctest's
+// fake clock, against a send that counts calls instead of the service.
+
+type fakeSend struct {
+	calls atomic.Int32
+	err   error
+}
+
+func (f *fakeSend) send(context.Context) error {
+	f.calls.Add(1)
+	return f.err
+}
+
+func heartbeatContext(t *testing.T) context.Context {
+	ctx := logdiag.InitContext(t.Context())
+	logdiag.SetCollect(ctx, true)
+	return ctx
+}
+
+func TestHeartbeatBeatsWhileIdle(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		f := &fakeSend{}
+		h := startHeartbeat(heartbeatContext(t), f.send)
+		defer h.Stop()
+
+		// CreateVersion counts as the first update, so nothing goes out before the interval.
+		time.Sleep(heartbeatInterval - time.Second)
+		synctest.Wait()
+		assert.Equal(t, int32(0), f.calls.Load())
+
+		time.Sleep(2 * time.Second)
+		synctest.Wait()
+		assert.Equal(t, int32(1), f.calls.Load())
+
+		time.Sleep(heartbeatInterval)
+		synctest.Wait()
+		assert.Equal(t, int32(2), f.calls.Load())
+	})
+}
+
+func TestHeartbeatSkippedAfterWrite(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		f := &fakeSend{}
+		h := startHeartbeat(heartbeatContext(t), f.send)
+		defer h.Stop()
+
+		time.Sleep(20 * time.Second)
+		h.renew()
+
+		// The write counts as an update, so the beat due at 30s moves to 50s.
+		time.Sleep(25 * time.Second)
+		synctest.Wait()
+		assert.Equal(t, int32(0), f.calls.Load())
+
+		time.Sleep(10 * time.Second)
+		synctest.Wait()
+		assert.Equal(t, int32(1), f.calls.Load())
+	})
+}
+
+func TestHeartbeatRetriesThenWaitsForInterval(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		ctx := heartbeatContext(t)
+		f := &fakeSend{err: errors.New("connection refused")}
+		h := startHeartbeat(ctx, f.send)
+		defer h.Stop()
+
+		// The beat at 30s and its retries 1s, 2s and 4s apart.
+		time.Sleep(heartbeatInterval + 8*time.Second)
+		synctest.Wait()
+		assert.Equal(t, int32(1+len(heartbeatRetryDelays)), f.calls.Load())
+
+		diags := logdiag.FlushCollected(ctx)
+		require.Len(t, diags, 1)
+		assert.Equal(t, "Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in 30s: connection refused", diags[0].Summary)
+
+		// Then nothing more until a full interval has passed.
+		time.Sleep(heartbeatInterval - 2*time.Second)
+		synctest.Wait()
+		assert.Equal(t, int32(1+len(heartbeatRetryDelays)), f.calls.Load())
+
+		// Checked before its first retry, 1s later.
+		time.Sleep(1500 * time.Millisecond)
+		synctest.Wait()
+		assert.Equal(t, int32(2+len(heartbeatRetryDelays)), f.calls.Load())
+	})
+}
+
+func TestHeartbeatStop(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		f := &fakeSend{}
+		h := startHeartbeat(heartbeatContext(t), f.send)
+
+		h.Stop()
+		h.Stop()
+		time.Sleep(10 * time.Minute)
+		synctest.Wait()
+		assert.Equal(t, int32(0), f.calls.Load())
+	})
+}
diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go
index efecb3a330a..cd80ab5e40b 100644
--- a/libs/dms/operation_buffer.go
+++ b/libs/dms/operation_buffer.go
@@ -21,6 +21,9 @@ type OperationBuffer struct {
 	deploymentID string
 	versionNum   int
 
+	// heartbeat is told of each successful write, which the service counts as a heartbeat.
+	heartbeat *Heartbeat
+
 	// The buffer. queue holds bundle state keys, and pending the newest update per key, so a
 	// second write for a resource replaces the first.
 	queue     chan string
@@ -53,10 +56,11 @@ type OperationBuffer struct {
 
 // StartOperationBuffer opens the buffer for the version the caller just created. The version
 // must already exist: operations record under it, and nothing here creates it.
-func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum int) *OperationBuffer {
+func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum int, heartbeat *Heartbeat) *OperationBuffer {
 	b := &OperationBuffer{
 		deploymentID: deploymentID,
 		versionNum:   versionNum,
+		heartbeat:    heartbeat,
 		queue:        make(chan string, bufferedOperations),
 		done:         make(chan struct{}),
 		pending:      make(map[string]OperationUpdate),
@@ -158,6 +162,7 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio
 
 	// The next write for this resource echoes the sequence id this one earned.
 	b.sequenceIDs[key] = result.SequenceId
+	b.heartbeat.renew()
 	return nil
 }
 

From 57340f2b39431cc6d69760a594a2cda05b915751 Mon Sep 17 00:00:00 2001
From: Sankalp-Mittal <sankalp.mittal@databricks.com>
Date: Wed, 4 Nov 2026 15:24:47 +0000
Subject: [PATCH 2/7] Simplify the heartbeat loop

Stop is a context cancel instead of a stop channel, a done channel and a
sync.OnceFunc guard. The quick retries after a failed heartbeat are gone:
the SDK already retries transient errors inside each call, and the next
attempt is at most 30s away against a 2-minute UI threshold.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 libs/dms/heartbeat.go      | 103 +++++++++++--------------------------
 libs/dms/heartbeat_test.go |  18 ++-----
 2 files changed, 35 insertions(+), 86 deletions(-)

diff --git a/libs/dms/heartbeat.go b/libs/dms/heartbeat.go
index 73a57b28a95..8a2da222ce6 100644
--- a/libs/dms/heartbeat.go
+++ b/libs/dms/heartbeat.go
@@ -3,12 +3,10 @@ package dms
 import (
 	"context"
 	"fmt"
-	"sync"
 	"time"
 
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/diag"
-	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/logdiag"
 	"github.com/databricks/databricks-sdk-go/service/bundledeployments"
 )
@@ -17,105 +15,64 @@ import (
 // The UI shows a version as interrupted once its last update is over 2 minutes old.
 const heartbeatInterval = 30 * time.Second
 
-// heartbeatRetryDelays are the quick retries after a failed heartbeat, before waiting a full interval.
-var heartbeatRetryDelays = []time.Duration{time.Second, 2 * time.Second, 4 * time.Second}
-
 // Heartbeat keeps an open deployment version's last update time fresh, which is how the UI tells
 // a running deploy from an interrupted one. CreateVersion and every successful write update it
 // too, so a heartbeat is sent only after an interval without one. Best effort: a failure can only
-// make the deploy show as interrupted, so it never fails the deploy.
+// make the deploy show as interrupted, so it is reported and the next interval tries again.
 type Heartbeat struct {
-	send func(ctx context.Context) error
-
-	renewed  chan struct{}
-	stop     chan struct{}
-	done     chan struct{}
-	stopOnce func()
+	renewed chan struct{}
+	cancel  context.CancelFunc
 }
 
 // StartHeartbeat starts heartbeating the version the caller just created. CreateVersion counts
 // as the first update, so the first heartbeat goes out an interval later.
 func StartHeartbeat(ctx context.Context, deploymentID string, versionNum int) *Heartbeat {
 	name := VersionName(deploymentID, versionNum)
-	send := func(ctx context.Context) error {
+	return startHeartbeat(ctx, func(ctx context.Context) error {
 		_, err := cmdctx.WorkspaceClient(ctx).BundleDeployments.Heartbeat(ctx, bundledeployments.HeartbeatRequest{Name: name})
 		return err
-	}
-	return startHeartbeat(ctx, send)
-}
-
-func startHeartbeat(ctx context.Context, send func(ctx context.Context) error) *Heartbeat {
-	h := &Heartbeat{
-		send:    send,
-		renewed: make(chan struct{}, 1),
-		stop:    make(chan struct{}),
-		done:    make(chan struct{}),
-	}
-	h.stopOnce = sync.OnceFunc(func() {
-		close(h.stop)
-		<-h.done
 	})
-	go h.run(ctx)
-	return h
 }
 
-// renew records that a successful write updated the version, which the service counts as a
-// heartbeat, so the next heartbeat is pushed back a full interval.
-func (h *Heartbeat) renew() {
-	select {
-	case h.renewed <- struct{}{}:
-	default:
-	}
+// startHeartbeat takes the send call so tests can count calls instead of reaching a server.
+func startHeartbeat(ctx context.Context, send func(context.Context) error) *Heartbeat {
+	ctx, cancel := context.WithCancel(ctx)
+	h := &Heartbeat{renewed: make(chan struct{}, 1), cancel: cancel}
+	go h.run(ctx, send)
+	return h
 }
 
-func (h *Heartbeat) run(ctx context.Context) {
-	defer close(h.done)
-
-	next := time.NewTimer(heartbeatInterval)
-	defer next.Stop()
+func (h *Heartbeat) run(ctx context.Context, send func(context.Context) error) {
+	timer := time.NewTimer(heartbeatInterval)
 	for {
 		select {
 		case <-ctx.Done():
 			return
-		case <-h.stop:
-			return
 		case <-h.renewed:
-			next.Reset(heartbeatInterval)
-		case <-next.C:
-			h.beat(ctx)
-			next.Reset(heartbeatInterval)
+		case <-timer.C:
+			// A request cut short by Stop is not a failure.
+			if err := send(ctx); err != nil && ctx.Err() == nil {
+				logdiag.LogDiag(ctx, diag.Diagnostic{
+					Severity: diag.Warning,
+					Summary:  fmt.Sprintf("Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err),
+				})
+			}
 		}
+		timer.Reset(heartbeatInterval)
 	}
 }
 
-// beat sends a heartbeat, retrying quickly a few times before leaving it to the next interval.
-func (h *Heartbeat) beat(ctx context.Context) {
-	err := h.send(ctx)
-	for _, delay := range heartbeatRetryDelays {
-		if err == nil || ctx.Err() != nil {
-			break
-		}
-		log.Warnf(ctx, "Failed to send deployment heartbeat, retrying in %s: %s", delay, err)
-		select {
-		case <-time.After(delay):
-		case <-ctx.Done():
-			return
-		case <-h.stop:
-			return
-		}
-		err = h.send(ctx)
-	}
-
-	// A cancelled run needs no warning.
-	if err != nil && ctx.Err() == nil {
-		logdiag.LogDiag(ctx, diag.Diagnostic{
-			Severity: diag.Warning,
-			Summary:  fmt.Sprintf("Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err),
-		})
+// renew records that a successful write updated the version, which the service counts as a
+// heartbeat, so the next heartbeat is pushed back a full interval. Non-blocking: one pending
+// signal is enough.
+func (h *Heartbeat) renew() {
+	select {
+	case h.renewed <- struct{}{}:
+	default:
 	}
 }
 
-// Stop ends the heartbeat and waits for it. Safe to call twice.
+// Stop ends the heartbeat, cutting short one in flight.
 func (h *Heartbeat) Stop() {
-	h.stopOnce()
+	h.cancel()
 }
diff --git a/libs/dms/heartbeat_test.go b/libs/dms/heartbeat_test.go
index b7d71c4b2f9..b7de3d8bb1d 100644
--- a/libs/dms/heartbeat_test.go
+++ b/libs/dms/heartbeat_test.go
@@ -73,31 +73,24 @@ func TestHeartbeatSkippedAfterWrite(t *testing.T) {
 	})
 }
 
-func TestHeartbeatRetriesThenWaitsForInterval(t *testing.T) {
+func TestHeartbeatFailureWarnsAndRetriesNextInterval(t *testing.T) {
 	synctest.Test(t, func(t *testing.T) {
 		ctx := heartbeatContext(t)
 		f := &fakeSend{err: errors.New("connection refused")}
 		h := startHeartbeat(ctx, f.send)
 		defer h.Stop()
 
-		// The beat at 30s and its retries 1s, 2s and 4s apart.
-		time.Sleep(heartbeatInterval + 8*time.Second)
+		time.Sleep(heartbeatInterval + time.Second)
 		synctest.Wait()
-		assert.Equal(t, int32(1+len(heartbeatRetryDelays)), f.calls.Load())
+		assert.Equal(t, int32(1), f.calls.Load())
 
 		diags := logdiag.FlushCollected(ctx)
 		require.Len(t, diags, 1)
 		assert.Equal(t, "Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in 30s: connection refused", diags[0].Summary)
 
-		// Then nothing more until a full interval has passed.
-		time.Sleep(heartbeatInterval - 2*time.Second)
-		synctest.Wait()
-		assert.Equal(t, int32(1+len(heartbeatRetryDelays)), f.calls.Load())
-
-		// Checked before its first retry, 1s later.
-		time.Sleep(1500 * time.Millisecond)
+		time.Sleep(heartbeatInterval)
 		synctest.Wait()
-		assert.Equal(t, int32(2+len(heartbeatRetryDelays)), f.calls.Load())
+		assert.Equal(t, int32(2), f.calls.Load())
 	})
 }
 
@@ -106,7 +99,6 @@ func TestHeartbeatStop(t *testing.T) {
 		f := &fakeSend{}
 		h := startHeartbeat(heartbeatContext(t), f.send)
 
-		h.Stop()
 		h.Stop()
 		time.Sleep(10 * time.Minute)
 		synctest.Wait()

From b99be930399f78e26325b5d7c55c2d693e1dfc37 Mon Sep 17 00:00:00 2001
From: Sankalp-Mittal <sankalp.mittal@databricks.com>
Date: Thu, 5 Nov 2026 10:57:50 +0000
Subject: [PATCH 3/7] Update heartbeat comments

Co-authored-by: Isaac <no-reply@databricks.com>
---
 libs/dms/heartbeat.go | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

diff --git a/libs/dms/heartbeat.go b/libs/dms/heartbeat.go
index 8a2da222ce6..70bcfcc69f2 100644
--- a/libs/dms/heartbeat.go
+++ b/libs/dms/heartbeat.go
@@ -11,21 +11,17 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/bundledeployments"
 )
 
-// heartbeatInterval is how long a version may go without an update before a heartbeat is sent.
-// The UI shows a version as interrupted once its last update is over 2 minutes old.
+// heartbeatInterval is the maximum time between consecutive updates.
 const heartbeatInterval = 30 * time.Second
 
-// Heartbeat keeps an open deployment version's last update time fresh, which is how the UI tells
-// a running deploy from an interrupted one. CreateVersion and every successful write update it
-// too, so a heartbeat is sent only after an interval without one. Best effort: a failure can only
-// make the deploy show as interrupted, so it is reported and the next interval tries again.
+// Heartbeat updates the update_time of an open deployment version. This is also updated by
+// CreateVersion and every successful write.
 type Heartbeat struct {
 	renewed chan struct{}
 	cancel  context.CancelFunc
 }
 
-// StartHeartbeat starts heartbeating the version the caller just created. CreateVersion counts
-// as the first update, so the first heartbeat goes out an interval later.
+// StartHeartbeat starts the heartbeat. We count CreateVersion as the first heartbeat.
 func StartHeartbeat(ctx context.Context, deploymentID string, versionNum int) *Heartbeat {
 	name := VersionName(deploymentID, versionNum)
 	return startHeartbeat(ctx, func(ctx context.Context) error {
@@ -62,9 +58,7 @@ func (h *Heartbeat) run(ctx context.Context, send func(context.Context) error) {
 	}
 }
 
-// renew records that a successful write updated the version, which the service counts as a
-// heartbeat, so the next heartbeat is pushed back a full interval. Non-blocking: one pending
-// signal is enough.
+// renew records that a successful write updated the version.
 func (h *Heartbeat) renew() {
 	select {
 	case h.renewed <- struct{}{}:
@@ -72,7 +66,7 @@ func (h *Heartbeat) renew() {
 	}
 }
 
-// Stop ends the heartbeat, cutting short one in flight.
+// Stop ends the heartbeat.
 func (h *Heartbeat) Stop() {
 	h.cancel()
 }

From 85b6141fd817fe89f7506b244dc29520123794ee Mon Sep 17 00:00:00 2001
From: Sankalp-Mittal <sankalp.mittal@databricks.com>
Date: Thu, 5 Nov 2026 15:38:22 +0000
Subject: [PATCH 4/7] Split the heartbeat send into separate statements

Co-authored-by: Isaac <no-reply@databricks.com>
---
 libs/dms/heartbeat.go | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/libs/dms/heartbeat.go b/libs/dms/heartbeat.go
index 70bcfcc69f2..b2dab0a47e3 100644
--- a/libs/dms/heartbeat.go
+++ b/libs/dms/heartbeat.go
@@ -46,13 +46,18 @@ func (h *Heartbeat) run(ctx context.Context, send func(context.Context) error) {
 			return
 		case <-h.renewed:
 		case <-timer.C:
-			// A request cut short by Stop is not a failure.
-			if err := send(ctx); err != nil && ctx.Err() == nil {
-				logdiag.LogDiag(ctx, diag.Diagnostic{
-					Severity: diag.Warning,
-					Summary:  fmt.Sprintf("Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err),
-				})
+			err := send(ctx)
+			if err == nil {
+				break
 			}
+			// We only get an error if Stop was called
+			if ctx.Err() != nil {
+				return
+			}
+			logdiag.LogDiag(ctx, diag.Diagnostic{
+				Severity: diag.Warning,
+				Summary:  fmt.Sprintf("Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err),
+			})
 		}
 		timer.Reset(heartbeatInterval)
 	}

From 8c36af64b6f6a8f86320d96022b8ea5746e20540 Mon Sep 17 00:00:00 2001
From: Sankalp-Mittal <sankalp.mittal@databricks.com>
Date: Mon, 9 Nov 2026 13:00:08 +0000
Subject: [PATCH 5/7] Log failed heartbeats at info level

A failed heartbeat is not config-related and needs no user action, so it
is logged with log.Infof instead of a warning diagnostic.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 libs/dms/heartbeat.go      |  9 ++-------
 libs/dms/heartbeat_test.go | 23 +++++------------------
 2 files changed, 7 insertions(+), 25 deletions(-)

diff --git a/libs/dms/heartbeat.go b/libs/dms/heartbeat.go
index b2dab0a47e3..c8602fc8b71 100644
--- a/libs/dms/heartbeat.go
+++ b/libs/dms/heartbeat.go
@@ -2,12 +2,10 @@ package dms
 
 import (
 	"context"
-	"fmt"
 	"time"
 
 	"github.com/databricks/cli/libs/cmdctx"
-	"github.com/databricks/cli/libs/diag"
-	"github.com/databricks/cli/libs/logdiag"
+	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go/service/bundledeployments"
 )
 
@@ -54,10 +52,7 @@ func (h *Heartbeat) run(ctx context.Context, send func(context.Context) error) {
 			if ctx.Err() != nil {
 				return
 			}
-			logdiag.LogDiag(ctx, diag.Diagnostic{
-				Severity: diag.Warning,
-				Summary:  fmt.Sprintf("Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err),
-			})
+			log.Infof(ctx, "Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err)
 		}
 		timer.Reset(heartbeatInterval)
 	}
diff --git a/libs/dms/heartbeat_test.go b/libs/dms/heartbeat_test.go
index b7de3d8bb1d..3b6dc320567 100644
--- a/libs/dms/heartbeat_test.go
+++ b/libs/dms/heartbeat_test.go
@@ -8,9 +8,7 @@ import (
 	"testing/synctest"
 	"time"
 
-	"github.com/databricks/cli/libs/logdiag"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
 // Timing is not reachable from acceptance tests, so the heartbeat is tested here on synctest's
@@ -26,16 +24,10 @@ func (f *fakeSend) send(context.Context) error {
 	return f.err
 }
 
-func heartbeatContext(t *testing.T) context.Context {
-	ctx := logdiag.InitContext(t.Context())
-	logdiag.SetCollect(ctx, true)
-	return ctx
-}
-
 func TestHeartbeatBeatsWhileIdle(t *testing.T) {
 	synctest.Test(t, func(t *testing.T) {
 		f := &fakeSend{}
-		h := startHeartbeat(heartbeatContext(t), f.send)
+		h := startHeartbeat(t.Context(), f.send)
 		defer h.Stop()
 
 		// CreateVersion counts as the first update, so nothing goes out before the interval.
@@ -56,7 +48,7 @@ func TestHeartbeatBeatsWhileIdle(t *testing.T) {
 func TestHeartbeatSkippedAfterWrite(t *testing.T) {
 	synctest.Test(t, func(t *testing.T) {
 		f := &fakeSend{}
-		h := startHeartbeat(heartbeatContext(t), f.send)
+		h := startHeartbeat(t.Context(), f.send)
 		defer h.Stop()
 
 		time.Sleep(20 * time.Second)
@@ -73,21 +65,16 @@ func TestHeartbeatSkippedAfterWrite(t *testing.T) {
 	})
 }
 
-func TestHeartbeatFailureWarnsAndRetriesNextInterval(t *testing.T) {
+func TestHeartbeatFailureRetriesNextInterval(t *testing.T) {
 	synctest.Test(t, func(t *testing.T) {
-		ctx := heartbeatContext(t)
 		f := &fakeSend{err: errors.New("connection refused")}
-		h := startHeartbeat(ctx, f.send)
+		h := startHeartbeat(t.Context(), f.send)
 		defer h.Stop()
 
 		time.Sleep(heartbeatInterval + time.Second)
 		synctest.Wait()
 		assert.Equal(t, int32(1), f.calls.Load())
 
-		diags := logdiag.FlushCollected(ctx)
-		require.Len(t, diags, 1)
-		assert.Equal(t, "Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in 30s: connection refused", diags[0].Summary)
-
 		time.Sleep(heartbeatInterval)
 		synctest.Wait()
 		assert.Equal(t, int32(2), f.calls.Load())
@@ -97,7 +84,7 @@ func TestHeartbeatFailureWarnsAndRetriesNextInterval(t *testing.T) {
 func TestHeartbeatStop(t *testing.T) {
 	synctest.Test(t, func(t *testing.T) {
 		f := &fakeSend{}
-		h := startHeartbeat(heartbeatContext(t), f.send)
+		h := startHeartbeat(t.Context(), f.send)
 
 		h.Stop()
 		time.Sleep(10 * time.Minute)

From e3885ec2b4425917f13f0ecab0f75f97568b62af Mon Sep 17 00:00:00 2001
From: Sankalp-Mittal <sankalp.mittal@databricks.com>
Date: Mon, 9 Nov 2026 13:38:52 +0000
Subject: [PATCH 6/7] Let the operation buffer own the heartbeat

DeploymentState held both the buffer and the heartbeat for the open
version and had to wire them together. The buffer now starts the
heartbeat itself and stops it in Stop, so DeploymentState keeps one
field. The heartbeat still runs past Drain until CompleteVersion.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 bundle/direct/dstate/state.go     | 11 +++--------
 libs/dms/heartbeat.go             |  4 ++--
 libs/dms/operation_buffer.go      | 12 ++++++++++--
 libs/dms/operation_buffer_test.go |  9 ++-------
 4 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go
index 94f85f4da67..44725de8203 100644
--- a/bundle/direct/dstate/state.go
+++ b/bundle/direct/dstate/state.go
@@ -107,9 +107,6 @@ type DeploymentState struct {
 	// history, in which case InitializeOperationBuffer installs it once the version exists.
 	operationBuffer *dms.OperationBuffer
 
-	// heartbeat keeps the open version's last update time fresh. Set together with operationBuffer.
-	heartbeat *dms.Heartbeat
-
 	// versionCompleted makes CompleteVersion a no-op after the first call, so a deferred safety-net
 	// completion after an explicit one does nothing.
 	versionCompleted bool
@@ -190,8 +187,7 @@ func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploy
 		BestEffort:     db.mode == WritingDeploymentHistoryShadow,
 		OperationCount: operationCount,
 	}
-	db.heartbeat = dms.StartHeartbeat(ctx, deploymentID, versionID)
-	db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options, db.heartbeat)
+	db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options)
 	db.DeploymentID = deploymentID
 	db.VersionID = versionID
 }
@@ -237,12 +233,11 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b
 	}
 	db.versionCompleted = true
 	deploymentID, versionID := db.DeploymentID, db.VersionID
-	heartbeat := db.heartbeat
 	db.reset()
 	db.mu.Unlock()
 
-	// Stopped first: a heartbeat racing the completion would be rejected and warn for nothing.
-	heartbeat.Stop()
+	// Stopped first: a heartbeat racing the completion would be rejected for nothing.
+	buf.Stop()
 
 	// A recording failure fails the version even when the caller counted the deploy a success: the
 	// service does not then hold everything the WAL does. Finalize already drained and surfaced it;
diff --git a/libs/dms/heartbeat.go b/libs/dms/heartbeat.go
index c8602fc8b71..3af601e1373 100644
--- a/libs/dms/heartbeat.go
+++ b/libs/dms/heartbeat.go
@@ -19,8 +19,8 @@ type Heartbeat struct {
 	cancel  context.CancelFunc
 }
 
-// StartHeartbeat starts the heartbeat. We count CreateVersion as the first heartbeat.
-func StartHeartbeat(ctx context.Context, deploymentID string, versionNum int) *Heartbeat {
+// startVersionHeartbeat starts the heartbeat. We count CreateVersion as the first heartbeat.
+func startVersionHeartbeat(ctx context.Context, deploymentID string, versionNum int) *Heartbeat {
 	name := VersionName(deploymentID, versionNum)
 	return startHeartbeat(ctx, func(ctx context.Context) error {
 		_, err := cmdctx.WorkspaceClient(ctx).BundleDeployments.Heartbeat(ctx, bundledeployments.HeartbeatRequest{Name: name})
diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go
index 9c62691e5e5..5a1edce6669 100644
--- a/libs/dms/operation_buffer.go
+++ b/libs/dms/operation_buffer.go
@@ -32,7 +32,8 @@ type OperationBuffer struct {
 	deploymentID string
 	versionNum   int
 
-	// heartbeat is told of each successful write, which the service counts as a heartbeat.
+	// heartbeat keeps the version's update_time fresh until Stop, past Drain. Each successful
+	// write renews it, since the service counts a write as a heartbeat.
 	heartbeat *Heartbeat
 
 	// The buffer. queue holds bundle state keys, and pending the newest update per key, so a
@@ -72,7 +73,9 @@ type OperationBuffer struct {
 
 // StartOperationBuffer opens the buffer for the version the caller just created. The version
 // must already exist: operations record under it, and nothing here creates it.
-func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum int, options BufferOptions, heartbeat *Heartbeat) *OperationBuffer {
+func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum int, options BufferOptions) *OperationBuffer {
+	// Started on ctx, not the writer's context below, so it keeps running after Drain.
+	heartbeat := startVersionHeartbeat(ctx, deploymentID, versionNum)
 	// History can stop independently without canceling the deployment.
 	ctx, cancel := context.WithCancel(ctx)
 	capacity := bufferedOperations
@@ -241,6 +244,11 @@ func (b *OperationBuffer) Drain(ctx context.Context) error {
 	}
 }
 
+// Stop ends the heartbeat, the buffer's only background work left after Drain.
+func (b *OperationBuffer) Stop() {
+	b.heartbeat.Stop()
+}
+
 // setErr keeps the first recording error for Drain to report.
 func (b *OperationBuffer) setErr(err error) {
 	b.mu.Lock()
diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go
index 546933abc0d..0e019284360 100644
--- a/libs/dms/operation_buffer_test.go
+++ b/libs/dms/operation_buffer_test.go
@@ -30,11 +30,6 @@ func queued() *OperationBuffer {
 	}
 }
 
-// idleHeartbeat is a heartbeat whose sends do nothing, for buffers that never write.
-func idleHeartbeat(ctx context.Context) *Heartbeat {
-	return startHeartbeat(ctx, func(context.Context) error { return nil })
-}
-
 func stateUpdate(t *testing.T, name string) OperationUpdate {
 	t.Helper()
 	update, err := NewStateUpdate("id-1", json.RawMessage(`{"state":{"name":"`+name+`"}}`), false)
@@ -79,7 +74,7 @@ func TestBufferFailsOnOversizedState(t *testing.T) {
 
 func TestDrainIsIdempotent(t *testing.T) {
 	// An idle writer must finish when draining begins.
-	b := StartOperationBuffer(t.Context(), "dep-1", 1, BufferOptions{}, idleHeartbeat(t.Context()))
+	b := StartOperationBuffer(t.Context(), "dep-1", 1, BufferOptions{})
 	require.NoError(t, b.Drain(t.Context()))
 	require.NoError(t, b.Drain(t.Context()))
 }
@@ -124,7 +119,7 @@ func TestHistoryQueueDoesNotBlockWhenWriterStops(t *testing.T) {
 		const operationCount = 1000
 		ctx, cancel := context.WithCancel(t.Context())
 		defer cancel()
-		b := StartOperationBuffer(ctx, "dep-1", 1, BufferOptions{BestEffort: true, OperationCount: operationCount}, idleHeartbeat(ctx))
+		b := StartOperationBuffer(ctx, "dep-1", 1, BufferOptions{BestEffort: true, OperationCount: operationCount})
 		synctest.Wait()
 		cancel()
 		<-b.done

From b7209141a3e4871ead149a80522ba9ade0e439d3 Mon Sep 17 00:00:00 2001
From: Sankalp-Mittal <sankalp.mittal@databricks.com>
Date: Mon, 9 Nov 2026 14:14:34 +0000
Subject: [PATCH 7/7] Reset the heartbeat when a write starts

Both UpdateOperation and Heartbeat update the version row, so running
them at the same time can conflict on the server. Renewing when a write
starts, instead of when it succeeds, means a heartbeat only overlaps a
write that has been in flight for a full interval.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 libs/dms/heartbeat.go        | 2 +-
 libs/dms/operation_buffer.go | 9 ++++++---
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/libs/dms/heartbeat.go b/libs/dms/heartbeat.go
index 3af601e1373..4c239bcc647 100644
--- a/libs/dms/heartbeat.go
+++ b/libs/dms/heartbeat.go
@@ -58,7 +58,7 @@ func (h *Heartbeat) run(ctx context.Context, send func(context.Context) error) {
 	}
 }
 
-// renew records that a successful write updated the version.
+// renew records that a write started, so the next heartbeat waits a full interval.
 func (h *Heartbeat) renew() {
 	select {
 	case h.renewed <- struct{}{}:
diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go
index 5a1edce6669..963e5ddfb39 100644
--- a/libs/dms/operation_buffer.go
+++ b/libs/dms/operation_buffer.go
@@ -32,8 +32,8 @@ type OperationBuffer struct {
 	deploymentID string
 	versionNum   int
 
-	// heartbeat keeps the version's update_time fresh until Stop, past Drain. Each successful
-	// write renews it, since the service counts a write as a heartbeat.
+	// heartbeat keeps the version's update_time fresh until Stop, past Drain. Each write renews
+	// it as it starts, since the service counts a write as a heartbeat.
 	heartbeat *Heartbeat
 
 	// The buffer. queue holds bundle state keys, and pending the newest update per key, so a
@@ -198,6 +198,10 @@ func (b *OperationBuffer) run(ctx context.Context) {
 
 // write sends one update, at the sequence id the resource is at.
 func (b *OperationBuffer) write(ctx context.Context, key string, update OperationUpdate) error {
+	// Renewed when the write starts, not when it succeeds, so a heartbeat only overlaps a write
+	// that has been in flight for a full interval. Both update the version row.
+	b.heartbeat.renew()
+
 	w := cmdctx.WorkspaceClient(ctx)
 	result, err := w.BundleDeployments.UpdateOperation(ctx, bundledeployments.UpdateOperationRequest{
 		Name:       OperationName(b.deploymentID, b.versionNum, key),
@@ -210,7 +214,6 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio
 
 	// The next write for this resource echoes the sequence id this one earned.
 	b.sequenceIDs[key] = result.SequenceId
-	b.heartbeat.renew()
 	return nil
 }