fix(backend): Kimi OAuth 轮询移出 DB 会话——只在最终持久化开短会话(CR-H1)

This commit is contained in:
Yaojia Wang
2026-07-08 10:37:42 +02:00
parent ed0a679322
commit ca8fbc89a1
2 changed files with 176 additions and 37 deletions

View File

@@ -15,6 +15,8 @@
from __future__ import annotations from __future__ import annotations
import os import os
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Any from typing import Any
import pytest import pytest
@@ -51,12 +53,65 @@ class _ScriptedHttp:
return None 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( def _make_client(
*, *,
store: FakeCredentialStore, store: FakeCredentialStore,
http: _ScriptedHttp, http: Any,
enc_key: str, enc_key: str,
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
session_factory: Any = None,
) -> TestClient: ) -> TestClient:
os.environ["CREDENTIAL_ENC_KEY"] = enc_key os.environ["CREDENTIAL_ENC_KEY"] = enc_key
from ww_config import get_settings from ww_config import get_settings
@@ -88,11 +143,13 @@ def _make_client(
job_repo = FakeJobRepo() job_repo = FakeJobRepo()
monkeypatch.setattr(job_runner, "SqlJobRepo", lambda _session: job_repo) monkeypatch.setattr(job_runner, "SqlJobRepo", lambda _session: job_repo)
sf = session_factory if session_factory is not None else FakeSessionFactory()
app = create_app() app = create_app()
app.dependency_overrides[get_credential_store] = lambda: store app.dependency_overrides[get_credential_store] = lambda: store
app.dependency_overrides[get_job_repo] = lambda: job_repo app.dependency_overrides[get_job_repo] = lambda: job_repo
app.dependency_overrides[get_session] = lambda: FakeSession() 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 app.dependency_overrides[routers.kimi_oauth._default_http_client] = lambda: http
return TestClient(app) 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 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: def test_disconnect_clears_credential(monkeypatch: pytest.MonkeyPatch) -> None:
enc_key = Fernet.generate_key().decode() enc_key = Fernet.generate_key().decode()
store = FakeCredentialStore() store = FakeCredentialStore()

View File

@@ -17,6 +17,7 @@ token 仅以 Fernet 密文存 `provider_credentials.oauth_enc`。
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import uuid
from typing import Annotated, Any from typing import Annotated, Any
import httpx import httpx
@@ -47,6 +48,7 @@ from ww_api.services.kimi_oauth import (
AuthorizationPending, AuthorizationPending,
DeviceAuth, DeviceAuth,
SlowDown, SlowDown,
TokenSet,
decrypt_oauth_bundle, decrypt_oauth_bundle,
encrypt_oauth_bundle, encrypt_oauth_bundle,
poll_token, poll_token,
@@ -74,47 +76,83 @@ def _default_http_client() -> AsyncHttpClient:
return httpx.AsyncClient(timeout=30.0) return httpx.AsyncClient(timeout=30.0)
def _make_poll_work(device: DeviceAuth) -> Any: # SlowDown 退避每次增大的间隔量(秒)。
"""构造后台轮询工作闭包:循环 poll token 直到成功/过期/拒绝。 _SLOW_DOWN_STEP_SECONDS = 5
`work(session)` 自建 httpx 客户端 + 凭据 store用 `run_job` 传入的独立 session
按 `interval` 轮询 token → 成功则加密存 `oauth_enc` 并返回非密 job 结果 → 过期/拒绝则 async def _poll_for_token(http: AsyncHttpClient, device: DeviceAuth) -> TokenSet:
抛 AppError`run_job` 置 job failed。**token 绝不进 job 结果/日志** """设备授权轮询循环(**纯 HTTP无 DB 会话**CR-H1直到拿到 token 或超时
`authorization_pending` → 继续;`slow_down` → 增大 interval成功 → 返回 `TokenSet`
次数耗尽 → 抛 `AppError`(视作过期,由调用方经 `run_job` 置 job failed
""" """
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) interval = max(1, device.interval)
http = _default_http_client()
try:
for _ in range(_MAX_POLL_ATTEMPTS): for _ in range(_MAX_POLL_ATTEMPTS):
await asyncio.sleep(interval) await asyncio.sleep(interval)
try: try:
token = await poll_token(http, device.device_code) return await poll_token(http, device.device_code)
except AuthorizationPending: except AuthorizationPending:
continue continue
except SlowDown: except SlowDown:
interval += 5 interval += _SLOW_DOWN_STEP_SECONDS
continue 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( raise AppError(
ErrorCode.LLM_UNAVAILABLE, ErrorCode.LLM_UNAVAILABLE,
"Kimi 设备授权轮询超时", "Kimi 设备授权轮询超时",
{"provider": KIMI_CODE_PROVIDER}, {"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)
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: finally:
# `httpx.AsyncClient` 有 `aclose`(最小 Protocol 无;测试 fake 也实现了它)。 # `httpx.AsyncClient` 有 `aclose`(最小 Protocol 无;测试 fake 也实现了它)。
aclose = getattr(http, "aclose", None) aclose = getattr(http, "aclose", None)
if aclose is not None: if aclose is not None:
await aclose() await aclose()
await run_job(session_factory, job_id, work, request_id=request_id)
return work
@router.post("/start", status_code=202) @router.post("/start", status_code=202)
@@ -139,11 +177,12 @@ async def start_oauth(
job = await job_repo.create(None, _JOB_KIND_KIMI_OAUTH) job = await job_repo.create(None, _JOB_KIND_KIMI_OAUTH)
await session.commit() # job 行需在 202 前持久化(供前端立即轮询)。 await session.commit() # job 行需在 202 前持久化(供前端立即轮询)。
# 轮询移出 DB 会话CR-H1后台先在会话外轮询 token仅最终持久化开一个短会话。
background_tasks.add_task( background_tasks.add_task(
run_job, _run_kimi_oauth_poll,
session_factory, session_factory,
job.id, job.id,
_make_poll_work(device), device,
request_id=request_id, request_id=request_id,
) )