atlas-iac/services/hermes/scripts/suite_multipass.py

"""Bounded whole-suite orchestration with one shared deadline and cost allowance."""
from __future__ import annotations

import hashlib
import math
import time

import suite_backends
from suite_assignments import decode
from suite_contract import (EXECUTION_REVISION, MAX_BODY, MAX_RESULT, MODELS, PROMPT_REVISION,
                            PROMPT_SHA256, REVISION, Problem, digest, encoded, reasoning_selection,
                            validate_partition)
from suite_policy import (BASE_NAME_LIMIT, MAX_GROUP, POLICY_REVISION, invocation,
                          validate_natural)
from suite_sizing import cap_families, review_summary
from suite_review_batches import review_stage
from suite_recovery import Checkpoints, MAX_INVOCATIONS, allocate, recover
from suite_hierarchical import MAX_LOGICAL_CALLS, FALLBACK_ERRORS, execute as hierarchical, selection
from collections import Counter

MAX_PASSES = 5
MAX_MODEL_CALLS = MAX_INVOCATIONS
CAPACITY_REVISION = "suite-context-turns-v1-20260930"


def ordered_request(request, alternative=False):
    """Use stable independent orders while retaining every complete case record."""
    def key(case):
        alias = case["alias"]
        return hashlib.sha256((POLICY_REVISION + ":order-b:" + alias).encode()).hexdigest() if alternative else alias
    cases = sorted(request["cases"], key=key)
    if alternative and len(cases) > 1 and [c["alias"] for c in cases] == sorted(c["alias"] for c in cases):
        cases.reverse()
    return {**request, "cases": cases}


def capacity(call, provider, count, family_count=None):
    """Check the actual complete pass, including proposal/review input and schema."""
    model = MODELS[provider]
    input_bytes = len(call["input"].encode()) + len(call["system"].encode()) + len(encoded(call["schema"]))
    if input_bytes > MAX_BODY:
        raise Problem("pass_request_too_large", 413, review_pass=call["stage"], input_bytes=input_bytes)
    if call["stage"] == "implementation_profiles":
        estimate = 1024 + 1400 * count
    elif family_count is None:
        estimate = 1024 + 128 * count
    else:
        estimate = 1024 + 48 * count + 384 * family_count
    reserve = estimate + (8192 if provider == "claude" else 0)
    bound = input_bytes + model["overhead"]
    ceiling, minimum = model.get("max_turns", 1), model.get("min_turns", 1)
    # Reserve a full output for every possible CLI turn, while retaining bounded
    # repair headroom. Do not reject a complete input merely to reserve unused turns.
    turns = min(ceiling, max(0, (model["context"] - bound) // model["output"]))
    details = {"review_pass": call["stage"], "input_bytes": input_bytes,
               "input_token_bound": bound, "context_limit": model["context"],
               "output_reservation_tokens": reserve, "max_output_tokens": model["output"],
               "configured_max_turns": ceiling, "minimum_max_turns": minimum,
               "available_max_turns": turns, "capacity_revision": CAPACITY_REVISION}
    if reserve > model["output"] or turns < minimum:
        raise Problem("pass_capacity", 422, **details, capacity_reason=
                      "output_reservation" if reserve > model["output"] else "context_reservation")
    return {"input_bytes": input_bytes, "input_token_bound": bound, "input_token_count": None,
            "max_turns": turns, "configured_max_turns": ceiling, "minimum_max_turns": minimum,
            "context_reserved_tokens": bound + model["output"] * turns,
            "context_headroom_tokens": model["context"] - bound - model["output"] * turns,
            "capacity_revision": CAPACITY_REVISION,
            "output_reservation_tokens": reserve, "output_reservation_verified": False,
            "input_count_method": "Complete UTF-8 input/system/schema byte bound plus harness overhead; not a tokenizer"}


def preflight_workflow(request):
    """Choose once, respecting permissions; recheck actual expanded inputs at every pass."""
    count = len(request["cases"])
    candidates = ["local"]
    if request["routing"]["allow_external"]:
        candidates += request["routing"]["allowed_external_providers"]
    reasons = {}
    for provider in candidates:
        model = MODELS[provider]
        if not model["enabled"]:
            reasons[provider] = "unverified_output_capacity"
            continue
        strategy = selection(request) if provider == "claude" else {"execution_mode": "direct"}
        try:
            if strategy["execution_mode"] == "direct":
                try:
                    initial = capacity(invocation("proposal_a", request), provider, count)
                    capacity(invocation("reconciliation", request), provider, count, 1)
                except Problem:
                    if provider != "claude":
                        raise
                    strategy.update(execution_mode="hierarchical", strategy_reasons=["direct_capacity"])
            if strategy["execution_mode"] == "hierarchical":
                from suite_profiles import batches, profile_call
                pieces = batches(request)
                checked = [capacity(profile_call(piece, i), provider, len(piece["cases"]))
                           for i, piece in enumerate(pieces, 1)]
                initial = checked[0]
                strategy["profile_batch_count"] = len(pieces)
        except Problem as exc:
            reasons[provider] = exc.code
            continue
        return {"provider": provider, **model, **initial, **strategy, **reasoning_selection(request, provider),
                "configuration_revision": REVISION,
                "prompt_revision": PROMPT_REVISION, "prompt_sha256": PROMPT_SHA256,
                "execution_revision": EXECUTION_REVISION, "policy_revision": POLICY_REVISION,
                "case_count": count, "source_sha256": digest(request["cases"]),
                "minimum_model_passes": 3, "maximum_model_passes": MAX_MODEL_CALLS,
                "logical_stages": MAX_PASSES, "automatic_review_batching": True,
                "max_final_group_cases": MAX_GROUP, "max_final_name_characters": 64,
                "cost_guard_enforced": False, "maximum_attempts_per_pass": 3,
                "checkpoint_storage": "active_job_memory", "checkpoint_restart_recovery": False,
                "later_pass_capacity_verified": False, "later_pass_checks": "before_each_invocation"}
    raise Problem("capacity_or_unsupported_backend", 422, candidates=reasons)


def sum_usage(records):
    """Aggregate measured counts only; keep an unknown counter unknown."""
    values = [r.get("usage") for r in records]
    if not values or any(not isinstance(v, dict) for v in values):
        return None
    keys = set().union(*(v.keys() for v in values))
    result = {}
    for key in keys:
        counts = [v.get(key) for v in values]
        if all(type(n) in (int, float) for n in counts):
            result[key] = sum(counts)
        elif all(isinstance(n, dict) for n in counts):
            result[key] = sum_usage([{"usage": n} for n in counts])
        else:
            result[key] = None
    return result


class Workflow:
    """One pinned provider, shared deadline, isolated attempts and scoped checkpoints."""

    def __init__(self, request, selected, cancel, client_ip, progress=None, *, started=None,
                 deadline=None, credential_scope=None):
        self.request, self.provider, self.cancel = request, selected["provider"], cancel
        self.reasoning = selected.get("reasoning", MODELS[self.provider]["reasoning"])
        self.client_ip, self.progress = client_ip, progress or (lambda _: None)
        self.started = time.monotonic() if started is None else started
        self.deadline = self.started + request["execution"]["max_seconds"] if deadline is None else deadline
        self.spent, self.cost_known = 0.0, True
        self.records, self.attempts, self.last_metadata = [], [], {}
        self.stage, self.max_calls = "proposal_a", MAX_LOGICAL_CALLS
        self.execution_mode = selected.get("execution_mode", "direct")
        self.strategy_events = [{"mode": self.execution_mode, "reason": "preflight"}]
        self.retries = Counter()
        self.profile_batch_count = self.cross_batch_review_count = self.source_review_batch_count = 0
        self.future_seconds = 0
        self.global_floor = min(500, max(10, request["execution"]["max_seconds"] / 8))
        self.checkpoints = Checkpoints(request, self.provider, credential_scope, self.deadline)

    def checkpoint(self):
        """Cancellation and time are hard limits; cost is informational only."""
        if self.cancel.is_set():
            raise Problem("cancelled", 409)
        if time.monotonic() >= self.deadline:
            raise Problem("job_time_budget_exhausted", 504)
        if len(self.attempts) >= MAX_INVOCATIONS:
            raise Problem("model_pass_limit", 502)

    def safe_state(self):
        """Never place cached outputs, profiles or rationales into status metadata."""
        return {"execution_mode": self.execution_mode, "strategy_events": self.strategy_events,
                "direct_pass_attempts": sum(a["execution_mode"] == "direct" for a in self.attempts),
                "retry_count": dict(self.retries), "checkpointed_passes": self.checkpoints.names(),
                "checkpoint_reused": list(self.checkpoints.reused),
                "profile_batch_count": self.profile_batch_count,
                "source_review_batch_count": self.source_review_batch_count,
                "cross_batch_review_count": self.cross_batch_review_count,
                "model_pass_count": len(self.attempts), "cost_guard_enforced": False,
                "cost_used_usd_estimate": round(self.spent, 8) if self.cost_known else None,
                "cost_limit_usd_estimate": None,
                "remaining_seconds": round(max(0, self.deadline-time.monotonic()), 3)}

    def call(self, stage, source, context=None, family_count=None, *, custom_call=None, validator=None):
        """Recover one logical pass without regenerating validated predecessors."""
        self.stage = stage
        self.checkpoint()
        if len(self.records) >= self.max_calls:
            raise Problem("model_pass_limit", 502)
        call = custom_call or invocation(stage, source, context)
        call["reasoning"] = self.reasoning
        if context and context.get("review_batch"):
            call["review_batch"] = context["review_batch"]
        started = time.monotonic()
        def report(activity=None):
            now = time.monotonic()
            self.progress({"current_pass": stage, "completed_model_passes": len(self.records),
                           "maximum_model_passes": MAX_INVOCATIONS, "passes": list(self.records),
                           "attempts": list(self.attempts), **self.safe_state(),
                           "review_batch": (context or {}).get("review_batch"),
                           "profile_batch_index": call.get("batch_index"),
                           "heartbeat_at": time.time(), "pass_elapsed_seconds": round(now-started, 1),
                           "job_elapsed_seconds": round(now-self.started, 1),
                           "job_remaining_seconds": round(max(0, self.deadline-now), 1),
                           **(activity or {})})
        report({"cli_running": False, "pass_state": "capacity_preflight"})
        limits = capacity(call, self.provider, len(source["cases"]), family_count)
        def validate(value):
            if validator:
                return validator(value)
            if stage in {"proposal_a", "proposal_b", "profile_proposal_a", "profile_proposal_b", "profile_reconciliation"}:
                return validate_partition(value, source, name_limit=BASE_NAME_LIMIT, unique_names=False)
            originals = None
            if stage in {"large_family_review", "decision_audit"}:
                originals = context.get("original_partition", context.get("natural_partition"))["groups"]
            return validate_natural(value, source, originals)
        before = len(self.records)
        value = recover(self, call, source, limits, validate, report)
        if len(self.records) == before and self.attempts:
            # Only a successfully validated output becomes a completed model pass.
            successful = next((a for a in reversed(self.attempts) if a["stage"] == stage and a["status"] == "generated"), None)
            if successful:
                successful["status"] = "validated"
                self.records.append(dict(successful))
        report({"cli_running": False, "pass_state": "validated"})
        return value

    def invoke(self, call, source, limits, attempt, report):
        """Account for successful and failed attempts before returning any output."""
        self.checkpoint()
        seconds = allocate(self)
        effective = {**source, "execution": {**source["execution"], "max_seconds": seconds,
                                            "max_cost_usd": None}}
        started = time.monotonic()
        record = {"stage": call["stage"], "attempt": attempt, "execution_mode": self.execution_mode,
                  "provider": self.provider, "model": MODELS[self.provider]["model"],
                  "reasoning": self.reasoning, **limits, "case_count": len(source["cases"]),
                  "review_batch": call.get("review_batch"), "profile_batch_index": call.get("batch_index"),
                  "allocated_seconds": seconds, "reserved_future_seconds": self.future_seconds,
                  "allocated_cost_usd": None, "system_sha256": hashlib.sha256(call["system"].encode()).hexdigest(),
                  "schema_sha256": digest(call["schema"]),
                  "case_order_sha256": digest([c["alias"] for c in source["cases"]])}
        metadata, error = {}, None
        try:
            if self.provider == "claude":
                value, metadata = suite_backends.claude_generate(effective, self.cancel, invocation=call,
                    progress=report, job_deadline=self.deadline)
            elif self.provider == "local":
                value, metadata = suite_backends.local_generate(effective, self.cancel, self.client_ip, invocation=call)
            else:
                raise Problem("unsupported_backend", 422)
            record["status"] = "generated"
        except Problem as exc:
            error = exc
            metadata = {**exc.details, "cli_diagnostics": exc.details}
            record.update(status="failed", error_code=exc.code)
        finally:
            cost = metadata.get("cost_usd_estimate") if self.provider == "claude" else 0.0
            if type(cost) in (float, int) and math.isfinite(cost) and cost >= 0:
                self.spent += cost
            else:
                self.cost_known = False
                cost = None
            record.update(wall_seconds=round(time.monotonic()-started, 3), usage=metadata.get("usage"),
                          cli_turns=metadata.get("turns"), duration_api_ms=metadata.get("duration_api_ms"),
                          cost_usd_estimate=cost, cli_diagnostics=metadata.get("cli_diagnostics"))
            self.attempts.append(record)
            self.last_metadata = metadata
        if error:
            raise error
        self.checkpoint()
        try:
            return value if call["stage"] == "implementation_profiles" else decode(value, call, source)
        except Problem as exc:
            record.update(status="failed", error_code=exc.code)
            raise

    def direct(self, source):
        """Reserve a viable time floor for every subsequent required stage."""
        self.future_seconds = 4*self.global_floor
        a = self.call("proposal_a", source)
        self.future_seconds = 3*self.global_floor
        b = self.call("proposal_b", ordered_request(source, True))
        self.future_seconds = 2*self.global_floor
        c = self.call("reconciliation", source, {"proposal_a": a, "proposal_b": b},
                      max(len(a["groups"]), len(b["groups"])))
        return a, b, c

    def execute(self):
        """Choose direct or hierarchical analysis, then review and size once."""
        source = ordered_request(self.request)
        if self.execution_mode == "direct":
            try:
                a, b, reconciled = self.direct(source)
            except Problem as exc:
                if self.provider != "claude" or exc.code not in FALLBACK_ERRORS:
                    raise
                self.strategy_events.append({"mode": "hierarchical", "reason": exc.code, "after_pass": self.stage})
                a, b, reconciled = hierarchical(self, source, capacity)
        else:
            a, b, reconciled = hierarchical(self, source, capacity)
        large = [g for g in reconciled["groups"] if len(g["members"]) > MAX_GROUP]
        reviewed = audited = None
        if large:
            self.future_seconds = self.global_floor
            reviewed = review_stage(self, "large_family_review", source, reconciled, None, capacity)
            self.future_seconds = 0
            audited = review_stage(self, "decision_audit", source, reconciled, reviewed, capacity)
        natural = audited or reconciled
        final, divisions = cap_families(natural, source)
        review = review_summary(a, b, reconciled, reviewed, audited, final, divisions)
        self.checkpoint()
        metadata = {**self.last_metadata, **self.safe_state(), "passes": self.records, "attempts": self.attempts,
                    "reasoning": self.reasoning, "usage": sum_usage(self.attempts),
                    "cli_diagnostics_scope": "last_model_attempt", "cost_usd_estimate": self.safe_state()["cost_used_usd_estimate"],
                    "turns": sum(a["cli_turns"] for a in self.attempts) if all(type(a["cli_turns"]) is int for a in self.attempts) else None,
                    "natural_family_count": len(natural["groups"]), "final_task_count": len(final["groups"]),
                    "singleton_statistics": {k:v for k,v in review["counts"].items() if "singleton" in k},
                    "policy_revision": POLICY_REVISION, "review_summary": review}
        if len(encoded({"result": final, **metadata})) > MAX_RESULT - 16384:
            raise Problem("response_too_large", 502)
        return final, metadata


def generate(request, selected, cancel, client_ip, progress=None, *, started=None, deadline=None, credential_scope=None):
    """No partial completion, provider fallback, or persistent source-derived cache."""
    workflow = Workflow(request, selected, cancel, client_ip, progress, started=started,
                        deadline=deadline, credential_scope=credential_scope)
    try:
        return workflow.execute()
    except Problem as exc:
        details = {"failure_stage": "multi_pass_orchestration", **exc.details, **workflow.safe_state(),
                   "review_pass": workflow.stage, "completed_model_passes": len(workflow.records),
                   "passes": workflow.records, "attempts_metadata": workflow.attempts,
                   "aggregate_usage": sum_usage(workflow.attempts),
                   "aggregate_cost_usd_estimate": workflow.safe_state()["cost_used_usd_estimate"]}
        raise Problem(exc.code, exc.status, **details) from None
    finally:
        workflow.checkpoints.clear()