writer-work-flow/apps/api/ww_api/routers/kimi_oauth.py

"""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
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,
    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


def _default_http_client() -> AsyncHttpClient:
    return httpx.AsyncClient(timeout=30.0)


def _make_poll_work(device: DeviceAuth) -> Any:
    """构造后台轮询工作闭包：循环 poll token 直到成功/过期/拒绝。

    `work(session)` 自建 httpx 客户端 + 凭据 store（用 `run_job` 传入的独立 session）→
    按 `interval` 轮询 token → 成功则加密存 `oauth_enc` 并返回非密 job 结果 → 过期/拒绝则
    抛 AppError（`run_job` 置 job failed）。**token 绝不进 job 结果/日志**。
    """

    async def work(session: AsyncSession) -> dict[str, Any]:
        enc_key = get_settings().credential_enc_key
        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()

    return work


@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(_default_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 前持久化（供前端立即轮询）。

    background_tasks.add_task(
        run_job,
        session_factory,
        job.id,
        _make_poll_work(device),
        request_id=request_id,
    )

    log.info(
        "kimi_oauth_started",
        request_id=request_id,
        job_id=str(job.id),
        user_code=device.user_code,
    )
    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,
) -> OAuthDisconnectResponse:
    """断开 Kimi Code：删除 OAuth 凭据行（token 一并消失）。"""
    request_id = getattr(request.state, "request_id", None)
    deleted = await store.delete_credential(STUB_OWNER_ID, KIMI_CODE_PROVIDER)
    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
    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())