修复 state.json 持久化竞态、负载膨胀与幂等重放 (PR-1)

src/agent_bridge/registry.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import contextlib
+import hashlib
 import json
 import logging
 import os
@@ -96,6 +97,18 @@
 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
+# Poll cadence while flush_state() waits on the shared flusher; small so a
+# terminal wait_task does not add up to a full second under save traffic.
+STATE_FLUSH_POLL_SEC = 0.05
+# 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
+# Persisted request_id bindings kept; bindings also drop with their task.
+REQUEST_KEEP_MAX = 1000
 
 
 def _new_id(prefix: str) -> str:
@@ -124,6 +137,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 +156,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 +177,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
         self._state_quarantined = False
 
@@ -216,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:
@@ -246,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
@@ -256,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:
@@ -264,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():
@@ -276,28 +331,62 @@ 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
+    async def flush_state(self) -> bool:
+        """Wait until every save() issued so far is on disk.
+
+        Generation-scoped and returns True iff that generation landed on
+        disk: 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.
+        Callers that cannot proceed on stale disk (start, stop) check the
+        return value; wait_task reports it as ``persisted``.
+        """
+        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=STATE_FLUSH_POLL_SEC)
+                if fresh and task.done():
+                    break  # this attempt's run finished; retry or give up
+            if self._flushed_gen >= target:
+                return True
+            if self._pending_state is None:
+                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,
+            )
+        return False
 
     def touch_activity(self) -> None:
         self._last_activity = time.monotonic()
@@ -403,6 +492,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)
         if self._state_quarantined:
             # The tasks map is empty only because state was quarantined, not
             # because the tasks are gone: sweeping now would delete result
@@ -413,10 +510,12 @@ async def start(self) -> None:
         # 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()
-        await self.flush_state()
+        if not await self.flush_state():
+            raise RuntimeError("could not persist state.json") from self._state_error
         self._prune()
         self.save()
-        await self.flush_state()
+        if not await self.flush_state():
+            raise RuntimeError("could not persist state.json") from self._state_error
         self.touch_activity()
         if self.config.server.idle_exit_sec > 0:
             self._watchdog = asyncio.create_task(
@@ -431,9 +530,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 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()
@@ -460,7 +565,8 @@ async def stop(self) -> None:
         except OSError:
             log.exception("could not flush transcripts during shutdown")
         self.save()
-        await self.flush_state()
+        if not await self.flush_state():
+            raise RuntimeError("could not persist state.json") from self._state_error
 
     def _adapter_for(self, session: Session) -> Adapter:
         existing = self._adapters.get(session.session_id)
@@ -611,8 +717,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 {
@@ -674,7 +780,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()
@@ -1045,6 +1151,11 @@ async def _idle() -> None:
                     await adapter.shutdown(current)
                 except Exception:
                     log.exception("idle unload failed for %s", session_id)
+                else:
+                    # Drop the dead adapter mapping: keeping it excludes the
+                    # session from inactive pruning (the filter requires
+                    # `not in _adapters`); revival rebuilds via _adapter_for.
+                    self._adapters.pop(session_id, None)
             current.proc_state = ProcState.idle_unloaded
             self.save()
 
@@ -1198,7 +1309,16 @@ 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)}
+        persisted = False
+        if task.status in TERMINAL_STATUSES:
+            # A terminal answer must also be on disk, not merely queued
+            # behind the background flusher.
+            persisted = await self.flush_state()
+        return {
+            "timed_out": False,
+            "persisted": persisted,
+            **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, allow_disk=True))