From ca8fbc89a10d2513b4f6a8a676b8e98337f366b4 Mon Sep 17 00:00:00 2001 From: Yaojia Wang Date: Wed, 8 Jul 2026 10:37:42 +0200 Subject: [PATCH] =?UTF-8?q?fix(backend):=20Kimi=20OAuth=20=E8=BD=AE?= =?UTF-8?q?=E8=AF=A2=E7=A7=BB=E5=87=BA=20DB=20=E4=BC=9A=E8=AF=9D=E2=80=94?= =?UTF-8?q?=E2=80=94=E5=8F=AA=E5=9C=A8=E6=9C=80=E7=BB=88=E6=8C=81=E4=B9=85?= =?UTF-8?q?=E5=8C=96=E5=BC=80=E7=9F=AD=E4=BC=9A=E8=AF=9D=EF=BC=88CR-H1?= =?UTF-8?q?=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/api/tests/test_kimi_oauth_endpoints.py | 104 ++++++++++++++++++- apps/api/ww_api/routers/kimi_oauth.py | 109 +++++++++++++------- 2 files changed, 176 insertions(+), 37 deletions(-) diff --git a/apps/api/tests/test_kimi_oauth_endpoints.py b/apps/api/tests/test_kimi_oauth_endpoints.py index 8a4fdd4..1e1887e 100644 --- a/apps/api/tests/test_kimi_oauth_endpoints.py +++ b/apps/api/tests/test_kimi_oauth_endpoints.py @@ -15,6 +15,8 @@ from __future__ import annotations import os +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from typing import Any import pytest @@ -51,12 +53,65 @@ class _ScriptedHttp: return None +class _TrackingSessionFactory: + """记录同一时刻在持会话数（open_count）+ 累计开启次数（total_opened）的 session 工厂替身。 + + CR-H1 断言用：轮询期间 open_count 应恒 0（无会话被长时间占用），持久化时 total_opened≥1。 + """ + + def __init__(self) -> None: + self.open_count = 0 + self.total_opened = 0 + + def __call__(self) -> Any: + from fakes_projects import FakeSession + + @asynccontextmanager + async def _cm() -> AsyncIterator[Any]: + self.open_count += 1 + self.total_opened += 1 + try: + yield FakeSession() + finally: + self.open_count -= 1 + + return _cm() + + +class _ObservingHttp: + """观测型 scripted http：每次 POST 到 TOKEN_URL（轮询）时记录当下在持会话数。aclose 无操作。""" + + def __init__( + self, + responses: list[_FakeResponse], + *, + factory: _TrackingSessionFactory, + token_url: str, + ) -> None: + self._responses = responses + self._factory = factory + self._token_url = token_url + self.calls: list[tuple[str, dict[str, str]]] = [] + self.open_during_poll: list[int] = [] + + async def post(self, url: str, *, data: dict[str, str]) -> _FakeResponse: + idx = len(self.calls) + self.calls.append((url, data)) + if url == self._token_url: # 轮询 POST（区别于 device_authorization POST）。 + self.open_during_poll.append(self._factory.open_count) + return self._responses[idx] + + async def aclose(self) -> None: + return None + + def _make_client( *, store: FakeCredentialStore, - http: _ScriptedHttp, + http: Any, enc_key: str, monkeypatch: pytest.MonkeyPatch, + session_factory: Any = None, ) -> TestClient: os.environ["CREDENTIAL_ENC_KEY"] = enc_key from ww_config import get_settings @@ -88,11 +143,13 @@ def _make_client( job_repo = FakeJobRepo() monkeypatch.setattr(job_runner, "SqlJobRepo", lambda _session: job_repo) + sf = session_factory if session_factory is not None else FakeSessionFactory() + app = create_app() app.dependency_overrides[get_credential_store] = lambda: store app.dependency_overrides[get_job_repo] = lambda: job_repo app.dependency_overrides[get_session] = lambda: FakeSession() - app.dependency_overrides[get_session_factory] = lambda: FakeSessionFactory() + app.dependency_overrides[get_session_factory] = lambda: sf app.dependency_overrides[routers.kimi_oauth._default_http_client] = lambda: http return TestClient(app) @@ -163,6 +220,49 @@ def test_start_polls_through_authorization_pending(monkeypatch: pytest.MonkeyPat assert (STUB_OWNER_ID, "kimi-code") in store.oauth +def test_poll_loop_does_not_hold_db_session(monkeypatch: pytest.MonkeyPatch) -> None: + """CR-H1：设备授权轮询期间**不得**占用 DB 会话（否则连接池被长时间空占→饥饿）。 + + 轮询循环只做 HTTP（无 DB）；仅最终持久化才开一个**短**会话。断言：轮询 POST 期间 + session_factory 未开会话（open_count==0），而持久化确开过短会话；凭据已落库、响应不泄 token。 + """ + from ww_api.services.kimi_oauth import TOKEN_URL + + enc_key = Fernet.generate_key().decode() + store = FakeCredentialStore() + factory = _TrackingSessionFactory() + # device auth → poll #1 pending → poll #2 success（≥2 次轮询 POST）。 + http = _ObservingHttp( + [ + _device_auth_response(), + _FakeResponse(400, {"error": "authorization_pending"}), + _token_response(), + ], + factory=factory, + token_url=TOKEN_URL, + ) + client = _make_client( + store=store, + http=http, + enc_key=enc_key, + monkeypatch=monkeypatch, + session_factory=factory, + ) + + resp = client.post("/settings/providers/kimi-code/oauth/start") + + assert resp.status_code == 202 + # 核心断言：轮询 POST 期间无 session_factory 会话在持。 + assert http.open_during_poll # 确有轮询 POST 发生 + assert all(c == 0 for c in http.open_during_poll) + # 持久化时确开过一个短会话。 + assert factory.total_opened >= 1 + # 凭据已落库；响应绝不含 token。 + assert (STUB_OWNER_ID, "kimi-code") in store.oauth + assert "secret-acc" not in resp.text + assert "secret-ref" not in resp.text + + def test_disconnect_clears_credential(monkeypatch: pytest.MonkeyPatch) -> None: enc_key = Fernet.generate_key().decode() store = FakeCredentialStore() diff --git a/apps/api/ww_api/routers/kimi_oauth.py b/apps/api/ww_api/routers/kimi_oauth.py index d68650e..619d52e 100644 --- a/apps/api/ww_api/routers/kimi_oauth.py +++ b/apps/api/ww_api/routers/kimi_oauth.py @@ -17,6 +17,7 @@ token 仅以 Fernet 密文存 `provider_credentials.oauth_enc`。 from __future__ import annotations import asyncio +import uuid from typing import Annotated, Any import httpx @@ -47,6 +48,7 @@ from ww_api.services.kimi_oauth import ( AuthorizationPending, DeviceAuth, SlowDown, + TokenSet, decrypt_oauth_bundle, encrypt_oauth_bundle, poll_token, @@ -74,49 +76,85 @@ def _default_http_client() -> AsyncHttpClient: return httpx.AsyncClient(timeout=30.0) -def _make_poll_work(device: DeviceAuth) -> Any: - """构造后台轮询工作闭包：循环 poll token 直到成功/过期/拒绝。 +# SlowDown 退避每次增大的间隔量（秒）。 +_SLOW_DOWN_STEP_SECONDS = 5 - `work(session)` 自建 httpx 客户端 + 凭据 store（用 `run_job` 传入的独立 session）→ - 按 `interval` 轮询 token → 成功则加密存 `oauth_enc` 并返回非密 job 结果 → 过期/拒绝则 - 抛 AppError（`run_job` 置 job failed）。**token 绝不进 job 结果/日志**。 + +async def _poll_for_token(http: AsyncHttpClient, device: DeviceAuth) -> TokenSet: + """设备授权轮询循环（**纯 HTTP，无 DB 会话**，CR-H1）：直到拿到 token 或超时。 + + `authorization_pending` → 继续；`slow_down` → 增大 interval；成功 → 返回 `TokenSet`； + 次数耗尽 → 抛 `AppError`（视作过期，由调用方经 `run_job` 置 job failed）。 + """ + interval = max(1, device.interval) + for _ in range(_MAX_POLL_ATTEMPTS): + await asyncio.sleep(interval) + try: + return await poll_token(http, device.device_code) + except AuthorizationPending: + continue + except SlowDown: + interval += _SLOW_DOWN_STEP_SECONDS + continue + raise AppError( + ErrorCode.LLM_UNAVAILABLE, + "Kimi 设备授权轮询超时", + {"provider": KIMI_CODE_PROVIDER}, + ) + + +def _make_persist_work(token: TokenSet) -> Any: + """构造**最终持久化**工作闭包：只做「加密 token → upsert oauth 凭据」的一次短事务。 + + `run_job` 用独立 session 调用此 work；轮询已在会话外完成，故此处会话极短。 + **明文 token 绝不进 job 结果/日志。** """ async def work(session: AsyncSession) -> dict[str, Any]: enc_key = get_settings().credential_enc_key.get_secret_value() store = SqlCredentialStore(session) - interval = max(1, device.interval) - - http = _default_http_client() - try: - for _ in range(_MAX_POLL_ATTEMPTS): - await asyncio.sleep(interval) - try: - token = await poll_token(http, device.device_code) - except AuthorizationPending: - continue - except SlowDown: - interval += 5 - continue - # 成功：加密 token 包入库（auth_type=oauth）。明文 token 不进结果/日志。 - blob = encrypt_oauth_bundle(token, key=enc_key) - await store.upsert_oauth_credential(STUB_OWNER_ID, KIMI_CODE_PROVIDER, blob) - return {"connected": True, "provider": KIMI_CODE_PROVIDER} - # 轮询次数耗尽（视作过期）。 - raise AppError( - ErrorCode.LLM_UNAVAILABLE, - "Kimi 设备授权轮询超时", - {"provider": KIMI_CODE_PROVIDER}, - ) - finally: - # `httpx.AsyncClient` 有 `aclose`（最小 Protocol 无；测试 fake 也实现了它）。 - aclose = getattr(http, "aclose", None) - if aclose is not None: - await aclose() + blob = encrypt_oauth_bundle(token, key=enc_key) + await store.upsert_oauth_credential(STUB_OWNER_ID, KIMI_CODE_PROVIDER, blob) + return {"connected": True, "provider": KIMI_CODE_PROVIDER} return work +def _raise_work(exc: Exception) -> Any: + """构造只把 `exc` 重抛的 work（轮询失败路径：交由 `run_job` 分类脱敏并置 job failed）。""" + + async def work(_session: AsyncSession) -> dict[str, Any]: + raise exc + + return work + + +async def _run_kimi_oauth_poll( + session_factory: SessionFactory, + job_id: uuid.UUID, + device: DeviceAuth, + *, + request_id: str | None = None, +) -> None: + """后台任务：**先在会话外**轮询 token，再经 `run_job` 开短会话持久化（CR-H1）。 + + 轮询期间不占用任何 DB 会话（防连接池饥饿）；成功 → 交由 `run_job` 落库并置 done； + 失败 → 把异常带进 `run_job`，走其 `_classify_job_error`/`_mark_failed` 脱敏置 failed。 + """ + http = _default_http_client() + try: + token = await _poll_for_token(http, device) + work = _make_persist_work(token) + except Exception as exc: # noqa: BLE001 — 失败经 run_job 统一脱敏置 job failed，不在此处冒泡。 + work = _raise_work(exc) + finally: + # `httpx.AsyncClient` 有 `aclose`（最小 Protocol 无；测试 fake 也实现了它）。 + aclose = getattr(http, "aclose", None) + if aclose is not None: + await aclose() + await run_job(session_factory, job_id, work, request_id=request_id) + + @router.post("/start", status_code=202) async def start_oauth( request: Request, @@ -139,11 +177,12 @@ async def start_oauth( job = await job_repo.create(None, _JOB_KIND_KIMI_OAUTH) await session.commit() # job 行需在 202 前持久化（供前端立即轮询）。 + # 轮询移出 DB 会话（CR-H1）：后台先在会话外轮询 token，仅最终持久化开一个短会话。 background_tasks.add_task( - run_job, + _run_kimi_oauth_poll, session_factory, job.id, - _make_poll_work(device), + device, request_id=request_id, )