#!/usr/bin/env python3 """Generate bounded, auditable task context packets. The selector applies byte and count limits before writing a packet. It never creates an over-budget packet for a later validator to reject. """ from __future__ import annotations import argparse import hashlib import json import sys from dataclasses import dataclass from pathlib import Path from typing import Any, Iterable PROFILES = { "doc": {"tokens": 2048, "bytes": 8192, "evidence_bytes": 2048, "files": 8}, "feature": {"tokens": 4096, "bytes": 16384, "evidence_bytes": 6144, "files": 16}, "contract": {"tokens": 6144, "bytes": 24576, "evidence_bytes": 10240, "files": 24}, "milestone": {"tokens": 8192, "bytes": 32768, "evidence_bytes": 14336, "files": 32}, } SEGMENTS = ("header", "baseline", "scope", "contracts", "evidence", "risks", "next", "forbidden") # Compact codes keep the audit trail inside the packet budget. The human # explanation is stable and can be rendered by the reader without repeating # large strings in every packet. EXCLUSION_REASONS = { "E1": "duplicate_path", "E2": "evidence_item_bytes_exceed_budget", "E3": "evidence_file_count_exceed_budget", "E4": "evidence_bytes_exceed_budget", "E5": "context_remaining_bytes_insufficient", } class PacketError(ValueError): """Raised when a packet cannot be generated without violating a budget.""" @dataclass(frozen=True) class Candidate: path: str size: int sha256: str | None kind: str priority: int required: bool def _json_bytes(value: Any) -> bytes: return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8") 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 _candidate(raw: dict[str, Any], base_dir: Path) -> Candidate: path_value = str(raw.get("path", "")) if not path_value: raise PacketError("candidate is missing path") path = Path(path_value) if "size" in raw: size = int(raw["size"]) else: file_path = path if path.is_absolute() else base_dir / path try: size = file_path.stat().st_size except OSError as exc: raise PacketError(f"cannot stat evidence {path_value}: {exc}") from exc if size < 0: raise PacketError(f"evidence size must be non-negative: {path_value}") sha256 = raw.get("sha256") if sha256 is None and raw.get("hash_file", False): file_path = path if path.is_absolute() else base_dir / path sha256 = _sha256_file(file_path) return Candidate( path=path_value, size=size, sha256=str(sha256) if sha256 is not None else None, kind=str(raw.get("kind", "artifact")), priority=int(raw.get("priority", 0)), required=bool(raw.get("required", False)), ) def _evidence_item(candidate: Candidate) -> dict[str, Any]: item: dict[str, Any] = {"path": candidate.path, "bytes": candidate.size, "kind": candidate.kind} if candidate.sha256: item["sha256"] = candidate.sha256 return item def _packet_bytes(packet: dict[str, Any]) -> int: return len(_json_bytes(packet)) def _context_bytes(packet: dict[str, Any], evidence_used: int) -> int: """Return the bytes consumed by the envelope and referenced evidence. Evidence is referenced by path in the JSON envelope, but the referenced payload is read into the task context by the consumer. Counting both here prevents a packet that is small on disk from expanding past the context limit when it is loaded. """ return _packet_bytes(packet) + evidence_used def _finalize_packet_size(packet: dict[str, Any]) -> int: """Store the self-consistent serialized size in the budget metadata.""" packet["budget"]["packet_bytes"] = 0 for _ in range(4): packet_bytes = _packet_bytes(packet) if packet["budget"]["packet_bytes"] == packet_bytes: return packet_bytes packet["budget"]["packet_bytes"] = packet_bytes raise PacketError("could not stabilize serialized packet size") def _base_packet(input_data: dict[str, Any], context_id: str, packet_id: str) -> dict[str, Any]: packet: dict[str, Any] = { "context_id": context_id, "packet_id": packet_id, "task": str(input_data.get("task", "")), "header": dict(input_data.get("header", {})), "baseline": dict(input_data.get("baseline", {})), "scope": dict(input_data.get("scope", {})), "contracts": list(input_data.get("contracts", [])), "evidence": {"selected": [], "excluded": []}, "risks": list(input_data.get("risks", [])), "next": list(input_data.get("next", []))[:3], "forbidden": list(input_data.get("forbidden", [])), } return packet def _budget_metadata(limits: dict[str, int], selected: list[dict[str, Any]], excluded: list[dict[str, Any]], evidence_used: int) -> dict[str, Any]: return { "profile": limits["profile"], "token_limit": limits["tokens"], "byte_limit": limits["bytes"], "evidence_byte_limit": limits["evidence_bytes"], "evidence_file_limit": limits["files"], "selected_evidence_bytes": evidence_used, "selected_evidence_files": len(selected), "excluded_evidence_files": len(excluded), # Reserve the complete decimal width of the configured limit while # selecting. This keeps the selection check conservative. "packet_bytes": limits["bytes"], } def generate_packet( input_data: dict[str, Any], *, context_id: str, packet_id: str, profile: str = "feature", total_bytes: int | None = None, evidence_bytes: int | None = None, max_files: int | None = None, base_dir: Path = Path("."), ) -> dict[str, Any]: """Select evidence under all budgets and return a packet within its limit.""" if profile not in PROFILES: raise PacketError(f"unknown profile: {profile}") limits = dict(PROFILES[profile]) limits["profile"] = profile if total_bytes is not None: limits["bytes"] = int(total_bytes) if evidence_bytes is not None: limits["evidence_bytes"] = int(evidence_bytes) if max_files is not None: limits["files"] = int(max_files) if any(limits[key] < 0 for key in ("tokens", "bytes", "evidence_bytes", "files")): raise PacketError("budgets must be non-negative") raw_candidates = input_data.get("candidates", []) if not isinstance(raw_candidates, list): raise PacketError("candidates must be a list") candidates = [] for raw in raw_candidates: if not isinstance(raw, dict): raise PacketError("each candidate must be an object") candidates.append(_candidate(raw, base_dir)) candidates.sort(key=lambda item: (-int(item.required), -item.priority, item.path)) packet = _base_packet(input_data, context_id, packet_id) skeleton_bytes = _packet_bytes(packet) if skeleton_bytes > limits["bytes"]: raise PacketError( f"context skeleton exceeds byte budget ({skeleton_bytes}>{limits['bytes']}); " "reduce task metadata or use a larger profile" ) selected: list[dict[str, Any]] = [] excluded: list[dict[str, Any]] = [] evidence_used = 0 seen: set[str] = set() for candidate in candidates: if candidate.path in seen: excluded.append({"path": candidate.path, "reason_code": "E1"}) continue seen.add(candidate.path) item = _evidence_item(candidate) reason: str | None = None if candidate.size > limits["evidence_bytes"]: reason = "E2" elif len(selected) >= limits["files"]: reason = "E3" elif evidence_used + candidate.size > limits["evidence_bytes"]: reason = "E4" else: trial = dict(packet) trial_selected = selected + [item] trial["evidence"] = {"selected": trial_selected, "excluded": excluded} trial["budget"] = _budget_metadata( limits, trial_selected, excluded, evidence_used + candidate.size ) # This check is intentionally performed before selection is committed. trial_context_bytes = _context_bytes(trial, evidence_used + candidate.size) if trial_context_bytes > limits["bytes"]: reason = "E5" if reason: excluded.append({"path": candidate.path, "reason_code": reason}) if candidate.required: raise PacketError(f"required evidence cannot fit ({reason}): {candidate.path}") continue selected.append(item) evidence_used += candidate.size # The exclusion audit trail is part of the packet, so it can consume the # remaining context space. Remove optional selections until the complete # packet (selected evidence, exclusions and budget metadata) fits. while True: packet["evidence"] = {"selected": selected, "excluded": excluded} packet["budget"] = _budget_metadata(limits, selected, excluded, evidence_used) packet_bytes = _finalize_packet_size(packet) context_bytes = packet_bytes + evidence_used if packet_bytes <= limits["bytes"] and context_bytes <= limits["bytes"]: break removable_index = next( (index for index in range(len(selected) - 1, -1, -1) if not any(item["path"] == selected[index]["path"] and candidate.required for candidate in candidates)), None, ) if removable_index is None: raise PacketError("context skeleton and required evidence exceed byte budget") removed = selected.pop(removable_index) evidence_used -= int(removed["bytes"]) excluded.append({ "path": removed["path"], "reason_code": "E5", }) return packet def _load(path: Path) -> dict[str, Any]: with path.open("r", encoding="utf-8") as handle: value = json.load(handle) if not isinstance(value, dict): raise PacketError("input must be a JSON object") return value def main(argv: Iterable[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--input", required=True, type=Path, help="candidate/context JSON") parser.add_argument("--output", required=True, type=Path) parser.add_argument("--context-id", required=True) parser.add_argument("--packet-id", required=True) parser.add_argument("--profile", choices=sorted(PROFILES), default="feature") parser.add_argument("--base-dir", type=Path, default=Path(".")) parser.add_argument("--total-bytes", type=int) parser.add_argument("--evidence-bytes", type=int) parser.add_argument("--max-files", type=int) args = parser.parse_args(list(argv) if argv is not None else None) try: packet = generate_packet( _load(args.input), context_id=args.context_id, packet_id=args.packet_id, profile=args.profile, total_bytes=args.total_bytes, evidence_bytes=args.evidence_bytes, max_files=args.max_files, base_dir=args.base_dir, ) args.output.parent.mkdir(parents=True, exist_ok=True) payload = _json_bytes(packet) # Do not expose a partially written packet to a concurrent validator. temporary = args.output.with_name(f".{args.output.name}.tmp") temporary.write_bytes(payload) if len(payload) != packet["budget"]["packet_bytes"]: raise PacketError("serialized packet size changed before write") temporary.replace(args.output) except (PacketError, OSError, json.JSONDecodeError) as exc: print(f"context packet generation failed: {exc}", file=sys.stderr) return 2 return 0 if __name__ == "__main__": raise SystemExit(main())