refactor(BA-5528): split chat tokens into deployment_chat_config; both stores as BaseModel

- Move user-managed API keys out of DeploymentChatCache and into a new
  DeploymentChatConfig store backed by ~/.backend.ai/deployment_chat_config.json.
  The cache is now strictly endpoint metadata that the manager owns;
  the config is strictly the user-supplied tokens.
- Make both DeploymentChatCache and DeploymentChatConfig Pydantic
  BaseModel subclasses. Drop the manual to_dict helpers in favor of
  model_dump_json for serialization. Per-entry validation in load
  functions still skips malformed records individually.
- Update chat / chat-config / clear / show to operate on both stores
  in tandem, preserving the single ./bai deployment chat-config UX.

Author: jopemachine

src/ai/backend/client/cli/v2/deployment/chat.py

@@ -18,6 +18,11 @@
     load_chat_cache,
     save_chat_cache,
 )
+from ai.backend.client.cli.v2.deployment_chat_config import (
+    IncompatibleChatConfigError,
+    load_chat_config,
+    save_chat_config,
+)
 from ai.backend.client.cli.v2.helpers import create_v2_registry, load_v2_config
 
 
@@ -75,12 +80,16 @@ def chat(
         DeploymentChatClient,
     )
 
-    config = load_v2_config()
+    connection = load_v2_config()
 
     try:
         cache = load_chat_cache()
     except IncompatibleChatCacheError as e:
         raise click.ClickException(str(e)) from e
+    try:
+        chat_config = load_chat_config()
+    except IncompatibleChatConfigError as e:
+        raise click.ClickException(str(e)) from e
 
     if not isinstance(params, dict):
         raise click.ClickException("--params must be a JSON object.")
@@ -90,7 +99,7 @@ def chat(
     async def _ensure_endpoint_entry() -> DeploymentChatCacheEntry:
         if entry is not None and entry.endpoint_url:
             return entry
-        registry = await create_v2_registry(config)
+        registry = await create_v2_registry(connection)
         try:
             deployment = await registry.deployment.get(deployment_id)
         finally:
@@ -128,9 +137,9 @@ async def _run() -> None:
             "model": request_model,
             "messages": [{"role": "user", "content": content}],
         }
-        api_key = cache.get_token(deployment_id)
+        api_key = chat_config.get_token(deployment_id)
         async with DeploymentChatClient(
-            skip_ssl_verification=config.skip_ssl_verification,
+            skip_ssl_verification=connection.skip_ssl_verification,
         ) as client:
             try:
                 response = await client.chat_completion(
@@ -139,8 +148,8 @@ async def _run() -> None:
                     body,
                 )
             except DeploymentChatAuthError as e:
-                cache.clear_token(deployment_id)
-                save_chat_cache(cache)
+                chat_config.clear_token(deployment_id)
+                save_chat_config(chat_config)
                 raise click.ClickException(
                     f"The inference endpoint rejected the configured API key for "
                     f"deployment {deployment_id}. The cached key has been cleared.\n"

src/ai/backend/client/cli/v2/deployment/chat_config.py

@@ -14,9 +14,14 @@
     DeploymentChatCacheEntry,
     IncompatibleChatCacheError,
     load_chat_cache,
-    mask_token,
     save_chat_cache,
 )
+from ai.backend.client.cli.v2.deployment_chat_config import (
+    IncompatibleChatConfigError,
+    load_chat_config,
+    mask_token,
+    save_chat_config,
+)
 from ai.backend.client.cli.v2.helpers import create_v2_registry, load_v2_config
 
 
@@ -81,23 +86,27 @@ def set_(
     if api_key and no_token:
         raise click.ClickException("--token and --no-token are mutually exclusive.")
 
-    config = load_v2_config()
+    connection = load_v2_config()
     try:
         cache = load_chat_cache()
     except IncompatibleChatCacheError as e:
         raise click.ClickException(str(e)) from e
+    try:
+        chat_config_store = load_chat_config()
+    except IncompatibleChatConfigError as e:
+        raise click.ClickException(str(e)) from e
 
-    existing = cache.get(deployment_id)
+    existing_entry = cache.get(deployment_id)
     resolved_key: str | None
     if no_token:
         resolved_key = None
     elif api_key is not None:
         resolved_key = api_key
     else:
-        resolved_key = cache.get_token(deployment_id)
+        resolved_key = chat_config_store.get_token(deployment_id)
 
     async def _run() -> None:
-        registry = await create_v2_registry(config)
+        registry = await create_v2_registry(connection)
         try:
             deployment = await registry.deployment.get(deployment_id)
         finally:
@@ -115,8 +124,8 @@ async def _run() -> None:
             served_model = await _discover_model(
                 endpoint_url,
                 resolved_key,
-                config.skip_ssl_verification,
-                existing.default_model if existing is not None else None,
+                connection.skip_ssl_verification,
+                existing_entry.default_model if existing_entry is not None else None,
             )
 
         cache.upsert(
@@ -127,11 +136,13 @@ async def _run() -> None:
                 last_synced_at=datetime.now(UTC),
             ),
         )
+        save_chat_cache(cache)
         if resolved_key is None:
-            cache.clear_token(deployment_id)
+            chat_config_store.clear_token(deployment_id)
         else:
-            cache.set_token(deployment_id, resolved_key)
-        save_chat_cache(cache)
+            chat_config_store.set_token(deployment_id, resolved_key)
+        save_chat_config(chat_config_store)
+
         click.echo(f"Updated chat cache entry for deployment {deployment_id}.")
         if served_model:
             click.echo(f"  default_model: {served_model}")
@@ -180,21 +191,25 @@ def show(deployment_id: UUID | None) -> None:
         cache = load_chat_cache()
     except IncompatibleChatCacheError as e:
         raise click.ClickException(str(e)) from e
+    try:
+        chat_config_store = load_chat_config()
+    except IncompatibleChatConfigError as e:
+        raise click.ClickException(str(e)) from e
 
     if deployment_id is not None:
         entry = cache.get(deployment_id)
-        token = cache.get_token(deployment_id)
+        token = chat_config_store.get_token(deployment_id)
         if entry is None and token is None:
             raise click.ClickException(f"No chat cache entry for deployment {deployment_id}.")
         _print_entry(deployment_id, entry, token)
         return
 
-    dep_ids = set(cache.entries) | set(cache.tokens)
+    dep_ids = set(cache.deployments) | set(chat_config_store.tokens)
     if not dep_ids:
         click.echo("No chat cache entries.")
         return
     for dep_id in dep_ids:
-        _print_entry(dep_id, cache.get(dep_id), cache.get_token(dep_id))
+        _print_entry(dep_id, cache.get(dep_id), chat_config_store.get_token(dep_id))
         click.echo("")
 
 
@@ -206,8 +221,18 @@ def clear(deployment_id: UUID) -> None:
         cache = load_chat_cache()
     except IncompatibleChatCacheError as e:
         raise click.ClickException(str(e)) from e
-    if cache.remove(deployment_id):
+    try:
+        chat_config_store = load_chat_config()
+    except IncompatibleChatConfigError as e:
+        raise click.ClickException(str(e)) from e
+
+    removed_entry = cache.remove(deployment_id)
+    removed_token = chat_config_store.clear_token(deployment_id)
+    if removed_entry:
         save_chat_cache(cache)
+    if removed_token:
+        save_chat_config(chat_config_store)
+    if removed_entry or removed_token:
         click.echo(f"Removed chat cache entry for deployment {deployment_id}.")
     else:
         click.echo(f"No chat cache entry for deployment {deployment_id}.")

src/ai/backend/client/cli/v2/deployment_chat_cache.py

@@ -1,13 +1,11 @@
-"""Local cache for ``./bai deployment chat`` per-deployment settings.
+"""Local cache for ``./bai deployment chat`` per-deployment endpoint metadata.
 
-Persists the manager-resolved ``endpoint_url`` and the served model name
-discovered from the inference endpoint, plus a separate map of API keys
-the user registered through ``./bai deployment chat-config set``. The
-endpoint entry is auto-managed (refetched when missing); the token is
-user-supplied and never auto-discovered.
+Stores the manager-resolved ``endpoint_url`` and the served model name
+discovered from the inference endpoint. Auto-managed: refetched on cache
+miss, never user-edited. The user-supplied API key lives in a separate
+file managed by ``deployment_chat_config``.
 
-Stored as a single JSON file at ``~/.backend.ai/deployment_chat.json``
-with ``0600`` permissions because the API keys are kept in plaintext.
+Persisted as a JSON file at ``~/.backend.ai/deployment_chat.json``.
 """
 
 from __future__ import annotations
@@ -16,13 +14,11 @@
 import os
 import stat
 import tempfile
-from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Any
 from uuid import UUID
 
-from pydantic import BaseModel, ConfigDict, ValidationError
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
 
 from ai.backend.client.cli.v2.helpers import CONFIG_DIR
 
@@ -40,60 +36,28 @@ class DeploymentChatCacheEntry(BaseModel):
     last_synced_at: datetime
 
 
-@dataclass
-class DeploymentChatCache:
-    """In-memory representation of the chat cache file.
+class DeploymentChatCache(BaseModel):
+    """In-memory representation of the chat cache file."""
 
-    ``entries`` is the auto-managed endpoint cache; ``tokens`` is the
-    user-managed API-key store. They are kept in the same file under
-    distinct top-level keys.
-    """
-
-    entries: dict[UUID, DeploymentChatCacheEntry] = field(default_factory=dict)
-    tokens: dict[UUID, str] = field(default_factory=dict)
+    schema_version: int = Field(default=CHAT_CACHE_SCHEMA_VERSION)
+    deployments: dict[UUID, DeploymentChatCacheEntry] = Field(default_factory=dict)
 
     def get(self, deployment_id: UUID) -> DeploymentChatCacheEntry | None:
-        return self.entries.get(deployment_id)
+        return self.deployments.get(deployment_id)
 
     def upsert(self, deployment_id: UUID, entry: DeploymentChatCacheEntry) -> None:
-        self.entries[deployment_id] = entry
+        self.deployments[deployment_id] = entry
 
     def remove(self, deployment_id: UUID) -> bool:
-        had_entry = self.entries.pop(deployment_id, None) is not None
-        had_token = self.tokens.pop(deployment_id, None) is not None
-        return had_entry or had_token
-
-    def get_token(self, deployment_id: UUID) -> str | None:
-        return self.tokens.get(deployment_id)
-
-    def set_token(self, deployment_id: UUID, token: str) -> None:
-        self.tokens[deployment_id] = token
-
-    def clear_token(self, deployment_id: UUID) -> bool:
-        return self.tokens.pop(deployment_id, None) is not None
-
-    def to_dict(self) -> dict[str, Any]:
-        return {
-            "schema_version": CHAT_CACHE_SCHEMA_VERSION,
-            "deployments": {
-                str(dep_id): entry.model_dump(mode="json") for dep_id, entry in self.entries.items()
-            },
-            "tokens": {str(dep_id): token for dep_id, token in self.tokens.items()},
-        }
+        return self.deployments.pop(deployment_id, None) is not None
 
 
 class IncompatibleChatCacheError(Exception):
     """Raised when the on-disk cache file uses a newer schema than this build."""
 
 
 def load_chat_cache(path: Path = CHAT_CACHE_FILE) -> DeploymentChatCache:
-    """Load the chat cache; return an empty cache when the file is absent or unreadable.
-
-    A corrupted JSON file or unreadable file is treated as an empty cache —
-    individual malformed entries are skipped rather than aborting the whole
-    load. A schema version newer than this build raises
-    :class:`IncompatibleChatCacheError` so the caller can warn the user.
-    """
+    """Load the chat cache; return an empty cache when the file is absent or unreadable."""
     if not path.exists():
         return DeploymentChatCache()
     try:
@@ -109,7 +73,7 @@ def load_chat_cache(path: Path = CHAT_CACHE_FILE) -> DeploymentChatCache:
             f"deployment_chat.json schema version {schema} is newer than supported "
             f"{CHAT_CACHE_SCHEMA_VERSION}; please upgrade the client."
         )
-    entries: dict[UUID, DeploymentChatCacheEntry] = {}
+    deployments: dict[UUID, DeploymentChatCacheEntry] = {}
     deployments_raw = raw.get("deployments") or {}
     if isinstance(deployments_raw, dict):
         for key, value in deployments_raw.items():
@@ -120,26 +84,16 @@ def load_chat_cache(path: Path = CHAT_CACHE_FILE) -> DeploymentChatCache:
             if not isinstance(value, dict):
                 continue
             try:
-                entries[dep_id] = DeploymentChatCacheEntry.model_validate(value)
+                deployments[dep_id] = DeploymentChatCacheEntry.model_validate(value)
             except ValidationError:
                 continue
-    tokens: dict[UUID, str] = {}
-    tokens_raw = raw.get("tokens") or {}
-    if isinstance(tokens_raw, dict):
-        for key, value in tokens_raw.items():
-            try:
-                dep_id = UUID(str(key))
-            except ValueError:
-                continue
-            if isinstance(value, str):
-                tokens[dep_id] = value
-    return DeploymentChatCache(entries=entries, tokens=tokens)
+    return DeploymentChatCache(deployments=deployments)
 
 
 def save_chat_cache(cache: DeploymentChatCache, path: Path = CHAT_CACHE_FILE) -> None:
-    """Atomically write the chat cache and enforce ``0600`` permissions."""
+    """Atomically write the chat cache."""
     path.parent.mkdir(parents=True, exist_ok=True)
-    payload = json.dumps(cache.to_dict(), indent=2, ensure_ascii=False)
+    payload = cache.model_dump_json(indent=2)
     fd, tmp_path_str = tempfile.mkstemp(
         prefix=path.name + ".",
         suffix=".tmp",
@@ -155,12 +109,3 @@ def save_chat_cache(cache: DeploymentChatCache, path: Path = CHAT_CACHE_FILE) ->
         if tmp_path.exists():
             tmp_path.unlink(missing_ok=True)
         raise
-
-
-def mask_token(token: str | None) -> str:
-    """Render a token as ``sk-***...***xxxx`` for diagnostic display."""
-    if token is None:
-        return "<unset>"
-    if len(token) <= 8:
-        return "***"
-    return f"{token[:3]}***...***{token[-4:]}"