😏
This commit is contained in:
Executable
+899
@@ -0,0 +1,899 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Low-level utilities shared by Codex MMO.
|
||||
|
||||
The Python control plane intentionally has no third-party package dependency;
|
||||
documented external executables still apply to the selected execution path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import datetime as dt
|
||||
import fcntl
|
||||
import hashlib
|
||||
import ipaddress
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import signal
|
||||
import socket
|
||||
import tempfile
|
||||
import time
|
||||
import tomllib
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from collections.abc import Iterable, Iterator, Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from mmo_version import MMO_SCHEMA_VERSION, PACKAGE_VERSION
|
||||
|
||||
ID_PATTERN = re.compile(r"[a-z][a-z0-9_.-]{1,63}")
|
||||
SAFE_JOB_ID = re.compile(r"[A-Za-z0-9_.-]{8,128}")
|
||||
|
||||
_URI_SCHEME = re.compile(r"[A-Za-z][A-Za-z0-9+.-]*", re.ASCII)
|
||||
_URI_UNRESERVED = frozenset("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~")
|
||||
_URI_SUB_DELIMS = frozenset("!$&'()*+,;=")
|
||||
_URI_HEXDIGITS = frozenset("0123456789ABCDEFabcdef")
|
||||
_URI_PCHAR = _URI_UNRESERVED | _URI_SUB_DELIMS | frozenset(":@")
|
||||
|
||||
|
||||
class _NoRedirectHandler(urllib.request.HTTPRedirectHandler):
|
||||
"""Keep readiness checks scoped to the configured endpoint response."""
|
||||
|
||||
def redirect_request(self, *args: Any, **kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return dt.datetime.now(dt.UTC).isoformat(timespec="milliseconds")
|
||||
|
||||
|
||||
def shell_exit_status(returncode: int) -> int:
|
||||
"""Translate a Python child return code to the shell's signal convention."""
|
||||
|
||||
return 128 - returncode if returncode < 0 else returncode
|
||||
|
||||
|
||||
def install_root() -> Path:
|
||||
override = os.environ.get("MMO_INSTALL_ROOT")
|
||||
if override:
|
||||
return Path(override).expanduser().resolve()
|
||||
return Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def package_version() -> str:
|
||||
return PACKAGE_VERSION
|
||||
|
||||
|
||||
def install_runtime_path() -> Path:
|
||||
return install_root() / "config" / "runtime.json"
|
||||
|
||||
|
||||
def load_install_runtime() -> dict[str, Any]:
|
||||
path = install_runtime_path()
|
||||
try:
|
||||
data = strict_json_loads(path.read_text(encoding="utf-8"))
|
||||
except FileNotFoundError as exc:
|
||||
raise RuntimeError(f"installed runtime configuration is missing: {path}") from exc
|
||||
if not isinstance(data, dict):
|
||||
raise RuntimeError(f"installed runtime configuration root must be an object: {path}")
|
||||
schema_version = data.get("schema_version")
|
||||
if (
|
||||
not isinstance(schema_version, int)
|
||||
or isinstance(schema_version, bool)
|
||||
or schema_version != MMO_SCHEMA_VERSION
|
||||
):
|
||||
raise RuntimeError(f"unsupported installed runtime schema in {path}")
|
||||
for key in ("install_root", "config_root", "state_root", "bin_dir", "python_bin"):
|
||||
value = data.get(key)
|
||||
if not isinstance(value, str) or not value:
|
||||
raise RuntimeError(f"installed runtime field {key!r} is invalid in {path}")
|
||||
for key in ("install_root", "config_root", "state_root", "bin_dir"):
|
||||
if not Path(data[key]).expanduser().is_absolute():
|
||||
raise RuntimeError(f"installed runtime path {key!r} is not absolute in {path}")
|
||||
return data
|
||||
|
||||
|
||||
def valid_http_header_value(value: str) -> bool:
|
||||
"""Match the byte validity rule used by Rust's ``http::HeaderValue``.
|
||||
|
||||
RFC 9110 permits HTAB plus visible/obsolete field-value bytes. The Rust
|
||||
HTTP stack used by both Codex and Switchyard rejects every other C0 control
|
||||
byte and DEL. Encoding first also rejects lone Python surrogates.
|
||||
"""
|
||||
|
||||
try:
|
||||
encoded = value.encode("utf-8")
|
||||
except UnicodeEncodeError:
|
||||
return False
|
||||
return all(byte == 0x09 or (byte >= 0x20 and byte != 0x7F) for byte in encoded)
|
||||
|
||||
|
||||
def _valid_uri_component(value: str, allowed: frozenset[str]) -> bool:
|
||||
"""Validate one RFC 3986 component, including percent triplets."""
|
||||
|
||||
index = 0
|
||||
while index < len(value):
|
||||
character = value[index]
|
||||
if character == "%":
|
||||
if (
|
||||
index + 2 >= len(value)
|
||||
or value[index + 1] not in _URI_HEXDIGITS
|
||||
or value[index + 2] not in _URI_HEXDIGITS
|
||||
):
|
||||
return False
|
||||
index += 3
|
||||
continue
|
||||
if character not in allowed:
|
||||
return False
|
||||
index += 1
|
||||
return True
|
||||
|
||||
|
||||
def _valid_uri_authority(authority: str) -> bool:
|
||||
"""Validate RFC 3986 authority syntax without performing DNS resolution."""
|
||||
|
||||
if authority.count("@") > 1:
|
||||
return False
|
||||
if "@" in authority:
|
||||
userinfo, host_port = authority.split("@", 1)
|
||||
if not _valid_uri_component(userinfo, _URI_UNRESERVED | _URI_SUB_DELIMS | frozenset(":")):
|
||||
return False
|
||||
else:
|
||||
host_port = authority
|
||||
|
||||
if host_port.startswith("["):
|
||||
closing = host_port.find("]")
|
||||
if closing < 0:
|
||||
return False
|
||||
literal = host_port[1:closing]
|
||||
remainder = host_port[closing + 1 :]
|
||||
if remainder and not remainder.startswith(":"):
|
||||
return False
|
||||
if remainder[1:] and (
|
||||
not remainder[1:].isascii()
|
||||
or not all("0" <= character <= "9" for character in remainder[1:])
|
||||
):
|
||||
return False
|
||||
ipv_future = re.fullmatch(
|
||||
r"[Vv][0-9A-Fa-f]+\.[A-Za-z0-9._~!$&'()*+,;=:-]+",
|
||||
literal,
|
||||
re.ASCII,
|
||||
)
|
||||
if ipv_future is None:
|
||||
try:
|
||||
ipaddress.IPv6Address(literal)
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
if "[" in host_port or "]" in host_port or host_port.count(":") > 1:
|
||||
return False
|
||||
if ":" in host_port:
|
||||
host, port = host_port.rsplit(":", 1)
|
||||
if port and (not port.isascii() or not all("0" <= character <= "9" for character in port)):
|
||||
return False
|
||||
else:
|
||||
host = host_port
|
||||
# A digit-and-dot string that is not an IPv4 address remains a valid
|
||||
# reg-name. RFC 3986 deliberately gives a valid IPv4 address first-match
|
||||
# precedence; it does not otherwise reserve that spelling.
|
||||
return _valid_uri_component(host, _URI_UNRESERVED | _URI_SUB_DELIMS)
|
||||
|
||||
|
||||
def valid_absolute_uri(value: str) -> bool:
|
||||
"""Return whether ``value`` is an absolute URI under RFC 3986.
|
||||
|
||||
This is syntax validation only: it neither resolves hostnames nor imposes
|
||||
scheme-specific semantics. URI references are intentionally rejected.
|
||||
"""
|
||||
|
||||
if not isinstance(value, str):
|
||||
return False
|
||||
colon = value.find(":")
|
||||
if colon <= 0 or _URI_SCHEME.fullmatch(value[:colon]) is None:
|
||||
return False
|
||||
remainder = value[colon + 1 :]
|
||||
|
||||
if "#" in remainder:
|
||||
hierarchical, fragment = remainder.split("#", 1)
|
||||
if not _valid_uri_component(fragment, _URI_PCHAR | frozenset("/?")):
|
||||
return False
|
||||
else:
|
||||
hierarchical = remainder
|
||||
if "?" in hierarchical:
|
||||
hierarchical, query = hierarchical.split("?", 1)
|
||||
if not _valid_uri_component(query, _URI_PCHAR | frozenset("/?")):
|
||||
return False
|
||||
|
||||
if hierarchical.startswith("//"):
|
||||
authority_and_path = hierarchical[2:]
|
||||
slash = authority_and_path.find("/")
|
||||
if slash < 0:
|
||||
authority, path = authority_and_path, ""
|
||||
else:
|
||||
authority, path = authority_and_path[:slash], authority_and_path[slash:]
|
||||
return _valid_uri_authority(authority) and _valid_uri_component(
|
||||
path, _URI_PCHAR | frozenset("/")
|
||||
)
|
||||
|
||||
# The remaining hier-part alternatives are path-absolute, path-rootless,
|
||||
# and path-empty. A leading "//" was handled above; every non-empty
|
||||
# rootless path therefore has the required non-empty first segment.
|
||||
return _valid_uri_component(hierarchical, _URI_PCHAR | frozenset("/"))
|
||||
|
||||
|
||||
def config_root(runtime: Mapping[str, Any] | None = None) -> Path:
|
||||
override = os.environ.get("MMO_CONFIG_ROOT")
|
||||
if override:
|
||||
return Path(override).expanduser().resolve()
|
||||
data = runtime or load_install_runtime()
|
||||
return Path(str(data["config_root"])).expanduser().resolve()
|
||||
|
||||
|
||||
def state_root(runtime: Mapping[str, Any] | None = None) -> Path:
|
||||
override = os.environ.get("MMO_STATE_ROOT")
|
||||
if override:
|
||||
return Path(override).expanduser().resolve()
|
||||
data = runtime or load_install_runtime()
|
||||
return Path(str(data["state_root"])).expanduser().resolve()
|
||||
|
||||
|
||||
def atomic_write_bytes(path: Path, data: bytes, mode: int = 0o600) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
|
||||
try:
|
||||
with os.fdopen(fd, "wb") as handle:
|
||||
handle.write(data)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.chmod(temporary, mode)
|
||||
os.replace(temporary, path)
|
||||
with contextlib.suppress(OSError):
|
||||
directory_fd = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY)
|
||||
try:
|
||||
os.fsync(directory_fd)
|
||||
finally:
|
||||
os.close(directory_fd)
|
||||
finally:
|
||||
with contextlib.suppress(FileNotFoundError):
|
||||
os.unlink(temporary)
|
||||
|
||||
|
||||
def atomic_write_text(path: Path, text: str, mode: int = 0o600) -> None:
|
||||
atomic_write_bytes(path, text.encode("utf-8"), mode)
|
||||
|
||||
|
||||
def atomic_write_json(path: Path, data: Any, mode: int = 0o600) -> None:
|
||||
atomic_write_text(
|
||||
path,
|
||||
json.dumps(data, ensure_ascii=False, indent=2, sort_keys=True, allow_nan=False) + "\n",
|
||||
mode,
|
||||
)
|
||||
|
||||
|
||||
def _reject_json_constant(value: str) -> Any:
|
||||
raise ValueError(f"non-standard JSON constant: {value}")
|
||||
|
||||
|
||||
def _finite_json_float(value: str) -> float:
|
||||
parsed = float(value)
|
||||
if not math.isfinite(parsed):
|
||||
raise ValueError(f"JSON number is outside the supported finite range: {value}")
|
||||
return parsed
|
||||
|
||||
|
||||
def _unique_json_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
|
||||
result: dict[str, Any] = {}
|
||||
for key, item in pairs:
|
||||
if key in result:
|
||||
raise ValueError(f"duplicate JSON object member: {key!r}")
|
||||
result[key] = item
|
||||
return result
|
||||
|
||||
|
||||
def validate_json_unicode(value: Any) -> None:
|
||||
"""Reject strings that contain lone UTF-16 surrogate code points."""
|
||||
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value.encode("utf-8")
|
||||
except UnicodeEncodeError as exc:
|
||||
raise ValueError("JSON strings must contain valid Unicode scalar values") from exc
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
validate_json_unicode(item)
|
||||
elif isinstance(value, Mapping):
|
||||
for key, item in value.items():
|
||||
validate_json_unicode(key)
|
||||
validate_json_unicode(item)
|
||||
|
||||
|
||||
def strict_json_decoder() -> json.JSONDecoder:
|
||||
"""Build the decoder shared by full-document and embedded JSON parsing."""
|
||||
|
||||
return json.JSONDecoder(
|
||||
parse_constant=_reject_json_constant,
|
||||
parse_float=_finite_json_float,
|
||||
object_pairs_hook=_unique_json_object,
|
||||
)
|
||||
|
||||
|
||||
def strict_json_loads(value: str | bytes | bytearray) -> Any:
|
||||
"""Parse interoperable RFC 8259 JSON without ambiguous extensions."""
|
||||
|
||||
parsed = json.loads(
|
||||
value,
|
||||
parse_constant=_reject_json_constant,
|
||||
parse_float=_finite_json_float,
|
||||
object_pairs_hook=_unique_json_object,
|
||||
)
|
||||
validate_json_unicode(parsed)
|
||||
return parsed
|
||||
|
||||
|
||||
def read_json(path: Path) -> Any:
|
||||
return strict_json_loads(path.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def read_json_object(path: Path, *, label: str = "JSON document") -> dict[str, Any]:
|
||||
value = read_json(path)
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError(f"{label} root must be an object: {path}")
|
||||
return value
|
||||
|
||||
|
||||
def read_toml(path: Path) -> dict[str, Any]:
|
||||
try:
|
||||
with path.open("rb") as handle:
|
||||
value = tomllib.load(handle)
|
||||
except tomllib.TOMLDecodeError as exc:
|
||||
raise ValueError(f"invalid TOML in {path}: {exc}") from exc
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError(f"TOML root must be a table: {path}")
|
||||
return value
|
||||
|
||||
|
||||
def canonical_json_bytes(value: Any) -> bytes:
|
||||
return json.dumps(
|
||||
value,
|
||||
ensure_ascii=False,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
allow_nan=False,
|
||||
).encode("utf-8")
|
||||
|
||||
|
||||
def sha256_bytes(data: bytes) -> str:
|
||||
return hashlib.sha256(data).hexdigest()
|
||||
|
||||
|
||||
def sha256_file(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def stable_hash(value: Any) -> str:
|
||||
return sha256_bytes(canonical_json_bytes(value))
|
||||
|
||||
|
||||
def validate_id(value: str, label: str = "identifier") -> str:
|
||||
if not isinstance(value, str) or not ID_PATTERN.fullmatch(value):
|
||||
raise ValueError(
|
||||
f"invalid {label} {value!r}; use 2-64 lowercase letters, digits, '.', '_' or '-'"
|
||||
)
|
||||
return value
|
||||
|
||||
|
||||
def safe_name(value: str) -> str:
|
||||
cleaned = re.sub(r"[^A-Za-z0-9_.-]+", "-", value).strip("-.")
|
||||
return cleaned[:80] or "item"
|
||||
|
||||
|
||||
def toml_quote(value: str) -> str:
|
||||
try:
|
||||
value.encode("utf-8")
|
||||
except UnicodeEncodeError as exc:
|
||||
raise ValueError("generated TOML strings must contain valid Unicode scalar values") from exc
|
||||
# JSON and TOML basic-string escapes overlap for the values emitted here,
|
||||
# except that JSON permits a raw DEL character while TOML forbids it.
|
||||
return json.dumps(value, ensure_ascii=False, allow_nan=False).replace("\x7f", "\\u007F")
|
||||
|
||||
|
||||
def _toml_scalar(value: Any) -> str:
|
||||
if isinstance(value, bool):
|
||||
return "true" if value else "false"
|
||||
if isinstance(value, int):
|
||||
return str(value)
|
||||
if isinstance(value, float):
|
||||
if value != value or value in (float("inf"), float("-inf")):
|
||||
raise ValueError("non-finite floats are not supported in generated TOML")
|
||||
return repr(value)
|
||||
if isinstance(value, str):
|
||||
return toml_quote(value)
|
||||
if isinstance(value, list):
|
||||
if any(isinstance(item, dict) for item in value):
|
||||
raise TypeError("array-of-table values are emitted separately")
|
||||
return "[" + ", ".join(_toml_scalar(item) for item in value) + "]"
|
||||
if value is None:
|
||||
raise TypeError("TOML has no null value")
|
||||
raise TypeError(f"unsupported TOML scalar: {type(value).__name__}")
|
||||
|
||||
|
||||
def toml_dumps(data: Mapping[str, Any]) -> str:
|
||||
"""Serialize the subset of TOML used by this project deterministically."""
|
||||
|
||||
lines: list[str] = []
|
||||
|
||||
def emit_table(table: Mapping[str, Any], path: tuple[str, ...], header: bool) -> None:
|
||||
scalar_items: list[tuple[str, Any]] = []
|
||||
child_tables: list[tuple[str, Mapping[str, Any]]] = []
|
||||
arrays_of_tables: list[tuple[str, list[Mapping[str, Any]]]] = []
|
||||
for key in sorted(table):
|
||||
value = table[key]
|
||||
if value is None:
|
||||
continue
|
||||
if isinstance(value, Mapping):
|
||||
child_tables.append((key, value))
|
||||
elif (
|
||||
isinstance(value, list)
|
||||
and value
|
||||
and all(isinstance(item, Mapping) for item in value)
|
||||
):
|
||||
arrays_of_tables.append((key, value))
|
||||
else:
|
||||
scalar_items.append((key, value))
|
||||
|
||||
if header:
|
||||
if lines and lines[-1] != "":
|
||||
lines.append("")
|
||||
lines.append("[" + ".".join(toml_quote(part) for part in path) + "]")
|
||||
for key, value in scalar_items:
|
||||
lines.append(f"{toml_quote(key)} = {_toml_scalar(value)}")
|
||||
|
||||
for key, child in child_tables:
|
||||
emit_table(child, path + (key,), True)
|
||||
for key, values in arrays_of_tables:
|
||||
for item in values:
|
||||
if lines and lines[-1] != "":
|
||||
lines.append("")
|
||||
item_path = path + (key,)
|
||||
lines.append("[[" + ".".join(toml_quote(part) for part in item_path) + "]]")
|
||||
item_scalars = {
|
||||
item_key: item_value
|
||||
for item_key, item_value in item.items()
|
||||
if not isinstance(item_value, Mapping)
|
||||
}
|
||||
nested = {
|
||||
item_key: item_value
|
||||
for item_key, item_value in item.items()
|
||||
if isinstance(item_value, Mapping)
|
||||
}
|
||||
for item_key in sorted(item_scalars):
|
||||
lines.append(f"{toml_quote(item_key)} = {_toml_scalar(item_scalars[item_key])}")
|
||||
for nested_key, nested_value in sorted(nested.items()):
|
||||
emit_table(nested_value, item_path + (nested_key,), True)
|
||||
|
||||
emit_table(data, (), False)
|
||||
return "\n".join(lines).rstrip() + "\n"
|
||||
|
||||
|
||||
def deep_merge(base: Mapping[str, Any], override: Mapping[str, Any]) -> dict[str, Any]:
|
||||
result: dict[str, Any] = dict(base)
|
||||
for key, value in override.items():
|
||||
existing = result.get(key)
|
||||
if isinstance(existing, Mapping) and isinstance(value, Mapping):
|
||||
result[key] = deep_merge(existing, value)
|
||||
else:
|
||||
result[key] = value
|
||||
return result
|
||||
|
||||
|
||||
def parse_env_file(path: Path) -> dict[str, str]:
|
||||
values: dict[str, str] = {}
|
||||
if not path.is_file():
|
||||
return values
|
||||
for number, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
|
||||
line = raw.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
if line.startswith("export "):
|
||||
line = line[7:].lstrip()
|
||||
if "=" not in line:
|
||||
raise ValueError(f"invalid environment assignment at {path}:{number}")
|
||||
key, value = line.split("=", 1)
|
||||
key = key.strip()
|
||||
if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", key):
|
||||
raise ValueError(f"invalid environment variable at {path}:{number}: {key!r}")
|
||||
value = value.strip()
|
||||
if len(value) >= 2 and value[0] == value[-1] and value[0] in {'"', "'"}:
|
||||
quote = value[0]
|
||||
value = value[1:-1]
|
||||
if quote == '"':
|
||||
decoded: list[str] = []
|
||||
index = 0
|
||||
escapes = {"n": "\n", "r": "\r", "t": "\t", '"': '"', "\\": "\\"}
|
||||
while index < len(value):
|
||||
if value[index] != "\\" or index + 1 >= len(value):
|
||||
decoded.append(value[index])
|
||||
index += 1
|
||||
continue
|
||||
escaped = value[index + 1]
|
||||
if escaped in escapes:
|
||||
decoded.append(escapes[escaped])
|
||||
else:
|
||||
# Preserve unknown escapes literally. This avoids the
|
||||
# lossy/non-ASCII behavior of ``unicode_escape`` while
|
||||
# remaining compatible with ordinary dotenv values.
|
||||
decoded.extend(("\\", escaped))
|
||||
index += 2
|
||||
value = "".join(decoded)
|
||||
values[key] = value
|
||||
return values
|
||||
|
||||
|
||||
def filtered_environment(
|
||||
*,
|
||||
allow_sensitive: Iterable[str] = (),
|
||||
extra: Mapping[str, str] | None = None,
|
||||
) -> dict[str, str]:
|
||||
sensitive = re.compile(
|
||||
r"(?:^|_)(?:API_?KEY|KEY|TOKEN|SECRET|PASSWORD|PASSWD|CREDENTIAL|COOKIE|AUTH|PRIVATE)(?:_|$)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
allowed = set(allow_sensitive)
|
||||
result: dict[str, str] = {}
|
||||
for key, value in os.environ.items():
|
||||
if key.startswith("MMO_"):
|
||||
continue
|
||||
if sensitive.search(key) and key not in allowed:
|
||||
continue
|
||||
result[key] = value
|
||||
if extra:
|
||||
result.update(extra)
|
||||
return result
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def file_lock(path: Path) -> Iterator[None]:
|
||||
path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||||
with path.open("a+", encoding="utf-8") as handle:
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
||||
|
||||
|
||||
def _proc_stat_fields(value: str) -> list[str]:
|
||||
"""Return Linux ``/proc/<pid>/stat`` fields beginning with process state.
|
||||
|
||||
The parenthesized command name may contain spaces and closing parentheses,
|
||||
so tokenizing the complete record is not safe. The kernel's fields after
|
||||
the final closing parenthesis have a stable whitespace-delimited layout.
|
||||
"""
|
||||
|
||||
closing = value.rfind(")")
|
||||
if closing < 0:
|
||||
return []
|
||||
return value[closing + 1 :].split()
|
||||
|
||||
|
||||
def process_alive(pid: int | None) -> bool:
|
||||
if not pid or pid <= 0:
|
||||
return False
|
||||
# kill(pid, 0) reports unreaped zombies as alive. On Linux this caused
|
||||
# cancellation and gateway shutdown to consume the full grace period even
|
||||
# though the process had already exited. Treat a /proc zombie as terminal.
|
||||
stat = Path(f"/proc/{pid}/stat")
|
||||
with contextlib.suppress(OSError):
|
||||
fields = _proc_stat_fields(stat.read_text(encoding="ascii", errors="replace"))
|
||||
if fields and fields[0] == "Z":
|
||||
return False
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except ProcessLookupError:
|
||||
return False
|
||||
except PermissionError:
|
||||
return True
|
||||
return True
|
||||
|
||||
|
||||
def process_start_token(pid: int | None) -> str | None:
|
||||
"""Return the Linux process start-time token used to detect PID reuse."""
|
||||
|
||||
if not pid or pid <= 0:
|
||||
return None
|
||||
try:
|
||||
value = Path(f"/proc/{pid}/stat").read_text(encoding="ascii", errors="replace")
|
||||
fields = _proc_stat_fields(value)
|
||||
return fields[19]
|
||||
except (OSError, IndexError):
|
||||
return None
|
||||
|
||||
|
||||
def process_matches(pid: int | None, start_token: object = None) -> bool:
|
||||
"""Return whether *pid* is alive and still denotes the recorded process."""
|
||||
|
||||
if start_token is None or not process_alive(pid):
|
||||
return False
|
||||
return process_start_token(pid) == str(start_token)
|
||||
|
||||
|
||||
def process_group_alive(pgid: int | None) -> bool:
|
||||
if not pgid or pgid <= 1:
|
||||
return False
|
||||
proc = Path("/proc")
|
||||
if proc.is_dir():
|
||||
inspected = False
|
||||
for stat_path in proc.glob("[0-9]*/stat"):
|
||||
try:
|
||||
value = stat_path.read_text(encoding="ascii", errors="replace")
|
||||
fields = _proc_stat_fields(value)
|
||||
inspected = True
|
||||
if int(fields[2]) == pgid and fields[0] != "Z":
|
||||
return True
|
||||
except (OSError, ValueError, IndexError):
|
||||
continue
|
||||
if inspected:
|
||||
return False
|
||||
try:
|
||||
os.killpg(pgid, 0)
|
||||
except ProcessLookupError:
|
||||
return False
|
||||
except PermissionError:
|
||||
return True
|
||||
return True
|
||||
|
||||
|
||||
def terminate_process_group(pgid: int, grace_seconds: float = 8.0) -> None:
|
||||
"""Terminate an isolated process group, including after its leader exits."""
|
||||
|
||||
if not process_group_alive(pgid):
|
||||
return
|
||||
with contextlib.suppress(ProcessLookupError, PermissionError):
|
||||
os.killpg(pgid, signal.SIGTERM)
|
||||
deadline = time.monotonic() + max(0.1, grace_seconds)
|
||||
while time.monotonic() < deadline:
|
||||
if not process_group_alive(pgid):
|
||||
return
|
||||
time.sleep(0.05)
|
||||
with contextlib.suppress(ProcessLookupError, PermissionError):
|
||||
os.killpg(pgid, signal.SIGKILL)
|
||||
# Give the kernel a bounded interval to retire the process group. This is
|
||||
# especially important for archive/release tests, which must not leave
|
||||
# detached descendants behind after a forced cancellation.
|
||||
kill_deadline = time.monotonic() + 2.0
|
||||
while time.monotonic() < kill_deadline:
|
||||
if not process_group_alive(pgid):
|
||||
return
|
||||
time.sleep(0.02)
|
||||
|
||||
|
||||
def terminate_process(pid: int, grace_seconds: float = 8.0) -> None:
|
||||
"""Terminate one process when it cannot safely own an isolated group."""
|
||||
|
||||
if not process_alive(pid):
|
||||
return
|
||||
with contextlib.suppress(ProcessLookupError, PermissionError):
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
deadline = time.monotonic() + max(0.1, grace_seconds)
|
||||
while time.monotonic() < deadline:
|
||||
if not process_alive(pid):
|
||||
return
|
||||
time.sleep(0.05)
|
||||
with contextlib.suppress(ProcessLookupError, PermissionError):
|
||||
os.kill(pid, signal.SIGKILL)
|
||||
kill_deadline = time.monotonic() + 2.0
|
||||
while time.monotonic() < kill_deadline:
|
||||
if not process_alive(pid):
|
||||
return
|
||||
time.sleep(0.02)
|
||||
|
||||
|
||||
def is_within(path: Path, root: Path) -> bool:
|
||||
try:
|
||||
path.relative_to(root)
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def resolve_inside(value: str | Path, root: Path, *, must_exist: bool = False) -> Path:
|
||||
candidate = Path(value).expanduser()
|
||||
if not candidate.is_absolute():
|
||||
candidate = root / candidate
|
||||
resolved = candidate.resolve(strict=must_exist)
|
||||
if not is_within(resolved, root):
|
||||
raise ValueError(f"path escapes allowed root {root}: {value}")
|
||||
return resolved
|
||||
|
||||
|
||||
def copy_tree_static(source: Path, destination: Path) -> None:
|
||||
"""Copy a static profile tree while rejecting links and special files."""
|
||||
|
||||
for path in sorted(source.rglob("*")):
|
||||
relative = path.relative_to(source)
|
||||
target = destination / relative
|
||||
if path.is_symlink():
|
||||
raise ValueError(f"symbolic links are not allowed in profile packs: {relative}")
|
||||
if path.is_dir():
|
||||
target.mkdir(parents=True, exist_ok=True, mode=0o755)
|
||||
continue
|
||||
if not path.is_file():
|
||||
raise ValueError(f"special files are not allowed in profile packs: {relative}")
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copyfile(path, target)
|
||||
os.chmod(target, 0o644)
|
||||
|
||||
|
||||
def make_tree_read_only(root: Path) -> None:
|
||||
for path in sorted(root.rglob("*"), reverse=True):
|
||||
if path.is_dir():
|
||||
os.chmod(path, 0o555)
|
||||
elif path.is_file():
|
||||
os.chmod(path, 0o444)
|
||||
os.chmod(root, 0o555)
|
||||
|
||||
|
||||
def http_json(url: str, *, timeout: float = 10.0) -> Any:
|
||||
request = urllib.request.Request(url, headers={"Accept": "application/json"})
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
return strict_json_loads(response.read().decode("utf-8"))
|
||||
|
||||
|
||||
def http_ready(url: str, *, timeout: float = 0.5) -> bool:
|
||||
try:
|
||||
request = urllib.request.Request(url, headers={"Accept": "application/json"})
|
||||
opener = urllib.request.build_opener(_NoRedirectHandler())
|
||||
with opener.open(request, timeout=timeout) as response:
|
||||
return 200 <= response.status < 300
|
||||
except urllib.error.HTTPError as exc:
|
||||
exc.close()
|
||||
return False
|
||||
except (OSError, urllib.error.URLError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def port_available(host: str, port: int) -> bool:
|
||||
family = socket.AF_INET6 if ":" in host else socket.AF_INET
|
||||
with socket.socket(family, socket.SOCK_STREAM) as sock:
|
||||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
try:
|
||||
sock.bind((host, port))
|
||||
except OSError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def append_jsonl(path: Path, value: Mapping[str, Any]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||||
line = json.dumps(value, ensure_ascii=False, separators=(",", ":"), allow_nan=False) + "\n"
|
||||
fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600)
|
||||
try:
|
||||
# O_APPEND makes each individual write atomic, but a short write would
|
||||
# otherwise let another writer splice a record between retry chunks.
|
||||
fcntl.flock(fd, fcntl.LOCK_EX)
|
||||
data = line.encode("utf-8")
|
||||
offset = 0
|
||||
while offset < len(data):
|
||||
written = os.write(fd, data[offset:])
|
||||
if written <= 0:
|
||||
raise OSError("unable to append complete JSONL record")
|
||||
offset += written
|
||||
os.fsync(fd)
|
||||
finally:
|
||||
with contextlib.suppress(OSError):
|
||||
fcntl.flock(fd, fcntl.LOCK_UN)
|
||||
os.close(fd)
|
||||
|
||||
|
||||
def event_usage(path: Path) -> dict[str, int]:
|
||||
"""Extract canonical token usage, preferring cumulative app-server totals."""
|
||||
|
||||
maxima: dict[str, int] = {}
|
||||
app_server_totals: dict[str, int] = {}
|
||||
event_count = 0
|
||||
interesting = {
|
||||
"input_tokens",
|
||||
"output_tokens",
|
||||
"cached_input_tokens",
|
||||
"cache_write_input_tokens",
|
||||
"reasoning_output_tokens",
|
||||
"reasoning_tokens",
|
||||
"total_tokens",
|
||||
}
|
||||
app_server_names = {
|
||||
"inputTokens": "input_tokens",
|
||||
"outputTokens": "output_tokens",
|
||||
"cachedInputTokens": "cached_input_tokens",
|
||||
"cacheWriteInputTokens": "cache_write_input_tokens",
|
||||
"reasoningOutputTokens": "reasoning_output_tokens",
|
||||
"totalTokens": "total_tokens",
|
||||
}
|
||||
|
||||
def visit(value: Any) -> None:
|
||||
if isinstance(value, Mapping):
|
||||
for key, child in value.items():
|
||||
if key in interesting and isinstance(child, int) and not isinstance(child, bool):
|
||||
maxima[key] = max(maxima.get(key, 0), child)
|
||||
visit(child)
|
||||
elif isinstance(value, list):
|
||||
for child in value:
|
||||
visit(child)
|
||||
|
||||
if path.is_file():
|
||||
with path.open("r", encoding="utf-8", errors="replace") as handle:
|
||||
for line in handle:
|
||||
try:
|
||||
value = strict_json_loads(line)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
continue
|
||||
if not isinstance(value, Mapping):
|
||||
continue
|
||||
event_count += 1
|
||||
message = value.get("message")
|
||||
if (
|
||||
isinstance(message, Mapping)
|
||||
and message.get("method") == "thread/tokenUsage/updated"
|
||||
):
|
||||
params = message.get("params")
|
||||
token_usage = params.get("tokenUsage") if isinstance(params, Mapping) else None
|
||||
total = token_usage.get("total") if isinstance(token_usage, Mapping) else None
|
||||
if isinstance(total, Mapping):
|
||||
# App-server reports a cumulative thread total. Notifications
|
||||
# may be repeated during recovery, so summing them would
|
||||
# double-count usage. Keep the greatest observed cumulative
|
||||
# value for each canonical category.
|
||||
for source, target in app_server_names.items():
|
||||
count = total.get(source)
|
||||
if isinstance(count, int) and not isinstance(count, bool):
|
||||
app_server_totals[target] = max(
|
||||
app_server_totals.get(target, 0), count
|
||||
)
|
||||
visit(value)
|
||||
maxima.update(app_server_totals)
|
||||
maxima["event_count"] = event_count
|
||||
return maxima
|
||||
|
||||
|
||||
def bounded_text(text: str, limit: int) -> tuple[str, bool]:
|
||||
limit = max(0, limit)
|
||||
if len(text) <= limit:
|
||||
return text, False
|
||||
marker = "\n\n...[truncated]...\n\n"
|
||||
if limit <= len(marker):
|
||||
return marker[:limit], True
|
||||
payload_limit = limit - len(marker)
|
||||
head = max(1, payload_limit * 2 // 3)
|
||||
tail = payload_limit - head
|
||||
return text[:head] + marker + (text[-tail:] if tail else ""), True
|
||||
|
||||
|
||||
def walk_files(root: Path) -> list[Path]:
|
||||
files: list[Path] = []
|
||||
for path in root.rglob("*"):
|
||||
if path.is_symlink():
|
||||
raise ValueError(f"symbolic link is not allowed in manifest tree: {path}")
|
||||
if path.is_file():
|
||||
files.append(path)
|
||||
elif not path.is_dir():
|
||||
raise ValueError(f"special file is not allowed in manifest tree: {path}")
|
||||
return sorted(files)
|
||||
|
||||
|
||||
def manifest_for_tree(root: Path, *, exclude: Sequence[str] = ()) -> dict[str, str]:
|
||||
excluded = set(exclude)
|
||||
result: dict[str, str] = {}
|
||||
for path in walk_files(root):
|
||||
relative = path.relative_to(root).as_posix()
|
||||
if relative in excluded:
|
||||
continue
|
||||
result[relative] = sha256_file(path)
|
||||
return result
|
||||
Reference in New Issue
Block a user