diff --git a/agent_reach/doctor.py b/agent_reach/doctor.py index 9ee98cf2..5295665f 100644 --- a/agent_reach/doctor.py +++ b/agent_reach/doctor.py @@ -4,45 +4,67 @@ Each channel knows how to check itself. Doctor just collects the results. """ -from typing import Dict +import threading +from concurrent.futures import ThreadPoolExecutor +from typing import Dict, Tuple from rich.markup import escape -from agent_reach.channels import get_all_channels +from agent_reach.channels import Channel, get_all_channels from agent_reach.config import Config from agent_reach.utils.text import scrub_url_credentials +# Channels are registry singletons and check() mutates ch.active_backend, so +# two overlapping doctor runs would cross-contaminate each other's results. +# ponytail: one global lock serializes whole doctor runs (each still probes +# its channels in parallel); per-channel locks if concurrent callers matter. +_doctor_lock = threading.Lock() -def check_all(config: Config) -> Dict[str, dict]: - """Check all channels and return status dict. + +def _check_one(ch: Channel, config: Config) -> Tuple[str, dict]: + """Probe a single channel; return (name, result). Never raises. A single misbehaving channel must never take the whole report down, so per-channel exceptions degrade to status="error". """ - results = {} - for ch in get_all_channels(): - try: - status, message = ch.check(config) - active = getattr(ch, "active_backend", None) - except Exception as e: # noqa: BLE001 — doctor must survive any channel - # Channels are registry singletons: a stale active_backend from a - # previous check must not leak into an errored result. - status = "error" - message = f"体检异常：{e}" - active = None - # Doctor is the final output boundary for both expected channel - # messages and unexpected exceptions. Upstream probe output can echo a - # configured URL, so scrub every path before JSON/text rendering. - message = scrub_url_credentials(message) - results[ch.name] = { - "status": status, - "name": ch.description, - "message": message, - "tier": ch.tier, - "backends": ch.backends, - "active_backend": active, - } - return results + try: + status, message = ch.check(config) + active = getattr(ch, "active_backend", None) + except Exception as e: # noqa: BLE001 — doctor must survive any channel + # Channels are registry singletons: a stale active_backend from a + # previous check must not leak into an errored result. + status, message, active = "error", f"体检异常：{e}", None + # Doctor is the final output boundary for both expected channel + # messages and unexpected exceptions. Upstream probe output can echo a + # configured URL, so scrub every path before JSON/text rendering. + return ch.name, { + "status": status, + "name": ch.description, + "message": scrub_url_credentials(message), + "tier": ch.tier, + "backends": ch.backends, + "active_backend": active, + } + + +def check_all(config: Config) -> Dict[str, dict]: + """Check all channels concurrently and return status dict. + + Each check() does real subprocess/network probing with multi-second + timeouts (e.g. `rdt status` waits up to 10s), so sequential probing makes + doctor needlessly slow — wall time becomes the *sum* of every timeout. + Probing is I/O-bound, so threads sidestep the GIL and wall time collapses + to roughly the single slowest channel. + + Result order follows the channel registry — ThreadPoolExecutor.map() + preserves input order — which format_report() relies on for tiered output. + """ + channels = get_all_channels() + if not channels: + return {} + with _doctor_lock: + with ThreadPoolExecutor(max_workers=min(len(channels), 16)) as pool: + return dict(pool.map(lambda ch: _check_one(ch, config), channels)) def _name_msg(r: dict, escape) -> str: diff --git a/tests/test_doctor.py b/tests/test_doctor.py index e181dd9a..0b38dea6 100644 --- a/tests/test_doctor.py +++ b/tests/test_doctor.py @@ -114,6 +114,76 @@ def test_format_report(self): assert "可选渠道可以解锁" in plain +def test_check_all_probes_channels_concurrently(monkeypatch): + """探测必须并发：所有渠道同时越过屏障才能通过，串行会触发屏障超时。 + + 不做墙钟计时断言（CI 上不稳定）；屏障要求 8 个 check() 同时在跑， + 串行执行时第一个 check() 会等满超时并报 BrokenBarrierError → status=error。 + 同时校验结果顺序仍跟随注册表 —— format_report 的分层渲染依赖此顺序。 + """ + import threading + + n = 8 + barrier = threading.Barrier(n, timeout=5) + + class _BarrierChannel: + tier = 0 + backends = [] + active_backend = None + + def __init__(self, i): + self.name = self.description = f"ch{i}" + + def check(self, config=None): + barrier.wait() + return "ok", "done" + + monkeypatch.setattr( + doctor, "get_all_channels", lambda: [_BarrierChannel(i) for i in range(n)] + ) + + results = doctor.check_all(config=None) + + assert [r["status"] for r in results.values()] == ["ok"] * n + assert list(results.keys()) == [f"ch{i}" for i in range(n)] + + +def test_concurrent_doctor_runs_do_not_cross_contaminate_active_backend(monkeypatch): + """渠道是注册表单例，check() 会改写 active_backend； + + 两个并发的 doctor 调用必须各自拿到自己那轮探测出的后端，不能互相污染。 + """ + import threading + import time + + class _SingletonChannel: + name = "single" + description = "单例渠道" + tier = 0 + backends = ["a", "b"] + active_backend = None + + def check(self, config=None): + self.active_backend = config["token"] + time.sleep(0.05) # 拉宽竞态窗口 + return "ok", "done" + + monkeypatch.setattr(doctor, "get_all_channels", lambda: [_SingletonChannel()]) + + seen = {} + + def run(token): + seen[token] = doctor.check_all({"token": token})["single"]["active_backend"] + + threads = [threading.Thread(target=run, args=(t,)) for t in ("t1", "t2")] + for t in threads: + t.start() + for t in threads: + t.join() + + assert seen == {"t1": "t1", "t2": "t2"} + + def test_stale_active_backend_does_not_leak_into_errored_result(monkeypatch): """渠道单例上一轮的 active_backend 不得泄漏进本轮异常结果(Codex review 发现)。""" from agent_reach import doctor