From 72c0bde8a914b1089c9be56fe84dc11dfec4c69b Mon Sep 17 00:00:00 2001
From: Michael <52305679+michaelxer@users.noreply.github.com>
Date: Wed, 24 Jun 2026 04:28:53 +0700
Subject: [PATCH] fix: use atomic write in APIKeyManager.save() to prevent
 credential data loss (#4591) (#4597)

* fix: use atomic write in APIKeyManager.save() to prevent data loss

Opening api_keys.json with 'w' truncates the file before writing, so a
crash, disk-full, or mid-write error leaves all stored provider API keys
corrupted. Switch to atomic write (temp file + fsync + os.replace) so
the original file is always intact on any failure.

Fixes #4591

* chore: trigger CI re-run

* chore: update PR description

* chore: fix how-to-test section for description check

---------

Co-authored-by: michaelxer <michaelxer@users.noreply.github.com>
---
 src/api_key_manager.py                    | 19 +++++-
 tests/test_api_key_manager_atomic_save.py | 79 +++++++++++++++++++++++
 2 files changed, 96 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_api_key_manager_atomic_save.py

diff --git a/src/api_key_manager.py b/src/api_key_manager.py
index b3cf9a7b6..337204144 100644
--- a/src/api_key_manager.py
+++ b/src/api_key_manager.py
@@ -81,11 +81,26 @@ class APIKeyManager:
         keys stay encrypted. Loading via load() first would decrypt them and
         write them back as plaintext, which then fails to decrypt on the next
         load() and silently drops those providers.
+
+        Uses atomic write (temp file + os.replace) so a crash, disk-full, or
+        mid-write error never truncates the existing keys file.
         """
         keys = self._load_raw()
         keys[provider] = self.encrypt_api_key(api_key)
-        with open(self.api_keys_file, 'w', encoding="utf-8") as f:
-            json.dump(keys, f)
+        tmp_file = self.api_keys_file + ".tmp"
+        try:
+            with open(tmp_file, 'w', encoding="utf-8") as f:
+                json.dump(keys, f)
+                f.flush()
+                os.fsync(f.fileno())
+            os.replace(tmp_file, self.api_keys_file)
+        except OSError:
+            # Clean up temp file on failure; re-raise so callers see the error
+            try:
+                os.remove(tmp_file)
+            except OSError:
+                pass
+            raise
 
     def load(self) -> Dict[str, str]:
         """Load and decrypt API keys"""
diff --git a/tests/test_api_key_manager_atomic_save.py b/tests/test_api_key_manager_atomic_save.py
new file mode 100644
index 000000000..5ae883b5e
--- /dev/null
+++ b/tests/test_api_key_manager_atomic_save.py
@@ -0,0 +1,79 @@
+"""Test that APIKeyManager.save() uses atomic write to prevent data loss."""
+import os
+import json
+import pytest
+from unittest.mock import patch, mock_open
+from src.api_key_manager import APIKeyManager
+
+
+def test_save_creates_atomic_tmp_file(tmp_path):
+    """Verify save() writes to a temp file and replaces atomically."""
+    mgr = APIKeyManager(str(tmp_path))
+    mgr.save("openai", "sk-test")
+
+    # The final file should exist with the correct content
+    assert os.path.exists(mgr.api_keys_file)
+    with open(mgr.api_keys_file, "r", encoding="utf-8") as f:
+        keys = json.load(f)
+    assert "openai" in keys
+
+    # The temp file should NOT remain after successful save
+    tmp_file = mgr.api_keys_file + ".tmp"
+    assert not os.path.exists(tmp_file)
+
+
+def test_save_preserves_existing_keys_atomically(tmp_path):
+    """Verify atomic save doesn't corrupt other providers' keys."""
+    mgr = APIKeyManager(str(tmp_path))
+    mgr.save("openai", "sk-openai")
+    mgr.save("anthropic", "sk-anthropic")
+
+    loaded = mgr.load()
+    assert loaded["openai"] == "sk-openai"
+    assert loaded["anthropic"] == "sk-anthropic"
+
+
+def test_save_preserves_original_on_write_failure(tmp_path):
+    """If the temp file write fails, the original keys file must survive intact."""
+    mgr = APIKeyManager(str(tmp_path))
+    mgr.save("openai", "sk-original")
+
+    # Now attempt a save that will fail during json.dump
+    with patch("builtins.open", side_effect=OSError("disk full")):
+        with pytest.raises(OSError, match="disk full"):
+            mgr.save("anthropic", "sk-new")
+
+    # Original file must still be intact with the original key
+    loaded = mgr.load()
+    assert loaded == {"openai": "sk-original"}
+    assert "anthropic" not in loaded
+
+
+def test_save_cleans_up_tmp_on_failure(tmp_path):
+    """Temp file should be removed if the write fails."""
+    mgr = APIKeyManager(str(tmp_path))
+    mgr.save("openai", "sk-original")
+
+    tmp_file = mgr.api_keys_file + ".tmp"
+
+    # Force a failure after the temp file is opened
+    original_open = open
+
+    def failing_open(*args, **kwargs):
+        f = original_open(*args, **kwargs)
+        if args and isinstance(args[0], str) and args[0].endswith(".tmp"):
+            # Close the file then raise
+            f.close()
+            raise OSError("simulated write failure")
+        return f
+
+    with patch("builtins.open", side_effect=failing_open):
+        with pytest.raises(OSError):
+            mgr.save("anthropic", "sk-new")
+
+    # Temp file should be cleaned up
+    assert not os.path.exists(tmp_file)
+
+    # Original should be intact
+    loaded = mgr.load()
+    assert loaded == {"openai": "sk-original"}