From 2f7cfe48522f942eb7250b00f5fd8c045de6e3dd Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Wed, 16 Sep 2026 14:21:33 +0200
Subject: [PATCH 1/8] Add libs/atomicfile and use it for in-place file saves

Introduce libs/atomicfile.Write (same-dir temp + chmod + rename) and route
the hand-rolled atomic writers and a few plain os.WriteFile state saves
through it. The result always gets the caller-specified mode; it does not
inherit the replaced file's permissions.

Pure refactor (already atomic, behavior unchanged):
- bundle/direct/dstate/state.go (deployment resources.json)
- libs/auth/storage/filestore.go (token cache)
- libs/aitools/installer/state.go
- libs/localenv/constraints.go
- libs/clicompat/clicompat.go
- experimental/ssh/internal/sshconfig/knownhosts.go
- cmd/genie/conversations.go (best-effort store; also drops a double-close)

Bug fixes:
- libs/cache/file_cache.go: removed the os.Remove-before-rename hack, which
  deleted the destination first and defeated atomicity (os.Rename already
  replaces an existing file on Windows).
- cmd/sandbox/state.go, cmd/sandbox/sshconfig.go: used a fixed "<path>.tmp"
  temp name, so concurrent CLI runs collided on it; CreateTemp avoids that.

New atomicity (was a plain os.WriteFile):
- bundle/statemgmt/state_pull.go: the state file pulled from the remote is
  now written atomically (closes the "TODO: write + rename").

Co-authored-by: Isaac <no-reply@databricks.com>
---
 bundle/direct/dstate/state.go                 | 22 +----
 bundle/statemgmt/state_pull.go                |  4 +-
 cmd/genie/conversations.go                    | 18 +---
 cmd/sandbox/sshconfig.go                      | 21 +----
 cmd/sandbox/state.go                          | 11 +--
 .../ssh/internal/sshconfig/knownhosts.go      | 20 +----
 libs/aitools/installer/state.go               | 24 +-----
 libs/atomicfile/atomicfile.go                 | 53 ++++++++++++
 libs/atomicfile/atomicfile_test.go            | 83 +++++++++++++++++++
 libs/auth/storage/filestore.go                | 32 +------
 libs/cache/file_cache.go                      | 34 +-------
 libs/clicompat/clicompat.go                   | 26 +-----
 libs/localenv/constraints.go                  | 30 +------
 13 files changed, 165 insertions(+), 213 deletions(-)
 create mode 100644 libs/atomicfile/atomicfile.go
 create mode 100644 libs/atomicfile/atomicfile_test.go

diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go
index aa2d53620d0..3bb30d25861 100644
--- a/bundle/direct/dstate/state.go
+++ b/bundle/direct/dstate/state.go
@@ -19,6 +19,7 @@ import (
 	"github.com/databricks/cli/bundle/deployplan"
 	"github.com/databricks/cli/bundle/statemgmt/resourcestate"
 	"github.com/databricks/cli/internal/build"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/dms"
 	"github.com/databricks/cli/libs/log"
@@ -1003,26 +1004,7 @@ func (db *DeploymentState) unlockedSave() error {
 		return fmt.Errorf("failed to create directory %#v: %w", dir, err)
 	}
 
-	// CreateTemp creates the file with mode 0o600, matching the state file.
-	tmp, err := os.CreateTemp(dir, "."+filepath.Base(db.Path)+".tmp-*")
-	if err != nil {
-		return fmt.Errorf("failed to create temp file for %#v: %w", db.Path, err)
-	}
-	tmpPath := tmp.Name()
-	// Cleans up the temp file on failure; a no-op once the rename succeeded.
-	defer os.Remove(tmpPath)
-
-	if _, err := tmp.Write(data); err != nil {
-		tmp.Close()
-		return fmt.Errorf("failed to write %#v: %w", tmpPath, err)
-	}
-
-	// Close before the rename: on Windows the file must not be open for writing.
-	if err := tmp.Close(); err != nil {
-		return fmt.Errorf("failed to close %#v: %w", tmpPath, err)
-	}
-
-	if err := os.Rename(tmpPath, db.Path); err != nil {
+	if err := atomicfile.Write(db.Path, data, 0o600); err != nil {
 		return fmt.Errorf("failed to save resources state to %#v: %w", db.Path, err)
 	}
 
diff --git a/bundle/statemgmt/state_pull.go b/bundle/statemgmt/state_pull.go
index b382e82dd1c..52f0d858435 100644
--- a/bundle/statemgmt/state_pull.go
+++ b/bundle/statemgmt/state_pull.go
@@ -16,6 +16,7 @@ import (
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config/engine"
 	"github.com/databricks/cli/bundle/deploy"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/log"
@@ -202,8 +203,7 @@ func PullResourcesState(ctx context.Context, b *bundle.Bundle, alwaysPull Always
 			return ctx, winner
 		}
 
-		// TODO: write + rename
-		err = os.WriteFile(localStatePath, winner.Content, 0o600)
+		err = atomicfile.Write(localStatePath, winner.Content, 0o600)
 		if err != nil {
 			logdiag.LogError(ctx, err)
 			return ctx, winner
diff --git a/cmd/genie/conversations.go b/cmd/genie/conversations.go
index cda7f170ec2..b60d9e0e4a5 100644
--- a/cmd/genie/conversations.go
+++ b/cmd/genie/conversations.go
@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"time"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/env"
 )
 
@@ -85,22 +86,7 @@ func saveStore(path string, store conversationStore) {
 	if err := os.MkdirAll(filepath.Dir(path), conversationDirPerm); err != nil {
 		return
 	}
-	tmp, err := os.CreateTemp(filepath.Dir(path), ".genie-conversations-*.tmp")
-	if err != nil {
-		return
-	}
-	defer os.Remove(tmp.Name())
-	defer tmp.Close()
-	if _, err := tmp.Write(raw); err != nil {
-		return
-	}
-	if err := tmp.Chmod(conversationFilePerm); err != nil {
-		return
-	}
-	if err := tmp.Close(); err != nil {
-		return
-	}
-	_ = os.Rename(tmp.Name(), path)
+	_ = atomicfile.Write(path, raw, conversationFilePerm)
 }
 
 // lookupConversationID returns the server conversation id mapped to sessionID on
diff --git a/cmd/sandbox/sshconfig.go b/cmd/sandbox/sshconfig.go
index 41680742c27..d4c54986aed 100644
--- a/cmd/sandbox/sshconfig.go
+++ b/cmd/sandbox/sshconfig.go
@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
 )
@@ -119,15 +120,7 @@ func writeManagedConfig(path, content string) error {
 	if existing, err := os.ReadFile(path); err == nil && bytes.Equal(existing, []byte(content)) {
 		return nil
 	}
-	tmp := path + ".tmp"
-	if err := os.WriteFile(tmp, []byte(content), 0o600); err != nil {
-		return fmt.Errorf("writing %s: %w", tmp, err)
-	}
-	if err := os.Rename(tmp, path); err != nil {
-		_ = os.Remove(tmp)
-		return fmt.Errorf("renaming %s to %s: %w", tmp, path, err)
-	}
-	return nil
+	return atomicfile.Write(path, []byte(content), 0o600)
 }
 
 // ensureMainIncludesManaged makes sure ~/.ssh/config begins with an
@@ -162,15 +155,7 @@ func ensureMainIncludesManaged(mainPath, managedPath string) error {
 		buf.Write(existing)
 	}
 
-	tmp := mainPath + ".tmp"
-	if err := os.WriteFile(tmp, buf.Bytes(), 0o600); err != nil {
-		return fmt.Errorf("writing %s: %w", tmp, err)
-	}
-	if err := os.Rename(tmp, mainPath); err != nil {
-		_ = os.Remove(tmp)
-		return fmt.Errorf("renaming %s to %s: %w", tmp, mainPath, err)
-	}
-	return nil
+	return atomicfile.Write(mainPath, buf.Bytes(), 0o600)
 }
 
 // hasOurMarkedBlock reports whether the given config text already has
diff --git a/cmd/sandbox/state.go b/cmd/sandbox/state.go
index bae2341cac2..3760e6d0381 100644
--- a/cmd/sandbox/state.go
+++ b/cmd/sandbox/state.go
@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"slices"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/env"
 )
 
@@ -96,15 +97,7 @@ func saveState(ctx context.Context, state *stateFile) error {
 		return err
 	}
 
-	tmp := path + ".tmp"
-	if err := os.WriteFile(tmp, data, 0o600); err != nil {
-		return fmt.Errorf("writing %s: %w", tmp, err)
-	}
-	if err := os.Rename(tmp, path); err != nil {
-		_ = os.Remove(tmp)
-		return fmt.Errorf("renaming %s to %s: %w", tmp, path, err)
-	}
-	return nil
+	return atomicfile.Write(path, data, 0o600)
 }
 
 func getDefault(ctx context.Context, profile string) string {
diff --git a/experimental/ssh/internal/sshconfig/knownhosts.go b/experimental/ssh/internal/sshconfig/knownhosts.go
index 8cfecc2a060..530c2c26af0 100644
--- a/experimental/ssh/internal/sshconfig/knownhosts.go
+++ b/experimental/ssh/internal/sshconfig/knownhosts.go
@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/env"
 	"golang.org/x/crypto/ssh"
 )
@@ -65,25 +66,8 @@ func PinHostKey(path, hostName string, publicKey []byte) error {
 
 	// Write and rename so a connection racing this one (every ssh invocation runs the
 	// ProxyCommand, which refreshes the pin) never reads a half-written file.
-	tmp, err := os.CreateTemp(dir, ".known-hosts-*.tmp")
-	if err != nil {
-		return fmt.Errorf("failed to create known hosts file: %w", err)
-	}
-	defer os.Remove(tmp.Name())
-
-	_, err = tmp.WriteString(line)
-	if err == nil {
-		err = tmp.Chmod(0o600)
-	}
-	if closeErr := tmp.Close(); err == nil {
-		err = closeErr
-	}
-	if err != nil {
+	if err := atomicfile.Write(path, []byte(line), 0o600); err != nil {
 		return fmt.Errorf("failed to write known hosts file: %w", err)
 	}
-
-	if err := os.Rename(tmp.Name(), path); err != nil {
-		return fmt.Errorf("failed to replace known hosts file: %w", err)
-	}
 	return nil
 }
diff --git a/libs/aitools/installer/state.go b/libs/aitools/installer/state.go
index c8461b586bf..a2feee3b010 100644
--- a/libs/aitools/installer/state.go
+++ b/libs/aitools/installer/state.go
@@ -10,6 +10,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/env"
 )
 
@@ -128,28 +129,7 @@ func SaveState(dir string, state *InstallState) error {
 	}
 	data = append(data, '\n')
 
-	// Atomic write: write to temp file in the same directory, then rename.
-	tmp, err := os.CreateTemp(dir, ".state-*.tmp")
-	if err != nil {
-		return fmt.Errorf("failed to create temp file: %w", err)
-	}
-	tmpName := tmp.Name()
-
-	if _, err := tmp.Write(data); err != nil {
-		tmp.Close()
-		os.Remove(tmpName)
-		return fmt.Errorf("failed to write temp file: %w", err)
-	}
-	if err := tmp.Close(); err != nil {
-		os.Remove(tmpName)
-		return fmt.Errorf("failed to close temp file: %w", err)
-	}
-
-	if err := os.Rename(tmpName, filepath.Join(dir, stateFileName)); err != nil {
-		os.Remove(tmpName)
-		return fmt.Errorf("failed to rename state file: %w", err)
-	}
-	return nil
+	return atomicfile.Write(filepath.Join(dir, stateFileName), data, 0o600)
 }
 
 // GlobalSkillsDir returns the path to the global skills directory (~/.databricks/aitools/skills/).
diff --git a/libs/atomicfile/atomicfile.go b/libs/atomicfile/atomicfile.go
new file mode 100644
index 00000000000..e1a9567f374
--- /dev/null
+++ b/libs/atomicfile/atomicfile.go
@@ -0,0 +1,53 @@
+// Package atomicfile writes a file so that readers and crashes never observe a
+// partial result: the data goes to a temporary file in the same directory and
+// is then renamed over the destination. The rename is atomic on a single
+// filesystem, which is why the temp file must live next to the target rather
+// than in os.TempDir (a cross-filesystem rename fails with EXDEV).
+package atomicfile
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+)
+
+// Write atomically writes data to path with mode perm. It creates a temp file
+// in path's directory, writes and chmods it, then renames it over path.
+//
+// The resulting file always has mode perm; it does not inherit the permissions
+// of a file it replaces (os.CreateTemp starts at 0600, and the rename replaces
+// the inode, so callers state the mode they want explicitly). The parent
+// directory must already exist.
+func Write(path string, data []byte, perm os.FileMode) error {
+	dir := filepath.Dir(path)
+
+	// Temp file in the same directory so the rename stays on one filesystem.
+	// The "." prefix hides a leftover temp file and the ".tmp" suffix lets
+	// callers sweep leftovers with a "*.tmp" glob.
+	tmp, err := os.CreateTemp(dir, "."+filepath.Base(path)+".*.tmp")
+	if err != nil {
+		return fmt.Errorf("create temp file for %s: %w", path, err)
+	}
+	tmpPath := tmp.Name()
+
+	if _, err := tmp.Write(data); err != nil {
+		tmp.Close()
+		os.Remove(tmpPath)
+		return fmt.Errorf("write %s: %w", tmpPath, err)
+	}
+	if err := tmp.Chmod(perm); err != nil {
+		tmp.Close()
+		os.Remove(tmpPath)
+		return fmt.Errorf("chmod %s: %w", tmpPath, err)
+	}
+	// Close before the rename: on Windows the file must not be open.
+	if err := tmp.Close(); err != nil {
+		os.Remove(tmpPath)
+		return fmt.Errorf("close %s: %w", tmpPath, err)
+	}
+	if err := os.Rename(tmpPath, path); err != nil {
+		os.Remove(tmpPath)
+		return fmt.Errorf("rename %s to %s: %w", tmpPath, path, err)
+	}
+	return nil
+}
diff --git a/libs/atomicfile/atomicfile_test.go b/libs/atomicfile/atomicfile_test.go
new file mode 100644
index 00000000000..3cef98e0264
--- /dev/null
+++ b/libs/atomicfile/atomicfile_test.go
@@ -0,0 +1,83 @@
+package atomicfile
+
+import (
+	"os"
+	"path/filepath"
+	"runtime"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWriteCreatesFile(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "out.json")
+
+	require.NoError(t, Write(path, []byte("hello"), 0o600))
+
+	got, err := os.ReadFile(path)
+	require.NoError(t, err)
+	assert.Equal(t, "hello", string(got))
+}
+
+func TestWriteUsesGivenMode(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("Windows does not honor unix file modes")
+	}
+	path := filepath.Join(t.TempDir(), "out")
+
+	require.NoError(t, Write(path, []byte("x"), 0o644))
+
+	info, err := os.Stat(path)
+	require.NoError(t, err)
+	assert.Equal(t, os.FileMode(0o644), info.Mode().Perm())
+}
+
+func TestWriteDoesNotPreserveReplacedMode(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("Windows does not honor unix file modes")
+	}
+	path := filepath.Join(t.TempDir(), "out")
+	require.NoError(t, os.WriteFile(path, []byte("old"), 0o600))
+
+	// Replacing a 0600 file with perm 0644 yields 0644, not the old mode.
+	require.NoError(t, Write(path, []byte("new"), 0o644))
+
+	info, err := os.Stat(path)
+	require.NoError(t, err)
+	assert.Equal(t, os.FileMode(0o644), info.Mode().Perm())
+}
+
+func TestWriteOverwrites(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "out")
+	require.NoError(t, os.WriteFile(path, []byte("old contents"), 0o600))
+
+	require.NoError(t, Write(path, []byte("new"), 0o600))
+
+	got, err := os.ReadFile(path)
+	require.NoError(t, err)
+	assert.Equal(t, "new", string(got))
+}
+
+func TestWriteLeavesNoTempFile(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "out")
+
+	require.NoError(t, Write(path, []byte("x"), 0o600))
+
+	entries, err := os.ReadDir(dir)
+	require.NoError(t, err)
+	require.Len(t, entries, 1)
+	assert.Equal(t, "out", entries[0].Name())
+}
+
+func TestWriteMissingDir(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "no-such-dir", "out")
+
+	err := Write(path, []byte("x"), 0o600)
+	require.Error(t, err)
+
+	// The failure must not leave the target behind.
+	_, statErr := os.Stat(path)
+	assert.ErrorIs(t, statErr, os.ErrNotExist)
+}
diff --git a/libs/auth/storage/filestore.go b/libs/auth/storage/filestore.go
index 328d7d466c6..7696c8c15db 100644
--- a/libs/auth/storage/filestore.go
+++ b/libs/auth/storage/filestore.go
@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"sync"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/env"
 	"golang.org/x/oauth2"
 )
@@ -206,33 +207,8 @@ func (c *fileStore) load() (*tokenStoreFile, error) {
 	return f, nil
 }
 
-// atomicWriteFile writes data to the file atomically by first writing to a
-// temporary file in the same directory and then renaming it to the target.
-// This prevents corruption from interrupted writes.
+// atomicWriteFile writes data to the token store file atomically, so an
+// interrupted write cannot corrupt existing tokens.
 func (c *fileStore) atomicWriteFile(data []byte) error {
-	tmp, err := c.writeTmpFile(data)
-	if err != nil {
-		return err
-	}
-	defer os.Remove(tmp)
-	return os.Rename(tmp, c.fileLocation)
-}
-
-func (c *fileStore) writeTmpFile(data []byte) (string, error) {
-	tmp, err := os.CreateTemp(filepath.Dir(c.fileLocation), ".token-cache-*.tmp")
-	if err != nil {
-		return "", fmt.Errorf("create temp file: %w", err)
-	}
-	defer tmp.Close()
-
-	if _, err := tmp.Write(data); err != nil {
-		return "", err
-	}
-	if err := tmp.Chmod(ownerReadWrite); err != nil {
-		return "", err
-	}
-	if err := tmp.Close(); err != nil {
-		return "", err
-	}
-	return tmp.Name(), nil
+	return atomicfile.Write(c.fileLocation, data, ownerReadWrite)
 }
diff --git a/libs/cache/file_cache.go b/libs/cache/file_cache.go
index 163cbf3bacc..cb9c5e433d7 100644
--- a/libs/cache/file_cache.go
+++ b/libs/cache/file_cache.go
@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/databricks/cli/internal/build"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
 )
@@ -278,38 +279,9 @@ func (fc *fileCache) readFromCacheJSON(ctx context.Context, cachePath string) ([
 }
 
 // writeToCacheJSON writes data to the cache file atomically.
-// Uses atomic write: writes to temp file first, then renames to actual cache file.
 func (fc *fileCache) writeToCacheJSON(ctx context.Context, cachePath string, data []byte) {
-	// Create temporary file in the same directory for atomic operation
-	tempFile, err := os.CreateTemp(fc.baseDir, ".cache-*.tmp")
-	if err != nil {
-		log.Debugf(ctx, "[Local Cache] failed to create temp cache file: %v", err)
-		return
-	}
-	tempPath := tempFile.Name()
-	defer func() {
-		_ = tempFile.Close()
-		_ = os.Remove(tempPath) // Clean up temp file if still exists
-	}()
-
-	// Write data to temp file
-	if _, err := tempFile.Write(data); err != nil {
-		log.Debugf(ctx, "[Local Cache] failed to write to temp cache file: %v", err)
-		return
-	}
-
-	if err := tempFile.Close(); err != nil {
-		log.Debugf(ctx, "[Local Cache] failed to close temp cache file: %v", err)
-		return
-	}
-
-	// On Windows, os.Rename fails if target exists, so remove it first
-	// This is a best-effort operation - if it fails because file doesn't exist, that's fine
-	_ = os.Remove(cachePath)
-
-	// Atomically rename temp file to actual cache file
-	if err := os.Rename(tempPath, cachePath); err != nil {
-		log.Debugf(ctx, "[Local Cache] failed to rename temp cache file: %v", err)
+	if err := atomicfile.Write(cachePath, data, 0o600); err != nil {
+		log.Debugf(ctx, "[Local Cache] failed to write cache file: %v", err)
 	}
 }
 
diff --git a/libs/clicompat/clicompat.go b/libs/clicompat/clicompat.go
index 49034a5ed0e..797c39b51ed 100644
--- a/libs/clicompat/clicompat.go
+++ b/libs/clicompat/clicompat.go
@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/databricks/cli/internal/build"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
 	"golang.org/x/mod/semver"
@@ -318,8 +319,7 @@ func readLocalManifest(path string) (cachedManifest, error) {
 	return cachedManifest{manifest: m, modTime: info.ModTime()}, nil
 }
 
-// writeLocalManifest writes the manifest to the local cache path using a
-// temp-file-then-rename pattern for atomicity.
+// writeLocalManifest writes the manifest to the local cache path atomically.
 func writeLocalManifest(ctx context.Context, path string, m Manifest) {
 	if path == "" {
 		return
@@ -334,26 +334,8 @@ func writeLocalManifest(ctx context.Context, path string, m Manifest) {
 		log.Warnf(ctx, "Failed to create cache directory %s: %v", dir, err)
 		return
 	}
-	tmp, err := os.CreateTemp(dir, ".compat-manifest-*.tmp")
-	if err != nil {
-		log.Warnf(ctx, "Failed to create temp cache file: %v", err)
-		return
-	}
-	tmpPath := tmp.Name()
-	defer func() {
-		_ = tmp.Close()
-		_ = os.Remove(tmpPath)
-	}()
-	if _, err := tmp.Write(data); err != nil {
-		log.Debugf(ctx, "Failed to write temp cache file: %v", err)
-		return
-	}
-	if err := tmp.Close(); err != nil {
-		log.Debugf(ctx, "Failed to close temp cache file: %v", err)
-		return
-	}
-	if err := os.Rename(tmpPath, path); err != nil {
-		log.Warnf(ctx, "Failed to rename temp cache file: %v", err)
+	if err := atomicfile.Write(path, data, 0o600); err != nil {
+		log.Warnf(ctx, "Failed to write cache file: %v", err)
 	}
 }
 
diff --git a/libs/localenv/constraints.go b/libs/localenv/constraints.go
index 7c26798412d..fe54ebfc0ec 100644
--- a/libs/localenv/constraints.go
+++ b/libs/localenv/constraints.go
@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/BurntSushi/toml"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
 )
@@ -94,38 +95,13 @@ func cacheFileName(envKey string) string {
 	return fmt.Sprintf("%s-%s.toml", slug, hex.EncodeToString(sum[:8]))
 }
 
-// writeCacheAtomic writes data to path via a temp file and rename, creating the
-// parent directory first. The rename is atomic on the same filesystem, so a
-// concurrent reader never observes a truncated or partial cache file (os.WriteFile
-// truncates in place, which a fallback reader could catch mid-write).
+// writeCacheAtomic writes data to path, creating the parent directory first.
 func writeCacheAtomic(path string, data []byte) error {
 	dir := filepath.Dir(path)
 	if err := os.MkdirAll(dir, 0o755); err != nil {
 		return err
 	}
-	tmp, err := os.CreateTemp(dir, ".constraints-*.tmp")
-	if err != nil {
-		return err
-	}
-	tmpName := tmp.Name()
-	if _, err := tmp.Write(data); err != nil {
-		tmp.Close()
-		os.Remove(tmpName)
-		return err
-	}
-	if err := tmp.Close(); err != nil {
-		os.Remove(tmpName)
-		return err
-	}
-	if err := os.Chmod(tmpName, 0o600); err != nil {
-		os.Remove(tmpName)
-		return err
-	}
-	if err := os.Rename(tmpName, path); err != nil {
-		os.Remove(tmpName)
-		return err
-	}
-	return nil
+	return atomicfile.Write(path, data, 0o600)
 }
 
 // FetchConstraints fetches the pyproject.toml for envKey from baseURL and caches it in

From a93cb7987fc2cc510aec662c10eecf994b95479f Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Wed, 16 Sep 2026 14:23:10 +0200
Subject: [PATCH 2/8] Add changelog fragment

Co-authored-by: Isaac <no-reply@databricks.com>
---
 .nextchanges/cli/atomic-file-writes.md | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 .nextchanges/cli/atomic-file-writes.md

diff --git a/.nextchanges/cli/atomic-file-writes.md b/.nextchanges/cli/atomic-file-writes.md
new file mode 100644
index 00000000000..4f7da6a9a1c
--- /dev/null
+++ b/.nextchanges/cli/atomic-file-writes.md
@@ -0,0 +1 @@
+* Write local state and cache files atomically so an interrupted or concurrent write cannot corrupt them. ([#6708](https://github.com/databricks/cli/pull/6708))

From 9671bc18aebd5eda5a95dbcc86dd2649d26edea6 Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Wed, 16 Sep 2026 14:51:27 +0200
Subject: [PATCH 3/8] atomicfile: add MkDir option, fold in redundant MkdirAll
 calls

Add a functional-options MkDir(perm) to atomicfile.Write so a caller can
create the parent directory (with its own deliberate mode: 0o700 for secrets,
0o755 otherwise) in the same call. Route the eight sites that did an explicit
MkdirAll right before the write through it, and inline localenv.writeCacheAtomic,
which became a single-caller pass-through.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 bundle/direct/dstate/state.go                 |  7 +---
 bundle/statemgmt/state_pull.go                | 10 +-----
 cmd/genie/conversations.go                    |  5 +--
 cmd/sandbox/state.go                          |  6 +---
 .../ssh/internal/sshconfig/knownhosts.go      |  7 +---
 libs/aitools/installer/state.go               |  6 +---
 libs/atomicfile/atomicfile.go                 | 34 +++++++++++++++++--
 libs/atomicfile/atomicfile_test.go            | 16 +++++++++
 libs/clicompat/clicompat.go                   |  7 +---
 libs/localenv/constraints.go                  | 11 +-----
 10 files changed, 56 insertions(+), 53 deletions(-)

diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go
index 3bb30d25861..046335e09e4 100644
--- a/bundle/direct/dstate/state.go
+++ b/bundle/direct/dstate/state.go
@@ -999,12 +999,7 @@ func (db *DeploymentState) unlockedSave() error {
 		return err
 	}
 
-	dir := filepath.Dir(db.Path)
-	if err := os.MkdirAll(dir, 0o755); err != nil {
-		return fmt.Errorf("failed to create directory %#v: %w", dir, err)
-	}
-
-	if err := atomicfile.Write(db.Path, data, 0o600); err != nil {
+	if err := atomicfile.Write(db.Path, data, 0o600, atomicfile.MkDir(0o755)); err != nil {
 		return fmt.Errorf("failed to save resources state to %#v: %w", db.Path, err)
 	}
 
diff --git a/bundle/statemgmt/state_pull.go b/bundle/statemgmt/state_pull.go
index 52f0d858435..cf5e15bb9fb 100644
--- a/bundle/statemgmt/state_pull.go
+++ b/bundle/statemgmt/state_pull.go
@@ -195,15 +195,7 @@ func PullResourcesState(ctx context.Context, b *bundle.Bundle, alwaysPull Always
 			localStatePath = localPathDirect
 		}
 
-		localStateDir := filepath.Dir(localStatePath)
-
-		err := os.MkdirAll(localStateDir, 0o700)
-		if err != nil {
-			logdiag.LogError(ctx, err)
-			return ctx, winner
-		}
-
-		err = atomicfile.Write(localStatePath, winner.Content, 0o600)
+		err := atomicfile.Write(localStatePath, winner.Content, 0o600, atomicfile.MkDir(0o700))
 		if err != nil {
 			logdiag.LogError(ctx, err)
 			return ctx, winner
diff --git a/cmd/genie/conversations.go b/cmd/genie/conversations.go
index b60d9e0e4a5..bfa86348ff8 100644
--- a/cmd/genie/conversations.go
+++ b/cmd/genie/conversations.go
@@ -83,10 +83,7 @@ func saveStore(path string, store conversationStore) {
 	if err != nil {
 		return
 	}
-	if err := os.MkdirAll(filepath.Dir(path), conversationDirPerm); err != nil {
-		return
-	}
-	_ = atomicfile.Write(path, raw, conversationFilePerm)
+	_ = atomicfile.Write(path, raw, conversationFilePerm, atomicfile.MkDir(conversationDirPerm))
 }
 
 // lookupConversationID returns the server conversation id mapped to sessionID on
diff --git a/cmd/sandbox/state.go b/cmd/sandbox/state.go
index 3760e6d0381..64294c17824 100644
--- a/cmd/sandbox/state.go
+++ b/cmd/sandbox/state.go
@@ -88,16 +88,12 @@ func saveState(ctx context.Context, state *stateFile) error {
 		return err
 	}
 
-	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
-		return err
-	}
-
 	data, err := json.MarshalIndent(state, "", "  ")
 	if err != nil {
 		return err
 	}
 
-	return atomicfile.Write(path, data, 0o600)
+	return atomicfile.Write(path, data, 0o600, atomicfile.MkDir(0o700))
 }
 
 func getDefault(ctx context.Context, profile string) string {
diff --git a/experimental/ssh/internal/sshconfig/knownhosts.go b/experimental/ssh/internal/sshconfig/knownhosts.go
index 530c2c26af0..f12749feb9e 100644
--- a/experimental/ssh/internal/sshconfig/knownhosts.go
+++ b/experimental/ssh/internal/sshconfig/knownhosts.go
@@ -59,14 +59,9 @@ func PinHostKey(path, hostName string, publicKey []byte) error {
 		return nil
 	}
 
-	dir := filepath.Dir(path)
-	if err := os.MkdirAll(dir, 0o700); err != nil {
-		return fmt.Errorf("failed to create known hosts directory: %w", err)
-	}
-
 	// Write and rename so a connection racing this one (every ssh invocation runs the
 	// ProxyCommand, which refreshes the pin) never reads a half-written file.
-	if err := atomicfile.Write(path, []byte(line), 0o600); err != nil {
+	if err := atomicfile.Write(path, []byte(line), 0o600, atomicfile.MkDir(0o700)); err != nil {
 		return fmt.Errorf("failed to write known hosts file: %w", err)
 	}
 	return nil
diff --git a/libs/aitools/installer/state.go b/libs/aitools/installer/state.go
index a2feee3b010..ddd51568bde 100644
--- a/libs/aitools/installer/state.go
+++ b/libs/aitools/installer/state.go
@@ -119,17 +119,13 @@ func migrateState(state *InstallState) {
 // SaveState writes install state to the given directory atomically.
 // Creates the directory if it does not exist.
 func SaveState(dir string, state *InstallState) error {
-	if err := os.MkdirAll(dir, 0o755); err != nil {
-		return fmt.Errorf("failed to create state directory: %w", err)
-	}
-
 	data, err := json.MarshalIndent(state, "", "  ")
 	if err != nil {
 		return fmt.Errorf("failed to marshal state: %w", err)
 	}
 	data = append(data, '\n')
 
-	return atomicfile.Write(filepath.Join(dir, stateFileName), data, 0o600)
+	return atomicfile.Write(filepath.Join(dir, stateFileName), data, 0o600, atomicfile.MkDir(0o755))
 }
 
 // GlobalSkillsDir returns the path to the global skills directory (~/.databricks/aitools/skills/).
diff --git a/libs/atomicfile/atomicfile.go b/libs/atomicfile/atomicfile.go
index e1a9567f374..a8a26cd7d57 100644
--- a/libs/atomicfile/atomicfile.go
+++ b/libs/atomicfile/atomicfile.go
@@ -11,16 +11,46 @@ import (
 	"path/filepath"
 )
 
+// Option configures a Write call.
+type Option func(*config)
+
+type config struct {
+	mkdirPerm os.FileMode
+	mkdir     bool
+}
+
+// MkDir makes Write create path's parent directory (and any missing parents)
+// with mode perm before writing, like os.MkdirAll. The directory mode is a
+// separate, deliberate choice from the file mode (e.g. 0o700 for a directory
+// holding secrets), so callers pass it explicitly rather than it defaulting.
+func MkDir(perm os.FileMode) Option {
+	return func(c *config) {
+		c.mkdir = true
+		c.mkdirPerm = perm
+	}
+}
+
 // Write atomically writes data to path with mode perm. It creates a temp file
 // in path's directory, writes and chmods it, then renames it over path.
 //
 // The resulting file always has mode perm; it does not inherit the permissions
 // of a file it replaces (os.CreateTemp starts at 0600, and the rename replaces
 // the inode, so callers state the mode they want explicitly). The parent
-// directory must already exist.
-func Write(path string, data []byte, perm os.FileMode) error {
+// directory must already exist unless the MkDir option is passed.
+func Write(path string, data []byte, perm os.FileMode, opts ...Option) error {
+	var cfg config
+	for _, opt := range opts {
+		opt(&cfg)
+	}
+
 	dir := filepath.Dir(path)
 
+	if cfg.mkdir {
+		if err := os.MkdirAll(dir, cfg.mkdirPerm); err != nil {
+			return fmt.Errorf("create directory %s: %w", dir, err)
+		}
+	}
+
 	// Temp file in the same directory so the rename stays on one filesystem.
 	// The "." prefix hides a leftover temp file and the ".tmp" suffix lets
 	// callers sweep leftovers with a "*.tmp" glob.
diff --git a/libs/atomicfile/atomicfile_test.go b/libs/atomicfile/atomicfile_test.go
index 3cef98e0264..9ad3ce9bae4 100644
--- a/libs/atomicfile/atomicfile_test.go
+++ b/libs/atomicfile/atomicfile_test.go
@@ -81,3 +81,19 @@ func TestWriteMissingDir(t *testing.T) {
 	_, statErr := os.Stat(path)
 	assert.ErrorIs(t, statErr, os.ErrNotExist)
 }
+
+func TestWriteMkDirCreatesParents(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "a", "b", "out")
+
+	require.NoError(t, Write(path, []byte("x"), 0o600, MkDir(0o700)))
+
+	got, err := os.ReadFile(path)
+	require.NoError(t, err)
+	assert.Equal(t, "x", string(got))
+
+	if runtime.GOOS != "windows" {
+		info, err := os.Stat(filepath.Dir(path))
+		require.NoError(t, err)
+		assert.Equal(t, os.FileMode(0o700), info.Mode().Perm())
+	}
+}
diff --git a/libs/clicompat/clicompat.go b/libs/clicompat/clicompat.go
index 797c39b51ed..71ffe75d36c 100644
--- a/libs/clicompat/clicompat.go
+++ b/libs/clicompat/clicompat.go
@@ -329,12 +329,7 @@ func writeLocalManifest(ctx context.Context, path string, m Manifest) {
 		log.Debugf(ctx, "Failed to marshal manifest for cache: %v", err)
 		return
 	}
-	dir := filepath.Dir(path)
-	if err := os.MkdirAll(dir, 0o700); err != nil {
-		log.Warnf(ctx, "Failed to create cache directory %s: %v", dir, err)
-		return
-	}
-	if err := atomicfile.Write(path, data, 0o600); err != nil {
+	if err := atomicfile.Write(path, data, 0o600, atomicfile.MkDir(0o700)); err != nil {
 		log.Warnf(ctx, "Failed to write cache file: %v", err)
 	}
 }
diff --git a/libs/localenv/constraints.go b/libs/localenv/constraints.go
index fe54ebfc0ec..29c66134ad6 100644
--- a/libs/localenv/constraints.go
+++ b/libs/localenv/constraints.go
@@ -95,15 +95,6 @@ func cacheFileName(envKey string) string {
 	return fmt.Sprintf("%s-%s.toml", slug, hex.EncodeToString(sum[:8]))
 }
 
-// writeCacheAtomic writes data to path, creating the parent directory first.
-func writeCacheAtomic(path string, data []byte) error {
-	dir := filepath.Dir(path)
-	if err := os.MkdirAll(dir, 0o755); err != nil {
-		return err
-	}
-	return atomicfile.Write(path, data, 0o600)
-}
-
 // FetchConstraints fetches the pyproject.toml for envKey from baseURL and caches it in
 // cacheDir. On a transport or non-404 HTTP failure it falls back to the cached copy if one
 // exists (E_FETCH otherwise). A 404 means the env key is not published (E_ENV_UNSUPPORTED)
@@ -139,7 +130,7 @@ func FetchConstraints(ctx context.Context, baseURL, envKey, cacheDir string, wri
 		// so a read-only cacheDir doesn't break the command. Skipped under a dry
 		// run so --dry-run performs no disk writes at all.
 		if writeCache {
-			if err := writeCacheAtomic(cachePath, data); err != nil {
+			if err := atomicfile.Write(cachePath, data, 0o600, atomicfile.MkDir(0o755)); err != nil {
 				log.Debugf(ctx, "failed to write constraint cache %s: %v", filepath.ToSlash(cachePath), err)
 			}
 		}

From 782cf85db400c67e0848755f419d06f3012ac607 Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Wed, 16 Sep 2026 15:09:59 +0200
Subject: [PATCH 4/8] Make user file and CLI state saves atomic

Route the remaining in-place os.WriteFile saves that overwrite a persistent
file (torn write would lose the prior content) through atomicfile.Write:

New atomicity for files a user owns:
- bundle/configsync/output.go: rewriting local source config files on sync
- libs/completion/{install,uninstall}.go: editing the user's shell RC file
  (mode preserved via the existing info.Mode()); install now rewrites instead
  of appending, matching uninstall
- libs/localenv/pipeline.go: overwriting pyproject.toml on merge

New atomicity for CLI-owned state/cache read back later:
- cmd/labs/project/{login,project}.go: labs auth and version files
- cmd/labs/localcache/jsonfile.go: labs JSON cache (write+mkdir+retry collapsed
  into one MkDir call)
- bundle/configsync/diff.go: local config snapshot used for the next diff

Left as-is: the aitools installer writes to a temp staging dir already promoted
by rename, .databrickscfg (go-ini SaveTo, needs its own approach), and the
experimental/ssh in-place config editors.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 .nextchanges/cli/atomic-file-writes.md |  2 +-
 bundle/configsync/diff.go              | 10 ++------
 bundle/configsync/output.go            | 10 ++------
 cmd/labs/localcache/jsonfile.go        | 15 ++----------
 cmd/labs/project/login.go              |  4 ++--
 cmd/labs/project/project.go            |  3 ++-
 libs/completion/install.go             | 32 +++++++-------------------
 libs/completion/uninstall.go           |  4 +++-
 libs/localenv/pipeline.go              |  3 ++-
 9 files changed, 24 insertions(+), 59 deletions(-)

diff --git a/.nextchanges/cli/atomic-file-writes.md b/.nextchanges/cli/atomic-file-writes.md
index 4f7da6a9a1c..dac5d484cfa 100644
--- a/.nextchanges/cli/atomic-file-writes.md
+++ b/.nextchanges/cli/atomic-file-writes.md
@@ -1 +1 @@
-* Write local state and cache files atomically so an interrupted or concurrent write cannot corrupt them. ([#6708](https://github.com/databricks/cli/pull/6708))
+* Write local state, cache, and config files atomically so an interrupted or concurrent write cannot corrupt them. ([#6708](https://github.com/databricks/cli/pull/6708))
diff --git a/bundle/configsync/diff.go b/bundle/configsync/diff.go
index 7c79456abae..8192f5e9b4c 100644
--- a/bundle/configsync/diff.go
+++ b/bundle/configsync/diff.go
@@ -7,7 +7,6 @@ import (
 	"io"
 	"io/fs"
 	"os"
-	"path/filepath"
 	"strings"
 
 	"github.com/databricks/cli/bundle"
@@ -17,6 +16,7 @@ import (
 	"github.com/databricks/cli/bundle/deployplan"
 	"github.com/databricks/cli/bundle/direct"
 	"github.com/databricks/cli/bundle/direct/dstate"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/cli/libs/dyn/convert"
 	"github.com/databricks/cli/libs/log"
@@ -260,13 +260,7 @@ func ensureSnapshotAvailable(ctx context.Context, b *bundle.Bundle, engine engin
 		return fmt.Errorf("reading snapshot content: %w", err)
 	}
 
-	localStateDir := filepath.Dir(localPathSnapshot)
-	err = os.MkdirAll(localStateDir, 0o700)
-	if err != nil {
-		return fmt.Errorf("creating snapshot directory: %w", err)
-	}
-
-	err = os.WriteFile(localPathSnapshot, content, 0o600)
+	err = atomicfile.Write(localPathSnapshot, content, 0o600, atomicfile.MkDir(0o700))
 	if err != nil {
 		return fmt.Errorf("writing snapshot file: %w", err)
 	}
diff --git a/bundle/configsync/output.go b/bundle/configsync/output.go
index 5f77220fec3..ac6e81bd66c 100644
--- a/bundle/configsync/output.go
+++ b/bundle/configsync/output.go
@@ -6,10 +6,9 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"os"
-	"path/filepath"
 
 	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/telemetry"
 	"github.com/databricks/cli/libs/telemetry/protos"
 )
@@ -91,12 +90,7 @@ func WriteResult(out io.Writer, jsonOutput bool, stats *Stats, files []FileChang
 // SaveFiles writes all file changes to disk.
 func SaveFiles(ctx context.Context, b *bundle.Bundle, files []FileChange) error {
 	for _, file := range files {
-		err := os.MkdirAll(filepath.Dir(file.Path), 0o755)
-		if err != nil {
-			return err
-		}
-
-		err = os.WriteFile(file.Path, []byte(file.ModifiedContent), 0o644)
+		err := atomicfile.Write(file.Path, []byte(file.ModifiedContent), 0o644, atomicfile.MkDir(0o755))
 		if err != nil {
 			return err
 		}
diff --git a/cmd/labs/localcache/jsonfile.go b/cmd/labs/localcache/jsonfile.go
index d86ff2dd9e1..a4fd8c0db39 100644
--- a/cmd/labs/localcache/jsonfile.go
+++ b/cmd/labs/localcache/jsonfile.go
@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"time"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/log"
 )
 
@@ -73,19 +74,7 @@ func (r *LocalCache[T]) writeCache(ctx context.Context, data T) (T, error) {
 		return r.zero, fmt.Errorf("json marshal: %w", err)
 	}
 	cacheFile := r.FileName()
-	err = os.WriteFile(cacheFile, raw, userRW)
-	if errors.Is(err, fs.ErrNotExist) {
-		cacheDir := filepath.Dir(cacheFile)
-		err := os.MkdirAll(cacheDir, ownerRWXworldRX)
-		if err != nil {
-			return r.zero, fmt.Errorf("create %s: %w", cacheDir, err)
-		}
-		err = os.WriteFile(cacheFile, raw, userRW)
-		if err != nil {
-			return r.zero, fmt.Errorf("retry save cache: %w", err)
-		}
-		return data, nil
-	} else if err != nil {
+	if err := atomicfile.Write(cacheFile, raw, userRW, atomicfile.MkDir(ownerRWXworldRX)); err != nil {
 		return r.zero, fmt.Errorf("save cache: %w", err)
 	}
 	return data, nil
diff --git a/cmd/labs/project/login.go b/cmd/labs/project/login.go
index 4631fc1de3e..39ccbf24540 100644
--- a/cmd/labs/project/login.go
+++ b/cmd/labs/project/login.go
@@ -4,9 +4,9 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"os"
 
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg/cfgpickers"
 	"github.com/databricks/cli/libs/log"
@@ -120,5 +120,5 @@ func (lc *loginConfig) save(ctx context.Context) error {
 		return err
 	}
 	log.Debugf(ctx, "Writing auth configuration to: %s", authFile)
-	return os.WriteFile(authFile, raw, ownerRW)
+	return atomicfile.Write(authFile, raw, ownerRW)
 }
diff --git a/cmd/labs/project/project.go b/cmd/labs/project/project.go
index 42eba593293..e40ef2a5501 100644
--- a/cmd/labs/project/project.go
+++ b/cmd/labs/project/project.go
@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/databricks/cli/cmd/labs/github"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
@@ -291,7 +292,7 @@ func (p *Project) writeVersionFile(ctx context.Context, ver string) error {
 		return err
 	}
 	log.Debugf(ctx, "Writing installed version info to: %s", versionFile)
-	return os.WriteFile(versionFile, raw, ownerRW)
+	return atomicfile.Write(versionFile, raw, ownerRW)
 }
 
 // checkUpdates is called before every command of an installed project,
diff --git a/libs/completion/install.go b/libs/completion/install.go
index 9ee6b49ae42..7514c99d18e 100644
--- a/libs/completion/install.go
+++ b/libs/completion/install.go
@@ -3,7 +3,8 @@ package completion
 import (
 	"context"
 	"os"
-	"path/filepath"
+
+	"github.com/databricks/cli/libs/atomicfile"
 )
 
 // Install configures shell completion for the given shell. homeDir is used
@@ -36,12 +37,7 @@ func Install(ctx context.Context, shell Shell, homeDir string) (filePath string,
 // The caller must check Status before calling this — existence checks are not
 // repeated here.
 func installFish(filePath string, shell Shell) (string, bool, error) {
-	dir := filepath.Dir(filePath)
-	if err := os.MkdirAll(dir, 0o755); err != nil {
-		return filePath, false, err
-	}
-
-	return filePath, false, os.WriteFile(filePath, []byte(ShimContent(shell)), 0o644)
+	return filePath, false, atomicfile.Write(filePath, []byte(ShimContent(shell)), 0o644, atomicfile.MkDir(0o755))
 }
 
 // installRC handles the RC file model for bash, zsh, and powershell.
@@ -59,27 +55,15 @@ func installRC(filePath string, shell Shell) (string, bool, error) {
 		}
 	}
 
-	// Create parent directory if needed (e.g. for PowerShell profiles).
-	dir := filepath.Dir(filePath)
-	if err := os.MkdirAll(dir, 0o755); err != nil {
-		return filePath, false, err
-	}
-
 	// Ensure a leading newline before the block if the file doesn't end with one.
 	shim := ShimContent(shell)
 	if len(content) > 0 && content[len(content)-1] != '\n' {
 		shim = "\n" + shim
 	}
 
-	f, err := os.OpenFile(filePath, os.O_WRONLY|os.O_CREATE|os.O_APPEND, perm)
-	if err != nil {
-		return filePath, false, err
-	}
-	defer f.Close()
-
-	if _, err := f.WriteString(shim); err != nil {
-		return filePath, false, err
-	}
-
-	return filePath, false, nil
+	// Append our block by rewriting the file atomically, so an interrupted
+	// install cannot leave a torn block in the user's RC file. The parent dir
+	// may be missing (e.g. for PowerShell profiles), so create it too.
+	newContent := append(content, []byte(shim)...)
+	return filePath, false, atomicfile.Write(filePath, newContent, perm, atomicfile.MkDir(0o755))
 }
diff --git a/libs/completion/uninstall.go b/libs/completion/uninstall.go
index 0e70be344da..00654fea06f 100644
--- a/libs/completion/uninstall.go
+++ b/libs/completion/uninstall.go
@@ -7,6 +7,8 @@ import (
 	"os"
 	"regexp"
 	"strings"
+
+	"github.com/databricks/cli/libs/atomicfile"
 )
 
 var multiBlankLine = regexp.MustCompile(`\n{3,}`)
@@ -89,5 +91,5 @@ func uninstallRC(filePath string) (string, bool, error) {
 	// Collapse double blank lines left by removal.
 	result = multiBlankLine.ReplaceAllString(result, "\n\n")
 
-	return filePath, true, os.WriteFile(filePath, []byte(result), info.Mode())
+	return filePath, true, atomicfile.Write(filePath, []byte(result), info.Mode())
 }
diff --git a/libs/localenv/pipeline.go b/libs/localenv/pipeline.go
index 1c855194679..bbde0ad12cf 100644
--- a/libs/localenv/pipeline.go
+++ b/libs/localenv/pipeline.go
@@ -10,6 +10,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/log"
 	"github.com/hexops/gotextdiff"
 	"github.com/hexops/gotextdiff/myers"
@@ -509,7 +510,7 @@ func (p *Pipeline) applyMerge(_ context.Context, mergedBytes []byte, greenfield
 		p.res.BackupPath = filepath.ToSlash(backup)
 	}
 
-	if err := os.WriteFile(pyproject, mergedBytes, 0o644); err != nil {
+	if err := atomicfile.Write(pyproject, mergedBytes, 0o644); err != nil {
 		code := ErrMerge
 		if greenfield {
 			code = ErrWrite

From 0adf30c7cc28e8d6a54143d490ca56b3c61fd378 Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Wed, 16 Sep 2026 21:54:34 +0200
Subject: [PATCH 5/8] Inline one-line atomicfile wrappers

Remove atomicWriteFile and installFish, which became single-expression
pass-throughs to atomicfile.Write, inlining them into their callers.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 libs/auth/storage/filestore.go | 10 ++--------
 libs/completion/install.go     | 10 ++--------
 2 files changed, 4 insertions(+), 16 deletions(-)

diff --git a/libs/auth/storage/filestore.go b/libs/auth/storage/filestore.go
index 7696c8c15db..d7f13d397d3 100644
--- a/libs/auth/storage/filestore.go
+++ b/libs/auth/storage/filestore.go
@@ -144,7 +144,7 @@ func (c *fileStore) write(f *tokenStoreFile) error {
 	if err != nil {
 		return fmt.Errorf("marshal: %w", err)
 	}
-	if err := c.atomicWriteFile(raw); err != nil {
+	if err := atomicfile.Write(c.fileLocation, raw, ownerReadWrite); err != nil {
 		return fmt.Errorf("error storing token in local cache: %w", err)
 	}
 	return nil
@@ -180,7 +180,7 @@ func (c *fileStore) init(ctx context.Context) error {
 		if err != nil {
 			return fmt.Errorf("marshal: %w", err)
 		}
-		if err := c.atomicWriteFile(raw); err != nil {
+		if err := atomicfile.Write(c.fileLocation, raw, ownerReadWrite); err != nil {
 			return fmt.Errorf("error creating token store file: %w", err)
 		}
 	}
@@ -206,9 +206,3 @@ func (c *fileStore) load() (*tokenStoreFile, error) {
 	}
 	return f, nil
 }
-
-// atomicWriteFile writes data to the token store file atomically, so an
-// interrupted write cannot corrupt existing tokens.
-func (c *fileStore) atomicWriteFile(data []byte) error {
-	return atomicfile.Write(c.fileLocation, data, ownerReadWrite)
-}
diff --git a/libs/completion/install.go b/libs/completion/install.go
index 7514c99d18e..76c2ef60bc9 100644
--- a/libs/completion/install.go
+++ b/libs/completion/install.go
@@ -27,19 +27,13 @@ func Install(ctx context.Context, shell Shell, homeDir string) (filePath string,
 		return filePath, true, nil
 	}
 
+	// Fish uses a file-drop model; any existing file was ruled out above.
 	if shell == Fish {
-		return installFish(filePath, shell)
+		return filePath, false, atomicfile.Write(filePath, []byte(ShimContent(shell)), 0o644, atomicfile.MkDir(0o755))
 	}
 	return installRC(filePath, shell)
 }
 
-// installFish handles the file-drop model for fish completions.
-// The caller must check Status before calling this — existence checks are not
-// repeated here.
-func installFish(filePath string, shell Shell) (string, bool, error) {
-	return filePath, false, atomicfile.Write(filePath, []byte(ShimContent(shell)), 0o644, atomicfile.MkDir(0o755))
-}
-
 // installRC handles the RC file model for bash, zsh, and powershell.
 // The caller must check Status before calling this — marker checks are not
 // repeated here.

From ad6dfdf89e71532aabec7a5e2c66bf8735a21494 Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Thu, 17 Sep 2026 09:50:18 +0200
Subject: [PATCH 6/8] Make terraform-engine state and SSH config saves atomic

Route the remaining in-place state and user-config writes through
atomicfile.Write:

Terraform-engine state (the direct engine already writes these atomically):
- bundle/deploy/state_update.go: deployment state written after deploy
- bundle/deploy/state_pull.go: state pulled from the remote (read the local
  copy for the staleness check, then write atomically instead of truncating
  the open handle)
- libs/sync/snapshot.go: the sync snapshot, rewritten every sync and read back
  to compute the file diff

experimental/ssh in-place config editors:
- internal/sshconfig/sshconfig.go: the user's ~/.ssh/config (Include-directive
  add and migration) and the per-host tunnel config
- internal/vscode/settings.go: the user's VSCode settings.json

Co-authored-by: Isaac <no-reply@databricks.com>
---
 bundle/deploy/state_pull.go                   | 24 ++++---------------
 bundle/deploy/state_update.go                 | 16 ++++---------
 .../ssh/internal/sshconfig/sshconfig.go       | 13 ++++------
 experimental/ssh/internal/vscode/settings.go  |  3 ++-
 libs/sync/snapshot.go                         | 10 ++------
 5 files changed, 17 insertions(+), 49 deletions(-)

diff --git a/bundle/deploy/state_pull.go b/bundle/deploy/state_pull.go
index 832fac87fbb..462981cf4e2 100644
--- a/bundle/deploy/state_pull.go
+++ b/bundle/deploy/state_pull.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/deploy/files"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/log"
@@ -44,37 +45,22 @@ func (s *statePull) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostic
 		return diag.FromErr(err)
 	}
 
-	local, err := os.OpenFile(statePath, os.O_CREATE|os.O_RDWR, 0o600)
-	if err != nil {
-		return diag.FromErr(err)
-	}
-	defer local.Close()
-
 	data := remote.Bytes()
 	err = validateRemoteStateCompatibility(bytes.NewReader(data))
 	if err != nil {
 		return diag.FromErr(err)
 	}
 
-	if !isLocalStateStale(local, bytes.NewReader(data)) {
+	// A missing or unreadable local file counts as stale, so we write the remote copy.
+	localData, _ := os.ReadFile(statePath)
+	if !isLocalStateStale(bytes.NewReader(localData), bytes.NewReader(data)) {
 		log.Infof(ctx, "Local deployment state is the same or newer, ignoring remote state")
 		return nil
 	}
 
-	// Truncating the file before writing
-	err = local.Truncate(0)
-	if err != nil {
-		return diag.FromErr(err)
-	}
-	_, err = local.Seek(0, 0)
-	if err != nil {
-		return diag.FromErr(err)
-	}
-
 	// Write file to disk.
 	log.Infof(ctx, "Writing remote deployment state file to local cache directory")
-	_, err = io.Copy(local, bytes.NewReader(data))
-	if err != nil {
+	if err := atomicfile.Write(statePath, data, 0o600); err != nil {
 		return diag.FromErr(err)
 	}
 
diff --git a/bundle/deploy/state_update.go b/bundle/deploy/state_update.go
index 55cf2393bf1..fe68a6e519b 100644
--- a/bundle/deploy/state_update.go
+++ b/bundle/deploy/state_update.go
@@ -1,17 +1,16 @@
 package deploy
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
-	"io"
 	"io/fs"
 	"os"
 	"time"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/internal/build"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/log"
 	"github.com/google/uuid"
@@ -56,21 +55,14 @@ func (s *stateUpdate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnost
 	if err != nil {
 		return diag.FromErr(err)
 	}
-	// Write the state back to the file.
-	f, err := os.OpenFile(statePath, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0o600)
-	if err != nil {
-		log.Infof(ctx, "Unable to open deployment state file: %s", err)
-		return diag.FromErr(err)
-	}
-	defer f.Close()
-
 	data, err := json.Marshal(state)
 	if err != nil {
 		return diag.FromErr(err)
 	}
 
-	_, err = io.Copy(f, bytes.NewReader(data))
-	if err != nil {
+	// Write the state back to the file.
+	if err := atomicfile.Write(statePath, data, 0o600); err != nil {
+		log.Infof(ctx, "Unable to write deployment state file: %s", err)
 		return diag.FromErr(err)
 	}
 
diff --git a/experimental/ssh/internal/sshconfig/sshconfig.go b/experimental/ssh/internal/sshconfig/sshconfig.go
index 884ce9287c8..ff9ad9ad69e 100644
--- a/experimental/ssh/internal/sshconfig/sshconfig.go
+++ b/experimental/ssh/internal/sshconfig/sshconfig.go
@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/experimental/ssh/internal/fileutil"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
 )
@@ -107,7 +108,7 @@ func EnsureIncludeDirective(ctx context.Context, configPath string) error {
 		if err := fileutil.BackupFile(ctx, configPath, content); err != nil {
 			return fmt.Errorf("failed to backup SSH config before migration: %w", err)
 		}
-		return os.WriteFile(configPath, replaceLine(content, oldIncludeLine, includeLine), 0o600)
+		return atomicfile.Write(configPath, replaceLine(content, oldIncludeLine, includeLine), 0o600)
 	}
 
 	if err := fileutil.BackupFile(ctx, configPath, content); err != nil {
@@ -119,7 +120,7 @@ func EnsureIncludeDirective(ctx context.Context, configPath string) error {
 	}
 	newContent += string(content)
 
-	err = os.WriteFile(configPath, []byte(newContent), 0o600)
+	err = atomicfile.Write(configPath, []byte(newContent), 0o600)
 	if err != nil {
 		return fmt.Errorf("failed to update SSH config file with Include directive: %w", err)
 	}
@@ -191,13 +192,7 @@ func CreateOrUpdateHostConfig(ctx context.Context, hostName, hostConfig string,
 		return false, nil
 	}
 
-	configDir := filepath.Dir(configPath)
-	err = os.MkdirAll(configDir, 0o700)
-	if err != nil {
-		return false, fmt.Errorf("failed to create config directory: %w", err)
-	}
-
-	err = os.WriteFile(configPath, []byte(hostConfig), 0o600)
+	err = atomicfile.Write(configPath, []byte(hostConfig), 0o600, atomicfile.MkDir(0o700))
 	if err != nil {
 		return false, fmt.Errorf("failed to write host config file: %w", err)
 	}
diff --git a/experimental/ssh/internal/vscode/settings.go b/experimental/ssh/internal/vscode/settings.go
index 5b4a0bc34fc..ddaf1ace4e7 100644
--- a/experimental/ssh/internal/vscode/settings.go
+++ b/experimental/ssh/internal/vscode/settings.go
@@ -12,6 +12,7 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/experimental/ssh/internal/fileutil"
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
@@ -417,7 +418,7 @@ func updateSettings(v *hujson.Value, connectionName string, missing *missingSett
 }
 
 func saveSettings(path string, v *hujson.Value) error {
-	if err := os.WriteFile(path, v.Pack(), 0o600); err != nil {
+	if err := atomicfile.Write(path, v.Pack(), 0o600); err != nil {
 		return fmt.Errorf("failed to write settings file: %w", err)
 	}
 	return nil
diff --git a/libs/sync/snapshot.go b/libs/sync/snapshot.go
index 1b338aab887..f8454d3c1d4 100644
--- a/libs/sync/snapshot.go
+++ b/libs/sync/snapshot.go
@@ -12,6 +12,7 @@ import (
 	"path/filepath"
 	"time"
 
+	"github.com/databricks/cli/libs/atomicfile"
 	"github.com/databricks/cli/libs/fileset"
 	"github.com/databricks/cli/libs/log"
 )
@@ -121,19 +122,12 @@ func newSnapshot(ctx context.Context, opts *SyncOptions) (*Snapshot, error) {
 }
 
 func (s *Snapshot) Save(ctx context.Context) error {
-	f, err := os.OpenFile(s.snapshotPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
-	if err != nil {
-		return fmt.Errorf("failed to create/open persisted sync snapshot file: %s", err)
-	}
-	defer f.Close()
-
 	// persist snapshot to disk
 	bytes, err := json.MarshalIndent(s, "", "  ")
 	if err != nil {
 		return fmt.Errorf("failed to json marshal in-memory snapshot: %s", err)
 	}
-	_, err = f.Write(bytes)
-	if err != nil {
+	if err := atomicfile.Write(s.snapshotPath, bytes, 0o644); err != nil {
 		return fmt.Errorf("failed to write sync snapshot to disk: %s", err)
 	}
 	return nil

From 05a7825a71e65f36a3bbbd5cf8e7b21190fcca5d Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Thu, 17 Sep 2026 10:45:27 +0200
Subject: [PATCH 7/8] atomicfile: test MkDir against an existing directory

Cover MkDir when the parent directory already exists, with a matching and a
differing mode, pinning down that os.MkdirAll leaves an existing directory's
mode unchanged (it never widens 0o700 to the requested 0o755).

Co-authored-by: Isaac <no-reply@databricks.com>
---
 libs/atomicfile/atomicfile_test.go | 44 ++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/libs/atomicfile/atomicfile_test.go b/libs/atomicfile/atomicfile_test.go
index 9ad3ce9bae4..59bbfac5676 100644
--- a/libs/atomicfile/atomicfile_test.go
+++ b/libs/atomicfile/atomicfile_test.go
@@ -97,3 +97,47 @@ func TestWriteMkDirCreatesParents(t *testing.T) {
 		assert.Equal(t, os.FileMode(0o700), info.Mode().Perm())
 	}
 }
+
+// MkDir wraps os.MkdirAll, which is a no-op on a directory that already exists
+// and never changes its mode. These two cases pin that down: whether the
+// requested mode matches the existing directory or not, the directory keeps the
+// mode it already had and the write still succeeds.
+func TestWriteMkDirExistingDirSameMode(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("Windows does not honor unix directory modes")
+	}
+	dir := t.TempDir()
+	require.NoError(t, os.Chmod(dir, 0o700))
+	path := filepath.Join(dir, "out")
+
+	require.NoError(t, Write(path, []byte("x"), 0o600, MkDir(0o700)))
+
+	got, err := os.ReadFile(path)
+	require.NoError(t, err)
+	assert.Equal(t, "x", string(got))
+
+	info, err := os.Stat(dir)
+	require.NoError(t, err)
+	assert.Equal(t, os.FileMode(0o700), info.Mode().Perm())
+}
+
+func TestWriteMkDirExistingDirDifferentMode(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("Windows does not honor unix directory modes")
+	}
+	dir := t.TempDir()
+	require.NoError(t, os.Chmod(dir, 0o700))
+	path := filepath.Join(dir, "out")
+
+	// Ask for 0o755 even though the directory already exists at 0o700.
+	require.NoError(t, Write(path, []byte("x"), 0o600, MkDir(0o755)))
+
+	got, err := os.ReadFile(path)
+	require.NoError(t, err)
+	assert.Equal(t, "x", string(got))
+
+	// The existing directory keeps 0o700; MkDir does not widen it to 0o755.
+	info, err := os.Stat(dir)
+	require.NoError(t, err)
+	assert.Equal(t, os.FileMode(0o700), info.Mode().Perm())
+}

From c4ad2279850ffb3750520cedf3eaeeaab1b6af98 Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Thu, 17 Sep 2026 10:46:16 +0200
Subject: [PATCH 8/8] update glob

---
 libs/atomicfile/atomicfile.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/libs/atomicfile/atomicfile.go b/libs/atomicfile/atomicfile.go
index a8a26cd7d57..09cf5a4b139 100644
--- a/libs/atomicfile/atomicfile.go
+++ b/libs/atomicfile/atomicfile.go
@@ -53,7 +53,7 @@ func Write(path string, data []byte, perm os.FileMode, opts ...Option) error {
 
 	// Temp file in the same directory so the rename stays on one filesystem.
 	// The "." prefix hides a leftover temp file and the ".tmp" suffix lets
-	// callers sweep leftovers with a "*.tmp" glob.
+	// callers sweep leftovers with a ".*.tmp" glob.
 	tmp, err := os.CreateTemp(dir, "."+filepath.Base(path)+".*.tmp")
 	if err != nil {
 		return fmt.Errorf("create temp file for %s: %w", path, err)