Files
wyndham-ARR/arr_processing/agent_trace.py
2026-07-31 15:11:42 +08:00

470 lines
16 KiB
Python

"""Privacy-minimized persistence and projection for SuperAgent trace SSE events."""
from __future__ import annotations
import json
import os
import re
import threading
from datetime import datetime, timezone
from hashlib import sha256
from pathlib import Path
from typing import Any, Dict, List, Mapping, Optional, Protocol, Union
from agent_integration.client import OpenAgentEvent
from arr_ingestion.contracts import OPAQUE_ID_RE
TRACE_RECORD_VERSION = "arr-agent-trace-event-1"
MAX_TRACE_EVENTS = 10_000
MAX_TRACE_FILE_BYTES = 8 * 1024 * 1024
MAX_TRACE_SUMMARY_CHARS = 1_200
_LEVELS = frozenset({"info", "success", "warning", "error"})
_CODE_RE = re.compile(r"^[A-Z][A-Z0-9_]{0,95}$")
_SAFE_DETAIL_KEYS = frozenset(
{
"attempt_no",
"remote_run_id",
"remote_status",
"trace_event",
"invalid_records",
"capture_error",
}
)
_SENSITIVE_ASSIGNMENT_RE = re.compile(
r"(?i)\b(submission[_ -]?grant|access[_ -]?key(?:[_ -]?(?:id|secret))?|"
r"api[_ -]?key|secret|token|password|authorization|cookie)\b"
r"\s*[:=]\s*(?:[\"']?)[^\s,;}\]]+"
)
_BEARER_RE = re.compile(r"(?i)\bbearer\s+[A-Za-z0-9._~+/=-]+")
_URL_RE = re.compile(r"(?i)\b(?:https?|oss)://[^\s<>{}\[\]\"']+")
_PATH_RE = re.compile(
r"(?<![A-Za-z0-9._-])/(?:Users|home|private|tmp|var|opt|etc)(?:/[^\s,;}\]]+)+"
)
_EMAIL_RE = re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b")
_LONG_NUMBER_RE = re.compile(r"(?<!\d)(?:\+?\d[\d ()-]{7,}\d)(?!\d)")
_XML_FRAGMENT_RE = re.compile(r"<[^>\r\n]{1,240}>")
class AgentTraceStoreError(RuntimeError):
"""Raised when the local trace store cannot safely persist a record."""
class AgentTraceSink(Protocol):
def append(self, job_id: str, record: Mapping[str, Any]) -> None:
...
class AgentTraceReader(Protocol):
def read(self, job_id: str) -> List[Dict[str, Any]]:
...
def _utc_now() -> datetime:
return datetime.now(timezone.utc)
def _timestamp(value: Any, *, fallback: Optional[datetime] = None) -> str:
if isinstance(value, datetime):
parsed = value
elif isinstance(value, str):
try:
parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
except ValueError:
parsed = fallback or _utc_now()
else:
parsed = fallback or _utc_now()
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=timezone.utc)
return parsed.astimezone(timezone.utc).isoformat()
def sanitize_trace_summary(value: Any) -> Optional[str]:
"""Keep the server's operational summary while removing common secret/PII forms."""
if not isinstance(value, str):
return None
text = " ".join(value.replace("\x00", " ").split())
if not text:
return None
text = _BEARER_RE.sub("Bearer [REDACTED]", text)
text = _SENSITIVE_ASSIGNMENT_RE.sub(lambda match: f"{match.group(1)}=[REDACTED]", text)
text = _URL_RE.sub("[URL]", text)
text = _PATH_RE.sub("[PATH]", text)
text = _EMAIL_RE.sub("[EMAIL]", text)
text = _LONG_NUMBER_RE.sub("[NUMBER]", text)
text = _XML_FRAGMENT_RE.sub("[XML]", text)
return text[:MAX_TRACE_SUMMARY_CHARS]
def trace_run_id(event: OpenAgentEvent) -> Optional[str]:
"""Return the run id only from the documented trace run-start event."""
payload = event.data
if event.event != "trace" or not isinstance(payload, Mapping):
return None
if payload.get("event") != "run.started":
return None
run_id = payload.get("run_id")
if not isinstance(run_id, str) or OPAQUE_ID_RE.fullmatch(run_id) is None:
return None
return run_id
def project_trace_event(
event: OpenAgentEvent,
*,
attempt_no: int,
sequence: int,
now: Optional[datetime] = None,
) -> Optional[Dict[str, Any]]:
"""Project one raw SSE event into the strict local diagnostic contract."""
fallback = now or _utc_now()
if event.event == "end":
return _record(
sequence=sequence,
timestamp=_timestamp(None, fallback=fallback),
level="info",
code="AGENT_TRACE_STREAM_ENDED",
title="SuperAgent trace 流已结束",
message="SuperAgent 已关闭本次运行的 trace 连接。",
details={"attempt_no": attempt_no, "trace_event": "end"},
)
if event.event != "trace" or not isinstance(event.data, Mapping):
return None
payload = event.data
name = payload.get("event")
if not isinstance(name, str) or not name:
return None
if name in {"message.delta", "message.final"}:
# Assistant text can contain result payloads, grants, object routes or PII.
return None
nested = payload.get("data")
summary = sanitize_trace_summary(
nested.get("summary") if isinstance(nested, Mapping) else None
)
run_id = payload.get("run_id")
if not isinstance(run_id, str) or OPAQUE_ID_RE.fullmatch(run_id) is None:
run_id = None
status = payload.get("status")
if not isinstance(status, str) or not status or len(status) > 32:
status = None
details: Dict[str, Any] = {
"attempt_no": attempt_no,
"trace_event": name,
"remote_run_id": run_id,
"remote_status": status,
}
timestamp = _timestamp(payload.get("ts"), fallback=fallback)
if name == "run.started":
return _record(
sequence=sequence,
timestamp=timestamp,
level="info",
code="AGENT_RUN_STARTED",
title="SuperAgent 运行已启动",
message="已建立包含任务 trace 的 SuperAgent 运行流。",
details=details,
)
if name == "task.updated":
if summary is None:
return None
return _record(
sequence=sequence,
timestamp=timestamp,
level="info",
code="AGENT_TASK_UPDATED",
title="Agent 任务状态更新",
message=summary,
details=details,
)
if name == "step.updated":
if summary is None:
return None
return _record(
sequence=sequence,
timestamp=timestamp,
level="info",
code="AGENT_STEP_UPDATED",
title="Agent 执行步骤更新",
message=summary,
details=details,
)
if name == "run.completed":
succeeded = status in {"success", "completed", "succeeded"}
return _record(
sequence=sequence,
timestamp=timestamp,
level="success" if succeeded else "warning",
code="AGENT_RUN_COMPLETED",
title="SuperAgent 运行已结束",
message=summary or "SuperAgent 已报告运行终态;业务成功仍以 ARR 数据库提交为准。",
details=details,
)
if name in {"run.failed", "error"}:
return _record(
sequence=sequence,
timestamp=timestamp,
level="error",
code="AGENT_RUN_FAILED",
title="SuperAgent 运行失败",
message=summary or "SuperAgent trace 报告运行失败。",
details=details,
)
if name in {"run.cancelled", "run.canceled"}:
return _record(
sequence=sequence,
timestamp=timestamp,
level="warning",
code="AGENT_RUN_CANCELLED",
title="SuperAgent 运行已取消",
message=summary or "SuperAgent trace 报告运行已取消。",
details=details,
)
if summary is None:
return None
return _record(
sequence=sequence,
timestamp=timestamp,
level="info",
code="AGENT_TRACE_EVENT",
title="Agent trace 事件",
message=summary,
details=details,
)
def capture_failure_record(
*,
attempt_no: int,
sequence: int,
error: BaseException,
now: Optional[datetime] = None,
) -> Dict[str, Any]:
"""Create a safe local event without serializing the exception message."""
error_name = type(error).__name__
if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,95}", error_name):
error_name = "Exception"
return _record(
sequence=sequence,
timestamp=_timestamp(now or _utc_now()),
level="warning",
code="AGENT_TRACE_CAPTURE_FAILED",
title="SuperAgent trace 采集中断",
message="远端运行可能仍在继续;当前日志连接未能完整消费。",
details={
"attempt_no": attempt_no,
"trace_event": "capture.failed",
"capture_error": error_name,
},
)
def _record(
*,
sequence: int,
timestamp: str,
level: str,
code: str,
title: str,
message: str,
details: Optional[Mapping[str, Any]] = None,
) -> Dict[str, Any]:
value: Dict[str, Any] = {
"version": TRACE_RECORD_VERSION,
"id": f"agent.{sequence:06d}.{code.lower()}",
"timestamp": timestamp,
"stage": "agent",
"level": level,
"code": code,
"title": title,
"message": message,
}
safe_details = _validated_details(details or {})
if safe_details:
value["details"] = safe_details
return value
def _validated_details(value: Mapping[str, Any]) -> Dict[str, Any]:
output: Dict[str, Any] = {}
for key, item in value.items():
if key not in _SAFE_DETAIL_KEYS or item is None or item == "":
continue
if key in {"attempt_no", "invalid_records"}:
if isinstance(item, int) and not isinstance(item, bool) and item >= 0:
output[key] = item
elif key == "remote_run_id":
if isinstance(item, str) and OPAQUE_ID_RE.fullmatch(item) is not None:
output[key] = item
elif key == "capture_error":
if isinstance(item, str) and re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,95}", item):
output[key] = item
elif isinstance(item, str) and re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._:-]{0,127}", item):
output[key] = item
return output
def validate_trace_record(record: Mapping[str, Any]) -> Dict[str, Any]:
required = {
"version",
"id",
"timestamp",
"stage",
"level",
"code",
"title",
"message",
}
optional = {"details"}
if not isinstance(record, Mapping) or not required <= set(record) <= required | optional:
raise AgentTraceStoreError("trace record field set is invalid")
if record.get("version") != TRACE_RECORD_VERSION or record.get("stage") != "agent":
raise AgentTraceStoreError("trace record version or stage is invalid")
if record.get("level") not in _LEVELS:
raise AgentTraceStoreError("trace record level is invalid")
code = record.get("code")
if not isinstance(code, str) or _CODE_RE.fullmatch(code) is None:
raise AgentTraceStoreError("trace record code is invalid")
timestamp = _bounded_text(record.get("timestamp"), 64, "timestamp")
try:
parsed_timestamp = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
except ValueError as error:
raise AgentTraceStoreError("trace record timestamp is invalid") from error
if parsed_timestamp.tzinfo is None:
raise AgentTraceStoreError("trace record timestamp is invalid")
title = sanitize_trace_summary(_bounded_text(record.get("title"), 160, "title"))
message = sanitize_trace_summary(
_bounded_text(record.get("message"), MAX_TRACE_SUMMARY_CHARS, "message")
)
if title is None or message is None:
raise AgentTraceStoreError("trace record text is invalid")
output: Dict[str, Any] = {
"version": TRACE_RECORD_VERSION,
"id": _bounded_text(record.get("id"), 180, "id"),
"timestamp": parsed_timestamp.astimezone(timezone.utc).isoformat(),
"stage": "agent",
"level": str(record["level"]),
"code": code,
"title": title,
"message": message,
}
details = record.get("details")
if details is not None:
if not isinstance(details, Mapping):
raise AgentTraceStoreError("trace record details are invalid")
safe_details = _validated_details(details)
if safe_details:
output["details"] = safe_details
return output
def _bounded_text(value: Any, maximum: int, label: str) -> str:
if not isinstance(value, str) or not value or len(value) > maximum:
raise AgentTraceStoreError(f"trace record {label} is invalid")
return value
class JsonlAgentTraceStore:
"""Append-only local trace store; raw SSE payloads are never written."""
def __init__(self, root: Union[str, Path]) -> None:
self.root = Path(root)
self._lock = threading.RLock()
def append(self, job_id: str, record: Mapping[str, Any]) -> None:
path = self._path(job_id)
value = validate_trace_record(record)
line = (
json.dumps(
value,
ensure_ascii=False,
allow_nan=False,
sort_keys=True,
separators=(",", ":"),
)
+ "\n"
).encode("utf-8")
try:
with self._lock:
if self.root.is_symlink() or (
self.root.exists() and not self.root.is_dir()
):
raise AgentTraceStoreError("agent trace root is invalid")
self.root.mkdir(mode=0o700, parents=True, exist_ok=True)
flags = os.O_WRONLY | os.O_CREAT | os.O_APPEND
if hasattr(os, "O_NOFOLLOW"):
flags |= os.O_NOFOLLOW
descriptor = os.open(path, flags, 0o600)
try:
remaining = memoryview(line)
while remaining:
written = os.write(descriptor, remaining)
if written <= 0:
raise OSError("trace append made no progress")
remaining = remaining[written:]
finally:
os.close(descriptor)
except (OSError, ValueError) as error:
raise AgentTraceStoreError("cannot append agent trace record") from error
def read(self, job_id: str) -> List[Dict[str, Any]]:
path = self._path(job_id)
try:
with self._lock:
if self.root.is_symlink():
raise AgentTraceStoreError("agent trace root is invalid")
if not path.exists():
return []
if path.is_symlink() or not path.is_file():
raise AgentTraceStoreError("agent trace path is invalid")
metadata = path.stat()
if metadata.st_size > MAX_TRACE_FILE_BYTES:
return [self._read_warning(metadata.st_mtime, 1, "file.too_large")]
modified_at = metadata.st_mtime
lines = path.read_text(encoding="utf-8").splitlines()
except AgentTraceStoreError:
raise
except (OSError, UnicodeError) as error:
raise AgentTraceStoreError("cannot read agent trace records") from error
records: List[Dict[str, Any]] = []
invalid = 0
for line in lines[:MAX_TRACE_EVENTS]:
try:
payload = json.loads(line)
records.append(validate_trace_record(payload))
except (ValueError, TypeError, AgentTraceStoreError):
invalid += 1
if len(lines) > MAX_TRACE_EVENTS:
invalid += len(lines) - MAX_TRACE_EVENTS
if invalid:
records.append(self._read_warning(modified_at, invalid, "records.invalid"))
return records
def _path(self, job_id: str) -> Path:
if not isinstance(job_id, str) or OPAQUE_ID_RE.fullmatch(job_id) is None:
raise AgentTraceStoreError("job id is invalid")
return self.root / f"{sha256(job_id.encode('utf-8')).hexdigest()}.jsonl"
@staticmethod
def _read_warning(modified_at: float, count: int, reason: str) -> Dict[str, Any]:
return _record(
sequence=MAX_TRACE_EVENTS + 1,
timestamp=datetime.fromtimestamp(modified_at, tz=timezone.utc).isoformat(),
level="warning",
code="AGENT_TRACE_RECORD_INVALID",
title="本地 Agent trace 记录不完整",
message="部分本地 trace 记录无法安全读取,已从任务日志中忽略。",
details={
"trace_event": reason,
"invalid_records": count,
},
)