fix(login): preserve runpodctl credentials in ~/.runpod/config.toml

src/runpod_flash/core/credentials.py

@@ -8,18 +8,26 @@
 
 import logging
 import os
+import re
 from pathlib import Path
 from typing import Optional
 
 import runpod.cli.groups.config.functions as _runpod_config
 
 from runpod.cli.groups.config.functions import (
     get_credentials,
-    set_credentials,
 )
 
 log = logging.getLogger(__name__)
 
+# runpodctl writes top-level `apikey`/`apiurl` keys into the same config.toml
+# that runpod-python uses for its `[default]` profile. We must preserve those
+# (and any other unrelated content) when updating flash's api_key, so flash
+# login does not clobber runpodctl's credentials.
+_DEFAULT_HEADER_RE = re.compile(r"^\s*\[default\]\s*$")
+_SECTION_HEADER_RE = re.compile(r"^\s*\[[^\]]+\]\s*$")
+_API_KEY_LINE_RE = re.compile(r"^\s*api_key\s*=")
+
 _OLD_XDG_PATH = Path.home() / ".config" / "runpod" / "credentials.toml"
 
 
@@ -50,7 +58,11 @@ def get_api_key() -> Optional[str]:
 
 
 def save_api_key(api_key: str) -> Path:
-    """Save API key to ~/.runpod/config.toml via runpod-python.
+    """Save API key into the [default] section of ~/.runpod/config.toml.
+
+    Updates only flash's `[default].api_key` value, preserving any other
+    content in the file (notably runpodctl's top-level `apikey`/`apiurl`
+    keys and other profile sections).
 
     Args:
         api_key: The API key to save.
@@ -59,14 +71,62 @@ def save_api_key(api_key: str) -> Path:
         Path to the credentials file.
     """
     path = get_credentials_path()
-    set_credentials(api_key, overwrite=True)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    existing = path.read_text(encoding="utf-8") if path.exists() else ""
+    new_content = _upsert_default_api_key(existing, api_key)
+    path.write_text(new_content, encoding="utf-8")
+
     try:
         os.chmod(path, 0o600)
     except OSError:
         pass
     return path
 
 
+def _toml_quote(value: str) -> str:
+    escaped = value.replace("\\", "\\\\").replace('"', '\\"')
+    return f'"{escaped}"'
+
+
+def _upsert_default_api_key(content: str, api_key: str) -> str:
+    """Update `[default].api_key` in TOML text, leaving the rest intact."""
+    new_line = f"api_key = {_toml_quote(api_key)}"
+
+    if not content:
+        return f"[default]\n{new_line}\n"
+
+    lines = content.splitlines(keepends=True)
+
+    default_start: Optional[int] = None
+    default_end = len(lines)
+    for i, line in enumerate(lines):
+        if _DEFAULT_HEADER_RE.match(line):
+            default_start = i
+            for j in range(i + 1, len(lines)):
+                if _SECTION_HEADER_RE.match(lines[j]):
+                    default_end = j
+                    break
+            break
+
+    if default_start is None:
+        suffix = "" if content.endswith("\n") else "\n"
+        separator = "\n" if content.strip() else ""
+        return f"{content}{suffix}{separator}[default]\n{new_line}\n"
+
+    for i in range(default_start + 1, default_end):
+        if _API_KEY_LINE_RE.match(lines[i]):
+            ending = "\n" if lines[i].endswith("\n") else ""
+            lines[i] = new_line + ending
+            return "".join(lines)
+
+    insert_idx = default_end
+    while insert_idx > default_start + 1 and lines[insert_idx - 1].strip() == "":
+        insert_idx -= 1
+    lines.insert(insert_idx, new_line + "\n")
+    return "".join(lines)
+
+
 def check_and_migrate_legacy_credentials() -> None:
     """Check for credentials at old XDG path and migrate if needed.
 

tests/unit/test_credentials.py

@@ -1,9 +1,15 @@
 """Unit tests for credential storage and retrieval."""
 
 import os
+import sys
 from pathlib import Path
 from unittest.mock import patch
 
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
 from runpod_flash.core.credentials import (
     get_api_key,
     get_credentials_path,
@@ -69,3 +75,56 @@ def test_sets_restrictive_permissions(self, isolate_credentials_file):
         save_api_key("secret")
         mode = oct(isolate_credentials_file.stat().st_mode & 0o777)
         assert mode == "0o600"
+
+    def test_preserves_runpodctl_top_level_keys(self, isolate_credentials_file):
+        isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True)
+        isolate_credentials_file.write_text(
+            "apikey = 'rpa_runpodctl_key'\n"
+            "apiurl = 'https://api.runpod.io/graphql'\n"
+            "\n"
+            "[default]\n"
+            'api_key = "old-flash-key"\n'
+        )
+        save_api_key("new-flash-key")
+        text = isolate_credentials_file.read_text()
+        assert "apikey = 'rpa_runpodctl_key'" in text
+        assert "apiurl = 'https://api.runpod.io/graphql'" in text
+        parsed = tomllib.loads(text)
+        assert parsed["apikey"] == "rpa_runpodctl_key"
+        assert parsed["apiurl"] == "https://api.runpod.io/graphql"
+        assert parsed["default"]["api_key"] == "new-flash-key"
+
+    def test_adds_default_section_when_missing(self, isolate_credentials_file):
+        isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True)
+        isolate_credentials_file.write_text(
+            "apikey = 'rpa_runpodctl_key'\napiurl = 'https://api.runpod.io/graphql'\n"
+        )
+        save_api_key("flash-key")
+        text = isolate_credentials_file.read_text()
+        parsed = tomllib.loads(text)
+        assert parsed["apikey"] == "rpa_runpodctl_key"
+        assert parsed["apiurl"] == "https://api.runpod.io/graphql"
+        assert parsed["default"]["api_key"] == "flash-key"
+
+    def test_preserves_other_profile_sections(self, isolate_credentials_file):
+        isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True)
+        isolate_credentials_file.write_text(
+            "[default]\n"
+            'api_key = "old"\n'
+            "\n"
+            "[staging]\n"
+            'api_key = "staging-key"\n'
+            'extra = "preserved"\n'
+        )
+        save_api_key("new-default")
+        parsed = tomllib.loads(isolate_credentials_file.read_text())
+        assert parsed["default"]["api_key"] == "new-default"
+        assert parsed["staging"]["api_key"] == "staging-key"
+        assert parsed["staging"]["extra"] == "preserved"
+
+    def test_creates_file_with_only_default_when_missing(
+        self, isolate_credentials_file
+    ):
+        save_api_key("first-key")
+        parsed = tomllib.loads(isolate_credentials_file.read_text())
+        assert parsed == {"default": {"api_key": "first-key"}}