diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..73de01e --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,38 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +permissions: + contents: read + +jobs: + test: + strategy: + fail-fast: false + matrix: + include: + - os: ubuntu-latest + python: "3.11" + - os: macos-latest + python: "3.13" + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v6 + with: + python-version: ${{ matrix.python }} + enable-cache: true + - run: uv sync --locked --all-groups + - run: uv run python -m compileall -q agentbar tests + - name: Core tests (Linux) + if: runner.os == 'Linux' + run: >- + uv run pytest + --ignore=tests/test_menubar.py + --ignore=tests/test_provider_window.py + - name: Full tests (macOS) + if: runner.os == 'macOS' + run: uv run pytest diff --git a/README.md b/README.md index dc74b77..216dbbf 100644 --- a/README.md +++ b/README.md @@ -24,19 +24,20 @@ macOS 状态栏（Menu Bar）AI Agent 调度器 —— 让 Claude Code、Codex | ---- | ---- | | 常驻 Menu Bar / 开机自启 | 原生 AppKit（NSStatusItem/NSMenu）+ `scripts/install-launch-agent.sh` | | 可操作的 Menu Bar | 每一行都可点（概览/任务/额度行点击即打开面板）；菜单只在展开时刷新（menuWillOpen）；所有动作毫秒级返回，绝不阻塞主线程；实时菜单状态导出 `menu-debug.json` 可核查 | -| 原生任务面板窗口 | 菜单点击直接弹出 AppKit 窗口（accessory 进程激活自身窗口不受 macOS 26 协作激活限制）：添加任务（Prompt/工具/模型/强度/权限/目录选择器）、队列排优先级（⇧置顶/↑/↓）、取消/重试/暂停派发，进程内直连调度器不经浏览器；web 面板降为次要入口（`🌐 在浏览器中打开面板` / 手机远程） | +| 原生任务面板窗口 | 菜单点击直接弹出 AppKit 窗口：添加任务（Prompt/工具/模型/强度/权限/目录）、队列排序、取消/重试/暂停；底部「额度设置」可配置 Claude / Codex / MyToken / Tokenverse | | 添加任务（Prompt+工具+目录） | 菜单栏快捷添加 / Web 面板 / `agentbar add` CLI | | 支持 Claude Code、Codex，可扩展 | Adapter 插件制，新 CLI ≈ 60 行代码 | | 串行 / 有限并行 | `max_parallel`（默认 1 串行）+ `per_tool_limit` | | 完整生命周期 | queued / running / succeeded / failed / **waiting_quota** / paused / cancelled | | 额度耗尽不判失败 | 识别限流报错 → `waiting_quota`，解析恢复时间或指数退避，**自动 resume 原会话续跑** | -| 额度状态可见（不伪造） | 使用当前 CLI 登录态读取的 usage 数据 + 调度观测 + 可选 ccusage；窗口、重置时间和来源都明确标注 | +| 额度状态可见（不伪造） | 额度来源显式 opt-in；手动选择 Claude 模型家族或 Codex `metered_feature`，输入 OAuth Access Token 后才请求 | +| 可控制刷新 | 默认只在启动、保存配置或手动点击时请求一次；任务结束不再强制刷新全部来源，重复点击会合并 | | 菜单栏双环额度图标 | 外圈 Claude、内圈 Codex 用量一眼可见（同 aiusagebar 的外围圈样式），不用点开菜单；环心徽标显示调度状态（实心点=运行中、双竖条=已暂停），菜单栏只有这一个图标；模板图自适配深浅色；无可信数据只画轨道不编造 | | 键盘快捷键 | 面板/对话窗口前台时：`⌃W`/`⌘W` 关闭当前窗口，`⌃Q`/`⌘Q` 退出（走完整清理：停隧道/调度器/服务器） | | 模型 / 强度选择 | 任务级保存模型与强度；Claude 用 `--model` / `--effort`，Codex 用 `--model` / `model_reasoning_effort` 配置覆盖 | | 实时查看运行/队列/日志/历史 | Web 面板 2s 自刷新 + 日志实时 tail | | 查看本机其他 CLI | 只读发现正在运行的 Claude/Codex 进程；不读取 Prompt/完整命令，也不会终止外部进程 | -| 重启恢复 | 每次状态变更原子落盘 `state.json`；崩溃时 RUNNING 任务自动重新入队（续会话） | +| 重启恢复 | 每次状态变更原子落盘；崩溃遗留的 RUNNING 任务安全暂停，人工确认旧进程后再恢复，避免重复 agent 同时改文件 | | 默认安全 | 三档权限（默认不开高权限）、仅 127.0.0.1 + token、无 shell 拼接、进程组隔离、超时兜底 | ## 安装 & 运行 @@ -44,16 +45,38 @@ macOS 状态栏（Menu Bar）AI Agent 调度器 —— 让 Claude Code、Codex ```bash # 依赖: macOS + uv (https://docs.astral.sh/uv/) cd agentbar -uv sync # installs the macOS Security bridge for Claude quota access +uv sync # 安装 macOS 原生 UI 与运行依赖 uv run agentbar run # 菜单栏模式（推荐） uv run agentbar run --headless # 无 GUI（服务器/调试） bash scripts/install-launch-agent.sh # 开机自启（登录时拉起） -bash scripts/agentbar-restart.sh # 代码更新后重启服务（干净杀进程组，避免孤儿占端口） +bash scripts/agentbar-restart.sh # 只重启默认状态目录对应的实例 bash scripts/uninstall-launch-agent.sh # 取消自启 ``` +这三个脚本服务于**源码开发安装**：生成的 LaunchAgent 会记录当前仓库和 `uv` +的绝对路径，并以 `uv run --locked` 启动；移动/删除仓库或让 `uv.lock` 与项目 +元数据失配后将无法启动。安装采用候选 plist 校验后原子替换； +若新进程未通过本机 `/api/ping` 健康检查，会恢复旧 plist。重启只向 +`~/.agentbar/runtime.json` 所属 PID 发送 `SIGTERM`，不会用进程名批量杀掉其他 +`--state-dir` 实例，也不会用 `SIGKILL` 留下正在工作的 Claude/Codex 孤儿进程。 +卸载只取消自启并保留 `~/.agentbar` 中的任务、配置和私有日志。 + +需要脱离源码目录安装时，请构建并拖拽 DMG（依赖 Homebrew framework Python +3.13；Intel Homebrew 与 Apple Silicon 路径都会自动识别）： + +```bash +bash scripts/build-dmg.sh +# 输出 dist/AgentBar.app、dist/AgentBar-<版本>.dmg 和 SHA-256 +``` + +构建使用 `uv.lock` 中的运行依赖和一次性 py2app 环境，随后检查签名、冻结后的 +AppKit/WebKit/登录/二维码/Markdown 模块、内置 Web 资源及 DMG 完整性。版本号必须 +在 `pyproject.toml` 与 `agentbar/__init__.py` 中一致，否则构建会直接失败。DMG +架构跟随构建机（Apple Silicon 产出 arm64，Intel 产出 x86_64），当前不生成 +Universal 2 包。 + 启动后点菜单栏 🤖 →「打开任务面板」，或： ```bash @@ -72,9 +95,9 @@ uv run agentbar pause / resume / cancel / log add ──▶ queued ──▶ running ──▶ succeeded ▲ │ ▲ │ │ │ │ └────────┼─┼──▶ failed（真实错误才算失败，可手动重试） - 重启恢复│ ▼ │ └──▶ waiting_quota ──（到点自动回 queued，resume 原会话） - (RUNNING│ paused ◀──────┘ │ - →queued) │ ▼ + 人工恢复│ ▼ │ └──▶ waiting_quota ──（到点自动回 queued，resume 原会话） + (崩溃遗留│ paused ◀──────┘ │ + RUNNING) │ ▼ └───── cancelled ◀──────────────┘（各状态均可取消） ``` @@ -92,25 +115,28 @@ uv run agentbar pause / resume / cancel / log | ⚠️ full | `--dangerously-skip-permissions` | `--dangerously-bypass-approvals-and-sandbox` | - **full 档默认禁用**：需在 `~/.agentbar/config.json` 设 `allow_full_profile: true` 并重启，UI 中也有显式警告。 -- API 仅绑定 `127.0.0.1`，所有写操作要求 token（`config.json`，0600），校验 Host 头防 DNS rebinding——防止恶意网页通过浏览器向本机调度器投毒任务。 +- API 所有写操作都要求 AgentBar token，并校验 Host 头防 DNS rebinding；手动额度凭据只接受本机回环地址提交，避免经局域网明文传输。 +- OAuth Token / Account ID 是只写字段：Web、状态快照、菜单调试文件均只返回“是否已配置”；落盘的 `config.json` 权限固定为 `0600`。 - 子进程以 argv 数组直接 exec，无 shell 拼接；prompt 走 stdin，杜绝 flag 注入。 - 每个任务独立进程组，取消/超时（默认 2h）时整组终止，不留孤儿进程。 ## 额度状态的数据来源（诚实降级） -订阅版 CLI 没有承诺稳定的公开额度查询 API。AgentBar 因此按以下顺序读取，任何一种拿不到都会如实降级： +订阅版 CLI 没有承诺稳定的公开额度查询 API。AgentBar 因此只请求用户显式配置的来源，任何一种拿不到都如实降级： -1. **usage API**（默认、120 秒刷新）：使用本机已有的 Claude OAuth / Codex 登录态读取其当前 usage 响应，显示窗口用量与重置时间。Claude Keychain 默认静默读取，绝不会在后台弹窗；首次在面板主动授权后，AgentBar 只会把短期 access token 和授权状态保存到仅当前用户可读的 `0600` 缓存。refresh token 始终留在 Keychain，且只由官方 Claude Code 续期；AgentBar 会在后台无弹窗地同步新 access token，不会与 Claude Code 抢占或轮换 refresh token。 +1. **usage API**（显式配置）：Claude / Codex 默认关闭。在「额度设置」中启用来源、选择额度模型，并输入对应的 **OAuth Access Token** 后才会请求。Codex 可选填 Account ID。普通 Anthropic/OpenAI API key 不等于订阅 usage 凭据，接口拒绝时会如实显示错误。 2. **observed**：调度器自身观测的最近成功执行、真实限流和恢复时间；它会优先标记已确认的限流。 3. **ccusage**（可选增强）：`npm i -g ccusage` 后补充 Claude 本地 5h 成本。 4. 无任何可用数据时显示「未知」，并显示失败原因。**不会估算或编造百分比。** -usage 响应不是稳定的公开契约，接口结构变化时会显示解析错误而非虚构数值。可在 `config.json` 里调整轮询频率： +usage 响应不是稳定的公开契约，接口结构变化时会显示解析错误而非虚构数值。默认不周期轮询；如需显式开启： ```jsonc -{ "usage_refresh_seconds": 120 } // 最小 30 秒 +{ "usage_auto_refresh": true, "usage_refresh_seconds": 120 } ``` +手动模式下，任务成功/限额观测只更新本地状态，不会触发上游网络刷新。 + ## 模型与强度 添加任务时，模型和强度均为任务级字段，写入 `state.json` 并在重试/重启恢复时保留。 @@ -136,7 +162,7 @@ agentbar.log # 调度器日志 ``` - 调度器退出（含 SIGTERM）：在跑的 CLI 进程被整组终止，任务放回队列并标记续会话。 -- 崩溃/断电：下次启动时 `state.json` 里的 RUNNING 任务自动重新入队。 +- 崩溃/断电：旧 CLI 进程可能仍在运行，因此下次启动会把 RUNNING 任务转为 `paused`；确认现场后手动恢复，避免重复执行有副作用的任务。 - `state.json` 损坏：自动备份为 `state.json.corrupt-*` 并从空状态启动，不会起不来。 - launchd 场景 PATH 被裁剪：自动经登录 shell（`zsh -lc`）解析 claude/codex 真实路径，nvm 安装也能找到；亦可在 `config.json` 的 `tool_paths` 手动指定。 @@ -152,8 +178,23 @@ agentbar.log # 调度器日志 "task_timeout_seconds": 7200, "backoff_minutes": [5, 15, 30, 60], "usage_refresh_seconds": 120, + "usage_auto_refresh": false, // false=默认手动；true=按上方间隔轮询 "tool_paths": {}, // {"claude": "/abs/path"} 手动覆盖 "title_provider": "claude", // 状态栏标题显示哪个 provider 的用量：claude/codex/mytoken/tokenverse + "quota_sources": { // 订阅额度：默认关闭，必须手动输入凭据 + "claude": { + "enabled": false, + "model": "sonnet", // opus / sonnet；留空=账户通用窗口 + "access_token": "", // OAuth Access Token，非普通 API key + "account_id": "" + }, + "codex": { + "enabled": false, + "model": "", // metered_feature；留空=账户总额度 + "access_token": "", // ChatGPT OAuth Access Token + "account_id": "" // 可选；JWT 不含账号时填写 + } + }, "providers": { // 快手内部额度 provider（默认关闭，可在面板一键导入 Cookie） "mytoken": { "enabled": false, @@ -180,7 +221,7 @@ agentbar.log # 调度器日志 | **MyToken** | `mytoken.corp.kuaishou.com` — `/api/auth/sso/user` → `/api/v1/billing/account` | corp SSO cookie + `kwaipilot-username` 头（自动带） | | **Tokenverse** | `tokenverse.corp.kuaishou.com` — `/api/coding-plan/status` + `/api/coding-plan/usage/summary` | corp SSO cookie | -启用方式：打开 AgentBar 面板 → **内部额度配置**，可直接点“从浏览器导入 Cookie”读取本机 Chrome / Edge / Brave 登录态，也可手动粘贴浏览器 DevTools 里的整行 `Cookie`。保存后会立即重建 provider 并刷新额度，不需要手写 JSON 或重启。每个 provider 可独立配置展示单位（`credits`/`percent`/`token`）与刷新间隔。诚实原则不变：cookie 缺失或接口失败时如实显示错误，绝不编造额度。接口契约参考 [ylab/aiusagebar](https://git.corp.kuaishou.com/ylab/aiusagebar)。 +启用方式：点击菜单栏 AgentBar 图标 → **额度设置…**（或原生任务面板底部「额度设置」）。MyToken / Tokenverse 即使尚未配置也会以“未配置”状态出现在菜单中。点「浏览器登录」后 AgentBar 启动临时 Chrome 会话；完成企业 SSO 后捕获并校验 Cookie，成功后保存、启用并刷新一次。「读取已有登录」与手动粘贴整行 `Cookie` 仍作为兜底。凭据缺失或接口失败时如实显示错误，不编造额度。 ## 扩展新的 AI CLI @@ -205,7 +246,8 @@ class GeminiAdapter(Adapter): ## 测试 ```bash -uv run pytest # 59 个测试：生命周期/额度退避/恢复/取消/API 安全/适配器/模型/真实额度解析/进程观测/Menu Bar 子菜单 +uv run pytest # 生命周期/刷新竞态/凭据脱敏/模型额度/API 安全/原生 UI/Menu Bar +bash -n scripts/*.sh # 安装、重启、卸载与 DMG 脚本语法 ``` 测试用 `AGENTBAR_ENABLE_FAKE=1` 注册的 fake CLI 模拟成功/失败/限流/慢任务，不消耗真实额度。 @@ -214,17 +256,17 @@ uv run pytest # 59 个测试：生命周期/额度退避/恢复/取消/AP 架构已按 API-first 设计（菜单栏和 Web 面板都是同一 HTTP API 的客户端），远程控制是加通道而非改架构： -1. **内网/Tailscale**（推荐第一步）：Tailscale 组网后手机浏览器直接访问 `http://:8737/?token=…`，Web 面板本身就是响应式的。需把监听地址改为可配置并强化 token 策略。 +1. **内网/Tailscale**：默认只监听 `127.0.0.1`。明确开启 `lan_access` 后，从菜单生成一次带 `#token=…` fragment 的二维码；凭据不会进入 HTTP 请求日志。局域网仍是明文 HTTP，不应提交额度凭据。 2. **IM Bot**：Telegram/Slack/Kim bot 进程调用同一 API（add/status/log/pause），推送任务完成/额度恢复通知——适合"下班路上派活"。 3. **PWA + 推送**：面板加 manifest + Web Push，任务完成/失败/等额度主动通知手机。 4. **中继模式**：Mac 出网受限时，经云端轻量 relay（WebSocket 反向连接）转发 API，手机端连 relay。 -## 已知限制（v0.3） +## 已知限制（当前版本） - usage 响应不是公开契约，接口结构变化时显示解析错误而非虚构数值。 - 修改 `config.json` 需重启生效（无热加载）。 - 任务级依赖（A 完成才跑 B）未实现，当前是 FIFO + 并发上限。 -- Menu Bar 标题为文本符号（◇◆◐Ⅱ + 用量百分比），自动适配深浅色菜单栏。 +- Menu Bar 使用原生矢量双环图标，自动适配深浅色；没有可信额度数据时只画轨道。 ## Menu Bar 架构备注（为什么不用 rumps） @@ -233,7 +275,7 @@ v0.1/v0.2 基于 rumps 时出现两类线上事故，v0.3 改为直接使用 App 1. **主线程阻塞**：菜单回调里同步等待 `/usr/bin/open`（最长 8s）或 webbrowser （macOS 上走 osascript/Apple Events，可能卡在 TCC 授权）→ 整个 App 卡死。 现在 GUI 一律 `open_url_async`（Popen fire-and-forget，实测 ~3ms 返回）， - Keychain 授权等慢操作全部丢后台线程。 + 浏览器登录等慢操作全部丢后台线程。 2. **定时重建打开中的菜单**：rumps.Timer 每 2s clear+rebuild 菜单导致点击落空。 现在菜单内容只在 `menuWillOpen`（AppKit 正统时机）重建，NSTimer 只改标题文本。 diff --git a/agentbar/__init__.py b/agentbar/__init__.py index 106642a..18563f9 100644 --- a/agentbar/__init__.py +++ b/agentbar/__init__.py @@ -1,3 +1,3 @@ """AgentBar — macOS menu bar scheduler for AI CLI agents.""" -__version__ = "0.10.2" +__version__ = "0.10.4" diff --git a/agentbar/browser.py b/agentbar/browser.py index 155d9f6..4db4731 100644 --- a/agentbar/browser.py +++ b/agentbar/browser.py @@ -28,7 +28,9 @@ def _sanitized(url: str) -> str: - return url.split("?")[0] # 日志里不落 token + # Legacy links used ?token= and current links bootstrap from #token=. + # Strip both components before logging so neither form can leak a secret. + return url.split("?", 1)[0].split("#", 1)[0] def open_url_async(url: str, on_result: OnResult | None = None) -> bool: diff --git a/agentbar/chrome_login.py b/agentbar/chrome_login.py new file mode 100644 index 0000000..d9b5df9 --- /dev/null +++ b/agentbar/chrome_login.py @@ -0,0 +1,294 @@ +"""Interactive corp SSO login through an isolated Chrome CDP session. + +This mirrors the reliable part of AIUsageBar's setup flow: launch a real Chrome +window with a temporary profile, let Kit/0Pass complete SSO, read the target +host's cookies through Chrome DevTools Protocol, then validate them against the +real provider API. No cookie is persisted until validation succeeds. +""" + +from __future__ import annotations + +import json +import shutil +import socket +import subprocess +import tempfile +import threading +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import Callable + +import websocket + +from .browser_cookies import ImportedCookie +from .config import PROVIDER_HOSTS, PROVIDER_LOGIN_URLS +from .usage import MyTokenUsageFetcher, TokenverseUsageFetcher + + +class ChromeLoginError(RuntimeError): + pass + + +_READY_COOKIE_NAMES = { + "mytoken": {"accessproxy_session", "JSESSIONID", "KP_SSO_SID", "ktrace-context"}, + "tokenverse": {"accessproxy_session"}, +} + + +def _chrome_binary() -> Path | None: + candidates = ( + Path("/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"), + Path.home() / "Applications/Google Chrome.app/Contents/MacOS/Google Chrome", + Path("/Applications/Google Chrome Beta.app/Contents/MacOS/Google Chrome Beta"), + ) + return next((path for path in candidates if path.is_file()), None) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _domain_matches(cookie_host: str, target_host: str) -> bool: + cookie_host = (cookie_host or "").lstrip(".").lower() + target_host = target_host.lower() + return target_host == cookie_host or target_host.endswith("." + cookie_host) + + +def _cookie_header(raw_cookies: list[dict], target_host: str) -> str: + """Build a deterministic Cookie header for target_host without logging values.""" + now = time.time() + candidates = [] + for entry in raw_cookies: + name = str(entry.get("name") or "").strip() + value = str(entry.get("value") or "") + domain = str(entry.get("domain") or "") + expires = entry.get("expires") + if not name or not value or not _domain_matches(domain, target_host): + continue + if isinstance(expires, (int, float)) and expires > 0 and expires <= now: + continue + path = str(entry.get("path") or "/") + exact = int(domain.lstrip(".").lower() == target_host.lower()) + candidates.append((exact, len(domain.lstrip(".")), len(path), name, value)) + candidates.sort(reverse=True) + seen = set() + parts = [] + for _, _, _, name, value in candidates: + if name in seen: + continue + seen.add(name) + parts.append(f"{name}={value}") + return "; ".join(parts) + + +def _validate_cookie(provider: str, header: str) -> bool: + if provider == "mytoken": + snap = MyTokenUsageFetcher(cookie=header).fetch() + elif provider == "tokenverse": + snap = TokenverseUsageFetcher(cookie=header).fetch() + else: + return False + return bool(snap and not snap.error and snap.windows) + + +class ChromeCDPLogin: + def __init__(self, provider: str): + if provider not in PROVIDER_HOSTS: + raise ChromeLoginError(f"未知 provider: {provider}") + self.provider = provider + self.host = PROVIDER_HOSTS[provider] + self.login_url = PROVIDER_LOGIN_URLS[provider] + self.port = _free_port() + self._cancel = threading.Event() + self._process: subprocess.Popen | None = None + self._profile_dir: Path | None = None + + def cancel(self) -> None: + self._cancel.set() + # UI 主线程只发终止信号，不在这里 wait/rmtree；run() 的后台线程 + # 会在 finally 里完成回收，避免“取消登录”按钮卡住菜单栏。 + process = self._process + if process and process.poll() is None: + try: + process.terminate() + except OSError: + pass + + def close(self) -> None: + process = self._process + self._process = None + profile_dir = self._profile_dir + self._profile_dir = None + try: + if process and process.poll() is None: + try: + process.terminate() + except OSError: + pass + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + try: + process.kill() + except OSError: + pass + try: + process.wait(timeout=3) + except subprocess.TimeoutExpired: + # Keep shutdown bounded even for an injected/broken + # process handle. The profile cleanup below is still an + # invariant and must not be skipped by this condition. + pass + finally: + if profile_dir: + shutil.rmtree(profile_dir, ignore_errors=True) + + def _chrome_args(self, chrome: Path) -> list[str]: + return [ + str(chrome), + f"--remote-debugging-port={self.port}", + "--remote-debugging-address=127.0.0.1", + f"--remote-allow-origins=http://127.0.0.1:{self.port}", + f"--user-data-dir={self._profile_dir}", + "--no-first-run", + "--no-default-browser-check", + "--new-window", + self.login_url, + ] + + def run( + self, + on_status: Callable[[str], None] | None = None, + timeout: float = 300, + ) -> ImportedCookie: + chrome = _chrome_binary() + if chrome is None: + raise ChromeLoginError("未检测到 Google Chrome / Chrome Beta") + status = on_status or (lambda _message: None) + self._profile_dir = Path(tempfile.mkdtemp(prefix="agentbar-chrome-login-")) + args = self._chrome_args(chrome) + status("正在启动独立 Chrome 登录窗口…") + try: + self._process = subprocess.Popen( + args, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + self._wait_debugger(timeout=20) + status("请在弹出的 Chrome 中完成企业 SSO 登录…") + deadline = time.time() + timeout + last_count = -1 + while time.time() < deadline: + self._check_cancelled() + raw = self._fetch_all_cookies() + header = _cookie_header(raw, self.host) + fields = { + part.split("=", 1)[0].strip() + for part in header.split(";") + if "=" in part + } + count = len(fields) + if count != last_count: + last_count = count + status( + f"已捕获 {count} 个 Cookie，等待 SSO 完成…" + if count + else "等待企业登录 Cookie…" + ) + ready = _READY_COOKIE_NAMES[self.provider].issubset(fields) + if ready: + status("已捕获完整 Cookie，正在用真实额度接口验证…") + if _validate_cookie(self.provider, header): + return ImportedCookie( + header=header, + source="Chrome 企业 SSO 登录", + count=count, + ) + status("Cookie 已捕获，但额度接口尚未通过，继续等待登录完成…") + self._cancel.wait(1.5) + raise ChromeLoginError("浏览器登录超时（5 分钟）") + except OSError as exc: + raise ChromeLoginError(f"启动 Chrome 失败：{exc}") from exc + finally: + self.close() + + def _check_cancelled(self) -> None: + if self._cancel.is_set(): + raise ChromeLoginError("已取消浏览器登录") + if self._process and self._process.poll() is not None: + raise ChromeLoginError("Chrome 登录窗口已关闭") + + def _json_get(self, path: str, timeout: float = 5) -> object: + request = urllib.request.Request( + f"http://127.0.0.1:{self.port}{path}", + headers={"Accept": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.loads(response.read().decode("utf-8")) + + def _wait_debugger(self, timeout: float) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + self._check_cancelled() + try: + if isinstance(self._json_get("/json/version", timeout=1), dict): + return + except (OSError, urllib.error.URLError, json.JSONDecodeError): + pass + self._cancel.wait(0.5) + raise ChromeLoginError("Chrome 调试端口启动超时") + + def _fetch_all_cookies(self) -> list[dict]: + try: + targets = self._json_get("/json/list") + except (OSError, urllib.error.URLError, json.JSONDecodeError): + return [] + if not isinstance(targets, list): + return [] + page = next( + ( + target + for target in targets + if target.get("type") == "page" and target.get("webSocketDebuggerUrl") + and self.host in str(target.get("url") or "") + ), + None, + ) + if page is None: + page = next( + ( + target + for target in targets + if target.get("type") == "page" and target.get("webSocketDebuggerUrl") + ), + None, + ) + if not page: + return [] + try: + connection = websocket.create_connection( + page["webSocketDebuggerUrl"], + timeout=6, + origin=f"http://127.0.0.1:{self.port}", + ) + except (OSError, ValueError, websocket.WebSocketException): + return [] + try: + connection.send(json.dumps({"id": 1, "method": "Network.enable"})) + connection.send(json.dumps({"id": 2, "method": "Network.getAllCookies"})) + deadline = time.time() + 6 + while time.time() < deadline: + message = json.loads(connection.recv()) + if message.get("id") == 2: + cookies = (message.get("result") or {}).get("cookies") or [] + return cookies if isinstance(cookies, list) else [] + except (OSError, ValueError, websocket.WebSocketException): + return [] + finally: + connection.close() + return [] diff --git a/agentbar/cli.py b/agentbar/cli.py index ca5f1dd..23fc1ce 100644 --- a/agentbar/cli.py +++ b/agentbar/cli.py @@ -7,12 +7,13 @@ import logging import os import signal +import stat import sys import threading -import time import urllib.error import urllib.request from pathlib import Path +from urllib.parse import urlencode from . import __version__ from .browser import open_panel_url @@ -24,11 +25,21 @@ def _setup_logging(settings: Settings) -> None: + log_path = settings.state_dir / "agentbar.log" + fd = os.open( + log_path, + os.O_WRONLY | os.O_CREAT | os.O_APPEND, + stat.S_IRUSR | stat.S_IWUSR, + ) + try: + os.fchmod(fd, stat.S_IRUSR | stat.S_IWUSR) + finally: + os.close(fd) logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s", handlers=[ - logging.FileHandler(settings.state_dir / "agentbar.log", encoding="utf-8"), + logging.FileHandler(log_path, encoding="utf-8"), logging.StreamHandler(), ], ) @@ -77,60 +88,132 @@ def _instance_alive(settings: Settings) -> bool: # ---------------- subcommands ---------------- -def cmd_run(args, settings: Settings) -> int: +def _cleanup_runtime(core, server, store, app=None) -> None: + """Run every shutdown stage even if an earlier one fails.""" + first_error: Exception | None = None + actions = [] + if app is not None: + begin_shutdown = getattr(app, "begin_shutdown", None) + if callable(begin_shutdown): + actions.append(("menu actions", begin_shutdown)) + actions.append(("HTTP server", server.stop)) + if app is not None: + actions.extend([ + ("public tunnel", app.tunnel.stop), + ("provider login", app.close_active_provider_logins), + ]) + actions.extend([ + ("scheduler", core.shutdown), + ("runtime file", store.clear_runtime), + ]) + for label, action in actions: + try: + action() + except Exception as exc: + if first_error is None: + first_error = exc + logging.getLogger("agentbar.cli").exception("failed to stop %s", label) + if first_error is not None: + raise first_error + + +def cmd_run(args, settings: Settings, *, instance_checked: bool = False) -> int: # 先应用端口覆盖再查重：实例身份 = state dir + 端口， # 否则 --port/--state-dir 启动的独立实例会被别的实例误判为"已在运行" if args.port is not None: settings.port = args.port - if _instance_alive(settings): + if not instance_checked and _instance_alive(settings): print("AgentBar 已在运行（用 `agentbar open` 打开面板，或先退出旧实例）") return 2 _setup_logging(settings) store = StateStore(settings.state_dir) core = Scheduler(settings, store) - server = ApiServer(core, settings) - server.start() - core.start() - store.write_runtime(server.port) - - print(f"AgentBar v{__version__} 已启动") - print(f" 状态目录: {settings.state_dir}") - print(f" 任务面板: {server.url(with_token=True)}") - print(f" 模式: {'headless' if args.headless else 'menu bar'}") - + server = None stop_evt = threading.Event() def _graceful(*_): stop_evt.set() - signal.signal(signal.SIGTERM, _graceful) - signal.signal(signal.SIGINT, _graceful) + try: + # Install termination admission before starting any thread/listener. + # Otherwise SIGTERM in the startup window takes the default hard-exit + # path and can strand a freshly dispatched CLI child. + signal.signal(signal.SIGTERM, _graceful) + signal.signal(signal.SIGINT, _graceful) + server = ApiServer(core, settings) + server.start() + core.start() + store.write_runtime(server.port) + except Exception: + try: + if server is not None: + _cleanup_runtime(core, server, store) + else: + # ApiServer construction itself may fail after a stale runtime + # file was left behind. Scheduler has not started, but shutdown + # is safe. + try: + core.shutdown() + finally: + store.clear_runtime() + except Exception: + # Cleanup failures are already logged; preserve the startup failure + # that explains why this rollback was necessary. + pass + raise + + if stop_evt.is_set(): + _cleanup_runtime(core, server, store) + return 0 + + try: + print(f"AgentBar v{__version__} 已启动") + print(f" 状态目录: {settings.state_dir}") + # 启动输出会被 launchd 持久化，绝不打印访问令牌。 + print(f" 任务面板: {server.url()}") + print(f" 模式: {'headless' if args.headless else 'menu bar'}") + except Exception: + try: + _cleanup_runtime(core, server, store) + except Exception: + pass + raise if args.headless: try: while not stop_evt.wait(0.5): pass finally: - core.shutdown() - server.stop() - store.clear_runtime() + _cleanup_runtime(core, server, store) return 0 # menu bar 模式：AppKit 事件循环占主线程；信号在后台线程清理后停掉主循环 - from .menubar import AgentBarApp # 延迟 import，headless/测试不依赖 pyobjc + try: + from .menubar import AgentBarApp # 延迟 import，headless/测试不依赖 pyobjc - app = AgentBarApp(core, settings, server) - server.hooks["dispatch"] = app.dispatch_async # /api/debug/dispatch 通道 + app = AgentBarApp(core, settings, server) + server.hooks["dispatch"] = app.dispatch_async # /api/debug/dispatch 通道 + except Exception: + try: + _cleanup_runtime(core, server, store) + except Exception: + pass + raise def _watch_signal(): stop_evt.wait() - core.shutdown() - server.stop() - store.clear_runtime() - app.stop_from_thread() + try: + _cleanup_runtime(core, server, store, app) + finally: + app.stop_from_thread() - threading.Thread(target=_watch_signal, daemon=True).start() - app.run() + try: + threading.Thread(target=_watch_signal, daemon=True).start() + app.run() + finally: + # Normal menu quit cleans up before terminating AppKit; this also covers + # an unexpected event-loop return. All component stops are idempotent. + _cleanup_runtime(core, server, store, app) return 0 @@ -181,11 +264,17 @@ def cmd_status(args, settings: Settings) -> int: def cmd_open(args, settings: Settings) -> int: base, token = _endpoint(settings) - url = f"{base}/?token={token}" - print(url) + # Fragment 只由浏览器端消费，不会进入 HTTP/proxy 日志。 + # 终端只显示无凭据基础 URL，避免 shell 历史/日志泄漏 token。 + url = f"{base}/#{urlencode({'token': token})}" + print(f"正在打开 AgentBar 面板：{base}/") if _instance_alive(settings): if not open_panel_url(url): - print("无法调用 macOS 浏览器；请复制上面的链接手动打开。", file=sys.stderr) + print( + "无法调用 macOS 浏览器；为避免泄漏访问令牌，" + "未在终端输出完整链接。请从 AgentBar 菜单栏打开面板。", + file=sys.stderr, + ) return 1 else: print("（调度器未运行，先执行 agentbar run）") @@ -226,7 +315,7 @@ def main(argv: list[str] | None = None) -> None: sp = sub.add_parser("run", help="启动调度器（默认带菜单栏）") sp.add_argument("--headless", action="store_true", help="无菜单栏模式（服务器/调试）") - sp.add_argument("--port", type=int, help=f"API 端口（默认取配置）") + sp.add_argument("--port", type=int, help="API 端口（默认取配置）") sp.set_defaults(fn=cmd_run) sp = sub.add_parser("add", help="添加任务") @@ -257,5 +346,17 @@ def main(argv: list[str] | None = None) -> None: sp.set_defaults(fn=cmd_log) args = p.parse_args(argv) - settings = load_settings(Path(args.state_dir) if args.state_dir else None) - sys.exit(args.fn(args, settings) or 0) + state_dir = Path(args.state_dir) if args.state_dir else None + # Client commands must retain the token understood by an older process that + # is still running. Secure schema/token migration belongs to the next server + # start, not to `status`, `open`, or another read/write client invocation. + provisional = load_settings(state_dir, migrate=False) + if args.cmd == "run": + if args.port is not None: + provisional.port = args.port + if _instance_alive(provisional): + print("AgentBar 已在运行（用 `agentbar open` 打开面板，或先退出旧实例）") + sys.exit(2) + settings = load_settings(state_dir, migrate=True) + sys.exit(cmd_run(args, settings, instance_checked=True) or 0) + sys.exit(args.fn(args, provisional) or 0) diff --git a/agentbar/config.py b/agentbar/config.py index 4760b42..9603879 100644 --- a/agentbar/config.py +++ b/agentbar/config.py @@ -3,14 +3,39 @@ from __future__ import annotations import copy +import fcntl import json +import math import os import secrets import stat +import tempfile +import threading from dataclasses import dataclass, field from pathlib import Path DEFAULT_PORT = 8737 +CONFIG_SCHEMA_VERSION = 2 +LEGACY_CLAUDE_CREDENTIAL_CACHE = "claude_credentials.json" +LEGACY_CLAUDE_CREDENTIAL_TEMP = "claude_credentials.tmp" + +# Claude/Codex 订阅额度默认不读取本机登录态。用户必须在额度设置中 +# 显式启用、选择模型并输入 OAuth access token。这样未配置的来源不会读 Keychain、 +# 不会扫描 CLI 凭据，也不会发起网络请求。 +DEFAULT_QUOTA_SOURCES: dict = { + "claude": { + "enabled": False, + "model": "", + "access_token": "", + "account_id": "", + }, + "codex": { + "enabled": False, + "model": "", + "access_token": "", + "account_id": "", + }, +} # 快手内部额度 provider（MyToken / Tokenverse）默认配置。 # 均为凭 corp SSO cookie 访问的信用额度（credits）接口，默认关闭—— @@ -20,18 +45,96 @@ "mytoken": {"enabled": False, "cookie": "", "unit": "credits", "refresh_seconds": 300}, "tokenverse": {"enabled": False, "cookie": "", "unit": "credits", "refresh_seconds": 300}, } +PROVIDER_HOSTS = { + "mytoken": "mytoken.corp.kuaishou.com", + "tokenverse": "tokenverse.corp.kuaishou.com", +} +PROVIDER_LOGIN_URLS = { + name: f"https://{host}/usage" for name, host in PROVIDER_HOSTS.items() +} PROVIDER_UNITS = ("credits", "percent", "token") +_SAVE_LOCK = threading.RLock() def _merge_providers(user: dict | None) -> dict: - """把用户 config 里的 providers 与内置默认深合并，补齐缺失字段、丢弃未知 provider。""" + """Return canonical provider config for any JSON-shaped input.""" merged = copy.deepcopy(DEFAULT_PROVIDERS) + user = user if isinstance(user, dict) else {} for name, defaults in merged.items(): - override = (user or {}).get(name) + override = user.get(name) if isinstance(override, dict): - for k, v in override.items(): - if k in defaults: - defaults[k] = v + enabled = override.get("enabled") + defaults["enabled"] = enabled if isinstance(enabled, bool) else False + + cookie = override.get("cookie") + defaults["cookie"] = ( + cookie.strip() + if isinstance(cookie, str) + and len(cookie) <= 65_536 + and not any(char in cookie for char in "\r\n\x00") + else "" + ) + + unit = override.get("unit") + if isinstance(unit, str) and unit in PROVIDER_UNITS: + defaults["unit"] = unit + + seconds = override.get("refresh_seconds", defaults["refresh_seconds"]) + try: + parsed_seconds = int(seconds) if not isinstance(seconds, bool) else 0 + except (TypeError, ValueError, OverflowError): + parsed_seconds = defaults["refresh_seconds"] + defaults["refresh_seconds"] = ( + min(86_400, max(60, parsed_seconds)) + if parsed_seconds > 0 + else defaults["refresh_seconds"] + ) + if defaults["enabled"] and not defaults["cookie"]: + # A provider without credentials is not runnable. Canonicalize + # the half-configured state instead of leaving a source that can + # only fail (and used to be retried after every restart). + defaults["enabled"] = False + return merged + + +def _merge_quota_sources(user: dict | None) -> dict: + """Return canonical quota-source config and migrate legacy ``api_key``.""" + merged = copy.deepcopy(DEFAULT_QUOTA_SOURCES) + user = user if isinstance(user, dict) else {} + for name, defaults in merged.items(): + override = user.get(name) + if isinstance(override, dict): + enabled = override.get("enabled") + defaults["enabled"] = enabled if isinstance(enabled, bool) else False + + for key, max_len in (("model", 160), ("account_id", 256)): + value = override.get(key) + defaults[key] = ( + value.strip() + if isinstance(value, str) + and len(value) <= max_len + and not any(char in value for char in "\r\n\x00") + else "" + ) + + # v0.10.3 预发布配置曾把 OAuth token 误称为 api_key。 + # 新字段缺失时无损迁移；显式 access_token 始终优先。 + token = ( + override.get("access_token") + if "access_token" in override + else override.get("api_key") + ) + defaults["access_token"] = ( + token.strip() + if isinstance(token, str) + and len(token) <= 16_384 + and not any(char in token for char in "\r\n\x00") + else "" + ) + # Invalid/missing credentials must never leave an enabled source + # that crashes fetcher construction or repeatedly reports errors. + if defaults["enabled"] and not defaults["access_token"]: + defaults["enabled"] = False return merged @@ -45,6 +148,7 @@ def default_state_dir() -> Path: @dataclass class Settings: state_dir: Path + config_schema_version: int = CONFIG_SCHEMA_VERSION port: int = DEFAULT_PORT max_parallel: int = 1 # 全局并行度，1 = 串行 per_tool_limit: int = 1 # 每个 AI CLI 的并行上限 @@ -53,14 +157,23 @@ class Settings: task_timeout_seconds: int = 7200 # 单任务运行上限 backoff_minutes: list[float] = field(default_factory=lambda: [5, 15, 30, 60]) usage_refresh_seconds: int = 120 # 订阅额度接口轮询间隔（最小 30 秒） + # 默认仅启动时、保存设置后或用户手动点击时刷新；避免后台反复请求。 + usage_auto_refresh: bool = False tick_seconds: float = 1.0 tool_paths: dict = field(default_factory=dict) # 手动指定 CLI 路径: {"claude": "/path"} - lan_access: bool = True # 绑定 0.0.0.0 供同一局域网的手机访问（API 仍需 token） + lan_access: bool = False # 默认只绑定本机；显式开启后才允许局域网手机访问 token: str = "" # 快手内部额度 provider 配置（见 DEFAULT_PROVIDERS）；cookie 空 / enabled=False 则不拉取。 providers: dict = field(default_factory=lambda: copy.deepcopy(DEFAULT_PROVIDERS)) + # Claude/Codex 额度查询：显式选择来源/模型，并配置 OAuth access token。 + quota_sources: dict = field(default_factory=lambda: copy.deepcopy(DEFAULT_QUOTA_SOURCES)) # 菜单栏标题显示哪个 provider 的用量百分比：claude / codex / mytoken / tokenverse title_provider: str = "claude" + _lock: threading.RLock = field( + default_factory=threading.RLock, + repr=False, + compare=False, + ) @property def config_path(self) -> Path: @@ -68,6 +181,7 @@ def config_path(self) -> Path: _PERSISTED_KEYS = ( + "config_schema_version", "port", "max_parallel", "per_tool_limit", @@ -76,15 +190,79 @@ def config_path(self) -> Path: "task_timeout_seconds", "backoff_minutes", "usage_refresh_seconds", + "usage_auto_refresh", "tool_paths", "lan_access", "token", "providers", + "quota_sources", "title_provider", ) -def load_settings(state_dir: Path | None = None) -> Settings: +def _normalize_loaded_settings(s: Settings) -> bool: + """Repair malformed user-edited values before worker threads consume them.""" + before = {key: copy.deepcopy(getattr(s, key)) for key in _PERSISTED_KEYS} + + def bounded_int(value, default: int, low: int, high: int) -> int: + if isinstance(value, bool): + return default + try: + parsed = int(value) + except (TypeError, ValueError, OverflowError): + return default + return parsed if low <= parsed <= high else default + + def bounded_float(value, default: float, low: float, high: float) -> float: + if isinstance(value, bool): + return default + try: + parsed = float(value) + except (TypeError, ValueError, OverflowError): + return default + return parsed if math.isfinite(parsed) and low <= parsed <= high else default + + s.port = bounded_int(s.port, DEFAULT_PORT, 1, 65_535) + s.config_schema_version = CONFIG_SCHEMA_VERSION + s.max_parallel = bounded_int(s.max_parallel, 1, 1, 64) + s.per_tool_limit = bounded_int(s.per_tool_limit, 1, 1, 64) + s.task_timeout_seconds = bounded_int(s.task_timeout_seconds, 7200, 1, 604_800) + s.usage_refresh_seconds = bounded_int(s.usage_refresh_seconds, 120, 30, 86_400) + s.tick_seconds = bounded_float(s.tick_seconds, 1.0, 0.05, 60.0) + s.allow_full_profile = s.allow_full_profile if isinstance(s.allow_full_profile, bool) else False + s.usage_auto_refresh = s.usage_auto_refresh if isinstance(s.usage_auto_refresh, bool) else False + s.lan_access = s.lan_access if isinstance(s.lan_access, bool) else False + s.default_cwd = s.default_cwd if isinstance(s.default_cwd, str) and s.default_cwd else str(Path.home()) + s.token = ( + s.token + if isinstance(s.token, str) + and 16 <= len(s.token) <= 512 + and not any(ord(char) < 32 or ord(char) == 127 for char in s.token) + else "" + ) + s.tool_paths = ( + {str(k): str(v) for k, v in s.tool_paths.items() if isinstance(k, str) and isinstance(v, str)} + if isinstance(s.tool_paths, dict) + else {} + ) + raw_backoff = s.backoff_minutes if isinstance(s.backoff_minutes, list) else [] + backoff = [] + for value in raw_backoff: + parsed = bounded_float(value, -1, 0, 10_080) + if parsed >= 0: + backoff.append(parsed) + s.backoff_minutes = backoff or [5, 15, 30, 60] + allowed_titles = {*DEFAULT_QUOTA_SOURCES, *DEFAULT_PROVIDERS} + if not isinstance(s.title_provider, str) or s.title_provider not in allowed_titles: + s.title_provider = "claude" + return before != {key: getattr(s, key) for key in _PERSISTED_KEYS} + + +def load_settings( + state_dir: Path | None = None, + *, + migrate: bool = True, +) -> Settings: sd = Path(state_dir).expanduser() if state_dir else default_state_dir() sd.mkdir(parents=True, exist_ok=True) try: @@ -92,38 +270,224 @@ def load_settings(state_dir: Path | None = None) -> Settings: except OSError: pass - s = Settings(state_dir=sd) cfg = sd / "config.json" + lock_path = cfg.with_name(".config.lock") + # A migrating load is a read/modify/write transaction: a missing or legacy + # config generates a token and commits it below. Keep an exclusive lock + # throughout. CLI client commands use migrate=False and a shared lock so + # merely querying an older, still-running AgentBar never rotates the token + # in its config behind the process's in-memory credentials. + with _SAVE_LOCK: + lock_fd = os.open( + lock_path, + os.O_RDWR | os.O_CREAT, + stat.S_IRUSR | stat.S_IWUSR, + ) + try: + os.fchmod(lock_fd, stat.S_IRUSR | stat.S_IWUSR) + fcntl.flock(lock_fd, fcntl.LOCK_EX if migrate else fcntl.LOCK_SH) + return _load_settings_holding_file_lock(sd, cfg, migrate=migrate) + finally: + try: + fcntl.flock(lock_fd, fcntl.LOCK_UN) + finally: + os.close(lock_fd) + + +def _load_settings_holding_file_lock( + sd: Path, + cfg: Path, + *, + migrate: bool, +) -> Settings: + """Load, normalize, and if needed commit while ``.config.lock`` is held.""" + s = Settings(state_dir=sd) data: dict = {} if cfg.exists(): + raw = cfg.read_bytes() try: - data = json.loads(cfg.read_text(encoding="utf-8")) - except (json.JSONDecodeError, OSError): - data = {} + parsed = json.loads(raw) + except (json.JSONDecodeError, UnicodeDecodeError): + parsed = None + if isinstance(parsed, dict): + data = parsed + elif migrate: + # Preserve the exact damaged input before replacing it with a + # canonical config. This also covers valid JSON scalars/arrays, + # which are not valid AgentBar configurations. + _backup_corrupt_config(cfg, raw) for key in _PERSISTED_KEYS: if key in data: setattr(s, key, data[key]) + # v1 exposed the admin API on the LAN by default. Because every old config + # persisted that default, merely changing the dataclass would leave existing + # installations exposed. Secure the one-time migration; users who really + # want LAN access can explicitly re-enable it afterwards. + try: + old_schema = int(data.get("config_schema_version", 1)) + except (TypeError, ValueError, OverflowError): + old_schema = 1 + if migrate and old_schema < 2: + s.lan_access = False + # Older releases printed token-bearing panel URLs to launchd stdout. + # Rotate once so a value copied into browser history or legacy /tmp + # logs cannot continue administering the upgraded instance. + s.token = secrets.token_urlsafe(24) + # providers 始终与内置默认深合并：老 config 缺字段时补齐，未知 provider 丢弃。 s.providers = _merge_providers(data.get("providers")) + s.quota_sources = _merge_quota_sources(data.get("quota_sources")) + normalized = _normalize_loaded_settings(s) + if not migrate and isinstance(data.get("token"), str): + # The provisional CLI client must authenticate exactly as an already + # running older process does, even if its manually edited token would be + # rotated by today's stricter length/control-character validation. + s.token = data["token"] - changed = not cfg.exists() or set(_PERSISTED_KEYS) - set(data) + raw_sources = data.get("quota_sources") + legacy_quota_key = isinstance(raw_sources, dict) and any( + isinstance(source, dict) and "api_key" in source + for source in raw_sources.values() + ) + changed = bool( + not cfg.exists() + or set(_PERSISTED_KEYS) - set(data) + or legacy_quota_key + or normalized + or data.get("providers") != s.providers + or data.get("quota_sources") != s.quota_sources + ) if not s.token: s.token = secrets.token_urlsafe(24) changed = True - if changed: - save_settings(s) + if changed and migrate: + _write_settings_file_holding_lock(s, cfg) + elif not changed: + # A user may have hand-edited or restored a valid config with a loose + # umask. Permissions are an invariant, not merely a side effect of writes. + try: + os.chmod(cfg, stat.S_IRUSR | stat.S_IWUSR) + except OSError: + pass + # Releases before v0.10.4 hard-coded this cache under ~/.agentbar even when + # AgentBar itself used a custom state directory. Delete only the two exact + # historical names in both locations; the .tmp variant can remain after an + # interrupted legacy write and contains the same OAuth token. + if migrate: + legacy_default_dir = Path.home() / ".agentbar" + legacy_paths = { + directory / name + for directory in (sd, legacy_default_dir) + for name in ( + LEGACY_CLAUDE_CREDENTIAL_CACHE, + LEGACY_CLAUDE_CREDENTIAL_TEMP, + ) + } + for legacy_path in legacy_paths: + try: + legacy_path.unlink(missing_ok=True) + except OSError: + pass return s -def save_settings(s: Settings) -> None: - cfg = s.config_path +def _backup_corrupt_config(cfg: Path, raw: bytes) -> Path: + """Create a private, non-overwriting forensic copy of a damaged config.""" + base = cfg.with_name(f"{cfg.name}.corrupt") + candidate = base + index = 0 + while True: + fd = -1 + try: + fd = os.open( + candidate, + os.O_WRONLY | os.O_CREAT | os.O_EXCL, + stat.S_IRUSR | stat.S_IWUSR, + ) + os.fchmod(fd, stat.S_IRUSR | stat.S_IWUSR) + with os.fdopen(fd, "wb") as handle: + fd = -1 + handle.write(raw) + handle.flush() + os.fsync(handle.fileno()) + return candidate + except FileExistsError: + index += 1 + candidate = cfg.with_name(f"{cfg.name}.corrupt.{index}") + except Exception: + # Never replace the original after creating only a partial backup. + try: + candidate.unlink(missing_ok=True) + except OSError: + pass + raise + finally: + if fd >= 0: + os.close(fd) + + +def _write_settings_file_holding_lock(s: Settings, cfg: Path) -> None: + """Atomically write ``s``; caller owns the process and config-file locks.""" payload = {k: getattr(s, k) for k in _PERSISTED_KEYS} - tmp = cfg.with_suffix(".json.tmp") - tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") - os.replace(tmp, cfg) + tmp_path: str | None = None + fd = -1 try: - os.chmod(cfg, stat.S_IRUSR | stat.S_IWUSR) # 0600：含 API token - except OSError: - pass + fd, tmp_path = tempfile.mkstemp( + prefix=".config.", suffix=".tmp", dir=str(cfg.parent) + ) + os.fchmod(fd, stat.S_IRUSR | stat.S_IWUSR) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + fd = -1 + handle.write(json.dumps(payload, ensure_ascii=False, indent=2)) + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp_path, cfg) + tmp_path = None + # The replacement already inherits the temp file's fchmod(0600). A + # redundant chmod must not turn a committed save into an apparent + # failure that makes callers roll their in-memory transaction back. + try: + os.chmod(cfg, stat.S_IRUSR | stat.S_IWUSR) + except OSError: + pass + finally: + if fd >= 0: + os.close(fd) + if tmp_path: + try: + os.unlink(tmp_path) + except OSError: + pass + + +def save_settings(s: Settings) -> None: + """Atomically persist settings without exposing or racing secret data. + + ``Settings._lock`` lets callers hold a read/modify/write transaction. The + module lock covers multiple Settings instances in one process; ``flock`` + covers accidental concurrent AgentBar processes. + """ + # Every settings read/modify/write path acquires the per-instance lock + # first. Keep that order here as well: callers are allowed to already hold + # ``s._lock``, and reversing the order would deadlock against such a caller + # while another thread is doing a plain save. + with s._lock, _SAVE_LOCK: + cfg = s.config_path + cfg.parent.mkdir(parents=True, exist_ok=True) + lock_path = cfg.with_name(".config.lock") + lock_fd = os.open( + lock_path, + os.O_RDWR | os.O_CREAT, + stat.S_IRUSR | stat.S_IWUSR, + ) + try: + os.fchmod(lock_fd, stat.S_IRUSR | stat.S_IWUSR) + fcntl.flock(lock_fd, fcntl.LOCK_EX) + _write_settings_file_holding_lock(s, cfg) + finally: + try: + fcntl.flock(lock_fd, fcntl.LOCK_UN) + finally: + os.close(lock_fd) diff --git a/agentbar/menu_spec.py b/agentbar/menu_spec.py index 6c4006a..e3b44ee 100644 --- a/agentbar/menu_spec.py +++ b/agentbar/menu_spec.py @@ -132,15 +132,24 @@ def _window_value(w: dict) -> str: def _quota_compact(qi: dict) -> str: windows = qi.get("windows") or [] if windows: - return " · ".join( + value = " · ".join( f"{w['label']} {_window_value(w)}" for w in windows[:2] ) + return value + ("（已过期）" if qi.get("stale") else "") return {"ok": "正常", "limited": "受限", "unknown": "未知"}.get(qi.get("state"), "未知") -def _quota_submenu(tool: str, qi: dict) -> dict: +def _quota_submenu(tool: str, qi: dict, quota_cfg: dict | None = None) -> dict: dot = {"ok": "🟢", "limited": "🟠"}.get(qi.get("state"), "⚪") children: list[dict] = [] + model = str(qi.get("model") or "").strip() + if model: + children.append(_info(f"模型 {model}")) + available_models = [ + str(value) for value in (qi.get("available_models") or []) if str(value).strip() + ] + if available_models: + children.append(_info(f"可选额度标识 {'、'.join(available_models[:4])}")) for w in qi.get("windows") or []: line = f"{w['label']} 已用 {_window_value(w)}" if w.get("unit") in ("credits", "token"): @@ -154,16 +163,81 @@ def _quota_submenu(tool: str, qi: dict) -> dict: meta.append(f"来源 {qi.get('source') or 'none'}") if qi.get("fetched_at"): meta.append(f"{_clock(qi['fetched_at'])} 更新") + if qi.get("stale"): + meta.append("数据已过期") children.append(_action(" · ".join(meta))) if not (qi.get("windows")): children.append(_action(qi.get("detail") or "暂无额度数据")) if qi.get("error"): children.append(_action(f"⚠ {qi['error'][:70]}")) children.append(_sep()) - children.append(_action("↻ 立即刷新额度", "refresh_quota")) - if tool == "claude" and "Keychain" in (qi.get("error") or ""): - children.append(_action("🔑 授权读取 Claude Keychain…", "authorize_keychain")) - return _submenu(f"{dot} {_provider_name(tool)} · {_quota_compact(qi)}", children) + if tool in ("claude", "codex"): + # 任务观测也可能产生 quota 行，但只有显式启用且已配 + # 手动密钥的来源才能请求上游；配置入口始终保留。 + quota_cfg = quota_cfg or {} + if quota_cfg.get("enabled") and quota_cfg.get("key_set"): + children.append(_action("↻ 立即刷新额度", f"refresh_quota:{tool}")) + children.append(_action(f"⚙ 配置 {_provider_name(tool)}…", "provider_settings")) + else: + children.append(_action("↻ 立即刷新额度", f"refresh_quota:{tool}")) + if tool in ("mytoken", "tokenverse"): + children.append(_action(f"⚙ 配置 {_provider_name(tool)}…", "provider_settings")) + model_label = f" · {model}" if model else "" + return _submenu( + f"{dot} {_provider_name(tool)}{model_label} · {_quota_compact(qi)}", + children, + ) + + +def _quota_source_setup_submenu(tool: str, cfg: dict) -> dict: + """Claude/Codex manual quota setup row when no live snapshot exists.""" + enabled = bool(cfg.get("enabled")) + key_set = bool(cfg.get("key_set")) + model = str(cfg.get("model") or "").strip() + if enabled and key_set: + state = "等待刷新" + detail = "已启用且额度凭据已配置。" + elif enabled: + state = "缺少密钥" + detail = "已启用，但尚未配置额度访问凭据。" + elif key_set: + state = "未启用" + detail = "额度凭据已保存；启用后才会请求额度。" + else: + state = "未配置" + detail = "点击配置，手动选择额度模型并输入访问凭据。" + children = [_info(detail)] + if model: + children.append(_info(f"模型 {model}")) + if enabled and key_set: + children.append(_action("↻ 立即刷新额度", f"refresh_quota:{tool}")) + children.append(_action(f"⚙ 配置 {_provider_name(tool)}…", "provider_settings")) + model_label = f" · {model}" if model else "" + return _submenu( + f"⚪ {_provider_name(tool)}{model_label} · {state}", + children, + ) + + +def _provider_setup_submenu(tool: str, cfg: dict) -> dict: + enabled = bool(cfg.get("enabled")) + cookie_set = bool(cfg.get("cookie_set")) + if enabled and cookie_set: + state = "等待刷新" + detail = "已启用且 Cookie 已配置，正在等待额度接口返回。" + elif enabled: + state = "缺少 Cookie" + detail = "已启用，但尚未配置企业 SSO Cookie。" + elif cookie_set: + state = "未启用" + detail = "Cookie 已保存；启用后才会请求并显示额度。" + else: + state = "未配置" + detail = "点击配置，从本机浏览器导入企业 SSO 登录态。" + return _submenu( + f"⚪ {_provider_name(tool)} · {state}", + [_info(detail), _action(f"⚙ 配置 {_provider_name(tool)}…", "provider_settings")], + ) def _mobile_submenu(t: dict) -> dict: @@ -215,6 +289,7 @@ def build_menu_spec(snapshot: dict) -> list[dict]: rows.append(_sep()) quota = snapshot.get("quota") or {} + quota_source_config = snapshot.get("quota_source_config") # claude/codex 优先，其余 corp provider（mytoken/tokenverse…）按 key 顺序附后 ordered = [t for t in ("claude", "codex") if t in quota] + [ t for t in quota if t not in ("claude", "codex") @@ -223,13 +298,36 @@ def build_menu_spec(snapshot: dict) -> list[dict]: for tool in ordered: qi = quota.get(tool) if qi: - rows.append(_quota_submenu(tool, qi)) + cfg = ( + quota_source_config.get(tool) or {} + if isinstance(quota_source_config, dict) and tool in ("claude", "codex") + else None + ) + rows.append(_quota_submenu(tool, qi, cfg)) shown = True + # Claude/Codex 额度源是显式 opt-in；即使未启用也保留配置入口。 + # 老快照不含 quota_source_config 时保持历史输出不变。 + if isinstance(quota_source_config, dict): + for tool in ("claude", "codex"): + if tool not in quota: + rows.append(_quota_source_setup_submenu( + tool, quota_source_config.get(tool) or {}, + )) + shown = True + # 未配置的内部 provider 也必须出现在真实菜单中，否则用户无从发现入口。 + # 老测试/第三方调用未提供 provider_config 时保持历史输出不变。 + provider_config = snapshot.get("provider_config") + if isinstance(provider_config, dict): + for tool in ("mytoken", "tokenverse"): + if tool not in quota: + rows.append(_provider_setup_submenu(tool, provider_config.get(tool) or {})) + shown = True if shown: rows.append(_sep()) rows.append(_action("↗ 打开任务面板", "open_panel")) rows.append(_action("＋ 快速添加任务…", "quick_add")) + rows.append(_action("⚙ 额度设置…", "provider_settings")) rows.append(_mobile_submenu(snapshot.get("tunnel") or {})) rows.append(_sep()) if snapshot.get("paused"): diff --git a/agentbar/menubar.py b/agentbar/menubar.py index 9d946c6..d957611 100644 --- a/agentbar/menubar.py +++ b/agentbar/menubar.py @@ -1,13 +1,13 @@ """macOS menu-bar frontend — raw AppKit (NSStatusItem + NSMenu), rumps removed. 为什么不用 rumps：在 macOS 26 上出现两类事故—— -1. 菜单回调在主线程做阻塞调用（等 `open`、Keychain+HTTP）→ 整个 App 卡死； +1. 菜单回调在主线程做阻塞调用（等 `open`、登录或 HTTP）→ 整个 App 卡死； 2. rumps.Timer 每 2s clear+rebuild 打开中的菜单 → 菜单项点击落空。 本实现的纪律： - 菜单内容只在 menuWillOpen（AppKit 正统时机）重建；NSTimer 只更新标题文本。 - 所有 action 回调毫秒级返回：浏览器用 open_url_async（fire-and-forget）， - Keychain 授权丢后台线程，结果经 AppHelper.callAfter 回主线程弹提示。 + 登录、Cookie 导入等慢操作都在后台线程执行。 - setAutoenablesItems(False) + 显式 setEnabled，杜绝系统校验导致的置灰。 - 每次重建后把真实 NSMenu 状态导出到 state_dir/menu-debug.json，可实证核查。 @@ -18,8 +18,12 @@ import json import logging +import os +import tempfile import threading import time +from html import escape +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit import objc from AppKit import ( @@ -60,6 +64,30 @@ ) +def _atomic_private_write(path, text: str) -> None: + """Atomically replace a sensitive debug artifact with mode 0600.""" + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp( + dir=str(path.parent), prefix=f".{path.name}.", suffix=".tmp" + ) + try: + os.fchmod(fd, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as stream: + fd = -1 # ownership transferred to ``stream`` + stream.write(text) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + except Exception: + if fd >= 0: + os.close(fd) + try: + os.unlink(temporary) + except FileNotFoundError: + pass + raise + + def _ring_icon(outer: float | None, inner: float | None, status: str = "idle") -> NSImage: """画双环额度图标：模板图（黑+透明度），菜单栏深浅色自动适配。 @@ -110,6 +138,9 @@ def _qr_page_html(url: str, mode: str = "局域网", note: str = "") -> str: img = qrcode.make(url, image_factory=qrcode.image.svg.SvgPathImage) svg = img.to_string().decode() + safe_url = escape(url) + safe_mode = escape(mode) + safe_note = escape(note or "链接含访问令牌，勿外传") return f""" -

手机扫码 · 查看/提交任务（{mode}）

-

{note or "链接含访问令牌，勿外传"}

+

手机扫码 · 查看/提交任务（{safe_mode}）

+

{safe_note}

{svg} -
{url}
+
{safe_url}
""" +def _token_fragment_url(url: str, token: str) -> str: + """Put the bootstrap token in a fragment, outside HTTP requests and logs.""" + parts = urlsplit(url) + query = [(key, value) for key, value in parse_qsl(parts.query) if key != "token"] + fragment = [ + (key, value) for key, value in parse_qsl(parts.fragment) if key != "token" + ] + fragment.append(("token", token)) + return urlunsplit(( + parts.scheme, + parts.netloc, + parts.path, + urlencode(query), + urlencode(fragment), + )) + + class _Bridge(NSObject): """Objective-C 桥：菜单 delegate + action/timer target（只在主线程被调用）。""" @@ -158,9 +206,14 @@ def __init__(self, core: Scheduler, settings: Settings, server: ApiServer): self._menu = None self._timer = None self._panel = None # 原生任务面板窗口（懒加载） + self._provider_panel = None # MyToken / Tokenverse 原生设置窗口（懒加载） self._ring_key = None # 双环图标缓存键（进度没变不重画，避免 2s 一次的无谓刷新） self._qr_window = None # 手机访问二维码窗口（懒加载） self._qr_webview = None # 二维码窗口里的 WKWebView（内容可切换 LAN/公网） + self._debug_lock = threading.Lock() + self._debug_pending: tuple[object, str] | None = None + self._debug_worker_running = False + self._closing = threading.Event() from .tunnel import TunnelManager self.tunnel = TunnelManager( server.port, on_up=server.allow_host, on_down=server.disallow_host @@ -196,8 +249,20 @@ def stop_from_thread(self) -> None: """信号处理线程调用：清理已在外部完成，只负责停掉主事件循环。""" AppHelper.callAfter(self._terminate) + def begin_shutdown(self) -> None: + """Close menu/debug action admission before retiring external resources.""" + self._closing.set() + + def close_active_provider_logins(self) -> None: + """Safe from shutdown threads; no-op until the provider window is opened.""" + panel = self._provider_panel + if panel is not None: + panel.close_active_logins() + def dispatch_async(self, action: str) -> None: """任意线程安全：把动作转投主线程，与真实菜单点击走同一 _dispatch。""" + if self._closing.is_set() or getattr(self.server, "stopping", False): + return AppHelper.callAfter(self._dispatch, action) def _terminate(self) -> None: @@ -319,16 +384,25 @@ def _render(self, menu, nodes: list[dict]) -> None: # ---------- actions（必须毫秒级返回，禁止任何阻塞） ---------- def _dispatch(self, action: str) -> None: + # An action may already be queued on AppKit when SIGTERM begins. Drop it + # at execution time too, so a late tunnel_start cannot recreate a child + # after shutdown has stopped and cleared the current tunnel. + if action != "quit" and ( + self._closing.is_set() or getattr(self.server, "stopping", False) + ): + return log.info("menu action: %s", action) try: if action == "open_panel": self._show_panel(False) elif action == "quick_add": self._show_panel(True) - elif action == "refresh_quota": - self.core.quota.refresh_now() # 异步：只置事件 - elif action == "authorize_keychain": - self._authorize_keychain_bg() + elif action == "provider_settings": + self._show_provider_settings() + elif action.startswith("refresh_quota:"): + tool = action.split(":", 1)[1].strip() + if tool: + self.core.quota.refresh_now(tool) elif action == "pause_all": self.core.pause_all() elif action == "resume_all": @@ -357,8 +431,20 @@ def _show_panel(self, focus_prompt: bool) -> None: ) self._panel.show_(focus_prompt) + def _show_provider_settings(self) -> None: + """原生内部额度设置：菜单里始终可发现，不再藏在 Web 面板。""" + if self._provider_panel is None: + from .provider_window import ProviderSettingsWindowController + + self._provider_panel = ( + ProviderSettingsWindowController.alloc().initWithCore_settings_( + self.core, self.settings + ) + ) + self._provider_panel.show_(None) + def _show_mobile_qr(self) -> None: - """局域网扫码：http://:/m?token=…（手机与 Mac 同一 Wi-Fi）。""" + """局域网扫码；访问令牌只放 URL fragment。""" url = self.server.mobile_url() if not url: self._alert( @@ -366,7 +452,11 @@ def _show_mobile_qr(self) -> None: "未获取到局域网 IP（Mac 未联网？），或 config.json 中 lan_access 已关闭。", ) return - self._show_qr_window(url, "局域网", "手机需与 Mac 连同一 Wi-Fi") + self._show_qr_window( + _token_fragment_url(url, self.settings.token), + "局域网", + "手机需与 Mac 连同一 Wi-Fi", + ) def _start_tunnel_bg(self) -> None: """开通公网隧道（cloudflared，阻塞 ~5-15s → 后台线程），成功后自动弹二维码。""" @@ -386,7 +476,7 @@ def _show_tunnel_qr(self) -> None: self._alert("公网访问", "隧道未开通（先点「开通公网访问」）。") return self._show_qr_window( - f"{url}/m?token={self.settings.token}", + _token_fragment_url(f"{url.rstrip('/')}/m", self.settings.token), "公网（Cloudflare Tunnel）", "任何网络可达 · 链接含令牌切勿外传 · 每次开通域名会变", ) @@ -422,22 +512,6 @@ def _show_qr_window(self, url: str, mode: str, note: str) -> None: log.exception("qr window failed") self._alert("手机访问", f"二维码窗口创建失败；手机浏览器直接打开：\n{url}") - def _authorize_keychain_bg(self) -> None: - def work(): - try: - ok = self.core.quota.authorize_claude_keychain() - except Exception: - log.exception("keychain authorize failed") - ok = False - msg = ( - "已取得 Keychain 授权，额度数据已刷新。" - if ok - else "未能读取 Keychain；请确认已登录 Claude Code 后重试。" - ) - AppHelper.callAfter(self._alert, "Claude 额度", msg) - - threading.Thread(target=work, name="agentbar-keychain", daemon=True).start() - def _alert(self, title: str, text: str) -> None: self._nsapp.activateIgnoringOtherApps_(True) alert = NSAlert.alloc().init() @@ -447,13 +521,28 @@ def _alert(self, title: str, text: str) -> None: def _quit(self) -> None: log.info("quit from menu") + self.begin_shutdown() + first_error = None try: - self.tunnel.stop() # 杀掉 cloudflared，避免孤儿进程占着公网域名 - self.core.shutdown() - self.server.stop() - self.core.store.clear_runtime() + for label, action in ( + # Stop admission first. Scheduler shutdown may wait for active + # CLI processes; HTTP cannot remain open during that window. + ("HTTP server", self.server.stop), + ("public tunnel", self.tunnel.stop), + ("provider login", self.close_active_provider_logins), + ("scheduler", self.core.shutdown), + ("runtime file", self.core.store.clear_runtime), + ): + try: + action() + except Exception as exc: + if first_error is None: + first_error = exc + log.exception("quit cleanup failed: %s", label) finally: self._terminate() + if first_error is not None: + raise first_error # ---------- evidence ---------- @@ -462,10 +551,36 @@ def _write_debug_dump(self) -> None: payload = {"ts": time.time(), "title": str(self._item.button().title()), "items": self._dump(self._menu)} path = self.settings.state_dir / "menu-debug.json" - path.write_text(json.dumps(payload, ensure_ascii=False, indent=1), - encoding="utf-8") + text = json.dumps(payload, ensure_ascii=False, indent=1) except Exception: log.debug("menu debug dump failed", exc_info=True) + return + with self._debug_lock: + # Menu open is a UI path. Coalesce rapid rebuilds and keep fsync off + # the AppKit thread; the newest snapshot always drains last. + self._debug_pending = (path, text) + if self._debug_worker_running: + return + self._debug_worker_running = True + threading.Thread( + target=self._drain_debug_dump, + name="agentbar-menu-debug", + daemon=True, + ).start() + + def _drain_debug_dump(self) -> None: + while True: + with self._debug_lock: + job = self._debug_pending + self._debug_pending = None + if job is None: + self._debug_worker_running = False + return + path, text = job + try: + _atomic_private_write(path, text) + except Exception: + log.debug("menu debug dump failed", exc_info=True) def _dump(self, menu) -> list[dict]: out = [] diff --git a/agentbar/models.py b/agentbar/models.py index 2315823..ead048c 100644 --- a/agentbar/models.py +++ b/agentbar/models.py @@ -29,8 +29,10 @@ class TaskState(str, Enum): # 权限档位：readonly=只读, edits=可编辑文件(默认), full=完全权限(必须显式开启) PROFILES = ("readonly", "edits", "full") -# 模型名称不硬编码：CLI 与账号可用模型会持续变化，留空时完全交给 CLI 默认。 -EFFORTS = ("low", "medium", "high", "xhigh", "max") +# CLI parser only needs the union. Runtime validation remains adapter-specific, +# so Codex can accept ``minimal`` while Claude can accept ``max`` without the +# command line maintaining a second, contradictory capability table. +EFFORTS = ("minimal", "low", "medium", "high", "xhigh", "max") def new_id() -> str: diff --git a/agentbar/panel_window.py b/agentbar/panel_window.py index f5f4eb9..b113d41 100644 --- a/agentbar/panel_window.py +++ b/agentbar/panel_window.py @@ -2,7 +2,9 @@ from __future__ import annotations +import html import logging +import threading import time import objc @@ -35,6 +37,7 @@ NSWindowStyleMaskTitled, ) from Foundation import NSObject, NSTimer +from PyObjCTools import AppHelper log = logging.getLogger("agentbar.panel") @@ -95,6 +98,13 @@ def initWithCore_settings_server_(self, core, settings, server): self._current_efforts = [] self._transcript_windows = {} self._transcript_meta = {} + # Transcript discovery/stat/parsing can touch large recursive session + # trees and multi-megabyte JSONL files. Keep one worker per task and let + # AppKit's main thread only apply an already-rendered result. + self._transcript_workers = {} + self._transcript_pending = set() + self._transcript_generation = {} + self._provider_settings_panel = None self._timer = NSTimer.scheduledTimerWithTimeInterval_target_selector_userInfo_repeats_( 2.0, self, "onTick:", None, True ) @@ -258,6 +268,11 @@ def _build(self): self.pause_btn.setAutoresizingMask_(NSViewMaxYMargin) v.addSubview_(self.pause_btn) + provider_btn = _button("⚙ 额度设置", W - PAD - 234, PAD + 26, 118, + self, "onProviderSettings:") + provider_btn.setAutoresizingMask_(NSViewMaxYMargin | 1) + v.addSubview_(provider_btn) + self.transcript_btn = _button("📄 查看对话", W - PAD - 108, PAD + 26, 108, self, "onShowTranscript:") self.transcript_btn.setAutoresizingMask_(NSViewMaxYMargin | 1) @@ -403,6 +418,13 @@ def onAdd_(self, _sender): except ValueError as e: self._alert("无法添加任务", str(e)) return + except OSError: + log.exception("persist added task failed") + self._alert( + "无法添加任务", + "本地状态保存失败，任务未添加；请检查磁盘后重试。", + ) + return self.prompt_view.setString_("") self.schedule_check.setState_(0) self.schedule_picker.setHidden_(True) @@ -466,10 +488,20 @@ def onShowTranscript_(self, _sender): with self.core._lock: live = self.core._tasks.get(t["id"]) if live and not live.session_id: + checkpoint = live.to_dict() live.session_id = sid - self.core._persist_locked() + try: + self.core._persist_locked() + except Exception: + self.core._restore_task_locked(live, checkpoint) + raise except Exception: log.exception("persist recovered sid failed") + self._alert( + "无法保存会话记录", + "本地状态文件写入失败，恢复结果未生效；请检查磁盘后重试。", + ) + return else: self._alert( "无会话记录", @@ -479,6 +511,17 @@ def onShowTranscript_(self, _sender): return self._open_transcript_window(t) + def onProviderSettings_(self, _sender): + if self._provider_settings_panel is None: + from .provider_window import ProviderSettingsWindowController + + self._provider_settings_panel = ( + ProviderSettingsWindowController.alloc().initWithCore_settings_( + self.core, self.settings + ) + ) + self._provider_settings_panel.show_(None) + @objc.python_method def _open_transcript_window(self, task_dict: dict): tid = task_dict["id"] @@ -515,7 +558,16 @@ def _open_transcript_window(self, task_dict: dict): wkview = self._make_webview(NSMakeRect(PAD, PAD, 780 - 2 * PAD, 580 - 2 * PAD - 36)) cv.addSubview_(wkview) - self._transcript_meta[tid] = {"tool": tool, "cwd": cwd, "sid": sid, "wkview": wkview} + generation = self._transcript_generation.get(tid, 0) + 1 + self._transcript_generation[tid] = generation + self._transcript_meta[tid] = { + "tool": tool, + "cwd": cwd, + "sid": sid, + "wkview": wkview, + "_generation": generation, + "_identity": (tool, cwd, sid), + } self._transcript_windows[tid] = tw self._refresh_transcript_window(tid) tw.makeKeyAndOrderFront_(None) @@ -524,8 +576,15 @@ def _open_transcript_window(self, task_dict: dict): def _make_webview(self, frame): """Create a WKWebView; fallback to NSTextView if WebKit is unavailable.""" try: - from WebKit import WKWebView, WKWebViewConfiguration + from WebKit import WKWebView, WKWebViewConfiguration, WKWebsiteDataStore cfg = WKWebViewConfiguration.alloc().init() + # 对话页是纯静态本地 HTML：禁用 JS 并使用非持久化存储， + # 即使渲染内容异常也不能执行脚本或污染其他 WebView 会话。 + try: + cfg.preferences().setJavaScriptEnabled_(False) + cfg.setWebsiteDataStore_(WKWebsiteDataStore.nonPersistentDataStore()) + except Exception: + pass wk = WKWebView.alloc().initWithFrame_configuration_(frame, cfg) wk.setAutoresizingMask_(NSViewWidthSizable | NSViewHeightSizable) return wk @@ -556,40 +615,134 @@ def _refresh_transcript_window(self, tid: str): pass if not sid: return - - from .transcript import find_session_file, to_html, parse_transcript - path = find_session_file(tool, cwd, sid) wkview = meta.get("wkview") if not wkview: return - # Skip reload if file hasn't changed since last load — preserves scroll position - cur_mtime = path.stat().st_mtime if path else None - if (cur_mtime is not None - and cur_mtime == meta.get("_mtime") - and sid == meta.get("_loaded_sid")): + identity = (tool, cwd, sid) + if identity != meta.get("_identity"): + generation = self._transcript_generation.get(tid, 0) + 1 + self._transcript_generation[tid] = generation + meta["_generation"] = generation + meta["_identity"] = identity + meta.pop("_loaded_key", None) + generation = meta.get("_generation", 0) + + # A timer tick while parsing merely requests one follow-up stat. This + # both coalesces rapid file appends and prevents concurrent full parses. + if tid in self._transcript_workers: + self._transcript_pending.add(tid) return - meta["_mtime"] = cur_mtime - meta["_loaded_sid"] = sid + # Avoid importing WebKit on every timer path. The view was constructed + # on the main thread already; capability detection is sufficient here. + rich = callable(getattr(wkview, "loadHTMLString_baseURL_", None)) + + self._transcript_workers[tid] = generation + self._transcript_pending.discard(tid) + worker = threading.Thread( + target=self._render_transcript_background, + args=( + tid, + generation, + identity, + meta.get("_loaded_key"), + rich, + ), + name=f"agentbar-transcript-{tid}", + daemon=True, + ) + try: + worker.start() + except Exception: + self._transcript_workers.pop(tid, None) + log.exception("failed to start transcript worker") + @objc.python_method + def _render_transcript_background( + self, tid, generation, identity, loaded_key, rich, + ): + """Discover and render a transcript without touching AppKit objects.""" + tool, cwd, sid = identity + key = None + content = None try: - from WebKit import WKWebView - if isinstance(wkview, WKWebView): + from .transcript import find_session_file, parse_transcript, to_html + + path = find_session_file(tool, cwd, sid) + if path: + stat_result = path.stat() + key = ( + sid, + str(path), + stat_result.st_mtime_ns, + stat_result.st_size, + rich, + ) + else: + key = (sid, None, None, None, rich) + if key != loaded_key: if path: - content = to_html(tool, path) + content = to_html(tool, path) if rich else parse_transcript(tool, path) + elif rich: + safe_sid = html.escape(str(sid)) + safe_tool = html.escape(str(tool)) + safe_cwd = html.escape(str(cwd)) + content = ( + "" + "" + "" + "

未找到会话文件
" + f"session_id: {safe_sid}
工具: {safe_tool}
" + f"目录: {safe_cwd}

" + ) else: - content = f"

未找到会话文件
session_id: {sid}
工具: {tool}
目录: {cwd}

" - content = f"{content}" - wkview.loadHTMLString_baseURL_(content, None) - return + content = f"[未找到会话文件]\nsession_id: {sid}\n工具: {tool}" except Exception: - pass - # Fallback: NSTextView plain text - if path: - text = parse_transcript(tool, path) - else: - text = f"[未找到会话文件]\nsession_id: {sid}\n工具: {tool}" - wkview.setString_(text) + log.exception("transcript rendering failed for task %s", tid) + # Keep the old successfully rendered document. The next timer tick + # may retry after a transient file replacement/read race. + key = None + content = None + AppHelper.callAfter( + self._finish_transcript_refresh, + tid, + generation, + identity, + key, + rich, + content, + ) + + @objc.python_method + def _finish_transcript_refresh( + self, tid, generation, identity, key, rich, content, + ): + """Apply one worker result on the AppKit main thread if still current.""" + if self._transcript_workers.get(tid) == generation: + self._transcript_workers.pop(tid, None) + meta = self._transcript_meta.get(tid) + window = self._transcript_windows.get(tid) + current = bool( + meta + and window + and window.isVisible() + and meta.get("_generation") == generation + and meta.get("_identity") == identity + ) + if current and key is not None and content is not None: + wkview = meta.get("wkview") + if wkview: + if rich: + wkview.loadHTMLString_baseURL_(content, None) + else: + wkview.setString_(content) + meta["_loaded_key"] = key + + pending = tid in self._transcript_pending + self._transcript_pending.discard(tid) + if pending and meta and window and window.isVisible(): + self._refresh_transcript_window(tid) def onCopyResumeCmd_(self, _sender): try: @@ -611,7 +764,15 @@ def _act_selected(self, action): tid = self._selected_id() if not tid: return - ok, msg = self.core.act(tid, action) + try: + ok, msg = self.core.act(tid, action) + except OSError: + log.exception("persist task action failed") + self._alert( + "操作失败", + "本地状态保存失败，操作未生效；请检查磁盘后重试。", + ) + return if not ok: self._alert("操作失败", msg) self.refresh() @@ -632,10 +793,18 @@ def onRetryTask_(self, _s): self._act_selected("retry") def onTogglePause_(self, _s): - if self.core.paused: - self.core.resume_all() - else: - self.core.pause_all() + try: + if self.core.paused: + self.core.resume_all() + else: + self.core.pause_all() + except OSError: + log.exception("persist pause toggle failed") + self._alert( + "操作失败", + "本地状态保存失败，暂停状态未更改；请检查磁盘后重试。", + ) + return self.refresh() @objc.python_method diff --git a/agentbar/processes.py b/agentbar/processes.py index 51b7178..fc29d79 100644 --- a/agentbar/processes.py +++ b/agentbar/processes.py @@ -12,10 +12,11 @@ import time from pathlib import Path -_CACHE_SECONDS = 3.0 +_CACHE_SECONDS = 5.0 _cache_lock = threading.Lock() _cache_at = 0.0 _cache_processes: dict[int, dict] = {} +_scan_inflight = False def _tool_for_executable(executable: str) -> str | None: @@ -58,14 +59,40 @@ def _scan_processes() -> dict[int, dict]: return rows +def _refresh_cache() -> None: + global _cache_at, _cache_processes, _scan_inflight + try: + scanned = _scan_processes() + completed_at = time.monotonic() + with _cache_lock: + _cache_processes = scanned + _cache_at = completed_at + finally: + with _cache_lock: + _scan_inflight = False + + def _processes() -> dict[int, dict]: - global _cache_at, _cache_processes + """Return cached metadata immediately and refresh it on one background worker.""" + global _scan_inflight now = time.monotonic() + worker = None with _cache_lock: - if now - _cache_at >= _CACHE_SECONDS: - _cache_processes = _scan_processes() - _cache_at = now - return dict(_cache_processes) + if now - _cache_at >= _CACHE_SECONDS and not _scan_inflight: + _scan_inflight = True + worker = threading.Thread( + target=_refresh_cache, + name="agentbar-process-scan", + daemon=True, + ) + cached = dict(_cache_processes) + if worker is not None: + try: + worker.start() + except Exception: + with _cache_lock: + _scan_inflight = False + return cached def _owner_for(pid: int, rows: dict[int, dict], owned: dict[int, dict]) -> dict | None: diff --git a/agentbar/provider_window.py b/agentbar/provider_window.py new file mode 100644 index 0000000..5d30b5e --- /dev/null +++ b/agentbar/provider_window.py @@ -0,0 +1,824 @@ +"""Native quota-source settings window. + +Claude/Codex use manually supplied OAuth access tokens. MyToken/Tokenverse use +corp browser cookies. Secret fields are always write-only: reopening the window +shows only whether a credential is configured, never the credential itself. +""" + +from __future__ import annotations + +import copy +import logging +import threading + +import objc +from AppKit import ( + NSAlert, + NSAlertFirstButtonReturn, + NSApp, + NSBackingStoreBuffered, + NSButton, + NSFont, + NSMakeRect, + NSPopUpButton, + NSSecureTextField, + NSTextField, + NSWindow, + NSWindowStyleMaskClosable, + NSWindowStyleMaskMiniaturizable, + NSWindowStyleMaskTitled, +) +from PyObjCTools import AppHelper +from Foundation import NSObject + +from .browser_cookies import CookieImportError, import_cookie_header +from .chrome_login import ChromeCDPLogin, ChromeLoginError +from .config import ( + DEFAULT_PROVIDERS, + DEFAULT_QUOTA_SOURCES, + PROVIDER_HOSTS, + PROVIDER_UNITS, + save_settings, +) + +log = logging.getLogger("agentbar.provider_window") + +W, H = 760, 760 +PAD = 18 +_QUOTA_SOURCES = ("claude", "codex") +_PROVIDERS = ("mytoken", "tokenverse") +_NAMES = { + "claude": "Claude", + "codex": "Codex", + "mytoken": "MyToken", + "tokenverse": "Tokenverse", +} +_TITLE_OPTIONS = ( + ("Claude", "claude"), + ("Codex", "codex"), + ("MyToken", "mytoken"), + ("Tokenverse", "tokenverse"), +) + + +def _label(text, x, y, w, h=18, *, bold=False, dim=False): + label = NSTextField.labelWithString_(text) + label.setFrame_(NSMakeRect(x, y, w, h)) + label.setFont_( + NSFont.boldSystemFontOfSize_(12) + if bold + else NSFont.systemFontOfSize_(11 if dim else 12) + ) + if dim: + label.setTextColor_(label.textColor().colorWithAlphaComponent_(0.58)) + return label + + +def _button(title, x, y, w, target, selector, h=27): + button = NSButton.alloc().initWithFrame_(NSMakeRect(x, y, w, h)) + button.setTitle_(title) + button.setBezelStyle_(1) + button.setTarget_(target) + button.setAction_(selector) + return button + + +class ProviderSettingsWindowController(NSObject): + + def initWithCore_settings_(self, core, settings): + self = objc.super(ProviderSettingsWindowController, self).init() + if self is None: + return None + self.core = core + self.settings = settings + self.window = None + self._enabled = {} + self._unit = {} + self._refresh = {} + self._cookie = {} + self._status = {} + self._source_model = {} + self._source_key = {} + self._source_account = {} + self._source_account_clear = {} + self._source_refresh_buttons = {} + self._import_buttons = {} + self._chrome_buttons = {} + self._chrome_logins = {} + self._chrome_login_threads = {} + self._chrome_login_lock = threading.Lock() + self._closing_logins = False + return self + + def show_(self, _sender): + if self.window is None: + self._build() + self._reload_controls() + try: + NSApp.activateIgnoringOtherApps_(True) + except Exception: + pass + self.window.makeKeyAndOrderFront_(None) + + def _build(self): + mask = ( + NSWindowStyleMaskTitled + | NSWindowStyleMaskClosable + | NSWindowStyleMaskMiniaturizable + ) + self.window = NSWindow.alloc().initWithContentRect_styleMask_backing_defer_( + NSMakeRect(0, 0, W, H), mask, NSBackingStoreBuffered, False + ) + self.window.setTitle_("AgentBar 额度来源设置") + self.window.setReleasedWhenClosed_(False) + self.window.center() + view = self.window.contentView() + + view.addSubview_(_label("额度来源设置", PAD, H - 42, 220, 22, bold=True)) + view.addSubview_(_label( + "默认关闭周期请求；保存时只刷新有变更的来源，也可单独点“刷新此来源”。", + PAD, H - 65, W - 2 * PAD, 18, dim=True, + )) + + view.addSubview_(_label("菜单栏数字", PAD, H - 103, 90)) + self.title_popup = NSPopUpButton.alloc().initWithFrame_pullsDown_( + NSMakeRect(PAD + 92, H - 109, 180, 27), False + ) + self.title_popup.addItemsWithTitles_([f"标题：{name}" for name, _ in _TITLE_OPTIONS]) + view.addSubview_(self.title_popup) + + self.auto_refresh_check = NSButton.alloc().initWithFrame_( + NSMakeRect(PAD + 290, H - 106, 255, 24) + ) + self.auto_refresh_check.setButtonType_(3) # NSSwitchButton + self.auto_refresh_check.setTitle_("周期自动刷新（默认关闭）") + view.addSubview_(self.auto_refresh_check) + + view.addSubview_(_label("订阅额度", PAD, H - 142, 220, 20, bold=True)) + view.addSubview_(_label( + "显式输入 OAuth Access Token；不会读取 Keychain 或 CLI 登录文件。", + PAD + 84, H - 142, W - 2 * PAD - 84, 18, dim=True, + )) + self._build_quota_source_row(view, "claude", H - 235) + self._build_quota_source_row(view, "codex", H - 345) + + view.addSubview_(_label("内部额度", PAD, H - 376, 220, 20, bold=True)) + view.addSubview_(_label( + "MyToken / Tokenverse 使用企业 SSO Cookie；浏览器登录和导入都只在显式点击后执行。", + PAD + 84, H - 376, W - 2 * PAD - 84, 18, dim=True, + )) + self._build_provider_row(view, "mytoken", H - 525) + self._build_provider_row(view, "tokenverse", H - 670) + + self._message = _label("", PAD, 23, W - 200, 20, dim=True) + view.addSubview_(self._message) + save = _button("保存配置", W - PAD - 176, 18, 176, self, "onSave:", 30) + save.setKeyEquivalent_("\r") + view.addSubview_(save) + + def _build_quota_source_row(self, view, source: str, y: int): + name = _NAMES[source] + enabled = NSButton.alloc().initWithFrame_(NSMakeRect(PAD, y + 72, 120, 24)) + enabled.setButtonType_(3) # NSSwitchButton + enabled.setTitle_(name) + enabled.setFont_(NSFont.boldSystemFontOfSize_(13)) + view.addSubview_(enabled) + self._enabled[source] = enabled + + status = _label("", PAD + 126, y + 74, W - PAD * 2 - 126, 20, dim=True) + view.addSubview_(status) + self._status[source] = status + + model_label = "额度模型" if source == "claude" else "metered_feature" + view.addSubview_(_label(model_label, PAD, y + 41, 105)) + model = NSTextField.alloc().initWithFrame_(NSMakeRect(PAD + 108, y + 37, 245, 24)) + model.setPlaceholderString_( + "opus / sonnet；留空=通用窗口" + if source == "claude" + else "留空=账户总额度；或填 metered_feature" + ) + view.addSubview_(model) + self._source_model[source] = model + + if source == "codex": + view.addSubview_(_label("Account ID", PAD + 370, y + 41, 78)) + account = NSTextField.alloc().initWithFrame_( + NSMakeRect(PAD + 450, y + 37, W - PAD * 2 - 450 - 82, 24) + ) + account.setPlaceholderString_("可选；留空保留已配置值") + view.addSubview_(account) + self._source_account[source] = account + clear_account = _button( + "清 Account", W - PAD - 78, y + 36, 78, + self, "onClearAccount:", 26, + ) + clear_account.setRepresentedObject_(source) + view.addSubview_(clear_account) + self._source_account_clear[source] = clear_account + + view.addSubview_(_label("OAuth Access Token", PAD, y + 9, 126)) + key = NSSecureTextField.alloc().initWithFrame_( + NSMakeRect(PAD + 130, y + 5, W - 2 * PAD - 130 - 174, 24) + ) + key.setPlaceholderString_("留空保留已配置的 Access Token") + view.addSubview_(key) + self._source_key[source] = key + + refresh = _button("刷新此来源", W - PAD - 166, y + 4, 92, self, "onRefreshSource:", 26) + refresh.setRepresentedObject_(source) + view.addSubview_(refresh) + self._source_refresh_buttons[source] = refresh + + clear = _button("清空", W - PAD - 68, y + 4, 68, self, "onClearSource:", 26) + clear.setRepresentedObject_(source) + view.addSubview_(clear) + + def _build_provider_row(self, view, provider: str, y: int): + name = _NAMES[provider] + enabled = NSButton.alloc().initWithFrame_(NSMakeRect(PAD, y + 92, 150, 24)) + enabled.setButtonType_(3) # NSSwitchButton + enabled.setTitle_(name) + enabled.setFont_(NSFont.boldSystemFontOfSize_(13)) + view.addSubview_(enabled) + self._enabled[provider] = enabled + + status = _label("", PAD + 158, y + 94, W - PAD * 2 - 158, 20, dim=True) + view.addSubview_(status) + self._status[provider] = status + + view.addSubview_(_label("展示单位", PAD, y + 57, 70)) + unit = NSPopUpButton.alloc().initWithFrame_pullsDown_( + NSMakeRect(PAD + 72, y + 52, 110, 27), False + ) + unit.addItemsWithTitles_(list(PROVIDER_UNITS)) + view.addSubview_(unit) + self._unit[provider] = unit + + view.addSubview_(_label("刷新秒数", PAD + 198, y + 57, 70)) + refresh = NSTextField.alloc().initWithFrame_(NSMakeRect(PAD + 270, y + 53, 76, 24)) + refresh.setPlaceholderString_("300") + view.addSubview_(refresh) + self._refresh[provider] = refresh + + login = _button("🌐 浏览器登录", PAD + 365, y + 51, 122, self, "onChromeLogin:") + login.setRepresentedObject_(provider) + view.addSubview_(login) + self._chrome_buttons[provider] = login + + import_btn = _button("读取已有登录", PAD + 495, y + 51, 122, self, "onImport:") + import_btn.setRepresentedObject_(provider) + view.addSubview_(import_btn) + self._import_buttons[provider] = import_btn + + view.addSubview_(_label("Cookie", PAD, y + 20, 55)) + cookie = NSSecureTextField.alloc().initWithFrame_( + NSMakeRect(PAD + 55, y + 16, W - 2 * PAD - 55 - 90, 24) + ) + cookie.setPlaceholderString_("可选：手动粘贴 Request Headers 的 Cookie；留空不修改") + view.addSubview_(cookie) + self._cookie[provider] = cookie + + clear = _button("清空", W - PAD - 82, y + 15, 82, self, "onClear:", 26) + clear.setRepresentedObject_(provider) + view.addSubview_(clear) + + @objc.python_method + def _reload_controls(self): + # Copy one coherent settings revision, then release the lock before + # touching AppKit controls. Writers replace/mutate these dictionaries + # under the same lock. + with self.settings._lock: + title_provider = self.settings.title_provider + usage_auto_refresh = self.settings.usage_auto_refresh + quota_sources = { + source: dict((self.settings.quota_sources or {}).get(source) or {}) + for source in _QUOTA_SOURCES + } + providers = { + provider: dict((self.settings.providers or {}).get(provider) or {}) + for provider in _PROVIDERS + } + title_values = [value for _, value in _TITLE_OPTIONS] + try: + self.title_popup.selectItemAtIndex_(title_values.index(title_provider)) + except ValueError: + self.title_popup.selectItemAtIndex_(1) + + self.auto_refresh_check.setState_(1 if usage_auto_refresh else 0) + for source in _QUOTA_SOURCES: + defaults = DEFAULT_QUOTA_SOURCES[source] + cfg = quota_sources[source] or defaults + self._enabled[source].setState_(1 if cfg.get("enabled") else 0) + self._source_model[source].setStringValue_(str(cfg.get("model") or "")) + # Credentials are write-only. Empty controls mean "keep existing" on save. + self._source_key[source].setStringValue_("") + account = self._source_account.get(source) + if account is not None: + account.setStringValue_("") + key_set = bool(str(cfg.get("access_token") or "").strip()) + if key_set: + state = "OAuth Access Token 已配置" + if source == "codex" and str(cfg.get("account_id") or "").strip(): + state += " · Account ID 已配置" + if not cfg.get("enabled"): + state += " · 已停用" + elif cfg.get("enabled"): + state = "已启用，但缺少 OAuth Access Token" + else: + state = "未配置 OAuth Access Token" + self._status[source].setStringValue_(state) + self._source_refresh_buttons[source].setEnabled_( + bool(cfg.get("enabled") and key_set) + ) + account_clear = self._source_account_clear.get(source) + if account_clear is not None: + account_clear.setEnabled_(bool(str(cfg.get("account_id") or "").strip())) + + for provider in _PROVIDERS: + defaults = DEFAULT_PROVIDERS[provider] + cfg = providers[provider] or defaults + self._enabled[provider].setState_(1 if cfg.get("enabled") else 0) + unit = cfg.get("unit") if cfg.get("unit") in PROVIDER_UNITS else defaults["unit"] + self._unit[provider].selectItemWithTitle_(unit) + self._refresh[provider].setStringValue_( + str(int(cfg.get("refresh_seconds") or defaults["refresh_seconds"])) + ) + self._cookie[provider].setStringValue_("") + cookie = str(cfg.get("cookie") or "") + if cookie: + # Cookie is write-only too. Even the text before the first '=' + # can be a pasted JWT/Bearer value rather than a safe cookie name, + # so the native status never echoes any portion of the secret. + state = "Cookie 已配置" + elif cfg.get("enabled"): + state = "已启用，但缺少 Cookie" + else: + state = "未配置 Cookie" + self._status[provider].setStringValue_(state) + + def onRefreshSource_(self, sender): + source = str(sender.representedObject() or "") + with self.settings._lock: + cfg = dict((self.settings.quota_sources or {}).get(source) or {}) + if source not in _QUOTA_SOURCES or not cfg.get("enabled") or not str( + cfg.get("access_token") or "" + ).strip(): + self._message.setStringValue_("请先保存并启用该来源的 OAuth Access Token。") + return + self.core.quota.refresh_now(source) + self._message.setStringValue_(f"{_NAMES[source]}：已触发一次手动额度刷新。") + + def onClearSource_(self, sender): + source = str(sender.representedObject() or "") + if source not in DEFAULT_QUOTA_SOURCES: + return + alert = NSAlert.alloc().init() + alert.setMessageText_(f"清空 {_NAMES[source]} OAuth Access Token？") + alert.setInformativeText_("将删除 AgentBar 保存的 Access Token 和 Account ID，并停用该额度来源。") + alert.addButtonWithTitle_("清空") + alert.addButtonWithTitle_("取消") + if alert.runModal() != NSAlertFirstButtonReturn: + return + try: + with self.settings._lock: + old_sources = copy.deepcopy(self.settings.quota_sources) + try: + cfg = self.settings.quota_sources.setdefault( + source, dict(DEFAULT_QUOTA_SOURCES[source]) + ) + cfg["access_token"] = "" + cfg["account_id"] = "" + cfg["enabled"] = False + save_settings(self.settings) + except Exception: + self.settings.quota_sources = old_sources + raise + except Exception: + log.exception("failed to clear %s quota credential", source) + self._reload_controls() + self._alert("无法清空额度凭据", "配置文件写入失败，原设置未更改。") + return + # Clearing disables the source; rebuild without issuing unrelated + # requests. There is nothing useful to refresh for the cleared source. + self.core.quota.reload_fetchers(refresh=False) + self._reload_controls() + self._message.setStringValue_(f"{_NAMES[source]} OAuth Access Token 已清空并停用。") + + def onClearAccount_(self, sender): + source = str(sender.representedObject() or "") + if source != "codex": + return + alert = NSAlert.alloc().init() + alert.setMessageText_("清空 Codex Account ID？") + alert.setInformativeText_("OAuth Access Token 和额度来源开关保持不变。") + alert.addButtonWithTitle_("清空") + alert.addButtonWithTitle_("取消") + if alert.runModal() != NSAlertFirstButtonReturn: + return + try: + with self.settings._lock: + old_sources = copy.deepcopy(self.settings.quota_sources) + try: + cfg = self.settings.quota_sources.setdefault( + source, dict(DEFAULT_QUOTA_SOURCES[source]) + ) + cfg["account_id"] = "" + save_settings(self.settings) + ready = bool(cfg.get("enabled") and cfg.get("access_token")) + except Exception: + self.settings.quota_sources = old_sources + raise + except Exception: + log.exception("failed to clear Codex account id") + self._reload_controls() + self._alert("Account ID 清空失败", "配置文件写入失败，原设置未更改。") + return + self.core.quota.reload_fetchers(refresh=False) + if ready: + self.core.quota.refresh_now(source) + self._reload_controls() + self._message.setStringValue_("Codex Account ID 已清空；Access Token 保持不变。") + + def onChromeLogin_(self, sender): + provider = str(sender.representedObject() or "") + with self._chrome_login_lock: + if self._closing_logins: + return + active = self._chrome_logins.get(provider) + if active is not None: + active.cancel() + self._status[provider].setStringValue_("正在取消浏览器登录…") + return + try: + login = ChromeCDPLogin(provider) + except ChromeLoginError as exc: + self._alert("无法启动浏览器登录", str(exc)) + return + sender.setTitle_("取消登录") + self._status[provider].setStringValue_("正在启动独立 Chrome 登录窗口…") + + def update_status(message): + AppHelper.callAfter(self._set_provider_status, provider, message) + + def work(): + try: + imported = login.run(on_status=update_status) + error = "" + except ChromeLoginError as exc: + imported = None + error = str(exc) + except Exception as exc: + log.exception("Chrome CDP login failed") + imported = None + error = str(exc) + AppHelper.callAfter(self._finish_chrome_login, provider, imported, error) + + worker = threading.Thread( + target=work, + name=f"agentbar-chrome-login-{provider}", + daemon=True, + ) + # Publish and start atomically with respect to close_active_logins(). A + # quit can therefore either miss this attempt entirely (before the click) + # or cancel and join a real started thread, never an unstarted object. + with self._chrome_login_lock: + if self._closing_logins: + rejected = True + else: + rejected = False + self._chrome_logins[provider] = login + self._chrome_login_threads[provider] = worker + worker.start() + if rejected: + login.cancel() + login.close() + + @objc.python_method + def close_active_logins(self): + """Cancel active SSO windows and synchronously reclaim Chrome/profile data.""" + with self._chrome_login_lock: + self._closing_logins = True + active = list(self._chrome_logins.items()) + threads = dict(self._chrome_login_threads) + self._chrome_logins.clear() + self._chrome_login_threads.clear() + for _provider, login in active: + try: + login.cancel() + login.close() + except Exception: + log.exception("failed to close active Chrome login") + for provider, thread in threads.items(): + if thread is threading.current_thread(): + continue + try: + thread.join(timeout=8) + if thread.is_alive(): + log.warning("Chrome login worker %s did not stop within 8s", provider) + except RuntimeError: + # Defensive for injected/test workers; production threads are + # started before publication under _chrome_login_lock. + pass + + @objc.python_method + def _set_provider_status(self, provider, message): + self._status[provider].setStringValue_(message) + + @objc.python_method + def _finish_chrome_login(self, provider, imported, error): + with self._chrome_login_lock: + self._chrome_logins.pop(provider, None) + self._chrome_login_threads.pop(provider, None) + closing = self._closing_logins + # A worker may queue this callback just before shutdown cancels it. + # Once close_active_logins() has linearized, never touch AppKit controls + # or persist a late credential from the retired login attempt. + if closing: + return + self._chrome_buttons[provider].setTitle_("🌐 浏览器登录") + if error or imported is None: + message = error or "未捕获到 Cookie" + self._status[provider].setStringValue_(f"浏览器登录失败：{message}") + if "已取消" not in message: + self._alert("浏览器登录失败", message) + return + self._save_imported_cookie(provider, imported) + + def onImport_(self, sender): + provider = str(sender.representedObject() or "") + host = PROVIDER_HOSTS.get(provider) + if not host: + return + sender.setEnabled_(False) + self._status[provider].setStringValue_("正在读取 Chrome / Edge / Brave 已有登录态…") + + def work(): + try: + imported = import_cookie_header(host) + error = "" + except CookieImportError as exc: + imported = None + error = str(exc) + except Exception as exc: # Keychain/SQLite 失败不能带崩 AppKit 主线程 + log.exception("provider cookie import failed") + imported = None + error = str(exc) + AppHelper.callAfter(self._finish_import, provider, imported, error) + + threading.Thread( + target=work, + name=f"agentbar-cookie-{provider}", + daemon=True, + ).start() + + @objc.python_method + def _finish_import(self, provider, imported, error): + self._import_buttons[provider].setEnabled_(True) + if error or imported is None: + message = error or "未导入到 Cookie" + self._status[provider].setStringValue_(f"导入失败：{message}") + self._alert("Cookie 导入失败", message) + return + self._save_imported_cookie(provider, imported) + + @objc.python_method + def _save_imported_cookie(self, provider, imported): + try: + with self.settings._lock: + old_providers = copy.deepcopy(self.settings.providers) + try: + cfg = self.settings.providers.setdefault( + provider, dict(DEFAULT_PROVIDERS[provider]) + ) + cfg["enabled"] = True + cfg["cookie"] = imported.header + save_settings(self.settings) + except Exception: + self.settings.providers = old_providers + raise + except Exception: + log.exception("failed to save imported %s cookie", provider) + self._reload_controls() + self._alert("Cookie 保存失败", "配置文件写入失败，原设置未更改。") + return + self.core.quota.reload_fetchers(refresh=False) + self.core.quota.refresh_now(provider) + self._reload_controls() + self._message.setStringValue_( + f"{_NAMES[provider]}：已从 {imported.source} 导入 {imported.count} 个 Cookie，并触发刷新" + ) + + def onSave_(self, _sender): + # Read and validate control values before taking the settings lock; AppKit + # calls must never sit inside a cross-thread configuration transaction. + source_inputs = {} + provider_inputs = {} + try: + for source in _QUOTA_SOURCES: + enabled = bool(self._enabled[source].state()) + model = self._validated_field( + self._source_model[source], "额度模型", 160 + ) + typed_key = self._validated_field( + self._source_key[source], "OAuth Access Token", 16_384 + ) + account = "" + if source == "codex": + account = self._validated_field( + self._source_account[source], "Codex Account ID", 256 + ) + source_inputs[source] = { + "enabled": enabled, + "model": model, + "typed_key": typed_key, + "account_id": account, + } + for provider in _PROVIDERS: + defaults = DEFAULT_PROVIDERS[provider] + unit = str( + self._unit[provider].titleOfSelectedItem() or defaults["unit"] + ) + try: + refresh_seconds = max( + 60, + int(str(self._refresh[provider].stringValue()).strip() or 300), + ) + except ValueError: + refresh_seconds = defaults["refresh_seconds"] + provider_inputs[provider] = { + "enabled": bool(self._enabled[provider].state()), + "unit": unit if unit in PROVIDER_UNITS else defaults["unit"], + "refresh_seconds": refresh_seconds, + "cookie": str(self._cookie[provider].stringValue()).strip(), + } + except ValueError as exc: + self._alert("无法保存额度设置", str(exc)) + return + + title_index = self.title_popup.indexOfSelectedItem() + selected_title = ( + _TITLE_OPTIONS[title_index][1] + if 0 <= title_index < len(_TITLE_OPTIONS) + else None + ) + selected_auto_refresh = bool(self.auto_refresh_check.state()) + + try: + # One settings transaction covers credential preservation, mutation, + # persistence and change detection. A concurrent Web PATCH can run + # wholly before or after this transaction, never interleave with it. + with self.settings._lock: + old_sources = copy.deepcopy(self.settings.quota_sources) + old_providers = copy.deepcopy(self.settings.providers) + old_title = self.settings.title_provider + old_auto_refresh = self.settings.usage_auto_refresh + before_sources = { + source: dict((old_sources or {}).get(source) or {}) + for source in _QUOTA_SOURCES + } + before_providers = { + provider: dict((old_providers or {}).get(provider) or {}) + for provider in _PROVIDERS + } + try: + for source, update in source_inputs.items(): + current = (self.settings.quota_sources or {}).get(source) or {} + effective_key = update["typed_key"] or str( + current.get("access_token") or "" + ).strip() + if update["enabled"] and not effective_key: + raise ValueError( + f"{_NAMES[source]} 启用前必须输入 OAuth Access Token" + ) + + for source, update in source_inputs.items(): + cfg = self.settings.quota_sources.setdefault( + source, dict(DEFAULT_QUOTA_SOURCES[source]) + ) + cfg["enabled"] = update["enabled"] + cfg["model"] = update["model"] + # Empty credential controls preserve the saved write-only value. + if update["typed_key"]: + cfg["access_token"] = update["typed_key"] + if source == "codex" and update["account_id"]: + cfg["account_id"] = update["account_id"] + + for provider, update in provider_inputs.items(): + cfg = self.settings.providers.setdefault( + provider, dict(DEFAULT_PROVIDERS[provider]) + ) + cfg["enabled"] = update["enabled"] + cfg["unit"] = update["unit"] + cfg["refresh_seconds"] = update["refresh_seconds"] + if update["cookie"]: + cfg["cookie"] = update["cookie"] + + if selected_title is not None: + self.settings.title_provider = selected_title + self.settings.usage_auto_refresh = selected_auto_refresh + save_settings(self.settings) + except Exception: + self.settings.quota_sources = old_sources + self.settings.providers = old_providers + self.settings.title_provider = old_title + self.settings.usage_auto_refresh = old_auto_refresh + raise + + changed = [ + source for source in _QUOTA_SOURCES + if before_sources[source] + != dict((self.settings.quota_sources or {}).get(source) or {}) + ] + changed.extend( + provider for provider in _PROVIDERS + if before_providers[provider] + != dict((self.settings.providers or {}).get(provider) or {}) + ) + ready_tools = [] + for tool in changed: + if tool in _QUOTA_SOURCES: + cfg = (self.settings.quota_sources or {}).get(tool) or {} + ready = cfg.get("enabled") and str( + cfg.get("access_token") or "" + ).strip() + else: + cfg = (self.settings.providers or {}).get(tool) or {} + ready = cfg.get("enabled") and str( + cfg.get("cookie") or "" + ).strip() + if ready: + ready_tools.append(tool) + auto_refresh_enabled = self.settings.usage_auto_refresh + auto_refresh_changed = old_auto_refresh != auto_refresh_enabled + except ValueError as exc: + self._alert("无法保存额度设置", str(exc)) + return + except Exception: + log.exception("failed to save quota settings") + self._reload_controls() + self._alert("无法保存额度设置", "配置文件写入失败，原设置未更改。") + return + + refreshed = [] + if changed or auto_refresh_changed: + self.core.quota.reload_fetchers(refresh=False) + if changed: + for tool in ready_tools: + self.core.quota.refresh_now(tool) + refreshed.append(_NAMES[tool]) + self._reload_controls() + mode = "已开启周期自动刷新" if auto_refresh_enabled else "之后仅手动刷新" + refresh_note = ( + f"已刷新变更来源：{'、'.join(refreshed)}" + if refreshed else "没有来源需要刷新" + ) + self._message.setStringValue_(f"配置已保存；{refresh_note}；{mode}。") + + def onClear_(self, sender): + provider = str(sender.representedObject() or "") + if provider not in DEFAULT_PROVIDERS: + return + alert = NSAlert.alloc().init() + alert.setMessageText_(f"清空 {_NAMES[provider]} 登录态？") + alert.setInformativeText_("将删除 AgentBar 保存的 Cookie 并停用该额度来源。") + alert.addButtonWithTitle_("清空") + alert.addButtonWithTitle_("取消") + if alert.runModal() != NSAlertFirstButtonReturn: + return + try: + with self.settings._lock: + old_providers = copy.deepcopy(self.settings.providers) + try: + cfg = self.settings.providers.setdefault( + provider, dict(DEFAULT_PROVIDERS[provider]) + ) + cfg["cookie"] = "" + cfg["enabled"] = False + save_settings(self.settings) + except Exception: + self.settings.providers = old_providers + raise + except Exception: + log.exception("failed to clear %s cookie", provider) + self._reload_controls() + self._alert("无法清空 Cookie", "配置文件写入失败，原设置未更改。") + return + self.core.quota.reload_fetchers(refresh=False) + self._reload_controls() + self._message.setStringValue_(f"{_NAMES[provider]} Cookie 已清空并停用。") + + @objc.python_method + def _validated_field(self, control, label, max_len): + value = str(control.stringValue() or "").strip() + if len(value) > max_len: + raise ValueError(f"{label} 过长") + if any(char in value for char in "\r\n\x00"): + raise ValueError(f"{label} 不能包含换行或 NUL") + return value + + @objc.python_method + def _alert(self, title, text): + alert = NSAlert.alloc().init() + alert.setMessageText_(title) + alert.setInformativeText_(text) + alert.runModal() diff --git a/agentbar/quota.py b/agentbar/quota.py index 1df08f0..9151858 100644 --- a/agentbar/quota.py +++ b/agentbar/quota.py @@ -1,8 +1,8 @@ """Quota status — honest, layered; never fabricated. 数据优先级（来源在 UI 明确标注，取不到就降级，绝不编造）： - 1. usage API — aiusagebar 同款官方接口（Claude OAuth usage / Codex wham usage）， - 后台定时刷新（默认 120s），给出各窗口用量百分比 + 重置时间 + 1. usage API — 用户显式配置的 Claude OAuth usage / Codex wham usage； + 默认仅启动、保存设置或手动点击时请求，周期刷新需主动开启 2. observed — 调度器观测事实：真实任务的限流失败（ground truth，优先于 API 展示） 3. ccusage — 本机装了 ccusage 时补充 5h 窗口成本 4. unknown — 都取不到，如实显示未知 @@ -12,6 +12,8 @@ import json import logging +import math +import re import shutil import subprocess import threading @@ -24,6 +26,10 @@ log = logging.getLogger("agentbar.quota") USAGE_STALE_SECONDS = 15 * 60 # usage API 结果超过该时长视为过期，不再参与判定 +# MyToken's token-unit fetch can perform three sequential 12s HTTP calls. Give +# the active request enough time to reach its own bounded timeout before stop() +# reports a lifecycle failure. +STOP_JOIN_SECONDS = 45 @dataclass @@ -37,6 +43,9 @@ class QuotaStatus: plan: str | None = None fetched_at: float | None = None error: str | None = None + model: str | None = None + available_models: list[str] = field(default_factory=list) + stale: bool = False def to_dict(self) -> dict: return { @@ -49,9 +58,20 @@ def to_dict(self) -> dict: "plan": self.plan, "fetched_at": self.fetched_at, "error": self.error, + "model": self.model, + "available_models": self.available_models, + "stale": self.stale, } +@dataclass(frozen=True) +class _RefreshBatch: + """An immutable claim on fetchers from one configuration generation.""" + + generation: int + fetchers: dict[str, object] + + def _clock(ts: float) -> str: return time.strftime("%H:%M", time.localtime(ts)) @@ -68,36 +88,125 @@ def __init__(self, settings: Settings): self._lock = threading.Lock() # {tool: {"last_success_at": float, "last_quota_at": float, "reset_at": float}} self._obs: dict[str, dict] = {} + # _usage 只保存最后一次成功结果；刷新错误单独保存，不能用一次临时 + # 网络失败覆盖仍在有效期内的 last-good 数据。 self._usage: dict[str, UsageSnapshot] = {} + self._usage_errors: dict[str, UsageSnapshot] = {} self._fetchers = get_usage_fetchers(settings) self._next_due: dict[str, float] = {} # 各 provider 的下次到期刷新时刻（per-provider 间隔） - self._force_refresh = False # refresh_now 触发时强制刷新全部 + self._scheduled_intervals: dict[str, float] = {} + self._pending_refresh: set[str] = set() # 手动刷新来源；Event + set 天然合并快速重复点击 + # tool -> fetch generation。旧请求结束时只能清理自己的 generation， + # 不能误删 reload 后同名来源的新请求标记。 + # tool -> 当前正在请求的 fetcher 对象。以对象身份而不是全局 + # generation 判重，使 reload 一个来源时不打断其他未变化来源的在途请求。 + self._refreshing: dict[str, object] = {} + self._fetch_generation = 0 self._ccusage: dict | None = None self._ccusage_bin = shutil.which("ccusage") self._stop = threading.Event() self._refresh_evt = threading.Event() self._thread: threading.Thread | None = None + self._lifecycle_lock = threading.Lock() def provider_tools(self) -> list[str]: """本监视器能上报额度的全部工具名（含 corp provider）。""" - return list(self._fetchers) + with self._lock: + return list(self._fetchers) - def reload_fetchers(self) -> None: - """Rebuild usage fetchers after settings.providers changes.""" + def observed_tools(self) -> list[str]: + """已有任务成功/限额观测的工具（无需配置上游额度来源）。""" + with self._lock: + return list(self._obs) + + @staticmethod + def _fetcher_signature(fetcher: object) -> tuple: + """Configuration identity without depending on a fetcher's implementation. + + Production fetchers expose one or more of these immutable constructor + fields. Unknown/testing fetchers fall back to object identity so replacing + one is conservatively treated as a real configuration change. + """ + names = ( + "access_token", "api_key", "account_id", "model", + "cookie", "unit", "refresh_seconds", + ) + values = tuple((name, getattr(fetcher, name)) for name in names if hasattr(fetcher, name)) + return (type(fetcher), values if values else id(fetcher)) + + def reload_fetchers(self, refresh: bool = True) -> None: + """Rebuild only changed sources and optionally refresh those sources. + + Saving an unrelated UI preference must not clear every displayed quota or + cause every configured provider to make another network request. + """ + # A settings writer mutates credentials and persists them under this + # lock. Never build fetchers from another thread's uncommitted revision: + # if that save later rolls back, the monitor could otherwise keep using a + # credential/model that the UI reported as not saved. + with self.settings._lock: + rebuilt = get_usage_fetchers(self.settings) + wake = False with self._lock: - self._fetchers = get_usage_fetchers(self.settings) - self._next_due = {} - self._usage = { - tool: snap for tool, snap in self._usage.items() - if tool in self._fetchers + previous = self._fetchers + all_tools = set(previous) | set(rebuilt) + changed = { + name for name in all_tools + if name not in previous + or name not in rebuilt + or self._fetcher_signature(previous[name]) + != self._fetcher_signature(rebuilt[name]) } - self.refresh_now() + if not changed: + # A settings save may only toggle automatic refresh or change + # the global interval, neither of which changes fetcher identity. + # Wake the loop so the new schedule takes effect immediately. + wake = not self._stop.is_set() + else: + self._fetch_generation += 1 + # Preserve the exact old object for unchanged sources. An in-flight + # result is therefore still valid, while a changed source's old result + # is rejected by the identity checks in _run_refresh_batch. + for name in set(previous) & set(rebuilt) - changed: + rebuilt[name] = previous[name] + self._fetchers = rebuilt + + for name in changed: + self._next_due.pop(name, None) + self._scheduled_intervals.pop(name, None) + self._pending_refresh.discard(name) + self._usage.pop(name, None) + self._usage_errors.pop(name, None) + if "claude" in changed: + self._ccusage = None + if refresh and not self._stop.is_set(): + self._pending_refresh.update(changed & set(rebuilt)) + # Even refresh=False must recompute the auto schedule (for + # example after disabling a source). + wake = not self._stop.is_set() + if wake: + self._refresh_evt.set() # ---------- persistence (由 scheduler 存进 state.json) ---------- def load(self, data: dict) -> None: + cleaned: dict[str, dict] = {} + if isinstance(data, dict): + for tool, raw in data.items(): + if not isinstance(tool, str) or not isinstance(raw, dict): + continue + entry = {} + for key in ("last_success_at", "last_quota_at", "reset_at"): + value = raw.get(key) + if isinstance(value, bool) or not isinstance(value, (int, float)): + continue + parsed = float(value) + if math.isfinite(parsed) and parsed >= 0: + entry[key] = parsed + if entry: + cleaned[tool] = entry with self._lock: - self._obs = dict(data or {}) + self._obs = cleaned def dump(self) -> dict: with self._lock: @@ -110,7 +219,6 @@ def record_success(self, tool: str) -> None: o = self._obs.setdefault(tool, {}) o["last_success_at"] = time.time() o.pop("reset_at", None) - self.refresh_now() def record_quota(self, tool: str, reset_at: float | None) -> None: with self._lock: @@ -118,11 +226,20 @@ def record_quota(self, tool: str, reset_at: float | None) -> None: o["last_quota_at"] = time.time() if reset_at: o["reset_at"] = reset_at - self.refresh_now() # ---------- cooldown（调度器据此暂缓派发该工具的任务） ---------- - def cooldown_until(self, tool: str) -> float | None: + @staticmethod + def _same_model(configured: str | None, task_model: str | None) -> bool: + if not configured: + return True # 账户级额度适用于该工具的全部模型 + if not task_model: + return False # 任务用 CLI 默认模型，无法证明它命中已选专属额度 + left = re.sub(r"[^a-z0-9]+", "", configured.casefold()) + right = re.sub(r"[^a-z0-9]+", "", task_model.casefold()) + return bool(left and right and (left == right or left in right or right in left)) + + def cooldown_until(self, tool: str, model: str | None = None) -> float | None: now = time.time() candidates: list[float] = [] with self._lock: @@ -134,7 +251,11 @@ def cooldown_until(self, tool: str) -> float | None: # usage API 显示某窗口已打满 → 主动冷却到重置时间（不用真跑一次失败） if snap and not snap.error and now - snap.fetched_at < USAGE_STALE_SECONDS: for w in snap.windows: - if (snap.limited or w.used_percent >= 99.9) and w.resets_at and w.resets_at > now: + # 一个 snapshot 可能同时含账户通用窗口和多个模型专属窗口。 + # 必须逐窗口匹配，不能用 UI 所选的 snapshot.model 把通用窗口一并过滤。 + if not self._same_model(w.model, model): + continue + if (w.limited or w.used_percent >= 99.9) and w.resets_at and w.resets_at > now: candidates.append(w.resets_at) future = [c for c in candidates if c > now] return max(future) if future else None @@ -145,7 +266,12 @@ def status(self, tool: str) -> QuotaStatus: with self._lock: o = dict(self._obs.get(tool) or {}) snap = self._usage.get(tool) + refresh_error = self._usage_errors.get(tool) cc = dict(self._ccusage) if (self._ccusage and tool == "claude") else None + # Compatibility for callers/tests that injected an error snapshot directly + # into _usage before errors were split from last-good data. + if snap and snap.error: + refresh_error, snap = snap, None now = time.time() lq, ls, obs_reset = o.get("last_quota_at"), o.get("last_success_at"), o.get("reset_at") observed_limited = bool(lq and (not ls or lq > ls)) @@ -156,11 +282,15 @@ def status(self, tool: str) -> QuotaStatus: worst_reset = None for w in snap.windows[:3]: seg = f"{w.label} {w.used_percent:.0f}%" - if w.used_percent >= 99.9 and w.resets_at: + if (w.limited or w.used_percent >= 99.9) and w.resets_at: seg += f"（{_clock_day(w.resets_at)} 重置）" worst_reset = max(worst_reset or 0, w.resets_at) parts.append(seg) - limited = observed_limited or snap.limited or any(w.used_percent >= 99.9 for w in snap.windows) + # snapshot.limited 只保留给旧 fetcher / 展示兼容；新 fetcher 在窗口 + # 上标注 limited，供模型级 cooldown 精确判定。 + limited = observed_limited or snap.limited or any( + w.limited or w.used_percent >= 99.9 for w in snap.windows + ) st = QuotaStatus( tool, "limited" if limited else "ok", @@ -184,9 +314,25 @@ def status(self, tool: str) -> QuotaStatus: st.windows = [w.to_dict() for w in snap.windows] st.plan = snap.plan st.fetched_at = snap.fetched_at - st.error = snap.error - if snap.error and st.source in ("none",): - st.detail += f"；usage API: {snap.error}" + st.model = snap.model + st.available_models = list(snap.available_models) + st.stale = now - snap.fetched_at >= USAGE_STALE_SECONDS + if st.stale and st.source == "none": + # Manual mode deliberately keeps the last-known windows visible, + # but expired data must never look current or drive cooldown. + st.detail = "上次额度数据已过期，请手动刷新" + metadata = snap or refresh_error + if metadata and not snap: + st.plan = metadata.plan + st.fetched_at = metadata.fetched_at + st.model = metadata.model + st.available_models = list(metadata.available_models) + if refresh_error and refresh_error.error: + st.error = refresh_error.error + if snap and snap.windows: + st.detail += f"；上次刷新失败：{refresh_error.error}" + elif st.source == "none": + st.detail += f"；usage API: {refresh_error.error}" if cc: st.detail += f"；5h 已用 ${cc['cost']:.2f}（ccusage）" st.source += "+ccusage" @@ -195,73 +341,219 @@ def status(self, tool: str) -> QuotaStatus: # ---------- background refresh ---------- def start_background(self) -> None: - self._thread = threading.Thread( - target=self._loop, name="agentbar-usage", daemon=True - ) - self._thread.start() + with self._lifecycle_lock: + if self._stop.is_set(): + return + if self._thread and self._thread.is_alive(): + return + self._thread = threading.Thread( + target=self._loop, name="agentbar-usage", daemon=True + ) + self._thread.start() def stop(self) -> None: self._stop.set() self._refresh_evt.set() - - def refresh_now(self) -> None: - """异步触发一次立即刷新（任务结束/用户点菜单时调用）——强制刷新全部 provider。""" - self._force_refresh = True - self._refresh_evt.set() - - def authorize_claude_keychain(self) -> bool: - """交互式读取 Keychain（允许系统弹窗；用户点"始终允许"后静默读取即长期可用）。""" - fetcher = self._fetchers.get("claude") - if not fetcher: - return False - snap = fetcher.fetch(interactive=True) - ok = bool(snap and not snap.error) - if snap: - with self._lock: - self._usage["claude"] = snap - return ok + with self._lifecycle_lock: + thread = self._thread + if thread and thread is not threading.current_thread(): + thread.join(timeout=STOP_JOIN_SECONDS) + if thread.is_alive(): + log.warning("usage monitor did not stop within %ss", STOP_JOIN_SECONDS) + + def refresh_now(self, tool: str | None = None) -> None: + """异步刷新已配置来源。快速重复请求会在 pending set 中合并。""" + if self._stop.is_set(): + return + wake = False + with self._lock: + candidates: set[str] + if tool is None: + candidates = set(self._fetchers) + elif tool in self._fetchers: + candidates = {tool} + else: + candidates = set() + # 同一来源正在请求时的额外点击直接合并，不在完成后紧接着再拉一次。 + queued = { + name for name in candidates + if self._refreshing.get(name) is not self._fetchers.get(name) + } + self._pending_refresh.update(queued) + wake = bool(queued) + if wake: + self._refresh_evt.set() def _loop(self) -> None: + if self._stop.is_set(): + return + # 启动时有效配置拉取一次；claim 会同时吸收启动前的重复点击。 self._refresh_all() - try: - interval = max(30, float(self.settings.usage_refresh_seconds)) - except (TypeError, ValueError): - interval = 120 - log.warning("invalid usage_refresh_seconds; falling back to %ss", interval) while not self._stop.is_set(): - self._refresh_evt.wait(interval) + # pending set 是事实来源，Event 只负责唤醒。必须先 clear 再检查 + # predicate，避免 wait() 超时与 refresh_now().set() 同时发生时丢唤醒。 self._refresh_evt.clear() + batch = self._claim_pending_refresh() + if batch.fetchers: + self._run_refresh_batch(batch) + continue + signaled = self._refresh_evt.wait(self._next_wait_seconds()) if self._stop.is_set(): return - self._refresh_all() + if signaled: + continue + elif self.settings.usage_auto_refresh: + self._refresh_all(respect_due=True) - def _refresh_all(self) -> None: - force = self._force_refresh - self._force_refresh = False + def _global_interval(self) -> float: + try: + interval = float(self.settings.usage_refresh_seconds) + except (TypeError, ValueError, OverflowError): + interval = 120.0 + if not math.isfinite(interval): + interval = 120.0 + return min(86_400.0, max(30.0, interval)) + + def _source_interval(self, fetcher: object) -> float: + value = getattr(fetcher, "refresh_seconds", None) + if value is None: + return self._global_interval() + try: + interval = float(value) + except (TypeError, ValueError, OverflowError): + return self._global_interval() + if not math.isfinite(interval): + return self._global_interval() + return min(86_400.0, max(30.0, interval)) + + def _next_wait_seconds(self) -> float: + """Return the nearest source deadline without refreshing in manual mode.""" + fallback = self._global_interval() + if not self.settings.usage_auto_refresh: + return fallback + now = time.time() + with self._lock: + if not self._fetchers: + return fallback + due_times = [] + for name, fetcher in self._fetchers.items(): + interval = self._source_interval(fetcher) + if self._scheduled_intervals.get(name) != interval: + # Runtime interval changes take effect from the settings + # revision that woke this loop, not from an obsolete due time. + self._scheduled_intervals[name] = interval + self._next_due[name] = now + interval + due_times.append(self._next_due.get(name, now)) + return max(0.0, min(due_times) - now) + + def _claim_locked( + self, + candidates: set[str], + *, + respect_due: bool, + ) -> _RefreshBatch: + """Bind tools to fetchers/generation and mark them in-flight atomically.""" + generation = self._fetch_generation + if self._stop.is_set(): + return _RefreshBatch(generation, {}) + candidates.intersection_update(self._fetchers) + # A direct/startup claim absorbs an already pending click for the same + # source. A current-generation in-flight marker absorbs duplicates too. + coalesced = { + name for name in candidates + if self._refreshing.get(name) is self._fetchers.get(name) + } now = time.time() - for tool, fetcher in self._fetchers.items(): - # per-provider 刷新间隔：设了 refresh_seconds 的（corp provider）未到期就跳过， - # 除非用户显式 refresh_now 强制刷新。claude/codex 无此属性 → 每轮都刷。 - interval = getattr(fetcher, "refresh_seconds", None) - if not force and interval and now < self._next_due.get(tool, 0): + claimed: dict[str, object] = {} + # Preserve configured/provider insertion order. Besides deterministic UI + # and tests, this makes stop semantics predictable within a claimed batch. + for name in (name for name in self._fetchers if name in candidates - coalesced): + fetcher = self._fetchers[name] + if ( + respect_due + and now < self._next_due.get(name, 0) + ): continue - if interval: - self._next_due[tool] = now + interval - try: - snap = fetcher.fetch() - except Exception as e: # 任何异常都不能带崩后台线程 - log.warning("usage fetch %s failed: %s", tool, e) - snap = UsageSnapshot(tool, source="usage_api", error=str(e)) - if snap: + claimed[name] = fetcher + self._refreshing[name] = fetcher + self._pending_refresh.difference_update(set(claimed) | coalesced) + return _RefreshBatch(generation, claimed) + + def _claim_refresh( + self, + tools: set[str] | None = None, + *, + respect_due: bool = False, + ) -> _RefreshBatch: + with self._lock: + candidates = set(self._fetchers) if tools is None else set(tools) + return self._claim_locked(candidates, respect_due=respect_due) + + def _claim_pending_refresh(self) -> _RefreshBatch: + with self._lock: + candidates = set(self._pending_refresh) + # Drop names disabled by a reload so they cannot remain pending + # forever. Valid names are removed by _claim_locked once claimed or + # coalesced with an in-flight request. + self._pending_refresh.intersection_update(self._fetchers) + return self._claim_locked(candidates, respect_due=False) + + def _refresh_all( + self, + tools: set[str] | None = None, + *, + respect_due: bool = False, + ) -> None: + batch = self._claim_refresh(tools, respect_due=respect_due) + if batch.fetchers: + self._run_refresh_batch(batch) + + def _run_refresh_batch(self, batch: _RefreshBatch) -> None: + """Run one previously claimed batch without changing its identity.""" + try: + for tool, fetcher in batch.fetchers.items(): + if self._stop.is_set(): + break with self._lock: - prev = self._usage.get(tool) - # 新结果拿不到数据时，保留旧的有效窗口（标过期由 STALE 判定），只更新错误信息 - if snap.error and prev and prev.windows: - prev.error = snap.error - else: - self._usage[tool] = snap - if self._ccusage_bin: - self._refresh_ccusage() + if self._fetchers.get(tool) is not fetcher: + continue + if self._stop.is_set(): + break + try: + snap = fetcher.fetch() + except Exception as e: # 任何异常都不能带崩后台线程 + log.warning("usage fetch %s failed: %s", tool, e) + snap = UsageSnapshot(tool, source="usage_api", error=str(e)) + if snap: + with self._lock: + # 配置切换/禁用时丢弃旧请求的迟到结果，避免把 A 模型数字误标为 B。 + if self._fetchers.get(tool) is not fetcher: + continue + if snap.error: + self._usage_errors[tool] = snap + else: + self._usage[tool] = snap + self._usage_errors.pop(tool, None) + with self._lock: + if self._fetchers.get(tool) is fetcher: + interval = self._source_interval(fetcher) + self._scheduled_intervals[tool] = interval + # Schedule from completion so a slow request cannot make + # the same source immediately due again. + self._next_due[tool] = time.time() + interval + with self._lock: + refresh_ccusage = ( + self._ccusage_bin + and "claude" in batch.fetchers + and self._fetchers.get("claude") is batch.fetchers["claude"] + ) + if refresh_ccusage and not self._stop.is_set(): + self._refresh_ccusage() + finally: + with self._lock: + for tool, fetcher in batch.fetchers.items(): + if self._refreshing.get(tool) is fetcher: + self._refreshing.pop(tool, None) def _refresh_ccusage(self) -> None: try: diff --git a/agentbar/scheduler.py b/agentbar/scheduler.py index e0f51cc..430bc56 100644 --- a/agentbar/scheduler.py +++ b/agentbar/scheduler.py @@ -7,6 +7,7 @@ from __future__ import annotations import logging +import math import os import random import re @@ -17,7 +18,7 @@ from . import __version__ from .adapters.base import Outcome, get_registry -from .config import Settings +from .config import DEFAULT_PROVIDERS, DEFAULT_QUOTA_SOURCES, Settings from .models import ( FINISHED_STATES, PROFILES, @@ -44,6 +45,47 @@ def _clock(ts: float) -> str: return time.strftime("%H:%M", time.localtime(ts)) +def _validated_title(value, prompt: str) -> str: + if value is None or value == "": + return default_title(prompt) + if not isinstance(value, str): + raise ValueError("标题必须是字符串") + if len(value) > 200 or any(char in value for char in "\r\n\x00"): + raise ValueError("标题无效（最多 200 个字符，不能包含换行或 NUL）") + return value.strip() or default_title(prompt) + + +def _validated_schedule(value) -> float | None: + if value is None or value == "": + return None + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError("定时执行时间必须是有限的 Unix 时间戳") + scheduled = float(value) + if not math.isfinite(scheduled): + raise ValueError("定时执行时间必须是有限的 Unix 时间戳") + return scheduled if scheduled > time.time() else None + + +def _optional_text(value, label: str, max_len: int) -> str | None: + if value is None or value == "": + return None + if not isinstance(value, str): + raise ValueError(f"{label}必须是字符串") + cleaned = value.strip() + if len(cleaned) > max_len or any(char in cleaned for char in "\r\n\x00"): + raise ValueError(f"{label}无效（最多 {max_len} 个字符，不能包含换行或 NUL）") + return cleaned or None + + +def _valid_optional_number(value) -> bool: + return value is None or ( + not isinstance(value, bool) + and isinstance(value, (int, float)) + and math.isfinite(float(value)) + and float(value) >= 0 + ) + + class _Run: """一个正在运行的任务的进程句柄与控制位。""" @@ -66,81 +108,208 @@ def __init__(self, settings: Settings, store: StateStore): self._tasks: dict[str, Task] = {} self._order: list[str] = [] self._running: dict[str, _Run] = {} + # A user-triggered "立即重试" bypasses the shared tool cooldown exactly + # once, without clearing the observation for other queued tasks. + self._cooldown_bypass_once: set[str] = set() self._paused = False self._stop = threading.Event() self._wake = threading.Event() self._tick_thread: threading.Thread | None = None + self._lifecycle_lock = threading.Lock() + self._started = False + self._shutdown_complete = False self._load() # ================= persistence & recovery ================= def _load(self) -> None: data = self.store.load() - for d in data.get("tasks", []): + raw_tasks = data.get("tasks", []) + if not isinstance(raw_tasks, list): + raw_tasks = [] + for d in raw_tasks: + if not isinstance(d, dict): + log.warning("skip non-object task record") + continue try: t = Task.from_dict(d) + # IDs become log filenames and API path components. Reject + # traversal/control characters before they enter scheduler state. + StateStore._safe_task_id(t.id) + if not all( + isinstance(value, str) and value + for value in (t.title, t.prompt, t.tool, t.cwd, t.profile) + ): + raise ValueError("invalid core task fields") + if t.model is not None and not isinstance(t.model, str): + raise ValueError("invalid task model") + if t.effort is not None and not isinstance(t.effort, str): + raise ValueError("invalid task effort") + if not isinstance(t.resume_next, bool): + raise ValueError("invalid task resume flag") + if any( + isinstance(value, bool) or not isinstance(value, int) or value < 0 + for value in (t.attempts, t.quota_waits) + ): + raise ValueError("invalid task counters") + if not all(_valid_optional_number(value) for value in ( + t.created_at, + t.scheduled_at, + t.started_at, + t.finished_at, + t.next_retry_at, + t.cost_usd, + )): + raise ValueError("invalid task timestamps/cost") + if t.exit_code is not None and ( + isinstance(t.exit_code, bool) or not isinstance(t.exit_code, int) + ): + raise ValueError("invalid task exit code") + if t.session_id is not None and ( + not isinstance(t.session_id, str) + or not re.fullmatch(r"[A-Za-z0-9_-]{1,128}", t.session_id) + ): + raise ValueError("invalid task session id") + if not isinstance(t.state_reason, str): + raise ValueError("invalid task reason") except (TypeError, ValueError): log.warning("skip unparsable task record: %r", d) continue + if t.id in self._tasks: + log.warning("skip duplicate task id: %s", t.id) + continue + if t.tool not in self.registry: + t.state = TaskState.FAILED + t.finished_at = time.time() + t.state_reason = f"恢复失败：未知工具 {t.tool!r}" + elif t.profile not in PROFILES: + t.state = TaskState.FAILED + t.finished_at = time.time() + t.state_reason = f"恢复失败：未知权限档位 {t.profile!r}" if t.state == TaskState.RUNNING: - # 上个进程死掉时正在运行的任务 → 重新排队（能续会话就续） - t.state = TaskState.QUEUED + # AgentBar 与 CLI 使用独立进程组。若 AgentBar 崩溃，旧 CLI + # 可能仍在运行；自动重派会产生两个 agent 同时改代码/发消息。 + # 因此恢复为人工暂停，用户确认现场后再显式恢复。 + t.state = TaskState.PAUSED t.resume_next = bool(t.session_id) - t.state_reason = "调度器重启，已重新排队" + ( - "（将恢复会话）" if t.session_id else "" + t.state_reason = "检测到上次异常退出，为避免重复执行已暂停；请确认旧进程后手动恢复" + ( + "（可尝试恢复会话）" if t.session_id else "" ) self._tasks[t.id] = t self._order.append(t.id) - self._paused = bool(data.get("paused", False)) + paused = data.get("paused", False) + self._paused = paused if isinstance(paused, bool) else False self.quota.load(data.get("quota", {})) def _persist_locked(self) -> None: finished = [ tid for tid in self._order if self._tasks[tid].state in FINISHED_STATES ] + prune: set[str] = set() if len(finished) > MAX_FINISHED_KEPT: - for tid in finished[: len(finished) - MAX_FINISHED_KEPT]: - self._order.remove(tid) - self._tasks.pop(tid, None) + prune = set(finished[: len(finished) - MAX_FINISHED_KEPT]) + persisted_order = [tid for tid in self._order if tid not in prune] self.store.save( { "version": __version__, - "tasks": [self._tasks[tid].to_dict() for tid in self._order], + "tasks": [self._tasks[tid].to_dict() for tid in persisted_order], "paused": self._paused, "quota": self.quota.dump(), } ) + # Pruning is part of the same commit: a failed save must not silently + # delete completed tasks from the live scheduler. + if prune: + self._order[:] = persisted_order + for tid in prune: + self._tasks.pop(tid, None) + try: + self.store.delete_log(tid) + except OSError as exc: + # The state commit already retired this task. A filesystem + # cleanup error must not turn the successful transaction + # into a caller-visible failure, but should remain visible + # for diagnostics and a later manual cleanup. + log.warning("failed to delete retired task log %s: %s", tid, exc) + + def _ensure_mutable_locked(self) -> None: + if self._stop.is_set(): + raise ValueError("调度器正在停止，拒绝修改任务") + + @staticmethod + def _restore_task_locked(task: Task, checkpoint: dict) -> None: + """Restore a Task in place so any live references remain valid.""" + restored = Task.from_dict(checkpoint) + vars(task).clear() + vars(task).update(vars(restored)) # ================= lifecycle ================= def start(self) -> None: - self.quota.start_background() - self._tick_thread = threading.Thread( - target=self._tick_loop, name="agentbar-tick", daemon=True - ) - self._tick_thread.start() + with self._lifecycle_lock: + if self._started or self._shutdown_complete: + return + self.quota.start_background() + self._tick_thread = threading.Thread( + target=self._tick_loop, name="agentbar-tick", daemon=True + ) + self._started = True + self._tick_thread.start() def shutdown(self) -> None: """优雅退出：终止在跑的 CLI 进程，把任务放回队列（带 resume），落盘。""" - self._stop.set() - self._wake.set() - with self._lock: - runs = list(self._running.items()) - for _tid, run in runs: - run.mode = "interrupt" - run.cancel.set() - for _tid, run in runs: - if run.thread: - # `_tick` stores the Thread immediately before `.start()`; - # shutdown can win that tiny race. Wait for the worker's first - # instruction before joining so Python never sees an unstarted - # Thread object. - run.started.wait(timeout=1) - if run.thread.ident is not None: - run.thread.join(timeout=20) - self.quota.stop() - with self._lock: - self._persist_locked() + with self._lifecycle_lock: + if self._shutdown_complete: + return + # Linearize shutdown with public mutations and dispatch. Once this + # block completes, no later mutating API can commit new state. + with self._lock: + self._stop.set() + self._wake.set() + + # First retire the dispatcher. _tick also re-checks _stop while it + # holds the scheduler lock, so a tick that was already waiting on the + # lock cannot launch work after shutdown has begun. + tick = self._tick_thread + if tick and tick is not threading.current_thread(): + tick.join(timeout=5) + if tick.is_alive(): + log.warning("scheduler tick thread did not stop within 5s") + + with self._lock: + runs = list(self._running.items()) + for _tid, run in runs: + run.mode = "interrupt" + run.cancel.set() + # Do not rely solely on workers polling the cancel flag. A CLI that + # never reads stdin can leave its worker blocked in pipe.write(), so + # shutdown itself must retire every already-published process. + for _tid, run in runs: + proc = run.proc + if proc is not None and proc.poll() is None: + self._terminate(proc) + for _tid, run in runs: + if run.thread and run.thread is not threading.current_thread(): + # `_tick` stores the Thread immediately before `.start()`; + # wait for the worker's first instruction before joining so + # Python never sees an unstarted Thread object. + run.started.wait(timeout=1) + if run.thread.ident is not None: + # Binary discovery may make two 15s login-shell probes + # before it observes cancellation; allow that bounded + # path to retire too. + run.thread.join(timeout=35) + if run.thread.is_alive(): + log.warning("task worker %s did not stop within 35s", _tid) + self.quota.stop() + try: + with self._lock: + self._persist_locked() + finally: + # The scheduler cannot be restarted after its threads/processes + # have been retired. A disk failure is still reported, but repeat + # shutdown calls must be idempotent and all mutations stay closed. + self._shutdown_complete = True # ================= public API ================= @@ -158,53 +327,72 @@ def add_task( model: str | None = None, effort: str | None = None, scheduled_at: float | None = None, + before_commit=None, ) -> Task: - prompt = (prompt or "").strip() + if not isinstance(prompt, str): + raise ValueError("prompt 必须是字符串") + prompt = prompt.strip() if not prompt: raise ValueError("prompt 不能为空") if len(prompt) > 100_000: raise ValueError("prompt 过长（>100KB）") - if tool not in self.registry: + if not isinstance(tool, str) or tool not in self.registry: raise ValueError(f"未知工具 {tool!r}，可用: {sorted(self.registry)}") - if profile not in PROFILES: + if not isinstance(profile, str) or profile not in PROFILES: raise ValueError(f"未知权限档位 {profile!r}，可用: {PROFILES}") if profile == "full" and not self.settings.allow_full_profile: raise ValueError( "高权限档位默认关闭。如确需开启，编辑 config.json 设置 " "allow_full_profile=true 后重启" ) - model = (model or "").strip() or None - if model and (len(model) > 120 or any(c in model for c in "\r\n\x00")): - raise ValueError("模型名称无效（最多 120 个字符，不能包含换行）") - effort = (effort or "").strip().lower() or None + model = _optional_text(model, "模型名称", 120) + effort = _optional_text(effort, "推理强度", 32) + effort = effort.lower() if effort else None adapter = self.registry[tool] if effort and adapter.effort_choices and effort not in adapter.effort_choices: raise ValueError( f"{tool} 不支持强度 {effort!r}，可用: {adapter.effort_choices}" ) + if cwd is not None and not isinstance(cwd, str): + raise ValueError("工作目录必须是字符串") cwd = os.path.abspath(os.path.expanduser(cwd or self.settings.default_cwd)) if not os.path.isdir(cwd): raise ValueError(f"工作目录不存在: {cwd}") + clean_title = _validated_title(title, prompt) + clean_schedule = _validated_schedule(scheduled_at) t = Task( id=new_id(), - title=(title or "").strip() or default_title(prompt), + title=clean_title, prompt=prompt, tool=tool, cwd=cwd, profile=profile, model=model, effort=effort, - scheduled_at=scheduled_at if scheduled_at and scheduled_at > time.time() else None, + scheduled_at=clean_schedule, ) with self._lock: + if before_commit is not None: + before_commit() + self._ensure_mutable_locked() self._tasks[t.id] = t self._order.append(t.id) - self._persist_locked() + try: + self._persist_locked() + except Exception: + # A caller that saw a persistence error must not leave a task + # behind that a later tick can execute. + self._tasks.pop(t.id, None) + try: + self._order.remove(t.id) + except ValueError: + pass + raise self._wake.set() log.info("task %s added (%s, %s)", t.id, tool, t.title) return t - def edit_task(self, task_id: str, changes: dict) -> Task: + def edit_task(self, task_id: str, changes: dict, before_commit=None) -> Task: """Edit a task that has not started running. A live CLI process is intentionally immutable: changing its prompt or @@ -212,7 +400,12 @@ def edit_task(self, task_id: str, changes: dict) -> Task: already executing. Users can cancel a running task and create/retry a replacement instead. """ + if not isinstance(changes, dict): + raise ValueError("任务修改必须是 JSON object") with self._lock: + if before_commit is not None: + before_commit() + self._ensure_mutable_locked() t = self._tasks.get(task_id) if not t: raise ValueError("任务不存在") @@ -223,50 +416,70 @@ def edit_task(self, task_id: str, changes: dict) -> Task: }: raise ValueError(f"状态 {t.state.value} 的任务不能编辑；请使用重试创建新运行") - prompt = (changes.get("prompt", t.prompt) or "").strip() + prompt_value = changes.get("prompt", t.prompt) + if not isinstance(prompt_value, str): + raise ValueError("prompt 必须是字符串") + prompt = prompt_value.strip() if not prompt: raise ValueError("prompt 不能为空") if len(prompt) > 100_000: raise ValueError("prompt 过长（>100KB）") tool = changes.get("tool", t.tool) - if tool not in self.registry: + if not isinstance(tool, str) or tool not in self.registry: raise ValueError(f"未知工具 {tool!r}，可用: {sorted(self.registry)}") profile = changes.get("profile", t.profile) - if profile not in PROFILES: + if not isinstance(profile, str) or profile not in PROFILES: raise ValueError(f"未知权限档位 {profile!r}，可用: {PROFILES}") if profile == "full" and not self.settings.allow_full_profile: raise ValueError("高权限档位默认关闭。如确需开启，编辑 config.json 设置 allow_full_profile=true 后重启") - model = (changes.get("model", t.model) or "").strip() or None - if model and (len(model) > 120 or any(c in model for c in "\r\n\x00")): - raise ValueError("模型名称无效（最多 120 个字符，不能包含换行）") - effort = (changes.get("effort", t.effort) or "").strip().lower() or None + model = _optional_text(changes.get("model", t.model), "模型名称", 120) + effort = _optional_text(changes.get("effort", t.effort), "推理强度", 32) + effort = effort.lower() if effort else None adapter = self.registry[tool] if effort and adapter.effort_choices and effort not in adapter.effort_choices: raise ValueError( f"{tool} 不支持强度 {effort!r}，可用: {adapter.effort_choices}" ) - cwd = os.path.abspath(os.path.expanduser(changes.get("cwd", t.cwd) or self.settings.default_cwd)) + cwd_value = changes.get("cwd", t.cwd) + if cwd_value is not None and not isinstance(cwd_value, str): + raise ValueError("工作目录必须是字符串") + cwd = os.path.abspath(os.path.expanduser(cwd_value or self.settings.default_cwd)) if not os.path.isdir(cwd): raise ValueError(f"工作目录不存在: {cwd}") - title = (changes.get("title", t.title) or "").strip() or default_title(prompt) + title = _validated_title(changes.get("title", t.title), prompt) + scheduled_at = ( + _validated_schedule(changes.get("scheduled_at")) + if "scheduled_at" in changes + else t.scheduled_at + ) + checkpoint = t.to_dict() t.prompt, t.tool, t.cwd = prompt, tool, cwd t.title, t.profile, t.model, t.effort = title, profile, model, effort + t.scheduled_at = scheduled_at if t.state == TaskState.WAITING_QUOTA: t.state = TaskState.QUEUED t.next_retry_at = None t.state_reason = "编辑后重新入队" - self._persist_locked() + try: + self._persist_locked() + except Exception: + self._restore_task_locked(t, checkpoint) + raise self._wake.set() log.info("task %s edited", task_id) return t - def act(self, task_id: str, action: str) -> tuple[bool, str]: + def act(self, task_id: str, action: str, before_commit=None) -> tuple[bool, str]: with self._lock: + if before_commit is not None: + before_commit() + if self._stop.is_set(): + return False, "调度器正在停止，拒绝修改任务" t = self._tasks.get(task_id) if not t: return False, "任务不存在" @@ -277,22 +490,50 @@ def act(self, task_id: str, action: str) -> tuple[bool, str]: if run: run.mode = "cancel" run.cancel.set() + proc = run.proc + if proc is not None and proc.poll() is None: + threading.Thread( + target=self._terminate, + args=(proc,), + name=f"agentbar-cancel-{task_id}", + daemon=True, + ).start() return True, "正在终止进程…" if s in (TaskState.QUEUED, TaskState.PAUSED, TaskState.WAITING_QUOTA): + checkpoint = t.to_dict() + cooldown_before = set(self._cooldown_bypass_once) + self._cooldown_bypass_once.discard(task_id) t.state = TaskState.CANCELLED t.finished_at = time.time() t.state_reason = "用户取消" - self._persist_locked() + try: + self._persist_locked() + except Exception: + self._restore_task_locked(t, checkpoint) + self._cooldown_bypass_once.clear() + self._cooldown_bypass_once.update(cooldown_before) + raise return True, "已取消" return False, f"状态 {s.value} 不可取消" if action == "pause": if s in (TaskState.QUEUED, TaskState.WAITING_QUOTA): + checkpoint = t.to_dict() + cooldown_before = set(self._cooldown_bypass_once) + self._cooldown_bypass_once.discard(task_id) t.state = TaskState.PAUSED t.state_reason = "人工暂停" - self._persist_locked() + try: + self._persist_locked() + except Exception: + self._restore_task_locked(t, checkpoint) + self._cooldown_bypass_once.clear() + self._cooldown_bypass_once.update(cooldown_before) + raise return True, "已暂停" return False, f"状态 {s.value} 不可暂停（运行中请用取消）" if action == "resume": + checkpoint = t.to_dict() + cooldown_before = set(self._cooldown_bypass_once) if s == TaskState.PAUSED: t.state = TaskState.QUEUED t.state_reason = "人工恢复" @@ -300,20 +541,32 @@ def act(self, task_id: str, action: str) -> tuple[bool, str]: t.state = TaskState.QUEUED t.next_retry_at = None t.state_reason = "人工触发立即重试" + self._cooldown_bypass_once.add(task_id) else: return False, f"状态 {s.value} 不可恢复" - self._persist_locked() + try: + self._persist_locked() + except Exception: + self._restore_task_locked(t, checkpoint) + self._cooldown_bypass_once.clear() + self._cooldown_bypass_once.update(cooldown_before) + raise self._wake.set() return True, "已恢复" if action == "retry": if s in FINISHED_STATES: + checkpoint = t.to_dict() t.state = TaskState.QUEUED t.finished_at = None t.exit_code = None t.next_retry_at = None t.resume_next = bool(t.session_id) t.state_reason = "人工重试" - self._persist_locked() + try: + self._persist_locked() + except Exception: + self._restore_task_locked(t, checkpoint) + raise self._wake.set() return True, "已重新入队" return False, f"状态 {s.value} 不可重试" @@ -329,6 +582,7 @@ def _reorder_locked(self, t: Task, action: str) -> tuple[bool, str]: tid for tid in self._order if self._tasks[tid].state == TaskState.QUEUED ] + order_before = list(self._order) i = queued.index(t.id) if action == "move_up" and i > 0: self._swap_order(queued[i], queued[i - 1]) @@ -341,7 +595,11 @@ def _reorder_locked(self, t: Task, action: str) -> tuple[bool, str]: i -= 1 else: return True, "已在该位置" - self._persist_locked() + try: + self._persist_locked() + except Exception: + self._order[:] = order_before + raise self._wake.set() return True, "已调整优先级" @@ -349,16 +607,32 @@ def _swap_order(self, a: str, b: str) -> None: ia, ib = self._order.index(a), self._order.index(b) self._order[ia], self._order[ib] = self._order[ib], self._order[ia] - def pause_all(self) -> None: + def pause_all(self, before_commit=None) -> None: with self._lock: + if before_commit is not None: + before_commit() + self._ensure_mutable_locked() + paused_before = self._paused self._paused = True - self._persist_locked() + try: + self._persist_locked() + except Exception: + self._paused = paused_before + raise log.info("pause_all") - def resume_all(self) -> None: + def resume_all(self, before_commit=None) -> None: with self._lock: + if before_commit is not None: + before_commit() + self._ensure_mutable_locked() + paused_before = self._paused self._paused = False - self._persist_locked() + try: + self._persist_locked() + except Exception: + self._paused = paused_before + raise self._wake.set() log.info("resume_all") @@ -381,6 +655,23 @@ def snapshot(self) -> dict: status = "waiting" else: status = "idle" + provider_config = {} + for name in DEFAULT_PROVIDERS: + cfg = (self.settings.providers or {}).get(name) or {} + provider_config[name] = { + "enabled": bool(cfg.get("enabled")), + "cookie_set": bool(str(cfg.get("cookie") or "").strip()), + "unit": cfg.get("unit") or DEFAULT_PROVIDERS[name]["unit"], + } + quota_source_config = {} + for name in DEFAULT_QUOTA_SOURCES: + cfg = (self.settings.quota_sources or {}).get(name) or {} + quota_source_config[name] = { + "enabled": bool(cfg.get("enabled")), + "key_set": bool(str(cfg.get("access_token") or "").strip()), + "model": str(cfg.get("model") or ""), + "account_id_set": bool(str(cfg.get("account_id") or "").strip()), + } return { "version": __version__, "status": status, @@ -391,17 +682,23 @@ def snapshot(self) -> dict: "tasks": tasks, "quota": { name: self.quota.status(name).to_dict() - # 任务 adapter（claude/codex）+ 仅额度的 corp provider（mytoken/tokenverse） + # 展示已启用的上游来源 + 已有真实任务观测的工具。 for name in dict.fromkeys( - list(self.registry) + self.quota.provider_tools() + self.quota.provider_tools() + self.quota.observed_tools() ) }, "title_provider": self.settings.title_provider, + # 菜单只需要无敏感信息的配置摘要，用于展示“未配置/待刷新”入口。 + # Cookie 永远不进入 snapshot、日志或前端状态转储。 + "provider_config": provider_config, + # access key/account id 永不进入 snapshot。 + "quota_source_config": quota_source_config, "cli_processes": self._cli_processes_snapshot(), "settings": { "max_parallel": self.settings.max_parallel, "per_tool_limit": self.settings.per_tool_limit, "usage_refresh_seconds": self.settings.usage_refresh_seconds, + "usage_auto_refresh": self.settings.usage_auto_refresh, "default_cwd": self.settings.default_cwd, "allow_full_profile": self.settings.allow_full_profile, "state_dir": str(self.settings.state_dir), @@ -456,6 +753,17 @@ def _tool_running(self, tool: str) -> int: def _tick(self) -> None: to_start: list[tuple[str, _Run]] = [] with self._lock: + if self._stop.is_set(): + return + # Task fields are mutated before persistence so state.json records + # RUNNING before a worker can launch. Keep an in-place checkpoint so + # a failed save can roll the tentative transition back completely. + changed_tasks: dict[str, dict] = {} + cooldown_before = set(self._cooldown_bypass_once) + + def checkpoint(task: Task) -> None: + changed_tasks.setdefault(task.id, task.to_dict()) + now = time.time() dirty = False # 1) 额度等待期结束 → 回到队列 @@ -465,12 +773,15 @@ def _tick(self) -> None: and t.next_retry_at and t.next_retry_at <= now ): + checkpoint(t) t.state = TaskState.QUEUED t.state_reason = "额度等待结束，重新排队" dirty = True # 2) FIFO 派发（跳过被额度冷却/并发上限/定时卡住的任务） if not self._paused: for t in self._iter_tasks(): + if self._stop.is_set(): + break if t.state != TaskState.QUEUED: continue if t.scheduled_at and t.scheduled_at > now: @@ -479,21 +790,37 @@ def _tick(self) -> None: break if self._tool_running(t.tool) >= self.settings.per_tool_limit: continue - cd = self.quota.cooldown_until(t.tool) - if cd and cd > now: - continue + bypass_cooldown = t.id in self._cooldown_bypass_once + if not bypass_cooldown: + cd = self.quota.cooldown_until(t.tool, t.model) + if cd and cd > now: + continue run = _Run() run.was_resume = bool(t.resume_next and t.session_id) + checkpoint(t) self._running[t.id] = run t.state = TaskState.RUNNING t.started_at = now t.finished_at = None t.attempts += 1 t.state_reason = "执行中" + self._cooldown_bypass_once.discard(t.id) to_start.append((t.id, run)) dirty = True if dirty: - self._persist_locked() + try: + self._persist_locked() + except Exception: + for tid, previous in changed_tasks.items(): + live = self._tasks.get(tid) + if live is not None: + self._restore_task_locked(live, previous) + for tid, run in to_start: + if self._running.get(tid) is run: + self._running.pop(tid, None) + self._cooldown_bypass_once.clear() + self._cooldown_bypass_once.update(cooldown_before) + raise for tid, run in to_start: th = threading.Thread( target=self._run_task, args=(tid, run), name=f"agentbar-run-{tid}", @@ -516,28 +843,52 @@ def _run_task(self, task_id: str, run: _Run) -> None: timeout = self.settings.task_timeout_seconds cwd = t.cwd - log_path = self.store.log_path(task_id) - binary = adapter.binary() - if not binary: + if self._stop.is_set() or run.cancel.is_set(): self._finalize( task_id, run, - Outcome("failure", - f"未找到 {t.tool} 可执行文件；请安装，或在 config.json 的 " - f"tool_paths 中指定绝对路径"), + Outcome("interrupted", "调度器退出，任务已回到队列"), None, None, "", ) return - argv = adapter.build_argv(t, resume=resume, binary=binary) - payload = adapter.stdin_payload(t, resume=resume) - env = os.environ.copy() - for k in _ENV_STRIP: - env.pop(k, None) - env = adapter.build_env(env) # ensure tool binary dir is on PATH + try: + binary = adapter.binary() + if not binary: + self._finalize( + task_id, run, + Outcome("failure", + f"未找到 {t.tool} 可执行文件；请安装，或在 config.json 的 " + f"tool_paths 中指定绝对路径"), + None, None, "", + ) + return + + if self._stop.is_set() or run.cancel.is_set(): + self._finalize( + task_id, run, + Outcome("interrupted", "调度器退出，任务已回到队列"), + None, None, "", + ) + return + + argv = adapter.build_argv(t, resume=resume, binary=binary) + payload = adapter.stdin_payload(t, resume=resume) + env = os.environ.copy() + for k in _ENV_STRIP: + env.pop(k, None) + env = adapter.build_env(env) # ensure tool binary dir is on PATH + except Exception as e: + self._finalize( + task_id, run, + Outcome("failure", f"启动准备失败: {e}"), + None, None, "", + ) + return rc: int | None = None + proc: subprocess.Popen | None = None try: - with open(log_path, "ab") as lf: + with self.store.open_log_append(task_id) as lf: shown = [a if len(a) < 200 else a[:200] + "…" for a in argv] header = ( f"\n===== attempt {t.attempts} @ " @@ -546,16 +897,31 @@ def _run_task(self, task_id: str, run: _Run) -> None: ) lf.write(header.encode()) lf.flush() - proc = subprocess.Popen( - argv, - cwd=cwd, - env=env, - stdout=lf, - stderr=subprocess.STDOUT, - stdin=subprocess.PIPE if payload is not None else subprocess.DEVNULL, - start_new_session=True, # 独立进程组，方便整组终止 - ) - run.proc = proc + # Pair the final stop check and Popen registration under the + # scheduler lock. Shutdown can then either prevent the launch or + # see the process in _running and cancel it; there is no invisible + # post-shutdown child-process window. + with self._lock: + if self._stop.is_set() or run.cancel.is_set(): + proc = None + else: + proc = subprocess.Popen( + argv, + cwd=cwd, + env=env, + stdout=lf, + stderr=subprocess.STDOUT, + stdin=subprocess.PIPE if payload is not None else subprocess.DEVNULL, + start_new_session=True, # 独立进程组，方便整组终止 + ) + run.proc = proc + if proc is None: + self._finalize( + task_id, run, + Outcome("interrupted", "调度器退出，任务已回到队列"), + None, None, "", + ) + return if payload is not None and proc.stdin: try: proc.stdin.write(payload.encode()) @@ -566,6 +932,21 @@ def _run_task(self, task_id: str, run: _Run) -> None: next_sid_probe = time.time() + 2.0 sid_captured = False while True: + if run.cancel.is_set(): + self._terminate(proc) + if run.mode == "interrupt": + self._finalize( + task_id, run, + Outcome("interrupted", "调度器退出，任务已回到队列"), + None, None, "", + ) + else: + self._finalize( + task_id, run, + Outcome("cancelled", "用户取消，进程已终止"), + None, None, "", + ) + return rc = proc.poll() if rc is not None: break @@ -583,21 +964,6 @@ def _run_task(self, task_id: str, run: _Run) -> None: if t_live and not t_live.session_id: t_live.session_id = sid_early self._persist_locked() - if run.cancel.is_set(): - self._terminate(proc) - if run.mode == "interrupt": - self._finalize( - task_id, run, - Outcome("interrupted", "调度器退出，任务已回到队列"), - None, None, "", - ) - else: - self._finalize( - task_id, run, - Outcome("cancelled", "用户取消，进程已终止"), - None, None, "", - ) - return if time.time() > deadline: self._terminate(proc) self._finalize( @@ -607,17 +973,34 @@ def _run_task(self, task_id: str, run: _Run) -> None: ) return time.sleep(0.2) - except Exception as e: # Popen 失败等 + except Exception as e: # Popen 或运行中日志/会话探测失败 + # Once Popen succeeds, every exit path must retire the process before + # _finalize removes it from _running. Otherwise an agent can continue + # editing after the UI has already reported the task as failed, and + # shutdown no longer has a handle with which to stop it. + if proc is not None and proc.poll() is None: + try: + self._terminate(proc) + except Exception: + log.exception("failed to terminate task %s after monitor error", task_id) + phase = "执行监控失败" if proc is not None else "启动失败" self._finalize( - task_id, run, Outcome("failure", f"启动失败: {e}"), None, None, "" + task_id, run, Outcome("failure", f"{phase}: {e}"), None, None, "" ) return - tail = self.store.read_log_tail(task_id) - outcome = adapter.classify(rc, tail) - # 头+尾都扫：codex 的 session id 在输出开头，长输出下 tail 看不到它 - head = self.store.read_log_head(task_id) - sid = adapter.extract_session_id(head + "\n" + tail) + try: + tail = self.store.read_log_tail(task_id) + outcome = adapter.classify(rc, tail) + # 头+尾都扫：codex 的 session id 在输出开头，长输出下 tail 看不到它 + head = self.store.read_log_head(task_id) + sid = adapter.extract_session_id(head + "\n" + tail) + except Exception as e: + self._finalize( + task_id, run, + Outcome("failure", f"执行结果处理失败: {e}"), rc, None, "", + ) + return self._finalize(task_id, run, outcome, rc, sid, tail) @staticmethod @@ -640,7 +1023,10 @@ def _terminate(proc: subprocess.Popen) -> None: os.killpg(pgid, signal.SIGKILL) except (ProcessLookupError, OSError): pass - proc.wait(timeout=5) + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + log.warning("process group %s did not exit after SIGKILL", pgid) def _backoff_seconds(self, quota_waits: int) -> float: mins = self.settings.backoff_minutes or [5] diff --git a/agentbar/server.py b/agentbar/server.py index 2c24ef7..8fe4bef 100644 --- a/agentbar/server.py +++ b/agentbar/server.py @@ -1,7 +1,7 @@ """Localhost HTTP API + task manager web UI (+ optional LAN access for mobile). 安全：默认绑 127.0.0.1；lan_access=true 时绑 0.0.0.0 供同局域网手机访问。 -所有 /api（除 /api/ping）要求 token（Header 或 query）。Host 头校验只放行 +所有 /api（除 /api/ping）要求 Header token。Host 头校验只放行 IP 字面量（DNS rebinding 必须借助域名，放行裸 IP 不破坏该防御）。 token 存于 state 目录 config.json（0600）。 """ @@ -12,33 +12,120 @@ import ipaddress import json import logging +import re import socket import threading +import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from importlib import resources from urllib.parse import parse_qs, urlencode, urlparse from . import __version__ from .browser_cookies import CookieImportError, import_cookie_header -from .config import DEFAULT_PROVIDERS, PROVIDER_UNITS, Settings, save_settings +from .config import ( + DEFAULT_QUOTA_SOURCES, + DEFAULT_PROVIDERS, + PROVIDER_HOSTS, + PROVIDER_UNITS, + Settings, + save_settings, +) from .scheduler import Scheduler log = logging.getLogger("agentbar.server") MAX_BODY = 200_000 -ALLOWED_HOSTS = {"127.0.0.1", "localhost"} -PROVIDER_HOSTS = { - "mytoken": "mytoken.corp.kuaishou.com", - "tokenverse": "tokenverse.corp.kuaishou.com", +REQUEST_IO_TIMEOUT_SECONDS = 15 +HANDLER_DRAIN_SECONDS = 5 +ALLOWED_HOSTS = {"127.0.0.1", "::1", "localhost"} +_PROXY_HEADERS = { + "forwarded", + "via", + "x-forwarded-for", + "x-forwarded-host", + "x-forwarded-proto", + "x-real-ip", + "true-client-ip", } +_HTML_CSP = ( + "default-src 'none'; base-uri 'none'; frame-ancestors 'none'; " + "form-action 'self'; script-src 'unsafe-inline'; style-src 'unsafe-inline'; " + "connect-src 'self'; img-src 'self' data:" +) + + +class _BodyError(ValueError): + """A safe client-facing request-body error with an explicit HTTP status.""" + + def __init__(self, status: int, message: str): + super().__init__(message) + self.status = status + + +class _AdmissionClosed(RuntimeError): + """Raised when a request reaches its commit point during server shutdown.""" + + +class _RequestAdmission: + """Close mutation admission instantly, then drain accepted handlers boundedly.""" + + def __init__(self) -> None: + self.stopping = threading.Event() + self._condition = threading.Condition() + self._active_handlers = 0 + + def enter_handler(self) -> bool: + with self._condition: + if self.stopping.is_set(): + return False + self._active_handlers += 1 + return True + + def leave_handler(self) -> None: + with self._condition: + self._active_handlers -= 1 + if self._active_handlers == 0: + self._condition.notify_all() + + def close(self) -> None: + # Event.set() does not wait for a long-running handler. Every mutation + # checks this predicate again immediately before touching core/settings. + self.stopping.set() + + def require_open(self) -> None: + if self.stopping.is_set(): + raise _AdmissionClosed("HTTP server is stopping") + + def wait_for_handlers(self, timeout: float) -> bool: + deadline = time.monotonic() + timeout + with self._condition: + while self._active_handlers: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + self._condition.wait(remaining) + return True def _host_of(header: str) -> str: """Extract host part from a Host header ([::1]:8737 / 10.1.2.3:8737 / localhost).""" - header = header or "" + header = (header or "").strip() + if not header or any(char.isspace() for char in header): + return "" if header.startswith("["): - return header[1:].split("]")[0].lower() - return header.split(":")[0].lower() + end = header.find("]") + if end < 0: + return "" + suffix = header[end + 1:] + if suffix and not (suffix.startswith(":") and suffix[1:].isdigit()): + return "" + return header[1:end].lower() + if header.count(":") > 1: # IPv6 Host 必须使用 [addr]:port + return "" + host, sep, port = header.partition(":") + if sep and not port.isdigit(): + return "" + return host.lower() def _is_ip_literal(host: str) -> bool: @@ -67,10 +154,15 @@ def _cookie_preview(cookie: str) -> str: return "" names = [] for part in cookie.split(";"): - name = part.strip().split("=", 1)[0].strip() - if name: + raw = part.strip() + if "=" not in raw: + continue + name = raw.split("=", 1)[0].strip() + # Never echo arbitrary user input. Only cookie-name tokens are safe to + # expose; a pasted JWT/Bearer without '=' must remain completely masked. + if re.fullmatch(r"[!#$%&'*+\-.^_`|~0-9A-Za-z]{1,128}", name): names.append(name) - shown = ", ".join(names[:4]) + shown = ", ".join(names[:4]) or "已配置" return shown + (" ..." if len(names) > 4 else "") @@ -90,35 +182,162 @@ def _provider_config_payload(settings: Settings) -> dict: return { "ok": True, "providers": providers, + "quota_sources": { + name: { + "enabled": bool(((settings.quota_sources or {}).get(name) or {}).get("enabled")), + "model": str(((settings.quota_sources or {}).get(name) or {}).get("model") or ""), + "key_set": bool(str( + ((settings.quota_sources or {}).get(name) or {}).get("access_token") or "" + ).strip()), + "account_id_set": bool(str( + ((settings.quota_sources or {}).get(name) or {}).get("account_id") or "" + ).strip()), + } + for name in DEFAULT_QUOTA_SOURCES + }, + "usage_auto_refresh": bool(settings.usage_auto_refresh), "title_provider": settings.title_provider, } -def _apply_provider_settings(settings: Settings, payload: dict) -> None: - providers = payload.get("providers") or {} - merged = json.loads(json.dumps(settings.providers or DEFAULT_PROVIDERS)) - for name, defaults in DEFAULT_PROVIDERS.items(): - incoming = providers.get(name) - if not isinstance(incoming, dict): - continue - cfg = merged.setdefault(name, dict(defaults)) - if "enabled" in incoming: - cfg["enabled"] = bool(incoming.get("enabled")) - if incoming.get("unit") in PROVIDER_UNITS: - cfg["unit"] = incoming["unit"] - if "refresh_seconds" in incoming: - try: - cfg["refresh_seconds"] = max(60, int(incoming.get("refresh_seconds") or 300)) - except (TypeError, ValueError): - cfg["refresh_seconds"] = defaults["refresh_seconds"] - if "cookie" in incoming: - # Missing cookie keeps the existing secret; explicit empty string clears it. - cfg["cookie"] = str(incoming.get("cookie") or "").strip() - title = payload.get("title_provider") - if title in {"claude", "codex", *DEFAULT_PROVIDERS.keys()}: - settings.title_provider = title - settings.providers = merged - save_settings(settings) +def _apply_provider_settings( + settings: Settings, + payload: dict, + before_commit=None, +) -> None: + if not isinstance(payload, dict): + raise ValueError("配置请求必须是 JSON object") + with settings._lock: + providers = payload.get("providers", {}) + if not isinstance(providers, dict): + raise ValueError("providers 必须是 JSON object") + merged = json.loads(json.dumps(settings.providers or DEFAULT_PROVIDERS)) + for name, defaults in DEFAULT_PROVIDERS.items(): + incoming = providers.get(name) + if incoming is None: + continue + if not isinstance(incoming, dict): + raise ValueError(f"providers.{name} 必须是 JSON object") + cfg = merged.setdefault(name, dict(defaults)) + if "enabled" in incoming: + cfg["enabled"] = _clean_bool_field( + incoming.get("enabled"), f"providers.{name}.enabled", + ) + if "unit" in incoming: + unit = incoming.get("unit") + if not isinstance(unit, str) or unit not in PROVIDER_UNITS: + raise ValueError( + f"providers.{name}.unit 必须是 " + f"{', '.join(PROVIDER_UNITS)} 之一" + ) + cfg["unit"] = unit + if "refresh_seconds" in incoming: + seconds = incoming.get("refresh_seconds") + if isinstance(seconds, bool) or not isinstance(seconds, int): + raise ValueError( + f"providers.{name}.refresh_seconds 必须是整数" + ) + cfg["refresh_seconds"] = min(86_400, max(60, seconds)) + if "cookie" in incoming: + # Missing cookie keeps the existing secret; explicit empty string clears it. + cfg["cookie"] = _clean_credential_field( + incoming.get("cookie"), "Cookie", 65_536, + ) + + next_title = settings.title_provider + if "title_provider" in payload: + title = payload.get("title_provider") + if not isinstance(title, str) or title not in { + "claude", "codex", *DEFAULT_PROVIDERS.keys(), + }: + raise ValueError("title_provider 无效") + next_title = title + source_payload = payload.get("quota_sources", {}) + if not isinstance(source_payload, dict): + raise ValueError("quota_sources 必须是 JSON object") + sources = json.loads(json.dumps(settings.quota_sources or DEFAULT_QUOTA_SOURCES)) + for name, defaults in DEFAULT_QUOTA_SOURCES.items(): + incoming = source_payload.get(name) + if incoming is None: + continue + if not isinstance(incoming, dict): + raise ValueError(f"quota_sources.{name} 必须是 JSON object") + cfg = sources.setdefault(name, dict(defaults)) + cfg.pop("api_key", None) # 旧字段只读迁移，永不再落盘。 + if "enabled" in incoming: + cfg["enabled"] = _clean_bool_field( + incoming.get("enabled"), f"quota_sources.{name}.enabled", + ) + if "model" in incoming: + cfg["model"] = _clean_credential_field(incoming.get("model"), "模型", 160) + token_submitted = "access_token" in incoming or "api_key" in incoming + if token_submitted: + # access_token 是唯一正式字段；api_key 仅接受旧 Web 客户端迁移。 + raw_token = ( + incoming.get("access_token") + if "access_token" in incoming + else incoming.get("api_key") + ) + cfg["access_token"] = _clean_credential_field( + raw_token, "OAuth Access Token", 16_384, + ) + if "account_id" in incoming: + cfg["account_id"] = _clean_credential_field( + incoming.get("account_id"), "Account ID", 256, + ) + if token_submitted and not cfg.get("access_token"): + # 显式清空凭据时同时停用，避免留下会持续报错的半配置。 + if incoming.get("enabled") is True: + raise ValueError(f"{name} 启用前必须输入 OAuth Access Token") + cfg["enabled"] = False + if cfg.get("enabled") and not str(cfg.get("access_token") or "").strip(): + raise ValueError(f"{name} 启用前必须输入 OAuth Access Token") + next_auto_refresh = settings.usage_auto_refresh + if "usage_auto_refresh" in payload: + next_auto_refresh = _clean_bool_field( + payload.get("usage_auto_refresh"), "usage_auto_refresh", + ) + + # 保存失败时恢复旧内存状态，避免 API 回报失败却部分生效。 + old = ( + settings.title_provider, + settings.providers, + settings.quota_sources, + settings.usage_auto_refresh, + ) + if before_commit is not None: + before_commit() + settings.title_provider = next_title + settings.providers = merged + settings.quota_sources = sources + settings.usage_auto_refresh = next_auto_refresh + try: + save_settings(settings) + except Exception: + ( + settings.title_provider, + settings.providers, + settings.quota_sources, + settings.usage_auto_refresh, + ) = old + raise + + +def _clean_credential_field(value, label: str, max_len: int) -> str: + if not isinstance(value, str): + raise ValueError(f"{label}必须是字符串") + value = value.strip() + if len(value) > max_len: + raise ValueError(f"{label} 过长") + if any(char in value for char in "\r\n\x00"): + raise ValueError(f"{label} 不能包含换行或 NUL") + return value + + +def _clean_bool_field(value, label: str) -> bool: + if not isinstance(value, bool): + raise ValueError(f"{label} 必须是 boolean") + return value class ApiServer: @@ -129,16 +348,30 @@ def __init__(self, core: Scheduler, settings: Settings): self.hooks: dict = {"dispatch": None} # 动态 Host 白名单（公网隧道域名启动后注册进来；其余域名一律 403） self.extra_hosts: set[str] = set() - handler = _make_handler(core, settings, self.hooks, self.extra_hosts) + self._admission = _RequestAdmission() + handler = _make_handler( + core, settings, self.hooks, self.extra_hosts, self._admission + ) bind = "0.0.0.0" if settings.lan_access else "127.0.0.1" self.httpd = ThreadingHTTPServer((bind, settings.port), handler) + # We provide our own bounded drain. ThreadingMixIn's default unbounded + # join can hang SIGTERM behind a slow Keychain/browser-cookie import. + # Lingering daemon handlers cannot commit after admission closes. self.httpd.daemon_threads = True + self.httpd.block_on_close = False self._thread: threading.Thread | None = None + self._lifecycle_lock = threading.Lock() + self._started = False + self._stopped = False @property def port(self) -> int: return self.httpd.server_address[1] + @property + def stopping(self) -> bool: + return self._admission.stopping.is_set() + def url(self, with_token: bool = False, **query: str) -> str: """Return a local panel URL, optionally carrying an authenticated UI intent. @@ -148,9 +381,10 @@ def url(self, with_token: bool = False, **query: str) -> str: """ base = f"http://127.0.0.1:{self.port}/" params = {key: str(value) for key, value in query.items() if value is not None} - if with_token: - params["token"] = self.settings.token - return base + (f"?{urlencode(params)}" if params else "") + query_string = f"?{urlencode(params)}" if params else "" + # Fragment 不会进入 HTTP request line / proxy log / Referer。 + fragment = f"#{urlencode({'token': self.settings.token})}" if with_token else "" + return base + query_string + fragment def allow_host(self, hostname: str) -> None: self.extra_hosts.add(hostname.lower()) @@ -165,33 +399,93 @@ def mobile_url(self) -> str | None: ip = lan_ip() if not ip: return None - return f"http://{ip}:{self.port}/m?token={self.settings.token}" + return f"http://{ip}:{self.port}/m#{urlencode({'token': self.settings.token})}" def start(self) -> None: - self._thread = threading.Thread( - target=self.httpd.serve_forever, name="agentbar-http", daemon=True - ) - self._thread.start() + with self._lifecycle_lock: + if self._started or self._stopped: + return + self._thread = threading.Thread( + target=self.httpd.serve_forever, name="agentbar-http", daemon=True + ) + self._thread.start() + self._started = True log.info("api server on %s", self.url()) def stop(self) -> None: - self.httpd.shutdown() - self.httpd.server_close() + # Hold the lifecycle lock through the drain so concurrent stop callers do + # not return before the owner has finished waiting for accepted handlers. + with self._lifecycle_lock: + if self._stopped: + return + self._stopped = True + # Linearization point: long-running accepted requests may finish + # their read/import work, but every commit path rejects afterwards. + self._admission.close() + thread = self._thread + if self._started: + self.httpd.shutdown() + # server_close is sufficient before start (shutdown would deadlock). + self.httpd.server_close() + if thread and thread is not threading.current_thread(): + thread.join(timeout=5) + if thread.is_alive(): + log.warning("HTTP server thread did not stop within 5s") + if not self._admission.wait_for_handlers(HANDLER_DRAIN_SECONDS): + log.warning( + "HTTP handlers did not drain within %ss; late mutations remain rejected", + HANDLER_DRAIN_SECONDS, + ) def _load_web(name: str) -> str: return resources.files("agentbar").joinpath(f"web/{name}").read_text("utf-8") -def _make_handler(core: Scheduler, settings: Settings, hooks: dict | None = None, - extra_hosts: set | None = None): +def _make_handler( + core: Scheduler, + settings: Settings, + hooks: dict | None = None, + extra_hosts: set | None = None, + admission: _RequestAdmission | None = None, +): hooks = hooks if hooks is not None else {} extra_hosts = extra_hosts if extra_hosts is not None else set() + admission = admission if admission is not None else _RequestAdmission() class Handler(BaseHTTPRequestHandler): server_version = f"AgentBar/{__version__}" + sys_version = "" # ---------- plumbing ---------- + def setup(self) -> None: + self._admission_entered = False + super().setup() + self._admission_entered = admission.enter_handler() + # Bound slow header/body clients too. In-process work is governed by + # the admission predicate and our separate bounded handler drain. + self.connection.settimeout(REQUEST_IO_TIMEOUT_SECONDS) + + def finish(self) -> None: + try: + super().finish() + finally: + if self._admission_entered: + self._admission_entered = False + admission.leave_handler() + + def _reject_if_stopping(self) -> bool: + if not self._admission_entered or admission.stopping.is_set(): + try: + self._json(503, { + "ok": False, + "error": "AgentBar 正在退出，拒绝新的操作", + }) + except OSError: + pass + return True + return False + def log_message(self, fmt, *args): # 安静，不刷 stderr log.debug("http: " + fmt, *args) @@ -201,6 +495,7 @@ def _json(self, code: int, obj: dict) -> None: self.send_header("Content-Type", "application/json; charset=utf-8") self.send_header("Content-Length", str(len(body))) self.send_header("Cache-Control", "no-store") + self._security_headers() self.end_headers() self.wfile.write(body) @@ -209,9 +504,19 @@ def _html(self, text: str) -> None: self.send_response(200) self.send_header("Content-Type", "text/html; charset=utf-8") self.send_header("Content-Length", str(len(body))) + self.send_header("Cache-Control", "no-store") + self.send_header("Content-Security-Policy", _HTML_CSP) + self._security_headers() self.end_headers() self.wfile.write(body) + def _security_headers(self) -> None: + self.send_header("X-Content-Type-Options", "nosniff") + self.send_header("X-Frame-Options", "DENY") + self.send_header("Referrer-Policy", "no-referrer") + self.send_header("Permissions-Policy", "camera=(), microphone=(), geolocation=()") + self.send_header("Cross-Origin-Opener-Policy", "same-origin") + def _host_ok(self) -> bool: host = _host_of(self.headers.get("Host")) if host in ALLOWED_HOSTS or host in extra_hosts: @@ -221,24 +526,77 @@ def _host_ok(self) -> bool: # 公网隧道域名走 extra_hosts 动态注册。 return settings.lan_access and _is_ip_literal(host) - def _authed(self, query: dict) -> bool: - token = self.headers.get("X-Agentbar-Token") or ( - query.get("token", [""])[0] - ) + def _authed(self, _query: dict) -> bool: + # Query token 会泄漏到访问日志/历史/Referer，只接受请求头。 + token = self.headers.get("X-Agentbar-Token") or "" return bool(token) and hmac.compare_digest(token, settings.token) def _body(self) -> dict: - n = int(self.headers.get("Content-Length") or 0) - if n <= 0 or n > MAX_BODY: - return {} try: - return json.loads(self.rfile.read(n).decode("utf-8")) - except (json.JSONDecodeError, UnicodeDecodeError): + n = int(self.headers.get("Content-Length") or 0) + except (TypeError, ValueError): + raise _BodyError(400, "Content-Length 无效") + if n < 0: + raise _BodyError(400, "Content-Length 无效") + if n == 0: return {} + if n > MAX_BODY: + # Consume one bounded body window before answering. Returning + # immediately can close the socket while a normal client is + # still sending a just-over-limit payload, turning the useful + # 413 into a platform-dependent connection reset. Never drain + # an attacker-declared unbounded length; socket timeout still + # caps slow senders. + try: + self.rfile.read(min(n, MAX_BODY + 1)) + except OSError: + pass + raise _BodyError(413, f"请求体过大（上限 {MAX_BODY} 字节）") + try: + value = json.loads(self.rfile.read(n).decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError, OSError): + raise _BodyError(400, "请求体必须是合法 JSON object") + if not isinstance(value, dict): + raise _BodyError(400, "请求体必须是 JSON object") + return value + + def _persistence_unavailable(self, operation: str) -> None: + # Never echo exception details: filesystem paths or secret-bearing + # payload fragments may be present in lower-layer errors. + log.exception("%s persistence failed", operation) + self._json(503, { + "ok": False, + "error": "本地状态保存失败，操作未生效；请检查磁盘后重试", + }) + + def _trusted_local(self) -> bool: + """Only a direct local request may touch credentials or local UI hooks. + + Cloudflared and other reverse proxies connect from 127.0.0.1 too, so + peer IP alone is not a trust boundary. Require a loopback Host and + reject all common proxy provenance headers as well. + """ + try: + peer_loopback = ipaddress.ip_address(self.client_address[0]).is_loopback + except (ValueError, IndexError, TypeError): + return False + if not peer_loopback or _host_of(self.headers.get("Host")) not in ALLOWED_HOSTS: + return False + for name in self.headers.keys(): + lowered = name.lower() + if ( + lowered in _PROXY_HEADERS + or lowered.startswith("cf-") + or lowered.startswith("x-forwarded-") + ): + return False + return True # ---------- routing ---------- def do_GET(self): + if self._reject_if_stopping(): + return if not self._host_ok(): self._json(403, {"ok": False, "error": "bad host"}) return @@ -275,12 +633,32 @@ def do_GET(self): "allow_full_profile": settings.allow_full_profile}) return if path == "/api/provider-config": + if not self._trusted_local(): + self._json(403, { + "ok": False, + "error": "额度与凭据配置只能在本机 AgentBar 面板中查看", + }) + return self._json(200, _provider_config_payload(settings)) return parts = path.split("/") if len(parts) == 5 and parts[1:3] == ["api", "tasks"] and parts[4] == "log": - tail = min(int(q.get("tail_bytes", ["30000"])[0]), 200_000) - text = core.store.read_log_tail(parts[3], tail) + task_id = parts[3] + with core._lock: + task_exists = task_id in core._tasks + if not task_exists: + self._json(404, {"ok": False, "error": "任务不存在"}) + return + try: + tail = int(q.get("tail_bytes", ["30000"])[0]) + except (TypeError, ValueError): + tail = 30_000 + tail = min(max(0, tail), 200_000) + try: + text = core.store.read_log_tail(task_id, tail) + except ValueError: + self._json(400, {"ok": False, "error": "任务 ID 无效"}) + return self._json(200, {"ok": True, "log": text}) return if len(parts) == 5 and parts[1:3] == ["api", "tasks"] and parts[4] == "transcript": @@ -300,8 +678,21 @@ def do_GET(self): if sid: with core._lock: if not t.session_id: + try: + admission.require_open() + except _AdmissionClosed: + self._reject_if_stopping() + return + checkpoint = t.to_dict() t.session_id = sid - core._persist_locked() + try: + core._persist_locked() + except OSError: + core._restore_task_locked(t, checkpoint) + self._persistence_unavailable( + "recovered transcript session" + ) + return if not t.session_id: self._json(200, {"ok": True, "transcript": "", "message": "该任务尚无会话 ID"}) return @@ -317,6 +708,8 @@ def do_GET(self): self._json(404, {"ok": False, "error": "not found"}) def do_POST(self): + if self._reject_if_stopping(): + return if not self._host_ok(): self._json(403, {"ok": False, "error": "bad host"}) return @@ -326,10 +719,26 @@ def do_POST(self): if not self._authed(q): self._json(401, {"ok": False, "error": "unauthorized"}) return - body = self._body() + # Reject secret/UI-local routes before reading their request body. + # A proxied browser must never upload a credential to Cloudflare or + # a LAN hop only to receive a 403 after the bytes have been consumed. + local_only = { + "/api/provider-config", + "/api/provider-config/import-cookie", + "/api/debug/dispatch", + } + if path in local_only and not self._trusted_local(): + self._json(403, {"ok": False, "error": "该操作仅允许本机直连"}) + return + try: + body = self._body() + except _BodyError as e: + self._json(e.status, {"ok": False, "error": str(e)}) + return if path == "/api/tasks": try: + admission.require_open() t = core.add_task( prompt=body.get("prompt", ""), tool=body.get("tool", "claude"), @@ -339,45 +748,105 @@ def do_POST(self): model=body.get("model"), effort=body.get("effort"), scheduled_at=body.get("scheduled_at"), + before_commit=admission.require_open, ) + except _AdmissionClosed: + self._reject_if_stopping() + return except ValueError as e: self._json(400, {"ok": False, "error": str(e)}) return + except OSError: + self._persistence_unavailable("add task") + return self._json(200, {"ok": True, "task": t.to_dict()}) return if path == "/api/pause-all": - core.pause_all() + try: + admission.require_open() + core.pause_all(before_commit=admission.require_open) + except _AdmissionClosed: + self._reject_if_stopping() + return + except ValueError as e: + self._json(400, {"ok": False, "error": str(e)}) + return + except OSError: + self._persistence_unavailable("pause all") + return self._json(200, {"ok": True}) return if path == "/api/resume-all": - core.resume_all() + try: + admission.require_open() + core.resume_all(before_commit=admission.require_open) + except _AdmissionClosed: + self._reject_if_stopping() + return + except ValueError as e: + self._json(400, {"ok": False, "error": str(e)}) + return + except OSError: + self._persistence_unavailable("resume all") + return self._json(200, {"ok": True}) return if path == "/api/quota/refresh": - core.quota.refresh_now() - self._json(202, {"ok": True, "message": "额度刷新已触发"}) - return - if path == "/api/quota/authorize-claude": - # 仅由用户在本机面板主动调用，后台刷新不会触发 Keychain 弹窗。 - ok = core.quota.authorize_claude_keychain() - self._json(200 if ok else 400, { - "ok": ok, - "message": "Claude Keychain 已授权并刷新" if ok else "未取得 Claude Keychain 授权", + tool = str(body.get("tool") or "").strip() + if not tool: + self._json(400, { + "ok": False, + "error": "必须指定一个已启用的额度来源", + }) + return + if tool not in core.quota.provider_tools(): + self._json(400, {"ok": False, "error": f"未启用的额度来源: {tool!r}"}) + return + try: + admission.require_open() + except _AdmissionClosed: + self._reject_if_stopping() + return + core.quota.refresh_now(tool) + self._json(202, { + "ok": True, + "message": f"{tool} 额度刷新已触发", }) return if path == "/api/provider-config": + if not self._trusted_local(): + self._json(403, { + "ok": False, + "error": "额度与凭据配置只能在本机 AgentBar 面板中保存", + }) + return try: - _apply_provider_settings(settings, body) + _apply_provider_settings( + settings, body, before_commit=admission.require_open, + ) core.quota.reload_fetchers() - except Exception as e: + except _AdmissionClosed: + self._reject_if_stopping() + return + except ValueError as e: self._json(400, {"ok": False, "error": str(e)}) return + except OSError: + self._persistence_unavailable("provider config") + return + except Exception: + log.exception("provider config save failed") + self._json(500, {"ok": False, "error": "配置保存失败"}) + return self._json(200, { **_provider_config_payload(settings), "message": "额度配置已保存并刷新", }) return if path == "/api/provider-config/import-cookie": + if not self._trusted_local(): + self._json(403, {"ok": False, "error": "Cookie 导入仅允许本机直连"}) + return provider = str(body.get("provider") or "") host = PROVIDER_HOSTS.get(provider) if not host: @@ -388,13 +857,44 @@ def do_POST(self): except CookieImportError as e: self._json(400, {"ok": False, "error": str(e)}) return - cfg = settings.providers.setdefault( - provider, dict(DEFAULT_PROVIDERS[provider]) - ) - cfg["enabled"] = True - cfg["cookie"] = imported.header - save_settings(settings) - core.quota.reload_fetchers() + try: + with settings._lock: + # Cookie/Keychain scanning is intentionally outside the + # settings lock and can outlive bounded server drain. + # Re-check at the real commit point so a late result can + # never resurrect credentials during/after shutdown. + admission.require_open() + old_providers = settings.providers + providers = json.loads(json.dumps( + settings.providers or DEFAULT_PROVIDERS + )) + cfg = providers.setdefault( + provider, dict(DEFAULT_PROVIDERS[provider]) + ) + cfg["enabled"] = True + cfg["cookie"] = _clean_credential_field( + imported.header, "Cookie", 65_536, + ) + settings.providers = providers + try: + save_settings(settings) + except Exception: + settings.providers = old_providers + raise + core.quota.reload_fetchers() + except _AdmissionClosed: + self._reject_if_stopping() + return + except ValueError as e: + self._json(400, {"ok": False, "error": str(e)}) + return + except OSError: + self._persistence_unavailable("provider cookie import") + return + except Exception: + log.exception("provider cookie import save failed") + self._json(500, {"ok": False, "error": "Cookie 保存失败"}) + return self._json(200, { **_provider_config_payload(settings), "message": ( @@ -406,27 +906,54 @@ def do_POST(self): if path == "/api/debug/dispatch": # 触发与真实菜单点击完全相同的 _dispatch 路径（主线程执行）， # 用于无 GUI 交互的端到端验证。白名单限定只读性动作。 + if not self._trusted_local(): + self._json(403, {"ok": False, "error": "本机 UI 调试通道仅允许本机直连"}) + return fn = hooks.get("dispatch") action = str(body.get("action") or "") if fn is None: self._json(404, {"ok": False, "error": "menu bar 未运行（headless 无此通道）"}) return - if action not in {"open_panel", "quick_add", "refresh_quota", - "tunnel_start", "tunnel_stop"}: + scoped_refresh = ( + action.startswith("refresh_quota:") + and bool(action.split(":", 1)[1].strip()) + ) + allowed = action in { + "open_panel", "quick_add", "provider_settings", + "tunnel_start", "tunnel_stop", + } or scoped_refresh + if not allowed: self._json(400, {"ok": False, "error": f"action 不在白名单: {action!r}"}) return + try: + admission.require_open() + except _AdmissionClosed: + self._reject_if_stopping() + return fn(action) self._json(202, {"ok": True, "action": action}) return parts = path.split("/") if len(parts) == 5 and parts[1:3] == ["api", "tasks"]: - ok, msg = core.act(parts[3], parts[4]) + try: + admission.require_open() + ok, msg = core.act( + parts[3], parts[4], before_commit=admission.require_open, + ) + except _AdmissionClosed: + self._reject_if_stopping() + return + except OSError: + self._persistence_unavailable("task action") + return self._json(200 if ok else 400, {"ok": ok, "message": msg, "error": msg}) return self._json(404, {"ok": False, "error": "not found"}) def do_PUT(self): + if self._reject_if_stopping(): + return if not self._host_ok(): self._json(403, {"ok": False, "error": "bad host"}) return @@ -442,10 +969,24 @@ def do_PUT(self): self._json(404, {"ok": False, "error": "not found"}) return try: - t = core.edit_task(parts[3], self._body()) + body = self._body() + except _BodyError as e: + self._json(e.status, {"ok": False, "error": str(e)}) + return + try: + admission.require_open() + t = core.edit_task( + parts[3], body, before_commit=admission.require_open, + ) + except _AdmissionClosed: + self._reject_if_stopping() + return except ValueError as e: self._json(400, {"ok": False, "error": str(e)}) return + except OSError: + self._persistence_unavailable("edit task") + return self._json(200, {"ok": True, "task": t.to_dict()}) _INDEX_HTML = _load_web("index.html") diff --git a/agentbar/store.py b/agentbar/store.py index 3ae5cca..4df99c9 100644 --- a/agentbar/store.py +++ b/agentbar/store.py @@ -2,13 +2,30 @@ from __future__ import annotations +import fcntl import json import os +import re +import stat +import tempfile import threading import time +from contextlib import contextmanager from pathlib import Path +_PRIVATE_STATE_FILES = { + "state.json", + "runtime.json", + "menu-debug.json", + "agentbar.log", + "launchd.stdout.log", + "launchd.stderr.log", +} +_TASK_LOG_NAME = re.compile(r"[A-Za-z0-9_-]{1,64}\.log") +_CORRUPT_STATE_NAME = re.compile(r"state\.json\.corrupt-[0-9]+") + + class StateStore: def __init__(self, state_dir: Path): self.state_dir = Path(state_dir) @@ -17,7 +34,84 @@ def __init__(self, state_dir: Path): self.logs_dir = self.state_dir / "logs" self.state_dir.mkdir(parents=True, exist_ok=True) self.logs_dir.mkdir(parents=True, exist_ok=True) - self._lock = threading.Lock() + for directory in (self.state_dir, self.logs_dir): + try: + os.chmod(directory, stat.S_IRWXU) + except OSError: + pass + self._repair_existing_permissions() + self._lock = threading.RLock() + + @staticmethod + def _repair_private_file(path: Path) -> None: + """Repair a known private regular file without following symlinks.""" + flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) + flags |= getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0) + try: + fd = os.open(path, flags) + except OSError: + return + try: + if stat.S_ISREG(os.fstat(fd).st_mode): + os.fchmod(fd, stat.S_IRUSR | stat.S_IWUSR) + except OSError: + pass + finally: + os.close(fd) + + def _repair_existing_permissions(self) -> None: + """Self-heal permissions left by older releases, preserving contents.""" + try: + state_entries = tuple(self.state_dir.iterdir()) + except OSError: + state_entries = () + for path in state_entries: + if path.name in _PRIVATE_STATE_FILES or _CORRUPT_STATE_NAME.fullmatch(path.name): + self._repair_private_file(path) + + try: + log_entries = tuple(self.logs_dir.iterdir()) + except OSError: + log_entries = () + for path in log_entries: + if _TASK_LOG_NAME.fullmatch(path.name): + self._repair_private_file(path) + + @staticmethod + def _safe_task_id(task_id: str) -> str: + if not isinstance(task_id, str) or not re.fullmatch( + r"[A-Za-z0-9_-]{1,64}", task_id, + ): + raise ValueError("invalid task id") + return task_id + + @staticmethod + def _write_private_json(path: Path, payload: dict) -> None: + fd, tmp = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(path.parent)) + try: + os.fchmod(fd, stat.S_IRUSR | stat.S_IWUSR) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + fd = -1 + json.dump(payload, handle, ensure_ascii=False, indent=1) + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp, path) + # mkstemp + fchmod already made the committed inode private. Keep + # this repair best-effort so a post-commit metadata error cannot be + # reported as an uncommitted write to transactional callers. + try: + os.chmod(path, stat.S_IRUSR | stat.S_IWUSR) + except OSError: + pass + tmp = "" + finally: + if fd >= 0: + os.close(fd) + if tmp: + try: + os.unlink(tmp) + except OSError: + pass # ---------- state.json ---------- @@ -26,49 +120,77 @@ def load(self) -> dict: if not self.state_path.exists(): return {} try: - return json.loads(self.state_path.read_text(encoding="utf-8")) + data = json.loads(self.state_path.read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError): - # 损坏的状态文件：备份后从空状态启动，不让整个调度器起不来 - backup = self.state_path.with_name( - f"state.json.corrupt-{int(time.time())}" - ) - try: - os.replace(self.state_path, backup) - except OSError: - pass - return {} + data = None + if isinstance(data, dict): + return data + # Invalid JSON and valid non-object JSON are both unusable scheduler + # states. Back either up, then boot empty instead of crashing in + # Scheduler._load() on a missing dict interface. + backup = self.state_path.with_name( + f"state.json.corrupt-{time.time_ns()}" + ) + try: + os.replace(self.state_path, backup) + except OSError: + pass + return {} def save(self, data: dict) -> None: with self._lock: - tmp = self.state_path.with_suffix(".json.tmp") - tmp.write_text( - json.dumps(data, ensure_ascii=False, indent=1), encoding="utf-8" - ) - os.replace(tmp, self.state_path) + self._write_private_json(self.state_path, data) # ---------- per-task logs ---------- def log_path(self, task_id: str) -> Path: - return self.logs_dir / f"{task_id}.log" + return self.logs_dir / f"{self._safe_task_id(task_id)}.log" + + def open_log_append(self, task_id: str): + """Open a task transcript without ever creating a world-readable file.""" + path = self.log_path(task_id) + flags = os.O_WRONLY | os.O_CREAT | os.O_APPEND | getattr(os, "O_CLOEXEC", 0) + flags |= getattr(os, "O_NOFOLLOW", 0) + fd = os.open( + path, + flags, + stat.S_IRUSR | stat.S_IWUSR, + ) + try: + if not stat.S_ISREG(os.fstat(fd).st_mode): + raise OSError("task log is not a regular file") + os.fchmod(fd, stat.S_IRUSR | stat.S_IWUSR) + except Exception: + os.close(fd) + raise + return os.fdopen(fd, "ab") + + def _open_log_read(self, task_id: str): + """Open an existing regular task log without following symlinks.""" + path = self.log_path(task_id) + flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) + flags |= getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0) + fd = os.open(path, flags) + try: + if not stat.S_ISREG(os.fstat(fd).st_mode): + raise OSError("task log is not a regular file") + except Exception: + os.close(fd) + raise + return os.fdopen(fd, "rb") def read_log_head(self, task_id: str, max_bytes: int = 16384) -> str: """输出开头（codex 的 session id 打在 header，长输出会被挤出 tail）。""" - p = self.log_path(task_id) - if not p.exists(): - return "" try: - with open(p, "rb") as f: + with self._open_log_read(task_id) as f: return f.read(max_bytes).decode("utf-8", errors="replace") except OSError: return "" def read_log_tail(self, task_id: str, max_bytes: int = 65536) -> str: - p = self.log_path(task_id) - if not p.exists(): - return "" try: - size = p.stat().st_size - with open(p, "rb") as f: + with self._open_log_read(task_id) as f: + size = os.fstat(f.fileno()).st_size if size > max_bytes: f.seek(size - max_bytes) data = f.read() @@ -76,24 +198,79 @@ def read_log_tail(self, task_id: str, max_bytes: int = 65536) -> str: except OSError: return "" + def delete_log(self, task_id: str) -> None: + """Delete one retired task log without following a malicious symlink.""" + path = self.log_path(task_id) + with self._lock: + # unlink(2) removes the directory entry itself. If an attacker has + # replaced it with a symlink, its target is never followed. + path.unlink(missing_ok=True) + # ---------- runtime.json (实际端口/PID，供 CLI 客户端发现) ---------- def write_runtime(self, port: int) -> None: - self.runtime_path.write_text( - json.dumps({"port": port, "pid": os.getpid(), "started_at": time.time()}), - encoding="utf-8", - ) + with self._lock: + with self._runtime_file_lock(): + self._write_private_json(self.runtime_path, { + "port": port, + "pid": os.getpid(), + "started_at": time.time(), + }) def read_runtime(self) -> dict | None: - if not self.runtime_path.exists(): - return None - try: - return json.loads(self.runtime_path.read_text(encoding="utf-8")) - except (json.JSONDecodeError, OSError): - return None + with self._lock: + if not self.runtime_path.exists(): + return None + try: + data = json.loads(self.runtime_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + return None + if not isinstance(data, dict): + return None + port = data.get("port") + if isinstance(port, bool) or not isinstance(port, int): + return None + if not 1 <= port <= 65_535: + return None + return data - def clear_runtime(self) -> None: + def clear_runtime(self, owner_pid: int | None = None) -> None: + """Delete only this instance's discovery record. + + During a restart, a replacement process may bind/write runtime.json while + the old process is still finishing scheduler shutdown. The cross-process + lock makes check+unlink atomic with writes, and the PID check prevents the + old process from erasing its replacement's record. + """ + expected_pid = os.getpid() if owner_pid is None else owner_pid + with self._lock: + with self._runtime_file_lock(): + try: + data = json.loads(self.runtime_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + return + if not isinstance(data, dict) or data.get("pid") != expected_pid: + return + try: + self.runtime_path.unlink(missing_ok=True) + except OSError: + pass + + @contextmanager + def _runtime_file_lock(self): + """Hold the runtime ownership lock without exposing its descriptor.""" + path = self.runtime_path.with_name(".runtime.lock") + fd = os.open( + path, + os.O_RDWR | os.O_CREAT, + stat.S_IRUSR | stat.S_IWUSR, + ) try: - self.runtime_path.unlink(missing_ok=True) - except OSError: - pass + os.fchmod(fd, stat.S_IRUSR | stat.S_IWUSR) + fcntl.flock(fd, fcntl.LOCK_EX) + yield + finally: + try: + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(fd) diff --git a/agentbar/transcript.py b/agentbar/transcript.py index 29844bc..f4a94a2 100644 --- a/agentbar/transcript.py +++ b/agentbar/transcript.py @@ -7,6 +7,7 @@ import json import os import re +import shlex import time from pathlib import Path @@ -14,12 +15,11 @@ _HTML_TEMPLATE = """\ - - + + {body}""" @@ -146,17 +148,19 @@ def _codex_cwd_matches(path: str, cwd: str) -> bool: def resume_command(tool: str, cwd: str, session_id: str) -> str: + safe_cwd = shlex.quote(cwd) + safe_sid = shlex.quote(session_id) if tool == "claude": - return f"cd {cwd} && claude --resume {session_id}" + return f"cd {safe_cwd} && claude --resume {safe_sid}" if tool == "codex": - return f"cd {cwd} && codex exec resume {session_id} -" + return f"cd {safe_cwd} && codex exec resume {safe_sid} -" return f"# unknown tool {tool}" # ---------- HTML generation (for WKWebView) ---------- def to_html(tool: str, path: Path) -> str: - """Parse session file and return a full HTML document with MathJax support.""" + """Parse a session file into a self-contained, script-free HTML document.""" try: with open(path, encoding="utf-8") as f: lines = f.readlines() @@ -334,7 +338,7 @@ def _codex_content_to_html(content) -> str: return "".join(parts) -# ---------- text → HTML with math passthrough ---------- +# ---------- text → offline HTML ---------- _MD = None @@ -344,15 +348,24 @@ def _markdown(): global _MD if _MD is None: import mistune + + class _SafeRenderer(mistune.HTMLRenderer): + def image(self, text: str, url: str, title: str | None = None) -> str: + # 对话内容不可触发任意远程图片请求（隐私信标/内网探测）。 + label = html.escape(text or "图片") + return f'[已隐藏外部图片：{label}]' + # escape=True：把消息里的原始 HTML 转义掉，只认 markdown 语法 _MD = mistune.create_markdown( - escape=True, plugins=["table", "strikethrough", "math", "url"] + escape=True, + renderer=_SafeRenderer(escape=True), + plugins=["table", "strikethrough", "math", "url"], ) return _MD def _text_to_html(text: str) -> str: - """Markdown → HTML（mistune；math 插件输出 \\(..\\)/\\[..\\] 交给 MathJax）。""" + """Markdown → HTML（mistune；公式仅保留为离线文本标记）。""" try: return _markdown()(text) except Exception: @@ -360,7 +373,7 @@ def _text_to_html(text: str) -> str: def _text_to_html_legacy(text: str) -> str: - """手写降级渲染（mistune 不可用时）：代码块/行内样式/数学穿透。""" + """手写降级渲染（mistune 不可用时）：代码块/行内样式。""" lines = text.splitlines() out: list[str] = [] in_code = False @@ -410,13 +423,15 @@ def _group_paragraphs(lines: list[str]) -> list[list[str]]: def _format_line(line: str) -> str: - """Format one line: preserve math delimiters, escape rest, apply inline markdown.""" + """Format one line while escaping every user-controlled fragment.""" # Split on display math first ($$...$$), then inline math ($...$) parts = re.split(r"(\$\$[^$]*?\$\$|\$(?!\$)[^$\n]*?\$)", line) result: list[str] = [] for p in parts: if p.startswith("$"): - result.append(p) # math: pass through for MathJax + # 降级路径也必须转义公式内容；过去为 MathJax 原样穿透 + # 会让 `$$` 变成可执行 HTML。 + result.append(html.escape(p)) else: result.append(_format_inline(p)) return "".join(result) diff --git a/agentbar/tunnel.py b/agentbar/tunnel.py index 331e173..ce8af00 100644 --- a/agentbar/tunnel.py +++ b/agentbar/tunnel.py @@ -14,6 +14,7 @@ import os import re import shutil +import signal import subprocess import threading from urllib.parse import urlparse @@ -58,7 +59,11 @@ def binary(self) -> str | None: def status(self) -> dict: with self._lock: # 进程意外退出 → 降级为 error（reader 线程也会置，这里兜底） - if self._state == "up" and self._proc and self._proc.poll() is not None: + if ( + self._state in ("starting", "up") + and self._proc + and self._proc.poll() is not None + ): self._mark_down_locked("隧道进程已退出") return {"state": self._state, "url": self._url, "error": self._error, "installed": self.binary() is not None} @@ -88,12 +93,26 @@ def start(self, timeout: float = START_TIMEOUT) -> bool: f"http://127.0.0.1:{self.port}"], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, stdin=subprocess.DEVNULL, + # cloudflared may spawn helpers. A private process group lets + # stop() close every descendant that inherited the output pipe; + # otherwise the reader (and start()) can remain blocked after + # the direct child exits. + start_new_session=True, ) except OSError as e: with self._lock: self._state, self._error = "error", f"启动失败: {e}" return False + # Publish the process immediately. stop() may run while cloudflared is + # still starting; delaying this assignment until the URL appears leaks + # an untracked tunnel process during app shutdown. + with self._lock: + if self._state != "starting": + self._terminate_process(proc) + return False + self._proc = proc + url_evt = threading.Event() def _reader(): @@ -102,29 +121,84 @@ def _reader(): m = _URL_RE.search(line) if m: with self._lock: - self._url = m.group(0) - url_evt.set() + accepted = ( + self._proc is proc and self._state == "starting" + ) + if accepted: + self._url = m.group(0) + if accepted: + url_evt.set() + # Also wake start() when stop() terminates a still-starting process. + url_evt.set() with self._lock: - if self._proc is proc and self._state == "up": + # EOF can race with start() between accepting a URL and + # publishing state=up. Retire both phases atomically; otherwise + # an already-dead process can leave status() stuck at starting. + if self._proc is proc and self._state in ("starting", "up"): self._mark_down_locked("隧道进程已退出") threading.Thread(target=_reader, name="agentbar-tunnel-io", daemon=True).start() if not url_evt.wait(timeout): - proc.terminate() + self._terminate_process(proc) with self._lock: - self._state = "error" - self._error = f"启动超时（{timeout:.0f}s，公司网络可能拦截 Cloudflare）" + if self._proc is proc: + self._proc = None + if self._state == "starting": + self._state = "error" + self._error = f"启动超时（{timeout:.0f}s，公司网络可能拦截 Cloudflare）" return False with self._lock: - self._proc = proc - self._host = urlparse(self._url).hostname - self._state = "up" - host = self._host + if self._proc is not proc or self._state != "starting": + should_stop = True + host = None + elif proc.poll() is not None: + self._mark_down_locked("隧道进程已退出") + should_stop = True + host = None + else: + should_stop = False + self._host = urlparse(self._url).hostname + self._state = "up" + host = self._host + if should_stop: + self._terminate_process(proc) + return False log.info("tunnel up: %s", self._url) + notified_up = False if self._on_up and host: - self._on_up(host) + try: + self._on_up(host) + notified_up = True + except Exception as exc: + # A callback can fail after partially mutating its allow-list. + # Compensate and retire the process; otherwise start() would + # raise while leaving an untracked public tunnel alive. + log.exception("tunnel on_up callback failed") + with self._lock: + if self._proc is proc and self._state == "up": + self._proc, self._url, self._host = None, None, None + self._state = "error" + self._error = f"隧道注册失败: {exc}" + self._notify_down(host) + self._terminate_process(proc) + return False + # stop() may run after we publish state=up but while an arbitrary on_up + # callback is still executing. Re-check after notification and compensate + # a late allow-list add; callbacks must never run while holding _lock. + with self._lock: + still_up = ( + self._proc is proc + and self._state == "up" + and self._host == host + and proc.poll() is None + ) + if not still_up: + if notified_up and self._on_down and host: + self._notify_down(host) + self._terminate_process(proc) + return False return True def stop(self) -> None: @@ -133,17 +207,37 @@ def stop(self) -> None: host = self._host self._state, self._url, self._host, self._error = "off", None, None, "" if proc and proc.poll() is None: - proc.terminate() - try: - proc.wait(5) - except subprocess.TimeoutExpired: - proc.kill() + self._terminate_process(proc) if self._on_down and host: - self._on_down(host) + self._notify_down(host) log.info("tunnel stopped") # ---------- internal ---------- + @staticmethod + def _terminate_process(proc: subprocess.Popen) -> None: + if proc.poll() is not None: + return + try: + os.killpg(proc.pid, signal.SIGTERM) + except ProcessLookupError: + return + except OSError: + proc.terminate() + try: + proc.wait(5) + except subprocess.TimeoutExpired: + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + return + except OSError: + proc.kill() + try: + proc.wait(3) + except subprocess.TimeoutExpired: + log.warning("cloudflared did not exit after SIGKILL") + def _mark_down_locked(self, reason: str) -> None: """caller must hold self._lock""" host = self._host @@ -151,5 +245,15 @@ def _mark_down_locked(self, reason: str) -> None: self._state, self._error = "error", reason if self._on_down and host: threading.Thread( - target=self._on_down, args=(host,), daemon=True + target=self._notify_down, args=(host,), daemon=True ).start() + + def _notify_down(self, host: str) -> None: + """Best-effort compensating callback; never obstruct process cleanup.""" + callback = self._on_down + if callback is None: + return + try: + callback(host) + except Exception: + log.exception("tunnel on_down callback failed for %s", host) diff --git a/agentbar/usage.py b/agentbar/usage.py index 6428be9..c658d33 100644 --- a/agentbar/usage.py +++ b/agentbar/usage.py @@ -1,16 +1,12 @@ """Real quota/usage fetchers — approach mirrored from ylab/aiusagebar (Swift). Claude: GET https://api.anthropic.com/api/oauth/usage - 凭据链: env CLAUDE_CODE_OAUTH_TOKEN → 本地缓存 ~/.agentbar/claude_credentials.json - → ~/.claude/.credentials.json - → macOS Keychain "Claude Code-credentials"（只缓存短期 access token； - refresh token 仍由 Claude Code 独占续期，避免竞态与反复弹窗） + AgentBar 运行时只使用用户在额度设置中显式保存的 OAuth Access Token； 响应: {five_hour|seven_day|seven_day_opus|seven_day_sonnet: {utilization: 0-100, resets_at: ISO8601}} Codex: GET https://chatgpt.com/backend-api/wham/usage - 凭据: $CODEX_HOME/auth.json（默认 ~/.codex/auth.json）tokens.access_token - + chatgpt-account-id（tokens.account_id 或 JWT claim） + AgentBar 运行时使用显式保存的 OAuth Access Token，以及可选 Account ID； 响应: {rate_limits: {primary|secondary: {used_percent: 0-100, resets_at: epoch_s, window_duration_mins}}} @@ -22,13 +18,12 @@ import base64 import json import logging -import os +import re import time import urllib.error import urllib.request from dataclasses import dataclass, field from datetime import datetime -from pathlib import Path log = logging.getLogger("agentbar.usage") @@ -44,12 +39,17 @@ class UsageWindow: used: float | None = None # 已用（credits 或 token 数） total: float | None = None # 总额度 unit: str | None = None # "credits" | "percent" | "token" + # None 表示账户级通用窗口；有值时只限对应模型家族/额度桶。 + model: str | None = None + limited: bool = False def to_dict(self) -> dict: return { "label": self.label, "used_percent": round(self.used_percent, 1), "resets_at": self.resets_at, + "model": self.model, + "limited": self.limited, "used": self.used, "total": self.total, "unit": self.unit, @@ -65,6 +65,8 @@ class UsageSnapshot: fetched_at: float = field(default_factory=time.time) error: str | None = None limited: bool = False + model: str | None = None + available_models: list[str] = field(default_factory=list) @property def primary(self) -> UsageWindow | None: @@ -79,6 +81,8 @@ def to_dict(self) -> dict: "fetched_at": self.fetched_at, "error": self.error, "limited": self.limited, + "model": self.model, + "available_models": self.available_models, } @@ -103,103 +107,14 @@ def _jwt_payload(token: str) -> dict: try: part = token.split(".")[1] part += "=" * (-len(part) % 4) - return json.loads(base64.urlsafe_b64decode(part)) + payload = json.loads(base64.urlsafe_b64decode(part)) + return payload if isinstance(payload, dict) else {} except Exception: return {} # ================= Claude ================= -KEYCHAIN_SERVICE = "Claude Code-credentials" - -# 本地凭据缓存：Keychain 成功读到一次后写回这里（0600）。只保存 -# 短期 access token 和“已授权”标记，故意不保存/使用 refresh token： -# Claude Code 的 refresh token 会轮换，多进程独立续期可能导致登录态失效。 -CLAUDE_CRED_CACHE = Path.home() / ".agentbar" / "claude_credentials.json" - - -def _cache_read_credentials() -> dict | None: - """读取本地凭据缓存；过期凭据也保留，用于记录授权状态。""" - try: - creds = json.loads(CLAUDE_CRED_CACHE.read_bytes()) - except (OSError, json.JSONDecodeError): - return None - if not isinstance(creds, dict): - return None - if not (creds.get("token") or "").strip(): - return None - return creds - - -def _cache_write_credentials(creds: dict) -> None: - """把凭据写入本地缓存文件（0600 权限，仅当前用户可读）。""" - try: - CLAUDE_CRED_CACHE.parent.mkdir(parents=True, exist_ok=True) - tmp = CLAUDE_CRED_CACHE.with_suffix(".tmp") - # 允许列表会把旧版缓存里的 refresh_token 一并清理掉。 - safe = { - key: creds.get(key) - for key in ("token", "expires_at", "plan", "scopes", "keychain_authorized") - if creds.get(key) is not None - } - tmp.write_text(json.dumps(safe), encoding="utf-8") - os.chmod(tmp, 0o600) - os.replace(tmp, CLAUDE_CRED_CACHE) - except OSError: - log.warning("写入 Claude 凭据缓存失败: %s", CLAUDE_CRED_CACHE, exc_info=True) - - -def _keychain_read(interactive: bool = False) -> bytes | None: - """静默读取 Keychain（interactive=True 允许系统弹窗授权，仅由用户显式触发）。""" - try: - import Security # pyobjc-framework-Security - except ImportError: - return None - query = { - Security.kSecClass: Security.kSecClassGenericPassword, - Security.kSecAttrService: KEYCHAIN_SERVICE, - Security.kSecMatchLimit: Security.kSecMatchLimitOne, - Security.kSecReturnData: True, - Security.kSecUseAuthenticationUI: ( - Security.kSecUseAuthenticationUIAllow - if interactive - else Security.kSecUseAuthenticationUIFail - ), - } - status, data = Security.SecItemCopyMatching(query, None) - if status != 0 or data is None: - return None - return bytes(data) - - -def _parse_claude_credentials(raw: bytes) -> dict | None: - try: - oauth = json.loads(raw.decode("utf-8")).get("claudeAiOauth") or {} - except (json.JSONDecodeError, UnicodeDecodeError): - return None - token = (oauth.get("accessToken") or "").strip() - if not token: - return None - expires_ms = oauth.get("expiresAt") - return { - "token": token, - "expires_at": (expires_ms / 1000.0) if expires_ms else None, - "plan": oauth.get("subscriptionType") or oauth.get("rateLimitTier"), - "scopes": oauth.get("scopes"), - "keychain_authorized": True, - } - - -def _credentials_expired(creds: dict) -> bool: - """凭据的 access token 是否已过期。""" - exp = creds.get("expires_at") - if not exp: - return False - try: - return float(exp) <= time.time() - except (TypeError, ValueError): - return True - class ClaudeUsageFetcher: tool = "claude" @@ -211,53 +126,21 @@ class ClaudeUsageFetcher: ("seven_day_sonnet", "7d Sonnet"), ] - def load_credentials(self, interactive: bool = False) -> dict | None: - for env_key in ("CLAUDE_CODE_OAUTH_TOKEN", "CODEXBAR_CLAUDE_OAUTH_TOKEN"): - token = (os.environ.get(env_key) or "").strip() - if token: - return {"token": token, "expires_at": None, "plan": None} - # 本地缓存未过期时直接使用，完全不碰 Keychain。 - cached = _cache_read_credentials() - if cached and not _credentials_expired(cached): - return cached - cred_file = Path.home() / ".claude" / ".credentials.json" - if cred_file.exists(): - try: - creds = _parse_claude_credentials(cred_file.read_bytes()) - except OSError: - creds = None - if creds and not _credentials_expired(creds): - _cache_write_credentials(creds) - return creds - # access token 过期后只做静默同步（UIFail）。Claude Code 自己运行 - # 时会按其内置锁与 Keychain 写回机制续期，AgentBar 不参与轮换。 - raw = _keychain_read(interactive=interactive) - if raw: - creds = _parse_claude_credentials(raw) - if creds: - _cache_write_credentials(creds) - return creds - # 已经授权过的旧缓存保留为状态信号：不再误报“需要授权”。 - return cached if cached and cached.get("keychain_authorized") else None + def __init__(self, access_token: str = "", model: str = ""): + self.access_token = (access_token or "").strip() + self.model = (model or "").strip() + + def load_credentials(self) -> dict | None: + if not self.access_token: + return None + return {"token": self.access_token, "plan": None} def fetch(self, interactive: bool = False) -> UsageSnapshot | None: - creds = self.load_credentials(interactive=interactive) + creds = self.load_credentials() if not creds: return UsageSnapshot( self.tool, source="oauth_api", - error=("未读到可用的 Claude 凭据（只需手动授权 Keychain 一次，" - "后续静默同步）"), - ) - if _credentials_expired(creds): - if creds.get("keychain_authorized"): - return UsageSnapshot( - self.tool, source="oauth_api", - error=("Claude access token 已过期；Claude Code 下次运行时会安全续期，" - "AgentBar 将静默同步"), - ) - return UsageSnapshot( - self.tool, source="oauth_api", - error="Claude OAuth 凭据已过期，请运行 claude 重新登录", + error="未配置 Claude OAuth Access Token", ) headers = { "Authorization": f"Bearer {creds['token']}", @@ -277,7 +160,17 @@ def fetch(self, interactive: bool = False) -> UsageSnapshot | None: def parse(self, data: dict, plan: str | None = None) -> UsageSnapshot: windows = [] + selected = self.model.casefold() + selected_family = ( + "opus" if "opus" in selected else "sonnet" if "sonnet" in selected else "" + ) for key, label in self._WINDOW_KEYS: + # 选中 Opus/Sonnet 时，保留账户通用窗口，只隐藏另一个 + # 模型家族的专属周窗口。其他模型仍展示账户通用额度。 + if key.startswith("seven_day_") and ( + not selected_family or not key.endswith(selected_family) + ): + continue w = data.get(key) if not isinstance(w, dict) or w.get("utilization") is None: continue @@ -285,8 +178,16 @@ def parse(self, data: dict, plan: str | None = None) -> UsageSnapshot: label=label, used_percent=max(0.0, min(100.0, float(w["utilization"]))), resets_at=_parse_iso(w.get("resets_at")), + model=selected_family if key.startswith("seven_day_") else None, )) - snap = UsageSnapshot(self.tool, windows=windows, plan=plan, source="oauth_api") + snap = UsageSnapshot( + self.tool, + windows=windows, + plan=plan, + source="oauth_api", + model=self.model or None, + available_models=["opus", "sonnet"], + ) if not windows: snap.error = "usage 接口未返回可识别的额度窗口" return snap @@ -299,46 +200,30 @@ class CodexUsageFetcher: tool = "codex" URL = "https://chatgpt.com/backend-api/wham/usage" - @staticmethod - def _auth_path() -> Path: - home = (os.environ.get("CODEX_HOME") or "").strip() - base = Path(home).expanduser() if home else Path.home() / ".codex" - return base / "auth.json" + def __init__(self, access_token: str = "", account_id: str = "", model: str = ""): + self.access_token = (access_token or "").strip() + self.account_id = (account_id or "").strip() + self.model = (model or "").strip() def load_credentials(self) -> dict | None: - p = self._auth_path() - if not p.exists(): + if not self.access_token: return None - try: - data = json.loads(p.read_text(encoding="utf-8")) - except (json.JSONDecodeError, OSError): - return None - tokens = data.get("tokens") or data - token = (tokens.get("access_token") or tokens.get("accessToken") or "").strip() - if not token: - return None - id_token = tokens.get("id_token") or "" - id_payload = _jwt_payload(id_token) - access_payload = _jwt_payload(token) - auth_claim = ( - id_payload.get("https://api.openai.com/auth") - or access_payload.get("https://api.openai.com/auth") - or {} - ) + access_payload = _jwt_payload(self.access_token) + auth_claim = access_payload.get("https://api.openai.com/auth") or {} account_id = ( - tokens.get("account_id") + self.account_id or auth_claim.get("chatgpt_account_id") or access_payload.get("chatgpt_account_id") or access_payload.get("account_id") ) plan = auth_claim.get("chatgpt_plan_type") or access_payload.get("chatgpt_plan_type") - return {"token": token, "account_id": account_id, "plan": plan} + return {"token": self.access_token, "account_id": account_id, "plan": plan} def fetch(self, interactive: bool = False) -> UsageSnapshot | None: creds = self.load_credentials() if not creds: return UsageSnapshot(self.tool, source="wham_api", - error="未读到 ~/.codex/auth.json（先运行 codex login）") + error="未配置 Codex OAuth Access Token") headers = { "Authorization": f"Bearer {creds['token']}", "Accept": "*/*", @@ -358,14 +243,18 @@ def fetch(self, interactive: bool = False) -> UsageSnapshot | None: return UsageSnapshot(self.tool, source="wham_api", error=f"网络错误: {e}") return self.parse(data, plan=creds.get("plan")) - def parse(self, data: dict, plan: str | None = None) -> UsageSnapshot: - """Parse both observed WHAM response shapes. + @staticmethod + def _model_key(value: str) -> str: + return re.sub(r"[^a-z0-9]+", "", (value or "").casefold()) - Older clients expose ``rate_limits.primary`` with minute windows, while - the current ChatGPT-backed response exposes ``rate_limit.primary_window`` - with second windows and a ``limit_reached`` boolean. - """ - limits = data.get("rate_limits") or data.get("rate_limit") or {} + @staticmethod + def _windows( + limits: dict, + *, + model: str | None = None, + limited: bool = False, + label_prefix: str = "", + ) -> list[UsageWindow]: windows = [] for keys, fallback_label in ( (("primary", "primary_window"), "5h"), @@ -384,20 +273,103 @@ def parse(self, data: dict, plan: str | None = None) -> UsageSnapshot: label = f"{round(mins / 60)}h" if mins < 2880 else f"{round(mins / 1440)}d" elif seconds: label = f"{round(seconds / 3600)}h" if seconds < 2880 * 60 else f"{round(seconds / 86400)}d" - reset = w.get("resets_at", w.get("resetsAt")) + reset = w.get( + "resets_at", + w.get("resetsAt", w.get("reset_at", w.get("resetAt"))), + ) if not reset and w.get("reset_after_seconds"): reset = time.time() + float(w["reset_after_seconds"]) windows.append(UsageWindow( - label=label, + label=f"{label_prefix}{label}", used_percent=max(0.0, min(100.0, float(used))), resets_at=float(reset) if reset else None, + model=model, + limited=limited, + )) + return windows + + @staticmethod + def _limited(limits: dict) -> bool: + return ( + bool(limits.get("limit_reached", limits.get("limitReached"))) + or limits.get("allowed") is False + ) + + def parse(self, data: dict, plan: str | None = None) -> UsageSnapshot: + """Parse both observed WHAM response shapes. + + Older clients expose ``rate_limits.primary`` with minute windows, while + the current ChatGPT-backed response exposes ``rate_limit.primary_window`` + with second windows and a ``limit_reached`` boolean. + """ + additional = [ + item for item in (data.get("additional_rate_limits") or []) + if isinstance(item, dict) + ] + # Persist/select the stable metered_feature whenever the endpoint exposes + # one. limit_name is presentation text and can be renamed independently. + available = [ + str(item.get("metered_feature") or item.get("limit_name") or "").strip() + for item in additional + ] + available = [name for name in available if name] + selected_item = None + if self.model: + wanted = self._model_key(self.model) + for item in additional: + names = (str(item.get("limit_name") or ""), str(item.get("metered_feature") or "")) + if wanted and wanted in {self._model_key(name) for name in names}: + selected_item = item + break + if selected_item is None: + choices = "、".join(available[:6]) or "暂无模型专属额度" + return UsageSnapshot( + self.tool, + plan=plan or data.get("plan_type"), + source="wham_api", + error=f"未找到模型额度 {self.model!r}；接口可用：{choices}", + model=self.model, + available_models=available, + ) + + account_limits = data.get("rate_limits") or data.get("rate_limit") or {} + account_limited = self._limited(account_limits) + selected_label = ( + str( + selected_item.get("limit_name") + or selected_item.get("metered_feature") + or self.model + ) + if selected_item is not None + else (self.model or None) + ) + if selected_item is not None: + selected_limits = selected_item.get("rate_limit") or {} + selected_limited = self._limited(selected_limits) + windows = self._windows( + account_limits, + limited=account_limited, + label_prefix="账户 ", + ) + windows.extend(self._windows( + selected_limits, + model=selected_label, + limited=selected_limited, + label_prefix=f"{selected_label} ", )) + else: + selected_limited = False + windows = self._windows(account_limits, limited=account_limited) snap = UsageSnapshot( self.tool, windows=windows, plan=plan or data.get("plan_type"), source="wham_api", - limited=bool(limits.get("limit_reached")), + # 保留 snapshot 级标记给旧 UI；调度决策应按 window.model + # 逐窗口判断，避免模型专属限额污染账户级窗口。 + limited=account_limited or selected_limited, + model=selected_label, + available_models=available, ) if not windows: snap.error = "usage 接口未返回 rate_limits" @@ -489,7 +461,6 @@ def _username(self) -> str: return name def _monthly_tokens(self, username: str) -> float | None: - start = _next_month_start() - 1 # 仅用于本月窗口，取月初到现在 month_start_ms = int(time.mktime(time.strptime(_month_start_str(), "%Y-%m-%d")) * 1000) now_ms = int(time.time() * 1000) url = (f"{self.BASE}/api/v1/billing/usage/token-summary" @@ -597,11 +568,27 @@ def fetch(self, interactive: bool = False) -> UsageSnapshot | None: def get_usage_fetchers(settings=None) -> dict[str, object]: - """内置 claude/codex，外加 config 里 enabled 且已配 cookie 的 corp provider。""" - fetchers: dict[str, object] = { - "claude": ClaudeUsageFetcher(), - "codex": CodexUsageFetcher(), - } + """只创建用户显式启用且已配凭据的额度来源。 + + 没有 Settings 也不创建隐式来源，确保任何调用路径都不会读取 CLI + 登录文件或系统凭据存储。 + """ + fetchers: dict[str, object] = {} + if settings is not None: + sources = getattr(settings, "quota_sources", None) or {} + claude = sources.get("claude") or {} + if claude.get("enabled") and (claude.get("access_token") or "").strip(): + fetchers["claude"] = ClaudeUsageFetcher( + access_token=claude.get("access_token", ""), + model=claude.get("model", ""), + ) + codex = sources.get("codex") or {} + if codex.get("enabled") and (codex.get("access_token") or "").strip(): + fetchers["codex"] = CodexUsageFetcher( + access_token=codex.get("access_token", ""), + account_id=codex.get("account_id", ""), + model=codex.get("model", ""), + ) providers = getattr(settings, "providers", None) or {} for name, cls in _CORP_FETCHERS.items(): cfg = providers.get(name) or {} diff --git a/agentbar/web/index.html b/agentbar/web/index.html index a1d9735..63d1c60 100644 --- a/agentbar/web/index.html +++ b/agentbar/web/index.html @@ -36,6 +36,8 @@ .providerTop input { min-width:0; flex:0; } .providerActions { display:flex; gap:8px; align-items:center; flex-wrap:wrap; margin-top:8px; } .providerCookie { min-height:46px; font-family:ui-monospace,Menlo,monospace; } + .sourceSecret { width:100%; font-family:ui-monospace,Menlo,monospace; } + .providerSubhead { margin:12px 0 7px; } .providerStatus { min-height:18px; margin-top:6px; } .meter { height:6px; border-radius:6px; overflow:hidden; background:#2a2f37; } .meter > i { display:block; height:100%; background:var(--blue); border-radius:inherit; } @@ -91,7 +93,6 @@ Claude: … Codex: …
- @@ -122,21 +123,24 @@

添加任务

-

订阅额度

真实接口数据，默认每 120 秒刷新；获取失败会明确显示，不估算或伪造。 -
+

订阅额度

真实接口数据；默认仅在启动、保存设置或手动点刷新时请求，不会反复轮询。
-
-

内部额度配置

MyToken / Tokenverse 使用本机浏览器登录态；Cookie 只保存在 ~/.agentbar/config.json。 +
+

额度来源配置

Claude / Codex 必须显式输入 OAuth Access Token；密钥不会回显。
+ - +
+
订阅来源 · 保存时仅刷新有变更且凭据就绪的来源；关闭周期刷新后不再自动轮询。
+
+
内部来源 · MyToken / Tokenverse 使用本机浏览器登录态；Cookie 只保存在 ~/.agentbar/config.json。
@@ -171,21 +175,31 @@

本机正在运行的 AI CLI

正常文字") assert "