From d58837c3db1f677f6a930c5d51bec15f7ffcf029 Mon Sep 17 00:00:00 2001 From: lulu feizhu <279467036+FeiZhuLulu@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:23:53 +0000 Subject: [PATCH 1/3] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=20state.json=20=E6=8C=81?= =?UTF-8?q?=E4=B9=85=E5=8C=96=E7=AB=9E=E6=80=81=E3=80=81=E8=B4=9F=E8=BD=BD?= =?UTF-8?q?=E8=86=A8=E8=83=80=E4=B8=8E=E5=B9=82=E7=AD=89=E9=87=8D=E6=94=BE?= =?UTF-8?q?=20(PR-1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - flush_state 改代际语义并降级不 raise，消除并发 save 触发的虚假 RuntimeError（S1） - 落盘失败最多重试 3 次后降级为 warning + _state_error（H-08） - state task 行瘦身 message<=2000/result_text<=6000，全文留 result/transcript（S2） - RESIDENT_SESSION_KEEP=50 限制常驻 worker 会话（S2，soak 实测 ~970 worker/20GB） - request_id 幂等键持久化到 state.json 的 requests 键，sha256 指纹，上限 1000（H-01） - wait_task 终态返回前 flush，新增 persisted 字段（B7） --- src/agent_bridge/registry.py | 164 +++++++++++++++++++++++++++++++---- tests/test_registry.py | 68 ++++++++++++++- tests/test_request_id.py | 21 +++++ 3 files changed, 234 insertions(+), 19 deletions(-) diff --git a/src/agent_bridge/registry.py b/src/agent_bridge/registry.py index ca10342..9ddb5fb 100644 --- a/src/agent_bridge/registry.py +++ b/src/agent_bridge/registry.py @@ -2,6 +2,7 @@ import asyncio import contextlib +import hashlib import json import logging import os @@ -96,6 +97,19 @@ STATE_WRITE_ATTEMPTS = 5 STATE_WRITE_RETRY_BASE_SEC = 0.05 STATE_LOCK_TIMEOUT_SEC = 5.0 +# flush_state() waits for the saves issued at call time; write failures +# degrade to a logged warning instead of raising into the caller. +STATE_FLUSH_ATTEMPTS = 3 +# state.json is an index, not a data store: task message/result bodies are +# persisted truncated so a long-lived Bridge keeps each write bounded. +STATE_MESSAGE_MAX = 2000 +STATE_RESULT_TEXT_MAX = 6000 +# Live resident (ready) sessions kept before force-unloading the least +# recently active ones. With idle_unload_sec=0 workers otherwise accumulate +# a process each until the host runs out of memory. +RESIDENT_SESSION_KEEP = 50 +# Persisted request_id bindings kept; bindings also drop with their task. +REQUEST_KEEP_MAX = 1000 def _new_id(prefix: str) -> str: @@ -124,6 +138,11 @@ def _tail(text: str, limit: int = RESULT_TAIL) -> str: return encoded[-limit:].decode("utf-8", errors="ignore") +def _request_fingerprint(request: tuple) -> str: + blob = json.dumps(list(request), ensure_ascii=False, default=str, separators=(",", ":")) + return hashlib.sha256(blob.encode("utf-8")).hexdigest() + + class Registry: def __init__( self, @@ -138,7 +157,9 @@ def __init__( self.config = config self.sessions: dict[str, Session] = {} self.tasks: dict[str, Task] = {} - self._requests: dict[str, tuple[tuple, str]] = {} + # request_id -> (request fingerprint, task_id). Fingerprints instead + # of request tuples so bindings can persist to state.json compactly. + self._requests: dict[str, tuple[str, str]] = {} self._adapters: dict[str, Adapter] = {} self._done: dict[str, asyncio.Event] = {} self._idle: dict[str, asyncio.Task[None]] = {} @@ -157,7 +178,10 @@ def __init__( self._quota_cache = QuotaCache(config.quota.cache_sec) self._stopping = False self._pending_state: dict[str, list[dict]] | None = None + self._pending_gen = 0 self._flush_task: asyncio.Task[None] | None = None + self._save_gen = 0 + self._flushed_gen = 0 self._state_error: Exception | None = None @classmethod @@ -215,10 +239,33 @@ def _merge_owned( merged.update(mine) return list(merged.values()) + @staticmethod + def _state_task_row(task: Task) -> dict: + row = task.model_dump(mode="json") + # Full bodies already live in result files and transcripts; writing + # them into state.json again grew the payload without bound. + if len(row["message"]) > STATE_MESSAGE_MAX: + row["message"] = row["message"][:STATE_MESSAGE_MAX] + row["result_text"] = _tail(row["result_text"], STATE_RESULT_TEXT_MAX) + return row + def _own_rows(self) -> dict[str, list[dict]]: + request_rows = [ + { + "request_id": rid, + "task_id": tid, + "fingerprint": fp, + # Owner stamps let _merge_owned drop rows this instance pruned; + # without them pruned bindings linger on disk forever. + "owner_pid": self._owner_pid, + "owner_create_time": self._owner_create_time, + } + for rid, (fp, tid) in self._requests.items() + ][-REQUEST_KEEP_MAX:] return { "sessions": [s.model_dump(mode="json") for s in self.sessions.values()], - "tasks": [t.model_dump(mode="json") for t in self.tasks.values()], + "tasks": [self._state_task_row(t) for t in self.tasks.values()], + "requests": request_rows, } def _write_state(self, own: dict[str, list[dict]]) -> None: @@ -245,6 +292,11 @@ def _write_state(self, own: dict[str, list[dict]]) -> None: {row["task_id"]: row for row in own["tasks"]}, "task_id", ), + "requests": self._merge_owned( + disk.get("requests") or [], + {row["request_id"]: row for row in own["requests"]}, + "request_id", + ), }, ) return @@ -255,6 +307,8 @@ def _write_state(self, own: dict[str, list[dict]]) -> None: def save(self) -> None: self._pending_state = self._own_rows() + self._save_gen += 1 + self._pending_gen = self._save_gen try: loop = asyncio.get_running_loop() except RuntimeError: @@ -263,8 +317,10 @@ def save(self) -> None: self._write_state(own) except Exception as exc: self._pending_state = own + self._pending_gen = self._save_gen self._state_error = exc raise + self._flushed_gen = self._save_gen self._state_error = None return if self._flush_task is None or self._flush_task.done(): @@ -275,28 +331,53 @@ async def _flush_state_loop(self) -> None: # save() cannot observe a still-running flush task and skip creating # a new one after this loop has already decided to exit. while self._pending_state is not None: - own, self._pending_state = self._pending_state, None + own, gen = self._pending_state, self._pending_gen + self._pending_state = None try: await asyncio.to_thread(self._write_state, own) except Exception as exc: # Keep the newest snapshot. A later save or flush can retry it. if self._pending_state is None: self._pending_state = own + self._pending_gen = gen self._state_error = exc log.exception("could not write state.json") return + self._flushed_gen = gen self._state_error = None async def flush_state(self) -> None: - """Wait until every save() so far is on disk.""" - task = self._flush_task - if self._pending_state is not None and (task is None or task.done()): - task = asyncio.create_task(self._flush_state_loop(), name="state-flush") - self._flush_task = task - if task is not None and not task.done(): - await task - if self._pending_state is not None: - raise RuntimeError("could not persist state.json") from self._state_error + """Wait until every save() issued so far is on disk. + + Generation-scoped: a save() issued *after* this call does not block + it, so concurrent saves can no longer trigger the spurious + "could not persist state.json" raise. Persistent write failures + degrade to a logged warning — a broken disk must not kill + coordination; the next save() retries the same rows because + snapshots are cumulative. + """ + target = self._save_gen + for _ in range(STATE_FLUSH_ATTEMPTS): + if self._flushed_gen >= target: + return + task = self._flush_task + if task is None or task.done(): + if self._pending_state is None: + break + task = asyncio.create_task(self._flush_state_loop(), name="state-flush") + self._flush_task = task + try: + await task + except Exception as exc: # flush loop itself died unexpectedly + self._state_error = exc + break + if self._flushed_gen < target: + log.warning( + "state.json persistence degraded (flushed %s of %s saves): %s", + self._flushed_gen, + target, + self._state_error, + ) def touch_activity(self) -> None: self._last_activity = time.monotonic() @@ -394,6 +475,14 @@ async def start(self) -> None: done = asyncio.Event() done.set() self._done[task.task_id] = done + for raw in payload.get("requests") or []: + if not isinstance(raw, dict): + continue + rid = raw.get("request_id") + tid = raw.get("task_id") + fp = raw.get("fingerprint") + if rid and fp and tid in self.tasks: + self._requests[rid] = (fp, tid) # Stamp adopted owners onto disk first so a following prune is not # undone by _merge_owned treating the old unowned rows as foreign. self.save() @@ -595,8 +684,8 @@ async def dispatch_task( request = (agent, message, str(cwd_path.resolve()), session_id, model, effort, title, user_requested) async with self._lock: if request_id is not None and request_id in self._requests: - previous, task_id = self._requests[request_id] - if previous != request: + previous_fp, task_id = self._requests[request_id] + if previous_fp != _request_fingerprint(request): raise ValueError("request_id is already bound to a different dispatch request") task = self.tasks[task_id] return { @@ -658,7 +747,7 @@ async def dispatch_task( self._stamp_owner(task) self.tasks[task.task_id] = task if request_id is not None: - self._requests[request_id] = (request, task.task_id) + self._requests[request_id] = (_request_fingerprint(request), task.task_id) self._done[task.task_id] = asyncio.Event() self._cancel_idle(session.session_id) self._prune() @@ -930,6 +1019,41 @@ def _prune_sessions(self) -> None: session.proc_state.value, session.last_active_at, ) + # Bound live resident sessions: the default idle_unload_sec=0 never + # unloads workers on its own, so `ready` sessions accumulate a worker + # process each until the host runs out of memory. Force-unload the + # least recently active ones beyond RESIDENT_SESSION_KEEP; they land + # as idle_unloaded and become regular prune candidates later. + resident = [ + session + for session in self.sessions.values() + if session.proc_state == ProcState.ready + and session.session_id in self._adapters + and self._busy_task(session.session_id) is None + ] + resident.sort(key=lambda item: _session_last_active_ts(item.last_active_at)) + for session in resident[: max(0, len(resident) - RESIDENT_SESSION_KEEP)]: + log.info("over-cap unload of resident session %s", session.session_id) + self._unload_over_cap(session.session_id) + + def _unload_over_cap(self, session_id: str) -> None: + self._cancel_idle(session_id) + + async def _unload() -> None: + session = self.sessions.get(session_id) + if session is None or self._busy_task(session_id): + return + adapter = self._adapters.get(session_id) + if adapter is not None: + try: + await adapter.shutdown(session) + except Exception: + log.exception("over-cap unload failed for %s", session_id) + return + session.proc_state = ProcState.idle_unloaded + self.save() + + self._idle[session_id] = asyncio.create_task(_unload(), name=f"overcap-{session_id}") def _prune_tasks(self) -> None: by_session: dict[str, list[Task]] = {} @@ -1114,7 +1238,15 @@ async def wait_task(self, task_id: str, timeout_sec: float = DEFAULT_WAIT_SEC) - await asyncio.wait_for(event.wait(), timeout=timeout_sec) except TimeoutError: return {"timed_out": True, **self._task_snapshot(self.tasks[task_id])} - return {"timed_out": False, **self._task_snapshot(self.tasks[task_id], include_result=True)} + if task.status in TERMINAL_STATUSES: + # A terminal answer must also be on disk, not merely queued + # behind the background flusher. + await self.flush_state() + return { + "timed_out": False, + "persisted": self._flushed_gen >= self._save_gen, + **self._task_snapshot(self.tasks[task_id], include_result=True), + } def check_task(self, task_id: str) -> dict: return self._task_snapshot(self._require_task(task_id)) diff --git a/tests/test_registry.py b/tests/test_registry.py index ec011eb..3cb5807 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -404,14 +404,76 @@ def fail_once(own): monkeypatch.setattr(registry, "_write_state", fail_once) registry.save() - with pytest.raises(RuntimeError, match=r"could not persist state\.json"): - await registry.flush_state() + # flush_state degrades instead of raising: the failed write is retried + # by the same call (STATE_FLUSH_ATTEMPTS) and the snapshot lands. + await registry.flush_state() + assert registry._pending_state is None + payload = read_json(state_path(bridge_home), {}) + assert {row["session_id"] for row in payload["sessions"]} == {"sess_retry"} + + +@pytest.mark.asyncio +async def test_persistent_state_write_failure_degrades_flush_without_raising( + bridge_home, tmp_path, monkeypatch +): + registry = Registry.create(bridge_home) + registry.sessions["sess_dead_disk"] = Session( + session_id="sess_dead_disk", + agent="fake", + cwd=str(tmp_path), + ) + original_write = registry._write_state + + def fail_always(own): + raise PermissionError(5, "synthetic sharing/access failure") + + monkeypatch.setattr(registry, "_write_state", fail_always) + registry.save() + + # Never raises: a broken disk must not kill the caller. The newest + # snapshot is kept for a later retry and surfaced via _state_error. + await registry.flush_state() + assert registry._state_error is not None assert registry._pending_state is not None + monkeypatch.setattr(registry, "_write_state", original_write) await registry.flush_state() assert registry._pending_state is None + assert registry._state_error is None payload = read_json(state_path(bridge_home), {}) - assert {row["session_id"] for row in payload["sessions"]} == {"sess_retry"} + assert {row["session_id"] for row in payload["sessions"]} == {"sess_dead_disk"} + + +@pytest.mark.asyncio +async def test_flush_state_ignores_saves_issued_during_the_wait( + bridge_home, tmp_path, monkeypatch +): + """A save() landing mid-flush must not trip a spurious persist error.""" + import threading + + registry = Registry.create(bridge_home) + registry.sessions["sess_race"] = Session( + session_id="sess_race", + agent="fake", + cwd=str(tmp_path), + ) + original_write = registry._write_state + started = threading.Event() + gate = threading.Event() + + def slow_write(own): + started.set() + gate.wait(10) + original_write(own) + + monkeypatch.setattr(registry, "_write_state", slow_write) + registry.save() + flush = asyncio.create_task(registry.flush_state()) + await asyncio.to_thread(started.wait, 10) # first write in-flight, blocked + registry.save() # lands while flush_state is awaiting the first write + gate.set() + await asyncio.wait_for(flush, timeout=10) # must not raise + assert registry._flushed_gen >= 1 def test_state_write_retries_bounded_permission_errors(bridge_home, monkeypatch): diff --git a/tests/test_request_id.py b/tests/test_request_id.py index d10dc71..14a13f0 100644 --- a/tests/test_request_id.py +++ b/tests/test_request_id.py @@ -88,6 +88,27 @@ async def test_pruning_drops_binding(bridge_home, tmp_path, monkeypatch): await registry.stop() +async def test_request_id_survives_restart(bridge_home, tmp_path): + """Bindings persist to state.json so a replay after restart reuses the task.""" + registry = Registry.create(bridge_home) + await registry.start() + request_id = str(uuid.uuid4()) + args = dict(agent="fake", message="hello", cwd=str(tmp_path), request_id=request_id) + first = await registry.dispatch_task(**args) + await registry.wait_task(first["task_id"], timeout_sec=5) + await registry.stop() + + registry2 = Registry.create(bridge_home) + await registry2.start() + try: + retry = await registry2.dispatch_task(**args) + assert retry == {**first, "reused": True} + with pytest.raises(ValueError, match="different dispatch request"): + await registry2.dispatch_task(**{**args, "message": "different"}) + finally: + await registry2.stop() + + async def test_uuid_validation_and_canonicalization(bridge_home, tmp_path): registry = Registry.create(bridge_home) await registry.start() From fc034ba011641814efa92130c1d34ea195232bd6 Mon Sep 17 00:00:00 2001 From: lulu feizhu <279467036+FeiZhuLulu@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:37:17 +0000 Subject: [PATCH 2/3] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=A4=8D=E5=AE=A1?= =?UTF-8?q?=E5=8F=91=E7=8E=B0=E7=9A=84=E4=B8=89=E5=A4=84=E7=BC=BA=E9=99=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - flush_state 不再 await 共享 flusher 全程：持续 save 流下按代际轮询，目标代际落盘即返回（wait_task 不再被拖死） - 抽出 _cap_resident_sessions，_run_task 完成即复查常驻上限——burst 完成不再全体滞留（此前只在 dispatch/start 时修剪） - unload 成功后弹出 _adapters 映射（over-cap 与 idle 路径一致），否则 unloaded 会话永不进入淘汰候选、状态索引持续增长 --- src/agent_bridge/registry.py | 41 +++++++++++++++++++++++++++--------- 1 file changed, 31 insertions(+), 10 deletions(-) diff --git a/src/agent_bridge/registry.py b/src/agent_bridge/registry.py index 9ddb5fb..bc6405b 100644 --- a/src/agent_bridge/registry.py +++ b/src/agent_bridge/registry.py @@ -358,18 +358,24 @@ async def flush_state(self) -> None: """ target = self._save_gen for _ in range(STATE_FLUSH_ATTEMPTS): + while self._flushed_gen < target: + task = self._flush_task + fresh = False + if task is None or task.done(): + if self._pending_state is None: + break + task = asyncio.create_task(self._flush_state_loop(), name="state-flush") + self._flush_task = task + fresh = True + # Poll the shared flusher instead of awaiting it to the end: + # under steady save() traffic the loop keeps draining newer + # snapshots, but our target generation is already on disk. + await asyncio.wait({task}, timeout=1.0) + if fresh and task.done(): + break # this attempt's run finished; retry or give up if self._flushed_gen >= target: return - task = self._flush_task - if task is None or task.done(): - if self._pending_state is None: - break - task = asyncio.create_task(self._flush_state_loop(), name="state-flush") - self._flush_task = task - try: - await task - except Exception as exc: # flush loop itself died unexpectedly - self._state_error = exc + if self._pending_state is None: break if self._flushed_gen < target: log.warning( @@ -922,6 +928,9 @@ async def _run_task(self, task_id: str) -> None: (task.error or "")[:500], ) self._schedule_idle(session.session_id) + # Re-check the resident cap: a burst of completions all land here + # and would otherwise stay resident until the next dispatch. + self._cap_resident_sessions() @staticmethod def _set_files_changed(task: Task, paths: list[str]) -> None: @@ -1019,6 +1028,9 @@ def _prune_sessions(self) -> None: session.proc_state.value, session.last_active_at, ) + self._cap_resident_sessions() + + def _cap_resident_sessions(self) -> None: # Bound live resident sessions: the default idle_unload_sec=0 never # unloads workers on its own, so `ready` sessions accumulate a worker # process each until the host runs out of memory. Force-unload the @@ -1050,6 +1062,11 @@ async def _unload() -> None: except Exception: log.exception("over-cap unload failed for %s", session_id) return + # Drop the dead adapter mapping: keeping it excludes the session + # from inactive pruning (filter requires `not in _adapters`), so + # unloaded sessions would persist in every state snapshot. Revival + # rebuilds via _adapter_for using session.native_session_id. + self._adapters.pop(session_id, None) session.proc_state = ProcState.idle_unloaded self.save() @@ -1110,6 +1127,10 @@ async def _idle() -> None: await adapter.shutdown(current) except Exception: log.exception("idle unload failed for %s", session_id) + else: + # Same as the over-cap path: drop the dead mapping so the + # session stays prunable and revival rebuilds the adapter. + self._adapters.pop(session_id, None) current.proc_state = ProcState.idle_unloaded self.save() From 158210994864d5e9ac8992030523d213ecc969b0 Mon Sep 17 00:00:00 2001 From: lulu feizhu <279467036+FeiZhuLulu@users.noreply.github.com> Date: Tue, 6 Oct 2026 07:41:35 +0000 Subject: [PATCH 3/3] =?UTF-8?q?shutdown=20=E7=AB=9E=E6=80=81=EF=BC=9A=5Fst?= =?UTF-8?q?opping=20=E6=97=B6=E4=B8=8D=E5=86=8D=E8=B0=83=E5=BA=A6=20over-c?= =?UTF-8?q?ap=20unload=EF=BC=9Bstop()=20await=20=E5=B7=B2=E5=8F=96?= =?UTF-8?q?=E6=B6=88=E7=9A=84=20idle/unload=20=E4=BB=BB=E5=8A=A1=E5=90=8E?= =?UTF-8?q?=E5=86=8D=E5=85=B3=20adapter?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/agent_bridge/registry.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/agent_bridge/registry.py b/src/agent_bridge/registry.py index bc6405b..335c642 100644 --- a/src/agent_bridge/registry.py +++ b/src/agent_bridge/registry.py @@ -510,9 +510,15 @@ async def stop(self) -> None: self._watchdog = None if watchdog is not None: watchdog.cancel() - for idle in list(self._idle.values()): + idles = list(self._idle.values()) + for idle in idles: idle.cancel() self._idle.clear() + if idles: + # Await cancellation so an in-flight over-cap/idle unload finishes + # unwinding before adapter shutdown — otherwise the loop may close + # with a worker half-reaped. + await asyncio.wait(idles, timeout=STOP_TASK_GRACE_SEC) bgs = [task for task in self._bg.values() if not task.done()] for bg in bgs: bg.cancel() @@ -1036,6 +1042,10 @@ def _cap_resident_sessions(self) -> None: # process each until the host runs out of memory. Force-unload the # least recently active ones beyond RESIDENT_SESSION_KEEP; they land # as idle_unloaded and become regular prune candidates later. + if self._stopping: + # stop() owns all worker teardown; an over-cap unload scheduled now + # would race its adapter shutdown and could orphan a worker. + return resident = [ session for session in self.sessions.values()