cleveragents-core/tools/controller/db/payload_guard.py

"""Input-payload size guard for workflow_attempts.

Plan v9: ``input_payload`` is capped at 4MB. Hitting the cap triggers
truncation in priority order:
    1. ``prior_attempts.older_summary`` (recomputed next attempt)
    2. Oldest verbatim entries in ``prior_attempts.verbatim``
    3. Advisory (non-RC) comments in ``pr_comments_since_last_attempt``
    4. Full diff (keep diff_summary only; drop ``full_diff``)
    5. CIFailure.raw_log_excerpt entries (smallest impact gain)
    6. Fail: STUCK with reason="input-too-large"

The truncation is applied IN PLACE on a dict (not a Pydantic model)
because the controller's master_prefetch path assembles the payload
incrementally before validating. After truncation, the caller
re-validates against the V1 contract.
"""

from __future__ import annotations

import json
from typing import Any

MAX_INPUT_PAYLOAD_BYTES = 4 * 1024 * 1024  # 4 MiB


class PayloadTooLargeError(ValueError):
    """Raised when truncation can't get the payload under the cap.

    Master catches this and transitions the workflow to STUCK with
    ``reason="input-too-large"``.
    """

    def __init__(self, final_size: int, attempted: list[str]):
        self.final_size = final_size
        self.attempted = attempted
        super().__init__(
            f"input_payload still {final_size} bytes after truncation "
            f"({', '.join(attempted)}); cap is {MAX_INPUT_PAYLOAD_BYTES}"
        )


def _size_bytes(payload: dict[str, Any]) -> int:
    return len(json.dumps(payload, default=str))


def enforce_input_payload_size(
    payload: dict[str, Any],
    *,
    cap_bytes: int = MAX_INPUT_PAYLOAD_BYTES,
) -> tuple[dict[str, Any], bool, list[str]]:
    """Truncate ``payload`` until it fits ``cap_bytes`` or raise.

    Returns ``(payload, was_truncated, truncations_applied)``.
    Caller should set ``workflow_attempts.input_payload_truncated``
    from the second return value.

    Mutates the input dict for efficiency; caller should pass a
    fresh dict if it needs the original.
    """
    truncations: list[str] = []
    current = _size_bytes(payload)
    if current <= cap_bytes:
        return payload, False, truncations

    # Step 1: drop prior_attempts.older_summary.
    if (
        "prior_attempts" in payload
        and isinstance(payload["prior_attempts"], dict)
        and payload["prior_attempts"].get("older_summary") is not None
    ):
        payload["prior_attempts"]["older_summary"] = None
        truncations.append("prior_attempts.older_summary")
        current = _size_bytes(payload)
        if current <= cap_bytes:
            return payload, True, truncations

    # Step 2: drop oldest verbatim prior_attempts (keep last 1).
    if (
        "prior_attempts" in payload
        and isinstance(payload["prior_attempts"], dict)
        and isinstance(payload["prior_attempts"].get("verbatim"), list)
        and len(payload["prior_attempts"]["verbatim"]) > 1
    ):
        # Drop everything but the most recent.
        kept = payload["prior_attempts"]["verbatim"][-1:]
        dropped = len(payload["prior_attempts"]["verbatim"]) - len(kept)
        payload["prior_attempts"]["verbatim"] = kept
        truncations.append(f"prior_attempts.verbatim (-{dropped})")
        current = _size_bytes(payload)
        if current <= cap_bytes:
            return payload, True, truncations

    # Step 3: drop advisory PR comments.
    if (
        isinstance(payload.get("pr_comments_since_last_attempt"), list)
        and payload["pr_comments_since_last_attempt"]
    ):
        dropped = len(payload["pr_comments_since_last_attempt"])
        payload["pr_comments_since_last_attempt"] = []
        truncations.append(f"pr_comments_since_last_attempt (-{dropped})")
        current = _size_bytes(payload)
        if current <= cap_bytes:
            return payload, True, truncations

    # Step 4: drop full_diff (keep diff_summary).
    if payload.get("full_diff"):
        payload["full_diff"] = None
        truncations.append("full_diff")
        current = _size_bytes(payload)
        if current <= cap_bytes:
            return payload, True, truncations

    # Step 5: shrink CIFailure.raw_log_excerpt across gates.
    if isinstance(payload.get("ci_summary"), dict):
        gates = payload["ci_summary"].get("gates") or []
        for gate in gates:
            if isinstance(gate, dict) and isinstance(gate.get("failure"), dict):
                failure = gate["failure"]
                excerpt = failure.get("raw_log_excerpt", "")
                if isinstance(excerpt, str) and len(excerpt) > 1024:
                    failure["raw_log_excerpt"] = excerpt[-1024:]
                    failure["log_excerpt_lines"] = excerpt[-1024:].count("\n") + 1
        truncations.append("ci_summary.gates.failure.raw_log_excerpt (→1KB)")
        current = _size_bytes(payload)
        if current <= cap_bytes:
            return payload, True, truncations

    # Step 6: failed to fit.
    raise PayloadTooLargeError(final_size=current, attempted=truncations)