feat: prepare ARR for controlled public deployment
This commit is contained in:
Wyndham ARR
488
arr_processing/runtime_writeback.py
Normal file
488
arr_processing/runtime_writeback.py
Normal file
@@ -0,0 +1,488 @@
|
||||
"""Trusted runtime adapter: publish Agent files, sign the result, call ARR."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import stat
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from datetime import date
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Mapping, Optional, Protocol, Sequence, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from arr_ingestion.contracts import OPAQUE_ID_RE, PROCESSOR_VERSION_RE, SHA256_RE, XLSX_MIME
|
||||
from arr_processing.config import ResultVerificationConfig
|
||||
from arr_processing.contracts import ProcessingResult
|
||||
from arr_processing.errors import ProcessingError
|
||||
from arr_storage.aliyun_oss_v2 import AliyunOssConfig, AliyunOssV2Client
|
||||
from arr_storage.exchange import OutputExchangeConfig, OutputExchangePublisher
|
||||
|
||||
|
||||
FROZEN_FIELDS = {
|
||||
"contract_version",
|
||||
"job_id",
|
||||
"source_file_id",
|
||||
"status",
|
||||
"business_date",
|
||||
"processor_result",
|
||||
"files",
|
||||
}
|
||||
PROCESSOR_RESULT_FIELDS = {
|
||||
"version",
|
||||
"status",
|
||||
"business_date",
|
||||
"message",
|
||||
"metrics",
|
||||
"outputs",
|
||||
"errors",
|
||||
}
|
||||
FILE_FIELDS = {"filename", "local_path"}
|
||||
FILE_ROLES = {"daily_report", "result_json", "structured_result", "exception_report"}
|
||||
ROLE_MIME = {
|
||||
"daily_report": XLSX_MIME,
|
||||
"result_json": "application/json",
|
||||
"structured_result_json": "application/json",
|
||||
"exception_report": XLSX_MIME,
|
||||
}
|
||||
RESULT_ROLE = {
|
||||
"daily_report": "daily_report",
|
||||
"result_json": "result_json",
|
||||
"structured_result": "structured_result_json",
|
||||
"exception_report": "exception_report",
|
||||
}
|
||||
SAFE_FILENAME_RE = re.compile(r"^[^/\\\x00-\x1f\x7f]{1,255}$")
|
||||
|
||||
|
||||
class RuntimeWritebackError(RuntimeError):
|
||||
def __init__(self, code: str, safe_message: str, *, retryable: bool = False) -> None:
|
||||
super().__init__(safe_message)
|
||||
self.code = code
|
||||
self.safe_message = safe_message
|
||||
self.retryable = retryable
|
||||
|
||||
|
||||
class OutputPublisher(Protocol):
|
||||
def publish(
|
||||
self,
|
||||
*,
|
||||
file_handle: str,
|
||||
source: Path,
|
||||
mime_type: str,
|
||||
metadata: Mapping[str, str],
|
||||
) -> Any:
|
||||
...
|
||||
|
||||
|
||||
class HttpClient(Protocol):
|
||||
def post(self, url: str, **kwargs: Any) -> Any:
|
||||
...
|
||||
|
||||
|
||||
def _strict_json(raw: bytes) -> Mapping[str, Any]:
|
||||
def pairs(values: Sequence[Tuple[str, Any]]) -> Dict[str, Any]:
|
||||
output: Dict[str, Any] = {}
|
||||
for key, value in values:
|
||||
if key in output:
|
||||
raise ValueError("duplicate JSON key")
|
||||
output[key] = value
|
||||
return output
|
||||
|
||||
try:
|
||||
value = json.loads(raw.decode("utf-8"), object_pairs_hook=pairs)
|
||||
except (UnicodeError, ValueError, json.JSONDecodeError):
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", "Agent result is not strict JSON"
|
||||
) from None
|
||||
if not isinstance(value, Mapping):
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", "Agent result must be a JSON object"
|
||||
)
|
||||
return value
|
||||
|
||||
|
||||
def _opaque(value: Any, label: str) -> str:
|
||||
if not isinstance(value, str) or not OPAQUE_ID_RE.fullmatch(value):
|
||||
raise RuntimeWritebackError("RUNTIME_RESULT_INVALID", f"{label} is invalid")
|
||||
return value
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FrozenFile:
|
||||
filename: str
|
||||
path: Path
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FrozenAgentResult:
|
||||
job_id: str
|
||||
source_file_id: str
|
||||
status: str
|
||||
business_date: Optional[date]
|
||||
files: Mapping[str, Optional[FrozenFile]]
|
||||
|
||||
@classmethod
|
||||
def parse(cls, raw: bytes, *, output_root: Path) -> "FrozenAgentResult":
|
||||
values = _strict_json(raw)
|
||||
if set(values) != FROZEN_FIELDS or values.get("contract_version") != "arr-opera-daily-agent-1":
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", "Agent result contract is invalid"
|
||||
)
|
||||
job_id = _opaque(values.get("job_id"), "job identifier")
|
||||
source_file_id = _opaque(values.get("source_file_id"), "source file identifier")
|
||||
status = values.get("status")
|
||||
if status not in {"success", "failed"}:
|
||||
raise RuntimeWritebackError("RUNTIME_RESULT_INVALID", "Agent status is invalid")
|
||||
raw_date = values.get("business_date")
|
||||
try:
|
||||
if raw_date == "":
|
||||
business_date = None
|
||||
elif isinstance(raw_date, str):
|
||||
business_date = date.fromisoformat(raw_date)
|
||||
else:
|
||||
raise ValueError
|
||||
except ValueError:
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", "Agent business date is invalid"
|
||||
) from None
|
||||
processor_business_date = business_date.isoformat() if business_date else None
|
||||
processor = values.get("processor_result")
|
||||
if not isinstance(processor, Mapping) or set(processor) != PROCESSOR_RESULT_FIELDS:
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", "processor result contract is invalid"
|
||||
)
|
||||
if (
|
||||
processor.get("version") != "3.0"
|
||||
or processor.get("status") != status
|
||||
or processor.get("business_date") != processor_business_date
|
||||
or not isinstance(processor.get("message"), str)
|
||||
or not isinstance(processor.get("metrics"), Mapping)
|
||||
or not isinstance(processor.get("outputs"), Mapping)
|
||||
or not isinstance(processor.get("errors"), list)
|
||||
):
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", "processor result identity is invalid"
|
||||
)
|
||||
raw_files = values.get("files")
|
||||
if not isinstance(raw_files, Mapping) or set(raw_files) != FILE_ROLES:
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", "Agent output file contract is invalid"
|
||||
)
|
||||
root = output_root.resolve(strict=True)
|
||||
if not root.is_dir():
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_OUTPUT_ROOT_INVALID", "Agent output root is unavailable"
|
||||
)
|
||||
files: Dict[str, Optional[FrozenFile]] = {}
|
||||
for role in sorted(FILE_ROLES):
|
||||
item = raw_files.get(role)
|
||||
if item is None:
|
||||
files[role] = None
|
||||
continue
|
||||
if not isinstance(item, Mapping) or set(item) != FILE_FIELDS:
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", f"{role} file contract is invalid"
|
||||
)
|
||||
filename = item.get("filename")
|
||||
local_path = item.get("local_path")
|
||||
if (
|
||||
not isinstance(filename, str)
|
||||
or SAFE_FILENAME_RE.fullmatch(filename) is None
|
||||
or not isinstance(local_path, str)
|
||||
or not Path(local_path).is_absolute()
|
||||
):
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", f"{role} file identity is invalid"
|
||||
)
|
||||
path = Path(local_path)
|
||||
try:
|
||||
details = path.lstat()
|
||||
resolved = path.resolve(strict=True)
|
||||
resolved.relative_to(root)
|
||||
except (OSError, ValueError):
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_OUTPUT_INVALID", f"{role} output is outside the approved root"
|
||||
) from None
|
||||
if path.is_symlink() or not stat.S_ISREG(details.st_mode) or resolved.name != filename:
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_OUTPUT_INVALID", f"{role} output is not an approved file"
|
||||
)
|
||||
expected_extension = ".json" if role in {"result_json", "structured_result"} else ".xlsx"
|
||||
if not filename.lower().endswith(expected_extension):
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_OUTPUT_INVALID", f"{role} output type is invalid"
|
||||
)
|
||||
files[role] = FrozenFile(filename, resolved)
|
||||
|
||||
if files["result_json"] is None or files["structured_result"] is None:
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", "required Agent result files are missing"
|
||||
)
|
||||
if status == "success":
|
||||
valid_shape = (
|
||||
business_date is not None
|
||||
and files["daily_report"] is not None
|
||||
and files["exception_report"] is None
|
||||
)
|
||||
else:
|
||||
valid_shape = files["daily_report"] is None and files["exception_report"] is not None
|
||||
if not valid_shape:
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", "Agent result file shape is invalid"
|
||||
)
|
||||
return cls(job_id, source_file_id, str(status), business_date, files)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentRuntimeWritebackAdapter:
|
||||
publisher: OutputPublisher
|
||||
signer: Any
|
||||
http_client: HttpClient
|
||||
callback_url: str
|
||||
processor_version: str
|
||||
rule_set_sha256: str
|
||||
max_http_attempts: int = 3
|
||||
retry_delay_seconds: float = 0.5
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
parsed = urlparse(self.callback_url)
|
||||
localhost = parsed.hostname in {"127.0.0.1", "localhost", "::1"}
|
||||
if (
|
||||
parsed.scheme not in ({"http", "https"} if localhost else {"https"})
|
||||
or not parsed.netloc
|
||||
or parsed.query
|
||||
or parsed.fragment
|
||||
):
|
||||
raise ValueError("ARR Agent callback URL must be an HTTPS URL")
|
||||
if not PROCESSOR_VERSION_RE.fullmatch(self.processor_version):
|
||||
raise ValueError("processor version is invalid")
|
||||
if not SHA256_RE.fullmatch(self.rule_set_sha256):
|
||||
raise ValueError("rule set identity is invalid")
|
||||
if self.max_http_attempts < 1 or self.retry_delay_seconds < 0:
|
||||
raise ValueError("runtime writeback retry policy is invalid")
|
||||
|
||||
def deliver(
|
||||
self,
|
||||
raw_frozen_result: bytes,
|
||||
*,
|
||||
output_root: Path,
|
||||
attempt_no: int,
|
||||
remote_run_id: str,
|
||||
) -> Mapping[str, Any]:
|
||||
frozen = FrozenAgentResult.parse(raw_frozen_result, output_root=output_root)
|
||||
remote_run_id = _opaque(remote_run_id, "remote run identifier")
|
||||
if not isinstance(attempt_no, int) or isinstance(attempt_no, bool) or not 1 <= attempt_no <= 9999:
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_RESULT_INVALID", "attempt number is invalid"
|
||||
)
|
||||
artifacts: Dict[str, Any] = {
|
||||
"daily_report": None,
|
||||
"result_json": None,
|
||||
"structured_result_json": None,
|
||||
"exception_report": None,
|
||||
}
|
||||
for frozen_role in sorted(FILE_ROLES):
|
||||
item = frozen.files[frozen_role]
|
||||
if item is None:
|
||||
continue
|
||||
role = RESULT_ROLE[frozen_role]
|
||||
digest, byte_size = self._identity(item.path)
|
||||
file_handle = self._file_handle(
|
||||
frozen.job_id,
|
||||
attempt_no,
|
||||
remote_run_id,
|
||||
role,
|
||||
digest,
|
||||
)
|
||||
mime_type = ROLE_MIME[role]
|
||||
metadata = {
|
||||
"arr-exchange-schema": "1.0",
|
||||
"arr-file-handle": file_handle,
|
||||
"arr-job-id": frozen.job_id,
|
||||
"arr-attempt-no": str(attempt_no),
|
||||
"arr-role": role,
|
||||
"arr-sha256": digest,
|
||||
"arr-byte-size": str(byte_size),
|
||||
"arr-mime-type": mime_type,
|
||||
}
|
||||
try:
|
||||
self.publisher.publish(
|
||||
file_handle=file_handle,
|
||||
source=item.path,
|
||||
mime_type=mime_type,
|
||||
metadata=metadata,
|
||||
)
|
||||
except Exception as error:
|
||||
if isinstance(error, RuntimeWritebackError):
|
||||
raise
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_OUTPUT_PUBLISH_FAILED",
|
||||
"Agent output could not be published",
|
||||
retryable=True,
|
||||
) from None
|
||||
artifacts[role] = {
|
||||
"file_handle": file_handle,
|
||||
"original_filename": item.filename,
|
||||
"sha256": digest,
|
||||
"byte_size": byte_size,
|
||||
"mime_type": mime_type,
|
||||
}
|
||||
result = ProcessingResult.from_dict(
|
||||
{
|
||||
"contract_version": "arr-opera-daily-result-1",
|
||||
"delivery_id": self._delivery_id(frozen.job_id, attempt_no, remote_run_id),
|
||||
"job_id": frozen.job_id,
|
||||
"attempt_no": attempt_no,
|
||||
"remote_run_id": remote_run_id,
|
||||
"status": frozen.status,
|
||||
"business_date": frozen.business_date.isoformat() if frozen.business_date else None,
|
||||
"processor_version": self.processor_version,
|
||||
"rule_set_sha256": self.rule_set_sha256,
|
||||
"result_schema_version": "3.0",
|
||||
"artifacts": artifacts,
|
||||
}
|
||||
)
|
||||
signed = self.signer.sign(result)
|
||||
return self._post_same_bytes(signed)
|
||||
|
||||
def _post_same_bytes(self, signed: bytes) -> Mapping[str, Any]:
|
||||
for attempt in range(1, self.max_http_attempts + 1):
|
||||
try:
|
||||
response = self.http_client.post(
|
||||
self.callback_url,
|
||||
content=signed,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "arr-agent-runtime-writeback/1.0",
|
||||
},
|
||||
)
|
||||
except Exception:
|
||||
if attempt >= self.max_http_attempts:
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_CALLBACK_UNAVAILABLE",
|
||||
"ARR callback did not respond",
|
||||
retryable=True,
|
||||
) from None
|
||||
time.sleep(self.retry_delay_seconds * attempt)
|
||||
continue
|
||||
status = int(getattr(response, "status_code", 0) or 0)
|
||||
if status == 200:
|
||||
try:
|
||||
payload = response.json()
|
||||
except Exception:
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_CALLBACK_PROTOCOL_INVALID",
|
||||
"ARR callback returned an invalid response",
|
||||
) from None
|
||||
if (
|
||||
not isinstance(payload, Mapping)
|
||||
or payload.get("ok") is not True
|
||||
or not isinstance(payload.get("data"), Mapping)
|
||||
):
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_CALLBACK_PROTOCOL_INVALID",
|
||||
"ARR callback returned an invalid response",
|
||||
)
|
||||
return dict(payload["data"])
|
||||
if status in {408, 429} or status >= 500:
|
||||
if attempt < self.max_http_attempts:
|
||||
time.sleep(self.retry_delay_seconds * attempt)
|
||||
continue
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_CALLBACK_UNAVAILABLE",
|
||||
"ARR callback is temporarily unavailable",
|
||||
retryable=True,
|
||||
)
|
||||
raise RuntimeWritebackError(
|
||||
"RUNTIME_CALLBACK_REJECTED",
|
||||
"ARR rejected the Agent result",
|
||||
)
|
||||
raise AssertionError("unreachable retry loop")
|
||||
|
||||
@staticmethod
|
||||
def _identity(path: Path) -> Tuple[str, int]:
|
||||
digest = hashlib.sha256()
|
||||
size = 0
|
||||
with path.open("rb") as source:
|
||||
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
||||
size += len(chunk)
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest(), size
|
||||
|
||||
@staticmethod
|
||||
def _file_handle(
|
||||
job_id: str,
|
||||
attempt_no: int,
|
||||
remote_run_id: str,
|
||||
role: str,
|
||||
sha256: str,
|
||||
) -> str:
|
||||
value = "\x1f".join((job_id, str(attempt_no), remote_run_id, role, sha256))
|
||||
return "arrout_" + hashlib.sha256(value.encode("utf-8")).hexdigest()[:48]
|
||||
|
||||
@staticmethod
|
||||
def _delivery_id(job_id: str, attempt_no: int, remote_run_id: str) -> str:
|
||||
value = "\x1f".join((job_id, str(attempt_no), remote_run_id, "delivery"))
|
||||
return "arrdel_" + hashlib.sha256(value.encode("utf-8")).hexdigest()[:48]
|
||||
|
||||
|
||||
def _parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(prog="arr-agent-writeback")
|
||||
parser.add_argument("--frozen-result", type=Path, required=True)
|
||||
parser.add_argument("--output-root", type=Path, required=True)
|
||||
parser.add_argument("--attempt-no", type=int, required=True)
|
||||
parser.add_argument("--remote-run-id", required=True)
|
||||
parser.add_argument("--processor-version", required=True)
|
||||
parser.add_argument("--rule-set-sha256", required=True)
|
||||
parser.add_argument(
|
||||
"--callback-url",
|
||||
default=os.environ.get("ARR_AGENT_CALLBACK_URL", ""),
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: Optional[Sequence[str]] = None) -> int:
|
||||
args = _parser().parse_args(argv)
|
||||
if not args.callback_url:
|
||||
print(json.dumps({"ok": False, "error": {"code": "CALLBACK_CONFIG_MISSING"}}))
|
||||
return 2
|
||||
try:
|
||||
import httpx
|
||||
|
||||
oss_client = AliyunOssV2Client(AliyunOssConfig.from_environment())
|
||||
oss_client.assert_immutable_writes_supported()
|
||||
publisher = OutputExchangePublisher(
|
||||
oss_client,
|
||||
OutputExchangeConfig.from_environment(),
|
||||
)
|
||||
signing = ResultVerificationConfig.from_environment().signer()
|
||||
with httpx.Client(timeout=httpx.Timeout(30.0, connect=10.0)) as http_client:
|
||||
receipt = AgentRuntimeWritebackAdapter(
|
||||
publisher=publisher,
|
||||
signer=signing,
|
||||
http_client=http_client,
|
||||
callback_url=args.callback_url,
|
||||
processor_version=args.processor_version,
|
||||
rule_set_sha256=args.rule_set_sha256,
|
||||
).deliver(
|
||||
args.frozen_result.read_bytes(),
|
||||
output_root=args.output_root,
|
||||
attempt_no=args.attempt_no,
|
||||
remote_run_id=args.remote_run_id,
|
||||
)
|
||||
oss_client.close()
|
||||
except (OSError, ValueError, ProcessingError, RuntimeWritebackError) as error:
|
||||
code = getattr(error, "code", "RUNTIME_WRITEBACK_FAILED")
|
||||
print(json.dumps({"ok": False, "error": {"code": code}}))
|
||||
return 3
|
||||
print(json.dumps({"ok": True, "data": receipt}, ensure_ascii=False))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user