|
|
@@ -0,0 +1,130 @@
|
|
|
+"""Startup verifier for deployment-owned durable Trace stores."""
|
|
|
+
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import hashlib
|
|
|
+import os
|
|
|
+from pathlib import Path
|
|
|
+from uuid import uuid4
|
|
|
+
|
|
|
+from agent import FileSystemTraceStore
|
|
|
+from agent.trace.models import Message, Trace
|
|
|
+
|
|
|
+from script_build_host.domain.errors import InvalidConfiguration
|
|
|
+
|
|
|
+
|
|
|
+class ApiRoundTripTraceStoreVerifier:
|
|
|
+ async def verify(self, trace_store: object, *, require_attachments: bool = False) -> None:
|
|
|
+ required = ("create_trace", "get_trace", "append_event", "get_events")
|
|
|
+ if any(not callable(getattr(trace_store, name, None)) for name in required):
|
|
|
+ raise InvalidConfiguration("trace store lacks durable read/write APIs")
|
|
|
+ if require_attachments and any(
|
|
|
+ not callable(getattr(trace_store, name, None))
|
|
|
+ for name in ("store_message_attachment", "read_message_attachment")
|
|
|
+ ):
|
|
|
+ raise InvalidConfiguration(
|
|
|
+ "image understanding requires durable Trace attachment storage"
|
|
|
+ )
|
|
|
+ trace_id = f"script-build-startup-probe-{uuid4()}"
|
|
|
+ trace = Trace(
|
|
|
+ trace_id=trace_id,
|
|
|
+ mode="agent",
|
|
|
+ task="durability probe",
|
|
|
+ agent_type="script_build_probe",
|
|
|
+ status="completed",
|
|
|
+ )
|
|
|
+ await trace_store.create_trace(trace) # type: ignore[attr-defined]
|
|
|
+ loaded = await trace_store.get_trace(trace_id) # type: ignore[attr-defined]
|
|
|
+ if loaded is None or loaded.trace_id != trace_id:
|
|
|
+ raise InvalidConfiguration("trace store failed the durable metadata round trip")
|
|
|
+ await trace_store.append_event(trace_id, "durability_probe", {"ok": True}) # type: ignore[attr-defined]
|
|
|
+ events = await trace_store.get_events(trace_id, 0) # type: ignore[attr-defined]
|
|
|
+ if not events:
|
|
|
+ raise InvalidConfiguration("trace store failed the durable event round trip")
|
|
|
+
|
|
|
+
|
|
|
+class FileSystemTraceStoreVerifier(ApiRoundTripTraceStoreVerifier):
|
|
|
+ """Verify the built-in store across a fresh adapter instance and fsync boundary."""
|
|
|
+
|
|
|
+ async def verify(self, trace_store: object, *, require_attachments: bool = False) -> None:
|
|
|
+ if not isinstance(trace_store, FileSystemTraceStore):
|
|
|
+ raise InvalidConfiguration("FileSystem Trace verifier received another store type")
|
|
|
+ trace_id = f"script-build-startup-probe-{uuid4()}"
|
|
|
+ trace = Trace(
|
|
|
+ trace_id=trace_id,
|
|
|
+ mode="agent",
|
|
|
+ task="durability probe",
|
|
|
+ agent_type="script_build_probe",
|
|
|
+ status="completed",
|
|
|
+ )
|
|
|
+ await trace_store.create_trace(trace)
|
|
|
+ message = Message.create(
|
|
|
+ trace_id=trace_id,
|
|
|
+ role="user",
|
|
|
+ sequence=1,
|
|
|
+ content="durability probe",
|
|
|
+ )
|
|
|
+ await trace_store.add_message(message)
|
|
|
+ await trace_store.update_trace(trace_id, head_sequence=1)
|
|
|
+ await trace_store.append_event(trace_id, "durability_probe", {"ok": True})
|
|
|
+
|
|
|
+ attachment_ref = None
|
|
|
+ attachment_content = b"script-build-trace-attachment-probe"
|
|
|
+ if require_attachments:
|
|
|
+ attachment_ref = await trace_store.store_message_attachment(
|
|
|
+ trace_id=trace_id,
|
|
|
+ message_id=message.message_id,
|
|
|
+ media_type="application/octet-stream",
|
|
|
+ content=attachment_content,
|
|
|
+ sha256=f"sha256:{hashlib.sha256(attachment_content).hexdigest()}",
|
|
|
+ )
|
|
|
+
|
|
|
+ _fsync_tree(trace_store.base_path / trace_id)
|
|
|
+ reloaded = FileSystemTraceStore(str(trace_store.base_path))
|
|
|
+ loaded_trace = await reloaded.get_trace(trace_id)
|
|
|
+ loaded_message = await reloaded.get_message(message.message_id)
|
|
|
+ loaded_events = await reloaded.get_events(trace_id, 0)
|
|
|
+ if (
|
|
|
+ loaded_trace is None
|
|
|
+ or loaded_trace.trace_id != trace_id
|
|
|
+ or loaded_trace.head_sequence != 1
|
|
|
+ or loaded_message is None
|
|
|
+ or loaded_message.trace_id != trace_id
|
|
|
+ or loaded_message.content != "durability probe"
|
|
|
+ or not any(item.get("event") == "durability_probe" for item in loaded_events)
|
|
|
+ ):
|
|
|
+ raise InvalidConfiguration("FileSystem Trace store failed the reopen round trip")
|
|
|
+ if attachment_ref is not None:
|
|
|
+ try:
|
|
|
+ loaded_content = await reloaded.read_message_attachment(attachment_ref)
|
|
|
+ except Exception as exc:
|
|
|
+ raise InvalidConfiguration(
|
|
|
+ "FileSystem Trace attachment failed the reopen round trip"
|
|
|
+ ) from exc
|
|
|
+ if loaded_content != attachment_content:
|
|
|
+ raise InvalidConfiguration(
|
|
|
+ "FileSystem Trace attachment changed across the reopen boundary"
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def _fsync_tree(root: Path) -> None:
|
|
|
+ if not root.is_dir():
|
|
|
+ raise InvalidConfiguration("FileSystem Trace probe directory was not created")
|
|
|
+ for path in sorted(root.rglob("*"), reverse=True):
|
|
|
+ if path.is_file():
|
|
|
+ with path.open("rb") as handle:
|
|
|
+ os.fsync(handle.fileno())
|
|
|
+ elif path.is_dir():
|
|
|
+ _fsync_directory(path)
|
|
|
+ _fsync_directory(root)
|
|
|
+
|
|
|
+
|
|
|
+def _fsync_directory(path: Path) -> None:
|
|
|
+ descriptor = os.open(path, os.O_RDONLY)
|
|
|
+ try:
|
|
|
+ os.fsync(descriptor)
|
|
|
+ finally:
|
|
|
+ os.close(descriptor)
|
|
|
+
|
|
|
+
|
|
|
+__all__ = ["ApiRoundTripTraceStoreVerifier", "FileSystemTraceStoreVerifier"]
|