Repository navigation
Expand file tree
/
Copy pathagent_state.py
More file actions
288 lines (231 loc) · 12.2 KB
/
Copy pathagent_state.py
File metadata and controls
288 lines (231 loc) · 12.2 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
"""Canonical lifecycle state and authority selection for coding agents.
This module intentionally has no FastAPI, PTY, or database dependency. The
server can use it for reports and heuristic fallbacks while tests exercise the
state contract in isolation.
"""
import re
from dataclasses import dataclass
from enum import StrEnum
from time import monotonic
class State(StrEnum):
UNKNOWN = "unknown"
IDLE = "idle"
WORKING = "working"
BLOCKED = "blocked"
DONE = "done"
EXITED = "exited"
VALID_BLOCKED_REASONS = frozenset({"approval", "question", "authentication", "tool_error", "external_wait", "unknown"})
# Larger numbers outrank lower numbers. Sources in the same tier use the most
# recent report, which makes replacing an integration implementation safe.
SOURCE_PRIORITY = {
"integration": 60,
"custom": 50,
"native": 40,
"manifest": 30,
"heuristic": 20,
}
@dataclass(frozen=True)
class Status:
"""The server-owned lifecycle state shown to clients."""
state: State
source: str = ""
blocked_reason: str | None = None
detail: str | None = None
def __post_init__(self) -> None:
if self.blocked_reason and self.blocked_reason not in VALID_BLOCKED_REASONS:
raise ValueError("invalid blocked reason")
if self.blocked_reason and self.state is not State.BLOCKED:
raise ValueError("blocked reason requires blocked state")
@dataclass(frozen=True)
class AgentSessionReference:
"""Validated native conversation information retained for a future resume."""
agent: str
session_id: str
resume_argv: tuple[str, ...]
@dataclass(frozen=True)
class DisplayMetadata:
"""Presentation-only metadata; it never participates in state resolution."""
title: str | None = None
display_agent: str | None = None
tokens: tuple[tuple[str, str], ...] = ()
state_label: str | None = None
@dataclass(frozen=True)
class Report:
"""A time-limited lifecycle assertion from one source."""
source: str
status: Status
expires_at: float
reported_at: float
@property
def category(self) -> str:
return self.source.split(":", 1)[0]
def expired(self, now: float | None = None) -> bool:
return self.expires_at <= (monotonic() if now is None else now)
def resolve(reports: list[Report], now: float | None = None) -> Status:
"""Return the highest-authority, newest non-expired report or unknown."""
current = monotonic() if now is None else now
active = [report for report in reports if not report.expired(current)]
if not active:
return Status(State.UNKNOWN)
selected = max(active, key=lambda report: (SOURCE_PRIORITY.get(report.category, 0), report.reported_at))
return selected.status
# --- report validation (POST /api/agent-reports) -------------------------------------------------
MAX_TTL_MS = 120_000
DEFAULT_TTL_MS = 30_000
MAX_TEXT = 120
MAX_TOKENS = 8
MAX_ARGV = 16
MAX_ARGV_BYTES = 1024
_KEY = re.compile(r"^[a-z][a-z0-9_.-]{0,31}$")
_EXE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.+-]{0,63}$")
_SOURCE = re.compile(r"^[a-z]+:[a-z0-9_.-]{1,40}$")
_CTRL = re.compile(r"[\x00-\x1f\x7f]")
RESUME_EXE = {"cursor": "cursor-agent", "antigravity": "agy"} # agents whose command isn't their kind
class ReportError(ValueError):
"""A rejected report; `str(e)` is a stable error code for integration authors."""
def _text(value, code: str) -> str | None:
if value is None:
return None
if not isinstance(value, str):
raise ReportError(code)
return _CTRL.sub(" ", value).strip()[:MAX_TEXT] or None
def parse_metadata(raw) -> DisplayMetadata:
if raw is None:
return DisplayMetadata()
if not isinstance(raw, dict):
raise ReportError("invalid_metadata")
tokens = raw.get("tokens") or {}
if not isinstance(tokens, dict) or len(tokens) > MAX_TOKENS:
raise ReportError("invalid_tokens")
pairs = []
for k, v in tokens.items():
if not isinstance(k, str) or not _KEY.match(k) or not isinstance(v, (str, int, float)):
raise ReportError("invalid_tokens")
pairs.append((k, _CTRL.sub(" ", str(v))[:40]))
return DisplayMetadata(_text(raw.get("title"), "invalid_metadata"), _text(raw.get("display_agent"), "invalid_metadata"),
tuple(pairs), _text(raw.get("state_label"), "invalid_metadata"))
def parse_resume_argv(raw, agent: str) -> tuple[str, ...]:
"""Only an executable name (no path) followed by bounded, control-free arguments."""
if not isinstance(raw, list) or not 1 <= len(raw) <= MAX_ARGV or not all(isinstance(a, str) for a in raw):
raise ReportError("invalid_resume_argv")
if not _EXE.match(raw[0]) or any(_CTRL.search(a) for a in raw) or sum(len(a.encode()) for a in raw) > MAX_ARGV_BYTES:
raise ReportError("invalid_resume_argv")
if raw[0].lower().removesuffix(".exe").removesuffix(".cmd") not in (agent, RESUME_EXE.get(agent)):
raise ReportError("resume_agent_mismatch")
return tuple(raw)
def parse_activity(raw) -> dict | None:
"""A finished tool call from a report: {command, ok} or {file, change}; bounded, control characters removed."""
if raw is None:
return None
if not isinstance(raw, dict):
raise ReportError("invalid_activity")
if isinstance(raw.get("command"), str) and raw["command"].strip():
return {"command": _CTRL.sub(" ", raw["command"]).strip()[:500], "ok": bool(raw.get("ok", True))}
if isinstance(raw.get("file"), str) and raw["file"].strip():
return {"file": _CTRL.sub("", raw["file"]).strip()[:300], "change": "write" if raw.get("change") == "write" else "edit"}
raise ReportError("invalid_activity")
def parse_report(payload: dict, agent: str, now: float | None = None) -> tuple[Report | None, DisplayMetadata, AgentSessionReference | None]:
"""Validate one report for a terminal whose running agent is `agent` (from process detection).
With no `state` it is a session-only report (native session id, no lifecycle claim): Report is None."""
if not isinstance(payload, dict):
raise ReportError("invalid_payload")
source = payload.get("source")
if not isinstance(source, str) or not _SOURCE.match(source):
raise ReportError("invalid_source")
category = source.split(":", 1)[0]
if category not in SOURCE_PRIORITY or category == "heuristic":
raise ReportError("invalid_source")
if payload.get("agent") != agent:
raise ReportError("agent_mismatch")
session_only = payload.get("state") is None
if session_only and payload.get("agent_session_id") is None:
raise ReportError("invalid_state")
try:
state = State(payload.get("state") or "unknown")
except ValueError:
raise ReportError("invalid_state") from None
reason = payload.get("blocked_reason")
if reason is not None and (state is not State.BLOCKED or reason not in VALID_BLOCKED_REASONS):
raise ReportError("invalid_blocked_reason")
ttl = payload.get("ttl_ms", DEFAULT_TTL_MS)
if isinstance(ttl, bool) or not isinstance(ttl, int) or ttl <= 0:
raise ReportError("invalid_ttl")
ttl = min(ttl, MAX_TTL_MS)
status = Status(state, source=source, blocked_reason=reason if state is State.BLOCKED else None,
detail=_text(payload.get("detail"), "invalid_detail"))
t = monotonic() if now is None else now
ref = None
native = payload.get("agent_session_id")
if native is not None or payload.get("resume_argv") is not None:
if not isinstance(native, str) or not 1 <= len(native) <= 200 or _CTRL.search(native):
raise ReportError("invalid_agent_session_id")
if category not in ("integration", "native"):
raise ReportError("resume_not_allowed") # only built-in integrations/log readers may issue resume commands
ref = AgentSessionReference(agent, native, parse_resume_argv(payload.get("resume_argv"), agent))
report = None if session_only else Report(source, status, t + ttl / 1000, t)
return report, parse_metadata(payload.get("metadata")), ref
class Registry:
"""In-memory reports per terminal; the server is the only owner of lifecycle state."""
def __init__(self) -> None:
self.reports: dict[str, dict[str, Report]] = {}
self.meta: dict[str, dict[str, tuple[DisplayMetadata, float]]] = {}
def put(self, sid: str, report: Report, meta: DisplayMetadata | None = None) -> None:
self.reports.setdefault(sid, {})[report.source] = report
if meta is not None and meta != DisplayMetadata():
self.meta.setdefault(sid, {})[report.source] = (meta, report.expires_at)
def put_heuristic(self, sid: str, report: Report) -> None:
"""The screen check's latest verdict replaces its previous one (heuristic:screen vs heuristic:activity), or a stale
question could outlive the answer: with Windows' ~16ms clock both can carry the same timestamp."""
live = self.reports.setdefault(sid, {})
for key in [k for k in live if k.startswith("heuristic:")]:
del live[key]
live[report.source] = report
def status(self, sid: str, now: float | None = None) -> Status:
t = monotonic() if now is None else now
live = {k: r for k, r in self.reports.get(sid, {}).items() if not r.expired(t)}
self.reports[sid] = live
return resolve(list(live.values()), t)
def metadata(self, sid: str, now: float | None = None) -> DisplayMetadata:
"""Metadata from the newest unexpired source; presentation only."""
t = monotonic() if now is None else now
live = {k: v for k, v in self.meta.get(sid, {}).items() if v[1] > t}
self.meta[sid] = live
return max(live.values(), key=lambda v: v[1])[0] if live else DisplayMetadata()
def forget(self, sid: str) -> None:
self.reports.pop(sid, None)
self.meta.pop(sid, None)
def heuristic(legacy: str, now: float, ttl: float = 10) -> Report:
"""Compatibility adapter: the screen/activity watcher's working|approval|idle as a report."""
if legacy == "approval":
status = Status(State.BLOCKED, source="heuristic:screen", blocked_reason="approval")
elif legacy == "working":
status = Status(State.WORKING, source="heuristic:activity")
else:
status = Status(State.IDLE, source="heuristic:activity")
return Report(status.source, status, now + ttl, now)
def to_dict(status: Status, meta: DisplayMetadata | None = None) -> dict:
out = {"state": str(status.state), "source": status.source.split(":", 1)[0] if status.source else "unknown",
"reason": status.blocked_reason, "detail": status.detail}
if meta and meta != DisplayMetadata():
out["meta"] = {"title": meta.title, "display_agent": meta.display_agent, "state_label": meta.state_label, "tokens": dict(meta.tokens)}
return out
NATIVE_WORKING_S = 15 # a log that stopped growing may be waiting on a prompt: let the screen decide again
NATIVE_DONE_S = 120
def native_report(agent: str, found: dict, now: float | None = None) -> Report | None:
"""A report from an agent's own log ({state, age}); None when nothing current to say."""
state, age = found.get("state"), float(found.get("age") or 0)
hold = {"working": NATIVE_WORKING_S, "done": NATIVE_DONE_S}.get(state or "")
if not hold or age >= hold:
return None
t = monotonic() if now is None else now
source = f"native:{agent}"
return Report(source, Status(State(state), source=source), t + hold - age, t - age)
def explain(reports: dict[str, Report], now: float | None = None) -> list[dict]:
"""Every report for a terminal, highest authority first, with age, time left and whether it decides the state."""
t = monotonic() if now is None else now
live = [r for r in reports.values() if not r.expired(t)]
winner = max(live, key=lambda r: (SOURCE_PRIORITY.get(r.category, 0), r.reported_at)) if live else None
rows = sorted(reports.values(), key=lambda r: (-SOURCE_PRIORITY.get(r.category, 0), -r.reported_at))
return [{"source": r.source, "priority": SOURCE_PRIORITY.get(r.category, 0), "state": str(r.status.state),
"reason": r.status.blocked_reason, "age_s": round(t - r.reported_at, 1), "expires_in_s": round(r.expires_at - t, 1),
"expired": r.expired(t), "decides": r is winner} for r in rows]