P0-1 SqlCredentialStore/save_draft 由自提交改 flush,端点/服务统一 commit (新增 CredentialStore.commit() 统一提交点;token 刷新落库显式提交); 补多凭据一请求中途失败整体回滚集成测试。 P0-2 启动校验 _fernet(enc_key) 快速失败 + catch-all Exception → ErrorEnvelope; credential_enc_key 改 SecretStr。 P0-3 run_job 异常分类:AppError 存 code+message,其余存通用文案不泄 str(exc)。 P0-4 评审/正文 SSE 失败先发 error 事件,尾部 commit 包 try/except。 P1-4 max_version 加 FOR UPDATE 行锁消除 TOCTOU。 P1-5 scan_overdue 谓词下推 + 批量 UPDATE RETURNING。 P1-10 移除 OAuth user_code 日志。 P2 provider_deps 改调网关 build_adapter;accept_service Committable Protocol; CORS 白名单收窄;request_id 安全字符集白名单;stdlib 日志接管;读端点 404 校验; httpx timeout;测试用合法 Fernet key;类型化响应模型(JobResponse/DimensionEntry/ ReviewConflictView/selling_points)+路由 ErrorEnvelope responses(供 codegen)。
252 lines
9.2 KiB
Python
252 lines
9.2 KiB
Python
"""凭据存储与提供商探测的接口 + SQLAlchemy/网关实现(ARCH §4.7)。
|
||
|
||
路由依赖这里的 **接口**(Protocol),测试注入内存替身;运行时用 SQLAlchemy/网关实现。
|
||
单用户原型:`owner_id` 用固定 stub(见 `STUB_OWNER_ID`),多租户化时改为按认证主体取。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import uuid
|
||
from dataclasses import dataclass
|
||
from typing import Protocol
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from ww_db.models import ProviderCredential, TierRouting
|
||
from ww_llm_gateway.adapters.base import Capabilities
|
||
|
||
# 单用户 stub owner(与网关 Scope.user_id 的 stub 约定一致:UUID(int=1))。
|
||
# 多租户化时此常量由认证主体替换(见 ARCH §4.7 隔离)。
|
||
STUB_OWNER_ID = uuid.UUID(int=1)
|
||
|
||
|
||
# 凭据认证类型(`provider_credentials.auth_type`,见 C2 扩 K1.1)。
|
||
AUTH_TYPE_API_KEY = "api_key"
|
||
AUTH_TYPE_OAUTH = "oauth"
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class StoredCredential:
|
||
"""存储层视图:含密文,绝不出 API 边界(路由仅取 provider 并掩码)。
|
||
|
||
一行二选一:`auth_type="api_key"` → `api_key_enc` 有值、`oauth_enc=None`;
|
||
`auth_type="oauth"`(Kimi Code device-flow,K1.3)→ `oauth_enc` 有值、`api_key_enc=None`
|
||
(持 Fernet 加密的 `{access_token,refresh_token,expires_at}` JSON 包)。
|
||
"""
|
||
|
||
provider: str
|
||
api_key_enc: bytes | None
|
||
auth_type: str = AUTH_TYPE_API_KEY
|
||
oauth_enc: bytes | None = None
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class StoredRouting:
|
||
tier: str
|
||
provider: str
|
||
model: str
|
||
fallback: list[str]
|
||
|
||
|
||
class CredentialStore(Protocol):
|
||
"""凭据 + 档位路由的读写接口(按 owner_id 隔离)。
|
||
|
||
写方法(upsert/delete)**只 flush 不 commit**——提交交调用方(端点/服务)经
|
||
`commit()` 统一一次,保证「多凭据一请求」的原子性(任一步失败整体回滚,不留半更新)。
|
||
无 session 句柄的服务侧调用方(如 token 刷新落库)则直接调 `commit()`。
|
||
"""
|
||
|
||
async def list_credentials(self, owner_id: uuid.UUID) -> list[StoredCredential]: ...
|
||
|
||
async def list_routing(self) -> list[StoredRouting]: ...
|
||
|
||
async def get_credential(
|
||
self, owner_id: uuid.UUID, provider: str
|
||
) -> StoredCredential | None: ...
|
||
|
||
async def upsert_credential(
|
||
self, owner_id: uuid.UUID, provider: str, api_key_enc: bytes
|
||
) -> None: ...
|
||
|
||
async def upsert_oauth_credential(
|
||
self, owner_id: uuid.UUID, provider: str, oauth_enc: bytes
|
||
) -> None: ...
|
||
|
||
async def delete_credential(self, owner_id: uuid.UUID, provider: str) -> bool: ...
|
||
|
||
async def upsert_routing(self, routing: StoredRouting) -> None: ...
|
||
|
||
async def commit(self) -> None: ...
|
||
|
||
|
||
class ProviderProbe(Protocol):
|
||
"""最小连通探测:验证 Key + 返回能力矩阵。测试注入假探测,绝不联网。"""
|
||
|
||
async def probe(self, owner_id: uuid.UUID, provider: str) -> Capabilities: ...
|
||
|
||
|
||
class SqlCredentialStore:
|
||
"""SQLAlchemy 实现:写 `provider_credentials` / `tier_routing`,幂等 upsert。"""
|
||
|
||
def __init__(self, session: AsyncSession) -> None:
|
||
self._session = session
|
||
|
||
async def list_credentials(self, owner_id: uuid.UUID) -> list[StoredCredential]:
|
||
rows = (
|
||
await self._session.execute(
|
||
select(ProviderCredential).where(ProviderCredential.owner_id == owner_id)
|
||
)
|
||
).scalars()
|
||
return [
|
||
StoredCredential(
|
||
provider=r.provider,
|
||
api_key_enc=r.api_key_enc,
|
||
auth_type=r.auth_type,
|
||
oauth_enc=r.oauth_enc,
|
||
)
|
||
for r in rows
|
||
]
|
||
|
||
async def list_routing(self) -> list[StoredRouting]:
|
||
rows = (await self._session.execute(select(TierRouting))).scalars()
|
||
return [
|
||
StoredRouting(
|
||
tier=r.tier, provider=r.provider, model=r.model, fallback=list(r.fallback)
|
||
)
|
||
for r in rows
|
||
]
|
||
|
||
async def get_credential(self, owner_id: uuid.UUID, provider: str) -> StoredCredential | None:
|
||
row = (
|
||
await self._session.execute(
|
||
select(ProviderCredential).where(
|
||
ProviderCredential.owner_id == owner_id,
|
||
ProviderCredential.provider == provider,
|
||
)
|
||
)
|
||
).scalar_one_or_none()
|
||
if row is None:
|
||
return None
|
||
return StoredCredential(
|
||
provider=row.provider,
|
||
api_key_enc=row.api_key_enc,
|
||
auth_type=row.auth_type,
|
||
oauth_enc=row.oauth_enc,
|
||
)
|
||
|
||
async def upsert_credential(
|
||
self, owner_id: uuid.UUID, provider: str, api_key_enc: bytes
|
||
) -> None:
|
||
# 显式 read-modify-write:唯一约束含可空 project_id,PG ON CONFLICT
|
||
# 在 NULL 上不去重(NULLS DISTINCT),故不用 on_conflict。
|
||
existing = (
|
||
await self._session.execute(
|
||
select(ProviderCredential).where(
|
||
ProviderCredential.owner_id == owner_id,
|
||
ProviderCredential.project_id.is_(None),
|
||
ProviderCredential.provider == provider,
|
||
)
|
||
)
|
||
).scalar_one_or_none()
|
||
if existing is None:
|
||
self._session.add(
|
||
ProviderCredential(
|
||
owner_id=owner_id,
|
||
project_id=None,
|
||
provider=provider,
|
||
api_key_enc=api_key_enc,
|
||
auth_type=AUTH_TYPE_API_KEY,
|
||
oauth_enc=None,
|
||
)
|
||
)
|
||
else:
|
||
existing.api_key_enc = api_key_enc
|
||
existing.auth_type = AUTH_TYPE_API_KEY
|
||
existing.oauth_enc = None
|
||
# 仓储只 flush,提交交调用方(端点/服务)统一一次——保证多凭据一请求的原子性。
|
||
await self._session.flush()
|
||
|
||
async def upsert_oauth_credential(
|
||
self, owner_id: uuid.UUID, provider: str, oauth_enc: bytes
|
||
) -> None:
|
||
"""写/更新 OAuth 凭据行(Kimi Code device-flow,K1.3)。
|
||
|
||
`auth_type="oauth"`、`oauth_enc=<Fernet 加密 token 包>`、`api_key_enc=None`。
|
||
显式 read-modify-write(同 `upsert_credential`:含可空 project_id 的唯一约束不能用
|
||
PG `ON CONFLICT`,见 memory/gotchas)。明文 token 绝不进此层(已加密)。
|
||
"""
|
||
existing = (
|
||
await self._session.execute(
|
||
select(ProviderCredential).where(
|
||
ProviderCredential.owner_id == owner_id,
|
||
ProviderCredential.project_id.is_(None),
|
||
ProviderCredential.provider == provider,
|
||
)
|
||
)
|
||
).scalar_one_or_none()
|
||
if existing is None:
|
||
self._session.add(
|
||
ProviderCredential(
|
||
owner_id=owner_id,
|
||
project_id=None,
|
||
provider=provider,
|
||
api_key_enc=None,
|
||
auth_type=AUTH_TYPE_OAUTH,
|
||
oauth_enc=oauth_enc,
|
||
)
|
||
)
|
||
else:
|
||
existing.api_key_enc = None
|
||
existing.auth_type = AUTH_TYPE_OAUTH
|
||
existing.oauth_enc = oauth_enc
|
||
# 仓储只 flush,提交交调用方统一一次。
|
||
await self._session.flush()
|
||
|
||
async def delete_credential(self, owner_id: uuid.UUID, provider: str) -> bool:
|
||
"""删除凭据行(OAuth disconnect / 撤销)。返回是否删到行。"""
|
||
existing = (
|
||
await self._session.execute(
|
||
select(ProviderCredential).where(
|
||
ProviderCredential.owner_id == owner_id,
|
||
ProviderCredential.project_id.is_(None),
|
||
ProviderCredential.provider == provider,
|
||
)
|
||
)
|
||
).scalar_one_or_none()
|
||
if existing is None:
|
||
return False
|
||
await self._session.delete(existing)
|
||
# 仓储只 flush,提交交调用方统一一次。
|
||
await self._session.flush()
|
||
return True
|
||
|
||
async def upsert_routing(self, routing: StoredRouting) -> None:
|
||
existing = (
|
||
await self._session.execute(
|
||
select(TierRouting).where(
|
||
TierRouting.project_id.is_(None),
|
||
TierRouting.tier == routing.tier,
|
||
)
|
||
)
|
||
).scalar_one_or_none()
|
||
if existing is None:
|
||
self._session.add(
|
||
TierRouting(
|
||
project_id=None,
|
||
tier=routing.tier,
|
||
provider=routing.provider,
|
||
model=routing.model,
|
||
fallback=routing.fallback,
|
||
)
|
||
)
|
||
else:
|
||
existing.provider = routing.provider
|
||
existing.model = routing.model
|
||
existing.fallback = routing.fallback
|
||
# 仓储只 flush,提交交调用方统一一次。
|
||
await self._session.flush()
|
||
|
||
async def commit(self) -> None:
|
||
"""统一提交点:端点/服务侧在一组 flush 后调一次,落库所有挂起写入。"""
|
||
await self._session.commit()
|