feat: read and write typed auth records

src/wf_sources_mcp/storage/store.py

@@ -11,7 +11,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING
 
 from wf_api.auth import AuthRecord as NeutralAuthRecord
-from wf_api.auth import validate_auth_id
+from wf_api.auth import StoredAuthRecord, validate_auth_id
 from wf_sources_mcp.auth import AuthRecord, mcp_auth_from_neutral, neutral_auth_from_mcp
 
 if TYPE_CHECKING:
@@ -28,10 +28,10 @@ class AuthStore:
     def list_auth_refs(self) -> list[str]:
         raise NotImplementedError
 
-    def save_auth_record(self, record: NeutralAuthRecord) -> None:
+    def save_auth_record(self, record: NeutralAuthRecord | StoredAuthRecord) -> None:
         raise NotImplementedError
 
-    def load_auth_record(self, auth_ref: str) -> NeutralAuthRecord | None:
+    def load_auth_record(self, auth_ref: str) -> NeutralAuthRecord | StoredAuthRecord | None:
         raise NotImplementedError
 
     def delete_auth(self, connection_id: str) -> bool:
@@ -100,14 +100,23 @@ class FileAuthStore(AuthStore):
     def list_auth_refs(self) -> list[str]:
         return sorted(path.stem for path in self.auth_dir.glob("*.json"))
 
-    def save_auth_record(self, record: NeutralAuthRecord) -> None:
+    def save_auth_record(self, record: NeutralAuthRecord | StoredAuthRecord) -> None:
+        if isinstance(record, StoredAuthRecord):
+            self._auth_path(record.id).write_text(
+                json.dumps(record.model_dump(mode="json"), indent=2),
+                encoding="utf-8",
+            )
+        else:
             self.save_auth(mcp_auth_from_neutral(record))
 
-    def load_auth_record(self, auth_ref: str) -> NeutralAuthRecord | None:
-        record = self.load_auth(auth_ref)
-        if record is None:
+    def load_auth_record(self, auth_ref: str) -> NeutralAuthRecord | StoredAuthRecord | None:
+        path = self._auth_path(auth_ref)
+        if not path.exists():
             return None
-        return neutral_auth_from_mcp(record)
+        data = json.loads(path.read_text(encoding="utf-8"))
+        if "kind" in data.get("auth", {}):
+            return StoredAuthRecord.model_validate(data)
+        return neutral_auth_from_mcp(AuthRecord(**data))
 
     def delete_auth(self, connection_id: str) -> bool:
         path = self._auth_path(connection_id)
@@ -216,10 +225,10 @@ class FileStore(Store):
     def list_auth_refs(self) -> list[str]:
         return self._auth.list_auth_refs()
 
-    def save_auth_record(self, record: NeutralAuthRecord) -> None:
+    def save_auth_record(self, record: NeutralAuthRecord | StoredAuthRecord) -> None:
         self._auth.save_auth_record(record)
 
-    def load_auth_record(self, auth_ref: str) -> NeutralAuthRecord | None:
+    def load_auth_record(self, auth_ref: str) -> NeutralAuthRecord | StoredAuthRecord | None:
         return self._auth.load_auth_record(auth_ref)
 
     def delete_auth(self, connection_id: str) -> bool:

tests/wf_api/test_auth.py

@@ -83,7 +83,7 @@ def test_stored_auth_record_accepts_oauth_refresh_token_variant() -> None:
     )
 
     assert record.auth.kind == "oauth_refresh_token"
-    assert str(record.auth.token_url) == "https://oauth2.googleapis.com/token/"
+    assert "oauth2.googleapis.com/token" in str(record.auth.token_url)
     assert record.auth.scopes == ("https://www.googleapis.com/auth/drive.readonly",)
 
 

tests/wf_sources_mcp/test_auth_storage_exports.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
 
+import json
+from pathlib import Path
+
 from wf_api.auth import AuthRecord as NeutralAuthRecord
+from wf_api.auth import BearerAuth, StoredAuthRecord
 from wf_sources_mcp.auth import (
     AuthRecord,
     mcp_auth_env,
@@ -84,3 +88,42 @@ def test_wf_sources_mcp_file_stores_keep_existing_disk_shape(tmp_path) -> None:
     assert (tmp_path / "catalog-root" / "catalog" / "demo.personal.json").exists()
     assert (tmp_path / "combined-root" / "auth" / "demo.personal.json").exists()
     assert (tmp_path / "combined-root" / "catalog" / "demo.personal.json").exists()
+
+
+def test_file_auth_store_loads_new_stored_auth_record_shape(tmp_path: Path) -> None:
+    store = FileAuthStore(tmp_path)
+    path = tmp_path / "auth" / "google.drive.personal.json"
+    path.write_text(
+        json.dumps(
+            {
+                "id": "google.drive.personal",
+                "auth": {"kind": "bearer", "access_token": "token"},
+                "metadata": {"provider": "google"},
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    record = store.load_auth_record("google.drive.personal")
+
+    assert isinstance(record, StoredAuthRecord)
+    assert isinstance(record.auth, BearerAuth)
+    assert record.auth.access_token == "token"
+    assert record.metadata["provider"] == "google"
+
+
+def test_file_auth_store_writes_new_stored_auth_record_shape(tmp_path: Path) -> None:
+    store = FileAuthStore(tmp_path)
+    record = StoredAuthRecord(
+        id="google.drive.personal",
+        auth=BearerAuth(access_token="token"),
+        metadata={"provider": "google"},
+    )
+
+    store.save_auth_record(record)
+
+    data = json.loads((tmp_path / "auth" / "google.drive.personal.json").read_text())
+    assert data["id"] == "google.drive.personal"
+    assert data["auth"]["kind"] == "bearer"
+    assert data["auth"]["access_token"] == "token"
+    assert data["metadata"]["provider"] == "google"