249 lines
9.6 KiB
Python
249 lines
9.6 KiB
Python
"""Kimi Code OAuth device-flow 端点（C3 扩 K1.3 / ARCH §7.2 / §7.4 jobs）。
|
||
|
||
订阅 plan device-flow 登录：作者在设置页点「连接 Kimi Code」→ 后端发起 device
|
||
authorization → 返回 202 `{job_id, user_code, verification_uri, ...}` → 前端展示 user_code
|
||
+ 打开授权页 + 轮询 `GET /jobs/{id}`；后端经 BackgroundTask `run_job` **后台轮询** token
|
||
端点直到授权成功（加密存 `oauth_enc` + job done）或过期/拒绝（job failed）。
|
||
|
||
三端点（挂 `kimi_oauth.router`，已注册）：
|
||
- `POST /settings/providers/kimi-code/oauth/start` → 202 `OAuthStartResponse`
|
||
- `POST /settings/providers/kimi-code/oauth/disconnect` → 200 `OAuthDisconnectResponse`
|
||
- `GET /settings/providers/kimi-code/oauth/status` → 200 `OAuthStatusResponse`
|
||
|
||
**token 绝不出边界**：响应/job 结果/日志只含 user_code/connected 等非密信息；access/refresh
|
||
token 仅以 Fernet 密文存 `provider_credentials.oauth_enc`。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import uuid
|
||
from collections.abc import AsyncIterator
|
||
from typing import Annotated, Any
|
||
|
||
import httpx
|
||
from fastapi import APIRouter, BackgroundTasks, Depends, Request, Response
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from ww_config import get_settings
|
||
from ww_core.domain import JobRepo
|
||
from ww_db import get_session
|
||
from ww_llm_gateway.adapters.kimi_code import KIMI_CODE_PROVIDER
|
||
from ww_shared import AppError, ErrorCode
|
||
|
||
from ww_api.logging_config import get_logger
|
||
from ww_api.schemas.kimi_oauth import (
|
||
OAuthDisconnectResponse,
|
||
OAuthStartResponse,
|
||
OAuthStatusResponse,
|
||
)
|
||
from ww_api.services.credentials import (
|
||
AUTH_TYPE_OAUTH,
|
||
STUB_OWNER_ID,
|
||
CredentialStore,
|
||
SqlCredentialStore,
|
||
)
|
||
from ww_api.services.foreshadow_scan import SessionFactory
|
||
from ww_api.services.job_runner import run_job
|
||
from ww_api.services.kimi_oauth import (
|
||
AsyncHttpClient,
|
||
AuthorizationPending,
|
||
DeviceAuth,
|
||
SlowDown,
|
||
TokenSet,
|
||
decrypt_oauth_bundle,
|
||
encrypt_oauth_bundle,
|
||
poll_token,
|
||
start_device_authorization,
|
||
)
|
||
from ww_api.services.project_deps import get_job_repo, get_session_factory
|
||
from ww_api.services.provider_deps import get_credential_store
|
||
|
||
log = get_logger("ww.api.kimi_oauth")
|
||
|
||
router = APIRouter(prefix="/settings/providers/kimi-code/oauth", tags=["kimi-oauth"])
|
||
|
||
CredentialStoreDep = Annotated[CredentialStore, Depends(get_credential_store)]
|
||
JobRepoDep = Annotated[JobRepo, Depends(get_job_repo)]
|
||
SessionDep = Annotated[AsyncSession, Depends(get_session)]
|
||
SessionFactoryDep = Annotated[SessionFactory, Depends(get_session_factory)]
|
||
|
||
_JOB_KIND_KIMI_OAUTH = "kimi_oauth"
|
||
|
||
#: device flow 轮询安全上限（防止 work 在异常 interval 下无限循环；按 expires_in 兜底）。
|
||
_MAX_POLL_ATTEMPTS = 200
|
||
|
||
|
||
#: 请求/后台 HTTP 客户端超时（秒）。
|
||
_HTTP_TIMEOUT_SECONDS = 30.0
|
||
|
||
|
||
def _default_http_client() -> AsyncHttpClient:
|
||
"""裸工厂：现建一个 httpx 客户端（**后台轮询**直接调用，自行在 finally 里 aclose）。"""
|
||
return httpx.AsyncClient(timeout=_HTTP_TIMEOUT_SECONDS)
|
||
|
||
|
||
async def _http_client() -> AsyncIterator[AsyncHttpClient]:
|
||
"""请求级 httpx 客户端**生成器依赖**：请求结束时经 `async with` 自动 aclose（CR-H2）。
|
||
|
||
普通函数依赖无 teardown，返回的 client 永不关闭 → 每次 `/start` 泄漏一个连接。故请求路径
|
||
用此生成器依赖（后台路径仍用裸工厂 `_default_http_client`，自管生命周期）。
|
||
"""
|
||
async with httpx.AsyncClient(timeout=_HTTP_TIMEOUT_SECONDS) as http:
|
||
yield http
|
||
|
||
|
||
# SlowDown 退避每次增大的间隔量（秒）。
|
||
_SLOW_DOWN_STEP_SECONDS = 5
|
||
|
||
|
||
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)
|
||
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,
|
||
response: Response,
|
||
background_tasks: BackgroundTasks,
|
||
job_repo: JobRepoDep,
|
||
session: SessionDep,
|
||
session_factory: SessionFactoryDep,
|
||
http: Annotated[AsyncHttpClient, Depends(_http_client)],
|
||
) -> OAuthStartResponse:
|
||
"""发起 device authorization：返回 202 + user_code/verification_uri，后台轮询 token。
|
||
|
||
创建一行 `jobs(kind="kimi_oauth")`（202 前持久化供轮询）→ 调度 BackgroundTask 后台
|
||
轮询。前端展示 user_code + 打开 verification_uri + 轮询 `GET /jobs/{id}`。
|
||
"""
|
||
request_id = getattr(request.state, "request_id", None)
|
||
|
||
device = await start_device_authorization(http)
|
||
|
||
job = await job_repo.create(None, _JOB_KIND_KIMI_OAUTH)
|
||
await session.commit() # job 行需在 202 前持久化（供前端立即轮询）。
|
||
|
||
# 轮询移出 DB 会话（CR-H1）：后台先在会话外轮询 token，仅最终持久化开一个短会话。
|
||
background_tasks.add_task(
|
||
_run_kimi_oauth_poll,
|
||
session_factory,
|
||
job.id,
|
||
device,
|
||
request_id=request_id,
|
||
)
|
||
|
||
# 不记 user_code：授权窗口内日志可见者可冒用授权意图（P1-10）。仅留 request_id/job_id 关联。
|
||
log.info(
|
||
"kimi_oauth_started",
|
||
request_id=request_id,
|
||
job_id=str(job.id),
|
||
)
|
||
response.status_code = 202
|
||
return OAuthStartResponse(
|
||
job_id=job.id,
|
||
user_code=device.user_code,
|
||
verification_uri=device.verification_uri,
|
||
verification_uri_complete=device.verification_uri_complete,
|
||
expires_in=device.expires_in,
|
||
interval=device.interval,
|
||
)
|
||
|
||
|
||
@router.post("/disconnect")
|
||
async def disconnect_oauth(
|
||
request: Request,
|
||
store: CredentialStoreDep,
|
||
session: SessionDep,
|
||
) -> OAuthDisconnectResponse:
|
||
"""断开 Kimi Code：删除 OAuth 凭据行（token 一并消失）。"""
|
||
request_id = getattr(request.state, "request_id", None)
|
||
deleted = await store.delete_credential(STUB_OWNER_ID, KIMI_CODE_PROVIDER)
|
||
# store.delete_credential 只 flush，端点统一提交。
|
||
await session.commit()
|
||
log.info("kimi_oauth_disconnected", request_id=request_id, deleted=deleted)
|
||
return OAuthDisconnectResponse(disconnected=deleted)
|
||
|
||
|
||
@router.get("/status")
|
||
async def oauth_status(store: CredentialStoreDep) -> OAuthStatusResponse:
|
||
"""连接状态：是否已连接 + access token 过期时刻（**无 token 本体**）。"""
|
||
enc_key = get_settings().credential_enc_key.get_secret_value()
|
||
cred = await store.get_credential(STUB_OWNER_ID, KIMI_CODE_PROVIDER)
|
||
if cred is None or cred.auth_type != AUTH_TYPE_OAUTH or cred.oauth_enc is None:
|
||
return OAuthStatusResponse(connected=False)
|
||
try:
|
||
token = decrypt_oauth_bundle(cred.oauth_enc, key=enc_key)
|
||
except Exception: # noqa: BLE001 — 解密失败视作未连接（不泄露原因到响应）
|
||
return OAuthStatusResponse(connected=False)
|
||
return OAuthStatusResponse(connected=True, expires_at=token.expires_at.isoformat())
|