feat: Codex account_name 透传并优化池列表 OAuth 标识与额度重置倒计时展示

(cherry picked from commit 1b9176fc4a)
2026-09-02 01:10:23 +08:00 · 2026-03-18 20:58:33 +08:00
parent 8d8cddcef6
commit 6e55968487
16 changed files with 888 additions and 251 deletions
--- a/tests/api/test_provider_oauth_batch_import.py
+++ b/tests/api/test_provider_oauth_batch_import.py
@@ -68,6 +68,7 @@ async def test_standard_batch_import_commits_successes_in_chunks(
    monkeypatch: pytest.MonkeyPatch,
 ) -> None:
    key_ids = count(1)
+    created_auth_configs: list[dict[str, object]] = []

    monkeypatch.setattr(
        oauthmod,
@@ -91,7 +92,9 @@ async def test_standard_batch_import_commits_successes_in_chunks(
        "_parse_standard_oauth_import_entries",
        lambda _raw: [{"refresh_token": f"r-{idx}" + ("x" * 120)} for idx in range(3)],
    )
-    monkeypatch.setattr(oauthmod, "_get_provider_api_formats", lambda _provider: ["responses"])
+    monkeypatch.setattr(
+        oauthmod, "_get_provider_api_formats", lambda _provider: ["responses"]
+    )
    monkeypatch.setattr(
        oauthmod,
        "_release_batch_import_db_connection_before_await",
@@ -113,16 +116,24 @@ async def test_standard_batch_import_commits_successes_in_chunks(
    async def _fake_enrich_auth_config(**kwargs: object) -> dict[str, object]:
        auth_config = dict(kwargs["auth_config"])  # type: ignore[call-overload]
        auth_config["email"] = f"user-{next(key_ids)}@example.com"
+        auth_config["account_name"] = "Workspace Alpha"
        return auth_config

    created_ids = count(1)
    monkeypatch.setattr(oauthmod, "post_oauth_token", _fake_post_oauth_token)
    monkeypatch.setattr(oauthmod, "enrich_auth_config", _fake_enrich_auth_config)
-    monkeypatch.setattr(oauthmod, "_check_duplicate_oauth_account", lambda *_args, **_kwargs: None)
+    monkeypatch.setattr(
+        oauthmod, "_check_duplicate_oauth_account", lambda *_args, **_kwargs: None
+    )
+
+    def _fake_create_oauth_key(*_args: object, **kwargs: object) -> SimpleNamespace:
+        created_auth_configs.append(dict(kwargs["auth_config"]))
+        return SimpleNamespace(id=f"key-{next(created_ids)}")
+
    monkeypatch.setattr(
        oauthmod,
        "_create_oauth_key",
-        lambda *_args, **_kwargs: SimpleNamespace(id=f"key-{next(created_ids)}"),
+        _fake_create_oauth_key,
    )

    db = MagicMock()
@@ -140,6 +151,7 @@ async def test_standard_batch_import_commits_successes_in_chunks(
    assert result.success == 3
    assert result.failed == 0
    assert db.commit.call_count == 2
+    assert created_auth_configs[0]["account_name"] == "Workspace Alpha"


@pytest.mark.asyncio
@@ -152,14 +164,20 @@ async def test_kiro_batch_import_releases_db_connection_before_refresh(
        def __init__(self, data: dict[str, object]) -> None:
            self._data = dict(data)
            self.provider_type = str(data.get("provider_type") or "")
-            self.email = data.get("email") if isinstance(data.get("email"), str) else None
+            self.email = (
+                data.get("email") if isinstance(data.get("email"), str) else None
+            )
            self.auth_method = (
-                data.get("auth_method") if isinstance(data.get("auth_method"), str) else "social"
+                data.get("auth_method")
+                if isinstance(data.get("auth_method"), str)
+                else "social"
            )
            self.refresh_token = str(data.get("refresh_token") or "")

        @staticmethod
-        def validate_required_fields(_cred: dict[str, object]) -> tuple[bool, str | None]:
+        def validate_required_fields(
+            _cred: dict[str, object],
+        ) -> tuple[bool, str | None]:
            return True, None

        @classmethod
@@ -185,7 +203,9 @@ async def test_kiro_batch_import_releases_db_connection_before_refresh(
        FakeKiroAuthConfig,
    )

-    async def _fake_refresh_access_token(*_args: object, **_kwargs: object) -> tuple[str, object]:
+    async def _fake_refresh_access_token(
+        *_args: object, **_kwargs: object
+    ) -> tuple[str, object]:
        raise RuntimeError("refresh token reused")

    monkeypatch.setattr(
--- a/tests/core/test_provider_oauth_utils_codex.py
+++ b/tests/core/test_provider_oauth_utils_codex.py
@@ -1,10 +1,15 @@
+# pyright: reportMissingImports=false
+
 from __future__ import annotations

 import json
+from unittest.mock import AsyncMock

 import jwt
+import pytest

-from src.core.provider_oauth_utils import parse_codex_id_token
+from src.core import provider_oauth_utils as module
+from src.core.provider_oauth_utils import enrich_auth_config, parse_codex_id_token


 def _encode_unsigned_jwt(payload: dict[str, object]) -> str:
@@ -21,7 +26,9 @@ def test_parse_codex_id_token_extracts_auth_claim_fields() -> None:
                "chatgpt_account_user_id": "user-1__acc-1",
                "chatgpt_plan_type": "team",
                "chatgpt_user_id": "user-1",
-                "organizations": [{"id": "org-1", "title": "Personal", "is_default": True}],
+                "organizations": [
+                    {"id": "org-1", "title": "Personal", "is_default": True}
+                ],
            },
        }
    )
@@ -72,3 +79,45 @@ def test_parse_codex_id_token_accepts_dict_payload() -> None:
        "plan_type": "enterprise",
        "user_id": "user-3",
    }
+
+
+@pytest.mark.asyncio
+async def test_enrich_auth_config_codex_adds_current_account_name() -> None:
+    from src.services.provider.envelope import ensure_providers_bootstrapped
+
+    access_token = _encode_unsigned_jwt(
+        {
+            "email": "u@example.com",
+            "https://api.openai.com/auth": {
+                "chatgpt_account_id": "acc-1",
+                "chatgpt_account_user_id": "user-1__acc-1",
+                "chatgpt_plan_type": "team",
+                "chatgpt_user_id": "user-1",
+            },
+        }
+    )
+
+    ensure_providers_bootstrapped()
+
+    fetch_account_name = AsyncMock(return_value="Workspace Alpha")
+    original = module.fetch_openai_account_name
+    module.fetch_openai_account_name = fetch_account_name
+    try:
+        out = await enrich_auth_config(
+            provider_type="codex",
+            auth_config={},
+            token_response={"access_token": access_token},
+            access_token=access_token,
+            proxy_config=None,
+        )
+    finally:
+        module.fetch_openai_account_name = original
+
+    fetch_account_name.assert_awaited_once_with(
+        access_token,
+        "acc-1",
+        proxy_config=None,
+        timeout_seconds=10.0,
+    )
+    assert out["account_id"] == "acc-1"
+    assert out["account_name"] == "Workspace Alpha"
--- a/tests/services/test_codex_request.py
+++ b/tests/services/test_codex_request.py
@@ -1,6 +1,7 @@
 from __future__ import annotations

 from types import SimpleNamespace
+from unittest.mock import AsyncMock

 import jwt
 import pytest
@@ -17,18 +18,28 @@ def test_codex_provider_behavior_has_no_runtime_envelope_or_variants() -> None:
    assert behavior.cross_format_variant is None


-def test_openai_cli_normalizer_request_from_internal_codex_variant_preserves_store() -> None:
-    from src.core.api_format.conversion.normalizers.openai_cli import OpenAICliNormalizer
+def test_openai_cli_normalizer_request_from_internal_codex_variant_preserves_store() -> (
+    None
+):
+    from src.core.api_format.conversion.normalizers.openai_cli import (
+        OpenAICliNormalizer,
+    )

    normalizer = OpenAICliNormalizer()
-    internal = normalizer.request_to_internal({"model": "gpt-test", "input": [], "store": True})
+    internal = normalizer.request_to_internal(
+        {"model": "gpt-test", "input": [], "store": True}
+    )
    out = normalizer.request_from_internal(internal, target_variant="codex")

    assert out["store"] is True


-def test_openai_cli_normalizer_request_from_internal_codex_variant_does_not_inject_store() -> None:
-    from src.core.api_format.conversion.normalizers.openai_cli import OpenAICliNormalizer
+def test_openai_cli_normalizer_request_from_internal_codex_variant_does_not_inject_store() -> (
+    None
+):
+    from src.core.api_format.conversion.normalizers.openai_cli import (
+        OpenAICliNormalizer,
+    )

    normalizer = OpenAICliNormalizer()
    internal = normalizer.request_to_internal({"model": "gpt-test", "input": []})
@@ -37,9 +48,13 @@ def test_openai_cli_normalizer_request_from_internal_codex_variant_does_not_inje
    assert "store" not in out


-def test_openai_cli_normalizer_codex_variant_keeps_instructions_missing_for_default_rule() -> None:
+def test_openai_cli_normalizer_codex_variant_keeps_instructions_missing_for_default_rule() -> (
+    None
+):
    from src.api.handlers.base.request_builder import apply_body_rules
-    from src.core.api_format.conversion.normalizers.openai_cli import OpenAICliNormalizer
+    from src.core.api_format.conversion.normalizers.openai_cli import (
+        OpenAICliNormalizer,
+    )
    from src.core.api_format.metadata import CODEX_DEFAULT_BODY_RULES

    normalizer = OpenAICliNormalizer()
@@ -53,7 +68,9 @@ def test_openai_cli_normalizer_codex_variant_keeps_instructions_missing_for_defa


 def test_openai_cli_normalizer_patch_for_codex_is_noop() -> None:
-    from src.core.api_format.conversion.normalizers.openai_cli import OpenAICliNormalizer
+    from src.core.api_format.conversion.normalizers.openai_cli import (
+        OpenAICliNormalizer,
+    )

    normalizer = OpenAICliNormalizer()
    out = normalizer.patch_for_variant(
@@ -73,7 +90,9 @@ def test_openai_cli_normalizer_patch_for_codex_is_noop() -> None:

 def test_codex_passthrough_builder_preserves_real_codex_headers() -> None:
    builder = PassthroughRequestBuilder()
-    endpoint = SimpleNamespace(api_family="openai", endpoint_kind="cli", header_rules=None)
+    endpoint = SimpleNamespace(
+        api_family="openai", endpoint_kind="cli", header_rules=None
+    )
    key = SimpleNamespace(api_key="unused")

    headers = builder.build_headers(
@@ -139,7 +158,9 @@ def _encode_unsigned_jwt(payload: dict[str, object]) -> str:


@pytest.mark.asyncio
-async def test_enrich_codex_uses_access_token_when_id_token_missing() -> None:
+async def test_enrich_codex_uses_access_token_when_id_token_missing(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
    from src.services.provider.adapters.codex.plugin import enrich_codex

    access_token = _encode_unsigned_jwt(
@@ -154,6 +175,10 @@ async def test_enrich_codex_uses_access_token_when_id_token_missing() -> None:
    )

    auth_config: dict[str, object] = {}
+    monkeypatch.setattr(
+        "src.core.provider_oauth_utils.fetch_openai_account_name",
+        AsyncMock(return_value="Workspace Alpha"),
+    )
    out = await enrich_codex(
        auth_config=auth_config,
        token_response={"access_token": access_token},
@@ -163,5 +188,6 @@ async def test_enrich_codex_uses_access_token_when_id_token_missing() -> None:

    assert out["email"] == "u@example.com"
    assert out["account_id"] == "acc-access"
+    assert out["account_name"] == "Workspace Alpha"
    assert out["plan_type"] == "team"
    assert out["user_id"] == "user-access"
--- a/tests/services/test_provider_auth_detached.py
+++ b/tests/services/test_provider_auth_detached.py
@@ -6,6 +6,7 @@ from datetime import datetime, timezone
 from types import SimpleNamespace
 from typing import Any

+import httpx
 import pytest

 from src.services.provider import auth as module
@@ -46,7 +47,9 @@ class _FakeSessionCtx:
        return False


-def _install_module(monkeypatch: pytest.MonkeyPatch, name: str, attrs: dict[str, Any]) -> None:
+def _install_module(
+    monkeypatch: pytest.MonkeyPatch, name: str, attrs: dict[str, Any]
+) -> None:
    fake_module = types.ModuleType(name)
    for key, value in attrs.items():
        setattr(fake_module, key, value)
@@ -88,7 +91,9 @@ def test_mark_refresh_token_invalid_persists_detached_key(
        module, "object_session", lambda _key: (_ for _ in ()).throw(RuntimeError())
    )
    _install_module(
-        monkeypatch, "src.database", {"create_session": lambda: _FakeSessionCtx(fake_db)}
+        monkeypatch,
+        "src.database",
+        {"create_session": lambda: _FakeSessionCtx(fake_db)},
    )
    _install_module(
        monkeypatch,
@@ -105,7 +110,9 @@ def test_mark_refresh_token_invalid_persists_detached_key(
    assert fake_db.committed is True
    assert key.oauth_invalid_at is not None
    assert row.oauth_invalid_at is not None
-    assert str(key.oauth_invalid_reason).startswith("[REFRESH_FAILED] Token 续期失败 (401)")
+    assert str(key.oauth_invalid_reason).startswith(
+        "[REFRESH_FAILED] Token 续期失败 (401)"
+    )
    assert "refresh_token_reused" in str(row.oauth_invalid_reason)


@@ -144,3 +151,66 @@ def test_persist_refreshed_token_clears_legacy_token_invalidated_account_block(
    assert key.auth_config == 'enc:{"refresh_token": "rt-2"}'
    assert key.oauth_invalid_at is None
    assert key.oauth_invalid_reason is None
+
+
+@pytest.mark.asyncio
+async def test_refresh_generic_oauth_token_persists_enriched_account_name(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    key = SimpleNamespace(id="key-1")
+    endpoint = SimpleNamespace()
+    template = SimpleNamespace(
+        oauth=SimpleNamespace(
+            token_url="https://example.com/oauth/token",
+            client_id="client-id",
+            client_secret=None,
+            scopes=[],
+        )
+    )
+    persisted: dict[str, Any] = {}
+
+    async def _fake_post_oauth_token(**_kwargs: Any) -> httpx.Response:
+        return httpx.Response(
+            200,
+            json={
+                "access_token": "new-token",
+                "refresh_token": "rt-2",
+                "expires_in": 3600,
+                "token_type": "Bearer",
+            },
+            request=httpx.Request("POST", "https://example.com/oauth/token"),
+        )
+
+    async def _fake_enrich_auth_config(**kwargs: Any) -> dict[str, Any]:
+        auth_config = dict(kwargs["auth_config"])
+        auth_config["account_name"] = "Workspace Alpha"
+        return auth_config
+
+    monkeypatch.setattr(module, "_get_proxy_config", lambda *_args: None)
+    monkeypatch.setattr(module, "post_oauth_token", _fake_post_oauth_token)
+    monkeypatch.setattr(module, "enrich_auth_config", _fake_enrich_auth_config)
+    monkeypatch.setattr(
+        module,
+        "_persist_refreshed_token",
+        lambda _key, _access_token, token_meta: persisted.update(
+            {"access_token": _access_token, "token_meta": dict(token_meta)}
+        ),
+    )
+
+    token_meta = {
+        "provider_type": "codex",
+        "refresh_token": "rt-1",
+    }
+
+    refreshed = await module._refresh_generic_oauth_token(
+        key,
+        endpoint,
+        template,
+        "codex",
+        "rt-1",
+        token_meta,
+    )
+
+    assert refreshed["account_name"] == "Workspace Alpha"
+    assert persisted["access_token"] == "new-token"
+    assert persisted["token_meta"]["account_name"] == "Workspace Alpha"
--- a/tests/unit/test_pool_oauth_expires.py
+++ b/tests/unit/test_pool_oauth_expires.py
@@ -38,3 +38,10 @@ def test_derive_oauth_expires_at_fallback_to_legacy_datetime() -> None:
    )

    assert pool_routes._derive_oauth_expires_at(key) == 1772586123
+
+
+def test_derive_oauth_account_name_from_auth_config() -> None:
+    assert (
+        pool_routes._derive_oauth_account_name({"account_name": " Workspace Alpha "})
+        == "Workspace Alpha"
+    )
--- a/tests/unit/test_provider_key_response_builder.py
+++ b/tests/unit/test_provider_key_response_builder.py
@@ -18,7 +18,7 @@ def test_build_key_response_includes_codex_identity_metadata(
        api_formats=["openai:chat"],
        auth_type="oauth",
        api_key="enc-access-token",
-        auth_config='{"email":"u@example.com","plan_type":"team","account_id":"acc-1","account_user_id":"user-1__acc-1","organizations":[{"id":"org-1","title":"Personal","is_default":true,"role":"owner"}],"expires_at":123456}',
+        auth_config='{"email":"u@example.com","plan_type":"team","account_id":"acc-1","account_name":"Workspace Alpha","account_user_id":"user-1__acc-1","organizations":[{"id":"org-1","title":"Personal","is_default":true,"role":"owner"}],"expires_at":123456}',
        name="codex-user",
    )
    now = datetime.now(timezone.utc)
@@ -49,6 +49,7 @@ def test_build_key_response_includes_codex_identity_metadata(
    assert result.oauth_email == "u@example.com"
    assert result.oauth_plan_type == "team"
    assert result.oauth_account_id == "acc-1"
+    assert result.oauth_account_name == "Workspace Alpha"
    assert result.oauth_account_user_id == "user-1__acc-1"
    assert len(result.oauth_organizations) == 1
    assert result.oauth_organizations[0].title == "Personal"