diff --git a/.gitignore b/.gitignore index c209091..bba3e52 100644 --- a/.gitignore +++ b/.gitignore @@ -49,6 +49,8 @@ build/ *.md !README.md !README_zh.md +!ACCOUNT_MANAGER_PLAN.md +!ACCOUNT_MANAGER_ACCEPTANCE.md # VS Code 扩展文件 *.vsix diff --git a/ACCOUNT_MANAGER_ACCEPTANCE.md b/ACCOUNT_MANAGER_ACCEPTANCE.md new file mode 100644 index 0000000..8aff7a6 --- /dev/null +++ b/ACCOUNT_MANAGER_ACCEPTANCE.md @@ -0,0 +1,130 @@ +# 账号管理改造验收记录 + +日期:2026-10-24。基线:`origin/main` / `83f60e90eccf5495356b97cb2def79533a38d9a5`。 +分支:`codex/account-session-manager`。方案见 `ACCOUNT_MANAGER_PLAN.md`。 + +## 实现与委派 + +三个子 agent 均按用户指定使用 GPT-6 Luna / max,并在共享分支中按文件边界并行实现,没有代替用户提交或推送。 + +| 负责人 | 交付 | 交回验证 | +| --- | --- | --- | +| A | 独立账号库、两步登录、原子持久化、刷新与重新关联的并发保护、失败回滚 | 最终定向组 31 passed | +| B | 后台登录任务、取消和关闭、容量限制、仅本机管理 API、响应字段过滤 | 12 passed | +| C | 管理页面、双语文档、包资源、页面控制器模拟 | 页面测试 1 passed;wheel 构建及页面内容比对通过 | +| 主 agent | 启动迁移和选择规则、请求入口账号快照、三协议和 billing 集成、独立审阅 | 集成测试 15 passed;全量与打包结果见下 | + +## 已确认行为 + +- `/admin` 与代理共用进程和端口;未关联账号也可启动,管理/模型查询可用,推理返回结构化 401。 +- 国内与 Global 分区,身份包含区域、endpoint 和 uid;不同身份保存到独立文件。索引只包含展示字段,凭据留在每账号 session。 +- 页面支持选择区域、打开独立认证窗口、查询登录进度、取消与重试、关联、切换及删除非当前账号。首次关联自动激活,后续关联保留当前选择。 +- 继续复用原有 state → token → account 协议和认证头。只有同时得到 token 和含 uid 的账户才写入;超时、取消不提交正式账号。 +- 管理请求校验实际 loopback 对端、Host、Origin 和端口;写操作要求专用请求头;管理响应 no-store,账号名以文本渲染。 +- 默认恢复上次 active。显式区域/endpoint/session-file 选择匹配账号;无匹配可空账号启动。默认旧会话首次迁移且源文件保留,删除后不自动重新导入;显式 session-file 可重新关联已删除身份。 +- CLI `--login` 在临时 session 中认证后导入,不会用另一个账号覆盖旧会话文件。 +- 账号在请求进入时、读取请求体之前固定。切换只影响后续请求;旧流、请求重试和 billing 401 刷新不会跨账号或跨 endpoint。 +- 单账号并发认证只刷新一次;旧刷新不能覆盖重新关联的凭据。删除后在途刷新不能重建已删除凭据文件。索引写入失败会回滚磁盘和缓存,回滚期间禁止其他 session writer 进入。 + +## 验证与证据 + +自动化使用模拟上游和临时账号目录,不调用真实 CodeBuddy 认证或推理。 + +```powershell +$env:PYTHONPATH='src' +$env:PYTHONUTF8='1' +.\.venv\Scripts\python.exe -m pytest -q --tb=short +uv build --wheel +git -c core.whitespace=cr-at-eol diff --check +``` + +集成阶段一次全量为 **324 passed / 1 skipped**。追加并发、回滚、页面和请求入口回归后的全量为 **332 passed / 2 failed / 1 skipped**;两项失败均来自未改动的 Windows 进程监控测试:CIM 查询超过 10 秒,以及退出时间与最后采样时间的顺序断言。 + +原始主干的独立 checkout 已复现 CIM 超时;退出时间测试在该 checkout 的单独复验通过。账号改造没有修改 `scripts/watch_processes.py` 或 `test_external_process_monitor.py`。此外修正了原有跨平台权限测试:只在 POSIX 断言 0600,Windows 验证可读写,不把 chmod 当作 Windows ACL 验收。 + +最终无并行构建复验为 **333 passed / 1 failed / 1 skipped**,仅剩同一项已在原始主干复现的 CIM 查询超时(`test_watch_discovers_real_temporary_module`);退出时间断言本轮通过。没有通过删除或跳过测试来掩盖这项失败,也没有扩展修改进程监控模块。 + +新增的六个账号相关测试文件单独复验全部通过,共 **44 passed**:`test_account_store.py`、`test_client_login_management.py`、`test_account_manager.py`、`test_admin_api.py`、`test_admin_ui.py`、`test_account_integration.py`。 + +`uv build --wheel` 成功;wheel 包含 account_store/account_manager/admin_api 和 admin.html,页面内容与源码逐字节一致。保留既有 2.0.7 版本号,未发布安装包。 + +Edge 浏览器实际检查通过:页面加载、切换 Global、选择 Global 和国内区域、弹出独立窗口、模拟关联成功、账号列表刷新且当前选择保持。修复并重新检查成功态隐藏重试/取消按钮。浏览器验收服务器使用独立临时目录和合成账号,模拟登录工厂只访问本机,不修改用户真实账号。验收后停止该临时服务。 + +截图:`F:/Codex/visualizations/2026/10/24/01a16c22-2245-7651-89a9-06eed118ee27/account-manager-preview.png`。 + +## 429 限流与自动切换补充验收 + +新增默认关闭的自动切换开关、每账号冷却状态与管理页倒计时,独立运行状态文件保存开关和冷却截止时间。Retry-After 支持秒数、HTTP 日期、零秒;缺失或无效时使用默认 60 秒。重启恢复状态;冷却过期只解除标记,不主动切回。 + +同一次推理请求可在收到 HTTP 429 且响应尚未开始时,依次尝试同区域且同 endpoint 的其他未冷却账号,每账号最多一次,数量受请求开始时账号数限制。全部限流仍返回最后一次实际 429 与它的原始 Retry-After。401/503 不触发轮换;已开始输出的 SSE 不重放;晚到的旧账号 429 不覆盖手动切换。Billing 仅记录与选择后续账号,不重放签到操作。 + +主 agent 负责主转发和 billing 集成;GPT-6 Luna / max 子 agent B 实现持久化和管理 API,C 实现页面与文档,A 独立新增代理链路验收。新增代理组 test_rate_limit_proxy.py 为 **34 passed**;manager/API 为 **21 passed**;七个账号相关测试文件合跑 **87 passed**。这些用例全部使用临时账号与模拟上游,不消耗真实账号额度。 + +Edge 浏览器已检查默认关闭、开关保存,并通过隔离代理的实际本机 HTTP 请求触发:国内演示账号返回模拟 HTTP 429 后,同一区域备用账号自动接手并返回 HTTP 200;页面显示原账号冷却倒计时和备用账号当前状态。修复并复验了前台刷新与后台轮询竞争导致按钮持续禁用的问题;页面倒计时按本机代理的服务器时间校准,使用本地每秒更新,后台可见页 15 秒同步、隐藏页 60 秒同步。设置保存结果不确定时立即 GET 校准服务器状态。最终主 agent 独立全量复验为 **377 passed / 1 skipped**(7 条既有弃用或测试返回值警告);本轮进程监控用例均通过。主 agent 重建 `dist/workbuddy2api-2.0.7-py3-none-any.whl`,逐字节核对 `__main__.py`、`billing_proxy.py`、`account_manager.py`、`admin_api.py`、`admin.html` 全部与当前源码一致;差异检查通过。截图位于 `F:/Codex/visualizations/2026/10/24/01a16c22-2245-7651-89a9-06eed118ee27/account-rate-limit-preview.png`。临时演示服务已停止,用户实际运行的代理未被替换。 + +## 每账号签到与后台自动调度补充验收 + +自动签到采用一个默认关闭、持久化的全局开关,覆盖全部已关联账号:国内与 Global、当前与非当前账号都独立使用自己的凭据,不改变代理的 active 账号。签到缓存只保存在内存;重启后重新向上游查询当天状态。 + +主 agent 继续将实现分给 GPT-6 Luna / max:A 实现签到服务、定时器与可控时钟测试;B 实现设置迁移、管理 API、签到响应过滤与故障回滚;C 实现每账号查询/签到按钮、开关、状态同步和双语说明。主 agent 独立集成进程生命周期、结构化与控制台日志,新增真实应用链路的模拟上游测试,并审阅最终实现。 + +- 每个账号可手动查询、手动签到;签到前必查,只有明确未签到才发送动作,动作后再次确认。状态未知、活动未开放或查询失败时不自动签到。 +- 按北京时间(UTC+8)每日 00:00 检查所有账号;未签账号各自在当日 00:00–00:10 内随机计划一次复查,仍未签才执行。随机时间以午夜为基准,查询耗时不额外延长窗口。 +- 从关闭切为开启、服务启动时已开启,均立即检查全部账号并按需补签。后台调度随代理运行,关闭管理页仍有效。 +- 关闭开关取消等待任务,发出签到请求前再次检查开关、日期和账号存在性;已发送的请求不重复。取消后清除页面 busy 状态。手动操作不受自动开关控制。 +- 页面自动行为只查询与同步,不自行发起签到 POST,避免绕过后台的随机时间和开关检查。只有手动按钮调用手动签到 API。 +- 同账号串行、全局限制并发;明确拒绝后复查确认未签才允许后续重试;不确定结果保留防重复标记。签到遇到 429 只记录该账号冷却,不触发账号切换。 +- 查询开始/结果、签到开始/确认/失败、跳过、随机计划均记录安全字段;日志不含 token、会话、原始上游响应。并发设置回调重读当前开关,防止过期通知覆盖最终值。 +- 额外修复删除账号时冷却文件写入失败的一致性,并通过两个故障注入用例确认回滚行为。 + +签到状态适配参考仓库贡献者的实际接口记录:`data.active` 和严格布尔值 `data.today_checked_in`;已签到提示码 10001 仍需要状态接口确认。来源:[PR #7 的接口验证记录](https://github.com/hawklithm/workbuddy2api/pull/7#issuecomment-6210892410)。本次没有使用真实账号调用签到上游。 + +最终主 agent 独立全量:**413 passed / 1 skipped / 7 warnings**,耗时 118.61 秒;既有 Windows 进程监控用例本轮通过。最终页面稳定后追加 UI/API/应用签到集成组复验:**14 passed**。子 agent 最终签到相关组 **78 passed**,管理/API 组 **27 passed**,UI/API 组 **8 passed**。 + +主 agent 构建 wheel 成功,逐字节核对 `__main__.py`、`billing_proxy.py`、`account_store.py`、`account_manager.py`、`account_checkin.py`、`admin_api.py`、`admin.html` 全部与源码一致。差异检查通过。 + +Edge 浏览器使用独立临时目录、三个合成账号、模拟上游完成:初始状态查询;对非当前 Global 账号手动签到且当前国内账号不变;打开全局开关后后台逐账号检查,未签到账号补签、已签到账号跳过;刷新后三个账号均显示已签到,active 保持原值。实际检查 JSON 日志含 enable 检查、跳过、动作和确认事件。截图:`F:/Codex/visualizations/2026/10/24/01a16c22-2245-7651-89a9-06eed118ee27/account-checkin-preview.png`。临时验收服务在验收后停止,未替换用户实际代理。 + +## 签到状态旧接口误报修复 + +2026-10-24 用户实际账号反馈后确认:旧上游 `POST /v2/billing/meter/checkin-status` 虽然 HTTP 200、code 0,但返回 active=false、today_checked_in=false、空活动时间;使用相同国内账号和认证信息查询当前 `POST /v2/billing/meter/checkin-activity-status`,返回 active=true、today_checked_in=true、有效活动时间。已读取本机安装的 WorkBuddy 客户端代码确认当前状态路径。旧接口的成功响应不足以证明当前活动状态。 + +管理后台签到服务与 `/v1/checkin-status` 均更新为当前上游接口;本机 API 地址不变,仍兼容 GET/POST,向上游发送 POST 和空 JSON。每日签到动作地址保持不变,不回退到会返回无效活动信息的旧状态接口。 + +新增 6 项独立集成回归:国内/Global 各覆盖管理 API、v1 GET、v1 POST,模拟旧接口 inactive、新接口已签,验证展示 checked、原始响应透传、目标账号身份、active 不变,且已签到账号的手动签到请求也不会调用 daily-checkin。同步更新原有请求路径契约与自动任务模拟。定向签到组 **63 passed**。 + +主 agent 使用修复后的实际 AccountCheckinService,仅查询当前国内账号(未启动自动调度),结果为 status=checked、checked=true、error=null。没有发起真实签到动作,没有修改账号设置或重启运行中的用户代理。最终全量回归 **419 passed / 1 skipped / 7 warnings**,耗时 159.09 秒;警告为既有弃用和测试返回值提示。 + +## 账号与模型分别限流补充验收 + +2026-10-24 根据用户补充,推理 429 改为“账号 + 实际请求模型”维度;同账号的其他模型仍可用。每个模型分别保存截止时间、上次限流时间和 Retry-After/默认估算来源,管理页逐项显示模型名和倒计时,单个模型到期不影响其他模型标记。 + +主 agent 负责三协议请求模型传递、日志、billing/签到无模型隔离、方案及独立验收;GPT-6 Luna / max 子 agent B 负责 manager/API/版本化状态与故障回滚,C 负责页面和双语说明,A 负责独立代理回归及只读审阅。模型名称按请求原样保留,不推测别名。 + +自动切换候选仍限定同区域与同 endpoint,但只排除本次请求模型仍冷却的账号。其他模型冷却不排除该候选;每请求每账号最多一次、手动切换保护、Retry-After 原样透传和已经开始输出的流不重放均保留。日志包含 account_id 和 model。没有模型的 billing/签到 429 不记录推理冷却、不切换当前账号。 + +运行状态沿用原文件位置,schema 升级为 v2,以账号/模型嵌套存储;严格验证旧 v1 文件后保留设置,舍弃无法确定模型归属的整账号冷却,不把它扩展到所有模型。只读加载不改写旧文件。保存失败回滚内存,删除失败恢复完整嵌套状态。管理 API 删除旧顶层账号冷却字段,按 `model_cooldowns` 数组逐项白名单输出;拒绝无效或过长模型名及控制字符,页面使用文本渲染。 + +定向验证:manager/API **31 passed**;独立代理链路 **44 passed**;主 agent 将 manager/API/UI/签到/状态透传合跑 **95 passed**。覆盖三协议 × 流式/非流式、同账号模型 A 限流而 B 成功、候选同模型与其他模型冷却的两方向选择、两个不同模型同时收到 429 后分别保留 19/73 秒等待与正确日志归属、无模型请求隔离、状态升级/重启恢复和故障注入。均使用临时目录和模拟上游。 + +Edge 实际浏览器验证同一合成国内账号下 `model-alpha` 与 `model-beta` 显示各自倒计时和上游等待时间来源;备用国内和 Global 账号无错误的整账号标记,自动切换说明明确限定请求模型。仅在隔离的本机演示服务中验证,未使用真实账号额度或操作真实账号设置。截图:`F:/Codex/visualizations/2026/10/24/01a16c22-2245-7651-89a9-06eed118ee27/account-model-cooldown-preview.png`。临时服务验收后停止。 + +主 agent 最终独立全量回归:**433 passed / 1 skipped / 7 warnings**,耗时 126.53 秒;既有监控测试本轮通过。wheel 构建通过,逐字节核对主入口、manager、管理 API、页面、billing 和签到服务均与当前源码一致。差异检查通过。当前真实代理未被重启,修复及模型级状态需在加载新代码后生效。 + +## 管理页中英文切换补充验收 + +页面语言范围包括静态说明、按钮、标题、无障碍标签、动态账号列表/模型冷却/签到状态、设置和网络提示、登录进度及已知后端错误。中文和 English 选择保存在当前浏览器;有效保存值优先,随后跟随浏览器语言。无法访问本地存储时切换仍可用。 + +语言切换仅更新展示,不重新加载页面或额外发请求,不改变 active、区域选择、自动开关、进行中的登录 job 或签到/设置操作。进行中的 busy 状态保留,操作完成后使用当前语言显示结果。用户账号昵称、UID、模型名与 endpoint 保留原文并安全作为文本显示;北京时间调度及时间基准不随显示语言变化。 + +沿用指定的 GPT-6 Luna / max 子 agent:C 实现页面,A 独立扩充实际控制器的 DOM 模拟回归,B 清点后端错误并只读复核国际化和状态保留;主 agent 负责文档、独立测试、浏览器与安装包验收。未将界面语言设置发往上游或修改用户代理配置。 + +主 agent 独立验证:中英文词典各 **182** 项,键集合和模板参数一致,英文词典无误留中文文本。控制器与管理 API 合跑 **8 passed**;追加 HTTP 503 错误通知切换测试,确认保留状态码、不因语言切换重试请求。最终全量 **433 passed / 1 skipped / 7 warnings**,耗时 83.32 秒。修复取消登录后的重试按钮状态回归;未知原型名称错误不误作翻译键,查询成功通知中的状态参数按当前语言解析。 + +Edge 在独立临时目录和合成账号中验证:默认跟随浏览器英文;切换中文后页面标题、无障碍标签、模型冷却、签到及设置文字同步变化;刷新保留中文,再切英文并刷新保留英文;中文查询账号状态后切英文,已有通知中的状态也翻译且账号名保留原文。英文截图:`F:/Codex/visualizations/2026/10/24/01a16c22-2245-7651-89a9-06eed118ee27/account-manager-english-preview.png`。页面测试未执行真实账号操作。wheel 重建并与源码核对,临时演示服务验收后停止。 + +## 边界 + +本次完成本地代码和离线/浏览器验收。真实国内与 Global 的账号认证需要实际用户完成登录,尚未作为上游验收;没有替换正在运行的用户代理,没有发布、提交或推送。 + +构建文件、运行日志和测试缓存保持忽略;所有代码、方案、验收文档均可在当前分支检查。原始工作区在改造前为干净状态。 diff --git a/ACCOUNT_MANAGER_PLAN.md b/ACCOUNT_MANAGER_PLAN.md new file mode 100644 index 0000000..fedcdae --- /dev/null +++ b/ACCOUNT_MANAGER_PLAN.md @@ -0,0 +1,55 @@ +# 账号与登录态管理升级方案 + +基线:2026-10-24,origin/main `83f60e90eccf5495356b97cb2def79533a38d9a5`。 +工作分支:`codex/account-session-manager`。 + +## 目标与用户流程 + +代理启动时同时提供 `/admin` 管理页;未关联账号也能启动。页面展示账号名、国内 / Global、当前生效状态,允许关联、切换、删除非当前账号。关联时先选择区域,在点击动作中打开独立浏览器窗口,然后跳转到原有登录协议返回的认证地址。代理继续使用原有 state → token → account 轮询协议,成功后保存独立会话;首次关联自动生效,已有账号时关联成功不隐式切换。 + +同一 uid 在不同区域或不同 endpoint 中是不同身份。同身份重新关联更新自身凭据,其他账号不会覆盖。取消、超时、失败均不落入正式账号库。账号切换持久保存,下一次启动恢复;显式 `--global` / endpoint / session-file 用于选择或导入相符账号,不能把不匹配区域的凭据发往另一地址。 + +## 数据与模块 + +1. 新增 `account_store.py`:默认目录 `~/.workbuddy2api/accounts`,可用 `--accounts-dir` 覆盖。版本化索引保存脱敏展示元数据与 active id,每账号独立会话文件。身份由 backend、规范化 endpoint、uid 共同确定;只使用受控 ID 生成路径。原子替换、线程锁、文件最小权限;损坏索引明确报错,绝不静默覆盖。旧会话仅在有效且身份可识别时复制导入,保留源文件。 +2. 在 `CodeBuddyClient` 抽出 `begin_login()` 与 `complete_login(challenge, timeout=300, cancel_event=None)`;`login()` 组合复用二者,保留 CLI 行为。等待 token 和 account 都受总超时、取消控制,account 必须含 uid。增加 `ensure_authenticated(..., allow_login=True)`,管理模式请求只允许刷新,不自动打开登录浏览器。会话写入使用原子替换。 +3. 新增 `account_manager.py`:持有 AccountStore、当前 client,后台登录任务有 pending / succeeded / failed / cancelled 状态;返回 job id、auth_url、简洁错误,绝不返回 token、完整 session、上游原文异常。限制并发登录并回收临时会话和终态任务;关闭时取消任务。切换一次替换整个 client。 +4. 新增 `admin_api.py` 与包内 `admin.html`:同端口服务页面和管理 API,无前端构建依赖。只允许 loopback 访问,并校验 Host / Origin;写操作要求专用请求头,拒绝跨站请求。页面使用 textContent 渲染账号名,管理响应 no-store,不输出凭据。 +5. 主入口集成:启动账号库、导入历史 domestic/global 和显式 session;保留 mock 模式。每次代理请求捕获一个 client 快照,已发送和进行中的流继续使用原账号,新请求使用新账号;billing 的 401 refresh 也锁定原 client。模型列表、health 随当前账号更新。无账号时模型/管理可用,推理返回结构化 401。 + +## 子任务与接口约定 + +- A(存储和认证,GPT-6 Luna / max):负责 account_store.py、codebuddy_client_demo.py、对应测试。AccountStore(root),list_accounts() -> list[dict](id, backend, endpoint, uid, nickname, active),active_account_id -> str | None,import_client(client, activate=False) -> str,get_client(id) -> CodeBuddyClient,set_active(id) -> CodeBuddyClient,active_client() -> client | None,delete_account(id)。begin_login() -> dict(auth_url, state);complete_login 接受该 dict,成功写入 client 自己的独立 session,异常为 CodeBuddyError。import_client 不修改源文件,第一次导入自动生效。 +- B(后台任务与 API,GPT-6 Luna / max):负责 account_manager.py、admin_api.py、对应测试。AccountManager(store),client -> 当前 client | None,snapshot() -> {accounts, active_account_id},activate(id),delete(id),start_login(backend) -> job dict,get_login(id) -> job dict,cancel_login(id),close()。router 工厂 create_admin_router(get_manager),管理前缀 `/admin/api`:GET /accounts,POST /accounts/{id}/activate,DELETE /accounts/{id},POST /logins {backend},GET/DELETE /logins/{id};页面 GET /admin(兼容 /admin/)。写操作 `X-Account-Manager: 1`。job: id, backend, status, auth_url, error, account_id;后台准备登录地址,GET job 用于轮询。 +- C(页面和用户文档,GPT-6 Luna / max):负责 admin.html、pyproject.toml 包资源、README.md、README_zh.md。页面只依赖上述协议;先打开窗口,再 start_login 并轮询 auth_url,成功刷新列表,可取消和重试,失败给出清楚状态,弹窗被拦截提供链接;不展示 session 或 token。 +- 主 agent:负责 __main__.py / billing_proxy.py 集成、启动兼容、跨模块测试、独立审阅和最终验收。子 agent 不提交,不改别人负责的文件,各自运行必要测试直到通过,并反馈命令、结果和未验证边界。 + +## 429 限流与自动切换补充(2026-10-24) + +管理页按“账号 + 模型”展示独立的限流标记和冷却倒计时,以及默认关闭的“遇 429 自动切换账号”开关。同账号某模型冷却不影响其他模型。开关和冷却截止时间写入独立版本化运行状态文件,不改账号索引或凭据结构。冷却优先读取 Retry-After(整数秒或 HTTP 日期);缺失或无效时采用 60 秒并注明默认估算。过期解除限流标记,不主动切回旧账号。 + +三种推理协议在上游真正返回 HTTP 429、客户端响应尚未开始时,可按开关切换同区域且同 endpoint、该请求模型未冷却且有可用凭据的账号,再尝试当前请求。每账号最多一次,尝试次数不超过请求开始时账号数量。候选用尽仍返回实际的 429 与对应 Retry-After。迟到的旧请求只能标记其原始账号,不覆盖期间的手动切换。已开始输出的流不重放。 + +签到等 billing 操作透传 Retry-After;这类请求没有模型维度,不记录推理模型冷却,也不切换账号或重放业务动作。运行状态沿用原文件路径,版本升级到 v2,按账号/模型嵌套保存。v1 开关继续保留;旧冷却因没有模型信息,在严格校验后不再沿用,不扩展为全账号限制。日志包含实际请求账号与模型。 + +分工沿用 GPT-6 Luna / max 子 agent:B 负责状态、持久化和 API;C 负责页面、文档、交互测试;A 负责独立代理链路测试;主 agent 实现 HTTP 集成并独立审阅和验收。 + +## 每账号每日签到与自动调度补充(2026-10-24) + +全局自动签到开关控制所有关联账号,不限当前生效账号。列表显示各账号今日状态,提供状态查询与签到按钮;每次使用指定账号自身 client 调用既有 billing 底层接口,不更换当前账号。状态查询/签到复用 POST checkin-activity-status / daily-checkin 和 401 刷新逻辑。根据原贡献 PR #7 的上游记录,以 code=0 且 data.today_checked_in 严格布尔值判断;未知、失败和活动未开放均不得当作未签到直接发动作。 + +自动签到默认关闭,持久化开关。北京时间(UTC+8)每日 00:00 对启用功能下的所有账号检查;未签到者在该午夜后 0–600 秒内安排独立随机时刻,再查未签到才执行签到。关闭开关取消待执行任务;实际 POST 前重检日期、账号是否仍存在和开关。开关从关变开立即检查并补签;服务启动时开关已开也立即检查并补签。 + +服务随代理进程启动和停止,不使用 Codex 应用提醒。每账号互斥并限制总并发,防止页面、启动、开关和定时任务重复执行。跨日状态不能当成今天;每个动作和结果、跳过、失败、调度时间记录到现有普通日志、JSONL 和控制台,只记录账号 ID/区域/日期等安全字段。 + +A 子 agent 负责签到 service/调度及单元测试,B 负责管理器设置迁移与 API,C 负责列表/开关/文档和交互测试;均沿用 GPT-6 Luna / max。主 agent 负责生命周期、日志接入、HTTP 集成测试、独立审阅和最终验收。 + +## 管理页中英文切换补充(2026-10-24) + +管理页顶部提供中文/English 选择,保存到当前浏览器 localStorage;有效保存值优先,其次跟随浏览器语言,存储不可用时仍可切换。页面标题、语言属性、无障碍标签、静态介绍、动态账号/签到/模型冷却状态、设置提示、登录进度和已知后端错误覆盖中英文。账号名、模型 ID、UID 和服务地址保留原值,以文本安全渲染。 + +语言切换只重新渲染展示,不刷新页面、不额外请求、不变更账号或开关,不取消或重复登录与签到任务;进行中的 busy 状态与操作结果按当前语言显示。页面内部显示明确的北京时间,不因显示语言改变后台调度时区。C 实现页面,A 独立扩充控制器回归,B 只读复核错误映射和并发操作;主 agent 负责文档、独立浏览器/测试/安装包验收。 + +## 验收 + +离线自动化覆盖:多账号不覆盖、两区域同 uid 隔离、原有会话导入、重启恢复、损坏数据、原子写入;登录 token/account 延迟、取消、超时及失败不写入;API 不泄漏凭据、跨站和非本机管理请求被拒;无账号启动、切换后 health / models / 三种协议新请求使用新账号、在途请求不串号、billing 刷新不串号。全量现有测试需与基线区分;包构建验证页面被包含。浏览器检查管理页加载、关联窗口、状态和布局。真实国内/国际交互登录依赖用户亲自认证,离线模拟通过不能代替真实上游验收。 diff --git a/README.md b/README.md index 959839e..7c23774 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,9 @@ For unexplained Windows process exits, use the [independent process monitor](doc - **DSML parsing** - Automatically detects and converts DeepSeek Markup Language (DSML) tool calls - **Streaming responses** - SSE streaming output, returning generated content in real time, with built-in 60-second timeout protection - **Domestic and international backends** - Uses domestic CodeBuddy by default and switches to the international CodeBuddy service with `--global`, with isolated sessions and model catalogs -- **Multi-account management** - Supports isolation of multiple login states for easy switching between work/personal accounts +- **Multi-account management** - Link domestic and Global accounts, switch the active account, and remove non-active accounts from the local admin page +- **Model-level 429 handling** - Track cooldowns per account and model, with optional same-region failover +- **Daily check-in** - Check in manually or enable background checks for all linked accounts --- @@ -51,12 +53,14 @@ uv sync uv run python -m codebuddy_proxy --desensitize ``` -On first use, when login is required: +On first use, start the proxy and open the local account manager: ```bash -uv run python -m codebuddy_proxy --login --desensitize +uv run python -m codebuddy_proxy --desensitize ``` +Then visit `http://127.0.0.1:8787/admin` and link a domestic or Global account. + ## Quick Start ### 1. Start the proxy @@ -65,8 +69,8 @@ uv run python -m codebuddy_proxy --login --desensitize # Use the latest version (recommended) uv run --with workbuddy2api python -m codebuddy_proxy --desensitize -# First use: log in and start -uv run --with workbuddy2api python -m codebuddy_proxy --login --desensitize +# First use: start, then link an account at http://127.0.0.1:8787/admin +uv run --with workbuddy2api python -m codebuddy_proxy --desensitize # International service: first login and start uv run --with workbuddy2api python -m codebuddy_proxy --global --login --desensitize @@ -82,12 +86,23 @@ Listens on `http://127.0.0.1:8787` by default. The proxy supports both CodeBuddy regions. Domestic mode remains the default, so existing startup commands continue to work unchanged. -| Mode | Startup option | Upstream endpoint | Default session file | -| --- | --- | --- | --- | -| Domestic (default) | none | `https://copilot.tencent.com` | `~/.codebuddy-session.json` | -| International | `--global` | `https://www.codebuddy.ai` | `~/.codebuddy-global-session.json` | +| Mode | Startup option | Upstream endpoint | +| --- | --- | --- | +| Domestic (default) | none | `https://copilot.tencent.com` | +| International | `--global` | `https://www.codebuddy.ai` | + +Linked accounts are stored separately by region in `~/.workbuddy2api/accounts`. +Use `--accounts-dir DIR` to choose another local directory. With no explicit +selector, the proxy restores the saved active account. On the first migration, +matching legacy session files are imported once and kept; deleting an imported +account does not import it again. Explicit `--global` or `--endpoint` selects a +matching account, while `--session-file` can import that file again. If there is +no match, the proxy starts without an active account. `--login` uses a temporary +session file to add or refresh an account without replacing a legacy file. -First login and daily startup for the international service: +The account manager can add accounts for either region while the proxy is +running. For a terminal-driven login, the existing `--login` option remains +available: ```bash # First login @@ -103,16 +118,18 @@ If the package is installed as a command, the equivalent short form is `workbuddy2api --global --login`. When running from source, use `uv run python -m codebuddy_proxy --global --login`. -`--global` changes only the upstream service. Codex CLI, Claude Code/CC Switch, -OpenCode, and other clients continue to use the same local proxy URLs. The -selected profile also controls the model catalog returned by `/v1/models`. +`--global` selects the Global backend and a matching account when one exists. +Codex CLI, Claude Code/CC Switch, OpenCode, and other clients continue to use +the same local proxy URLs. The selected profile also controls the model catalog +returned by `/v1/models`. -Domestic and international sessions are intentionally isolated. An explicit -`--session-file` takes precedence, but its saved backend and endpoint must match -the current startup options. The proxy refuses mismatched credentials and does -not modify the original file. +Domestic and international accounts remain isolated. An explicit +`--session-file` is imported only when its saved backend and endpoint match the +current startup options. The proxy refuses mismatched credentials and does not +modify the original file. -For a custom international endpoint, always use a dedicated session file: +For a custom international endpoint, use a dedicated session file when +importing or creating a login from the command line: ```bash uv run --with workbuddy2api python -m codebuddy_proxy \ @@ -309,6 +326,7 @@ Select the model inside OMP with `/model codebuddy/hy3` (or set it as the defaul --global Use the international CodeBuddy backend (default domestic) --endpoint ENDPOINT CodeBuddy backend address (overrides the profile default) --session-file PATH Session file path (isolated by profile by default) +--accounts-dir DIR Account library directory (default ~/.workbuddy2api/accounts) --log-file PATH JSONL log file (default ~/.workbuddy2api/codebuddy-proxy.jsonl) --desensitize Enable desensitization (recommended) --optimize-context Enable message compression (recommended for Codex CLI) @@ -369,15 +387,57 @@ uv run --with workbuddy2api python -m codebuddy_proxy --desensitize --optimize-c ### Multi-account switching -```bash -# Account 1 -uv run --with workbuddy2api python -m codebuddy_proxy --session-file ~/.codebuddy-work.json --login \ - --log-file "$HOME/.workbuddy2api/codebuddy-proxy.jsonl" - -# Account 2 -uv run --with workbuddy2api python -m codebuddy_proxy --session-file ~/.codebuddy-personal.json --login \ - --log-file "$HOME/.workbuddy2api/codebuddy-proxy.jsonl" -``` +Start the proxy and open [http://127.0.0.1:8787/admin](http://127.0.0.1:8787/admin). +Choose **Domestic** or **Global**, then link the account in the browser window. +The first linked account becomes active automatically; linking another account +does not change the account currently in use. Use **Set active** to switch, or +remove a non-active account from the same page. The selection is saved for the +next proxy start. Switching affects new requests; requests already in progress +continue with the account selected when they began. + +![Account manager with check-in controls and redacted account identifiers](docs/images/account-manager-checkin.png) + +The screenshot uses synthetic demo accounts. Account UIDs are covered, and no +credentials or session contents are shown. Accounts are kept separately by +backend, normalized endpoint, and user ID in `~/.workbuddy2api/accounts`; use +`--accounts-dir DIR` to choose another local directory. The manager remains +available before any account is linked. Its page and API are restricted to +loopback access, and API responses do not include tokens or session data. + +Use the language selector at the top of the page to switch between Chinese and +English. The choice is remembered in this browser; without a saved preference, +the page follows the browser language. Switching languages keeps ongoing login +and account operations running. + +Automatic account switching after an upstream HTTP `429` is off by default. +When enabled, a request can fail over only before any response has been sent, +and only to accounts with the same region and endpoint. Each account is tried +at most once per request. The proxy skips a candidate only when the request's +model is cooling down there; other models on that account remain eligible. A +missing or invalid model does not create a model cooldown or trigger failover. +If every eligible account returns `429`, the original `429` is returned. +Cooldowns are tracked per account and model, using upstream `Retry-After` or a +default 60 seconds when that value is missing or invalid. The setting and model +cooldowns persist across restarts. Old v1 account-wide cooldowns had no model +identity and are not carried forward during upgrade. The page shows a separate +remaining countdown and source for each model, and refreshes account state +periodically and immediately when a hidden page becomes visible again. + +The account list also shows each linked account's check-in status for today +(Beijing time). The **Check in** action is available for every account, including +the active one; it checks the account first and does nothing when today's check-in +is already complete. The global automatic check-in setting is off by default and +applies to all linked accounts without changing the active account. The page saves +this setting and displays status; the proxy performs automatic checks and check-ins +in the background, including when the management page is closed. When enabled, the +proxy checks all linked accounts immediately and on startup. At Beijing midnight it +checks them, waits a random 0–10 minutes, then checks again before signing; disabling +it cancels pending scheduled work. Unknown status is never treated as permission to +sign. The setting persists; after restart, check-in status is confirmed by querying +the upstream service again. + +The command-line `--login` flow and `--session-file` option remain available for +terminal use and importing an existing session. ### Listening on all interfaces (LAN sharing) @@ -393,11 +453,20 @@ All endpoints do not require an extra token in the request by default; the proxy | Method | Path | Purpose | | --- | --- | --- | | GET | `/health` | Query local service and auth status | +| GET | `/admin` | Open the local account manager | | GET | `/v1/models` | Query the CodeBuddy model list | | POST | `/v1/chat/completions` | OpenAI Chat Completions, supports tools and streaming | | POST | `/v1/responses` | Responses API, compatible with Codex CLI | | POST | `/v1/messages` | Anthropic Messages API, compatible with Claude Code / CC Switch | +Account-management endpoints live under `/admin/api` and are restricted to +loopback clients. Write requests require the `X-Account-Manager: 1` header, and +cross-origin requests are rejected. `GET /admin/api/accounts` reports active +account, cooldown, and cached check-in state. `GET /admin/api/accounts/{id}/checkin-status` +checks that account against the upstream service; `POST /admin/api/accounts/{id}/checkin` +checks first and signs only when needed. `GET` and `PATCH /admin/api/settings` read +and save the automatic-switch and automatic-check-in preferences independently. + ### `/health` - health check ```bash diff --git a/README_zh.md b/README_zh.md index 38020a7..84985ea 100644 --- a/README_zh.md +++ b/README_zh.md @@ -15,7 +15,9 @@ Windows 代理无异常记录却自动退出时,使用[独立进程监控](doc - **DSML 解析** - 自动识别和转换 DeepSeek 标记语言(DSML)格式的工具调用 - **流式响应** - 支持 SSE 流式输出,实时返回生成内容,内置 60 秒超时保护 - **国内版与国际版后端** - 默认连接国内版 CodeBuddy,通过 `--global` 切换国际版,并隔离登录态和模型列表 -- **多账号管理** - 支持多个登录态隔离,方便工作/个人账号切换 +- **多账号管理** - 在本机管理页关联国内版和 Global 账号、切换当前账号或删除非当前账号 +- **模型级 429 冷却** - 按账号和模型记录冷却状态,并可选启用同区域账号自动切换 +- **每日签到** - 可手动签到,也可开启覆盖全部关联账号的后台自动检查 --- ## 安装 @@ -50,12 +52,14 @@ uv sync uv run python -m codebuddy_proxy --desensitize ``` -首次使用需要登录时: +首次使用时,启动代理并打开本机账号管理页: ```bash -uv run python -m codebuddy_proxy --login --desensitize +uv run python -m codebuddy_proxy --desensitize ``` +然后访问 `http://127.0.0.1:8787/admin`,关联国内版或 Global 账号。 + ## 快速开始 ### 1. 启动 proxy @@ -64,8 +68,8 @@ uv run python -m codebuddy_proxy --login --desensitize # 使用最新版本(推荐) uv run --with workbuddy2api python -m codebuddy_proxy --desensitize -# 首次使用:登录并启动 -uv run --with workbuddy2api python -m codebuddy_proxy --login --desensitize +# 首次使用:启动后访问 http://127.0.0.1:8787/admin 关联账号 +uv run --with workbuddy2api python -m codebuddy_proxy --desensitize # 国际版:首次登录并启动 uv run --with workbuddy2api python -m codebuddy_proxy --global --login --desensitize @@ -81,12 +85,20 @@ uv run --with workbuddy2api python -m codebuddy_proxy --global --desensitize 代理同时支持国内版和国际版 CodeBuddy。国内版仍是默认模式,因此已有启动命令 无需修改。 -| 模式 | 启动参数 | 上游 endpoint | 默认 session 文件 | -| --- | --- | --- | --- | -| 国内版(默认) | 无 | `https://copilot.tencent.com` | `~/.codebuddy-session.json` | -| 国际版 | `--global` | `https://www.codebuddy.ai` | `~/.codebuddy-global-session.json` | +| 模式 | 启动参数 | 上游 endpoint | +| --- | --- | --- | +| 国内版(默认) | 无 | `https://copilot.tencent.com` | +| 国际版 | `--global` | `https://www.codebuddy.ai` | + +关联账号按区域分别保存在 `~/.workbuddy2api/accounts`。可以通过 +`--accounts-dir DIR` 指定其他本机目录。未指定账号选择参数时,代理恢复上次选择的当前账号。 +首次迁移时会导入匹配的旧 session 文件, +并保留源文件;之后删除已导入账号不会再次自动导入。显式 `--global` 或 `--endpoint` +会选择相符账号,显式 `--session-file` 可再次导入该文件。找不到匹配账号时,代理会在无当前账号的状态下启动。 +`--login` 使用临时 session 文件关联或更新账号,不会覆盖旧 session 文件。 -国际版首次登录和日常启动: +代理运行期间可直接通过管理页关联任一区域的账号。也可以继续使用命令行 +`--login` 执行终端登录: ```bash # 首次登录 @@ -101,14 +113,13 @@ uv run --with workbuddy2api python -m codebuddy_proxy \ 如果已将包安装为命令,可使用简写 `workbuddy2api --global --login`。从本地源码 运行时,可使用 `uv run python -m codebuddy_proxy --global --login`。 -`--global` 只切换上游服务。Codex CLI、Claude Code/CC Switch、OpenCode 等客户端 -仍然连接相同的本地代理地址;`/v1/models` 返回的模型列表会随 profile 切换。 +`--global` 选择 Global 后端及匹配的账号。Codex CLI、Claude Code/CC Switch、OpenCode +等客户端仍然连接相同的本地代理地址;`/v1/models` 返回的模型列表会随 profile 切换。 -国内版和国际版登录态默认相互隔离。显式 `--session-file` 的优先级更高,但文件中 -保存的 backend 和 endpoint 必须与当前启动参数匹配。代理会拒绝不匹配的凭据, -并保留原 session 文件不变。 +国内版和国际版账号保持隔离。显式 `--session-file` 仅在其中保存的 backend 和 +endpoint 与当前启动参数匹配时导入。代理会拒绝不匹配的凭据,并保留原 session 文件。 -使用自定义国际版 endpoint 时,应同时指定独立的 session 文件: +从命令行导入或创建自定义国际版 endpoint 的登录态时,应使用独立的 session 文件: ```bash uv run --with workbuddy2api python -m codebuddy_proxy \ @@ -321,6 +332,7 @@ providers: --global 使用国际版 CodeBuddy(默认国内版) --endpoint ENDPOINT CodeBuddy 后端地址(覆盖 profile 默认值) --session-file PATH 会话文件路径(默认按 profile 隔离) +--accounts-dir DIR 账号库目录(默认 ~/.workbuddy2api/accounts) --log-file PATH JSONL 日志文件(默认 ~/.workbuddy2api/codebuddy-proxy.jsonl) --desensitize 启用脱敏处理(推荐) --optimize-context 启用消息压缩优化(Codex CLI 推荐) @@ -381,15 +393,34 @@ uv run --with workbuddy2api python -m codebuddy_proxy --desensitize --optimize-c ### 多账号切换 -```bash -# 账号 1 -uv run --with workbuddy2api python -m codebuddy_proxy --session-file ~/.codebuddy-work.json --login \ - --log-file "$HOME/.workbuddy2api/codebuddy-proxy.jsonl" +启动代理后访问 [http://127.0.0.1:8787/admin](http://127.0.0.1:8787/admin),选择 +**国内**或 **Global**,再在新打开的浏览器窗口中完成关联。首次关联的账号会自动设为当前账号; +之后关联其他账号不会隐式切换。可在列表中切换当前账号或删除非当前账号,当前选择会在下次启动时恢复。 +切换只影响之后开始的新请求;已经进行中的请求会继续使用进入时固定的账号。 -# 账号 2 -uv run --with workbuddy2api python -m codebuddy_proxy --session-file ~/.codebuddy-personal.json --login \ - --log-file "$HOME/.workbuddy2api/codebuddy-proxy.jsonl" -``` +![账号管理页面,展示签到控件并遮挡账号标识](docs/images/account-manager-checkin.png) + +截图中的账号为合成演示数据,UID 已遮挡,不包含凭据或 session 内容。账号按 backend、规范化 endpoint +和用户 ID 区分,分别保存在 `~/.workbuddy2api/accounts`;可用 `--accounts-dir DIR` 更换本机目录。 +尚未关联任何账号时,管理页也可以打开。页面和管理 API 仅允许 loopback 本机访问,API 响应不会包含 token 或 session 数据。 +页面顶部可切换中文和 English,选择会保存在当前浏览器;没有保存偏好时跟随浏览器语言。 +切换即时生效,不会中断进行中的登录、签到或账号操作。 + +遇到上游 HTTP `429` 时自动切换账号的功能默认关闭。开启后,仅当尚未向客户端发送响应时才会重试, +且只会尝试区域和 endpoint 都相同的账号;每个账号在一次请求中最多尝试一次。仅当该账号的“本次请求模型”仍在冷却时才会跳过,同一账号的其他模型仍可用。 +模型名缺失或无效时不会记录模型冷却,也不会触发自动切换。 +所有候选账号都返回 `429` 时,仍返回原始 `429`。冷却按账号和模型分别记录,遵循上游 `Retry-After`; +该值缺失或无效时,默认冷却 60 秒。自动切换设置和模型冷却会在重启后保留。 +旧版 v1 的整账号冷却没有模型信息,升级后不会沿用。管理页分别显示每个模型的剩余冷却时间 +及其来源(`Retry-After` 或默认值),并定期同步账号状态;页面从后台重新显示时会立即同步。 + +账号列表还会显示每个关联账号当天(北京时间)的签到状态。每个账号都有“查询签到”和“签到”按钮, +当前账号也可操作;签到前会先检查,已签到时不会重复签到。全局自动签到默认关闭,适用于所有关联账号, +不会切换当前账号。管理页只保存开关并显示状态,自动检查和签到由代理在后台执行,即使关闭管理页也会继续。 +开启后代理会立即检查全部账号,启动时也会重新检查;此外北京时间每日 00:00 检查所有账号,随机等待 0 至 10 分钟后再检查,确认仍未签到才会执行; +关闭开关会取消等待中的计划。未知状态不会自动签到。开关会持久化;重启后,签到状态以重新查询上游的结果为准。 + +命令行 `--login` 和 `--session-file` 仍可用于终端登录和导入已有 session。 ### 监听所有网卡(局域网共享) @@ -405,11 +436,18 @@ uv run --with workbuddy2api python -m codebuddy_proxy --host 0.0.0.0 --desensiti | 方法 | 路径 | 用途 | | --- | --- | --- | | GET | `/health` | 查询本地服务和认证状态 | +| GET | `/admin` | 打开本机账号管理页 | | GET | `/v1/models` | 查询 CodeBuddy 模型列表 | | POST | `/v1/chat/completions` | OpenAI Chat Completions,支持 tools 和流式响应 | | POST | `/v1/responses` | Responses API,兼容 Codex CLI | | POST | `/v1/messages` | Anthropic Messages API,兼容 Claude Code / CC Switch | +账号管理接口位于 `/admin/api`,只允许本机回环地址访问。写请求必须携带 +`X-Account-Manager: 1`,跨源请求会被拒绝。`GET /admin/api/accounts` 返回当前账号和冷却状态; +每个账号还包含缓存的当日签到状态。`GET /admin/api/accounts/{id}/checkin-status` 会向上游查询; +`POST /admin/api/accounts/{id}/checkin` 会先检查再执行必要的签到。`GET` 与 `PATCH /admin/api/settings` +用于分别读取和保存自动切换、全局自动签到设置。 + ### `/health` - 健康检查 ```bash diff --git a/docs/images/account-manager-checkin.png b/docs/images/account-manager-checkin.png new file mode 100644 index 0000000..d0be4c0 Binary files /dev/null and b/docs/images/account-manager-checkin.png differ diff --git a/pyproject.toml b/pyproject.toml index 6692d38..f518f9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ build-backend = "setuptools.build_meta" where = ["src"] [tool.setuptools.package-data] -codebuddy_proxy = ["models_config.*.json"] +codebuddy_proxy = ["models_config.*.json", "admin.html"] [dependency-groups] dev = [ diff --git a/src/codebuddy_proxy/__main__.py b/src/codebuddy_proxy/__main__.py index ed3fe35..0cdcff4 100755 --- a/src/codebuddy_proxy/__main__.py +++ b/src/codebuddy_proxy/__main__.py @@ -15,7 +15,9 @@ import argparse import asyncio +import copy from contextlib import asynccontextmanager, nullcontext +from contextvars import ContextVar import base64 import hashlib import io @@ -29,6 +31,7 @@ import platform import sys import time +import tempfile import uuid from contextlib import AsyncExitStack from typing import Any, Optional @@ -44,6 +47,8 @@ BackendProfile, ResolvedBackend, DOMESTIC_PROFILE, + GLOBAL_PROFILE, + resolve_backend, resolve_backend_from_args, ) from codebuddy_proxy.codebuddy_client_demo import CodeBuddyClient, CodeBuddyError @@ -51,6 +56,7 @@ from codebuddy_proxy.upstream_request import ChatRequestPreparer from codebuddy_proxy.protocol_response import format_nonstream_response from codebuddy_proxy.process_diagnostics import ProcessDiagnostics +from codebuddy_proxy.admin_api import create_admin_router # 尝试导入高级功能模块(可选) try: @@ -180,8 +186,12 @@ def __init__( verbose_llm: bool = False, logger: logging.Logger | None = None, json_logger: logging.Logger | None = None, + account_manager=None, ): - self.client = client + self._client = client + self.account_manager = account_manager + self.account_runtime = account_manager + self.managed_auth = account_manager is not None self.mock_dir = mock_dir self.log_file = log_file self.enable_desensitize = enable_desensitize @@ -191,13 +201,47 @@ def __init__( self.json_logger = json_logger self.runtime_info = get_runtime_info() self.started_at = time.time() + checkins = getattr(account_manager, "checkins", None) + if checkins is not None: + checkins.configure_logging(self.log_checkin) + + def log_checkin(self, event: str, **fields) -> None: + """Send safe check-in events to existing files and the process console.""" + self.write_log(event, **fields) + summary = json.dumps(fields, ensure_ascii=False) + if self.logger is not None: + self.logger.info("%s: %s", event, summary) + try: + print(f"[Checkin] {event} {summary}", flush=True) + except (OSError, UnicodeError): + pass + + @property + def client(self): + if self.account_manager is not None: + return self.account_manager.client or self._client + return self._client + + @client.setter + def client(self, value): + self._client = value + + def for_request(self): + """Pin the account before any await, including refresh and billing retries.""" + snapshot = copy.copy(self) + snapshot._client = self.client + snapshot.account_manager = None + return snapshot def ensure_auth(self) -> None: """确保已认证;认证失败(token 过期/网络错误/登录未完成)返回结构化 401 而非 500。""" if self.mock_dir is not None: return try: - self.client.ensure_authenticated() + if self.managed_auth: + self.client.ensure_authenticated(open_browser=False, allow_login=False) + else: + self.client.ensure_authenticated() except HTTPException: raise # 已是结构化异常,原样透传(如 503 proxy not initialized) except Exception as exc: @@ -205,9 +249,9 @@ def ensure_auth(self) -> None: status_code=401, detail={ "error": { - "message": "认证失败:token 无效或已过期,请重新登录(--login)", + "message": "认证失败:请在 /admin 关联或切换账号,或重新登录(--login)", "type": "authentication_error", - "details": str(exc)[:200], + "details": type(exc).__name__ if self.managed_auth else str(exc)[:200], } }, ) @@ -257,6 +301,21 @@ def write_body_log(self, event: str, body: bytes, **kwargs) -> None: # 全局状态(在 main() 中初始化) proxy_state: ProxyState | None = None process_diagnostics: ProcessDiagnostics | None = None +_request_state: ContextVar[Any] = ContextVar("account_request_state", default=None) + + +@app.middleware("http") +async def pin_account_at_request_entry(request: Request, call_next): + # Capture before body parsing or other awaits. Each downstream task also + # inherits this context, including streaming responses. + if proxy_state is None or not request.url.path.startswith("/v1/"): + return await call_next(request) + state = proxy_state.for_request() if isinstance(proxy_state, ProxyState) else proxy_state + token = _request_state.set(state) + try: + return await call_next(request) + finally: + _request_state.reset(token) async def install_process_asyncio_diagnostics() -> None: @@ -269,12 +328,47 @@ async def install_process_asyncio_diagnostics() -> None: app.router.add_event_handler("startup", install_process_asyncio_diagnostics) +async def start_account_checkin() -> None: + if isinstance(proxy_state, ProxyState) and proxy_state.account_manager is not None: + await proxy_state.account_manager.checkins.start() + + +app.router.add_event_handler("startup", start_account_checkin) + + +async def close_account_manager() -> None: + if isinstance(proxy_state, ProxyState) and proxy_state.account_manager is not None: + await proxy_state.account_manager.checkins.close() + proxy_state.account_manager.close() + + +app.router.add_event_handler("shutdown", close_account_manager) + + def get_state() -> ProxyState: if proxy_state is None: raise HTTPException(status_code=503, detail={"error": {"message": "proxy not initialized", "type": "internal_error"}}) return proxy_state +def get_request_state(): + pinned = _request_state.get() + if pinned is not None: + return pinned + state = get_state() + return state.for_request() if isinstance(state, ProxyState) else state + + +def get_account_manager(): + manager = getattr(get_state(), "account_manager", None) + if manager is None: + raise HTTPException(status_code=503, detail="Account management is unavailable in this mode") + return manager + + +app.include_router(create_admin_router(get_account_manager)) + + # ============================================================================ # 辅助函数 # ============================================================================ @@ -857,7 +951,7 @@ async def health(): async def checkin_status(): """Accept GET/POST locally; the real upstream billing status endpoint requires POST.""" return await forward_billing( - get_state(), "POST", "/v2/billing/meter/checkin-status", diagnostic + get_request_state(), "POST", "/v2/billing/meter/checkin-activity-status", diagnostic ) @@ -865,7 +959,7 @@ async def checkin_status(): async def checkin(): """POST daily check-in: preserve the upstream HTTP status and raw body.""" return await forward_billing( - get_state(), "POST", "/v2/billing/meter/daily-checkin", diagnostic + get_request_state(), "POST", "/v2/billing/meter/daily-checkin", diagnostic ) @@ -935,7 +1029,7 @@ async def parse_request_body(request: Request) -> Any: @app.post("/v1/chat/completions") async def chat_completions(request: Request): - state = get_state() + state = get_request_state() body = await parse_request_body(request) log_client_request("POST", "/v1/chat/completions", body) @@ -950,7 +1044,7 @@ async def chat_completions(request: Request): @app.post("/v1/responses") async def create_response(request: Request): - state = get_state() + state = get_request_state() body = await parse_request_body(request) log_client_request("POST", "/v1/responses", body) @@ -1013,7 +1107,7 @@ async def create_response(request: Request): @app.post("/v1/messages") async def create_message(request: Request): - state = get_state() + state = get_request_state() body = await parse_request_body(request) log_client_request("POST", "/v1/messages", body) @@ -1052,14 +1146,71 @@ def _normalize_tool_choice(tool_choice: Any) -> Any: return tool_choice +def _retry_after_header(response) -> dict[str, str]: + retry_after = getattr(response, "headers", {}).get("Retry-After") + return {"Retry-After": retry_after} if retry_after is not None else {} + + async def forward_chat( body: dict[str, Any], protocol: str, original: dict[str, Any] | None = None ) -> StreamingResponse | JSONResponse: - """转发 chat 请求到 CodeBuddy 上游,支持流式和非流式。""" - state = get_state() - state.ensure_auth() + """Forward with bounded account failover before sending a client response.""" + state = get_request_state() + manager = getattr(state, "account_runtime", None) + attempted: set[str] = set() + max_attempts = 1 + if manager is not None: + max_attempts = max(1, len(manager.snapshot()["accounts"])) + for attempt in range(max_attempts): + response = await _forward_chat_once(body, protocol, original, state) + if response.status_code != 429: + return response + retry_state = await rate_limited_account( + state, response.headers.get("Retry-After"), attempted, model=body.get("model"), + ) + if retry_state is None or attempt + 1 == max_attempts: + return response + state = retry_state + return response + + +async def rate_limited_account(state, retry_after, attempted, *, model): + """Record the original identity; late failures cannot override manual switches.""" + manager = getattr(state, "account_runtime", None) + if manager is None: + return None + try: + account_id = manager.account_id_for_client(state.client) + if account_id is not None: + attempted.add(account_id) + candidate = await asyncio.to_thread( + manager.handle_rate_limit, state.client, retry_after, + model=model, excluded_account_ids=set(attempted), + ) + next_id = manager.account_id_for_client(candidate) if candidate is not None else None + state.write_log( + "account_rate_limited", account_id=account_id, model=model, switched_to=next_id, + ) + if candidate is None or next_id is None or next_id in attempted: + return None + retry_state = copy.copy(state) + retry_state._client = candidate + retry_state.account_manager = None + return retry_state + except Exception as exc: + # Persistence or selection failures must not hide the provider's 429. + state.write_log("account_rate_limit_record_failed", error_type=type(exc).__name__) + return None + + +async def _forward_chat_once(body, protocol, original, state): + """One account attempt; any returned stream has already passed HTTP preflight.""" + if isinstance(state, ProxyState) and state.managed_auth: + await asyncio.to_thread(state.ensure_auth) + else: + state.ensure_auth() diagnostic("upstream_request", protocol=protocol, **body_summary(body)) @@ -1138,7 +1289,8 @@ async def open_preflight_response(): if protocol == "anthropic" else {"error": error_info} ) return JSONResponse( - status_code=upstream_response.status_code, content=error_payload + status_code=upstream_response.status_code, content=error_payload, + headers=_retry_after_header(upstream_response), ) finally: await upstream_response.aclose() @@ -1155,7 +1307,14 @@ async def open_preflight_response(): ) # 非流式:聚合后返回 - collected = await collect_upstream(url, headers, upstream_body, protocol) + try: + collected = await collect_upstream(url, headers, upstream_body, protocol) + except HTTPException as exc: + if exc.status_code != 429: + raise + return JSONResponse( + status_code=429, content={"detail": exc.detail}, headers=exc.headers, + ) return JSONResponse(content=convert_nonstream(collected, protocol, original)) @@ -1812,7 +1971,8 @@ async def collect_upstream( error_body = await resp.aread() raise HTTPException( status_code=resp.status_code, - detail={"error": {"message": error_body.decode("utf-8", "replace")[:500], "type": "upstream_error"}} + detail={"error": {"message": error_body.decode("utf-8", "replace")[:500], "type": "upstream_error"}}, + headers=_retry_after_header(resp), ) async for line in resp.aiter_lines(): @@ -1973,6 +2133,11 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--session-file", type=pathlib.Path, help="会话文件路径(默认按 backend profile 隔离)") + parser.add_argument( + "--accounts-dir", type=pathlib.Path, + default=pathlib.Path.home() / ".workbuddy2api" / "accounts", + help="独立账号库目录;管理页位于 /admin", + ) parser.add_argument("--mock-dir", type=pathlib.Path, help="只使用指定目录中的真实响应 fixture,不访问 CodeBuddy 后端") default_log_file = pathlib.Path( @@ -2017,11 +2182,92 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: return args +def initialize_account_manager(args, resolved_backend): + """Import legacy identities once and restore a persisted active account. + + Reading legacy files never refreshes or rewrites them. A forced login + always uses a temporary session, so --login cannot overwrite an old user. + """ + from codebuddy_proxy.account_store import AccountStore + from codebuddy_proxy.account_manager import AccountManager + + store = AccountStore(args.accounts_dir) + initial_import = not store.index_file.exists() + explicit = bool(args.global_mode or args.endpoint is not None + or args.session_file is not None + or os.getenv("CODEBUDDY_ENDPOINT")) + candidates = [resolved_backend] + for profile in (DOMESTIC_PROFILE, GLOBAL_PROFILE): + candidate = resolve_backend(profile=profile, environment={}) + if candidate not in candidates: + candidates.append(candidate) + selected_uid = None + for candidate in candidates: + # Migration is one-time. Re-reading a legacy file after its account + # was deleted must not silently associate that identity again. + if not initial_import and not (args.session_file is not None + and candidate is resolved_backend): + continue + try: + old = CodeBuddyClient(backend=candidate) + except CodeBuddyError: + if args.session_file is not None and candidate is resolved_backend: + raise + print(f"Skipped incompatible legacy session for {candidate.profile.key}", file=sys.stderr) + continue + account = old.session.get("account") or {} + auth = old.session.get("auth") or {} + if not isinstance(account, dict) or not isinstance(auth, dict): + continue + uid = account.get("uid") + if not uid or not (auth.get("accessToken") or auth.get("refreshToken")): + continue + if candidate is resolved_backend: + selected_uid = str(uid).strip() + known = any( + row["backend"] == old.profile.key and row["endpoint"] == old.endpoint + and str(row["uid"]) == str(uid) for row in store.list_accounts() + ) + if not known: + store.import_client(old) + if args.login: + with tempfile.TemporaryDirectory(prefix="login-", dir=store.root) as temporary: + login_backend = resolve_backend( + profile=resolved_backend.profile, endpoint=resolved_backend.endpoint, + session_file=pathlib.Path(temporary) / "session.json", environment={}, + ) + login_client = CodeBuddyClient(backend=login_backend) + login_client.login(open_browser=not args.no_browser) + store.import_client(login_client, activate=True) + elif explicit: + matching = [row for row in store.list_accounts() + if row["backend"] == resolved_backend.profile.key + and row["endpoint"] == resolved_backend.endpoint + and (args.session_file is None or str(row["uid"]) == selected_uid)] + if matching: + chosen = next((row for row in matching if row["active"]), matching[0]) + store.set_active(chosen["id"]) + else: + store.clear_active() + manager = AccountManager(store) + # The empty fallback preserves the requested profile for health/models. + # Managed requests cannot start browser login or write to a legacy file. + fallback = CodeBuddyClient(backend=resolve_backend( + profile=resolved_backend.profile, endpoint=resolved_backend.endpoint, + session_file=store.root / ".unlinked-session.json", environment={}, + )) + return manager, manager.client or fallback + + def _run_proxy(args: argparse.Namespace) -> int: global proxy_state try: resolved_backend = resolve_backend_from_args(args) - client = CodeBuddyClient(backend=resolved_backend) + manager = None + if args.mock_dir is not None: + client = CodeBuddyClient(backend=resolved_backend) + else: + manager, client = initialize_account_manager(args, resolved_backend) except (CodeBuddyError, ValueError) as exc: print(f"错误: {exc}", file=sys.stderr) return 1 @@ -2031,10 +2277,6 @@ def _run_proxy(args: argparse.Namespace) -> int: logger = setup_logging(log_dir) json_logger = setup_json_logging(args.log_file) - # 处理登录 - if args.login: - client.login(open_browser=not args.no_browser) - # 创建全局状态 proxy_state = ProxyState( client=client, @@ -2045,6 +2287,7 @@ def _run_proxy(args: argparse.Namespace) -> int: verbose_llm=args.verbose_llm, logger=logger, json_logger=json_logger, + account_manager=manager, ) proxy_state.write_log( "startup", @@ -2086,13 +2329,19 @@ def _run_proxy(args: argparse.Namespace) -> int: # 启动信息输出到 stdout print(f"CodeBuddy proxy listening on http://{args.host}:{args.port}") print("Endpoints: /v1/models /v1/chat/completions /v1/responses /v1/messages /health") + if manager is not None: + print(f"Account manager: http://127.0.0.1:{args.port}/admin (local access only)") # 同时记录到日志 logger.info(f"CodeBuddy proxy listening on http://{args.host}:{args.port}") logger.info("Endpoints: /v1/models /v1/chat/completions /v1/responses /v1/messages /health") # 启动 uvicorn - run_http_server(host=args.host, port=args.port, windows_loop=args.windows_loop) + try: + run_http_server(host=args.host, port=args.port, windows_loop=args.windows_loop) + finally: + if manager is not None: + manager.close() return 0 diff --git a/src/codebuddy_proxy/account_checkin.py b/src/codebuddy_proxy/account_checkin.py new file mode 100644 index 0000000..da0371a --- /dev/null +++ b/src/codebuddy_proxy/account_checkin.py @@ -0,0 +1,816 @@ +"""Per-account daily check-in automation with fail-closed status handling.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from datetime import date, datetime, time, timedelta, timezone +import json +import logging +import random +import threading +import time as time_module +from typing import Any + +from fastapi.responses import Response + +from codebuddy_proxy.billing_proxy import forward_billing + + +BEIJING = timezone(timedelta(hours=8), name="Asia/Shanghai") +# The legacy checkin-status endpoint can return an empty, inactive activity. +STATUS_PATH = "/v2/billing/meter/checkin-activity-status" +CHECKIN_PATH = "/v2/billing/meter/daily-checkin" +_SAFE_ERRORS = frozenset( + { + "签到活动未开放", + "签到状态查询失败", + "签到请求失败", + "签到结果未确认", + "账号不存在", + } +) +_REASONS = frozenset({"manual", "startup", "enable", "midnight", "timer"}) + + +class _BillingAdapterState: + """Bind billing_proxy to one saved client without its global failover hook.""" + + managed_auth = True + account_runtime = None + logger = None + json_logger = None + + def __init__(self, client: Any, diagnostic: Callable[..., Any]) -> None: + self.client = client + self._diagnostic = diagnostic + + def ensure_auth(self) -> None: + self.client.ensure_authenticated(open_browser=False, allow_login=False) + + +Requester = Callable[[Any, str, str], Awaitable[Response]] +Clock = Callable[[], datetime | float | int] +Sleep = Callable[[float], Awaitable[Any]] + + +class AccountCheckinService: + """Check and optionally sign each account without changing the active one. + + ``clock`` returns a datetime or Unix timestamp. ``random_delay`` returns a + delay in seconds in the midnight window. Both hooks, plus ``sleep`` and + ``requester``, exist so timer behavior can be tested without waiting. + """ + + def __init__( + self, + manager: Any, + *, + clock: Clock | None = None, + random_delay: Callable[[], float] | None = None, + sleep: Sleep | None = None, + requester: Requester | None = None, + max_concurrent: int = 4, + ) -> None: + if max_concurrent < 1: + raise ValueError("max_concurrent must be positive") + self.manager = manager + self._clock = clock or self._system_clock + self._random_delay = random_delay or (lambda: random.uniform(0.0, 600.0)) + self._sleep = sleep or asyncio.sleep + self._requester = requester or self._forward_request + self._semaphore = asyncio.Semaphore(max_concurrent) + + self._state_guard = threading.RLock() + self._states: dict[str, dict[str, Any]] = {} + self._account_locks: dict[str, asyncio.Lock] = {} + self._post_started: set[tuple[str, str]] = set() + self._logging_callback: Callable[..., Any] | None = None + + self._loop: asyncio.AbstractEventLoop | None = None + self._started = False + self._closed = False + self._stop_event = threading.Event() + self._scheduler_task: asyncio.Task[Any] | None = None + self._tasks: set[asyncio.Task[Any]] = set() + self._task_keys: dict[asyncio.Task[Any], tuple[str, str] | None] = {} + self._random_tasks: dict[tuple[str, str], asyncio.Task[Any]] = {} + + @staticmethod + def _system_clock() -> datetime: + return datetime.now(timezone.utc).astimezone(BEIJING) + + def _now(self) -> datetime: + value = self._clock() + if isinstance(value, datetime): + if value.tzinfo is None: + return value.replace(tzinfo=BEIJING) + return value.astimezone(BEIJING) + return datetime.fromtimestamp(float(value), BEIJING) + + def _epoch(self) -> float: + return self._now().timestamp() + + def _today(self) -> str: + return self._now().date().isoformat() + + def configure_logging(self, callback: Callable[..., Any] | None) -> None: + self._logging_callback = callback + + def _log(self, event: str, **fields: Any) -> None: + safe = { + key: value + for key, value in fields.items() + if key in { + "account_id", + "date", + "reason", + "status", + "http_status", + "next_checkin_at", + } + and (value is None or isinstance(value, (str, int, float, bool))) + } + try: + if self._logging_callback is not None: + self._logging_callback(event, **safe) + else: + logging.getLogger(__name__).info("%s %s", event, safe) + except Exception: + # Diagnostics must never change authentication or check-in behavior. + logging.getLogger(__name__).debug("check-in diagnostic callback failed") + + async def _forward_request(self, client: Any, method: str, path: str) -> Response: + state = _BillingAdapterState(client, self._log) + return await forward_billing(state, method, path, self._log) + + def snapshot(self, account_id: str) -> dict[str, Any]: + """Return the cached state only; never make an upstream request here.""" + + today = self._today() + with self._state_guard: + current = self._states.get(account_id) + if current is None or current["date"] != today: + return self._empty_state(today) + return dict(current) + + @staticmethod + def _empty_state(target_date: str) -> dict[str, Any]: + return { + "date": target_date, + "status": "unknown", + "checked": None, + "last_checked_at": None, + "last_checkin_at": None, + "error": None, + "next_checkin_at": None, + } + + def _state(self, account_id: str) -> dict[str, Any]: + today = self._today() + with self._state_guard: + current = self._states.get(account_id) + if current is None or current["date"] != today: + current = self._empty_state(today) + self._states[account_id] = current + return current + + def _account_lock(self, account_id: str) -> asyncio.Lock: + with self._state_guard: + lock = self._account_locks.get(account_id) + if lock is None: + lock = asyncio.Lock() + self._account_locks[account_id] = lock + return lock + + def _has_account(self, account_id: str) -> bool: + has_account = getattr(self.manager, "has_account", None) + if callable(has_account): + try: + return bool(has_account(account_id)) + except Exception: + return False + try: + return any( + isinstance(record, dict) and record.get("id") == account_id + for record in self.manager.store.list_accounts() + ) + except Exception: + return False + + def _account_ids(self) -> list[str]: + try: + records = self.manager.store.list_accounts() + except Exception: + return [] + result = [] + for record in records: + if isinstance(record, dict): + account_id = record.get("id") + if isinstance(account_id, str) and account_id: + result.append(account_id) + return result + + def _get_client(self, account_id: str) -> Any | None: + if not self._has_account(account_id): + return None + try: + return self.manager.store.get_client(account_id) + except Exception: + return None + + def _enabled(self) -> bool: + try: + return self.manager.get_settings().get("auto_checkin_enabled") is True + except Exception: + return False + + @staticmethod + def _reason(reason: str) -> str: + return reason if isinstance(reason, str) and reason in _REASONS else "manual" + + async def check_status(self, account_id: str, reason: str = "manual") -> dict[str, Any]: + async with self._account_lock(account_id): + return await self._check_status_locked(account_id, self._reason(reason)) + + async def _check_status_locked(self, account_id: str, reason: str) -> dict[str, Any]: + state = self._state(account_id) + client = self._get_client(account_id) + if client is None: + self._set_error(state, "账号不存在") + state["next_checkin_at"] = None + self._cancel_random_task(account_id, state["date"]) + self._log( + "account_checkin_status_skipped", + account_id=account_id, + date=state["date"], + reason=reason, + status="account_missing", + ) + return dict(state) + + previous = dict(state) + state.update(status="checking", checked=None, error=None) + self._log( + "account_checkin_status_check_started", + account_id=account_id, + date=state["date"], + reason=reason, + status="checking", + ) + try: + response = await self._request( + client, "POST", STATUS_PATH, account_id, reason + ) + except asyncio.CancelledError: + if state["date"] == self._today(): + stable = previous + if stable.get("status") in {"checking", "signing"}: + stable = self._empty_state(state["date"]) + state.clear() + state.update(stable) + state["next_checkin_at"] = None + raise + now = self._epoch() + state["last_checked_at"] = now + status_result = self._parse_status(response) + if status_result[0] == "checked": + state.update( + status="checked", checked=True, error=None, next_checkin_at=None + ) + self._cancel_random_task(account_id, state["date"]) + elif status_result[0] == "unchecked": + state.update(status="unchecked", checked=False, error=None) + elif status_result[0] == "inactive": + self._set_error(state, "签到活动未开放") + state["next_checkin_at"] = None + self._cancel_random_task(account_id, state["date"]) + else: + self._set_error(state, "签到状态查询失败") + self._log( + "account_checkin_status_checked", + account_id=account_id, + date=state["date"], + reason=reason, + status=state["status"], + http_status=response.status_code, + ) + return dict(state) + + async def check_in(self, account_id: str, reason: str = "manual") -> dict[str, Any]: + async with self._account_lock(account_id): + return await self._check_in_locked( + account_id, + self._reason(reason), + target_date=self._today(), + require_enabled=False, + ) + + async def _check_in_locked( + self, + account_id: str, + reason: str, + *, + target_date: str, + require_enabled: bool, + ) -> dict[str, Any]: + state = self._state(account_id) + if self._today() != target_date: + self._log( + "account_checkin_action_skipped", + account_id=account_id, + date=target_date, + reason=reason, + status="date_expired", + ) + return dict(state) + status = await self._check_status_locked(account_id, reason) + state = self._state(account_id) + if status["status"] != "unchecked": + self._log( + "account_checkin_action_skipped", + account_id=account_id, + date=target_date, + reason=reason, + status=status["status"], + ) + return dict(state) + if not self._before_post_allowed(account_id, target_date, require_enabled): + self._log( + "account_checkin_action_skipped", + account_id=account_id, + date=target_date, + reason=reason, + status="disabled_or_expired", + ) + return dict(state) + marker = (account_id, target_date) + with self._state_guard: + if marker in self._post_started: + self._set_error(state, "签到结果未确认") + self._log( + "account_checkin_action_skipped", + account_id=account_id, + date=target_date, + reason=reason, + status="already_attempted", + ) + return dict(state) + # Mark before yielding to the requester. A timeout or cancellation + # after this point must not send the same daily action twice. + self._post_started.add(marker) + + client = self._get_client(account_id) + if client is None: + with self._state_guard: + self._post_started.discard(marker) + self._set_error(state, "账号不存在") + return dict(state) + + state.update(status="signing", checked=False, error=None, next_checkin_at=None) + try: + response = await self._request( + client, + "POST", + CHECKIN_PATH, + account_id, + reason, + before_request=lambda: self._before_post_allowed( + account_id, target_date, require_enabled + ), + on_start=lambda: self._log( + "account_checkin_action_started", + account_id=account_id, + date=target_date, + reason=reason, + status="signing", + ), + ) + except asyncio.CancelledError: + if state["date"] == self._today(): + state.update( + status="unknown", + checked=None, + error=None, + next_checkin_at=None, + ) + raise + if response is None: + with self._state_guard: + self._post_started.discard(marker) + state.update(status="unchecked", checked=False, error=None) + self._log( + "account_checkin_action_skipped", + account_id=account_id, + date=target_date, + reason=reason, + status="disabled_or_expired", + ) + return dict(state) + post_kind, _ = self._parse_business_response(response) + if post_kind not in {"ok", "already"}: + if self._explicit_post_rejection(response): + confirmed = await self._check_status_locked(account_id, reason) + state = self._state(account_id) + if confirmed["status"] == "checked": + return dict(state) + if confirmed["status"] == "unchecked": + with self._state_guard: + self._post_started.discard(marker) + self._set_error(state, "签到请求失败") + self._log( + "account_checkin_failed", + account_id=account_id, + date=target_date, + reason=reason, + status="error", + http_status=response.status_code, + ) + return dict(state) + + # Even a successful POST must be confirmed by today's status endpoint; + # code 10001 is only an already-signed hint, never proof by itself. + confirmed = await self._check_status_locked(account_id, reason) + state = self._state(account_id) + if confirmed["status"] == "checked": + state["last_checkin_at"] = self._epoch() + self._log( + "account_checkin_confirmed", + account_id=account_id, + date=target_date, + reason=reason, + status="checked", + http_status=response.status_code, + ) + elif confirmed["status"] != "error": + self._set_error(state, "签到结果未确认") + return dict(state) + + def _before_post_allowed( + self, account_id: str, target_date: str, require_enabled: bool + ) -> bool: + return ( + not self._stop_event.is_set() + and self._today() == target_date + and self._has_account(account_id) + and (not require_enabled or self._enabled()) + ) + + @staticmethod + def _set_error(state: dict[str, Any], error: str) -> None: + state.update( + status="error", + checked=None, + error=error if error in _SAFE_ERRORS else "签到请求失败", + ) + + def _cancel_random_task(self, account_id: str, target_date: str) -> None: + key = (account_id, target_date) + task = self._random_tasks.pop(key, None) + if task is not None and task is not asyncio.current_task() and not task.done(): + task.cancel() + + @classmethod + def _explicit_post_rejection(cls, response: Response) -> bool: + if response.status_code == 429 or 400 <= response.status_code < 500: + return True + value = cls._response_json(response) + if value is None: + return False + code = value.get("code") + return type(code) is int and code not in {0, 10001} + + async def _request( + self, + client: Any, + method: str, + path: str, + account_id: str, + reason: str, + *, + before_request: Callable[[], bool] | None = None, + on_start: Callable[[], Any] | None = None, + ) -> Response | None: + async with self._semaphore: + if before_request is not None and not before_request(): + return None + if on_start is not None: + on_start() + try: + response = await self._requester(client, method, path) + except asyncio.CancelledError: + raise + except Exception: + self._log( + "account_checkin_upstream_failed", + account_id=account_id, + date=self._today(), + reason=reason, + status="error", + ) + response = Response(status_code=502, content=b"") + if response.status_code == 429: + await self._record_rate_limit(client, response.headers.get("Retry-After")) + return response + + async def _record_rate_limit(self, client: Any, retry_after: str | None) -> None: + handler = getattr(self.manager, "handle_rate_limit", None) + if not callable(handler): + return + try: + result = await asyncio.to_thread( + handler, client, retry_after, model=None, allow_switch=False + ) + if hasattr(result, "__await__"): + await result + except Exception: + self._log("account_checkin_rate_limit_record_failed") + + @staticmethod + def _response_json(response: Response) -> dict[str, Any] | None: + if not 200 <= int(response.status_code) < 300: + return None + try: + value = json.loads(response.body) + except (TypeError, ValueError, UnicodeDecodeError): + return None + if not isinstance(value, dict): + return None + # Some client wrappers retain the upstream envelope under this key; + # still validate its own business code and data strictly. + upstream = value.get("upstream") + if isinstance(upstream, dict): + value = upstream + return value + + @classmethod + def _parse_business_response(cls, response: Response) -> tuple[str, dict[str, Any] | None]: + value = cls._response_json(response) + if value is None: + return "failure", None + code = value.get("code") + if type(code) is int and code == 0: + data = value.get("data") + return ("ok", data if isinstance(data, dict) else None) + if type(code) is int and code == 10001: + return "already", None + return "failure", None + + @classmethod + def _parse_status(cls, response: Response) -> tuple[str, bool | None]: + value = cls._response_json(response) + if value is None or type(value.get("code")) is not int or value["code"] != 0: + return "unknown", None + data = value.get("data") + if not isinstance(data, dict): + return "unknown", None + if data.get("active") is False: + return "inactive", None + if "today_checked_in" in data: + result = data.get("today_checked_in") + else: + # ``checked`` is retained only for older, explicit boolean fixtures. + result = data.get("checked") + if type(result) is not bool: + return "unknown", None + return ("checked" if result else "unchecked"), result + + async def start(self) -> None: + """Register timers and return immediately; network work runs in tasks.""" + + if self._started or self._closed: + return + self._loop = asyncio.get_running_loop() + self._started = True + if self._enabled(): + self._schedule_catchup("startup") + self._scheduler_task = self._loop.create_task(self._scheduler_loop()) + + def settings_changed(self, enabled: bool) -> None: + """Apply a setting transition safely from either loop or worker thread.""" + + loop = self._loop + if loop is None or loop.is_closed() or not self._started or self._closed: + return + loop.call_soon_threadsafe(self._apply_settings_change, bool(enabled)) + + def _apply_settings_change(self, enabled: bool) -> None: + if self._closed or self._stop_event.is_set(): + return + # Callbacks can arrive out of order from concurrent settings writers. + # The persisted setting at execution time is authoritative. + enabled = self._enabled() + if not enabled: + for key, task in list(self._random_tasks.items()): + if not task.done(): + task.cancel() + self._random_tasks.pop(key, None) + account_id, target_date = key + state = self._state(account_id) + if state["date"] == target_date and state["status"] == "unchecked": + state["next_checkin_at"] = None + # Cancel other automatic work before its POST gate. A task whose + # POST has already begun is allowed to finish its single request. + for task in list(self._tasks): + key = self._task_keys.get(task) + if key is not None and key in self._post_started: + continue + if task is not asyncio.current_task() and not task.done(): + task.cancel() + return + self._schedule_catchup("enable") + + def _schedule_catchup(self, reason: str) -> None: + if not self._enabled() or self._stop_event.is_set() or self._loop is None: + return + target_date = self._today() + for account_id in self._account_ids(): + self._spawn( + self._automatic_checkin(account_id, target_date, reason), + key=(account_id, target_date), + ) + + async def _automatic_checkin( + self, account_id: str, target_date: str, reason: str + ) -> None: + async with self._account_lock(account_id): + if not self._before_post_allowed(account_id, target_date, True): + self._log( + "account_checkin_action_skipped", + account_id=account_id, + date=target_date, + reason=reason, + status="disabled_or_expired", + ) + return + await self._check_in_locked( + account_id, + self._reason(reason), + target_date=target_date, + require_enabled=True, + ) + + async def _prepare_midnight( + self, account_id: str, target_date: str, midnight: datetime + ) -> None: + if not self._before_post_allowed(account_id, target_date, True): + self._log( + "account_checkin_midnight_skipped", + account_id=account_id, + date=target_date, + reason="midnight", + status="disabled_or_expired", + ) + return + result = await self.check_status(account_id, reason="midnight") + if result["status"] != "unchecked": + self._log( + "account_checkin_midnight_skipped", + account_id=account_id, + date=target_date, + reason="midnight", + status=result["status"], + ) + return + if not self._before_post_allowed(account_id, target_date, True): + self._log( + "account_checkin_midnight_skipped", + account_id=account_id, + date=target_date, + reason="midnight", + status="disabled_or_expired", + ) + return + try: + offset = float(self._random_delay()) + except (TypeError, ValueError, OverflowError): + offset = 0.0 + if not (offset >= 0.0 and offset <= 600.0): + offset = min(600.0, max(0.0, offset)) if offset == offset else 0.0 + due = midnight + timedelta(seconds=offset) + due_epoch = due.timestamp() + state = self._state(account_id) + state["next_checkin_at"] = due_epoch + self._log( + "account_checkin_scheduled", + account_id=account_id, + date=target_date, + reason="midnight", + status="unchecked", + next_checkin_at=due_epoch, + ) + key = (account_id, target_date) + old = self._random_tasks.get(key) + if old is not None and not old.done(): + old.cancel() + task = self._spawn( + self._random_followup(account_id, target_date, due), key=key + ) + self._random_tasks[key] = task + + async def _random_followup( + self, account_id: str, target_date: str, due: datetime + ) -> None: + delay = max(0.0, due.timestamp() - self._epoch()) + await self._sleep(delay) + key = (account_id, target_date) + self._random_tasks.pop(key, None) + state = self._state(account_id) + if not self._before_post_allowed(account_id, target_date, True): + if state["date"] == target_date and state["status"] == "unchecked": + state["next_checkin_at"] = None + if self._today() != target_date: + skip_status = "date_expired" + elif not self._has_account(account_id): + skip_status = "account_missing" + else: + skip_status = "disabled_or_expired" + self._log( + "account_checkin_action_skipped", + account_id=account_id, + date=target_date, + reason="timer", + status=skip_status, + ) + return + state["next_checkin_at"] = None + await self._automatic_checkin(account_id, target_date, "timer") + + async def _run_midnight(self, target_date: date) -> None: + target = target_date.isoformat() + midnight = datetime.combine(target_date, time.min, tzinfo=BEIJING) + for account_id in self._account_ids(): + self._spawn( + self._prepare_midnight(account_id, target, midnight), + key=(account_id, target), + ) + + async def _scheduler_loop(self) -> None: + try: + while not self._stop_event.is_set(): + now = self._now() + next_day = now.date() + timedelta(days=1) + midnight = datetime.combine(next_day, time.min, tzinfo=BEIJING) + await self._sleep(max(0.0, (midnight - now).total_seconds())) + if self._stop_event.is_set(): + break + if self._enabled(): + await self._run_midnight(next_day) + except asyncio.CancelledError: + raise + + def _spawn( + self, + coroutine: Awaitable[Any], + *, + key: tuple[str, str] | None = None, + ) -> asyncio.Task[Any]: + if self._loop is None: + self._loop = asyncio.get_running_loop() + task = self._loop.create_task(coroutine) + self._tasks.add(task) + self._task_keys[task] = key + + def done(completed: asyncio.Task[Any]) -> None: + self._tasks.discard(completed) + task_key = self._task_keys.pop(completed, None) + if task_key is not None and self._random_tasks.get(task_key) is completed: + self._random_tasks.pop(task_key, None) + try: + completed.exception() + except (asyncio.CancelledError, Exception): + pass + + task.add_done_callback(done) + return task + + def stop(self) -> None: + """Synchronously signal shutdown and schedule cancellation on the loop.""" + + self._stop_event.set() + loop = self._loop + if loop is not None and not loop.is_closed(): + loop.call_soon_threadsafe(self._cancel_all_on_loop) + + def _cancel_all_on_loop(self) -> None: + if self._scheduler_task is not None and not self._scheduler_task.done(): + self._scheduler_task.cancel() + for task in list(self._tasks): + if not task.done(): + task.cancel() + self._random_tasks.clear() + + async def close(self) -> None: + if self._closed: + return + self._closed = True + self.stop() + current = asyncio.current_task() + tasks = [task for task in self._tasks if task is not current] + if self._scheduler_task is not None and self._scheduler_task is not current: + tasks.append(self._scheduler_task) + for task in tasks: + if not task.done(): + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) diff --git a/src/codebuddy_proxy/account_manager.py b/src/codebuddy_proxy/account_manager.py new file mode 100644 index 0000000..94e937b --- /dev/null +++ b/src/codebuddy_proxy/account_manager.py @@ -0,0 +1,947 @@ +"""Account switching and bounded background login management. + +Login credentials are first written to a private temporary session file. The +temporary client is imported into the account store only after authentication +has completed and the job still owns the commit boundary. +""" + +from __future__ import annotations + +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass, field +from datetime import timezone +from email.utils import parsedate_to_datetime +import ipaddress +import json +import math +import pathlib +import re +import tempfile +import threading +import time +import unicodedata +import uuid +from typing import Any, Callable +from urllib.parse import urlsplit + +from codebuddy_proxy.account_store import _atomic_json +from codebuddy_proxy.backend_profile import profile_for_key, resolve_backend +from codebuddy_proxy.codebuddy_client_demo import CodeBuddyClient, CodeBuddyError + + +class AccountManagerError(CodeBuddyError): + """An expected account-management error safe to show to a local user.""" + + def __init__(self, message: str, *, status_code: int = 400) -> None: + super().__init__(message) + self.message = message + self.status_code = status_code + + +_TERMINAL = frozenset({"succeeded", "failed", "cancelled"}) +_PUBLIC_JOB_FIELDS = ("id", "backend", "status", "auth_url", "error", "account_id") +_PUBLIC_ACCOUNT_FIELDS = ( + "id", + "backend", + "endpoint", + "uid", + "nickname", + "active", + "model_cooldowns", + "checkin", +) +_RUNTIME_STATE_VERSION = 2 +_LEGACY_RUNTIME_STATE_VERSION = 1 +_RUNTIME_STATE_FILENAME = "account-manager-state-v1.json" +_DEFAULT_COOLDOWN_SECONDS = 60 +_RETRY_AFTER_SECONDS_RE = re.compile(r"^[0-9]+$") + + +@dataclass(slots=True) +class _LoginJob: + id: str + backend: str + created_at: float + status: str = "pending" + auth_url: str | None = None + error: str | None = None + account_id: str | None = None + updated_at: float = field(default_factory=time.monotonic) + cancel_event: threading.Event = field(default_factory=threading.Event) + lock: threading.RLock = field(default_factory=threading.RLock) + future: Future[Any] | None = None + + +class AccountManager: + """Coordinate the account store and bounded, cancellable login jobs.""" + + def __init__( + self, + store: Any, + *, + max_concurrent_logins: int = 2, + retained_job_limit: int = 100, + job_ttl_seconds: float = 15 * 60, + login_timeout: int = 300, + client_factory: Callable[[str, pathlib.Path], Any] | None = None, + ) -> None: + if max_concurrent_logins < 1: + raise ValueError("max_concurrent_logins must be positive") + if retained_job_limit < 1: + raise ValueError("retained_job_limit must be positive") + self.store = store + self._max_concurrent_logins = max_concurrent_logins + self._retained_job_limit = retained_job_limit + self._job_ttl_seconds = max(0.0, job_ttl_seconds) + self._login_timeout = max(1, int(login_timeout)) + self._client_factory = client_factory or self._create_login_client + self._lock = threading.RLock() + self._store_lock = threading.RLock() + self._runtime_lock = threading.RLock() + self._runtime_state_path = self._find_runtime_state_path() + self._settings = { + "auto_switch_on_429": False, + "default_cooldown_seconds": _DEFAULT_COOLDOWN_SECONDS, + "auto_checkin_enabled": False, + } + self._cooldowns: dict[str, dict[str, dict[str, Any]]] = {} + self._load_runtime_state() + self._closed = False + self._jobs: dict[str, _LoginJob] = {} + self._executor = ThreadPoolExecutor( + max_workers=max_concurrent_logins, + thread_name_prefix="codebuddy-login", + ) + from codebuddy_proxy.account_checkin import AccountCheckinService + + self.checkins = AccountCheckinService(self) + + @property + def client(self) -> Any | None: + """Return the currently active client, if the store has one.""" + + with self._store_lock: + return self.store.active_client() + + def snapshot(self) -> dict[str, Any]: + """Return safe account metadata, cooldowns, settings, and server time.""" + + with self._store_lock: + accounts = self.store.list_accounts() + active_id = self.store.active_account_id + server_time = time.time() + with self._runtime_lock: + safe_accounts = [] + for account in accounts: + if not isinstance(account, dict): + continue + safe = { + key: account[key] + for key in _PUBLIC_ACCOUNT_FIELDS[:6] + if key in account + } + account_id = account.get("id") + if not isinstance(account_id, str): + account_id = None + cooldowns = self._cooldowns.get(account_id, {}) + model_cooldowns = [] + for model, cooldown in sorted(cooldowns.items()): + deadline = cooldown.get("deadline") + active_cooldown = ( + isinstance(deadline, (int, float)) + and deadline > server_time + ) + model_cooldowns.append( + { + "model": model, + "rate_limited": active_cooldown, + # Keep the original deadline and source as + # history after expiry; only the active flag + # and remaining duration change. + "cooldown_until": deadline, + "cooldown_remaining_seconds": ( + max(0, math.ceil(deadline - server_time)) + if active_cooldown + else 0 + ), + "last_rate_limited_at": cooldown.get( + "last_rate_limited_at" + ), + "cooldown_source": cooldown.get("source"), + } + ) + safe.update( + { + "model_cooldowns": model_cooldowns, + } + ) + safe_accounts.append(safe) + settings = dict(self._settings) + for account in safe_accounts: + account_id = account.get("id") + account["checkin"] = self.checkins.snapshot( + account_id if isinstance(account_id, str) else "" + ) + return { + "accounts": safe_accounts, + "active_account_id": active_id, + "settings": settings, + "server_time": server_time, + } + + def activate(self, account_id: str) -> None: + """Make an existing account active and persist the choice in the store.""" + + with self._store_lock: + if not self._has_account_locked(account_id): + raise AccountManagerError("未找到该账号", status_code=404) + try: + self.store.set_active(account_id) + except (KeyError, FileNotFoundError): + raise AccountManagerError("未找到该账号", status_code=404) from None + + def has_account(self, account_id: str) -> bool: + """Return whether an account id is present in the persistent store.""" + + with self._store_lock: + return self._has_account_locked(account_id) + + def account_id_for_client(self, client: Any) -> str | None: + """Find a stored account by its backend, endpoint, and uid identity.""" + + identity = _client_identity(client) + if identity is None: + return None + backend, endpoint, uid = identity + with self._store_lock: + return self._account_id_for_identity_locked(backend, endpoint, uid) + + def handle_rate_limit( + self, + client: Any, + retry_after: Any, + *, + model: str | None = None, + excluded_account_ids: set[str] | None = None, + allow_switch: bool = True, + ) -> Any | None: + """Record a model-scoped 429 and optionally switch within the same region.""" + + model_key = _valid_model_name(model) + if model_key is None: + # Billing and other model-less endpoints must not affect chat + # failover or impose a cooldown on every model for this account. + return None + identity = _client_identity(client) + if identity is None: + return None + backend, endpoint, uid = identity + with self._store_lock: + account_id = self._account_id_for_identity_locked(backend, endpoint, uid) + if account_id is None: + return None + now = time.time() + delay, source = _cooldown_delay(retry_after, now) + deadline = now + delay + with self._runtime_lock: + account_cooldowns = self._cooldowns.get(account_id) + had_account_cooldowns = account_cooldowns is not None + if account_cooldowns is None: + account_cooldowns = {} + self._cooldowns[account_id] = account_cooldowns + had_previous = model_key in account_cooldowns + previous = account_cooldowns.get(model_key) + if previous is not None and previous.get("deadline", 0) > now: + if previous["deadline"] > deadline: + deadline = previous["deadline"] + source = previous["source"] + account_cooldowns[model_key] = { + "deadline": deadline, + "last_rate_limited_at": now, + "source": source, + } + try: + self._save_runtime_state_locked() + except Exception: + if had_previous: + account_cooldowns[model_key] = previous + else: + account_cooldowns.pop(model_key, None) + if not account_cooldowns and not had_account_cooldowns: + self._cooldowns.pop(account_id, None) + raise + + excluded = { + item for item in (excluded_account_ids or set()) if isinstance(item, str) + } + excluded.add(account_id) + with self._store_lock: + with self._runtime_lock: + if ( + not allow_switch + or not self._settings["auto_switch_on_429"] + or self.store.active_account_id != account_id + ): + return None + accounts = self.store.list_accounts() + for candidate in accounts: + if not isinstance(candidate, dict): + continue + candidate_id = candidate.get("id") + if ( + not isinstance(candidate_id, str) + or candidate_id in excluded + or candidate.get("backend") != backend + or candidate.get("endpoint") != endpoint + ): + continue + cooldown = self._cooldowns.get(candidate_id, {}).get(model_key) + if ( + cooldown is not None + and cooldown.get("deadline", 0) > time.time() + ): + continue + try: + candidate_client = self.store.get_client(candidate_id) + except Exception: + # A broken/missing session is not an eligible failover. + continue + if not _client_has_usable_auth(candidate_client): + continue + # This is serialized with explicit activate() calls under + # _store_lock; a manual selection can never be displaced + # by a delayed response from a previously active client. + return self.store.set_active(candidate_id) + return None + + def get_settings(self) -> dict[str, Any]: + with self._runtime_lock: + return dict(self._settings) + + def set_settings( + self, + auto_switch_on_429: bool | None = None, + *, + auto_checkin_enabled: bool | None = None, + ) -> dict[str, Any]: + if auto_switch_on_429 is None and auto_checkin_enabled is None: + raise AccountManagerError("至少需要设置一个账号管理开关", status_code=400) + if auto_switch_on_429 is not None and type(auto_switch_on_429) is not bool: + raise AccountManagerError("auto_switch_on_429 必须为布尔值", status_code=400) + if auto_checkin_enabled is not None and type(auto_checkin_enabled) is not bool: + raise AccountManagerError("auto_checkin_enabled 必须为布尔值", status_code=400) + with self._runtime_lock: + previous = dict(self._settings) + if auto_switch_on_429 is not None: + self._settings["auto_switch_on_429"] = auto_switch_on_429 + if auto_checkin_enabled is not None: + self._settings["auto_checkin_enabled"] = auto_checkin_enabled + try: + self._save_runtime_state_locked() + except Exception: + self._settings = previous + raise + updated = dict(self._settings) + checkin_transition = ( + previous["auto_checkin_enabled"] + != updated["auto_checkin_enabled"] + ) + if checkin_transition: + self.checkins.settings_changed(updated["auto_checkin_enabled"]) + return updated + + def delete(self, account_id: str) -> None: + """Delete an inactive account.""" + + with self._store_lock: + if not self._has_account_locked(account_id): + raise AccountManagerError("未找到该账号", status_code=404) + if self.store.active_account_id == account_id: + raise AccountManagerError("请先切换到其他账号,再删除此账号", status_code=409) + + # Remove persisted runtime state first. If this write fails, the + # account remains untouched; otherwise a deleted identity could + # retain a cooldown on disk and inherit it if re-imported later. + removed_cooldown: dict[str, Any] | None = None + with self._runtime_lock: + removed_cooldown = self._cooldowns.pop(account_id, None) + if removed_cooldown is not None: + try: + self._save_runtime_state_locked() + except Exception: + self._cooldowns[account_id] = removed_cooldown + raise + + try: + self.store.delete_account(account_id) + except (KeyError, FileNotFoundError): + self._restore_deleted_account_cooldown(account_id, removed_cooldown) + raise AccountManagerError("未找到该账号", status_code=404) from None + except Exception: + self._restore_deleted_account_cooldown(account_id, removed_cooldown) + raise + + def _restore_deleted_account_cooldown( + self, account_id: str, cooldown: dict[str, Any] | None + ) -> None: + if cooldown is None: + return + with self._runtime_lock: + self._cooldowns[account_id] = cooldown + self._save_runtime_state_locked() + + def _account_id_for_identity_locked( + self, backend: str, endpoint: str, uid: str + ) -> str | None: + for account in self.store.list_accounts(): + if ( + isinstance(account, dict) + and account.get("backend") == backend + and account.get("endpoint") == endpoint + and str(account.get("uid", "")).strip() == uid + ): + account_id = account.get("id") + if isinstance(account_id, str): + return account_id + return None + + def _find_runtime_state_path(self) -> pathlib.Path | None: + root = getattr(self.store, "root", None) + if root is None: + return None + try: + return pathlib.Path(root).expanduser() / _RUNTIME_STATE_FILENAME + except (TypeError, ValueError, OSError) as exc: + raise AccountManagerError("账号管理状态目录无效,未修改任何数据", status_code=503) from exc + + def _load_runtime_state(self) -> None: + path = self._runtime_state_path + if path is None: + return + if path.is_symlink(): + raise AccountManagerError("账号管理状态文件无效,未修改任何数据", status_code=503) + try: + value = json.loads( + path.read_text(encoding="utf-8"), + object_pairs_hook=_unique_runtime_object, + ) + except FileNotFoundError: + return + except ( + OSError, + json.JSONDecodeError, + UnicodeDecodeError, + _DuplicateRuntimeStateKey, + RecursionError, + ): + raise AccountManagerError("账号管理状态文件损坏,未修改任何数据", status_code=503) from None + if not isinstance(value, dict) or set(value) != {"version", "settings", "cooldowns"}: + raise AccountManagerError("账号管理状态文件损坏,未修改任何数据", status_code=503) + version = value.get("version") + if type(version) is not int or version not in ( + _LEGACY_RUNTIME_STATE_VERSION, + _RUNTIME_STATE_VERSION, + ): + raise AccountManagerError("账号管理状态文件版本不受支持,未修改任何数据", status_code=503) + settings = value.get("settings") + setting_fields = set(settings) if isinstance(settings, dict) else set() + legacy_fields = {"auto_switch_on_429", "default_cooldown_seconds"} + current_fields = legacy_fields | {"auto_checkin_enabled"} + if ( + not isinstance(settings, dict) + or setting_fields not in ( + (legacy_fields, current_fields) + if version == _LEGACY_RUNTIME_STATE_VERSION + else (current_fields,) + ) + or type(settings.get("auto_switch_on_429")) is not bool + or type(settings.get("default_cooldown_seconds")) is not int + or settings.get("default_cooldown_seconds") != _DEFAULT_COOLDOWN_SECONDS + or ( + "auto_checkin_enabled" in settings + and type(settings.get("auto_checkin_enabled")) is not bool + ) + ): + raise AccountManagerError("账号管理状态文件损坏,未修改任何数据", status_code=503) + cooldowns = value.get("cooldowns") + if not isinstance(cooldowns, dict): + raise AccountManagerError("账号管理状态文件损坏,未修改任何数据", status_code=503) + loaded: dict[str, dict[str, dict[str, Any]]] = {} + if version == _LEGACY_RUNTIME_STATE_VERSION: + # A legacy account-wide cooldown cannot safely be assigned to a + # particular model. Validate its full shape, then discard it. + for account_id, record in cooldowns.items(): + if ( + not isinstance(account_id, str) + or not account_id + or len(account_id) > 256 + or not isinstance(record, dict) + or set(record) != {"deadline", "last_rate_limited_at", "source"} + or not _finite_number(record.get("deadline")) + or not _finite_number(record.get("last_rate_limited_at")) + or record.get("source") not in ("retry_after", "default") + ): + raise AccountManagerError( + "账号管理状态文件损坏,未修改任何数据", status_code=503 + ) + else: + for account_id, model_records in cooldowns.items(): + if ( + not isinstance(account_id, str) + or not account_id + or len(account_id) > 256 + or not isinstance(model_records, dict) + ): + raise AccountManagerError( + "账号管理状态文件损坏,未修改任何数据", status_code=503 + ) + loaded_records: dict[str, dict[str, Any]] = {} + for model, record in model_records.items(): + if ( + _valid_model_name(model) != model + or not isinstance(record, dict) + or set(record) + != {"deadline", "last_rate_limited_at", "source"} + or not _finite_number(record.get("deadline")) + or not _finite_number(record.get("last_rate_limited_at")) + or record.get("source") not in ("retry_after", "default") + ): + raise AccountManagerError( + "账号管理状态文件损坏,未修改任何数据", status_code=503 + ) + loaded_records[model] = { + "deadline": float(record["deadline"]), + "last_rate_limited_at": float(record["last_rate_limited_at"]), + "source": record["source"], + } + if loaded_records: + loaded[account_id] = loaded_records + self._settings = { + "auto_switch_on_429": settings["auto_switch_on_429"], + "default_cooldown_seconds": _DEFAULT_COOLDOWN_SECONDS, + "auto_checkin_enabled": settings.get("auto_checkin_enabled", False), + } + self._cooldowns = loaded + + def _save_runtime_state_locked(self) -> None: + path = self._runtime_state_path + if path is None: + return + payload = { + "version": _RUNTIME_STATE_VERSION, + "settings": dict(self._settings), + "cooldowns": { + account_id: { + model: dict(record) + for model, record in model_records.items() + } + for account_id, model_records in self._cooldowns.items() + if model_records + }, + } + try: + _atomic_json(path, payload) + except Exception: + # Keep private paths and low-level storage details out of manager + # responses. The caller sees only a fixed error message. + raise AccountManagerError("无法保存账号管理状态", status_code=500) from None + + + def start_login(self, backend: str) -> dict[str, Any]: + """Start a login flow and return its opaque job identifier.""" + + if backend not in {"domestic", "global"}: + raise AccountManagerError("backend 必须为 domestic 或 global", status_code=400) + + now = time.monotonic() + with self._lock: + if self._closed: + raise AccountManagerError("账号管理器已关闭", status_code=503) + self._reap_jobs_locked(now) + active_workers = sum( + job.future is None or not job.future.done() + for job in self._jobs.values() + ) + if active_workers >= self._max_concurrent_logins: + raise AccountManagerError("当前登录任务较多,请稍后重试", status_code=429) + + job = _LoginJob(id=uuid.uuid4().hex, backend=backend, created_at=now) + self._jobs[job.id] = job + try: + job.future = self._executor.submit(self._run_login, job) + except RuntimeError: + self._jobs.pop(job.id, None) + raise AccountManagerError("账号管理器已关闭", status_code=503) from None + return self._public_job(job) + + def get_login(self, job_id: str) -> dict[str, Any]: + job = self._find_job(job_id) + return self._public_job(job) + + def cancel_login(self, job_id: str) -> dict[str, Any]: + job = self._find_job(job_id) + with job.lock: + if job.status == "pending": + # This transition is the cancellation linearization point. A + # worker that finishes after it cannot import into the store. + job.cancel_event.set() + job.status = "cancelled" + job.error = None + job.updated_at = time.monotonic() + return self._public_job_unlocked(job) + + def close(self) -> None: + """Cancel outstanding jobs and stop the bounded worker pool.""" + + checkins = getattr(self, "checkins", None) + if checkins is not None: + checkins.stop() + with self._lock: + if self._closed: + return + # The manager lock also guards the commit boundary. If a worker + # already owns it, its commit finishes first; otherwise close wins. + self._closed = True + jobs = list(self._jobs.values()) + + for job in jobs: + with job.lock: + if job.status == "pending": + job.cancel_event.set() + job.status = "cancelled" + job.error = None + job.updated_at = time.monotonic() + + self._executor.shutdown(wait=True, cancel_futures=True) + + def _has_account_locked(self, account_id: str) -> bool: + return any( + isinstance(account, dict) and account.get("id") == account_id + for account in self.store.list_accounts() + ) + + def _create_login_client(self, backend: str, session_file: pathlib.Path) -> CodeBuddyClient: + profile = profile_for_key(backend) + resolved = resolve_backend( + profile=profile, + session_file=session_file, + environment={}, + ) + return CodeBuddyClient(backend=resolved) + + def _temporary_directory(self) -> tempfile.TemporaryDirectory[str]: + root = getattr(self.store, "root", None) + temp_parent: str | None = None + if root is not None: + try: + root_path = pathlib.Path(root).expanduser() + if root_path.is_dir(): + temp_parent = str(root_path) + except (OSError, TypeError, ValueError): + temp_parent = None + return tempfile.TemporaryDirectory(prefix="workbuddy-login-", dir=temp_parent) + + def _run_login(self, job: _LoginJob) -> None: + with job.lock: + if job.status != "pending" or job.cancel_event.is_set(): + return + + try: + with self._temporary_directory() as temp_directory: + session_file = pathlib.Path(temp_directory) / "session.json" + client = self._client_factory(job.backend, session_file) + with job.lock: + if job.status != "pending" or job.cancel_event.is_set(): + return + challenge = client.begin_login() + auth_url, state = self._read_challenge(challenge) + + with job.lock: + if job.status != "pending" or job.cancel_event.is_set(): + return + job.auth_url = auth_url + job.updated_at = time.monotonic() + + client.complete_login( + {"auth_url": auth_url, "state": state}, + timeout=self._login_timeout, + cancel_event=job.cancel_event, + ) + + session = getattr(client, "session", None) + account = session.get("account") if isinstance(session, dict) else None + if not isinstance(account, dict) or not account.get("uid"): + raise ValueError("login did not produce a valid account") + + self._commit_login(job, client) + except Exception as exc: + self._finish_failure(job, exc) + + @staticmethod + def _read_challenge(challenge: Any) -> tuple[str, str]: + if not isinstance(challenge, dict): + raise ValueError("invalid login challenge") + auth_url = challenge.get("auth_url") + state = challenge.get("state") + if ( + not isinstance(auth_url, str) + or len(auth_url) > 4096 + or any(ord(char) < 0x21 or ord(char) == 0x7F for char in auth_url) + ): + raise ValueError("invalid login URL") + parsed = urlsplit(auth_url) + if parsed.scheme.lower() not in {"http", "https"} or not parsed.hostname: + raise ValueError("invalid login URL") + if parsed.username is not None or parsed.password is not None: + raise ValueError("invalid login URL") + port = parsed.port + if port is not None and not 1 <= port <= 65535: + raise ValueError("invalid login URL") + if not isinstance(state, str) or not state or len(state) > 4096: + raise ValueError("invalid login state") + return auth_url, state + + def _commit_login(self, job: _LoginJob, client: Any) -> None: + with job.lock: + if job.status != "pending" or job.cancel_event.is_set(): + return + # Holding the manager lock through import gives close() and + # cancellation a deterministic boundary relative to persistence. + with self._lock: + if self._closed or job.cancel_event.is_set(): + self._mark_cancelled_unlocked(job) + return + with self._store_lock: + should_activate = not bool(self.store.list_accounts()) + account_id = self.store.import_client(client, activate=should_activate) + if not isinstance(account_id, str) or not account_id: + raise ValueError("account store returned an invalid identifier") + job.account_id = account_id + job.status = "succeeded" + job.error = None + job.updated_at = time.monotonic() + + def _finish_failure(self, job: _LoginJob, exc: Exception) -> None: + with job.lock: + if job.status != "pending": + return + if job.cancel_event.is_set(): + self._mark_cancelled_unlocked(job) + return + if isinstance(exc, TimeoutError): + message = "登录超时,请重新尝试" + elif isinstance(exc, (ValueError, TypeError)): + message = "登录信息无效,请重新尝试" + else: + # Do not expose provider text, paths, response bodies, or + # exception representations through the local management API. + message = "登录失败,请稍后重新尝试" + job.status = "failed" + job.error = message + job.updated_at = time.monotonic() + + @staticmethod + def _mark_cancelled_unlocked(job: _LoginJob) -> None: + job.cancel_event.set() + job.status = "cancelled" + job.error = None + job.updated_at = time.monotonic() + + def _find_job(self, job_id: str) -> _LoginJob: + with self._lock: + job = self._jobs.get(job_id) + if job is None: + raise AccountManagerError("未找到该登录任务", status_code=404) + return job + + def _reap_jobs_locked(self, now: float) -> None: + # Job state is monotonic: pending can only become terminal. Read-only + # status/timestamp inspection here is safe while the registry lock is + # held and avoids inverting the per-job -> registry lock order. + for job_id, job in list(self._jobs.items()): + future_done = job.future is not None and job.future.done() + if ( + job.status in _TERMINAL + and future_done + and now - job.updated_at >= self._job_ttl_seconds + ): + self._jobs.pop(job_id, None) + + terminal = sorted( + ( + job + for job in self._jobs.values() + if job.status in _TERMINAL + and job.future is not None + and job.future.done() + ), + key=lambda job: job.updated_at, + ) + while len(terminal) > self._retained_job_limit: + self._jobs.pop(terminal.pop(0).id, None) + + max_jobs = self._retained_job_limit + self._max_concurrent_logins + while len(self._jobs) >= max_jobs: + oldest = next((job for job in terminal if job.id in self._jobs), None) + if oldest is None: + break + self._jobs.pop(oldest.id, None) + terminal.remove(oldest) + + def _public_job(self, job: _LoginJob) -> dict[str, Any]: + with job.lock: + return self._public_job_unlocked(job) + + @staticmethod + def _public_job_unlocked(job: _LoginJob) -> dict[str, Any]: + values = { + "id": job.id, + "backend": job.backend, + "status": job.status, + "auth_url": job.auth_url, + "error": job.error, + "account_id": job.account_id, + } + return {key: values[key] for key in _PUBLIC_JOB_FIELDS} + + +class _DuplicateRuntimeStateKey(ValueError): + pass + + +def _unique_runtime_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + value: dict[str, Any] = {} + for key, item in pairs: + if key in value: + raise _DuplicateRuntimeStateKey("duplicate runtime state key") + value[key] = item + return value + + +def _finite_number(value: Any) -> bool: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return False + try: + return math.isfinite(float(value)) + except (OverflowError, TypeError, ValueError): + return False + + +def _valid_model_name(model: Any) -> str | None: + """Return a model identifier only when it is safe to key and persist.""" + + if not isinstance(model, str) or not model or len(model) > 256: + return None + if not model.strip() or any( + unicodedata.category(char) in {"Cc", "Cs"} for char in model + ): + return None + return model + + +def _client_identity(client: Any) -> tuple[str, str, str] | None: + if client is None: + return None + profile = getattr(client, "profile", None) + backend = getattr(profile, "key", None) + if not isinstance(backend, str): + backend_value = getattr(client, "backend", None) + backend = ( + backend_value + if isinstance(backend_value, str) + else getattr(getattr(backend_value, "profile", None), "key", None) + ) + endpoint = getattr(client, "endpoint", None) + try: + session = _client_session(client) + except Exception: + return None + account = session.get("account") if isinstance(session, dict) else None + uid = account.get("uid") if isinstance(account, dict) else None + if ( + not isinstance(backend, str) + or not backend + or not isinstance(endpoint, str) + or not endpoint + or uid is None + ): + return None + uid_text = str(uid).strip() + if not uid_text: + return None + return backend, endpoint, uid_text + + +def _client_session(client: Any) -> dict[str, Any]: + snapshot = getattr(client, "_snapshot_session", None) + if callable(snapshot): + value = snapshot() + if isinstance(value, tuple) and value: + value = value[0] + else: + value = getattr(client, "session", None) + return value if isinstance(value, dict) else {} + + +def _client_has_usable_auth(client: Any) -> bool: + try: + session = _client_session(client) + except Exception: + return False + auth = session.get("auth") + if not isinstance(auth, dict): + return False + access_token = auth.get("accessToken") + if isinstance(access_token, str) and access_token.strip(): + expires_at = auth.get("expiresAt") + try: + expires_ms = int(expires_at or 0) + except (TypeError, ValueError, OverflowError): + expires_ms = 1 + if not expires_ms or expires_ms > int(time.time() * 1000) + 60_000: + return True + refresh_token = auth.get("refreshToken") + return isinstance(refresh_token, str) and bool(refresh_token.strip()) + + +def _cooldown_delay(retry_after: Any, now: float) -> tuple[float, str]: + seconds: float | None = None + if isinstance(retry_after, int) and not isinstance(retry_after, bool) and retry_after >= 0: + try: + seconds = float(retry_after) + except OverflowError: + seconds = None + elif isinstance(retry_after, str): + value = retry_after.strip() + if _RETRY_AFTER_SECONDS_RE.fullmatch(value): + try: + seconds = float(int(value)) + except (OverflowError, ValueError): + seconds = None + else: + try: + parsed = parsedate_to_datetime(value) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + target = parsed.astimezone(timezone.utc).timestamp() + if math.isfinite(target): + seconds = max(0.0, target - now) + except (AttributeError, TypeError, ValueError, OverflowError, OSError): + seconds = None + if seconds is None or not math.isfinite(seconds): + return float(_DEFAULT_COOLDOWN_SECONDS), "default" + return seconds, "retry_after" + + +def is_loopback_address(value: str | None) -> bool: + """Return whether a peer address is a literal loopback IP address.""" + + if not value: + return False + try: + address = ipaddress.ip_address(value) + if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped is not None: + return address.ipv4_mapped.is_loopback + return address.is_loopback + except ValueError: + return False diff --git a/src/codebuddy_proxy/account_store.py b/src/codebuddy_proxy/account_store.py new file mode 100644 index 0000000..6438b81 --- /dev/null +++ b/src/codebuddy_proxy/account_store.py @@ -0,0 +1,513 @@ +"""Persistent, credential-isolated storage for CodeBuddy accounts.""" + +from __future__ import annotations + +import hashlib +import json +import os +import pathlib +import re +import tempfile +import threading +from contextlib import nullcontext +from typing import Any +from urllib.parse import urlsplit + +from codebuddy_proxy.backend_profile import ( + normalize_endpoint, + profile_for_key, + resolve_backend, +) +from codebuddy_proxy.codebuddy_client_demo import CodeBuddyClient, CodeBuddyError + + +_INDEX_VERSION = 1 +_ACCOUNT_ID_RE = re.compile(r"^acct_[0-9a-f]{64}$") +_INDEX_FIELDS = {"id", "backend", "endpoint", "uid", "nickname"} +_ROOT_LOCKS: dict[str, Any] = {} +_ROOT_LOCKS_GUARD = threading.Lock() + + +def _root_lock(root: pathlib.Path): + key = os.fspath(root) + with _ROOT_LOCKS_GUARD: + lock = _ROOT_LOCKS.get(key) + if lock is None: + lock = threading.RLock() + _ROOT_LOCKS[key] = lock + return lock + + +def _safe_endpoint(value: str) -> str: + try: + endpoint = normalize_endpoint(value) + parsed = urlsplit(endpoint) + except (TypeError, ValueError) as exc: + raise CodeBuddyError("账号 endpoint 无效") from exc + if parsed.username or parsed.password or parsed.query or parsed.fragment: + raise CodeBuddyError("账号 endpoint 不得在地址中包含凭据或查询参数") + return endpoint + + +def _account_id(backend: str, endpoint: str, uid: str) -> str: + identity = json.dumps( + [backend, endpoint, uid], ensure_ascii=False, separators=(",", ":") + ).encode("utf-8") + return f"acct_{hashlib.sha256(identity).hexdigest()}" + + +def _has_auth_token(auth: Any) -> bool: + return isinstance(auth, dict) and any( + isinstance(auth.get(name), str) and bool(auth[name].strip()) + for name in ("accessToken", "refreshToken") + ) + + +def _atomic_json(path: pathlib.Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) + try: + os.chmod(path.parent, 0o700) + except OSError: + pass + temp_name: str | None = None + try: + fd, temp_name = tempfile.mkstemp( + prefix=f".{path.name}.", suffix=".tmp", dir=path.parent + ) + with os.fdopen(fd, "w", encoding="utf-8") as stream: + json.dump(payload, stream, ensure_ascii=False, indent=2) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temp_name, path) + temp_name = None + try: + os.chmod(path, 0o600) + except OSError: + pass + except OSError as exc: + raise CodeBuddyError(f"无法原子写入账号文件: {path.name}") from exc + finally: + if temp_name is not None: + try: + os.unlink(temp_name) + except OSError: + pass + + +def _capture_file(path: pathlib.Path) -> bytes | None: + if path.is_symlink(): + raise CodeBuddyError("账号文件路径无效") + try: + return path.read_bytes() + except FileNotFoundError: + return None + except OSError as exc: + raise CodeBuddyError("无法读取账号文件进行安全更新") from exc + + +def _restore_file(path: pathlib.Path, contents: bytes | None) -> None: + if path.is_symlink(): + raise CodeBuddyError("账号文件路径无效") + if contents is None: + try: + path.unlink() + except FileNotFoundError: + pass + except OSError as exc: + raise CodeBuddyError("无法清理账号事务文件") from exc + return + temp_name: str | None = None + try: + fd, temp_name = tempfile.mkstemp( + prefix=f".{path.name}.", suffix=".rollback.tmp", dir=path.parent + ) + with os.fdopen(fd, "wb") as stream: + stream.write(contents) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temp_name, path) + temp_name = None + try: + os.chmod(path, 0o600) + except OSError: + pass + except OSError as exc: + raise CodeBuddyError("无法恢复账号文件") from exc + finally: + if temp_name is not None: + try: + os.unlink(temp_name) + except OSError: + pass + + +class AccountStore: + """Store account metadata separately from each credential-bearing session.""" + + def __init__(self, root: pathlib.Path | str | None = None) -> None: + self.root = pathlib.Path(root or (pathlib.Path.home() / ".workbuddy2api" / "accounts")) + self.root = self.root.expanduser().resolve() + self.index_file = self.root / "index.json" + self._lock = _root_lock(self.root) + self._clients: dict[str, CodeBuddyClient] = {} + with self._lock: + try: + self.root.mkdir(parents=True, exist_ok=True, mode=0o700) + os.chmod(self.root, 0o700) + except OSError as exc: + raise CodeBuddyError("无法创建私有账号目录") from exc + if self.index_file.is_symlink(): + raise CodeBuddyError("账号索引文件路径无效") + self._read_index_locked() + + def _session_file(self, account_id: str) -> pathlib.Path: + if not _ACCOUNT_ID_RE.fullmatch(account_id): + raise CodeBuddyError("账号 id 无效") + return self.root / f"account-{account_id}.json" + + def _read_index_locked(self) -> dict[str, Any]: + try: + raw = self.index_file.read_text(encoding="utf-8") + except FileNotFoundError: + return {"version": _INDEX_VERSION, "active_account_id": None, "accounts": []} + except OSError as exc: + raise CodeBuddyError("无法读取账号索引文件") from exc + try: + value = json.loads(raw) + except json.JSONDecodeError as exc: + raise CodeBuddyError("账号索引文件损坏,未修改任何账号数据") from exc + if ( + not isinstance(value, dict) + or set(value) != {"version", "active_account_id", "accounts"} + or value.get("version") != _INDEX_VERSION + ): + raise CodeBuddyError("账号索引文件格式或版本无效,未修改任何账号数据") + accounts = value.get("accounts") + if not isinstance(accounts, list): + raise CodeBuddyError("账号索引中的 accounts 字段无效") + normalized_accounts: list[dict[str, Any]] = [] + seen: set[str] = set() + for record in accounts: + if not isinstance(record, dict) or set(record) != _INDEX_FIELDS: + raise CodeBuddyError("账号索引包含无效字段,未修改任何账号数据") + backend = record.get("backend") + endpoint = record.get("endpoint") + uid = record.get("uid") + account_id = record.get("id") + nickname = record.get("nickname") + if not all(isinstance(item, str) and item for item in (backend, endpoint, uid)): + raise CodeBuddyError("账号索引包含不完整的身份信息") + try: + profile_for_key(backend) + canonical_endpoint = _safe_endpoint(endpoint) + except (ValueError, CodeBuddyError) as exc: + raise CodeBuddyError("账号索引包含无效的 backend 或 endpoint") from exc + if canonical_endpoint != endpoint or _account_id(backend, endpoint, uid) != account_id: + raise CodeBuddyError("账号索引身份校验失败,未修改任何账号数据") + if nickname is not None and not isinstance(nickname, str): + raise CodeBuddyError("账号索引包含无效的 nickname") + if account_id in seen: + raise CodeBuddyError("账号索引包含重复账号 id") + seen.add(account_id) + normalized_accounts.append( + { + "id": account_id, + "backend": backend, + "endpoint": endpoint, + "uid": uid, + "nickname": nickname, + } + ) + active_id = value.get("active_account_id") + if active_id is not None and active_id not in seen: + raise CodeBuddyError("账号索引的 active_account_id 不存在") + return { + "version": _INDEX_VERSION, + "active_account_id": active_id, + "accounts": normalized_accounts, + } + + def _write_index_locked(self, index: dict[str, Any]) -> None: + metadata = [ + {key: record[key] for key in ("id", "backend", "endpoint", "uid", "nickname")} + for record in index["accounts"] + ] + _atomic_json( + self.index_file, + { + "version": _INDEX_VERSION, + "active_account_id": index["active_account_id"], + "accounts": metadata, + }, + ) + + @staticmethod + def _find_account(index: dict[str, Any], account_id: str) -> dict[str, Any]: + if not isinstance(account_id, str) or not _ACCOUNT_ID_RE.fullmatch(account_id): + raise CodeBuddyError("账号 id 无效") + for record in index["accounts"]: + if record["id"] == account_id: + return record + raise CodeBuddyError("账号不存在") + + @staticmethod + def _validated_client_session( + client: CodeBuddyClient, + ) -> tuple[dict[str, Any], dict[str, Any]]: + if not isinstance(client, CodeBuddyClient): + raise CodeBuddyError("只能导入 CodeBuddyClient 会话") + session, _ = client._snapshot_session() + account = session.get("account") + auth = session.get("auth") + if ( + not isinstance(account, dict) + or account.get("uid") is None + or not str(account.get("uid")).strip() + ): + raise CodeBuddyError("账号 session 缺少 uid,不能导入") + if not _has_auth_token(auth): + raise CodeBuddyError("账号 session 缺少认证 token,不能导入") + try: + client._validate_session_metadata(session) + backend = client.profile.key + profile_for_key(backend) + endpoint = _safe_endpoint(client.endpoint) + except (ValueError, CodeBuddyError) as exc: + raise CodeBuddyError("账号 session 的 backend 或 endpoint 无效") from exc + uid = str(account["uid"]).strip() + if not uid: + raise CodeBuddyError("账号 session 缺少 uid,不能导入") + session["backend"] = backend + session["endpoint"] = endpoint + session_account = dict(account) + session_account["uid"] = uid + session["account"] = session_account + metadata = { + "id": _account_id(backend, endpoint, uid), + "backend": backend, + "endpoint": endpoint, + "uid": uid, + "nickname": ( + str(account.get("nickname"))[:200] + if account.get("nickname") is not None + else None + ), + } + return session, metadata + + def _new_client(self, record: dict[str, Any]) -> CodeBuddyClient: + session_file = self._session_file(record["id"]) + if session_file.is_symlink(): + raise CodeBuddyError("账号 session 路径无效") + profile = profile_for_key(record["backend"]) + client = CodeBuddyClient( + backend=resolve_backend( + profile=profile, + endpoint=record["endpoint"], + session_file=session_file, + environment={}, + ) + ) + session, _ = client._snapshot_session() + identity = self._identity_from_session(client, session) + if identity != record["id"]: + raise CodeBuddyError("账号 session 身份与索引不一致") + if not _has_auth_token(session.get("auth")): + raise CodeBuddyError("账号 session 缺少认证 token") + return client + + @staticmethod + def _identity_from_session(client: CodeBuddyClient, session: dict[str, Any]) -> str: + account = session.get("account") or {} + uid = account.get("uid") if isinstance(account, dict) else None + if uid is None or not str(uid).strip(): + raise CodeBuddyError("账号 session 缺少 uid") + endpoint = _safe_endpoint(client.endpoint) + return _account_id(client.profile.key, endpoint, str(uid).strip()) + + def _get_client_locked(self, record: dict[str, Any]) -> CodeBuddyClient: + client = self._clients.get(record["id"]) + if client is None: + client = self._new_client(record) + self._clients[record["id"]] = client + return client + + def list_accounts(self) -> list[dict[str, Any]]: + with self._lock: + index = self._read_index_locked() + active_id = index["active_account_id"] + return [ + {**record, "active": record["id"] == active_id} + for record in index["accounts"] + ] + + @property + def active_account_id(self) -> str | None: + with self._lock: + return self._read_index_locked()["active_account_id"] + + def import_client(self, client: CodeBuddyClient, activate: bool = False) -> str: + session, metadata = self._validated_client_session(client) + with self._lock: + index = self._read_index_locked() + existing = next( + (record for record in index["accounts"] if record["id"] == metadata["id"]), + None, + ) + record = metadata + session_file = self._session_file(metadata["id"]) + if session_file.is_symlink(): + raise CodeBuddyError("账号 session 路径无效") + cached = self._clients.get(metadata["id"]) + # Keep cached-client session writers outside the full disk/cache + # transaction. This covers the rollback interval too, so no writer + # can persist a newer session between restoring the file and cache. + session_guard = cached._session_lock if cached is not None else nullcontext() + with session_guard: + index_backup = _capture_file(self.index_file) + session_backup: bytes | None = None + cached_snapshot: dict[str, Any] | None = None + session_written = False + index_write_attempted = False + try: + if cached is None: + session_backup = _capture_file(session_file) + _atomic_json(session_file, session) + session_written = True + persistent_client = self._new_client(record) + else: + persistent_client = cached + for _ in range(5): + previous_session, generation = cached._snapshot_session() + candidate_backup = _capture_file(session_file) + if cached._save_session( + session, expected_generation=generation + ): + session_backup = candidate_backup + cached_snapshot = previous_session + session_written = True + break + else: + raise CodeBuddyError( + "账号 session 正在更新,请稍后重试关联" + ) + + if existing is None: + index["accounts"].append(record) + else: + position = index["accounts"].index(existing) + index["accounts"][position] = record + if activate or (existing is None and len(index["accounts"]) == 1): + index["active_account_id"] = metadata["id"] + index_write_attempted = True + self._write_index_locked(index) + except Exception as original_error: + rollback_errors: list[Exception] = [] + if session_written: + try: + _restore_file(session_file, session_backup) + except Exception as exc: + rollback_errors.append(exc) + if cached is not None and cached_snapshot is not None: + try: + cached._restore_session_snapshot(cached_snapshot) + except Exception as exc: + rollback_errors.append(exc) + if index_write_attempted: + try: + _restore_file(self.index_file, index_backup) + except Exception as exc: + rollback_errors.append(exc) + if rollback_errors: + raise CodeBuddyError( + "账号关联失败,且事务回滚未能完整恢复" + ) from rollback_errors[0] + raise original_error + self._clients[metadata["id"]] = persistent_client + return metadata["id"] + + def get_client(self, account_id: str) -> CodeBuddyClient: + with self._lock: + index = self._read_index_locked() + record = self._find_account(index, account_id) + return self._get_client_locked(record) + + def set_active(self, account_id: str) -> CodeBuddyClient: + with self._lock: + index = self._read_index_locked() + record = self._find_account(index, account_id) + client = self._get_client_locked(record) + index["active_account_id"] = account_id + self._write_index_locked(index) + return client + + def clear_active(self) -> None: + """Persist an empty active selection while retaining stored accounts.""" + with self._lock: + index = self._read_index_locked() + index["active_account_id"] = None + self._write_index_locked(index) + + def active_client(self) -> CodeBuddyClient | None: + with self._lock: + index = self._read_index_locked() + active_id = index["active_account_id"] + if active_id is None: + return None + record = self._find_account(index, active_id) + return self._get_client_locked(record) + + def delete_account(self, account_id: str) -> None: + with self._lock: + index = self._read_index_locked() + record = self._find_account(index, account_id) + if index["active_account_id"] == account_id: + raise CodeBuddyError("不能删除当前生效账号") + session_file = self._session_file(record["id"]) + if session_file.is_symlink(): + raise CodeBuddyError("账号 session 路径无效") + cached = self._clients.get(account_id) + if cached is not None: + cached._invalidate_persistence() + saved_bytes: bytes | None = None + try: + if session_file.exists(): + saved_bytes = session_file.read_bytes() + session_file.unlink() + except OSError as exc: + if cached is not None: + cached._enable_persistence() + raise CodeBuddyError("无法删除账号 session 文件") from exc + index["accounts"] = [ + item for item in index["accounts"] if item["id"] != account_id + ] + try: + self._write_index_locked(index) + except Exception: + if saved_bytes is not None: + temp_name: str | None = None + try: + fd, temp_name = tempfile.mkstemp( + prefix=f".{session_file.name}.", + suffix=".restore.tmp", + dir=self.root, + ) + with os.fdopen(fd, "wb") as stream: + stream.write(saved_bytes) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temp_name, session_file) + temp_name = None + except OSError: + pass + finally: + if temp_name is not None: + try: + os.unlink(temp_name) + except OSError: + pass + if cached is not None: + cached._enable_persistence() + raise + self._clients.pop(account_id, None) diff --git a/src/codebuddy_proxy/admin.html b/src/codebuddy_proxy/admin.html new file mode 100644 index 0000000..19ade72 --- /dev/null +++ b/src/codebuddy_proxy/admin.html @@ -0,0 +1,1883 @@ + + + + + + + CodeBuddy 账号管理 + + + +
+
+
+
CodeBuddy · 本机代理
+ +
+

账号管理

+

在这里关联、切换和管理 CodeBuddy 登录账号。账号凭据由本机代理保存,页面不会显示会话内容。

+
+ +
+ +
+
+
+
+

关联账号

+

选择账号所属区域,登录完成后会保存到本机账号库。

+
+
+
+ + +
+
+ +

登录页面会在新窗口打开;请在该窗口完成 CodeBuddy 登录。

+
+ +
+ +
+
+
+

已关联账号

+

当前账号会用于新发起的代理请求。

+
+
+ +
+
+ +
+
+ 遇到 429 自动切换账号 +

仅排除同一区域、同一服务地址中“本次请求模型仍在冷却”的账号;同一账号的其他模型仍可用。每个请求对每个账号最多尝试一次。默认关闭。

+
+ +
+

正在读取自动切换设置…

+
+
+ 全局自动签到 +

开启后由代理立即检查全部关联账号,并按北京时间计划自动签到。未知状态不会自动签到;关闭后会取消待执行计划。默认关闭。

+
+ +
+

正在读取自动签到设置…

+ + +
+
+
+ + + + diff --git a/src/codebuddy_proxy/admin_api.py b/src/codebuddy_proxy/admin_api.py new file mode 100644 index 0000000..9f9938e --- /dev/null +++ b/src/codebuddy_proxy/admin_api.py @@ -0,0 +1,547 @@ +"""Loopback-only admin page and account-management API.""" + +from __future__ import annotations + +import ipaddress +import json +import math +import re +import time +import unicodedata +from datetime import date as calendar_date +from importlib import resources +from typing import Any, Callable +from urllib.parse import urlsplit + +from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi.responses import HTMLResponse, JSONResponse, Response +from fastapi.routing import APIRoute +from fastapi.exceptions import RequestValidationError + +from codebuddy_proxy.account_manager import ( + AccountManager, + AccountManagerError, + is_loopback_address, +) + + +_WRITE_HEADER = "X-Account-Manager" +_MAX_LOGIN_BODY = 4096 +_MAX_SETTINGS_BODY = 4096 +_PUBLIC_ACCOUNT_FIELDS = ( + "id", + "backend", + "endpoint", + "uid", + "nickname", + "active", + "model_cooldowns", + "checkin", +) +_PUBLIC_JOB_FIELDS = ("id", "backend", "status", "auth_url", "error", "account_id") +_CHECKIN_STATUSES = frozenset( + {"unknown", "unchecked", "checked", "checking", "signing", "error"} +) +_SAFE_CHECKIN_ERRORS = frozenset( + { + "签到活动未开放", + "签到状态查询失败", + "签到请求失败", + "签到结果未确认", + "账号不存在", + } +) +_CHECKIN_DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$") + + +class _NoStoreRoute(APIRoute): + """Keep admin route failures generic and mark every routed response private.""" + + def get_route_handler(self): + original = super().get_route_handler() + + async def no_store_handler(request: Request) -> Response: + try: + response = await original(request) + except HTTPException as exc: + detail = exc.detail if isinstance(exc.detail, str) else "请求被拒绝" + response = JSONResponse({"detail": detail}, status_code=exc.status_code) + except RequestValidationError: + response = JSONResponse({"detail": "请求格式无效"}, status_code=400) + except Exception: + # Never serialize an exception, provider response, or traceback + # into a local administrative response. + response = JSONResponse({"detail": "管理请求失败"}, status_code=500) + response.headers["Cache-Control"] = "no-store" + response.headers["Pragma"] = "no-cache" + return response + + return no_store_handler + + +def _one_header(request: Request, name: str) -> str | None: + values = request.headers.getlist(name) + if len(values) != 1: + return None + return values[0] + + +def _host_identity(raw_host: str, scheme: str) -> tuple[str, int] | None: + if not raw_host or len(raw_host) > 320 or any(ch.isspace() for ch in raw_host): + return None + try: + parts = urlsplit("//" + raw_host) + if parts.username is not None or parts.password is not None: + return None + if parts.path or parts.query or parts.fragment or not parts.hostname: + return None + hostname = parts.hostname.lower() + if hostname.endswith("."): + hostname = hostname[:-1] + if hostname.endswith("."): + return None + try: + address = ipaddress.ip_address(hostname) + if not address.is_loopback: + return None + canonical_host = address.compressed.lower() + except ValueError: + if hostname != "localhost": + return None + canonical_host = hostname + port = parts.port + if port is None: + port = 443 if scheme.lower() == "https" else 80 + if not 1 <= port <= 65535: + return None + return canonical_host, port + except ValueError: + return None + + +def _origin_identity(raw_origin: str) -> tuple[str, str, int] | None: + if not raw_origin or len(raw_origin) > 512: + return None + try: + parts = urlsplit(raw_origin) + scheme = parts.scheme.lower() + if scheme not in {"http", "https"}: + return None + if not parts.netloc or parts.path or parts.query or parts.fragment: + return None + if parts.username is not None or parts.password is not None or not parts.hostname: + return None + host = parts.hostname.lower() + if host.endswith("."): + host = host[:-1] + if host.endswith("."): + return None + try: + address = ipaddress.ip_address(host) + if not address.is_loopback: + return None + host = address.compressed.lower() + except ValueError: + if host != "localhost": + return None + port = parts.port + if port is None: + port = 443 if scheme == "https" else 80 + if not 1 <= port <= 65535: + return None + return scheme, host, port + except ValueError: + return None + + +def _check_admin_request(request: Request) -> None: + peer = request.client.host if request.client is not None else None + if not is_loopback_address(peer): + raise HTTPException(status_code=403, detail="仅允许本机访问管理页面") + + raw_host = _one_header(request, "host") + host = _host_identity(raw_host or "", request.url.scheme) + if host is None: + raise HTTPException(status_code=403, detail="管理页面 Host 无效") + + raw_origin = _one_header(request, "origin") + if raw_origin is None: + # Missing or duplicated Origin is allowed for reads. Mutations are + # separately protected by the dedicated header below. + if request.headers.getlist("origin"): + raise HTTPException(status_code=403, detail="管理页面 Origin 无效") + return + + origin = _origin_identity(raw_origin) + expected = (request.url.scheme.lower(), host[0], host[1]) + if origin != expected: + raise HTTPException(status_code=403, detail="跨站管理请求已拒绝") + + +def _require_write_header(request: Request) -> None: + if _one_header(request, _WRITE_HEADER) != "1": + raise HTTPException(status_code=403, detail="缺少管理请求标记") + + +def _safe_job(value: Any) -> dict[str, Any]: + if not isinstance(value, dict): + return {key: None for key in _PUBLIC_JOB_FIELDS} + return {key: value.get(key) for key in _PUBLIC_JOB_FIELDS} + + +def _safe_snapshot(value: Any) -> dict[str, Any]: + if not isinstance(value, dict): + return { + "accounts": [], + "active_account_id": None, + "settings": _safe_settings(None), + "server_time": time.time(), + } + accounts = value.get("accounts") + safe_accounts = [] + if isinstance(accounts, list): + for account in accounts: + if isinstance(account, dict): + safe = { + key: account[key] + for key in _PUBLIC_ACCOUNT_FIELDS[:6] + if key in account + } + safe["model_cooldowns"] = _safe_model_cooldowns( + account.get("model_cooldowns") + ) + safe["checkin"] = _safe_checkin(account.get("checkin")) + safe_accounts.append(safe) + server_time = value.get("server_time") + return { + "accounts": safe_accounts, + "active_account_id": value.get("active_account_id"), + "settings": _safe_settings(value.get("settings")), + "server_time": float(server_time) if _finite_number(server_time) else time.time(), + } + + +def _finite_number(value: Any) -> bool: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return False + try: + return math.isfinite(float(value)) + except (OverflowError, TypeError, ValueError): + return False + + +def _safe_model_name(value: Any) -> str | None: + if ( + not isinstance(value, str) + or not value + or len(value) > 256 + or not value.strip() + or any(unicodedata.category(char) in {"Cc", "Cs"} for char in value) + ): + return None + return value + + +def _safe_model_cooldowns(value: Any) -> list[dict[str, Any]]: + if not isinstance(value, list): + return [] + result = [] + for record in value: + if not isinstance(record, dict): + continue + model = _safe_model_name(record.get("model")) + if model is None: + continue + result.append( + { + "model": model, + "rate_limited": ( + record.get("rate_limited") + if type(record.get("rate_limited")) is bool + else False + ), + "cooldown_until": ( + float(record["cooldown_until"]) + if _finite_number(record.get("cooldown_until")) + else None + ), + "cooldown_remaining_seconds": ( + record["cooldown_remaining_seconds"] + if type(record.get("cooldown_remaining_seconds")) is int + and record["cooldown_remaining_seconds"] >= 0 + else 0 + ), + "last_rate_limited_at": ( + float(record["last_rate_limited_at"]) + if _finite_number(record.get("last_rate_limited_at")) + else None + ), + "cooldown_source": ( + record.get("cooldown_source") + if record.get("cooldown_source") in ("retry_after", "default") + else None + ), + } + ) + return result + + +def _safe_settings(value: Any) -> dict[str, Any]: + auto_switch = value.get("auto_switch_on_429") if isinstance(value, dict) else None + auto_checkin = value.get("auto_checkin_enabled") if isinstance(value, dict) else None + return { + "auto_switch_on_429": auto_switch if type(auto_switch) is bool else False, + "auto_checkin_enabled": auto_checkin if type(auto_checkin) is bool else False, + "default_cooldown_seconds": 60, + } + + +def _safe_checkin(value: Any) -> dict[str, Any]: + """Return only the stable, non-sensitive check-in response fields.""" + + if not isinstance(value, dict): + value = {} + raw_date = value.get("date") + checkin_date = time.strftime("%Y-%m-%d", time.gmtime(time.time() + 8 * 60 * 60)) + if isinstance(raw_date, str) and _CHECKIN_DATE_RE.fullmatch(raw_date): + try: + calendar_date.fromisoformat(raw_date) + checkin_date = raw_date + except ValueError: + pass + status = value.get("status") + if not isinstance(status, str) or status not in _CHECKIN_STATUSES: + status = "error" + error = value.get("error") + if not isinstance(error, str) or error not in _SAFE_CHECKIN_ERRORS: + error = "签到状态查询失败" if status == "error" else None + return { + "date": checkin_date, + "status": status, + "checked": value.get("checked") if type(value.get("checked")) is bool else None, + "last_checked_at": ( + float(value["last_checked_at"]) + if _finite_number(value.get("last_checked_at")) + else None + ), + "last_checkin_at": ( + float(value["last_checkin_at"]) + if _finite_number(value.get("last_checkin_at")) + else None + ), + "error": error, + "next_checkin_at": ( + float(value["next_checkin_at"]) + if _finite_number(value.get("next_checkin_at")) + else None + ), + } + + +class _DuplicateJSONKey(ValueError): + pass + + +def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise _DuplicateJSONKey("duplicate JSON key") + result[key] = value + return result + + +async def _read_json_object(request: Request, max_bytes: int) -> dict[str, Any]: + chunks: list[bytes] = [] + length = 0 + async for chunk in request.stream(): + length += len(chunk) + if length > max_bytes: + raise HTTPException(status_code=413, detail="请求内容过大") + chunks.append(chunk) + try: + payload = json.loads(b"".join(chunks) or b"{}", object_pairs_hook=_unique_object) + except (json.JSONDecodeError, UnicodeDecodeError, _DuplicateJSONKey): + raise HTTPException(status_code=400, detail="请求 JSON 无效") from None + if not isinstance(payload, dict): + raise HTTPException(status_code=400, detail="请求 JSON 必须为对象") + return payload + + +def _raise_manager_error(exc: Exception) -> None: + if isinstance(exc, AccountManagerError): + raise HTTPException(status_code=exc.status_code, detail=exc.message) from None + if isinstance(exc, KeyError): + raise HTTPException(status_code=404, detail="未找到请求的账号或任务") from None + if isinstance(exc, ValueError): + raise HTTPException(status_code=400, detail="请求参数无效") from None + raise HTTPException(status_code=500, detail="管理请求失败") from None + + +def _manager_has_account(manager: Any, account_id: str) -> bool: + has_account = getattr(manager, "has_account", None) + if callable(has_account): + try: + return bool(has_account(account_id)) + except Exception: + return False + try: + accounts = manager.snapshot().get("accounts", []) + except Exception: + return False + return any( + isinstance(account, dict) and account.get("id") == account_id + for account in accounts + ) + + +def create_admin_router(get_manager: Callable[[], AccountManager]) -> APIRouter: + """Create the loopback-only `/admin` page and `/admin/api` routes.""" + + router = APIRouter(route_class=_NoStoreRoute, dependencies=[Depends(_check_admin_request)]) + + @router.get("/admin", response_class=HTMLResponse) + @router.get("/admin/", response_class=HTMLResponse, include_in_schema=False) + async def admin_page() -> Response: + try: + html = resources.files("codebuddy_proxy").joinpath("admin.html").read_text( + encoding="utf-8" + ) + except (FileNotFoundError, ModuleNotFoundError, OSError): + raise HTTPException(status_code=503, detail="管理页面暂不可用") from None + return HTMLResponse(html) + + @router.get("/admin/api/accounts") + async def list_accounts() -> dict[str, Any]: + try: + return _safe_snapshot(get_manager().snapshot()) + except Exception as exc: + _raise_manager_error(exc) + + @router.get("/admin/api/settings") + async def get_settings() -> dict[str, Any]: + try: + return _safe_settings(get_manager().get_settings()) + except Exception as exc: + _raise_manager_error(exc) + + @router.patch("/admin/api/settings") + async def patch_settings(request: Request) -> dict[str, Any]: + _require_write_header(request) + payload = await _read_json_object(request, _MAX_SETTINGS_BODY) + allowed_settings = {"auto_switch_on_429", "auto_checkin_enabled"} + if ( + not payload + or not set(payload).issubset(allowed_settings) + or any(type(value) is not bool for value in payload.values()) + ): + raise HTTPException( + status_code=400, + detail="仅接受非空布尔值账号管理设置", + ) + try: + updated = get_manager().set_settings( + payload.get("auto_switch_on_429"), + auto_checkin_enabled=payload.get("auto_checkin_enabled"), + ) + return _safe_settings(updated) + except Exception as exc: + _raise_manager_error(exc) + + @router.get( + "/admin/api/accounts/{account_id}/checkin-status", response_model=None + ) + async def checkin_status(account_id: str) -> dict[str, Any] | JSONResponse: + manager = get_manager() + if not _manager_has_account(manager, account_id): + raise HTTPException(status_code=404, detail="未找到该账号") + try: + result = await manager.checkins.check_status( + account_id, reason="manual_status" + ) + return _safe_checkin(result) + except Exception: + return JSONResponse( + _safe_checkin({"status": "error", "error": "签到状态查询失败"}), + status_code=502, + ) + + @router.post("/admin/api/accounts/{account_id}/checkin", response_model=None) + async def checkin(account_id: str, request: Request) -> dict[str, Any] | JSONResponse: + _require_write_header(request) + manager = get_manager() + if not _manager_has_account(manager, account_id): + raise HTTPException(status_code=404, detail="未找到该账号") + try: + result = await manager.checkins.check_in(account_id, reason="manual") + return _safe_checkin(result) + except Exception: + return JSONResponse( + _safe_checkin({"status": "error", "error": "签到请求失败"}), + status_code=502, + ) + + @router.post("/admin/api/accounts/{account_id}/activate") + async def activate_account(account_id: str, request: Request) -> dict[str, bool]: + _require_write_header(request) + try: + get_manager().activate(account_id) + return {"ok": True} + except Exception as exc: + _raise_manager_error(exc) + + @router.delete("/admin/api/accounts/{account_id}") + async def delete_account(account_id: str, request: Request) -> dict[str, bool]: + _require_write_header(request) + try: + get_manager().delete(account_id) + return {"ok": True} + except Exception as exc: + _raise_manager_error(exc) + + @router.post("/admin/api/logins", status_code=202) + async def start_login(request: Request) -> JSONResponse: + _require_write_header(request) + chunks: list[bytes] = [] + length = 0 + async for chunk in request.stream(): + length += len(chunk) + if length > _MAX_LOGIN_BODY: + raise HTTPException(status_code=413, detail="请求内容过大") + chunks.append(chunk) + try: + payload = json.loads(b"".join(chunks) or b"{}") + except (json.JSONDecodeError, UnicodeDecodeError): + raise HTTPException(status_code=400, detail="请求 JSON 无效") from None + if not isinstance(payload, dict) or not isinstance(payload.get("backend"), str): + raise HTTPException(status_code=400, detail="backend 必须为 domestic 或 global") + try: + job = _safe_job(get_manager().start_login(payload["backend"])) + return JSONResponse(job, status_code=202) + except Exception as exc: + _raise_manager_error(exc) + + @router.get("/admin/api/logins/{job_id}") + async def get_login(job_id: str) -> dict[str, Any]: + try: + return _safe_job(get_manager().get_login(job_id)) + except Exception as exc: + _raise_manager_error(exc) + + @router.delete("/admin/api/logins/{job_id}") + async def cancel_login(job_id: str, request: Request) -> dict[str, Any]: + _require_write_header(request) + try: + return _safe_job(get_manager().cancel_login(job_id)) + except Exception as exc: + _raise_manager_error(exc) + + @router.api_route( + "/admin/{unmatched_path:path}", + methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"], + include_in_schema=False, + ) + async def unknown_admin_path() -> Response: + return JSONResponse({"detail": "未找到管理接口"}, status_code=404) + + return router diff --git a/src/codebuddy_proxy/billing_proxy.py b/src/codebuddy_proxy/billing_proxy.py index af268c8..8b0d4de 100644 --- a/src/codebuddy_proxy/billing_proxy.py +++ b/src/codebuddy_proxy/billing_proxy.py @@ -26,6 +26,9 @@ def _upstream_response(response: httpx.Response) -> Response: content_type = response.headers.get("content-type") if content_type: headers["content-type"] = content_type + retry_after = response.headers.get("Retry-After") + if retry_after is not None: + headers["Retry-After"] = retry_after return Response( content=response.content, status_code=response.status_code, @@ -52,25 +55,31 @@ async def forward_billing( another concurrent request already updated the access token). If refresh fails, preserve the original upstream 401 instead of blindly retrying. """ - state.ensure_auth() - url = state.client.build_url(path) + # Capture once: an account switch while awaiting a response must never + # refresh the new account or send its credentials to the old endpoint. + account_client = state.client + if getattr(state, "managed_auth", False): + await asyncio.to_thread(state.ensure_auth) + else: + state.ensure_auth() + url = account_client.build_url(path) method = method.upper() kwargs = {"json": {}} if method == "POST" else {} async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client: try: - headers = state.client.auth_headers() + headers = account_client.auth_headers() response = await client.request(method, url, headers=headers, **kwargs) if response.status_code == 401: async with _refresh_lock: - latest_headers = state.client.auth_headers() + latest_headers = account_client.auth_headers() old_token = headers.get("Authorization") new_token = latest_headers.get("Authorization") if new_token == old_token: try: - refreshed = await asyncio.to_thread(state.client.refresh) + refreshed = await asyncio.to_thread(account_client.refresh) except Exception as exc: diagnostic( "billing_auth_refresh_failed", @@ -80,7 +89,7 @@ async def forward_billing( if not refreshed: diagnostic("billing_auth_refresh_failed", path=path, reason="not_refreshed") return _upstream_response(response) - latest_headers = state.client.auth_headers() + latest_headers = account_client.auth_headers() new_token = latest_headers.get("Authorization") if not new_token or new_token == old_token: @@ -96,4 +105,16 @@ async def forward_billing( diagnostic("billing_upstream_failed", path=path, reason=type(exc).__name__) return _transport_error(502, "upstream_connection_error") + if response.status_code == 429: + manager = getattr(state, "account_runtime", None) + if manager is not None: + try: + # Billing has no model dimension: it must neither cool down + # inference models nor change the selected account. + await asyncio.to_thread( + manager.handle_rate_limit, account_client, + response.headers.get("Retry-After"), model=None, allow_switch=False, + ) + except Exception as exc: + diagnostic("billing_rate_limit_record_failed", error_type=type(exc).__name__) return _upstream_response(response) diff --git a/src/codebuddy_proxy/codebuddy_client_demo.py b/src/codebuddy_proxy/codebuddy_client_demo.py index e4d2d09..582fea8 100644 --- a/src/codebuddy_proxy/codebuddy_client_demo.py +++ b/src/codebuddy_proxy/codebuddy_client_demo.py @@ -9,10 +9,13 @@ from __future__ import annotations import argparse +import copy import json import os import pathlib import sys +import tempfile +import threading import time import urllib.error import urllib.parse @@ -99,6 +102,10 @@ def __init__( self.platform = platform self.prefix = self.profile.auth_prefix self.session_file = resolved_backend.session_file + self._session_lock = threading.RLock() + self._auth_lock = threading.RLock() + self._session_generation = 0 + self._persistence_enabled = True self.session: dict[str, Any] = self._load_session() def _load_session(self) -> dict[str, Any]: @@ -169,23 +176,77 @@ def _validate_session_metadata(self, session: dict[str, Any]) -> None: "原文件未修改,请使用匹配的 --session-file。" ) - def _save_session(self, session: dict[str, Any]) -> None: + def _snapshot_session(self) -> tuple[dict[str, Any], int]: + with self._session_lock: + return copy.deepcopy(self.session), self._session_generation + + def _save_session( + self, + session: dict[str, Any], + *, + expected_generation: int | None = None, + ) -> bool: persisted_session = dict(session) persisted_session["backend"] = self.profile.key persisted_session["endpoint"] = self.endpoint - self.session_file.parent.mkdir(parents=True, exist_ok=True) - flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC - fd = os.open(self.session_file, flags, 0o600) - try: - with os.fdopen(fd, "w", encoding="utf-8") as stream: - json.dump(persisted_session, stream, ensure_ascii=False, indent=2) - stream.write("\n") - finally: - try: - os.chmod(self.session_file, 0o600) - except OSError: - pass - self.session = persisted_session + with self._session_lock: + if not self._persistence_enabled: + return False + if ( + expected_generation is not None + and self._session_generation != expected_generation + ): + return False + self.session_file.parent.mkdir(parents=True, exist_ok=True) + temp_name: str | None = None + try: + fd, temp_name = tempfile.mkstemp( + prefix=f".{self.session_file.name}.", + suffix=".tmp", + dir=self.session_file.parent, + ) + with os.fdopen(fd, "w", encoding="utf-8") as stream: + json.dump(persisted_session, stream, ensure_ascii=False, indent=2) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temp_name, self.session_file) + temp_name = None + try: + os.chmod(self.session_file, 0o600) + except OSError: + pass + except OSError as exc: + raise CodeBuddyError( + f"无法原子写入 session 文件: {self.session_file}" + ) from exc + finally: + if temp_name is not None: + try: + os.unlink(temp_name) + except OSError: + pass + self.session = persisted_session + self._session_generation += 1 + return True + + def _invalidate_persistence(self) -> None: + """Prevent an old client reference from recreating a removed session.""" + with self._session_lock: + self._session_generation += 1 + self._persistence_enabled = False + + def _enable_persistence(self) -> None: + """Restore writes if a store deletion transaction was rolled back.""" + with self._session_lock: + self._session_generation += 1 + self._persistence_enabled = True + + def _restore_session_snapshot(self, session: dict[str, Any]) -> None: + """Restore in-memory state after an account-store transaction fails.""" + with self._session_lock: + self.session = copy.deepcopy(session) + self._session_generation += 1 def build_url(self, path: str) -> str: """Build a URL from the resolved endpoint and a relative path.""" @@ -259,19 +320,27 @@ def _request( def auth_headers(self, *, access: bool = True, refresh: bool = False) -> dict[str, str]: """构造认证 headers,模拟插件行为""" - account = self.session.get("account") or {} - auth = self.session.get("auth") or {} + self._get_machine_id() + session, _ = self._snapshot_session() + return self._auth_headers_for_session(session, access=access, refresh=refresh) + + def _auth_headers_for_session( + self, + session: dict[str, Any], + *, + access: bool = True, + refresh: bool = False, + ) -> dict[str, str]: + account = session.get("account") or {} + auth = session.get("auth") or {} headers: dict[str, str] = { - # 🔥 关键标识:所有请求都需要这个 header "X-Product-Code": "codebuddy", - # 🔥 IDE 识别信息(服务端用于识别客户端类型) "X-IDE-Type": "vscode", "X-IDE-Name": "Visual Studio Code", - "X-IDE-Version": "1.70.2", # VSCode 最低版本要求 - "X-Product-Version": "4.10.33259736", # 插件版本 - "X-Machine-Id": self._get_machine_id(), + "X-IDE-Version": "1.70.2", + "X-Product-Version": "4.10.33259736", + "X-Machine-Id": str(session.get("machineId") or ""), } - if account.get("uid"): headers["X-User-Id"] = str(account["uid"]) if access and auth.get("accessToken"): @@ -282,12 +351,9 @@ def auth_headers(self, *, access: bool = True, refresh: bool = False) -> dict[st headers["X-Enterprise-Id"] = str(account["enterpriseId"]) headers["X-Tenant-Id"] = str(account["enterpriseId"]) if account.get("departmentInfo"): - # 🔥 新增:部门信息(如果有) headers["X-Department-Info"] = str(account["departmentInfo"]) domain = self._domain_for_token(auth) if domain: - # The extension calls this the domain header. The server accepts - # X-Domain for the plugin protocol. headers["X-Domain"] = domain return headers @@ -343,28 +409,22 @@ def _get_machine_id(self) -> str: """获取或生成机器ID(模拟VSCode的machineId)""" import platform import uuid - - # 尝试从session读取已保存的machineId - saved_machine_id = self.session.get("machineId") - if saved_machine_id: - return str(saved_machine_id) - - # 生成一个稳定的machineId(基于主机名和用户) - try: - hostname = platform.node() - username = os.environ.get("USER") or os.environ.get("USERNAME") or "unknown" - seed = f"{hostname}-{username}" - machine_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, seed)) - except Exception: - # 如果失败,生成随机ID - machine_id = str(uuid.uuid4()) - - # 保存到session - self.session["machineId"] = machine_id - self._save_session(self.session) - return machine_id + with self._session_lock: + saved_machine_id = self.session.get("machineId") + if saved_machine_id: + return str(saved_machine_id) + try: + hostname = platform.node() + username = os.environ.get("USER") or os.environ.get("USERNAME") or "unknown" + seed = f"{hostname}-{username}" + machine_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, seed)) + except Exception: + machine_id = str(uuid.uuid4()) + self._save_session({**self.session, "machineId": machine_id}) + return machine_id - def login(self, *, open_browser: bool = True, timeout: int = 300) -> None: + def begin_login(self) -> dict[str, str]: + """Request an authorization URL and state without opening a browser.""" no_auth = { "X-No-Authorization": "true", "X-No-User-Id": "true", @@ -385,35 +445,92 @@ def login(self, *, open_browser: bool = True, timeout: int = 300) -> None: state = state_payload.get("state") if not state: raise CodeBuddyError("登录状态响应缺少 state") - print(f"请在浏览器中完成登录:\n{auth_url}") - if open_browser: - webbrowser.open(auth_url) - deadline = time.monotonic() + timeout + return {"auth_url": auth_url, "state": str(state)} + + @staticmethod + def _check_login_cancelled(cancel_event: Any | None) -> None: + if cancel_event is not None and cancel_event.is_set(): + raise CodeBuddyError("登录已取消") + + @classmethod + def _wait_for_login_poll(cls, deadline: float, cancel_event: Any | None) -> None: + cls._check_login_cancelled(cancel_event) + remaining = deadline - time.monotonic() + if remaining <= 0: + return + pause = min(1.0, remaining) + if cancel_event is not None and hasattr(cancel_event, "wait"): + if cancel_event.wait(pause): + raise CodeBuddyError("登录已取消") + else: + time.sleep(pause) + cls._check_login_cancelled(cancel_event) + + @staticmethod + def _account_has_uid(account: Any) -> bool: + return ( + isinstance(account, dict) + and account.get("uid") is not None + and bool(str(account.get("uid")).strip()) + ) + + def complete_login( + self, + challenge: dict[str, Any], + timeout: float = 300, + cancel_event: Any | None = None, + ) -> None: + """Poll the challenge until a token and uid-bearing account are ready. + + One deadline covers both polling stages. No session state is saved + unless a token and an account with a uid have both been returned. + """ + if not isinstance(challenge, dict) or not challenge.get("state"): + raise CodeBuddyError("登录 challenge 缺少 state") + try: + timeout_value = float(timeout) + except (TypeError, ValueError) as exc: + raise CodeBuddyError("登录超时时间无效") from exc + if timeout_value <= 0: + raise CodeBuddyError("登录超时时间必须大于 0") + state = str(challenge["state"]) + state_query = urllib.parse.quote(state, safe="") + deadline = time.monotonic() + timeout_value + _, initial_generation = self._snapshot_session() + no_auth = { + "X-No-Authorization": "true", + "X-No-User-Id": "true", + "X-No-Enterprise-Id": "true", + "X-No-Department-Info": "true", + } while time.monotonic() < deadline: - time.sleep(1) + self._check_login_cancelled(cancel_event) + remaining = deadline - time.monotonic() + if remaining <= 0: + break try: token = self._unwrap( self._request( "GET", - f"{self._auth_path('/auth/token')}?state={urllib.parse.quote(str(state), safe='')}", + f"{self._auth_path('/auth/token')}?state={state_query}", headers={**self._enterprise_headers(), **no_auth}, + timeout=min(30.0, remaining), ) ) except CodeBuddyError: - continue + token = None if isinstance(token, dict) and token.get("accessToken"): - # 🔥 关键修复:/login/account 也需要轮询(模拟 VSIX loopGetAccount) - # 服务端在用户浏览器登录后,账户信息可能需要异步准备 account = None - account_deadline = time.monotonic() + 60 # 账户轮询最多60秒 - - while time.monotonic() < account_deadline: - time.sleep(1) # 每秒轮询一次 + while time.monotonic() < deadline: + self._check_login_cancelled(cancel_event) + remaining = deadline - time.monotonic() + if remaining <= 0: + break try: account = self._unwrap( self._request( "GET", - f"{self._auth_path('/login/account')}?state={urllib.parse.quote(str(state), safe='')}", + f"{self._auth_path('/login/account')}?state={state_query}", headers={ # 关键修复:必须包含完整的 enterprise headers(包括 X-Domain) **self._enterprise_headers(token), @@ -422,21 +539,48 @@ def login(self, *, open_browser: bool = True, timeout: int = 300) -> None: "X-No-Enterprise-Id": "true", "X-No-Department-Info": "true", }, + timeout=min(30.0, remaining), ) ) - if isinstance(account, dict) and account.get("uid"): - # 账户信息获取成功 + self._check_login_cancelled(cancel_event) + if time.monotonic() >= deadline: + account = None + break + if self._account_has_uid(account): break + account = None except CodeBuddyError: - # 401 或其他错误:账户信息可能尚未就绪,继续轮询 - continue - - if not account or not isinstance(account, dict): - raise CodeBuddyError("账户信息获取超时") - - self._save_session({**self.session, "auth": token, "account": account}) - print(f"登录成功,用户: {account.get('nickname') or account.get('uid', '')}") - return + account = None + if account is not None: + break + self._wait_for_login_poll(deadline, cancel_event) + + if self._account_has_uid(account): + self._check_login_cancelled(cancel_event) + if time.monotonic() >= deadline: + raise CodeBuddyError("登录超时:未能获取包含 uid 的账户信息") + session, _ = self._snapshot_session() + self._save_session( + {**session, "auth": token, "account": account}, + expected_generation=initial_generation, + ) + return + raise CodeBuddyError("登录超时:未能获取包含 uid 的账户信息") + + self._wait_for_login_poll(deadline, cancel_event) + + self._check_login_cancelled(cancel_event) + raise CodeBuddyError("登录超时:未能获取访问令牌") + + def login(self, *, open_browser: bool = True, timeout: int = 300) -> None: + challenge = self.begin_login() + auth_url = challenge["auth_url"] + print(f"请在浏览器中完成登录:\n{auth_url}") + if open_browser: + webbrowser.open(auth_url) + self.complete_login(challenge, timeout=timeout) + account = self.session.get("account") or {} + print(f"登录成功,用户: {account.get('nickname') or account.get('uid', '')}") @staticmethod def _token_headers(token: dict[str, Any]) -> dict[str, str]: @@ -448,7 +592,18 @@ def _token_headers(token: dict[str, Any]) -> dict[str, str]: return headers def refresh(self) -> bool: - auth = self.session.get("auth") or {} + with self._auth_lock: + return self._refresh_locked() + + def _refresh_locked(self) -> bool: + session, generation = self._snapshot_session() + auth = session.get("auth") or {} + refresh_token = auth.get("refreshToken") + if not refresh_token: + return False + self._get_machine_id() + session, generation = self._snapshot_session() + auth = session.get("auth") or {} refresh_token = auth.get("refreshToken") if not refresh_token: return False @@ -457,7 +612,7 @@ def refresh(self) -> bool: "POST", self._auth_path("/auth/token/refresh"), headers={ - **self.auth_headers(access=False, refresh=True), + **self._auth_headers_for_session(session, access=False, refresh=True), "X-Auth-Refresh-Source": "plugin", }, body={}, @@ -465,19 +620,41 @@ def refresh(self) -> bool: ) if not isinstance(payload, dict) or not payload.get("accessToken"): return False - self._save_session({**self.session, "auth": payload}) - return True - - def ensure_authenticated(self, *, open_browser: bool = True) -> None: - auth = self.session.get("auth") or {} - now_ms = int(time.time() * 1000) - expires_at = int(auth.get("expiresAt") or 0) - if auth.get("accessToken") and (not expires_at or expires_at > now_ms + 60_000): - return - if self.refresh(): - print("access token 已刷新") - return - self.login(open_browser=open_browser) + saved = self._save_session( + {**session, "auth": payload}, expected_generation=generation + ) + if saved: + return True + # A newer login or import completed while refresh was in flight. + current_auth = self._snapshot_session()[0].get("auth") or {} + current_expiry = int(current_auth.get("expiresAt") or 0) + return bool( + current_auth.get("accessToken") + and (not current_expiry or current_expiry > int(time.time() * 1000)) + ) + + def ensure_authenticated( + self, + *, + open_browser: bool = True, + allow_login: bool = True, + ) -> None: + with self._auth_lock: + # A preceding request may have refreshed the token while this + # caller waited for the per-client authentication lock. + auth = self._snapshot_session()[0].get("auth") or {} + now_ms = int(time.time() * 1000) + expires_at = int(auth.get("expiresAt") or 0) + if auth.get("accessToken") and ( + not expires_at or expires_at > now_ms + 60_000 + ): + return + if self.refresh(): + print("access token 已刷新") + return + if not allow_login: + raise CodeBuddyError("当前账号需要重新登录;此请求禁止自动打开登录流程") + self.login(open_browser=open_browser) def stream_chat( self, diff --git a/test_account_checkin.py b/test_account_checkin.py new file mode 100644 index 0000000..39b7d33 --- /dev/null +++ b/test_account_checkin.py @@ -0,0 +1,825 @@ +"""Offline behavior tests for per-account scheduled daily check-in.""" + +import asyncio +from datetime import date, datetime, timedelta, timezone +import json +from types import SimpleNamespace + +import httpx +import pytest +from fastapi.responses import Response + +import codebuddy_proxy.billing_proxy as billing +from codebuddy_proxy.account_checkin import ( + AccountCheckinService, + CHECKIN_PATH, + STATUS_PATH, +) + + +BEIJING = timezone(timedelta(hours=8)) + + +class FakeClient: + def __init__(self, account_id): + self.account_id = account_id + + +class FakeStore: + def __init__(self, account_ids=("acct-a",)): + self.clients = {item: FakeClient(item) for item in account_ids} + + def list_accounts(self): + return [{"id": item} for item in self.clients] + + def get_client(self, account_id): + return self.clients[account_id] + + +class FakeManager: + def __init__(self, account_ids=("acct-a",), *, enabled=False, active="acct-a"): + self.store = FakeStore(account_ids) + self.settings = {"auto_checkin_enabled": enabled} + self.active_account_id = active + self.rate_limits = [] + + def get_settings(self): + return dict(self.settings) + + def has_account(self, account_id): + return account_id in self.store.clients + + def handle_rate_limit(self, client, retry_after, *, model=None, allow_switch=True): + self.rate_limits.append((client.account_id, retry_after, model, allow_switch)) + return None + + +def response(payload, status=200, headers=None): + if isinstance(payload, bytes): + content = payload + else: + content = json.dumps(payload).encode("utf-8") + return Response(content=content, status_code=status, headers=headers or {}) + + +def status_body(checked, *, active=True): + return { + "code": 0, + "msg": "OK", + "data": {"active": active, "today_checked_in": checked}, + } + + +async def wait_event(event): + await asyncio.wait_for(event.wait(), timeout=1.5) + + +def run(coro): + return asyncio.run(coro) + + +class MutableClock: + def __init__(self, value): + self.value = value + + def __call__(self): + return self.value + + def advance(self, seconds): + self.value += timedelta(seconds=seconds) + + +async def never_sleep(_seconds): + await asyncio.Event().wait() + + +def test_startup_when_enabled_checks_every_account_and_supplements_only_unchecked(): + async def scenario(): + manager = FakeManager(("acct-a", "acct-b"), enabled=True, active="acct-a") + calls = [] + completed = asyncio.Event() + status_counts = {} + + async def requester(client, method, path): + calls.append((client.account_id, method, path)) + if path == STATUS_PATH: + count = status_counts.get(client.account_id, 0) + 1 + status_counts[client.account_id] = count + if client.account_id == "acct-a" or count > 1: + if client.account_id == "acct-b" and count > 1: + completed.set() + return response(status_body(True)) + return response(status_body(False)) + return response({"code": 0, "msg": "ok", "data": {}}) + + service = AccountCheckinService( + manager, requester=requester, sleep=never_sleep + ) + await service.start() + await wait_event(completed) + assert status_counts == {"acct-a": 1, "acct-b": 2} + assert sum(path == CHECKIN_PATH for _, _, path in calls) == 1 + assert service.snapshot("acct-a")["status"] == "checked" + assert service.snapshot("acct-b")["status"] == "checked" + assert service.snapshot("acct-a")["last_checkin_at"] is None + assert service.snapshot("acct-b")["last_checkin_at"] is not None + assert manager.active_account_id == "acct-a" + await service.close() + + run(scenario()) + + +def test_enable_transition_immediately_checks_and_supplements(): + async def scenario(): + manager = FakeManager(enabled=False) + calls = [] + completed = asyncio.Event() + + async def requester(client, method, path): + calls.append(path) + if path == STATUS_PATH: + if calls.count(STATUS_PATH) == 2: + completed.set() + return response(status_body(True)) + return response(status_body(False)) + return response({"code": 0, "data": {}}) + + service = AccountCheckinService( + manager, requester=requester, sleep=never_sleep + ) + await service.start() + await asyncio.sleep(0) + assert calls == [] + manager.settings["auto_checkin_enabled"] = True + service.settings_changed(True) + await wait_event(completed) + assert calls.count(CHECKIN_PATH) == 1 + await service.close() + + run(scenario()) + + +def test_midnight_flow_checks_then_rechecks_at_absolute_random_time(): + async def scenario(): + manager = FakeManager(enabled=True) + clock = MutableClock(datetime(2026, 10, 24, 0, 1, tzinfo=BEIJING)) + calls = [] + check_count = 0 + delays = [] + completed = asyncio.Event() + + async def requester(client, method, path): + nonlocal check_count + calls.append(path) + if path == STATUS_PATH: + check_count += 1 + if check_count < 3: + return response(status_body(False)) + completed.set() + return response(status_body(True)) + return response({"code": 0, "data": {}}) + + async def advancing_sleep(seconds): + delays.append(seconds) + clock.advance(seconds) + + service = AccountCheckinService( + manager, + clock=clock, + random_delay=lambda: 120, + sleep=advancing_sleep, + requester=requester, + ) + midnight = datetime(2026, 10, 24, tzinfo=BEIJING) + await service._prepare_midnight("acct-a", "2026-10-24", midnight) + await wait_event(completed) + assert delays == pytest.approx([60.0]) + assert calls == [STATUS_PATH, STATUS_PATH, CHECKIN_PATH, STATUS_PATH] + assert service.snapshot("acct-a")["status"] == "checked" + await service.close() + + run(scenario()) + + +def test_slow_midnight_query_does_not_push_random_window_past_600_seconds(): + async def scenario(): + manager = FakeManager(enabled=True) + clock = MutableClock(datetime(2026, 10, 24, 0, 7, tzinfo=BEIJING)) + delays = [] + + async def requester(client, method, path): + return response(status_body(False)) + + async def capture_sleep(seconds): + delays.append(seconds) + # Hold this task until the assertion has observed the due delay. + await asyncio.Event().wait() + + service = AccountCheckinService( + manager, + clock=clock, + random_delay=lambda: 60, + sleep=capture_sleep, + requester=requester, + ) + await service._prepare_midnight( + "acct-a", "2026-10-24", datetime(2026, 10, 24, tzinfo=BEIJING) + ) + for _ in range(20): + if delays: + break + await asyncio.sleep(0) + assert delays == [0.0] + assert service.snapshot("acct-a")["next_checkin_at"] == datetime( + 2026, 10, 24, 0, 1, tzinfo=BEIJING + ).timestamp() + await service.close() + + run(scenario()) + + +def test_midnight_skips_already_checked_account_without_random_task(): + async def scenario(): + manager = FakeManager(enabled=True) + random_calls = [] + paths = [] + + async def requester(client, method, path): + paths.append(path) + return response(status_body(True)) + + service = AccountCheckinService( + manager, + clock=lambda: datetime(2026, 10, 24, 1, tzinfo=BEIJING), + random_delay=lambda: random_calls.append(True) or 60, + sleep=never_sleep, + requester=requester, + ) + await service._prepare_midnight( + "acct-a", "2026-10-24", datetime(2026, 10, 24, tzinfo=BEIJING) + ) + await asyncio.sleep(0) + assert paths == [STATUS_PATH] + assert random_calls == [] + assert service.snapshot("acct-a")["next_checkin_at"] is None + await service.close() + + run(scenario()) + + +def test_disabling_cancels_pending_random_timer_and_clears_due_time(): + async def scenario(): + manager = FakeManager(enabled=False) + sleep_started = asyncio.Event() + requested = [] + + async def requester(client, method, path): + requested.append(path) + return response(status_body(False)) + + async def blocked_sleep(seconds): + if seconds <= 600: + sleep_started.set() + await asyncio.Event().wait() + + service = AccountCheckinService( + manager, + clock=lambda: datetime(2026, 10, 24, 0, 1, tzinfo=BEIJING), + random_delay=lambda: 120, + sleep=blocked_sleep, + requester=requester, + ) + await service.start() + manager.settings["auto_checkin_enabled"] = True + await service._prepare_midnight( + "acct-a", "2026-10-24", datetime(2026, 10, 24, tzinfo=BEIJING) + ) + await wait_event(sleep_started) + due = service.snapshot("acct-a")["next_checkin_at"] + assert due is not None + await service.check_status("acct-a") + assert service.snapshot("acct-a")["next_checkin_at"] == due + manager.settings["auto_checkin_enabled"] = False + service.settings_changed(False) + await asyncio.sleep(0) + assert service.snapshot("acct-a")["next_checkin_at"] is None + assert requested == [STATUS_PATH, STATUS_PATH] + await service.close() + assert service._random_tasks == {} + assert service._tasks == set() + + run(scenario()) + + +def test_stop_synchronously_signals_and_cancels_pending_random_work(): + async def scenario(): + manager = FakeManager(enabled=False) + sleep_started = asyncio.Event() + paths = [] + + async def requester(client, method, path): + paths.append(path) + return response(status_body(False)) + + async def blocked_sleep(seconds): + if seconds <= 600: + sleep_started.set() + await asyncio.Event().wait() + + service = AccountCheckinService( + manager, + clock=lambda: datetime(2026, 10, 24, 0, 1, tzinfo=BEIJING), + random_delay=lambda: 120, + sleep=blocked_sleep, + requester=requester, + ) + await service.start() + manager.settings["auto_checkin_enabled"] = True + await service._prepare_midnight( + "acct-a", "2026-10-24", datetime(2026, 10, 24, tzinfo=BEIJING) + ) + await wait_event(sleep_started) + service.stop() + assert service._stop_event.is_set() + await asyncio.sleep(0) + assert paths == [STATUS_PATH] + await service.close() + assert service._tasks == set() + assert service._random_tasks == {} + + run(scenario()) + + +def test_two_concurrent_manual_checkins_send_only_one_post(): + async def scenario(): + manager = FakeManager() + status_calls = 0 + post_calls = 0 + post_started = asyncio.Event() + release_post = asyncio.Event() + + async def requester(client, method, path): + nonlocal status_calls, post_calls + if path == STATUS_PATH: + status_calls += 1 + return response(status_body(status_calls > 1)) + post_calls += 1 + post_started.set() + await release_post.wait() + return response({"code": 0, "data": {}}) + + service = AccountCheckinService(manager, requester=requester) + first = asyncio.create_task(service.check_in("acct-a")) + second = asyncio.create_task(service.check_in("acct-a")) + await wait_event(post_started) + release_post.set() + results = await asyncio.gather(first, second) + assert post_calls == 1 + assert status_calls == 3 # Before POST, post confirmation, serialized caller. + assert all(item["status"] == "checked" for item in results) + await service.close() + + run(scenario()) + + +def test_explicit_rejection_confirmed_unchecked_can_be_retried_once_later(): + async def scenario(): + manager = FakeManager() + status_calls = 0 + post_calls = 0 + + async def requester(client, method, path): + nonlocal status_calls, post_calls + if path == STATUS_PATH: + status_calls += 1 + return response(status_body(status_calls == 4)) + post_calls += 1 + if post_calls == 1: + return response({"code": 429, "msg": "limited"}, status=429) + return response({"code": 0, "data": {}}) + + service = AccountCheckinService(manager, requester=requester) + failed = await service.check_in("acct-a") + assert failed["status"] == "error" + assert failed["error"] == "签到请求失败" + assert ("acct-a", "2026-10-24") not in service._post_started + succeeded = await service.check_in("acct-a") + assert succeeded["status"] == "checked" + assert post_calls == 2 + assert status_calls == 4 + assert manager.rate_limits == [("acct-a", None, None, False)] + await service.close() + + run(scenario()) + + +def test_disable_while_post_waits_for_concurrency_slot_blocks_requester(): + async def scenario(): + manager = FakeManager(enabled=True) + service = None + hold_started = asyncio.Event() + release_hold = asyncio.Event() + post_calls = [] + + async def hold_slot_after_status(): + await service._semaphore.acquire() + hold_started.set() + try: + await release_hold.wait() + finally: + service._semaphore.release() + + async def requester(client, method, path): + if path == STATUS_PATH: + holder = asyncio.create_task(hold_slot_after_status()) + await asyncio.sleep(0) # Queue the holder ahead of the next request. + return response(status_body(False)) + post_calls.append(path) + return response({"code": 0, "data": {}}) + + service = AccountCheckinService( + manager, + clock=lambda: datetime(2026, 10, 24, 1, tzinfo=BEIJING), + requester=requester, + max_concurrent=1, + ) + task = service._spawn( + service._automatic_checkin("acct-a", "2026-10-24", "startup"), + key=("acct-a", "2026-10-24"), + ) + await wait_event(hold_started) + await asyncio.sleep(0) + assert ("acct-a", "2026-10-24") in service._post_started + manager.settings["auto_checkin_enabled"] = False + service._apply_settings_change(True) # Stale enable callback reads persisted off. + release_hold.set() + await task + assert post_calls == [] + assert ("acct-a", "2026-10-24") not in service._post_started + assert service.snapshot("acct-a")["status"] == "unchecked" + await service.close() + + run(scenario()) + + +def test_stale_settings_callback_uses_current_setting_to_cancel_random_task(): + async def scenario(): + manager = FakeManager(enabled=False) + sleep_started = asyncio.Event() + requested = [] + + async def requester(client, method, path): + requested.append(path) + return response(status_body(False)) + + async def blocked_sleep(seconds): + if seconds <= 600: + sleep_started.set() + await asyncio.Event().wait() + + service = AccountCheckinService( + manager, + clock=lambda: datetime(2026, 10, 24, 0, 1, tzinfo=BEIJING), + random_delay=lambda: 120, + sleep=blocked_sleep, + requester=requester, + ) + await service.start() + manager.settings["auto_checkin_enabled"] = True + await service._prepare_midnight( + "acct-a", "2026-10-24", datetime(2026, 10, 24, tzinfo=BEIJING) + ) + await wait_event(sleep_started) + assert service.snapshot("acct-a")["next_checkin_at"] is not None + manager.settings["auto_checkin_enabled"] = False + service._apply_settings_change(True) # Delayed callback from older state. + await asyncio.sleep(0) + assert service.snapshot("acct-a")["next_checkin_at"] is None + assert requested == [STATUS_PATH] + await service.close() + + run(scenario()) + + +def test_stale_off_callback_cannot_cancel_schedule_when_saved_setting_is_on(): + async def scenario(): + manager = FakeManager(enabled=False) + sleep_started = asyncio.Event() + scheduled = [] + + async def requester(client, method, path): + return response(status_body(False)) + + async def blocked_sleep(seconds): + if seconds <= 600: + sleep_started.set() + await asyncio.Event().wait() + + service = AccountCheckinService( + manager, + clock=lambda: datetime(2026, 10, 24, 0, 1, tzinfo=BEIJING), + random_delay=lambda: 120, + sleep=blocked_sleep, + requester=requester, + ) + await service.start() + service._schedule_catchup = lambda reason: scheduled.append(reason) + manager.settings["auto_checkin_enabled"] = True + await service._prepare_midnight( + "acct-a", "2026-10-24", datetime(2026, 10, 24, tzinfo=BEIJING) + ) + await wait_event(sleep_started) + before = service.snapshot("acct-a")["next_checkin_at"] + service._apply_settings_change(False) # Delayed off callback from before the final on. + assert scheduled == ["enable"] + assert service.snapshot("acct-a")["next_checkin_at"] == before + await service.close() + + run(scenario()) + + +def test_safety_logging_records_query_action_schedule_and_skips_without_payloads(): + async def scenario(): + manager = FakeManager(("acct-a", "acct-b"), enabled=True) + events = [] + status_calls = {"acct-a": 0, "acct-b": 0} + + async def requester(client, method, path): + if path == STATUS_PATH: + status_calls[client.account_id] += 1 + return response( + status_body( + client.account_id == "acct-a" + and status_calls[client.account_id] > 1 + ) + ) + return response({"code": 0, "data": {}}) + + service = AccountCheckinService( + manager, + clock=lambda: datetime(2026, 10, 24, 0, 1, tzinfo=BEIJING), + random_delay=lambda: 600, + sleep=never_sleep, + requester=requester, + ) + service.configure_logging(lambda event, **fields: events.append((event, fields))) + await service._prepare_midnight( + "acct-b", + "2026-10-24", + datetime(2026, 10, 24, tzinfo=BEIJING), + ) + await service.check_in("acct-a") + await service.check_in("acct-a") # Checked status is logged as an action skip. + names = [event for event, _ in events] + assert "account_checkin_status_check_started" in names + assert "account_checkin_scheduled" in names + assert "account_checkin_action_started" in names + assert "account_checkin_confirmed" in names + assert "account_checkin_action_skipped" in names + assert all( + set(fields).issubset( + {"account_id", "date", "reason", "status", "http_status", "next_checkin_at"} + ) + for _, fields in events + ) + await service.close() + + run(scenario()) + + +@pytest.mark.parametrize( + "payload,expected_error", + [ + ({"code": 10001, "msg": "already", "data": None}, "签到状态查询失败"), + ({"code": False, "data": {"today_checked_in": False}}, "签到状态查询失败"), + ({"code": 0, "data": {"today_checked_in": "false"}}, "签到状态查询失败"), + ( + {"code": 0, "data": {"today_checked_in": "false", "checked": True}}, + "签到状态查询失败", + ), + ( + {"code": 0, "data": {"active": False, "today_checked_in": False}}, + "签到活动未开放", + ), + ], +) +def test_unknown_or_inactive_status_never_posts(payload, expected_error): + async def scenario(): + manager = FakeManager() + paths = [] + + async def requester(client, method, path): + paths.append(path) + return response(payload) + + service = AccountCheckinService(manager, requester=requester) + result = await service.check_in("acct-a") + assert result["status"] == "error" + assert result["error"] == expected_error + assert result["checked"] is None + assert paths == [STATUS_PATH] + await service.close() + + run(scenario()) + + +def test_boolean_checked_is_strict_and_legacy_boolean_fixture_is_supported(): + async def scenario(): + manager = FakeManager() + bodies = [ + {"code": 0, "data": {"checked": True}}, + {"code": 0, "data": {"checked": 1}}, + ] + paths = [] + + async def requester(client, method, path): + paths.append(path) + return response(bodies.pop(0)) + + service = AccountCheckinService(manager, requester=requester) + old_fixture = await service.check_status("acct-a") + invalid_fixture = await service.check_status("acct-a") + assert old_fixture["status"] == "checked" + assert invalid_fixture["status"] == "error" + assert paths == [STATUS_PATH, STATUS_PATH] + await service.close() + + run(scenario()) + + +def test_disable_during_status_get_prevents_automatic_post(): + async def scenario(): + manager = FakeManager(enabled=True) + get_started = asyncio.Event() + release_get = asyncio.Event() + paths = [] + + async def requester(client, method, path): + paths.append(path) + if path == STATUS_PATH: + get_started.set() + await release_get.wait() + return response(status_body(False)) + return response({"code": 0, "data": {}}) + + service = AccountCheckinService(manager, requester=requester) + task = asyncio.create_task( + service._automatic_checkin("acct-a", "2026-10-24", "startup") + ) + await wait_event(get_started) + manager.settings["auto_checkin_enabled"] = False + release_get.set() + await task + assert paths == [STATUS_PATH] + await service.close() + + run(scenario()) + + +def test_cancelled_automatic_status_query_clears_busy_state_and_allows_manual_checkin(): + async def scenario(): + manager = FakeManager(enabled=False) + get_started = asyncio.Event() + status_calls = 0 + post_calls = 0 + + async def requester(client, method, path): + nonlocal status_calls, post_calls + if path == STATUS_PATH: + status_calls += 1 + if status_calls == 1: + get_started.set() + await asyncio.Event().wait() + return response(status_body(status_calls > 2)) + post_calls += 1 + return response({"code": 0, "data": {}}) + + service = AccountCheckinService( + manager, + clock=lambda: datetime(2026, 10, 24, 1, tzinfo=BEIJING), + requester=requester, + sleep=never_sleep, + ) + await service.start() + manager.settings["auto_checkin_enabled"] = True + task = service._spawn( + service._automatic_checkin("acct-a", "2026-10-24", "startup"), + key=("acct-a", "2026-10-24"), + ) + await wait_event(get_started) + assert service.snapshot("acct-a")["status"] == "checking" + manager.settings["auto_checkin_enabled"] = False + service.settings_changed(False) + await asyncio.gather(task, return_exceptions=True) + cancelled_state = service.snapshot("acct-a") + assert cancelled_state["status"] == "unknown" + assert cancelled_state["next_checkin_at"] is None + assert post_calls == 0 + + manual = await service.check_in("acct-a") + assert manual["status"] == "checked" + assert post_calls == 1 + await service.close() + + run(scenario()) + + +@pytest.mark.parametrize("invalidate", ["deleted", "next-day"]) +def test_random_followup_rechecks_account_and_target_date(invalidate): + async def scenario(): + manager = FakeManager(enabled=True) + sleep_started = asyncio.Event() + release_sleep = asyncio.Event() + paths = [] + clock = MutableClock(datetime(2026, 10, 24, 0, 1, tzinfo=BEIJING)) + + async def requester(client, method, path): + paths.append(path) + return response(status_body(False)) + + async def gated_sleep(_seconds): + sleep_started.set() + await release_sleep.wait() + + service = AccountCheckinService( + manager, + clock=clock, + random_delay=lambda: 120, + sleep=gated_sleep, + requester=requester, + ) + await service._prepare_midnight( + "acct-a", "2026-10-24", datetime(2026, 10, 24, tzinfo=BEIJING) + ) + await wait_event(sleep_started) + if invalidate == "deleted": + manager.store.clients.pop("acct-a") + else: + clock.advance(24 * 60 * 60) + release_sleep.set() + if invalidate == "deleted": + release_sleep.set() + await asyncio.sleep(0) + await asyncio.sleep(0) + assert paths == [STATUS_PATH] + if invalidate == "next-day": + assert service.snapshot("acct-a")["date"] == "2026-10-25" + assert service.snapshot("acct-a")["status"] == "unknown" + await service.close() + + run(scenario()) + + +def test_default_billing_requester_is_bound_to_target_and_records_429_without_switch(monkeypatch): + manager = FakeManager(("acct-a", "acct-b"), active="acct-a") + target = manager.store.clients["acct-b"] + auth_options = [] + + class BillingClient(FakeClient): + endpoint = "https://billing.example" + + def build_url(self, path): + return self.endpoint + path + + def ensure_authenticated(self, **kwargs): + auth_options.append((self.account_id, kwargs)) + + def auth_headers(self): + return {"Authorization": f"Bearer {self.account_id}"} + + manager.store.clients["acct-b"] = BillingClient("acct-b") + client = manager.store.clients["acct-b"] + real_async_client = httpx.AsyncClient + requests = [] + + def handler(request): + requests.append(request) + return httpx.Response( + 429, + content=b'{"code":429,"msg":"slow down"}', + headers={"content-type": "application/json", "Retry-After": "37"}, + ) + + transport = httpx.MockTransport(handler) + monkeypatch.setattr( + billing.httpx, + "AsyncClient", + lambda **kwargs: real_async_client(transport=transport, **kwargs), + ) + + async def scenario(): + service = AccountCheckinService(manager) + result = await service.check_status("acct-b") + assert result["status"] == "error" + assert result["error"] == "签到状态查询失败" + await service.close() + + run(scenario()) + assert len(requests) == 1 + assert requests[0].url.path == STATUS_PATH + assert requests[0].headers["authorization"] == "Bearer acct-b" + assert auth_options == [ + ("acct-b", {"open_browser": False, "allow_login": False}) + ] + assert manager.rate_limits == [("acct-b", "37", None, False)] + assert manager.active_account_id == "acct-a" diff --git a/test_account_checkin_integration.py b/test_account_checkin_integration.py new file mode 100644 index 0000000..6c5fe4c --- /dev/null +++ b/test_account_checkin_integration.py @@ -0,0 +1,240 @@ +"""Exercise per-account check-in and automatic triggers through the real app.""" +import asyncio +import json +import threading +import time + +import httpx +import pytest +from fastapi.testclient import TestClient + +import codebuddy_proxy.__main__ as proxy +from codebuddy_proxy.account_manager import AccountManager +from codebuddy_proxy.account_store import AccountStore +from test_account_integration import make_client +from codebuddy_proxy.backend_profile import DOMESTIC_PROFILE, GLOBAL_PROFILE + +ADMIN_HEADERS = {"X-Account-Manager": "1"} + + +@pytest.fixture +def checkin_environment(tmp_path, monkeypatch): + store = AccountStore(tmp_path / "accounts") + ids = {} + for profile, uid in [(DOMESTIC_PROFILE, "first"), (GLOBAL_PROFILE, "second")]: + source = make_client(tmp_path, profile, uid, token="PRIVATE-" + uid) + ids[uid] = store.import_client(source) + manager = AccountManager(store) + state = proxy.ProxyState(manager.client, None, None, account_manager=manager) + monkeypatch.setattr(proxy, "proxy_state", state) + real_async_client = httpx.AsyncClient + seen = [] + signed = set() + + async def handler(request): + uid = request.headers["X-User-Id"] + seen.append((uid, request.url.host, request.url.path, request.method, request.content)) + assert request.content == b"{}" + if request.url.path == "/v2/billing/meter/checkin-activity-status": + return httpx.Response(200, json={ + "code": 0, "msg": "OK", + "data": {"active": True, "today_checked_in": uid in signed}, + }) + assert request.url.path.endswith("/daily-checkin") + signed.add(uid) + return httpx.Response(200, json={"code": 0, "msg": "OK", "data": {}}) + + def install(custom_handler=None): + monkeypatch.setattr( + httpx, "AsyncClient", + lambda **kw: real_async_client( + transport=httpx.MockTransport(custom_handler or handler), **kw, + ), + ) + install() + yield manager, store, ids, seen, signed, install, handler + manager.close() + + +def app_client(): + return TestClient(proxy.app, base_url="http://127.0.0.1", client=("127.0.0.1", 50500)) + + +def test_checkin_of_inactive_account_preserves_active_and_uses_its_own_credentials( + checkin_environment, capsys +): + manager, store, ids, seen, signed, _, _ = checkin_environment + active = store.active_account_id + with app_client() as app: + response = app.get(f"/admin/api/accounts/{ids['second']}/checkin-status") + assert response.status_code == 200 + assert response.json()["checked"] is False + response = app.post(f"/admin/api/accounts/{ids['second']}/checkin", headers=ADMIN_HEADERS) + assert response.status_code == 200 + assert response.json()["status"] == "checked" + assert response.json()["checked"] is True + row = next(row for row in app.get("/admin/api/accounts").json()["accounts"] + if row["id"] == ids["second"]) + assert row["checkin"]["checked"] is True + assert store.active_account_id == active + assert signed == {"second"} + assert {uid for uid, *_ in seen} == {"second"} + assert {host for _, host, *_ in seen} == {"www.codebuddy.ai"} + assert sum(path.endswith("/daily-checkin") for _, _, path, *_ in seen) == 1 + output = capsys.readouterr().out + assert "[Checkin]" in output + assert "PRIVATE-" not in output + + +@pytest.mark.parametrize("uid", ["first", "second"], ids=["domestic", "global"]) +@pytest.mark.parametrize("entrypoint", ["admin", "v1-get", "v1-post"]) +def test_current_activity_status_reports_already_signed_without_daily_post( + checkin_environment, uid, entrypoint +): + """Legacy status may say inactive while the current activity is signed.""" + manager, store, ids, seen, _, install, _ = checkin_environment + if entrypoint != "admin": + manager.activate(ids[uid]) + active = store.active_account_id + current_payload = { + "code": 0, + "msg": "OK", + "data": { + "active": True, + "today_checked_in": True, + "streak_days": 1, + "start_time": "2026-09-30 00:00:00", + "end_time": "2026-12-31 23:59:59", + "season": 10, + }, + } + + async def activity_status(request): + seen.append((request.headers["X-User-Id"], request.url.host, + request.url.path, request.method, request.content)) + assert request.method == "POST" + assert request.content == b"{}" + if request.url.path == "/v2/billing/meter/checkin-status": + return httpx.Response(200, json={ + "code": 0, "msg": "OK", + "data": {"active": False, "today_checked_in": False, + "start_time": "", "end_time": "", "season": 0}, + }) + assert request.url.path == "/v2/billing/meter/checkin-activity-status" + return httpx.Response(200, json=current_payload) + + install(activity_status) + with app_client() as app: + if entrypoint == "admin": + result = app.get(f"/admin/api/accounts/{ids[uid]}/checkin-status") + assert result.status_code == 200 + assert result.json()["status"] == "checked" + assert result.json()["checked"] is True + assert result.json()["error"] is None + row = next(a for a in app.get("/admin/api/accounts").json()["accounts"] + if a["id"] == ids[uid]) + assert row["checkin"]["status"] == "checked" + # A manual sign request must also recognize today's existing sign. + result = app.post(f"/admin/api/accounts/{ids[uid]}/checkin", + headers=ADMIN_HEADERS) + assert result.status_code == 200 + assert result.json()["status"] == "checked" + else: + method = "GET" if entrypoint == "v1-get" else "POST" + result = app.request(method, "/v1/checkin-status") + assert result.status_code == 200 + assert result.json() == current_payload + assert store.active_account_id == active + + assert seen + assert {user for user, *_ in seen} == {uid} + expected_host = "copilot.tencent.com" if uid == "first" else "www.codebuddy.ai" + assert {host for _, host, *_ in seen} == {expected_host} + assert {path for _, _, path, *_ in seen} == { + "/v2/billing/meter/checkin-activity-status" + } + + +def test_signed_account_is_not_signed_again(checkin_environment): + _, _, ids, seen, signed, _, _ = checkin_environment + signed.add("first") + with app_client() as app: + response = app.post(f"/admin/api/accounts/{ids['first']}/checkin", headers=ADMIN_HEADERS) + assert response.status_code == 200 + assert response.json()["checked"] is True + assert not any(path.endswith("/daily-checkin") for _, _, path, *_ in seen) + + +@pytest.mark.parametrize("startup", [False, True], ids=["enable-immediately", "startup-enabled"]) +def test_automatic_trigger_checks_all_accounts_including_inactive( + checkin_environment, startup +): + manager, store, ids, seen, signed, _, _ = checkin_environment + active = store.active_account_id + if startup: + manager.set_settings(auto_checkin_enabled=True) + with app_client() as app: + if not startup: + response = app.patch("/admin/api/settings", headers=ADMIN_HEADERS, + json={"auto_checkin_enabled": True}) + assert response.status_code == 200 + assert response.json()["auto_checkin_enabled"] is True + assert response.json()["auto_switch_on_429"] is False + deadline = time.monotonic() + 5 + while time.monotonic() < deadline and signed != {"first", "second"}: + time.sleep(0.01) + assert signed == {"first", "second"} + assert store.active_account_id == active + assert sum(path.endswith("/daily-checkin") for _, _, path, *_ in seen) == 2 + + +def test_auto_switch_429_does_not_change_active_for_per_account_checkin(checkin_environment, tmp_path): + manager, store, ids, seen, _, install, _ = checkin_environment + active = store.active_account_id + store.import_client(make_client(tmp_path, DOMESTIC_PROFILE, "backup")) + manager.set_settings(auto_switch_on_429=True) + + async def limited(request): + seen.append((request.headers["X-User-Id"], request.url.host, + request.url.path, request.method, request.content)) + return httpx.Response(429, json={"code": 429}, headers={"Retry-After": "19"}) + + install(limited) + with app_client() as app: + response = app.get(f"/admin/api/accounts/{ids['first']}/checkin-status") + assert response.status_code == 200 + assert response.json()["checked"] is None + assert response.json()["status"] == "error" + assert store.active_account_id == active + first = next(row for row in manager.snapshot()["accounts"] if row["id"] == active) + assert first["model_cooldowns"] == [] + assert len(seen) == 1 + + +def test_disabling_while_status_is_pending_prevents_auto_sign(checkin_environment): + manager, _, _, seen, signed, install, original = checkin_environment + status_entered = threading.Event() + release_status = threading.Event() + + async def pending(request): + if request.url.path == "/v2/billing/meter/checkin-activity-status": + status_entered.set() + while not release_status.is_set(): + await asyncio.sleep(0.01) + return await original(request) + + install(pending) + try: + with app_client() as app: + app.patch("/admin/api/settings", headers=ADMIN_HEADERS, + json={"auto_checkin_enabled": True}) + assert status_entered.wait(3) + response = app.patch("/admin/api/settings", headers=ADMIN_HEADERS, + json={"auto_checkin_enabled": False}) + assert response.status_code == 200 + release_status.set() + time.sleep(0.1) + assert not signed + finally: + release_status.set() + assert not any(path.endswith("/daily-checkin") for _, _, path, *_ in seen) diff --git a/test_account_integration.py b/test_account_integration.py new file mode 100644 index 0000000..3cd4564 --- /dev/null +++ b/test_account_integration.py @@ -0,0 +1,289 @@ +"""Account changes through the real proxy and startup boundary (offline).""" + +import asyncio +import json +from pathlib import Path + +import httpx +import pytest +from fastapi.testclient import TestClient + +import codebuddy_proxy.__main__ as proxy +import codebuddy_proxy.billing_proxy as billing +from codebuddy_proxy.account_manager import AccountManager +from codebuddy_proxy.account_store import AccountStore +from codebuddy_proxy.backend_profile import DOMESTIC_PROFILE, GLOBAL_PROFILE, resolve_backend +from codebuddy_proxy.codebuddy_client_demo import CodeBuddyClient + + +def make_client(tmp_path, profile, uid, token=None): + client = CodeBuddyClient(backend=resolve_backend( + profile=profile, session_file=tmp_path / f"{profile.key}-{uid}.json", + )) + client._save_session({ + "auth": {"accessToken": token or f"token-{profile.key}-{uid}"}, + "account": {"uid": uid, "nickname": f"User {uid}"}, + "machineId": "test-machine", + }) + return client + + +@pytest.fixture +def managed(tmp_path, monkeypatch): + store = AccountStore(tmp_path / "accounts") + domestic = make_client(tmp_path, DOMESTIC_PROFILE, "same-uid") + global_client = make_client(tmp_path, GLOBAL_PROFILE, "same-uid") + first = store.import_client(domestic) + second = store.import_client(global_client) + manager = AccountManager(store) + state = proxy.ProxyState(domestic, None, None, account_manager=manager) + monkeypatch.setattr(proxy, "proxy_state", state) + yield store, manager, state, first, second + manager.close() + + +def test_switch_updates_health_models_and_preserves_captured_request(managed): + store, manager, state, first, second = managed + snapshot = state.for_request() + with TestClient(proxy.app) as app: + before = app.get("/v1/models").json() + assert app.get("/health").json()["backend"] == "domestic" + manager.activate(second) + assert app.get("/health").json()["backend"] == "global" + assert app.get("/v1/models").json() != before + assert snapshot.client.profile.key == "domestic" + assert state.for_request().client.profile.key == "global" + assert snapshot.client.session["auth"]["accessToken"] == "token-domestic-same-uid" + assert AccountStore(store.root).active_account_id == second + + +@pytest.mark.parametrize("path,body", [ + ("/v1/chat/completions", {"messages": [{"role": "user", "content": "hello"}]}), + ("/v1/responses", {"input": "hello"}), + ("/v1/messages", {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 10}), +]) +def test_three_protocols_switch_endpoint_and_credentials(managed, monkeypatch, path, body): + store, manager, state, first, second = managed + seen = [] + real_client = httpx.AsyncClient + + async def upstream(request): + seen.append((request.url.host, request.headers["authorization"])) + data = {"id": "test", "model": "test", "choices": [ + {"index": 0, "delta": {"content": "OK"}, "finish_reason": None}, + ]} + return httpx.Response(200, content=("data: " + json.dumps(data) + "\n\ndata: [DONE]\n\n").encode(), + headers={"content-type": "text/event-stream"}) + + monkeypatch.setattr(proxy.httpx, "AsyncClient", lambda **kw: real_client( + transport=httpx.MockTransport(upstream), **kw, + )) + with TestClient(proxy.app) as app: + assert app.post(path, json={"model": "test", "stream": False, **body}).status_code == 200 + manager.activate(second) + assert app.post(path, json={"model": "test", "stream": False, **body}).status_code == 200 + assert seen == [("copilot.tencent.com", "Bearer token-domestic-same-uid"), + ("www.codebuddy.ai", "Bearer token-global-same-uid")] + + +def test_billing_refresh_during_switch_stays_with_original_account(managed, monkeypatch): + store, manager, state, first, second = managed + old_client = state.client + seen = [] + refreshed = [] + real_client = httpx.AsyncClient + + def refresh(): + refreshed.append(old_client.profile.key) + old_client._save_session({**old_client.session, "auth": {"accessToken": "refreshed-domestic"}}) + return True + + monkeypatch.setattr(old_client, "refresh", refresh) + + async def upstream(request): + seen.append((request.url.host, request.headers["authorization"])) + if len(seen) == 1: + manager.activate(second) + return httpx.Response(401, json={"error": "expired"}) + return httpx.Response(200, json={"ok": True}) + + monkeypatch.setattr(billing.httpx, "AsyncClient", lambda **kw: real_client( + transport=httpx.MockTransport(upstream), **kw, + )) + with TestClient(proxy.app) as app: + assert app.post("/v1/checkin").status_code == 200 + assert refreshed == ["domestic"] + assert seen == [("copilot.tencent.com", "Bearer token-domestic-same-uid"), + ("copilot.tencent.com", "Bearer refreshed-domestic")] + assert state.client.profile.key == "global" + + +def test_account_switch_while_streaming_preserves_original_request(managed, monkeypatch): + store, manager, state, first, second = managed + seen = [] + real_client = httpx.AsyncClient + + class SwitchingStream(httpx.AsyncByteStream): + async def __aiter__(self): + chunk = {"id": "stream-test", "model": "test", "choices": [ + {"index": 0, "delta": {"content": "old-account"}, "finish_reason": None}, + ]} + yield ("data: " + json.dumps(chunk) + "\n\n").encode() + manager.activate(second) + yield b'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n' + yield b"data: [DONE]\n\n" + + async def upstream(request): + seen.append((request.url.host, request.headers["authorization"])) + return httpx.Response(200, stream=SwitchingStream(), headers={"content-type": "text/event-stream"}) + + monkeypatch.setattr(proxy.httpx, "AsyncClient", lambda **kw: real_client( + transport=httpx.MockTransport(upstream), **kw, + )) + with TestClient(proxy.app) as app: + body = {"model": "test", "stream": True, "messages": [{"role": "user", "content": "hi"}]} + response = app.post("/v1/chat/completions", json=body) + assert response.status_code == 200 + assert "old-account" in response.text + assert "[DONE]" in response.text + assert state.client.profile.key == "global" + assert app.post("/v1/chat/completions", json=body).status_code == 200 + assert seen == [("copilot.tencent.com", "Bearer token-domestic-same-uid"), + ("www.codebuddy.ai", "Bearer token-global-same-uid")] + + +@pytest.mark.parametrize("path,body", [ + ("/v1/chat/completions", {"messages": [{"role": "user", "content": "hello"}]}), + ("/v1/responses", {"input": "hello"}), + ("/v1/messages", {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 10}), +]) +def test_request_pins_account_before_body_parsing(managed, monkeypatch, path, body): + store, manager, state, first, second = managed + original_parser = proxy.parse_request_body + seen = [] + real_client = httpx.AsyncClient + + async def switch_during_body_read(request): + manager.activate(second) + return await original_parser(request) + + async def upstream(request): + seen.append((request.url.host, request.headers["authorization"])) + return httpx.Response(200, content=b'data: {"choices":[{"index":0,"delta":{"content":"OK"}}]}\n\ndata: [DONE]\n\n') + + monkeypatch.setattr(proxy, "parse_request_body", switch_during_body_read) + monkeypatch.setattr(proxy.httpx, "AsyncClient", lambda **kw: real_client( + transport=httpx.MockTransport(upstream), **kw, + )) + with TestClient(proxy.app) as app: + assert app.post(path, json={"model": "test", "stream": False, **body}).status_code == 200 + assert seen == [("copilot.tencent.com", "Bearer token-domestic-same-uid")] + assert manager.client.profile.key == "global" + + +def test_empty_account_starts_and_requests_do_not_open_browser(tmp_path, monkeypatch): + monkeypatch.setattr(Path, "home", lambda: tmp_path) + monkeypatch.delenv("CODEBUDDY_ENDPOINT", raising=False) + args = proxy.parse_args(["--accounts-dir", str(tmp_path / "accounts")]) + manager, fallback = proxy.initialize_account_manager(args, resolve_backend()) + monkeypatch.setattr(proxy, "proxy_state", proxy.ProxyState(fallback, None, None, account_manager=manager)) + monkeypatch.setattr(CodeBuddyClient, "login", lambda *a, **kw: pytest.fail("unexpected browser login")) + with TestClient(proxy.app) as app: + assert app.get("/health").status_code == 200 + assert app.get("/health").json()["authenticated"] is False + assert app.get("/v1/models").status_code == 200 + response = app.post("/v1/chat/completions", json={"model": "test", "messages": []}) + assert response.status_code == 401 + assert "authentication_error" in response.text + + +def test_legacy_import_once_restore_and_explicit_region(tmp_path, monkeypatch): + monkeypatch.setattr(Path, "home", lambda: tmp_path) + monkeypatch.delenv("CODEBUDDY_ENDPOINT", raising=False) + legacy = CodeBuddyClient(backend=resolve_backend()) + legacy._save_session({"auth": {"accessToken": "legacy"}, "account": {"uid": "old"}}) + original = legacy.session_file.read_bytes() + args = proxy.parse_args(["--accounts-dir", str(tmp_path / "accounts")]) + manager, client = proxy.initialize_account_manager(args, resolve_backend()) + assert client.session["auth"]["accessToken"] == "legacy" + client._save_session({**client.session, "auth": {"accessToken": "fresh"}}) + manager.close() + manager, client = proxy.initialize_account_manager(args, resolve_backend()) + assert client.session["auth"]["accessToken"] == "fresh" + assert legacy.session_file.read_bytes() == original + manager.close() + global_args = proxy.parse_args(["--global", "--accounts-dir", str(tmp_path / "accounts")]) + manager, client = proxy.initialize_account_manager(global_args, resolve_backend(global_mode=True)) + assert manager.client is None + assert client.profile.key == "global" + assert len(manager.snapshot()["accounts"]) == 1 + manager.close() + + +def test_cli_login_preserves_legacy_and_existing_accounts(tmp_path, monkeypatch): + monkeypatch.setattr(Path, "home", lambda: tmp_path) + monkeypatch.delenv("CODEBUDDY_ENDPOINT", raising=False) + legacy = CodeBuddyClient(backend=resolve_backend()) + legacy._save_session({"auth": {"accessToken": "old-token"}, "account": {"uid": "old"}}) + original = legacy.session_file.read_bytes() + + def login(client, **kwargs): + client._save_session({"auth": {"accessToken": "new-token"}, "account": {"uid": "new"}}) + + monkeypatch.setattr(CodeBuddyClient, "login", login) + args = proxy.parse_args(["--login", "--no-browser", "--accounts-dir", str(tmp_path / "accounts")]) + manager, client = proxy.initialize_account_manager(args, resolve_backend()) + assert client.session["account"]["uid"] == "new" + assert len(manager.snapshot()["accounts"]) == 2 + assert legacy.session_file.read_bytes() == original + manager.close() + + +def test_cli_login_with_no_legacy_session_or_existing_directory(tmp_path, monkeypatch): + monkeypatch.setattr(Path, "home", lambda: tmp_path) + monkeypatch.delenv("CODEBUDDY_ENDPOINT", raising=False) + + def login(client, **kwargs): + client._save_session({"auth": {"accessToken": "first-token"}, "account": {"uid": "first"}}) + + monkeypatch.setattr(CodeBuddyClient, "login", login) + args = proxy.parse_args(["--login", "--accounts-dir", str(tmp_path / "new" / "accounts")]) + manager, client = proxy.initialize_account_manager(args, resolve_backend()) + assert client.session["account"]["uid"] == "first" + manager.close() + + +def test_deleted_legacy_account_is_not_reimported_on_restart(tmp_path, monkeypatch): + monkeypatch.setattr(Path, "home", lambda: tmp_path) + monkeypatch.delenv("CODEBUDDY_ENDPOINT", raising=False) + legacy = CodeBuddyClient(backend=resolve_backend()) + legacy._save_session({"auth": {"accessToken": "old-token"}, "account": {"uid": "old"}}) + args = proxy.parse_args(["--accounts-dir", str(tmp_path / "accounts")]) + manager, _ = proxy.initialize_account_manager(args, resolve_backend()) + old_id = manager.store.active_account_id + other = make_client(tmp_path, GLOBAL_PROFILE, "other") + new_id = manager.store.import_client(other, activate=True) + manager.delete(old_id) + manager.close() + manager, client = proxy.initialize_account_manager(args, resolve_backend()) + assert manager.store.active_account_id == new_id + assert len(manager.snapshot()["accounts"]) == 1 + assert client.profile.key == "global" + manager.close() + + +def test_explicit_endpoint_ignores_unmatched_default_legacy_session(tmp_path, monkeypatch): + monkeypatch.setattr(Path, "home", lambda: tmp_path) + monkeypatch.delenv("CODEBUDDY_ENDPOINT", raising=False) + legacy = CodeBuddyClient(backend=resolve_backend()) + legacy._save_session({"auth": {"accessToken": "old-token"}, "account": {"uid": "old"}}) + original = legacy.session_file.read_bytes() + args = proxy.parse_args(["--endpoint", "https://staging.example.test", + "--accounts-dir", str(tmp_path / "accounts")]) + manager, fallback = proxy.initialize_account_manager(args, proxy.resolve_backend_from_args(args)) + assert manager.client is None + assert fallback.endpoint == "https://staging.example.test" + assert not fallback.session + assert len(manager.snapshot()["accounts"]) == 1 + assert legacy.session_file.read_bytes() == original + manager.close() diff --git a/test_account_manager.py b/test_account_manager.py new file mode 100644 index 0000000..d96a4b6 --- /dev/null +++ b/test_account_manager.py @@ -0,0 +1,955 @@ +import threading +import time +import json +from datetime import datetime, timezone +from email.utils import format_datetime +from types import SimpleNamespace + +import pytest + +from codebuddy_proxy.account_manager import AccountManager, AccountManagerError +from codebuddy_proxy.account_store import AccountStore +from codebuddy_proxy.codebuddy_client_demo import CodeBuddyClient + + +class FakeStore: + def __init__(self, root): + self.root = root + self._accounts = [] + self._clients = {} + self._active_id = None + self.import_calls = [] + self.set_active_calls = [] + + @property + def active_account_id(self): + return self._active_id + + def list_accounts(self): + return [ + {key: account[key] for key in ("id", "backend", "endpoint", "uid", "nickname", "active")} + for account in self._accounts + ] + + def active_client(self): + return self._clients.get(self._active_id) + + def get_client(self, account_id): + return self._clients[account_id] + + def import_client(self, client, *, activate=False): + self.import_calls.append(activate) + profile = client.backend + endpoint = client.endpoint + account = client.session["account"] + existing = next( + ( + item + for item in self._accounts + if item["backend"] == profile + and item["endpoint"] == endpoint + and item["uid"] == str(account["uid"]) + ), + None, + ) + if existing is None: + account_id = f"account-{len(self._accounts) + 1}" + existing = { + "id": account_id, + "backend": profile, + "endpoint": endpoint, + "uid": str(account["uid"]), + "nickname": account.get("nickname"), + "active": False, + } + self._accounts.append(existing) + else: + account_id = existing["id"] + self._clients[account_id] = client + if activate or self._active_id is None: + self.set_active(account_id) + return account_id + + def set_active(self, account_id): + if account_id not in self._clients: + raise KeyError(account_id) + self.set_active_calls.append(account_id) + self._active_id = account_id + for account in self._accounts: + account["active"] = account["id"] == account_id + return self._clients[account_id] + + def delete_account(self, account_id): + if account_id not in self._clients: + raise KeyError(account_id) + self._accounts = [item for item in self._accounts if item["id"] != account_id] + self._clients.pop(account_id) + + +class FakeLoginClient: + def __init__(self, backend, session_file, *, uid, begin_gate=None, complete_gate=None): + self.backend = backend + self.endpoint = f"https://{backend}.example.test" + self.session_file = session_file + self.uid = uid + self.begin_gate = begin_gate + self.complete_gate = complete_gate + self.begin_started = threading.Event() + self.complete_started = threading.Event() + self.session = {} + self.completed_challenge = None + self.cancel_event = None + + def begin_login(self): + self.begin_started.set() + if self.begin_gate is not None: + self.begin_gate.wait(5) + return {"auth_url": "https://login.example.test/authorize?flow=abc", "state": "opaque-state"} + + def complete_login(self, challenge, *, timeout, cancel_event): + self.complete_started.set() + self.completed_challenge = challenge + self.cancel_event = cancel_event + if self.complete_gate is not None: + self.complete_gate.wait(5) + # Deliberately finish successfully even if cancellation won. The + # manager must still prevent this temporary session from being saved. + self.session = { + "auth": {"accessToken": "private-token-value"}, + "account": {"uid": self.uid, "nickname": f"User {self.uid}"}, + } + + +class StoredClient: + def __init__(self, backend, uid, *, endpoint=None, auth=None): + self.backend = backend + self.profile = SimpleNamespace(key=backend) + self.endpoint = endpoint or f"https://{backend}.example.test" + self.session = { + "auth": auth if auth is not None else {"accessToken": f"token-{uid}"}, + "account": {"uid": uid, "nickname": f"User {uid}"}, + } + + def _snapshot_session(self): + return self.session, 1 + + +def _wait_for(manager, job_id, status): + deadline = time.monotonic() + 3 + while time.monotonic() < deadline: + result = manager.get_login(job_id) + if result["status"] == status: + return result + time.sleep(0.01) + raise AssertionError(f"login did not reach {status}: {manager.get_login(job_id)}") + + +def _model_cooldown(account, model): + return next( + (item for item in account.get("model_cooldowns", []) if item["model"] == model), + None, + ) + + +def test_login_imports_only_after_success_and_does_not_switch_existing_account(tmp_path): + store = FakeStore(tmp_path) + made_clients = [] + + def make_client(backend, session_file): + client = FakeLoginClient(backend, session_file, uid=f"u{len(made_clients) + 1}") + made_clients.append(client) + return client + + manager = AccountManager(store, client_factory=make_client) + try: + first = manager.start_login("domestic") + first_result = _wait_for(manager, first["id"], "succeeded") + assert first_result["account_id"] == "account-1" + assert manager.snapshot()["active_account_id"] == "account-1" + assert manager.cancel_login(first["id"])["status"] == "succeeded" + + second = manager.start_login("global") + second_result = _wait_for(manager, second["id"], "succeeded") + assert second_result["account_id"] == "account-2" + assert manager.snapshot()["active_account_id"] == "account-1" + assert store.import_calls == [True, False] + assert manager.client is store.active_client() + assert all("state" not in result for result in (first_result, second_result)) + assert all("session" not in result and "token" not in result for result in (first_result, second_result)) + assert tmp_path in made_clients[0].session_file.parents + assert not made_clients[0].session_file.exists() + finally: + manager.close() + + +def test_cancel_wins_if_it_happens_before_login_commit(tmp_path): + store = FakeStore(tmp_path) + completion_gate = threading.Event() + client = FakeLoginClient( + "domestic", tmp_path / "unused", uid="cancelled-user", complete_gate=completion_gate + ) + manager = AccountManager(store, client_factory=lambda backend, path: client) + try: + job = manager.start_login("domestic") + assert client.complete_started.wait(2) + cancelled = manager.cancel_login(job["id"]) + assert cancelled["status"] == "cancelled" + assert client.cancel_event.is_set() + + completion_gate.set() + time.sleep(0.05) + assert manager.get_login(job["id"])["status"] == "cancelled" + assert store.import_calls == [] + assert manager.snapshot()["accounts"] == [] + finally: + completion_gate.set() + manager.close() + + +def test_close_cancels_active_login_and_prevents_store_import(tmp_path): + store = FakeStore(tmp_path) + client = FakeLoginClient("domestic", tmp_path / "unused", uid="closed-user") + + def wait_until_cancelled(challenge, *, timeout, cancel_event): + client.complete_started.set() + client.completed_challenge = challenge + client.cancel_event = cancel_event + assert cancel_event.wait(2) + # Even a completion racing with shutdown must not enter the real store. + client.session = { + "auth": {"accessToken": "private-token-value"}, + "account": {"uid": client.uid}, + } + + client.complete_login = wait_until_cancelled + manager = AccountManager(store, client_factory=lambda backend, path: client) + job = manager.start_login("domestic") + assert client.complete_started.wait(2) + manager.close() + + assert manager.get_login(job["id"])["status"] == "cancelled" + assert store.import_calls == [] + + +def test_login_failure_does_not_expose_exception_text(tmp_path): + store = FakeStore(tmp_path) + + class FailingClient(FakeLoginClient): + def begin_login(self): + raise RuntimeError("provider returned private-token-value and raw payload") + + client = FailingClient("domestic", tmp_path / "unused", uid="never-saved") + manager = AccountManager(store, client_factory=lambda backend, path: client) + try: + job = manager.start_login("domestic") + result = _wait_for(manager, job["id"], "failed") + assert result["error"] == "登录失败,请稍后重新尝试" + assert "private-token-value" not in repr(result) + assert "raw payload" not in repr(result) + assert store.import_calls == [] + finally: + manager.close() + + +def test_login_concurrency_is_bounded_and_invalid_backend_is_rejected(tmp_path): + store = FakeStore(tmp_path) + begin_gate = threading.Event() + client = FakeLoginClient("domestic", tmp_path / "unused", uid="u1", begin_gate=begin_gate) + manager = AccountManager( + store, + max_concurrent_logins=1, + client_factory=lambda backend, path: client, + ) + try: + job = manager.start_login("domestic") + assert client.begin_started.wait(2) + with pytest.raises(AccountManagerError) as capacity_error: + manager.start_login("global") + assert capacity_error.value.status_code == 429 + with pytest.raises(AccountManagerError): + manager.start_login("unknown") + manager.cancel_login(job["id"]) + finally: + begin_gate.set() + manager.close() + + +def test_cancelled_but_running_factory_keeps_its_concurrency_slot(tmp_path): + store = FakeStore(tmp_path) + factory_started = threading.Event() + factory_gate = threading.Event() + made_clients = [] + + def blocking_factory(backend, session_file): + factory_started.set() + factory_gate.wait(5) + client = FakeLoginClient(backend, session_file, uid="u1") + made_clients.append(client) + return client + + manager = AccountManager( + store, + max_concurrent_logins=1, + client_factory=blocking_factory, + ) + try: + job = manager.start_login("domestic") + assert factory_started.wait(2) + assert manager.cancel_login(job["id"])["status"] == "cancelled" + + # Cancellation marks the public state immediately, but the blocked + # worker still occupies the only slot until its future actually ends. + for _ in range(3): + with pytest.raises(AccountManagerError) as capacity_error: + manager.start_login("global") + assert capacity_error.value.status_code == 429 + + factory_gate.set() + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + with manager._lock: + original = manager._jobs[job["id"]] + if original.future.done(): + break + time.sleep(0.01) + assert original.future.done() + assert len(made_clients) == 1 + assert not made_clients[0].begin_started.is_set() + assert store.import_calls == [] + finally: + factory_gate.set() + manager.close() + + +def test_login_manager_imports_a_real_codebuddy_client_into_account_store(tmp_path, monkeypatch): + store = AccountStore(tmp_path / "accounts") + + def begin_login(_client): + return {"auth_url": "https://login.example.test/authorize", "state": "test-state"} + + def complete_login(self, challenge, *, timeout, cancel_event): + assert challenge == {"auth_url": "https://login.example.test/authorize", "state": "test-state"} + assert timeout == 300 + assert not cancel_event.is_set() + self._save_session( + { + **self.session, + "auth": {"accessToken": "test-token"}, + "account": {"uid": "integration-user", "nickname": "Integration user"}, + } + ) + + monkeypatch.setattr(CodeBuddyClient, "begin_login", begin_login) + monkeypatch.setattr(CodeBuddyClient, "complete_login", complete_login) + + manager = AccountManager(store) + try: + job = manager.start_login("global") + result = _wait_for(manager, job["id"], "succeeded") + assert result["account_id"] + accounts = store.list_accounts() + assert len(accounts) == 1 + assert accounts[0]["backend"] == "global" + assert accounts[0]["uid"] == "integration-user" + assert accounts[0]["active"] is True + assert manager.client.session["auth"]["accessToken"] == "test-token" + assert "test-token" not in repr(result) + finally: + manager.close() + + +def test_rate_limit_snapshot_defaults_and_runtime_state_persist_without_credentials(tmp_path): + store = FakeStore(tmp_path) + client = StoredClient("domestic", "rate-limited") + store.import_client(client, activate=True) + manager = AccountManager(store) + try: + account_id = manager.account_id_for_client(client) + assert account_id == "account-1" + assert manager.get_settings() == { + "auto_switch_on_429": False, + "default_cooldown_seconds": 60, + "auto_checkin_enabled": False, + } + manager.set_settings(True) + + before = time.time() + assert manager.handle_rate_limit( + client, "invalid retry value", model="model-a" + ) is None + snapshot = manager.snapshot() + limited = snapshot["accounts"][0] + assert snapshot["settings"]["auto_switch_on_429"] is True + assert snapshot["server_time"] >= before + assert "rate_limited" not in limited + assert "cooldown_until" not in limited + cooldown = _model_cooldown(limited, "model-a") + assert cooldown is not None + assert cooldown["rate_limited"] is True + assert cooldown["cooldown_source"] == "default" + assert cooldown["cooldown_until"] == pytest.approx(before + 60, abs=1) + assert 1 <= cooldown["cooldown_remaining_seconds"] <= 60 + assert cooldown["last_rate_limited_at"] == pytest.approx(before, abs=1) + + state_file = tmp_path / "account-manager-state-v1.json" + raw_state = state_file.read_text(encoding="utf-8") + parsed_state = json.loads(raw_state) + assert parsed_state["version"] == 2 + assert parsed_state["cooldowns"][account_id]["model-a"]["source"] == "default" + assert "token-rate-limited" not in raw_state + assert "accessToken" not in raw_state + finally: + manager.close() + + reopened = AccountManager(store) + try: + assert reopened.get_settings()["auto_switch_on_429"] is True + restored = reopened.snapshot()["accounts"][0] + restored_cooldown = _model_cooldown(restored, "model-a") + assert restored_cooldown is not None + assert restored_cooldown["rate_limited"] is True + assert restored_cooldown["cooldown_source"] == "default" + assert restored_cooldown["cooldown_until"] == pytest.approx( + cooldown["cooldown_until"] + ) + finally: + reopened.close() + + +def test_retry_after_zero_and_http_date_are_supported_and_expire_without_poll_writes( + tmp_path, monkeypatch +): + store = FakeStore(tmp_path) + zero_client = StoredClient("domestic", "zero") + date_client = StoredClient("domestic", "date") + store.import_client(zero_client, activate=True) + store.import_client(date_client) + manager = AccountManager(store) + try: + fixed_now = time.time() + zero_deadline = fixed_now + date_deadline = fixed_now + 75 + date_value = format_datetime( + datetime.fromtimestamp(date_deadline, timezone.utc), usegmt=True + ) + monkeypatch.setattr("codebuddy_proxy.account_manager.time.time", lambda: fixed_now) + assert manager.handle_rate_limit(zero_client, "0", model="model-zero") is None + assert manager.handle_rate_limit( + date_client, date_value, model="model-date" + ) is None + + rows = {item["uid"]: item for item in manager.snapshot()["accounts"]} + zero_cooldown = _model_cooldown(rows["zero"], "model-zero") + date_cooldown = _model_cooldown(rows["date"], "model-date") + assert zero_cooldown["rate_limited"] is False + assert zero_cooldown["cooldown_source"] == "retry_after" + assert zero_cooldown["cooldown_until"] == pytest.approx(fixed_now) + assert zero_cooldown["cooldown_remaining_seconds"] == 0 + assert zero_cooldown["last_rate_limited_at"] == pytest.approx(fixed_now) + assert date_cooldown["rate_limited"] is True + assert date_cooldown["cooldown_source"] == "retry_after" + assert date_cooldown["cooldown_until"] == pytest.approx(date_deadline, abs=1) + + persisted_before_expiry = (tmp_path / "account-manager-state-v1.json").read_text( + encoding="utf-8" + ) + monkeypatch.setattr( + "codebuddy_proxy.account_manager.time.time", lambda: fixed_now + 76 + ) + expired = {item["uid"]: item for item in manager.snapshot()["accounts"]}["date"] + expired_cooldown = _model_cooldown(expired, "model-date") + assert expired_cooldown is not None + assert expired_cooldown["rate_limited"] is False + assert expired_cooldown["cooldown_until"] == pytest.approx(date_deadline, abs=1) + assert expired_cooldown["cooldown_remaining_seconds"] == 0 + assert expired_cooldown["cooldown_source"] == "retry_after" + assert expired_cooldown["last_rate_limited_at"] == pytest.approx(fixed_now) + assert ( + tmp_path / "account-manager-state-v1.json" + ).read_text(encoding="utf-8") == persisted_before_expiry + finally: + manager.close() + + +def test_auto_switch_is_opt_in_region_scoped_and_respects_exclusions(tmp_path): + store = FakeStore(tmp_path) + original = StoredClient("domestic", "original") + excluded = StoredClient("domestic", "excluded") + candidate = StoredClient("domestic", "candidate") + other_region = StoredClient("global", "global") + original_id = store.import_client(original, activate=True) + excluded_id = store.import_client(excluded) + candidate_id = store.import_client(candidate) + store.import_client(other_region) + manager = AccountManager(store) + try: + assert manager.handle_rate_limit(original, "30", model="model-a") is None + assert store.active_account_id == original_id + + manager.set_settings(True) + switched = manager.handle_rate_limit( + original, + "30", + model="model-a", + excluded_account_ids={excluded_id}, + ) + assert switched is candidate + assert store.active_account_id == candidate_id + + # A delayed 429 from the old request records the longer deadline but + # cannot rotate away from the account that is now active. + latest = manager.handle_rate_limit(original, "5", model="model-a") + assert latest is None + assert store.active_account_id == candidate_id + row = next(item for item in manager.snapshot()["accounts"] if item["id"] == original_id) + cooldown = _model_cooldown(row, "model-a") + assert cooldown["cooldown_remaining_seconds"] > 0 + assert cooldown["cooldown_source"] == "retry_after" + finally: + manager.close() + + +def test_cooldowns_are_isolated_by_exact_model_and_candidate_model(tmp_path): + store = FakeStore(tmp_path) + original = StoredClient("domestic", "original") + blocked_for_model_a = StoredClient("domestic", "blocked-a") + cooled_for_model_b = StoredClient("domestic", "cooled-b") + original_id = store.import_client(original, activate=True) + blocked_id = store.import_client(blocked_for_model_a) + model_b_id = store.import_client(cooled_for_model_b) + manager = AccountManager(store) + try: + # Inactive 429s only mark their request model. The candidate cooled for + # model-b must remain eligible for a model-a failover. + manager.handle_rate_limit(blocked_for_model_a, "90", model="model-a") + manager.handle_rate_limit(cooled_for_model_b, "90", model="model-b") + manager.set_settings(True) + + switched = manager.handle_rate_limit(original, "45", model="model-a") + assert switched is cooled_for_model_b + assert store.active_account_id == model_b_id + + rows = {row["id"]: row for row in manager.snapshot()["accounts"]} + assert "rate_limited" not in rows[original_id] + assert _model_cooldown(rows[original_id], "model-a")["rate_limited"] is True + assert _model_cooldown(rows[blocked_id], "model-a")["rate_limited"] is True + assert _model_cooldown(rows[model_b_id], "model-b")["rate_limited"] is True + assert _model_cooldown(rows[model_b_id], "model-a") is None + finally: + manager.close() + + +def test_model_keys_preserve_case_and_model_less_or_invalid_events_are_noops( + tmp_path, +): + store = FakeStore(tmp_path) + client = StoredClient("domestic", "model-keys") + account_id = store.import_client(client, activate=True) + store.import_client(StoredClient("domestic", "model-keys-backup")) + manager = AccountManager(store) + try: + manager.handle_rate_limit(client, "30", model="Model-X") + manager.handle_rate_limit(client, "45", model="model-x") + manager.set_settings(True) + snapshot = manager.snapshot() + row = snapshot["accounts"][0] + assert [item["model"] for item in row["model_cooldowns"]] == [ + "Model-X", + "model-x", + ] + state_file = tmp_path / "account-manager-state-v1.json" + before = state_file.read_text(encoding="utf-8") + + for invalid_model in (None, "", " ", "bad\nmodel", "x" * 257, 7): + assert manager.handle_rate_limit( + client, "120", model=invalid_model + ) is None + + after = state_file.read_text(encoding="utf-8") + assert after == before + assert store.active_account_id == account_id + assert [ + item["model"] for item in manager.snapshot()["accounts"][0]["model_cooldowns"] + ] == ["Model-X", "model-x"] + finally: + manager.close() + + +def test_failed_model_cooldown_write_rolls_back_nested_memory_and_disk( + tmp_path, monkeypatch +): + import codebuddy_proxy.account_manager as account_manager_module + + store = FakeStore(tmp_path) + client = StoredClient("domestic", "rollback") + store.import_client(client, activate=True) + manager = AccountManager(store) + try: + manager.handle_rate_limit(client, "60", model="model-a") + state_file = tmp_path / "account-manager-state-v1.json" + original_text = state_file.read_text(encoding="utf-8") + original_deadline = _model_cooldown( + manager.snapshot()["accounts"][0], "model-a" + )["cooldown_until"] + + def failed_write(_path, _payload): + raise OSError("private runtime-state path") + + monkeypatch.setattr(account_manager_module, "_atomic_json", failed_write) + with pytest.raises(AccountManagerError) as failure: + manager.handle_rate_limit(client, "120", model="model-a") + assert failure.value.message == "无法保存账号管理状态" + + row = manager.snapshot()["accounts"][0] + assert _model_cooldown(row, "model-a")["cooldown_until"] == original_deadline + assert _model_cooldown(row, "model-b") is None + assert state_file.read_text(encoding="utf-8") == original_text + finally: + manager.close() + + +def test_concurrent_429s_from_same_active_account_only_switch_once(tmp_path): + from concurrent.futures import ThreadPoolExecutor + + store = FakeStore(tmp_path) + original = StoredClient("domestic", "original") + candidate = StoredClient("domestic", "candidate") + original_id = store.import_client(original, activate=True) + candidate_id = store.import_client(candidate) + manager = AccountManager(store) + manager.set_settings(True) + barrier = threading.Barrier(3) + + def rate_limit(): + barrier.wait(timeout=2) + return manager.handle_rate_limit(original, "45", model="model-a") + + try: + with ThreadPoolExecutor(max_workers=2) as executor: + results = [executor.submit(rate_limit) for _ in range(2)] + barrier.wait(timeout=2) + clients = [future.result(timeout=3) for future in results] + assert clients.count(candidate) == 1 + assert clients.count(None) == 1 + assert store.active_account_id == candidate_id + assert store.set_active_calls.count(candidate_id) == 1 + row = next(item for item in manager.snapshot()["accounts"] if item["id"] == original_id) + assert _model_cooldown(row, "model-a")["rate_limited"] is True + finally: + manager.close() + + +def test_cooldown_max_deadline_manual_switch_and_refreshable_candidate(tmp_path, monkeypatch): + store = FakeStore(tmp_path) + original = StoredClient("domestic", "original") + expired_candidate = StoredClient( + "domestic", + "expired", + auth={"accessToken": "expired", "expiresAt": 1}, + ) + refreshable = StoredClient( + "domestic", + "refreshable", + auth={"refreshToken": "refresh-only"}, + ) + original_id = store.import_client(original, activate=True) + store.import_client(expired_candidate) + refreshable_id = store.import_client(refreshable) + manager = AccountManager(store) + try: + manager.set_settings(True) + fixed_now = 1_800_000_000.0 + monkeypatch.setattr("codebuddy_proxy.account_manager.time.time", lambda: fixed_now) + assert manager.handle_rate_limit(original, "120", model="model-a") is refreshable + assert store.active_account_id == refreshable_id + + # The later shorter default cooldown must not reduce the original 120s. + assert manager.handle_rate_limit(original, None, model="model-a") is None + original_row = next( + item for item in manager.snapshot()["accounts"] if item["id"] == original_id + ) + original_cooldown = _model_cooldown(original_row, "model-a") + assert original_cooldown["cooldown_until"] == fixed_now + 120 + assert original_cooldown["cooldown_source"] == "retry_after" + + # Even with automation enabled, an explicit selection is preserved if + # a previously active client's 429 arrives afterward. + manager.activate(original_id) + assert manager.handle_rate_limit(refreshable, "10", model="model-a") is None + assert store.active_account_id == original_id + finally: + manager.close() + + +def test_cooldown_is_cleared_when_account_is_deleted_and_bad_runtime_file_is_preserved(tmp_path): + store = FakeStore(tmp_path) + active = StoredClient("domestic", "active") + inactive = StoredClient("domestic", "inactive") + store.import_client(active, activate=True) + inactive_id = store.import_client(inactive) + manager = AccountManager(store) + try: + manager.handle_rate_limit(inactive, "45", model="model-a") + state_file = tmp_path / "account-manager-state-v1.json" + manager.delete(inactive_id) + state = json.loads(state_file.read_text(encoding="utf-8")) + assert inactive_id not in state["cooldowns"] + assert all(row["id"] != inactive_id for row in manager.snapshot()["accounts"]) + finally: + manager.close() + + state_file = tmp_path / "account-manager-state-v1.json" + corrupted_values = ( + "{broken", + '{"version":2,"settings":{"auto_switch_on_429":false,"default_cooldown_seconds":60},"cooldowns":{}}', + '{"version":true,"settings":{"auto_switch_on_429":false,"default_cooldown_seconds":60,"auto_checkin_enabled":false},"cooldowns":{}}', + '{"version":1,"settings":{"auto_switch_on_429":"true","default_cooldown_seconds":60},"cooldowns":{}}', + '{"version":1,"settings":{"auto_switch_on_429":false,"default_cooldown_seconds":60,"auto_checkin_enabled":1},"cooldowns":{}}', + '{"version":1,"settings":{"auto_switch_on_429":false,"default_cooldown_seconds":60,"unexpected":false},"cooldowns":{}}', + '{"version":1,"version":1,"settings":{"auto_switch_on_429":false,"default_cooldown_seconds":60},"cooldowns":{}}', + '{"version":2,"settings":{"auto_switch_on_429":false,"default_cooldown_seconds":60,"auto_checkin_enabled":false},"cooldowns":{"account-1":{"\n":{"deadline":1,"last_rate_limited_at":1,"source":"default"}}}}', + '{"version":2,"settings":{"auto_switch_on_429":false,"default_cooldown_seconds":60,"auto_checkin_enabled":false},"cooldowns":{"account-1":{"model":{"deadline":1,"last_rate_limited_at":1,"source":"unknown"}}}}', + '{"version":2,"settings":{"auto_switch_on_429":false,"default_cooldown_seconds":60,"auto_checkin_enabled":false},"cooldowns":{"account-1":{"model":{"deadline":1,"last_rate_limited_at":1}}}}', + ) + corrupted_values += ( + json.dumps( + { + "version": 2, + "settings": { + "auto_switch_on_429": False, + "default_cooldown_seconds": 60, + "auto_checkin_enabled": False, + }, + "cooldowns": { + "account-1": { + "m" * 257: { + "deadline": 1, + "last_rate_limited_at": 1, + "source": "default", + } + } + }, + } + ), + ) + for corrupted in corrupted_values: + state_file.write_text(corrupted, encoding="utf-8") + with pytest.raises(AccountManagerError) as failure: + AccountManager(store) + assert failure.value.message in { + "账号管理状态文件损坏,未修改任何数据", + "账号管理状态文件版本不受支持,未修改任何数据", + } + assert state_file.read_text(encoding="utf-8") == corrupted + + +def test_failed_cooldown_persistence_aborts_account_deletion(tmp_path, monkeypatch): + import codebuddy_proxy.account_manager as account_manager_module + + store = FakeStore(tmp_path) + active = StoredClient("domestic", "active") + inactive = StoredClient("domestic", "inactive") + store.import_client(active, activate=True) + inactive_id = store.import_client(inactive) + manager = AccountManager(store) + try: + manager.handle_rate_limit(inactive, "45", model="model-a") + state_file = tmp_path / "account-manager-state-v1.json" + original_state = state_file.read_text(encoding="utf-8") + + def failed_write(_path, _payload): + raise OSError("private runtime-state path") + + monkeypatch.setattr(account_manager_module, "_atomic_json", failed_write) + with pytest.raises(AccountManagerError) as failure: + manager.delete(inactive_id) + + assert failure.value.message == "无法保存账号管理状态" + assert manager.has_account(inactive_id) + row = next( + account for account in manager.snapshot()["accounts"] + if account["id"] == inactive_id + ) + assert _model_cooldown(row, "model-a")["rate_limited"] is True + assert state_file.read_text(encoding="utf-8") == original_state + finally: + manager.close() + + +def test_failed_store_deletion_restores_persisted_cooldown(tmp_path, monkeypatch): + store = FakeStore(tmp_path) + active = StoredClient("domestic", "active") + inactive = StoredClient("domestic", "inactive") + store.import_client(active, activate=True) + inactive_id = store.import_client(inactive) + manager = AccountManager(store) + try: + manager.handle_rate_limit(inactive, "45", model="model-a") + state_file = tmp_path / "account-manager-state-v1.json" + original_state = json.loads(state_file.read_text(encoding="utf-8")) + + def failed_delete(_account_id): + raise OSError("simulated account-store failure") + + monkeypatch.setattr(store, "delete_account", failed_delete) + with pytest.raises(OSError, match="simulated account-store failure"): + manager.delete(inactive_id) + + assert manager.has_account(inactive_id) + row = next( + account for account in manager.snapshot()["accounts"] + if account["id"] == inactive_id + ) + assert _model_cooldown(row, "model-a")["rate_limited"] is True + restored_state = json.loads(state_file.read_text(encoding="utf-8")) + assert restored_state["cooldowns"][inactive_id] == original_state["cooldowns"][inactive_id] + finally: + manager.close() + + +def test_failed_settings_persistence_rolls_back_memory_and_uses_safe_error(tmp_path, monkeypatch): + import codebuddy_proxy.account_manager as account_manager_module + + manager = AccountManager(FakeStore(tmp_path)) + try: + manager.set_settings(True) + state_file = tmp_path / "account-manager-state-v1.json" + original = state_file.read_text(encoding="utf-8") + + def failed_write(_path, _payload): + raise OSError("private path and token value") + + monkeypatch.setattr(account_manager_module, "_atomic_json", failed_write) + with pytest.raises(AccountManagerError) as failure: + manager.set_settings(False) + assert failure.value.message == "无法保存账号管理状态" + assert "private path" not in str(failure.value) + assert manager.get_settings()["auto_switch_on_429"] is True + assert state_file.read_text(encoding="utf-8") == original + finally: + manager.close() + + +def test_fake_store_without_root_keeps_runtime_settings_in_memory(tmp_path): + store = FakeStore(tmp_path) + del store.root + manager = AccountManager(store) + try: + assert manager.set_settings(True)["auto_switch_on_429"] is True + assert manager.get_settings()["auto_switch_on_429"] is True + finally: + manager.close() + + +@pytest.mark.parametrize( + "legacy_settings", + ( + {"auto_switch_on_429": True, "default_cooldown_seconds": 60}, + { + "auto_switch_on_429": True, + "default_cooldown_seconds": 60, + "auto_checkin_enabled": False, + }, + ), +) +def test_v1_runtime_settings_migrate_on_write_and_discard_ambiguous_cooldowns( + tmp_path, legacy_settings +): + store = FakeStore(tmp_path) + client = StoredClient("domestic", "legacy") + store.import_client(client, activate=True) + deadline = time.time() + 120 + state_path = tmp_path / "account-manager-state-v1.json" + state_path.write_text( + json.dumps( + { + "version": 1, + "settings": legacy_settings, + "cooldowns": { + "account-1": { + "deadline": deadline, + "last_rate_limited_at": deadline - 10, + "source": "retry_after", + } + }, + } + ), + encoding="utf-8", + ) + + manager = AccountManager(store) + try: + assert manager.get_settings() == { + "auto_switch_on_429": True, + "default_cooldown_seconds": 60, + "auto_checkin_enabled": False, + } + original_text = state_path.read_text(encoding="utf-8") + assert manager.snapshot()["accounts"][0]["model_cooldowns"] == [] + assert state_path.read_text(encoding="utf-8") == original_text + + manager.set_settings(auto_checkin_enabled=True) + persisted = json.loads(state_path.read_text(encoding="utf-8")) + assert persisted["version"] == 2 + assert persisted["settings"] == { + "auto_switch_on_429": True, + "default_cooldown_seconds": 60, + "auto_checkin_enabled": True, + } + assert persisted["cooldowns"] == {} + finally: + manager.close() + + +def test_checkin_setting_callback_only_follows_successful_transitions(tmp_path, monkeypatch): + import codebuddy_proxy.account_manager as account_manager_module + + manager = AccountManager(FakeStore(tmp_path)) + callbacks = [] + manager.checkins.settings_changed = callbacks.append + try: + manager.set_settings(auto_switch_on_429=True) + manager.set_settings(auto_checkin_enabled=False) + assert callbacks == [] + + manager.set_settings(auto_checkin_enabled=True) + assert callbacks == [True] + manager.set_settings(auto_checkin_enabled=True) + assert callbacks == [True] + + state_path = tmp_path / "account-manager-state-v1.json" + original = state_path.read_text(encoding="utf-8") + + def failed_write(_path, _payload): + raise OSError("private path and token value") + + monkeypatch.setattr(account_manager_module, "_atomic_json", failed_write) + with pytest.raises(AccountManagerError) as failure: + manager.set_settings(auto_checkin_enabled=False) + assert failure.value.message == "无法保存账号管理状态" + assert manager.get_settings()["auto_checkin_enabled"] is True + assert callbacks == [True] + assert state_path.read_text(encoding="utf-8") == original + finally: + manager.close() + + +def test_model_less_rate_limit_does_not_record_or_switch(tmp_path): + store = FakeStore(tmp_path) + checked_client = StoredClient("domestic", "checked-client") + backup_client = StoredClient("domestic", "backup-client") + original_id = store.import_client(checked_client, activate=True) + store.import_client(backup_client) + manager = AccountManager(store) + try: + manager.set_settings(True) + assert manager.handle_rate_limit(checked_client, "45") is None + assert store.active_account_id == original_id + row = next( + account + for account in manager.snapshot()["accounts"] + if account["id"] == original_id + ) + assert row["model_cooldowns"] == [] + finally: + manager.close() diff --git a/test_account_store.py b/test_account_store.py new file mode 100644 index 0000000..0586913 --- /dev/null +++ b/test_account_store.py @@ -0,0 +1,320 @@ +import json +import threading + +import pytest + +import codebuddy_proxy.account_store as account_store_module +from codebuddy_proxy.account_store import AccountStore +from codebuddy_proxy.backend_profile import DOMESTIC_PROFILE, GLOBAL_PROFILE, resolve_backend +from codebuddy_proxy.codebuddy_client_demo import CodeBuddyClient, CodeBuddyError + + +def _client(path, *, profile=DOMESTIC_PROFILE, endpoint=None, uid="user-1", access="access-1"): + client = CodeBuddyClient( + backend=resolve_backend( + profile=profile, + endpoint=endpoint, + session_file=path, + environment={}, + ) + ) + client._save_session( + { + "auth": {"accessToken": access, "refreshToken": f"refresh-{access}"}, + "account": {"uid": uid, "nickname": f"Name {uid}"}, + "machineId": "machine-1", + } + ) + return client + + +def test_import_isolated_and_reassociation_updates_cached_client(tmp_path): + store = AccountStore(tmp_path / "accounts") + source_file = tmp_path / "source.json" + source = _client(source_file) + source_before = source_file.read_bytes() + + account_id = store.import_client(source) + assert store.active_account_id == account_id + cached = store.get_client(account_id) + assert cached.session_file != source_file + assert source_file.read_bytes() == source_before + + replacement = _client(tmp_path / "replacement.json", access="access-2") + replacement_before = replacement.session_file.read_bytes() + assert store.import_client(replacement) == account_id + assert replacement.session_file.read_bytes() == replacement_before + assert store.get_client(account_id) is cached + assert cached.session["auth"]["accessToken"] == "access-2" + + index_text = (store.root / "index.json").read_text(encoding="utf-8") + assert "access-1" not in index_text + assert "access-2" not in index_text + assert "refresh-access-2" not in index_text + assert set(json.loads(index_text)["accounts"][0]) == { + "id", "backend", "endpoint", "uid", "nickname" + } + assert store.list_accounts() == [ + { + "id": account_id, + "backend": "domestic", + "endpoint": DOMESTIC_PROFILE.default_endpoint, + "uid": "user-1", + "nickname": "Name user-1", + "active": True, + } + ] + + +def test_identity_includes_backend_and_endpoint_and_clear_is_persistent(tmp_path): + root = tmp_path / "accounts" + store = AccountStore(root) + domestic = _client(tmp_path / "domestic.json") + global_client = _client( + tmp_path / "global.json", profile=GLOBAL_PROFILE, uid="user-1", access="global-access" + ) + custom = _client( + tmp_path / "custom.json", + endpoint="https://staging.example.test", + uid="user-1", + access="custom-access", + ) + + domestic_id = store.import_client(domestic) + global_id = store.import_client(global_client) + custom_id = store.import_client(custom) + assert len({domestic_id, global_id, custom_id}) == 3 + assert store.active_account_id == domestic_id + + active_global = store.set_active(global_id) + assert active_global.session["auth"]["accessToken"] == "global-access" + with pytest.raises(CodeBuddyError, match="不能删除当前生效账号"): + store.delete_account(global_id) + + store.clear_active() + reopened = AccountStore(root) + assert reopened.active_account_id is None + assert reopened.active_client() is None + # Re-importing an existing account after an explicit clear keeps the store empty. + assert reopened.import_client(global_client) == global_id + assert reopened.active_account_id is None + reopened.delete_account(global_id) + assert all(item["id"] != global_id for item in reopened.list_accounts()) + + +def test_corrupt_index_is_rejected_without_overwrite(tmp_path): + root = tmp_path / "accounts" + root.mkdir() + index = root / "index.json" + index.write_text("{broken", encoding="utf-8") + + with pytest.raises(CodeBuddyError, match="账号索引文件损坏"): + AccountStore(root) + assert index.read_text(encoding="utf-8") == "{broken" + + +@pytest.mark.parametrize( + "session", + [ + {"auth": {"accessToken": "token"}, "account": {"nickname": "No uid"}}, + {"auth": {}, "account": {"uid": "user-1"}}, + ], +) +def test_import_rejects_legacy_session_without_uid_or_token(tmp_path, session): + client = CodeBuddyClient(session_file=tmp_path / "legacy.json") + client._save_session(session) + with pytest.raises(CodeBuddyError, match="缺少"): + AccountStore(tmp_path / "accounts").import_client(client) + assert not (tmp_path / "accounts" / "index.json").exists() + + +def test_inflight_old_refresh_cannot_overwrite_reassociated_token(tmp_path, monkeypatch): + store = AccountStore(tmp_path / "accounts") + account_id = store.import_client(_client(tmp_path / "first.json", access="first-access")) + cached = store.get_client(account_id) + request_started = threading.Event() + release_request = threading.Event() + refresh_result = [] + + def delayed_refresh(method, path, *, headers=None, body=None, timeout=30): + request_started.set() + assert release_request.wait(5) + return {"accessToken": "stale-refreshed-access", "refreshToken": "stale-refresh"} + + monkeypatch.setattr(cached, "_request", delayed_refresh) + thread = threading.Thread(target=lambda: refresh_result.append(cached.refresh())) + thread.start() + assert request_started.wait(5) + + assert store.import_client( + _client(tmp_path / "new-login.json", access="new-login-access") + ) == account_id + release_request.set() + thread.join(5) + + assert not thread.is_alive() + assert refresh_result == [True] + assert cached.session["auth"]["accessToken"] == "new-login-access" + saved = json.loads(cached.session_file.read_text(encoding="utf-8")) + assert saved["auth"]["accessToken"] == "new-login-access" + + +def test_failed_first_import_rolls_back_orphan_session_and_index(tmp_path, monkeypatch): + store = AccountStore(tmp_path / "accounts") + source = _client(tmp_path / "first-import.json", access="new-token") + _, metadata = store._validated_client_session(source) + session_file = store._session_file(metadata["id"]) + write_index = store._write_index_locked + + def fail_index_write(index): + write_index(index) + raise OSError("injected index write failure") + + monkeypatch.setattr(store, "_write_index_locked", fail_index_write) + with pytest.raises(OSError, match="injected index write failure"): + store.import_client(source) + + assert not session_file.exists() + assert not (store.root / "index.json").exists() + assert store.list_accounts() == [] + + +def test_failed_reassociation_restores_session_file_and_cached_state(tmp_path, monkeypatch): + store = AccountStore(tmp_path / "accounts") + account_id = store.import_client( + _client(tmp_path / "old-login.json", access="old-token") + ) + cached = store.get_client(account_id) + previous_session, _ = cached._snapshot_session() + session_file = cached.session_file + previous_file = session_file.read_bytes() + index_file = store.index_file + previous_index = index_file.read_bytes() + replacement = _client(tmp_path / "new-login.json", access="new-token") + write_index = store._write_index_locked + + def fail_index_write(index): + write_index(index) + raise OSError("injected index write failure") + + monkeypatch.setattr(store, "_write_index_locked", fail_index_write) + with pytest.raises(OSError, match="injected index write failure"): + store.import_client(replacement) + + assert session_file.read_bytes() == previous_file + assert index_file.read_bytes() == previous_index + assert cached.session == previous_session + assert store.get_client(account_id) is cached + assert store.list_accounts()[0]["nickname"] == "Name user-1" + assert "old-token" in previous_file.decode("utf-8") + assert "new-token" not in session_file.read_text(encoding="utf-8") + + +def test_import_rollback_holds_session_lock_against_concurrent_writer( + tmp_path, monkeypatch +): + store = AccountStore(tmp_path / "accounts") + account_id = store.import_client(_client(tmp_path / "first.json", access="old")) + cached = store.get_client(account_id) + session_file = cached.session_file + original_restore = account_store_module._restore_file + writer_session, _ = cached._snapshot_session() + writer_session["auth"]["accessToken"] = "writer-access" + writer_session["auth"]["refreshToken"] = "writer-refresh" + rollback_paused = threading.Event() + finish_rollback = threading.Event() + writer_call_entered = threading.Event() + writer_finished = threading.Event() + import_errors = [] + writer_results = [] + + def pause_after_session_restore(path, contents): + original_restore(path, contents) + if path == session_file: + rollback_paused.set() + assert finish_rollback.wait(5) + + monkeypatch.setattr(account_store_module, "_restore_file", pause_after_session_restore) + original_save = cached._save_session + + def observe_writer_call(session, *, expected_generation=None): + if threading.current_thread().name == "session-writer": + writer_call_entered.set() + return original_save(session, expected_generation=expected_generation) + + monkeypatch.setattr(cached, "_save_session", observe_writer_call) + original_index_write = store._write_index_locked + + def fail_after_index_write(index): + original_index_write(index) + raise OSError("injected index write failure") + + monkeypatch.setattr(store, "_write_index_locked", fail_after_index_write) + + def import_replacement(): + try: + store.import_client(_client(tmp_path / "replacement.json", access="new")) + except Exception as exc: # captured for assertions in the test thread + import_errors.append(exc) + + def write_concurrently(): + try: + writer_results.append(cached._save_session(writer_session)) + finally: + writer_finished.set() + + importer = threading.Thread(target=import_replacement) + importer.start() + assert rollback_paused.wait(5) + + writer = threading.Thread(target=write_concurrently, name="session-writer") + writer.start() + assert writer_call_entered.wait(5) + assert not writer_finished.wait(0.2) + + finish_rollback.set() + importer.join(5) + writer.join(5) + + assert not importer.is_alive() + assert not writer.is_alive() + assert len(import_errors) == 1 + assert isinstance(import_errors[0], OSError) + assert writer_results == [True] + persisted = json.loads(session_file.read_text(encoding="utf-8")) + assert persisted == cached.session + assert persisted["auth"]["accessToken"] == "writer-access" + + +def test_deleting_inactive_account_prevents_inflight_refresh_from_recreating_file( + tmp_path, monkeypatch +): + store = AccountStore(tmp_path / "accounts") + removed_id = store.import_client(_client(tmp_path / "removed.json", access="removed")) + keeper_id = store.import_client( + _client(tmp_path / "keeper.json", uid="user-2", access="keeper") + ) + store.set_active(keeper_id) + cached = store.get_client(removed_id) + session_file = cached.session_file + request_started = threading.Event() + release_request = threading.Event() + results = [] + + def delayed_refresh(method, path, *, headers=None, body=None, timeout=30): + request_started.set() + assert release_request.wait(5) + return {"accessToken": "late-refresh", "refreshToken": "late-refresh-token"} + + monkeypatch.setattr(cached, "_request", delayed_refresh) + thread = threading.Thread(target=lambda: results.append(cached.refresh())) + thread.start() + assert request_started.wait(5) + store.delete_account(removed_id) + release_request.set() + thread.join(5) + + assert not thread.is_alive() + assert results == [True] + assert not session_file.exists() + assert all(item["id"] != removed_id for item in store.list_accounts()) diff --git a/test_admin_api.py b/test_admin_api.py new file mode 100644 index 0000000..7566400 --- /dev/null +++ b/test_admin_api.py @@ -0,0 +1,451 @@ +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from codebuddy_proxy import admin_api +from codebuddy_proxy.admin_api import create_admin_router + + +class FakeManager: + def __init__(self): + self.login_calls = [] + self.activated = [] + self.deleted = [] + self.settings = { + "auto_switch_on_429": False, + "default_cooldown_seconds": 60, + "auto_checkin_enabled": False, + } + self.checkins = FakeCheckins() + + def has_account(self, account_id): + return account_id == "account-1" + + def snapshot(self): + return { + "accounts": [ + { + "id": "account-1", + "backend": "domestic", + "endpoint": "https://copilot.tencent.com", + "uid": "user-1", + "nickname": "Local user", + "active": True, + "model_cooldowns": [ + { + "model": "model-a", + "rate_limited": True, + "cooldown_until": 1234.5, + "cooldown_remaining_seconds": 7, + "last_rate_limited_at": 1227.5, + "cooldown_source": "retry_after", + "session": {"accessToken": "private-token-value"}, + }, + { + "model": "bad\nmodel", + "rate_limited": True, + "cooldown_source": "default", + }, + { + "model": "x" * 257, + "rate_limited": True, + "cooldown_source": "default", + }, + ], + # Previous account-wide cooldown fields must not leak into + # the public response even if a stale caller provides them. + "rate_limited": True, + "cooldown_until": 9999.0, + "checkin": { + "date": "2026-10-24", + "status": "checked", + "checked": True, + "last_checked_at": 1220, + "last_checkin_at": 1210, + "next_checkin_at": 1729766400, + "error": None, + "session": {"accessToken": "private-token-value"}, + }, + "session": {"auth": {"accessToken": "private-token-value"}}, + } + ], + "active_account_id": "account-1", + "settings": self.settings, + "server_time": 1227.5, + "raw_session": "private-token-value", + } + + def get_settings(self): + return dict(self.settings) + + def set_settings(self, auto_switch_on_429=None, *, auto_checkin_enabled=None): + if auto_switch_on_429 is not None: + self.settings["auto_switch_on_429"] = auto_switch_on_429 + if auto_checkin_enabled is not None: + self.settings["auto_checkin_enabled"] = auto_checkin_enabled + return dict(self.settings) + + def activate(self, account_id): + self.activated.append(account_id) + + def delete(self, account_id): + self.deleted.append(account_id) + + def start_login(self, backend): + self.login_calls.append(backend) + return { + "id": "job-1", + "backend": backend, + "status": "pending", + "auth_url": "https://login.example.test/auth", + "error": None, + "account_id": None, + "state": "private-state", + "session": {"auth": {"accessToken": "private-token-value"}}, + } + + def get_login(self, job_id): + return { + "id": job_id, + "backend": "global", + "status": "failed", + "auth_url": None, + "error": "登录失败,请重试", + "account_id": None, + "exception": "private-token-value", + } + + def cancel_login(self, job_id): + return {"id": job_id, "backend": "global", "status": "cancelled"} + + +class FakeCheckins: + def __init__(self): + self.calls = [] + self.status_result = { + "date": "2026-10-24", + "status": "unchecked", + "checked": False, + "last_checked_at": 1220.0, + "last_checkin_at": None, + "next_checkin_at": 1729766400.0, + "error": None, + "session": {"accessToken": "private-token-value"}, + } + self.checkin_result = { + "date": "2026-10-24", + "status": "checked", + "checked": True, + "last_checked_at": 1230.0, + "last_checkin_at": 1230.0, + "next_checkin_at": 1729766400.0, + "error": None, + "raw_response": "private-token-value", + } + + async def check_status(self, account_id, *, reason): + self.calls.append(("status", account_id, reason)) + return self.status_result + + async def check_in(self, account_id, *, reason): + self.calls.append(("checkin", account_id, reason)) + return self.checkin_result + + +def _app(manager): + app = FastAPI() + app.include_router(create_admin_router(lambda: manager)) + return app + + +def _client(manager): + return TestClient( + _app(manager), + base_url="http://localhost:8123", + client=("127.0.0.1", 50000), + ) + + +def test_account_and_job_responses_are_sanitized_and_uncached(): + manager = FakeManager() + with _client(manager) as client: + accounts = client.get("/admin/api/accounts") + assert accounts.status_code == 200 + assert accounts.headers["cache-control"] == "no-store" + assert accounts.json() == { + "accounts": [ + { + "id": "account-1", + "backend": "domestic", + "endpoint": "https://copilot.tencent.com", + "uid": "user-1", + "nickname": "Local user", + "active": True, + "model_cooldowns": [ + { + "model": "model-a", + "rate_limited": True, + "cooldown_until": 1234.5, + "cooldown_remaining_seconds": 7, + "last_rate_limited_at": 1227.5, + "cooldown_source": "retry_after", + } + ], + "checkin": { + "date": "2026-10-24", + "status": "checked", + "checked": True, + "last_checked_at": 1220.0, + "last_checkin_at": 1210.0, + "error": None, + "next_checkin_at": 1729766400.0, + }, + } + ], + "active_account_id": "account-1", + "settings": { + "auto_switch_on_429": False, + "auto_checkin_enabled": False, + "default_cooldown_seconds": 60, + }, + "server_time": 1227.5, + } + + job = client.get("/admin/api/logins/job-1") + assert job.status_code == 200 + assert job.headers["cache-control"] == "no-store" + assert "private-token-value" not in job.text + assert "exception" not in job.json() + + +def test_write_routes_require_header_and_same_origin_with_port_check(): + manager = FakeManager() + with _client(manager) as client: + denied = client.post("/admin/api/logins", json={"backend": "global"}) + assert denied.status_code == 403 + assert manager.login_calls == [] + + cross_port = client.post( + "/admin/api/logins", + json={"backend": "global"}, + headers={ + "X-Account-Manager": "1", + "Origin": "http://localhost:8124", + }, + ) + assert cross_port.status_code == 403 + assert manager.login_calls == [] + + accepted = client.post( + "/admin/api/logins", + json={"backend": "global"}, + headers={ + "X-Account-Manager": "1", + "Origin": "http://localhost:8123", + }, + ) + assert accepted.status_code == 202 + assert accepted.headers["cache-control"] == "no-store" + assert accepted.json() == { + "id": "job-1", + "backend": "global", + "status": "pending", + "auth_url": "https://login.example.test/auth", + "error": None, + "account_id": None, + } + assert manager.login_calls == ["global"] + + originless = client.post( + "/admin/api/accounts/account-1/activate", + headers={"X-Account-Manager": "1"}, + ) + assert originless.status_code == 200 + assert manager.activated == ["account-1"] + + +def test_settings_read_write_is_protected_strict_and_no_store(): + manager = FakeManager() + with _client(manager) as client: + read = client.get("/admin/api/settings") + assert read.status_code == 200 + assert read.headers["cache-control"] == "no-store" + assert read.json() == { + "auto_switch_on_429": False, + "auto_checkin_enabled": False, + "default_cooldown_seconds": 60, + } + + denied = client.patch( + "/admin/api/settings", + json={"auto_switch_on_429": True}, + ) + assert denied.status_code == 403 + assert manager.settings["auto_switch_on_429"] is False + + cross_port = client.patch( + "/admin/api/settings", + json={"auto_switch_on_429": True}, + headers={ + "X-Account-Manager": "1", + "Origin": "http://localhost:8124", + }, + ) + assert cross_port.status_code == 403 + assert manager.settings["auto_switch_on_429"] is False + + invalid_payloads = ( + {"auto_switch_on_429": 1}, + {"auto_switch_on_429": True, "unknown": "value"}, + {}, + {"auto_checkin_enabled": "true"}, + ) + for payload in invalid_payloads: + response = client.patch( + "/admin/api/settings", + json=payload, + headers={"X-Account-Manager": "1"}, + ) + assert response.status_code == 400 + duplicate_key = client.patch( + "/admin/api/settings", + content='{"auto_switch_on_429":false,"auto_switch_on_429":true}', + headers={"X-Account-Manager": "1", "Content-Type": "application/json"}, + ) + assert duplicate_key.status_code == 400 + assert manager.settings["auto_switch_on_429"] is False + + updated = client.patch( + "/admin/api/settings", + json={"auto_switch_on_429": True}, + headers={"X-Account-Manager": "1"}, + ) + assert updated.status_code == 200 + assert updated.headers["cache-control"] == "no-store" + assert updated.json() == { + "auto_switch_on_429": True, + "auto_checkin_enabled": False, + "default_cooldown_seconds": 60, + } + assert client.get("/admin/api/settings").json()["auto_switch_on_429"] is True + + checkin_enabled = client.patch( + "/admin/api/settings", + json={"auto_checkin_enabled": True}, + headers={"X-Account-Manager": "1"}, + ) + assert checkin_enabled.status_code == 200 + assert checkin_enabled.json() == { + "auto_switch_on_429": True, + "auto_checkin_enabled": True, + "default_cooldown_seconds": 60, + } + + +def test_checkin_routes_keep_shape_sanitize_errors_and_never_change_active(): + manager = FakeManager() + with _client(manager) as client: + status = client.get("/admin/api/accounts/account-1/checkin-status") + assert status.status_code == 200 + assert status.headers["cache-control"] == "no-store" + assert status.json() == { + "date": "2026-10-24", + "status": "unchecked", + "checked": False, + "last_checked_at": 1220.0, + "last_checkin_at": None, + "error": None, + "next_checkin_at": 1729766400.0, + } + assert manager.checkins.calls == [("status", "account-1", "manual_status")] + + denied = client.post("/admin/api/accounts/account-1/checkin") + assert denied.status_code == 403 + assert len(manager.checkins.calls) == 1 + + cross_port = client.post( + "/admin/api/accounts/account-1/checkin", + headers={ + "X-Account-Manager": "1", + "Origin": "http://localhost:8124", + }, + ) + assert cross_port.status_code == 403 + assert len(manager.checkins.calls) == 1 + + checked = client.post( + "/admin/api/accounts/account-1/checkin", + headers={"X-Account-Manager": "1", "Origin": "http://localhost:8123"}, + ) + assert checked.status_code == 200 + assert checked.json()["status"] == "checked" + assert checked.json()["checked"] is True + assert "private-token-value" not in checked.text + assert manager.checkins.calls[-1] == ("checkin", "account-1", "manual") + assert manager.snapshot()["active_account_id"] == "account-1" + + manager.checkins.status_result = { + "date": "2026-10-24", + "status": "error", + "error": "upstream token private-token-value", + "exception": "secret", + } + sanitized = client.get("/admin/api/accounts/account-1/checkin-status") + assert sanitized.status_code == 200 + assert sanitized.json()["error"] == "签到状态查询失败" + assert "private-token-value" not in sanitized.text + + unknown = client.get("/admin/api/accounts/missing/checkin-status") + assert unknown.status_code == 404 + assert client.post( + "/admin/api/accounts/missing/checkin", + headers={"X-Account-Manager": "1"}, + ).status_code == 404 + + +def test_cross_site_origin_host_and_non_loopback_peer_are_rejected(): + manager = FakeManager() + with _client(manager) as client: + cross_site = client.get( + "/admin/api/accounts", + headers={"Origin": "http://evil.example"}, + ) + assert cross_site.status_code == 403 + + bad_host = client.get("/admin/api/accounts", headers={"Host": "evil.example"}) + assert bad_host.status_code == 403 + + remote = TestClient( + _app(manager), + base_url="http://localhost:8123", + client=("192.0.2.4", 50000), + ) + with remote: + response = remote.get("/admin/api/accounts") + assert response.status_code == 403 + assert manager.login_calls == [] + + +def test_admin_page_is_loaded_as_a_no_store_package_resource(monkeypatch): + class Resource: + def joinpath(self, _name): + return self + + def read_text(self, *, encoding): + assert encoding == "utf-8" + return "Account manager" + + monkeypatch.setattr(admin_api.resources, "files", lambda _package: Resource()) + with _client(FakeManager()) as client: + for path in ("/admin", "/admin/"): + response = client.get(path) + assert response.status_code == 200 + assert response.headers["cache-control"] == "no-store" + assert "Account manager" in response.text + + +def test_unmatched_admin_path_is_no_store_and_has_generic_error(): + with _client(FakeManager()) as client: + response = client.get("/admin/api/not-real") + assert response.status_code == 404 + assert response.headers["cache-control"] == "no-store" + assert "private-token-value" not in response.text diff --git a/test_admin_ui.py b/test_admin_ui.py new file mode 100644 index 0000000..290d074 --- /dev/null +++ b/test_admin_ui.py @@ -0,0 +1,1198 @@ +"""Small browser-free checks for the packaged admin page's inline UI logic.""" + +from __future__ import annotations + +import json +import re +import shutil +import subprocess +from html.parser import HTMLParser +from pathlib import Path + +import pytest + + +ADMIN_HTML = Path(__file__).parent / "src" / "codebuddy_proxy" / "admin.html" + + +class _InlineScriptParser(HTMLParser): + def __init__(self) -> None: + super().__init__() + self._in_script = False + self.scripts: list[str] = [] + + def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: + if tag == "script": + self._in_script = True + self.scripts.append("") + + def handle_endtag(self, tag: str) -> None: + if tag == "script": + self._in_script = False + + def handle_data(self, data: str) -> None: + if self._in_script: + self.scripts[-1] += data + + +class _TranslationManifestParser(HTMLParser): + """Collect real translated static nodes for the browser-free DOM harness.""" + + def __init__(self) -> None: + super().__init__() + self.nodes: list[dict[str, object]] = [] + self._stack: list[tuple[str, dict[str, object] | None]] = [] + + def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: + attributes = {key: value or "" for key, value in attrs} + translated = any(key.startswith("data-i18n") for key in attributes) + node: dict[str, object] | None = None + if translated: + node = {"tag": tag, "attributes": attributes, "text": ""} + self.nodes.append(node) + if tag not in {"area", "base", "br", "col", "embed", "hr", "img", "input", "link", "meta", "param", "source", "track", "wbr"}: + self._stack.append((tag, node)) + + def handle_endtag(self, tag: str) -> None: + for index in range(len(self._stack) - 1, -1, -1): + if self._stack[index][0] == tag: + del self._stack[index:] + break + + def handle_data(self, data: str) -> None: + for _, node in self._stack: + if node is not None: + node["text"] = str(node["text"]) + data + + +def _inline_script() -> str: + parser = _InlineScriptParser() + parser.feed(ADMIN_HTML.read_text(encoding="utf-8")) + assert parser.scripts, "admin page must have its self-contained controller script" + return parser.scripts[-1] + + +_DOM_HARNESS = r""" +const assert = require('node:assert/strict'); +const fs = require('node:fs'); +const vm = require('node:vm'); +const translationManifest = JSON.parse(fs.readFileSync(process.argv[3], 'utf8')); + +class Element { + constructor(tag = 'div', id = '') { + this.tagName = tag.toUpperCase(); + this.id = id; + this.children = []; + this.dataset = {}; + this.attributes = {}; + this.listeners = {}; + this.className = ''; + this.textContent = ''; + this.value = ''; + this.lang = ''; + this.hidden = false; + this.disabled = false; + this.href = ''; + this.style = {}; + this.classList = { add: (name) => { + const classes = new Set(this.className.split(/\s+/).filter(Boolean)); + classes.add(name); + this.className = [...classes].join(' '); + }}; + } + append(...nodes) { this.children.push(...nodes); } + prepend(node) { this.children.unshift(node); } + replaceChildren(...nodes) { this.children = [...nodes]; } + addEventListener(name, fn) { (this.listeners[name] ||= []).push(fn); } + async dispatch(name) { + for (const fn of this.listeners[name] || []) await fn({ target: this }); + } + querySelector(selector) { + return selector === 'span' ? this.children.find((child) => child.tagName === 'SPAN') || null : null; + } + querySelectorAll(selector) { return walk(this).slice(1).filter((node) => matchesSelector(node, selector)); } + setAttribute(name, value) { + this.attributes[name] = String(value); + if (name.startsWith('data-')) this.dataset[name.slice(5).replace(/-([a-z])/g, (_, c) => c.toUpperCase())] = String(value); + if (name === 'lang') this.lang = String(value); + if (name === 'id') this.id = String(value); + } + getAttribute(name) { return this.attributes[name] ?? null; } + hasAttribute(name) { return Object.hasOwn(this.attributes, name); } + removeAttribute(name) { delete this.attributes[name]; } +} + +function makeTimers() { + let nextId = 0; + const timeouts = new Map(); + const intervals = new Map(); + return { + timeouts, + intervals, + setTimeout(fn, delay) { const id = ++nextId; timeouts.set(id, { fn, delay }); return id; }, + clearTimeout(id) { timeouts.delete(id); }, + setInterval(fn, delay) { const id = ++nextId; intervals.set(id, { fn, delay }); return id; }, + clearInterval(id) { intervals.delete(id); }, + }; +} + +function makePageElements() { + const ids = [ + 'notice', 'account-list', 'refresh-button', 'associate-button', 'login-progress', + 'progress-title', 'progress-copy', 'progress-spinner', 'cancel-button', + 'retry-button', 'manual-auth-link', 'auto-switch-toggle', 'settings-status', + 'auto-checkin-toggle', 'checkin-settings-status', 'sync-status', 'language-select' + ]; + const elements = Object.fromEntries(ids.map((id) => [id, new Element( + id === 'associate-button' ? 'button' : id === 'language-select' ? 'select' : 'div', id + )])); + const label = new Element('span'); + label.textContent = '关联所选区域'; + elements['associate-button'].append(label); + return elements; +} + +function makeStaticNodes(elements = {}) { + return translationManifest.map((row) => { + const element = row.attributes && row.attributes.id && elements[row.attributes.id] + ? elements[row.attributes.id] + : new Element(row.tag); + for (const [name, value] of Object.entries(row.attributes || {})) element.setAttribute(name, value); + if (!row.attributes || !row.attributes.id || !elements[row.attributes.id]) element.textContent = row.text || ''; + return element; + }); +} + +function matchesSelector(node, selector) { + if (!node || typeof node.hasAttribute !== 'function') return false; + return selector.split(',').some((part) => { + const attributeChecks = [...part.matchAll(/\[([^\]=]+)(?:=["']([^"']*)["'])?\]/g)]; + if (!attributeChecks.length) return false; + return attributeChecks.every((match) => { + const name = match[1].trim(); + const expected = match[2]; + return node.hasAttribute(name) && (expected === undefined || node.getAttribute(name) === expected); + }); + }); +} + +function makeDocument(elements, staticNodes = []) { + const regions = ['domestic', 'global'].map((backend) => { + const button = new Element('button'); + button.dataset.backend = backend; + button.setAttribute('data-backend', backend); + return button; + }); + const listeners = {}; + const documentElement = new Element('html'); + documentElement.lang = 'zh-CN'; + const allElements = () => [...new Set([ + ...Object.values(elements).flatMap(walk), + ...regions, + ...staticNodes, + documentElement, + ])]; + return { + regions, + staticNodes, + documentElement, + title: 'CodeBuddy 账号管理', + visibilityState: 'visible', + getElementById: (id) => elements[id] || null, + querySelectorAll: (selector) => allElements().filter((node) => matchesSelector(node, selector)), + createElement: (tag) => new Element(tag), + addEventListener(name, fn) { (listeners[name] ||= []).push(fn); }, + async dispatch(name) { for (const fn of listeners[name] || []) await fn(); }, + }; +} + +function makeStorage(initial = {}, { getThrows = false, setThrows = false } = {}) { + const values = new Map(Object.entries(initial)); + return { + values, + getItem(key) { + if (getThrows) throw new Error('storage read unavailable'); + return values.has(key) ? values.get(key) : null; + }, + setItem(key, value) { + if (setThrows) throw new Error('storage write unavailable'); + values.set(key, String(value)); + }, + }; +} + +function makeWindow(timers, now = () => Date.now(), options = {}) { + const popup = { + closed: false, + opener: null, + location: { href: '' }, + document: { title: '', body: { textContent: '', style: { cssText: '' } } }, + close() { this.closed = true; }, + }; + const localStorage = Object.hasOwn(options, 'localStorage') ? options.localStorage : makeStorage(); + const navigator = Object.hasOwn(options, 'navigator') + ? options.navigator + : { language: 'zh-CN', languages: ['zh-CN'] }; + return { + popup, + localStorage, + navigator, + window: { + location: { href: 'http://127.0.0.1:8765/admin' }, + confirm: () => true, + open: () => popup, + setTimeout: timers.setTimeout, + clearTimeout: timers.clearTimeout, + setInterval: timers.setInterval, + clearInterval: timers.clearInterval, + localStorage, + navigator, + }, + }; +} + +function walk(node) { + return [node, ...(node.children || []).flatMap(walk)]; +} + +const accountFixture = (overrides = {}) => ({ + id: 'existing', backend: 'domestic', endpoint: 'https://example.test', uid: 'u1', nickname: '旧账号', + model_cooldowns: [], + checkin: { date: null, status: 'unchecked', checked: false, last_checked_at: null, last_checkin_at: null, next_checkin_at: null, error: null }, + ...overrides, +}); +const settingsFixture = (autoSwitch = false, autoCheckin = false) => ({ + auto_switch_on_429: autoSwitch, auto_checkin_enabled: autoCheckin, default_cooldown_seconds: 60 +}); +const snapshotFixture = (accounts = [accountFixture()], autoSwitch = false, serverTime = Math.floor(Date.now() / 1000), autoCheckin = false) => { + const date = new Date(serverTime * 1000 + 8 * 60 * 60 * 1000).toISOString().slice(0, 10); + return { + accounts: accounts.map((account) => ({ ...account, checkin: { ...account.checkin, date: account.checkin.date || date } })), + active_account_id: 'existing', settings: settingsFixture(autoSwitch, autoCheckin), server_time: serverTime, + }; +}; + +function startPage(document, window, fetch, clock = Date) { + const script = fs.readFileSync(process.argv[2], 'utf8'); + const DateType = typeof clock === 'function' ? clock : class TestDate extends Date { static now() { return clock.now(); } }; + vm.runInNewContext(script, { + document, window, fetch, Headers, URL, AbortController, Promise, Date: DateType, Error, + localStorage: window.localStorage, + navigator: window.navigator, + }); +} + +async function flush() { await new Promise((resolve) => setImmediate(resolve)); } + +async function verifyTerminalState(status) { + const elements = makePageElements(); + const document = makeDocument(elements); + const timers = makeTimers(); + const { window, popup } = makeWindow(timers); + const writes = []; + async function fetch(url, options = {}) { + const method = options.method || 'GET'; + if (method !== 'GET') writes.push({ url, headers: options.headers }); + let payload; + if (url.endsWith('/accounts')) { + payload = snapshotFixture(); + } else if (url.endsWith('/settings') && method === 'GET') { + payload = settingsFixture(); + } else if (url.endsWith('/logins') && method === 'POST') { + payload = { id: 'job-1', status, account_id: status === 'succeeded' ? 'new' : null, error: status === 'failed' ? '模拟登录失败' : null }; + } else { + throw new Error(`Unexpected API request: ${method} ${url}`); + } + return { ok: true, status: 200, headers: { get: () => 'application/json' }, json: async () => payload }; + } + startPage(document, window, fetch); + await flush(); + await elements['associate-button'].dispatch('click'); + assert.equal(elements['cancel-button'].hidden, true, `${status}: cancel must be hidden after terminal result`); + assert.equal(elements['cancel-button'].disabled, true, `${status}: cancel must be disabled after terminal result`); + assert.equal(elements['retry-button'].hidden, status === 'succeeded', `${status}: retry visibility must match terminal result`); + assert.equal(writes.length, 1); + assert.equal(writes[0].headers.get('X-Account-Manager'), '1'); + assert.equal(popup.opener, null, 'auth popup must not retain access to the admin window'); +} + +async function verifyCancelFailureRecovery() { + const elements = makePageElements(); + const document = makeDocument(elements); + const timers = makeTimers(); + const { window, popup } = makeWindow(timers); + const writes = []; + let deleteCount = 0; + let loginReads = 0; + async function fetch(url, options = {}) { + const method = options.method || 'GET'; + if (method !== 'GET') writes.push({ url, headers: options.headers }); + let response = { ok: true, status: 200, payload: null }; + if (url.endsWith('/accounts')) { + response.payload = snapshotFixture(); + } else if (url.endsWith('/settings') && method === 'GET') { + response.payload = settingsFixture(); + } else if (url.endsWith('/logins') && method === 'POST') { + response.payload = { id: 'job-cancel', status: 'pending' }; + } else if (url.endsWith('/logins/job-cancel') && method === 'GET') { + loginReads += 1; + response.payload = { id: 'job-cancel', status: 'pending' }; + } else if (url.endsWith('/logins/job-cancel') && method === 'DELETE') { + deleteCount += 1; + if (deleteCount === 1) response = { ok: false, status: 503, payload: { detail: '模拟取消失败' } }; + else response.payload = { id: 'job-cancel', status: 'cancelled' }; + } else { + throw new Error(`Unexpected API request: ${method} ${url}`); + } + return { ok: response.ok, status: response.status, headers: { get: () => 'application/json' }, json: async () => response.payload }; + } + startPage(document, window, fetch); + await flush(); + await elements['associate-button'].dispatch('click'); + assert.equal(elements['cancel-button'].hidden, false, 'pending login must expose cancel'); + await elements['cancel-button'].dispatch('click'); + assert.equal(deleteCount, 1); + assert.equal(elements['cancel-button'].hidden, false, 'failed cancellation must leave cancel available'); + assert.equal(elements['cancel-button'].disabled, false, 'failed cancellation must re-enable cancel'); + assert.match(elements.notice.textContent, /模拟取消失败/); + assert.ok(loginReads >= 1, 'pending job remains tracked after cancel failure'); + await elements['cancel-button'].dispatch('click'); + assert.equal(deleteCount, 2, 'user can retry cancellation'); + assert.equal(elements['cancel-button'].hidden, true, 'confirmed cancellation hides cancel'); + assert.equal(elements['retry-button'].hidden, false, 'confirmed cancellation exposes retry'); + assert.ok(writes.every((write) => write.headers.get('X-Account-Manager') === '1')); +} + +async function verifyAutoSwitchSettings() { + const elements = makePageElements(); + const document = makeDocument(elements); + const timers = makeTimers(); + const { window } = makeWindow(timers); + let remoteSetting = true; + let patchCount = 0; + let settingsReads = 0; + const writes = []; + async function fetch(url, options = {}) { + const method = options.method || 'GET'; + if (method !== 'GET') writes.push({ url, headers: options.headers, body: options.body }); + if (url.endsWith('/accounts')) return jsonResponse(snapshotFixture([accountFixture()], remoteSetting)); + if (url.endsWith('/settings') && method === 'GET') { + settingsReads += 1; + return jsonResponse(settingsFixture(remoteSetting)); + } + if (url.endsWith('/settings') && method === 'PATCH') { + patchCount += 1; + const desired = JSON.parse(options.body).auto_switch_on_429; + if (patchCount === 3) return jsonResponse({ detail: '模拟保存失败' }, false, 503); + remoteSetting = desired; + return jsonResponse(settingsFixture(remoteSetting)); + } + throw new Error(`Unexpected API request: ${method} ${url}`); + } + startPage(document, window, fetch); + await flush(); + assert.equal(elements['auto-switch-toggle'].checked, true, 'initial setting must be restored from the server'); + + elements['auto-switch-toggle'].checked = false; + await elements['auto-switch-toggle'].dispatch('change'); + assert.equal(remoteSetting, false, 'turning the switch off must persist'); + assert.equal(elements['auto-switch-toggle'].checked, false); + + elements['auto-switch-toggle'].checked = true; + await elements['auto-switch-toggle'].dispatch('change'); + assert.equal(remoteSetting, true, 'turning the switch on must persist'); + assert.equal(elements['auto-switch-toggle'].checked, true); + + elements['auto-switch-toggle'].checked = false; + await elements['auto-switch-toggle'].dispatch('change'); + assert.equal(elements['auto-switch-toggle'].checked, true, 'failed save must restore the last server value'); + assert.equal(elements['auto-switch-toggle'].disabled, false, 'failed save must re-enable the setting'); + assert.match(elements.notice.textContent, /保存结果未确认/); + assert.match(elements.notice.textContent, /模拟保存失败/); + assert.ok(settingsReads >= 2, 'an unclear save response must trigger a settings resync'); + assert.ok(writes.every((write) => write.headers.get('X-Account-Manager') === '1')); + assert.deepEqual(writes.map((write) => JSON.parse(write.body).auto_switch_on_429), [false, true, false]); +} + +function jsonResponse(payload, ok = true, status = 200) { + return { ok, status, headers: { get: () => 'application/json' }, json: async () => payload }; +} + +async function verifyCooldownPollingAndLocalExpiry() { + const elements = makePageElements(); + const document = makeDocument(elements); + const timers = makeTimers(); + const { window } = makeWindow(timers); + let nowMs = 1900000000000; + let accountReads = 0; + const unsafeModel = 'model-'; + const initialModelCooldowns = [ + { model: 'model-a', rate_limited: true, cooldown_until: 1700000003, cooldown_remaining_seconds: 3, last_rate_limited_at: 1699999997, cooldown_source: 'default' }, + { model: unsafeModel, rate_limited: true, cooldown_until: 1700000008, cooldown_remaining_seconds: 8, last_rate_limited_at: 1699999990, cooldown_source: 'retry_after' }, + ]; + const refreshedModelCooldowns = [ + { model: 'model-a', rate_limited: true, cooldown_until: 1700000010, cooldown_remaining_seconds: 7, last_rate_limited_at: 1699999997, cooldown_source: 'default' }, + { model: unsafeModel, rate_limited: false, cooldown_until: null, cooldown_remaining_seconds: 0, last_rate_limited_at: 1699999990, cooldown_source: 'retry_after' }, + ]; + async function fetch(url, options = {}) { + const method = options.method || 'GET'; + if (url.endsWith('/accounts')) { + accountReads += 1; + if (accountReads === 1) return jsonResponse(snapshotFixture([accountFixture()], false, 1700000000)); + if (accountReads === 2) { + return jsonResponse(snapshotFixture([accountFixture({ model_cooldowns: initialModelCooldowns })], true, 1700000000)); + } + return jsonResponse(snapshotFixture([accountFixture({ model_cooldowns: refreshedModelCooldowns })], true, 1700000003)); + } + if (url.endsWith('/settings') && method === 'GET') return jsonResponse(settingsFixture(false)); + throw new Error(`Unexpected API request: ${method} ${url}`); + } + startPage(document, window, fetch, { now: () => nowMs }); + await flush(); + assert.equal(elements['auto-switch-toggle'].checked, false); + + document.visibilityState = 'hidden'; + await document.dispatch('visibilitychange'); + assert.equal(accountReads, 1, 'hiding the page must not trigger an immediate poll'); + assert.ok([...timers.timeouts.values()].some((timer) => timer.delay === 60000), 'hidden page must use a slower refresh interval'); + + document.visibilityState = 'visible'; + await document.dispatch('visibilitychange'); + await flush(); + await flush(); + assert.equal(accountReads, 2, 'returning to the page must refresh immediately'); + assert.equal(elements['auto-switch-toggle'].checked, true, 'poll must pick up settings changed in another page'); + let badges = walk(elements['account-list']).filter((node) => node.className.split(/\s+/).includes('model-cooldown')); + assert.equal(badges.length, 2); + const modelABadge = badges.find((badge) => badge.textContent.includes('model-a')); + const unsafeModelBadge = badges.find((badge) => badge.textContent.includes(unsafeModel)); + assert.ok(modelABadge && unsafeModelBadge); + assert.equal(modelABadge.hidden, false); + assert.match(modelABadge.textContent, /冷却 00:03.*默认估算/); + assert.match(unsafeModelBadge.textContent, /冷却 00:08.*上游指定/); + assert.equal(walk(elements['account-list']).some((node) => node.tagName === 'IMG'), false, 'model names are rendered as text, never markup'); + + const requestsBeforeTick = accountReads; + nowMs += 1000; + for (const timer of timers.intervals.values()) timer.fn(); + assert.match(modelABadge.textContent, /冷却 00:02/); + assert.match(unsafeModelBadge.textContent, /冷却 00:07/); + nowMs += 3000; + for (const timer of timers.intervals.values()) timer.fn(); + assert.equal(modelABadge.hidden, true, 'one model badge expires independently'); + assert.equal(unsafeModelBadge.hidden, false, 'a later model cooldown remains visible'); + assert.equal(accountReads, requestsBeforeTick, 'countdown ticks must not send network requests'); + + await document.dispatch('visibilitychange'); + await flush(); + badges = walk(elements['account-list']).filter((node) => node.className.split(/\s+/).includes('model-cooldown')); + assert.equal(accountReads, 3, 'a regular snapshot refresh updates cooldowns without extra requests'); + assert.equal(badges.length, 2); + assert.match(badges[0].textContent, /model-a.*冷却 00:07/); + assert.equal(badges[1].hidden, true, 'a refreshed non-limited model is hidden'); + const requestsAfterRefresh = accountReads; + for (const timer of timers.intervals.values()) timer.fn(); + assert.equal(accountReads, requestsAfterRefresh, 'refreshed one-second countdown also remains local'); +} + +async function verifyOlderSnapshotCannotResetNewSetting() { + const elements = makePageElements(); + const document = makeDocument(elements); + const timers = makeTimers(); + const { window } = makeWindow(timers); + let accountReads = 0; + let releaseStaleSnapshot; + async function fetch(url, options = {}) { + const method = options.method || 'GET'; + if (url.endsWith('/accounts')) { + accountReads += 1; + if (accountReads === 1) return jsonResponse(snapshotFixture([accountFixture()], false)); + if (accountReads === 2) return new Promise((resolve) => { releaseStaleSnapshot = () => resolve(jsonResponse(snapshotFixture([accountFixture()], false))); }); + return jsonResponse(snapshotFixture([accountFixture()], false)); + } + if (url.endsWith('/settings') && method === 'GET') return jsonResponse(settingsFixture(false)); + if (url.endsWith('/settings') && method === 'PATCH') return jsonResponse(settingsFixture(true)); + throw new Error(`Unexpected API request: ${method} ${url}`); + } + startPage(document, window, fetch); + await flush(); + const olderRefresh = elements['refresh-button'].dispatch('click'); + await flush(); + assert.equal(typeof releaseStaleSnapshot, 'function'); + elements['auto-switch-toggle'].checked = true; + await elements['auto-switch-toggle'].dispatch('change'); + releaseStaleSnapshot(); + await olderRefresh; + assert.equal(elements['auto-switch-toggle'].checked, true, 'an older in-flight account snapshot must not undo a newer saved setting'); +} + +async function verifyForegroundRefreshOwnsBusyState() { + const elements = makePageElements(); + const document = makeDocument(elements); + const timers = makeTimers(); + const { window } = makeWindow(timers); + let accountReads = 0; + let releaseForegroundSnapshot; + async function fetch(url, options = {}) { + const method = options.method || 'GET'; + if (url.endsWith('/accounts')) { + accountReads += 1; + if (accountReads === 1) return jsonResponse(snapshotFixture()); + if (accountReads === 2) { + return new Promise((resolve) => { + releaseForegroundSnapshot = () => resolve(jsonResponse(snapshotFixture())); + }); + } + return jsonResponse(snapshotFixture([accountFixture({ nickname: '后台同步账号' })])); + } + if (url.endsWith('/settings') && method === 'GET') return jsonResponse(settingsFixture(false)); + throw new Error(`Unexpected API request: ${method} ${url}`); + } + startPage(document, window, fetch); + await flush(); + + const foregroundRefresh = elements['refresh-button'].dispatch('click'); + await flush(); + assert.equal(typeof releaseForegroundSnapshot, 'function'); + assert.equal(elements['refresh-button'].disabled, true, 'manual refresh must mark the button busy'); + + document.visibilityState = 'visible'; + await document.dispatch('visibilitychange'); + await flush(); + await flush(); + assert.equal(accountReads, 3, 'visibility sync must run while a manual refresh is pending'); + assert.equal(elements['refresh-button'].disabled, true, 'background sync must not clear foreground busy state'); + + releaseForegroundSnapshot(); + await foregroundRefresh; + assert.equal(elements['refresh-button'].disabled, false, 'foreground completion must clear busy state after a newer background read'); +} + +function findAccountCard(elements, id) { + return walk(elements['account-list']).find((node) => node.tagName === 'ARTICLE' && node.dataset.accountId === id); +} + +function checkinButtons(card) { + const buttons = walk(card).filter((node) => node.tagName === 'BUTTON'); + return { + query: buttons.find((node) => node.textContent === '查询签到'), + sign: buttons.find((node) => node.textContent === '签到'), + }; +} + +async function flushMany(count = 8) { for (let i = 0; i < count; i += 1) await flush(); } + +function makeLanguagePage(config = {}) { + const elements = makePageElements(); + const document = makeDocument(elements, makeStaticNodes(elements)); + const timers = makeTimers(); + const windowOptions = {}; + if (Object.hasOwn(config, 'localStorage')) windowOptions.localStorage = config.localStorage; + if (Object.hasOwn(config, 'navigator')) windowOptions.navigator = config.navigator; + const { window } = makeWindow(timers, () => Date.now(), windowOptions); + const fetches = []; + const accounts = config.accounts || [accountFixture()]; + const serverTime = Math.floor(Date.now() / 1000); + async function fetch(url, options = {}) { + fetches.push({ url, method: options.method || 'GET' }); + if (url.endsWith('/accounts')) return jsonResponse(snapshotFixture(accounts, false, serverTime)); + if (url.endsWith('/settings') && (options.method || 'GET') === 'GET') return jsonResponse(settingsFixture(false)); + if (url.endsWith('/checkin-status')) return jsonResponse(config.checkinStatus || { status: 'unchecked', checked: false }); + throw new Error(`Unexpected API request: ${options.method || 'GET'} ${url}`); + } + startPage(document, window, fetch); + return { elements, document, window, fetches, localStorage: window.localStorage }; +} + +function staticTranslation(document, key) { + return document.staticNodes.find((node) => node.getAttribute('data-i18n') === key); +} + +function assertEnglishShell(page) { + const { elements, document } = page; + assert.equal(elements['language-select'].value, 'en'); + assert.equal(document.documentElement.lang, 'en'); + assert.match(document.title, /CodeBuddy.*account/i); + assert.match(elements['language-select'].getAttribute('aria-label'), /language/i); + assert.equal(staticTranslation(document, 'hero.title').textContent, 'Account manager'); + assert.equal(staticTranslation(document, 'button.refresh').textContent, 'Refresh list'); + assert.equal(staticTranslation(document, 'button.refresh').getAttribute('aria-label'), 'Refresh account list'); +} + +function assertChineseShell(page) { + const { elements, document } = page; + assert.equal(elements['language-select'].value, 'zh-CN'); + assert.equal(document.documentElement.lang, 'zh-CN'); + assert.equal(document.title, 'CodeBuddy 账号管理'); + assert.equal(elements['language-select'].getAttribute('aria-label'), '页面语言'); + assert.equal(staticTranslation(document, 'hero.title').textContent, '账号管理'); + assert.equal(staticTranslation(document, 'button.refresh').textContent, '刷新列表'); + assert.equal(staticTranslation(document, 'button.refresh').getAttribute('aria-label'), '刷新账号列表'); +} + +async function verifyLanguagePreferenceAndDynamicTranslations() { + const key = 'codebuddy.admin.language'; + const storage = makeStorage({ [key]: 'en' }); + const serverTime = Math.floor(Date.now() / 1000); + const accounts = [ + accountFixture({ + id: 'existing', nickname: '账号甲', uid: 'uid-alpha', + model_cooldowns: [ + { model: 'model-alpha', rate_limited: true, cooldown_until: serverTime + 120, cooldown_remaining_seconds: 120, last_rate_limited_at: serverTime - 1, cooldown_source: 'retry_after' }, + { model: 'embedding-v2', rate_limited: true, cooldown_until: serverTime + 240, cooldown_remaining_seconds: 240, last_rate_limited_at: serverTime - 1, cooldown_source: 'default' }, + ], + checkin: { date: null, status: 'unchecked', checked: false, error: null }, + }), + accountFixture({ + id: 'closed', nickname: '签到活动账号', uid: 'uid-closed', + checkin: { date: null, status: 'error', checked: null, error: '签到活动未开放' }, + }), + accountFixture({ + id: 'unknown-error', nickname: '未知错误账号', uid: 'uid-unknown', + checkin: { date: null, status: 'error', checked: null, error: '__proto__' }, + }), + ]; + const page = makeLanguagePage({ + localStorage: storage, + navigator: { language: 'zh-CN', languages: ['zh-CN'] }, + accounts, + }); + await flushMany(12); + assertEnglishShell(page); + assert.equal(storage.getItem(key), 'en', 'saved language takes priority over browser preference'); + + const accountText = (id) => walk(findAccountCard(page.elements, id)).map((node) => node.textContent).join(' '); + const englishChecked = accountText('existing'); + assert.match(englishChecked, /Today's check-in: Not checked in/); + assert.match(englishChecked, /Model model-alpha .*cooling down .*upstream specified/); + assert.match(englishChecked, /Model embedding-v2 .*estimated default/); + assert.match(englishChecked, /账号甲/, 'account names remain user data and are not translated'); + assert.match(englishChecked, /uid-alpha/); + const englishClosed = accountText('closed'); + assert.match(englishClosed, /Check-in is not available yet/); + assert.match(accountText('unknown-error'), /__proto__/); + assert.doesNotMatch(accountText('unknown-error'), /\[object Object\]|undefined/); + + const requestCount = page.fetches.length; + const originalUrl = page.window.location.href; + page.elements['language-select'].value = 'zh-CN'; + await page.elements['language-select'].dispatch('change'); + await flushMany(); + assertChineseShell(page); + assert.equal(storage.getItem(key), 'zh-CN', 'language selection is persisted'); + assert.equal(page.fetches.length, requestCount, 'switching language must not issue API requests'); + assert.equal(page.window.location.href, originalUrl, 'switching language must not reload or navigate'); + const chineseChecked = accountText('existing'); + assert.match(chineseChecked, /今日签到:未签到/); + assert.match(chineseChecked, /模型 model-alpha · 冷却/); + assert.match(chineseChecked, /模型 embedding-v2 · 冷却/); + assert.match(accountText('closed'), /签到活动未开放/); + + const reloaded = makeLanguagePage({ + localStorage: storage, + navigator: { language: 'en-US', languages: ['en-US'] }, + accounts, + }); + await flushMany(12); + assertChineseShell(reloaded); +} + +async function verifyCheckinQueryNoticeRelocalizes() { + const account = accountFixture({ + id: 'notice-account', + nickname: '签到显示账号', + uid: 'uid-notice', + checkin: { date: null, status: 'unchecked', checked: false }, + }); + const page = makeLanguagePage({ + navigator: { language: 'zh-CN', languages: ['zh-CN'] }, + accounts: [account], + checkinStatus: { status: 'unchecked', checked: false }, + }); + await flushMany(12); + const card = findAccountCard(page.elements, 'notice-account'); + assert.ok(card, 'account card should render before querying its check-in state'); + const queryButton = checkinButtons(card).query; + assert.ok(queryButton, 'account card should expose a manual check-in query'); + await queryButton.dispatch('click'); + assert.match(page.elements.notice.textContent, /签到显示账号.*今日签到状态:未签到/); + + page.elements['language-select'].value = 'en'; + await page.elements['language-select'].dispatch('change'); + assert.match(page.elements.notice.textContent, /签到显示账号.*today's check-in status is Not checked in/i); + + page.elements['language-select'].value = 'zh-CN'; + await page.elements['language-select'].dispatch('change'); + assert.match(page.elements.notice.textContent, /签到显示账号.*今日签到状态:未签到/); +} + +async function verifyHttpErrorNoticeRelocalizes() { + const elements = makePageElements(); + const document = makeDocument(elements); + const timers = makeTimers(); + const { window } = makeWindow(timers); + window.localStorage.setItem('codebuddy.admin.language', 'en'); + const requests = []; + async function fetch(url) { + requests.push(url); + if (url.endsWith('/accounts')) return jsonResponse(null, false, 503); + if (url.endsWith('/settings')) return jsonResponse(settingsFixture()); + throw new Error(`Unexpected API request: ${url}`); + } + startPage(document, window, fetch); + await flush(); + assert.match(elements.notice.textContent, /Request failed.*503/); + elements['language-select'].value = 'zh-CN'; + await elements['language-select'].dispatch('change'); + assert.match(elements.notice.textContent, /请求失败.*503/); + elements['language-select'].value = 'en'; + await elements['language-select'].dispatch('change'); + assert.match(elements.notice.textContent, /Request failed.*503/); + assert.equal(requests.length, 2, 'language changes must not retry failed requests'); +} + +async function verifyLanguageFallbackAndUnavailableStorage() { + const key = 'codebuddy.admin.language'; + const cases = [ + { saved: '__proto__', navigator: { language: 'zh-TW', languages: ['zh-TW'] }, expected: 'zh-CN' }, + { saved: 'fr', navigator: { language: 'en-US', languages: ['en-US'] }, expected: 'en' }, + { saved: 'invalid-language', navigator: { languages: ['zh-Hans-CN'] }, expected: 'zh-CN' }, + { saved: null, navigator: { language: 'de-DE' }, expected: 'en' }, + { saved: null, navigator: undefined, expected: 'zh-CN' }, + ]; + for (const scenario of cases) { + const storage = scenario.saved === null ? makeStorage() : makeStorage({ [key]: scenario.saved }); + const page = makeLanguagePage({ localStorage: storage, navigator: scenario.navigator }); + await flushMany(10); + assert.equal(page.elements['language-select'].value, scenario.expected, `preference ${scenario.saved} should fall back correctly`); + assert.equal(page.document.documentElement.lang, scenario.expected); + } + + const failedReadStorage = makeStorage({}, { getThrows: true }); + const readFailurePage = makeLanguagePage({ + localStorage: failedReadStorage, + navigator: { language: 'en-US', languages: ['en-US'] }, + }); + await flushMany(10); + assertEnglishShell(readFailurePage); + readFailurePage.elements['language-select'].value = 'zh-CN'; + await readFailurePage.elements['language-select'].dispatch('change'); + assertChineseShell(readFailurePage); + assert.equal(failedReadStorage.values.get(key), 'zh-CN', 'a read failure must not disable later writes or switching'); + + const failedWriteStorage = makeStorage({}, { setThrows: true }); + const writeFailurePage = makeLanguagePage({ + localStorage: failedWriteStorage, + navigator: { language: 'zh-CN', languages: ['zh-CN'] }, + }); + await flushMany(10); + assertChineseShell(writeFailurePage); + writeFailurePage.elements['language-select'].value = 'en'; + await writeFailurePage.elements['language-select'].dispatch('change'); + assertEnglishShell(writeFailurePage); + + const absentStoragePage = makeLanguagePage({ + localStorage: null, + navigator: { language: 'zh-CN', languages: ['zh-CN'] }, + }); + await flushMany(10); + absentStoragePage.elements['language-select'].value = 'en'; + await absentStoragePage.elements['language-select'].dispatch('change'); + assertEnglishShell(absentStoragePage); +} + +async function verifyLanguageSwitchDuringPendingOperations() { + const elements = makePageElements(); + const document = makeDocument(elements, makeStaticNodes(elements)); + const timers = makeTimers(); + const storage = makeStorage(); + const { window, popup } = makeWindow(timers, () => Date.now(), { + localStorage: storage, + navigator: { language: 'zh-CN', languages: ['zh-CN'] }, + }); + const fetches = []; + let releaseLoginRead; + let releaseSettingsSave; + async function fetch(url, options = {}) { + const method = options.method || 'GET'; + fetches.push({ url, method }); + if (url.endsWith('/accounts')) return jsonResponse(snapshotFixture([accountFixture()], false)); + if (url.endsWith('/settings') && method === 'GET') return jsonResponse(settingsFixture(false)); + if (url.endsWith('/settings') && method === 'PATCH') { + return new Promise((resolve) => { releaseSettingsSave = () => resolve(jsonResponse(settingsFixture(true))); }); + } + if (url.endsWith('/logins') && method === 'POST') return jsonResponse({ id: 'pending-language', status: 'pending' }); + if (url.endsWith('/logins/pending-language') && method === 'GET') { + return new Promise((resolve) => { + releaseLoginRead = () => resolve(jsonResponse({ id: 'pending-language', status: 'failed', error: '模拟登录失败' })); + }); + } + throw new Error(`Unexpected API request: ${method} ${url}`); + } + startPage(document, window, fetch); + await flushMany(12); + assertChineseShell({ elements, document }); + + const loginStart = elements['associate-button'].dispatch('click'); + await flushMany(12); + await loginStart; + assert.equal(typeof releaseLoginRead, 'function', 'pending login polling must stay active'); + assert.equal(elements['associate-button'].disabled, true); + assert.equal(popup.closed, false); + + elements['auto-switch-toggle'].checked = true; + const settingsSave = elements['auto-switch-toggle'].dispatch('change'); + await flushMany(8); + assert.equal(typeof releaseSettingsSave, 'function', 'settings write must be pending'); + assert.equal(elements['auto-switch-toggle'].disabled, true); + assert.match(elements['settings-status'].textContent, /正在保存设置/); + + const requestCount = fetches.length; + elements['language-select'].value = 'en'; + await elements['language-select'].dispatch('change'); + await flushMany(); + assertEnglishShell({ elements, document }); + assert.equal(fetches.length, requestCount, 'a pending operation is not restarted by translation'); + assert.equal(window.location.href, 'http://127.0.0.1:8765/admin'); + assert.equal(elements['associate-button'].disabled, true, 'login remains busy after switching language'); + assert.equal(elements['auto-switch-toggle'].disabled, true, 'setting remains busy after switching language'); + assert.match(elements['progress-title'].textContent, /sign-in/i); + assert.match(elements['progress-copy'].textContent, /proxy|address/i); + assert.match(elements['settings-status'].textContent, /Saving setting/i); + + releaseLoginRead(); + await flushMany(12); + assert.equal(elements['associate-button'].disabled, false); + assert.equal(popup.closed, true, 'the terminal login result closes its popup normally'); + assert.match(elements['progress-title'].textContent, /Linking failed/); + assert.match(elements['progress-copy'].textContent, /模拟登录失败/); + + releaseSettingsSave(); + await settingsSave; + assert.equal(elements['auto-switch-toggle'].checked, true); + assert.equal(elements['auto-switch-toggle'].disabled, false); + assert.match(elements['settings-status'].textContent, /Setting loaded/i); + assert.match(elements.notice.textContent, /Automatic account switching on 429 is enabled/); + + elements['language-select'].value = 'zh-CN'; + await elements['language-select'].dispatch('change'); + assertChineseShell({ elements, document }); + assert.match(elements['progress-title'].textContent, /关联失败/); + assert.match(elements['notice'].textContent, /已开启遇到 429 时自动切换账号/); +} + +async function verifyIndependentSettingsWrites() { + const elements = makePageElements(); + const document = makeDocument(elements); + const timers = makeTimers(); + const { window } = makeWindow(timers); + const remote = { autoSwitch: true, autoCheckin: false }; + const writes = []; + let releaseStaleSwitch; + let checkinPatchCount = 0; + async function fetch(url, options = {}) { + const method = options.method || 'GET'; + if (url.endsWith('/accounts')) return jsonResponse(snapshotFixture([accountFixture()], remote.autoSwitch, undefined, remote.autoCheckin)); + if (url.endsWith('/settings') && method === 'GET') return jsonResponse(settingsFixture(remote.autoSwitch, remote.autoCheckin)); + if (url.endsWith('/settings') && method === 'PATCH') { + const body = JSON.parse(options.body); + writes.push({ body, headers: options.headers }); + if (Object.hasOwn(body, 'auto_switch_on_429')) { + remote.autoSwitch = body.auto_switch_on_429; + const delayedResponse = settingsFixture(remote.autoSwitch, remote.autoCheckin); + return new Promise((resolve) => { releaseStaleSwitch = () => resolve(jsonResponse(delayedResponse)); }); + } + if (Object.hasOwn(body, 'auto_checkin_enabled')) { + checkinPatchCount += 1; + if (checkinPatchCount === 1) remote.autoCheckin = body.auto_checkin_enabled; + if (checkinPatchCount === 2) return jsonResponse({ detail: '模拟自动签到保存失败' }, false, 503); + return jsonResponse(settingsFixture(remote.autoSwitch, remote.autoCheckin)); + } + } + throw new Error(`Unexpected API request: ${method} ${url}`); + } + startPage(document, window, fetch); + await flushMany(); + assert.equal(elements['auto-switch-toggle'].checked, true); + assert.equal(elements['auto-checkin-toggle'].checked, false); + + elements['auto-switch-toggle'].checked = false; + const switchSave = elements['auto-switch-toggle'].dispatch('change'); + await flush(); + assert.equal(elements['auto-switch-toggle'].disabled, true); + assert.equal(elements['auto-checkin-toggle'].disabled, false, 'one setting must not lock the other control'); + + elements['auto-checkin-toggle'].checked = true; + await elements['auto-checkin-toggle'].dispatch('change'); + assert.equal(remote.autoCheckin, true); + assert.equal(typeof releaseStaleSwitch, 'function'); + releaseStaleSwitch(); + await switchSave; + assert.equal(elements['auto-switch-toggle'].checked, false); + assert.equal(elements['auto-checkin-toggle'].checked, true, 'a stale switch response must not overwrite the independent check-in setting'); + + elements['auto-checkin-toggle'].checked = false; + await elements['auto-checkin-toggle'].dispatch('change'); + assert.equal(elements['auto-checkin-toggle'].checked, true, 'failed setting save must restore the value confirmed by GET'); + assert.equal(elements['auto-switch-toggle'].checked, false, 'check-in save failure must preserve the 429 setting'); + assert.match(elements.notice.textContent, /模拟自动签到保存失败/); + assert.ok(writes.every((write) => write.headers.get('X-Account-Manager') === '1')); + assert.deepEqual(writes.map((write) => Object.keys(write.body)), [['auto_switch_on_429'], ['auto_checkin_enabled'], ['auto_checkin_enabled']]); +} + +async function verifyCheckinDiscoveryAndActions() { + const elements = makePageElements(); + const document = makeDocument(elements); + const timers = makeTimers(); + const { window } = makeWindow(timers); + const serverTime = Math.floor(Date.now() / 1000); + const today = new Date(serverTime * 1000 + 8 * 60 * 60 * 1000).toISOString().slice(0, 10); + const accounts = [ + accountFixture({ id: 'existing', nickname: '当前账号', checkin: { date: null, status: 'unknown', checked: null } }), + accountFixture({ id: 'second', nickname: '第二账号', checkin: { date: null, status: 'unknown', checked: null } }), + accountFixture({ id: 'third', nickname: '第三账号', checkin: { date: null, status: 'unknown', checked: null } }), + ]; + const serverCheckins = { + existing: { date: today, status: 'unchecked', checked: false }, + second: { date: today, status: 'checked', checked: true, last_checkin_at: serverTime, next_checkin_at: serverTime + 86400 }, + third: { date: today, status: 'error', checked: null, error: '' }, + }; + const statusReads = []; + const writes = []; + let concurrentReads = 0; + let maxConcurrentReads = 0; + let activations = 0; + let snapshotHasFreshCheckins = false; + async function fetch(url, options = {}) { + const method = options.method || 'GET'; + if (method !== 'GET') writes.push({ url, method, headers: options.headers }); + if (url.endsWith('/accounts') && method === 'GET') { + const rows = snapshotHasFreshCheckins ? accounts.map((account) => ({ ...account, checkin: serverCheckins[account.id] })) : accounts; + snapshotHasFreshCheckins = true; + return jsonResponse(snapshotFixture(rows, false, serverTime)); + } + if (url.endsWith('/settings') && method === 'GET') return jsonResponse(settingsFixture(false)); + const match = url.match(/\/accounts\/([^/]+)\/checkin-status$/); + if (match && method === 'GET') { + const id = decodeURIComponent(match[1]); + statusReads.push(id); + concurrentReads += 1; + maxConcurrentReads = Math.max(maxConcurrentReads, concurrentReads); + await new Promise((resolve) => setImmediate(resolve)); + concurrentReads -= 1; + return jsonResponse(serverCheckins[id]); + } + const signMatch = url.match(/\/accounts\/([^/]+)\/checkin$/); + if (signMatch && method === 'POST') { + const id = decodeURIComponent(signMatch[1]); + serverCheckins[id] = { date: today, status: 'checked', checked: true, last_checkin_at: serverTime }; + return jsonResponse(serverCheckins[id]); + } + if (url.endsWith('/activate')) activations += 1; + throw new Error(`Unexpected API request: ${method} ${url}`); + } + startPage(document, window, fetch); + await flushMany(12); + assert.deepEqual([...statusReads].sort(), ['existing', 'second', 'third']); + assert.ok(maxConcurrentReads <= 2, 'initial account checks must use at most two concurrent requests'); + assert.equal(writes.filter((write) => write.method === 'POST' && write.url.endsWith('/checkin')).length, 0, + 'initial discovery and snapshot synchronization must never auto-sign'); + + const activeCard = findAccountCard(elements, 'existing'); + const secondCard = findAccountCard(elements, 'second'); + const thirdCard = findAccountCard(elements, 'third'); + assert.ok(activeCard && secondCard && thirdCard); + assert.ok(checkinButtons(activeCard).query && checkinButtons(activeCard).sign, 'the active account also has check-in controls'); + assert.equal(checkinButtons(activeCard).sign.disabled, false, 'a confirmed unchecked account can be signed'); + assert.equal(checkinButtons(secondCard).sign.disabled, true, 'a checked account cannot be signed again'); + assert.match(walk(secondCard).map((node) => node.textContent).join(' '), /下次计划/); + assert.ok(walk(thirdCard).some((node) => typeof node.textContent === 'string' && node.textContent.includes('')), 'upstream error text must be rendered as text'); + assert.equal(checkinButtons(thirdCard).sign.disabled, false, 'an error state remains actionable so the server can check first'); + + await checkinButtons(activeCard).sign.dispatch('click'); + assert.equal(serverCheckins.existing.status, 'checked'); + assert.equal(checkinButtons(findAccountCard(elements, 'existing')).sign.disabled, true); + await checkinButtons(findAccountCard(elements, 'existing')).query.dispatch('click'); + assert.equal(statusReads.filter((id) => id === 'existing').length, 2, 'manual query calls the live status endpoint'); + document.visibilityState = 'visible'; + await document.dispatch('visibilitychange'); + await flushMany(); + assert.match(walk(findAccountCard(elements, 'existing')).map((node) => node.textContent).join(' '), /今日签到:已签到/, + 'a later account snapshot must reflect server-side status changes'); + assert.equal(activations, 0, 'check-in actions never switch the current account'); + assert.ok(writes.some((write) => write.method === 'POST' && write.url.endsWith('/checkin'))); + assert.ok(writes.every((write) => write.headers.get('X-Account-Manager') === '1')); + + const readsBeforePoll = statusReads.length; + await document.dispatch('visibilitychange'); + await flushMany(); + assert.equal(statusReads.length, readsBeforePoll, '15-second account snapshots must not re-query known check-in states'); +} + +async function verifyCheckinRechecksAfterBeijingMidnight() { + const elements = makePageElements(); + const document = makeDocument(elements); + const timers = makeTimers(); + const { window } = makeWindow(timers); + let nowMs = Date.UTC(2026, 9, 23, 15, 59, 50); + let accountReads = 0; + let statusReads = 0; + let serverCheckin = { date: '2026-10-23', status: 'unknown', checked: null }; + async function fetch(url, options = {}) { + const method = options.method || 'GET'; + if (url.endsWith('/accounts') && method === 'GET') { + accountReads += 1; + return jsonResponse(snapshotFixture([accountFixture({ checkin: serverCheckin })], false, Math.floor(nowMs / 1000))); + } + if (url.endsWith('/settings') && method === 'GET') return jsonResponse(settingsFixture(false)); + if (url.endsWith('/accounts/existing/checkin-status') && method === 'GET') { + statusReads += 1; + const date = statusReads === 1 ? '2026-10-23' : '2026-10-24'; + serverCheckin = { date, status: 'unchecked', checked: false }; + return jsonResponse(serverCheckin); + } + throw new Error(`Unexpected API request: ${method} ${url}`); + } + startPage(document, window, fetch, { now: () => nowMs }); + await flushMany(); + assert.equal(statusReads, 1); + nowMs += 20000; + document.visibilityState = 'visible'; + await document.dispatch('visibilitychange'); + await flushMany(); + assert.ok(accountReads >= 2); + assert.equal(statusReads, 2, 'Beijing date change must trigger one fresh underlying status check'); + await document.dispatch('visibilitychange'); + await flushMany(); + assert.equal(statusReads, 2, 'same-day account polling must not repeat underlying status checks'); + assert.match(walk(elements['account-list']).map((node) => node.textContent).join(' '), /未签到/); +} + +async function verifyAutoCheckinIsBackendOwned() { + const elements = makePageElements(); + const document = makeDocument(elements); + const timers = makeTimers(); + const { window } = makeWindow(timers); + let nowMs = Date.UTC(2026, 9, 23, 15, 59, 50); + let setting = false; + const statusReads = []; + const signWrites = []; + const serverCheckins = { + existing: { date: null, status: 'unknown', checked: null }, + second: { date: null, status: 'unknown', checked: null }, + }; + const accounts = [ + accountFixture({ id: 'existing', checkin: serverCheckins.existing }), + accountFixture({ id: 'second', nickname: '第二账号', checkin: serverCheckins.second }), + ]; + const currentBeijingDate = () => new Date(nowMs + 8 * 60 * 60 * 1000).toISOString().slice(0, 10); + async function fetch(url, options = {}) { + const method = options.method || 'GET'; + if (url.endsWith('/accounts') && method === 'GET') { + return jsonResponse(snapshotFixture(accounts.map((account) => ({ ...account, checkin: serverCheckins[account.id] })), false, Math.floor(nowMs / 1000), setting)); + } + if (url.endsWith('/settings') && method === 'GET') return jsonResponse(settingsFixture(false, setting)); + if (url.endsWith('/settings') && method === 'PATCH') { + const body = JSON.parse(options.body); + assert.deepEqual(Object.keys(body), ['auto_checkin_enabled']); + setting = body.auto_checkin_enabled; + return jsonResponse(settingsFixture(false, setting)); + } + const statusMatch = url.match(/\/accounts\/([^/]+)\/checkin-status$/); + if (statusMatch && method === 'GET') { + const id = decodeURIComponent(statusMatch[1]); + statusReads.push(id); + const result = id === 'existing' + ? { date: currentBeijingDate(), status: 'unchecked', checked: false } + : { date: currentBeijingDate(), status: 'checked', checked: true }; + serverCheckins[id] = result; + return jsonResponse(result); + } + const signMatch = url.match(/\/accounts\/([^/]+)\/checkin$/); + if (signMatch && method === 'POST') { + const id = decodeURIComponent(signMatch[1]); + signWrites.push(id); + serverCheckins[id] = { date: currentBeijingDate(), status: 'checked', checked: true }; + return jsonResponse(serverCheckins[id]); + } + throw new Error(`Unexpected API request: ${method} ${url}`); + } + startPage(document, window, fetch, { now: () => nowMs }); + await flushMany(); + assert.equal(elements['auto-checkin-toggle'].checked, false, 'global auto check-in defaults to off'); + assert.deepEqual([...statusReads].sort(), ['existing', 'second'], 'initial unknown states are queried once per account'); + assert.deepEqual(signWrites, [], 'initial status discovery only reads; it does not sign'); + + elements['auto-checkin-toggle'].checked = true; + await elements['auto-checkin-toggle'].dispatch('change'); + assert.equal(setting, true); + assert.match(elements.notice.textContent, /代理正在检查全部关联账号/); + assert.deepEqual(signWrites, [], 'enabling the setting only saves it; the proxy owns automatic check-ins'); + assert.equal(elements['auto-checkin-toggle'].checked, true); + + // Simulate an automatic check-in performed by the backend, then let the page synchronize its snapshot. + serverCheckins.existing = { date: currentBeijingDate(), status: 'checked', checked: true }; + document.visibilityState = 'visible'; + await document.dispatch('visibilitychange'); + await flushMany(); + assert.match(walk(findAccountCard(elements, 'existing')).map((node) => node.textContent).join(' '), /今日签到:已签到/); + assert.deepEqual(signWrites, [], 'background snapshot synchronization never sends a check-in write'); + + // A Beijing date change performs one status GET per newly unknown account, but no automatic POST. + nowMs += 20000; + serverCheckins.existing = { date: currentBeijingDate() === '2026-10-24' ? '2026-10-23' : currentBeijingDate(), status: 'checked', checked: true }; + document.visibilityState = 'visible'; + await document.dispatch('visibilitychange'); + await flushMany(); + assert.equal(statusReads.filter((id) => id === 'existing').length, 2, 'a new Beijing day performs one status check'); + assert.deepEqual(signWrites, [], 'cross-day rediscovery only reads; automatic writes remain backend-owned'); + + // Manual per-account action still uses the write endpoint. + serverCheckins.existing = { date: currentBeijingDate(), status: 'unchecked', checked: false }; + await checkinButtons(findAccountCard(elements, 'existing')).sign.dispatch('click'); + assert.deepEqual(signWrites, ['existing']); +} + +(async () => { + for (const status of ['succeeded', 'cancelled', 'failed']) await verifyTerminalState(status); + await verifyCancelFailureRecovery(); + await verifyAutoSwitchSettings(); + await verifyLanguagePreferenceAndDynamicTranslations(); + await verifyHttpErrorNoticeRelocalizes(); + await verifyCheckinQueryNoticeRelocalizes(); + await verifyLanguageFallbackAndUnavailableStorage(); + await verifyLanguageSwitchDuringPendingOperations(); + await verifyCooldownPollingAndLocalExpiry(); + await verifyOlderSnapshotCannotResetNewSetting(); + await verifyForegroundRefreshOwnsBusyState(); + await verifyIndependentSettingsWrites(); + await verifyCheckinDiscoveryAndActions(); + await verifyCheckinRechecksAfterBeijingMidnight(); + await verifyAutoCheckinIsBackendOwned(); +})().catch((error) => { console.error(error); process.exitCode = 1; }); +""" + + +def test_admin_inline_script_syntax_and_terminal_controls(tmp_path: Path) -> None: + node = shutil.which("node") + if node is None: + pytest.skip("Node.js is needed to execute the admin page's browser-free mock") + + html = ADMIN_HTML.read_text(encoding="utf-8") + assert re.search(r"(?m)^\s*\[hidden\]\s*\{\s*display:\s*none\s*!important;\s*\}", html) + script_path = tmp_path / "admin-page.js" + script_path.write_text(_inline_script(), encoding="utf-8") + translation_parser = _TranslationManifestParser() + translation_parser.feed(html) + translations_path = tmp_path / "admin-translations.json" + translations_path.write_text(json.dumps(translation_parser.nodes), encoding="utf-8") + + subprocess.run([node, "--check", str(script_path)], check=True, capture_output=True, text=True) + harness_path = tmp_path / "admin-page-mock.cjs" + harness_path.write_text(_DOM_HARNESS, encoding="utf-8") + subprocess.run( + [node, str(harness_path), str(script_path), str(translations_path)], + check=True, + capture_output=True, + text=True, + ) diff --git a/test_backend_profiles.py b/test_backend_profiles.py index 1ae67db..98bd539 100644 --- a/test_backend_profiles.py +++ b/test_backend_profiles.py @@ -1,5 +1,6 @@ import json import importlib.resources +import os import stat from pathlib import Path @@ -108,7 +109,14 @@ def test_session_profiles_are_isolated_and_saved_with_metadata(tmp_path): assert saved["backend"] == "global" assert saved["endpoint"] == "https://www.codebuddy.ai" assert saved["machineId"] == "machine-1" - assert stat.S_IMODE(saved_file.stat().st_mode) == 0o600 + if os.name == "posix": + assert stat.S_IMODE(saved_file.stat().st_mode) == 0o600 + else: + # Windows does not expose POSIX mode bits as an ACL check. Verify the + # replacement is readable and the atomic writer left no temporary file. + assert os.access(saved_file, os.R_OK) + assert json.loads(saved_file.read_text(encoding="utf-8")) == saved + assert not list(saved_file.parent.glob(f".{saved_file.name}.*.tmp")) mismatch_file = tmp_path / "mismatch.json" mismatch_file.write_text( diff --git a/test_checkin_proxy.py b/test_checkin_proxy.py index b55103b..70555d1 100644 --- a/test_checkin_proxy.py +++ b/test_checkin_proxy.py @@ -59,8 +59,8 @@ def install(handler): @pytest.mark.parametrize("method,path,upstream_method,upstream_path", [ - ("GET", "/v1/checkin-status", "POST", "/v2/billing/meter/checkin-status"), - ("POST", "/v1/checkin-status", "POST", "/v2/billing/meter/checkin-status"), + ("GET", "/v1/checkin-status", "POST", "/v2/billing/meter/checkin-activity-status"), + ("POST", "/v1/checkin-status", "POST", "/v2/billing/meter/checkin-activity-status"), ("POST", "/v1/checkin", "POST", "/v2/billing/meter/daily-checkin"), ]) def test_forward_method_url_headers_and_raw_business_body(environment, method, path, upstream_method, upstream_path): diff --git a/test_client_login_management.py b/test_client_login_management.py new file mode 100644 index 0000000..da071ff --- /dev/null +++ b/test_client_login_management.py @@ -0,0 +1,176 @@ +import threading + +import pytest + +import codebuddy_proxy.codebuddy_client_demo as client_module +from codebuddy_proxy.codebuddy_client_demo import CodeBuddyClient, CodeBuddyError + + +class _Clock: + def __init__(self): + self.now = 0.0 + + def monotonic(self): + return self.now + + def sleep(self, seconds): + self.now += seconds + + +def test_begin_and_complete_login_are_separate_and_require_uid(tmp_path, monkeypatch): + client = CodeBuddyClient(session_file=tmp_path / "login.json") + requests = [] + + def request(method, path, *, headers=None, body=None, timeout=30): + requests.append((method, path, headers or {}, timeout)) + if "/auth/state" in path: + return {"authUrl": "https://login.example.test/start", "state": "state /?"} + if "/auth/token" in path: + return {"accessToken": "access-token", "refreshToken": "refresh-token"} + if "/login/account" in path: + return {"uid": "u-1", "nickname": "Test User"} + raise AssertionError(path) + + monkeypatch.setattr(client, "_request", request) + challenge = client.begin_login() + assert challenge == { + "auth_url": "https://login.example.test/start", + "state": "state /?", + } + assert not client.session_file.exists() + + client.complete_login(challenge, timeout=2) + assert client.session["account"]["uid"] == "u-1" + assert client.session["auth"]["accessToken"] == "access-token" + assert any("state=state%20%2F%3F" in path for _, path, _, _ in requests) + account_request = next(item for item in requests if "/login/account" in item[1]) + assert account_request[2]["Authorization"] == "Bearer access-token" + assert "X-Domain" in account_request[2] + assert all(item[3] <= 2 for item in requests if "/auth/state" not in item[1]) + + +def test_complete_login_total_timeout_covers_account_polling(tmp_path, monkeypatch): + client = CodeBuddyClient(session_file=tmp_path / "never-saved.json") + clock = _Clock() + monkeypatch.setattr(client_module, "time", clock) + account_calls = 0 + request_timeouts = [] + + def request(method, path, *, headers=None, body=None, timeout=30): + nonlocal account_calls + request_timeouts.append(timeout) + if "/auth/token" in path: + return {"accessToken": "access-token"} + if "/login/account" in path: + account_calls += 1 + return {"nickname": "not ready"} + raise AssertionError(path) + + monkeypatch.setattr(client, "_request", request) + with pytest.raises(CodeBuddyError, match="登录超时.*uid"): + client.complete_login({"auth_url": "https://login.test", "state": "s"}, timeout=2) + + assert account_calls == 2 + assert request_timeouts == [2, 2, 1] + assert clock.now == 2 + assert not client.session_file.exists() + + +def test_complete_login_cancellation_during_account_poll_does_not_save(tmp_path, monkeypatch): + client = CodeBuddyClient(session_file=tmp_path / "cancelled.json") + cancel = threading.Event() + + def request(method, path, *, headers=None, body=None, timeout=30): + if "/auth/token" in path: + return {"accessToken": "access-token"} + if "/login/account" in path: + cancel.set() + return {"nickname": "not ready"} + raise AssertionError(path) + + monkeypatch.setattr(client, "_request", request) + with pytest.raises(CodeBuddyError, match="登录已取消"): + client.complete_login({"state": "s"}, timeout=30, cancel_event=cancel) + assert not client.session_file.exists() + + +def test_ensure_authenticated_can_forbid_interactive_login(tmp_path, monkeypatch): + client = CodeBuddyClient(session_file=tmp_path / "empty.json") + monkeypatch.setattr(client, "refresh", lambda: False) + login_calls = [] + monkeypatch.setattr(client, "login", lambda **kwargs: login_calls.append(kwargs)) + + with pytest.raises(CodeBuddyError, match="禁止自动打开登录流程"): + client.ensure_authenticated(open_browser=False, allow_login=False) + assert login_calls == [] + assert not client.session_file.exists() + + +def test_concurrent_ensure_authentication_refreshes_once(tmp_path, monkeypatch): + client = CodeBuddyClient(session_file=tmp_path / "refresh-once.json") + client._save_session( + { + "auth": { + "accessToken": "expired-access", + "refreshToken": "shared-refresh-token", + "expiresAt": 1, + }, + "account": {"uid": "u-1"}, + "machineId": "machine-1", + } + ) + request_started = threading.Event() + release_request = threading.Event() + calls = [] + errors = [] + barrier = threading.Barrier(3) + + def refresh_request(method, path, *, headers=None, body=None, timeout=30): + calls.append((method, path, headers)) + request_started.set() + assert release_request.wait(5) + return { + "accessToken": "fresh-access", + "refreshToken": "rotated-refresh-token", + "expiresAt": int(client_module.time.time() * 1000) + 3_600_000, + } + + monkeypatch.setattr(client, "_request", refresh_request) + + def ensure(): + barrier.wait(5) + try: + client.ensure_authenticated(allow_login=False) + except Exception as exc: # collected for assertion in the test thread + errors.append(exc) + + workers = [threading.Thread(target=ensure) for _ in range(2)] + for worker in workers: + worker.start() + barrier.wait(5) + assert request_started.wait(5) + release_request.set() + for worker in workers: + worker.join(5) + + assert all(not worker.is_alive() for worker in workers) + assert errors == [] + assert len(calls) == 1 + assert client.session["auth"]["accessToken"] == "fresh-access" + assert client.session["auth"]["refreshToken"] == "rotated-refresh-token" + + +def test_login_cli_entry_reuses_begin_and_complete(tmp_path, monkeypatch): + client = CodeBuddyClient(session_file=tmp_path / "cli.json") + seen = [] + monkeypatch.setattr( + client, "begin_login", lambda: {"auth_url": "https://login.test", "state": "s"} + ) + monkeypatch.setattr(client, "complete_login", lambda challenge, timeout: seen.append((challenge, timeout))) + opened = [] + monkeypatch.setattr(client_module.webbrowser, "open", lambda url: opened.append(url)) + + client.login(open_browser=True, timeout=37) + + assert opened == ["https://login.test"] + assert seen == [({"auth_url": "https://login.test", "state": "s"}, 37)] diff --git a/test_rate_limit_proxy.py b/test_rate_limit_proxy.py new file mode 100644 index 0000000..2bf3199 --- /dev/null +++ b/test_rate_limit_proxy.py @@ -0,0 +1,806 @@ +"""Offline end-to-end checks for account-aware HTTP 429 handling.""" + +import asyncio +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta, timezone +import json +import threading +from email.utils import format_datetime + +import httpx +import pytest +from fastapi.testclient import TestClient + +import codebuddy_proxy.__main__ as proxy +from codebuddy_proxy.account_manager import AccountManager +from codebuddy_proxy.account_store import AccountStore +from codebuddy_proxy.backend_profile import DOMESTIC_PROFILE, GLOBAL_PROFILE, resolve_backend +from codebuddy_proxy.codebuddy_client_demo import CodeBuddyClient + + +MODES = [ + ("/v1/chat/completions", {"messages": [{"role": "user", "content": "hello"}]}), + ("/v1/responses", {"input": "hello"}), + ( + "/v1/messages", + {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 16}, + ), +] +MODEL_A = "rate-limit-test" +MODEL_B = "other-model" +_FUTURE_HTTP_DATE = format_datetime( + datetime.now(timezone.utc) + timedelta(minutes=5), usegmt=True +) +RETRY_AFTER_CASES = [ + pytest.param("17", "17", id="seconds"), + pytest.param(_FUTURE_HTTP_DATE, _FUTURE_HTTP_DATE, id="http-date"), + pytest.param(None, None, id="absent"), +] + + +def _add_account(store, tmp_path, *, uid, token=None, profile=DOMESTIC_PROFILE, endpoint=None): + source = CodeBuddyClient( + backend=resolve_backend( + profile=profile, + endpoint=endpoint, + session_file=tmp_path / f"source-{profile.key}-{uid}.json", + environment={}, + ) + ) + access_token = token or f"token-{profile.key}-{uid}" + source._save_session( + { + "auth": {"accessToken": access_token, "refreshToken": f"refresh-{access_token}"}, + "account": {"uid": uid, "nickname": f"User {uid}"}, + "machineId": "rate-limit-test-machine", + } + ) + account_id = store.import_client(source) + return account_id, store.get_client(account_id) + + +def _environment(tmp_path, monkeypatch, accounts): + store = AccountStore(tmp_path / "accounts") + clients = {} + account_ids = {} + for specification in accounts: + account_id, client = _add_account(store, tmp_path, **specification) + clients[specification["uid"]] = client + account_ids[specification["uid"]] = account_id + manager = AccountManager(store) + state = proxy.ProxyState( + store.active_client(), None, None, account_manager=manager + ) + monkeypatch.setattr(proxy, "proxy_state", state) + return store, manager, state, clients, account_ids + + +def _install_transport(monkeypatch, handler): + real_async_client = httpx.AsyncClient + + def mock_async_client(*args, **kwargs): + kwargs["transport"] = httpx.MockTransport(handler) + return real_async_client(*args, **kwargs) + + monkeypatch.setattr(proxy.httpx, "AsyncClient", mock_async_client) + + +def _request_body(path, body, stream): + return {"model": MODEL_A, "stream": stream, **body} + + +def _request_body_for_model(model, body, stream=False): + return {"model": model, "stream": stream, **body} + + +def _model_cooldown(account, model): + for item in account.get("model_cooldowns", []): + if isinstance(item, dict) and item.get("model") == model: + return item + return None + + +def _is_model_limited(account, model): + cooldown = _model_cooldown(account, model) + return cooldown is not None and cooldown.get("rate_limited") is True + + +def _success_response(content="recovered"): + first = { + "id": "rate-limit-response", + "object": "chat.completion.chunk", + "created": 1, + "model": "rate-limit-test", + "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": None}], + } + final = { + "id": "rate-limit-response", + "object": "chat.completion.chunk", + "created": 1, + "model": "rate-limit-test", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + } + payload = ( + f"data: {json.dumps(first)}\n\n" + f"data: {json.dumps(final)}\n\n" + "data: [DONE]\n\n" + ) + return httpx.Response( + 200, + content=payload.encode(), + headers={"content-type": "text/event-stream"}, + ) + + +def _error_response(uid, *, status=429, retry_after="60", body=None): + headers = {"content-type": "application/json"} + if retry_after is not None: + headers["Retry-After"] = retry_after + content = body if body is not None else json.dumps({"msg": f"limited-{uid}"}).encode() + return httpx.Response(status, content=content, headers=headers) + + +@pytest.mark.parametrize("path,body", MODES) +@pytest.mark.parametrize("stream", [False, True], ids=["non-stream", "stream"]) +@pytest.mark.parametrize("retry_after,expected_header", RETRY_AFTER_CASES) +def test_retry_after_is_preserved_for_all_protocol_modes_with_default_off( + tmp_path, monkeypatch, path, body, stream, retry_after, expected_header +): + store, manager, state, clients, account_ids = _environment( + tmp_path, + monkeypatch, + [{"uid": "a"}, {"uid": "b"}], + ) + seen = [] + rate_limit_logs = [] + state.write_log = lambda event, **fields: rate_limit_logs.append((event, fields)) + + async def upstream(request): + seen.append(request.headers["authorization"]) + return _error_response("a", retry_after=retry_after) + + _install_transport(monkeypatch, upstream) + try: + assert manager.snapshot()["settings"]["auto_switch_on_429"] is False + with TestClient(proxy.app, raise_server_exceptions=False) as app: + response = app.post(path, json=_request_body(path, body, stream)) + assert response.status_code == 429 + assert seen == ["Bearer token-domestic-a"] + if expected_header is None: + assert "retry-after" not in response.headers + else: + assert response.headers.get("retry-after") == expected_header + + snapshot = manager.snapshot() + limited = next(account for account in snapshot["accounts"] if account["uid"] == "a") + assert isinstance(snapshot["server_time"], float) + cooldown = _model_cooldown(limited, MODEL_A) + assert cooldown is not None + assert "cooldown_until" in cooldown + assert isinstance(cooldown["cooldown_remaining_seconds"], int) + assert cooldown["cooldown_remaining_seconds"] >= 0 + assert cooldown["last_rate_limited_at"] is not None + assert cooldown["rate_limited"] is True + assert cooldown["cooldown_source"] == ( + "retry_after" if retry_after is not None else "default" + ) + assert "rate_limited" not in limited + assert snapshot["active_account_id"] == account_ids["a"] + assert any( + fields.get("account_id") == account_ids["a"] + and fields.get("model") == MODEL_A + for _, fields in rate_limit_logs + ) + assert clients["a"].session["auth"]["accessToken"] == "token-domestic-a" + assert clients["b"].session["auth"]["accessToken"] == "token-domestic-b" + assert len(store.list_accounts()) == 2 + finally: + manager.close() + + +@pytest.mark.parametrize("path,body", MODES) +@pytest.mark.parametrize("stream", [False, True], ids=["non-stream", "stream"]) +def test_same_region_429_retries_once_on_distinct_uid_for_each_protocol_mode( + tmp_path, monkeypatch, path, body, stream +): + store, manager, _, clients, account_ids = _environment( + tmp_path, + monkeypatch, + [{"uid": "a"}, {"uid": "b"}], + ) + manager.set_settings(auto_switch_on_429=True) + seen = [] + + async def upstream(request): + authorization = request.headers["authorization"] + seen.append(authorization) + if authorization == "Bearer token-domestic-a": + return _error_response("a", retry_after="120") + return _success_response("recovered-from-b") + + _install_transport(monkeypatch, upstream) + try: + with TestClient(proxy.app, raise_server_exceptions=False) as app: + response = app.post(path, json=_request_body(path, body, stream)) + assert response.status_code == 200 + assert seen == ["Bearer token-domestic-a", "Bearer token-domestic-b"] + assert "recovered-from-b" in response.text + assert manager.store.active_account_id == account_ids["b"] + + snapshot = manager.snapshot() + by_uid = {account["uid"]: account for account in snapshot["accounts"]} + assert set(by_uid) == {"a", "b"} + assert _is_model_limited(by_uid["a"], MODEL_A) + assert not _is_model_limited(by_uid["b"], MODEL_A) + assert clients["a"].session["auth"]["accessToken"] == "token-domestic-a" + assert clients["b"].session["auth"]["accessToken"] == "token-domestic-b" + assert clients["a"].session_file != clients["b"].session_file + assert len(store.list_accounts()) == 2 + finally: + manager.close() + + +@pytest.mark.parametrize( + "candidate", + [ + pytest.param({"uid": "other-region", "profile": GLOBAL_PROFILE}, id="other-region"), + pytest.param( + { + "uid": "other-endpoint", + "endpoint": "https://domestic-alt.example.test", + }, + id="other-endpoint", + ), + ], +) +def test_429_never_uses_candidate_from_another_region_or_endpoint( + tmp_path, monkeypatch, candidate +): + store, manager, _, _, account_ids = _environment( + tmp_path, monkeypatch, [{"uid": "a"}, candidate] + ) + manager.set_settings(auto_switch_on_429=True) + seen = [] + + async def upstream(request): + seen.append((request.url.host, request.headers["authorization"])) + return _error_response("a", retry_after="60") + + _install_transport(monkeypatch, upstream) + try: + with TestClient(proxy.app, raise_server_exceptions=False) as app: + response = app.post( + "/v1/chat/completions", + json={"model": "rate-limit-test", "stream": False, + "messages": [{"role": "user", "content": "hello"}]}, + ) + assert response.status_code == 429 + assert seen == [("copilot.tencent.com", "Bearer token-domestic-a")] + assert manager.store.active_account_id == account_ids["a"] + limited = manager.snapshot()["accounts"][0] + assert _is_model_limited(limited, MODEL_A) + finally: + manager.close() + + +def test_all_same_scope_accounts_are_attempted_at_most_once_and_last_429_is_returned( + tmp_path, monkeypatch +): + _, manager, _, _, account_ids = _environment( + tmp_path, monkeypatch, [{"uid": "a"}, {"uid": "b"}, {"uid": "c"}] + ) + manager.set_settings(auto_switch_on_429=True) + retry_after_by_token = { + "Bearer token-domestic-a": "11", + "Bearer token-domestic-b": "22", + "Bearer token-domestic-c": "33", + } + uid_by_token = { + "Bearer token-domestic-a": "a", + "Bearer token-domestic-b": "b", + "Bearer token-domestic-c": "c", + } + seen = [] + + async def upstream(request): + token = request.headers["authorization"] + seen.append(token) + uid = uid_by_token[token] + return _error_response(uid, retry_after=retry_after_by_token[token]) + + _install_transport(monkeypatch, upstream) + try: + with TestClient(proxy.app, raise_server_exceptions=False) as app: + response = app.post( + "/v1/chat/completions", + json={"model": "rate-limit-test", "stream": False, + "messages": [{"role": "user", "content": "hello"}]}, + ) + assert response.status_code == 429 + assert len(seen) == 3 + assert len(set(seen)) == 3 + assert set(seen) == set(retry_after_by_token) + assert response.headers.get("retry-after") == retry_after_by_token[seen[-1]] + assert f"limited-{uid_by_token[seen[-1]]}" in response.text + assert manager.store.active_account_id in account_ids.values() + assert all( + _is_model_limited(account, MODEL_A) + for account in manager.snapshot()["accounts"] + ) + finally: + manager.close() + + +def test_when_all_other_candidates_are_cooling_request_returns_original_429_once( + tmp_path, monkeypatch +): + _, manager, _, clients, account_ids = _environment( + tmp_path, monkeypatch, [{"uid": "a"}, {"uid": "b"}, {"uid": "c"}] + ) + manager.set_settings(auto_switch_on_429=True) + all_ids = set(account_ids.values()) + for uid in ("b", "c"): + manager.handle_rate_limit( + clients[uid], "300", model=MODEL_A, excluded_account_ids=all_ids + ) + manager.store.set_active(account_ids["a"]) + seen = [] + + async def upstream(request): + seen.append(request.headers["authorization"]) + return _error_response("a", retry_after="30") + + _install_transport(monkeypatch, upstream) + try: + with TestClient(proxy.app, raise_server_exceptions=False) as app: + response = app.post( + "/v1/chat/completions", + json={"model": "rate-limit-test", "stream": False, + "messages": [{"role": "user", "content": "hello"}]}, + ) + assert response.status_code == 429 + assert seen == ["Bearer token-domestic-a"] + assert manager.store.active_account_id == account_ids["a"] + limited = {a["uid"]: a for a in manager.snapshot()["accounts"]} + assert all( + _is_model_limited(limited[uid], MODEL_A) for uid in ("a", "b", "c") + ) + finally: + manager.close() + + +@pytest.mark.parametrize("path,body", MODES) +@pytest.mark.parametrize("stream", [False, True], ids=["non-stream", "stream"]) +def test_same_account_model_cooldowns_are_isolated_across_all_protocol_modes( + tmp_path, monkeypatch, path, body, stream +): + _, manager, state, _, account_ids = _environment( + tmp_path, monkeypatch, [{"uid": "a"}, {"uid": "b"}] + ) + seen = [] + logs = [] + state.write_log = lambda event, **fields: logs.append((event, fields)) + + async def upstream(request): + payload = json.loads(request.content) + model = payload["model"] + seen.append((model, request.headers["authorization"])) + if model == MODEL_A: + return _error_response("a", retry_after="90") + return _success_response("model-b-works") + + _install_transport(monkeypatch, upstream) + manager.set_settings(auto_switch_on_429=False) + try: + with TestClient(proxy.app, raise_server_exceptions=False) as app: + limited = app.post( + path, + json=_request_body_for_model(MODEL_A, body, stream), + ) + recovered = app.post( + path, + json=_request_body_for_model(MODEL_B, body, stream), + ) + assert limited.status_code == 429 + assert recovered.status_code == 200 + assert "model-b-works" in recovered.text + assert seen == [ + (MODEL_A, "Bearer token-domestic-a"), + (MODEL_B, "Bearer token-domestic-a"), + ] + snapshot = manager.snapshot() + by_uid = {item["uid"]: item for item in snapshot["accounts"]} + assert _is_model_limited(by_uid["a"], MODEL_A) + assert not _is_model_limited(by_uid["a"], MODEL_B) + assert not by_uid["b"]["model_cooldowns"] + assert manager.store.active_account_id == account_ids["a"] + assert any( + fields.get("account_id") == account_ids["a"] + and fields.get("model") == MODEL_A + for _, fields in logs + ) + finally: + manager.close() + + +@pytest.mark.parametrize("candidate_cooling_model", [MODEL_A, MODEL_B]) +def test_candidate_is_excluded_only_when_requested_model_is_cooling( + tmp_path, monkeypatch, candidate_cooling_model +): + _, manager, _, clients, account_ids = _environment( + tmp_path, monkeypatch, [{"uid": "a"}, {"uid": "b"}, {"uid": "c"}] + ) + manager.set_settings(auto_switch_on_429=True) + manager.handle_rate_limit( + clients["b"], "300", model=candidate_cooling_model, allow_switch=False + ) + seen = [] + + if candidate_cooling_model == MODEL_A: + requested_model = MODEL_A + expected_retry_uid = "c" + else: + requested_model = MODEL_B + # A is cooled for MODEL_A only and must remain eligible for MODEL_B. + manager.handle_rate_limit(clients["a"], "300", model=MODEL_A, allow_switch=False) + manager.store.set_active(account_ids["b"]) + expected_retry_uid = "a" + + async def upstream(request): + token = request.headers["authorization"] + seen.append(token) + if token == f"Bearer token-domestic-{expected_retry_uid}": + return _success_response("eligible-model-candidate") + return _error_response(token.rsplit("-", 1)[-1], retry_after="30") + + _install_transport(monkeypatch, upstream) + try: + body = { + "model": requested_model, + "stream": False, + "messages": [{"role": "user", "content": "model candidate"}], + } + with TestClient(proxy.app, raise_server_exceptions=False) as app: + result = app.post("/v1/chat/completions", json=body) + starting_uid = "a" if candidate_cooling_model == MODEL_A else "b" + assert result.status_code == 200 + assert "eligible-model-candidate" in result.text + assert seen == [ + f"Bearer token-domestic-{starting_uid}", + f"Bearer token-domestic-{expected_retry_uid}", + ] + assert manager.store.active_account_id == account_ids[expected_retry_uid] + finally: + manager.close() + + +def test_model_none_does_not_record_cooldown_or_switch(tmp_path, monkeypatch): + _, manager, _, clients, account_ids = _environment( + tmp_path, monkeypatch, [{"uid": "a"}, {"uid": "b"}] + ) + manager.set_settings(auto_switch_on_429=True) + try: + result = manager.handle_rate_limit(clients["a"], "60", model=None) + assert result is None + assert manager.store.active_account_id == account_ids["a"] + assert all( + not account["model_cooldowns"] for account in manager.snapshot()["accounts"] + ) + finally: + manager.close() + + +def test_delayed_old_account_429_does_not_replace_manual_active_account( + tmp_path, monkeypatch +): + _, manager, _, _, account_ids = _environment( + tmp_path, + monkeypatch, + [ + {"uid": "a"}, + {"uid": "b"}, + {"uid": "manual-global", "profile": GLOBAL_PROFILE}, + ], + ) + manager.set_settings(auto_switch_on_429=True) + seen = [] + + async def upstream(request): + authorization = request.headers["authorization"] + seen.append(authorization) + if authorization == "Bearer token-domestic-a": + # The request has already pinned A. A human changes the active + # account while A's upstream response is being processed. + manager.activate(account_ids["manual-global"]) + return _error_response("a", retry_after="90") + if authorization == "Bearer token-domestic-b": + return _success_response("retried-on-b") + pytest.fail(f"ineligible account was used: {authorization}") + + _install_transport(monkeypatch, upstream) + try: + with TestClient(proxy.app, raise_server_exceptions=False) as app: + response = app.post( + "/v1/chat/completions", + json={"model": "rate-limit-test", "stream": False, + "messages": [{"role": "user", "content": "hello"}]}, + ) + assert response.status_code == 429 + assert "limited-a" in response.text + assert seen == ["Bearer token-domestic-a"] + assert manager.store.active_account_id == account_ids["manual-global"] + states = {account["uid"]: account for account in manager.snapshot()["accounts"]} + assert _is_model_limited(states["a"], MODEL_A) + assert not _is_model_limited(states["b"], MODEL_A) + assert not _is_model_limited(states["manual-global"], MODEL_A) + finally: + manager.close() + + +def test_billing_429_preserves_status_body_and_retry_after_without_replay( + tmp_path, monkeypatch +): + _, manager, _, _, account_ids = _environment( + tmp_path, monkeypatch, [{"uid": "a"}, {"uid": "b"}] + ) + manager.set_settings(auto_switch_on_429=True) + body = b'{"billing":"slow down","request":"single-checkin"}\n' + seen = [] + + async def upstream(request): + seen.append((request.method, request.url.host, request.headers["authorization"])) + return httpx.Response( + 429, + content=body, + headers={"content-type": "application/problem+json", "Retry-After": "45"}, + ) + + _install_transport(monkeypatch, upstream) + try: + with TestClient(proxy.app, raise_server_exceptions=False) as app: + response = app.post("/v1/checkin") + assert response.status_code == 429 + assert response.content == body + assert response.headers.get("retry-after") == "45" + assert response.headers.get("content-type") == "application/problem+json" + assert seen == [ + ("POST", "copilot.tencent.com", "Bearer token-domestic-a") + ] + assert manager.store.active_account_id == account_ids["a"] + limited = {a["uid"]: a for a in manager.snapshot()["accounts"]} + assert limited["a"]["model_cooldowns"] == [] + assert limited["b"]["model_cooldowns"] == [] + finally: + manager.close() + + +@pytest.mark.parametrize("status", [401, 503]) +def test_non_429_statuses_never_rotate_accounts(tmp_path, monkeypatch, status): + _, manager, _, _, account_ids = _environment( + tmp_path, monkeypatch, [{"uid": "a"}, {"uid": "b"}] + ) + manager.set_settings(auto_switch_on_429=True) + seen = [] + + async def upstream(request): + seen.append(request.headers["authorization"]) + return httpx.Response(status, json={"msg": f"status-{status}"}) + + _install_transport(monkeypatch, upstream) + try: + with TestClient(proxy.app, raise_server_exceptions=False) as app: + response = app.post( + "/v1/chat/completions", + json={"model": "rate-limit-test", "stream": False, + "messages": [{"role": "user", "content": "hello"}]}, + ) + assert response.status_code == status + assert seen == ["Bearer token-domestic-a"] + assert manager.store.active_account_id == account_ids["a"] + assert all(not account["model_cooldowns"] for account in manager.snapshot()["accounts"]) + finally: + manager.close() + + +def test_started_sse_is_never_replayed_after_upstream_stream_failure( + tmp_path, monkeypatch +): + _, manager, _, _, account_ids = _environment( + tmp_path, monkeypatch, [{"uid": "a"}, {"uid": "b"}] + ) + manager.set_settings(auto_switch_on_429=True) + seen = [] + + class BrokenAfterFirstChunk(httpx.AsyncByteStream): + async def __aiter__(self): + chunk = { + "id": "partial-stream", + "object": "chat.completion.chunk", + "created": 1, + "model": "rate-limit-test", + "choices": [ + {"index": 0, "delta": {"content": "stream-started"}, + "finish_reason": None} + ], + } + yield f"data: {json.dumps(chunk)}\n\n".encode() + raise httpx.ReadError("injected failure after first SSE chunk") + + async def aclose(self): + return None + + async def upstream(request): + seen.append(request.headers["authorization"]) + return httpx.Response( + 200, + stream=BrokenAfterFirstChunk(), + headers={"content-type": "text/event-stream"}, + ) + + _install_transport(monkeypatch, upstream) + try: + with TestClient(proxy.app, raise_server_exceptions=False) as app: + response = app.post( + "/v1/chat/completions", + json={"model": "rate-limit-test", "stream": True, + "messages": [{"role": "user", "content": "hello"}]}, + ) + assert response.status_code == 200 + assert "stream-started" in response.text + assert seen == ["Bearer token-domestic-a"] + assert manager.store.active_account_id == account_ids["a"] + assert all( + not account["model_cooldowns"] + for account in manager.snapshot()["accounts"] + ) + finally: + manager.close() + + +def test_concurrent_requests_keep_their_body_and_account_snapshot_isolated( + tmp_path, monkeypatch +): + _, manager, _, _, account_ids = _environment( + tmp_path, monkeypatch, [{"uid": "a"}, {"uid": "b"}, {"uid": "c"}] + ) + manager.set_settings(auto_switch_on_429=True) + first_a_started = threading.Event() + release_first_a = threading.Event() + seen = [] + + async def upstream(request): + authorization = request.headers["authorization"] + payload = json.loads(request.content) + prompt = payload["messages"][-1]["content"] + seen.append((prompt, authorization)) + if prompt == "request-one" and authorization == "Bearer token-domestic-a": + first_a_started.set() + await asyncio.to_thread(release_first_a.wait, 5) + return _error_response("a", retry_after="90") + return _success_response(f"answer-{prompt}") + + _install_transport(monkeypatch, upstream) + body_one = { + "model": "rate-limit-test", "stream": False, + "messages": [{"role": "user", "content": "request-one"}], + } + body_two = { + "model": "rate-limit-test", "stream": False, + "messages": [{"role": "user", "content": "request-two"}], + } + try: + with TestClient(proxy.app, raise_server_exceptions=False) as app: + with ThreadPoolExecutor(max_workers=1) as pool: + first_future = pool.submit(app.post, "/v1/chat/completions", json=body_one) + try: + assert first_a_started.wait(5) + manager.activate(account_ids["b"]) + second_response = app.post("/v1/chat/completions", json=body_two) + finally: + release_first_a.set() + first_response = first_future.result(timeout=10) + + assert second_response.status_code == 200 + assert "answer-request-two" in second_response.text + assert first_response.status_code == 429 + assert "limited-a" in first_response.text + first_tokens = [token for prompt, token in seen if prompt == "request-one"] + second_tokens = [token for prompt, token in seen if prompt == "request-two"] + assert first_tokens[0] == "Bearer token-domestic-a" + assert first_tokens == ["Bearer token-domestic-a"] + assert second_tokens == ["Bearer token-domestic-b"] + assert manager.store.active_account_id == account_ids["b"] + finally: + release_first_a.set() + manager.close() + + +def test_concurrent_model_429s_keep_cooldowns_and_logs_isolated( + tmp_path, monkeypatch +): + _, manager, state, _, account_ids = _environment( + tmp_path, monkeypatch, [{"uid": "a"}, {"uid": "b"}] + ) + manager.set_settings(auto_switch_on_429=False) + entered = {MODEL_A: threading.Event(), MODEL_B: threading.Event()} + both_entered = threading.Event() + release = threading.Event() + seen = [] + logs = [] + seen_lock = threading.Lock() + logs_lock = threading.Lock() + + def record_log(event, **fields): + if event == "account_rate_limited": + with logs_lock: + logs.append(fields) + + state.write_log = record_log + + async def upstream(request): + model = json.loads(request.content)["model"] + with seen_lock: + seen.append((model, request.headers["authorization"])) + entered[model].set() + if all(event.is_set() for event in entered.values()): + both_entered.set() + await asyncio.to_thread(release.wait, 5) + retry_after = "19" if model == MODEL_A else "73" + return _error_response("a", retry_after=retry_after) + + _install_transport(monkeypatch, upstream) + body_a = { + "model": MODEL_A, + "stream": False, + "messages": [{"role": "user", "content": "parallel model A"}], + } + body_b = { + "model": MODEL_B, + "stream": False, + "messages": [{"role": "user", "content": "parallel model B"}], + } + try: + with TestClient(proxy.app, raise_server_exceptions=False) as app: + with ThreadPoolExecutor(max_workers=2) as pool: + future_a = pool.submit(app.post, "/v1/chat/completions", json=body_a) + future_b = pool.submit(app.post, "/v1/chat/completions", json=body_b) + assert both_entered.wait(5), "both model requests must reach upstream before release" + release.set() + response_a = future_a.result(timeout=10) + response_b = future_b.result(timeout=10) + + assert response_a.status_code == response_b.status_code == 429 + assert response_a.headers.get("retry-after") == "19" + assert response_b.headers.get("retry-after") == "73" + assert len(seen) == 2 + assert set(seen) == { + (MODEL_A, "Bearer token-domestic-a"), + (MODEL_B, "Bearer token-domestic-a"), + } + assert manager.store.active_account_id == account_ids["a"] + + snapshot = manager.snapshot() + by_uid = {account["uid"]: account for account in snapshot["accounts"]} + cooldowns = { + entry["model"]: entry + for entry in by_uid["a"]["model_cooldowns"] + } + assert set(cooldowns) == {MODEL_A, MODEL_B} + assert cooldowns[MODEL_A]["rate_limited"] is True + assert cooldowns[MODEL_B]["rate_limited"] is True + assert cooldowns[MODEL_A]["cooldown_source"] == "retry_after" + assert cooldowns[MODEL_B]["cooldown_source"] == "retry_after" + # The requests are released together; the deadline gap should reflect + # the 54-second difference between their independent Retry-After values. + assert 40 <= cooldowns[MODEL_B]["cooldown_until"] - cooldowns[MODEL_A]["cooldown_until"] <= 68 + assert not by_uid["b"]["model_cooldowns"] + assert len(logs) == 2 + assert { + (entry.get("account_id"), entry.get("model")) for entry in logs + } == { + (account_ids["a"], MODEL_A), + (account_ids["a"], MODEL_B), + } + finally: + release.set() + manager.close()