A2A のタスクが working のまま静かに滞留していたとき — 状態遷移を計測して詰まりを名指しする運用メモ

# 環境変数: A2A_STALL_FACTOR="3"
import asyncio
import time
import uuid
from dataclasses import dataclass, field
from collections import defaultdict
from antigravity import A2AClient, TaskRequest, Message
 
TERMINAL = {"completed", "failed", "cancelled"}
 
@dataclass
class Transition:
    task_id: str
    agent: str
    from_state: str
    to_state: str
    at: float
    dwell: float
 
@dataclass
class TraceStore:
    events: list[Transition] = field(default_factory=list)
    _last: dict = field(default_factory=dict)  # task_id -> (state, at)
 
    def record(self, task_id: str, agent: str, to_state: str) -> None:
        now = time.monotonic()
        prev_state, prev_at = self._last.get(task_id, ("<start>", now))
        self.events.append(Transition(
            task_id=task_id, agent=agent,
            from_state=prev_state, to_state=to_state,
            at=now, dwell=now - prev_at))
        self._last[task_id] = (to_state, now)
 
    def open_tasks(self) -> list[tuple[str, str, str, float]]:
        """非終端状態でとどまっているタスクを (task_id, agent, state, 経過秒) で返す"""
        now = time.monotonic()
        out = []
        for task_id, (state, at) in self._last.items():
            if state not in TERMINAL:
                agent = next((e.agent for e in reversed(self.events)
                              if e.task_id == task_id), "?")
                out.append((task_id, agent, state, now - at))
        return out
 
    def edge_p95(self) -> dict[tuple[str, str], float]:
        """辺ごとの dwell time の p95。滞留閾値の逆算に使う"""
        buckets = defaultdict(list)
        for e in self.events:
            buckets[(e.from_state, e.to_state)].append(e.dwell)
        p95 = {}
        for edge, vals in buckets.items():
            vals.sort()
            p95[edge] = vals[max(0, int(len(vals) * 0.95) - 1)]
        return p95

store = TraceStore()
 
class MonitoredClient(A2AClient):
    async def send_task(self, target, task: TaskRequest):
        store.record(task.id, self.my_agent_id, "submitted")
        resp = await super().send_task(target, task)
        store.record(task.id, target, "working")
        return resp
 
    async def on_state(self, task_id: str, agent: str, state: str):
        store.record(task_id, agent, state)
 
async def watchdog(stall_factor: float = 3.0, interval: float = 10.0):
    """滞留を辺の p95 × factor で判定し、名指しで警告する"""
    while True:
        await asyncio.sleep(interval)
        p95 = store.edge_p95()
        for task_id, agent, state, elapsed in store.open_tasks():
            # その状態から出ていく辺の p95 を基準にする
            budgets = [v for (f, _), v in p95.items() if f == state]
            budget = max(budgets) * stall_factor if budgets else 60.0
            if elapsed > budget:
                print(f"⚠️ STALL task={task_id[:8]} agent={agent} "
                      f"state={state} elapsed={elapsed:.0f}s budget={budget:.0f}s")