fix(backend): Kimi OAuth 轮询移出 DB 会话——只在最终持久化开短会话(CR-H1)
This commit is contained in:
@@ -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()
|
||||||
|
|||||||
@@ -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,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user