band-aider-pub/band_aider/judgement/stabilization_decision.py

# SPDX-FileCopyrightText: 2026 Art Sea LLC
#
# SPDX-License-Identifier: AGPL-3.0-only

"""Shared base for request-scoped stabilization judgements."""

from __future__ import annotations

from band_aider.judgement.work_definition_constants import ResponseValidation

from abc import abstractmethod
from dataclasses import replace
import json
import logging
from typing import Callable, cast

from band_aider.ai_client.completion_client import CompletionClient
from band_aider.ai_client.stabilization_escalation import StabilizationEscalation
from band_aider.ai_client import llm_logging
from band_aider.context.provenance import FutureGuidance
from band_aider.context.validation_utils import (
    SemanticValidationError,
    ValidationError,
    ValidationLocation,
    ValidationResult,
    validate_shape,
)
from band_aider.judgement.edit_request import AiderEditRequest
from band_aider.judgement.contextual_cot_renderer import ContextualCotRenderer
from band_aider.judgement.edit_request import AiderEditRequestMask, AiderFailureMode
from band_aider.judgement.file_manager import FileManager
from band_aider.judgement.vibe_check import VibeCodeCheck


class StabilizationDecision(VibeCodeCheck):
    """Carry the exact request whose local stabilization evidence is being judged.

    ``_client`` is the normal runtime client for rendering and helper construction.
    Judgment calls use ``escalation`` and its dumb clients capped at TokenJacker.
    """

    def __init__(
        self,
        *,
        client: CompletionClient,
        file_manager: FileManager,
        request: AiderEditRequest,
        current_cot: str | None = None,
        logger: logging.Logger | None = None,
        prompt: str = "",
        cot_projection: Callable[[str], str] | None = None,
        current_cot_in_prompt: bool = False,
        disambiguation: str | None = None,
    ) -> None:
        """Initialize one decision without creating or transforming workflow evidence."""

        self.escalation = StabilizationEscalation.for_stabilization(
            client, check_name=type(self).__name__,
            prompt_formatter_builder=file_manager.prompt_formatter_builder,
        )
        super().__init__(
            client=client,
            file_manager=file_manager,
            prompt=prompt,
        )
        self.request = request
        self.escalation.files = tuple(map(file_manager.repo_location, request.editable_files))
        self.disambiguation = disambiguation
        self.logger = logger or llm_logging.current_logger() or logging.getLogger(__name__)
        self._future_guidance_composed = False
        self._rendered_prompt: tuple[str, str] | None = None
        # Keep the caller's current authority when prior records need contextual selection.
        resolved_current_cot = (
            current_cot
            if current_cot is not None
            else (
                file_manager.prompt
                if request.previous_cot_records
                else ""
            )
        )
        self._contextual_cot_renderer = ContextualCotRenderer(
            client=client,
            current_cot=cot_projection(resolved_current_cot) if cot_projection else resolved_current_cot,
            task_projection=cot_projection,
            current_cot_in_prompt=current_cot_in_prompt,
            prompt_formatter_builder=file_manager.prompt_formatter_builder,
            logger=self.logger.getChild(
                "contextual-cot-renderer"
            ),
        )

    def render_prompt_for_turn(self, prompt: str) -> str:
        """Render prior COT authority for one exact decision attempt."""

        base_prompt = self.prompt
        work_logs = self.request.work_logs
        # Constructors that already render backlog guidance must not duplicate retained history.
        if work_logs and work_logs.strip() not in base_prompt:
            base_prompt += "\n\n" + self._prompt_formatter_builder.generate("BACKLOG GUIDANCE", escape=True)(work_logs)
        # Reuse authority while the input is stationary; refreshed source must remain visible.
        if self._rendered_prompt is None or self._rendered_prompt[0] != base_prompt:
            rendered = self._contextual_cot_renderer.render(
                AiderEditRequestMask(self.request, AiderFailureMode.SEMANTIC), base_prompt,
            )
            self._rendered_prompt = (base_prompt, rendered)
            self.escalation.ask = rendered
        original = self._rendered_prompt[1] + prompt.removeprefix(self.prompt)
        # Localize the completed invocation, including this attempt's immediate correction.
        # Import here because localization reuses this parent's validation loop without rendering recursively.
        from band_aider.judgement.prompt_localization import PromptLocalization

        return PromptLocalization(client=self._client, file_manager=self._file_manager,
            request=self.request, owner=type(self).__name__, logger=self.logger,
            current_cot=self._contextual_cot_renderer.current_cot or self._file_manager.prompt,
        ).localize(original)

    @property
    @abstractmethod
    def valid_response_references(self) -> list[object]:
        """Return every JSON reference that this decision accepts as structurally valid."""

    @abstractmethod
    def _decide(
        self,
        response: object | None = None,
        validation: ValidationResult | None = None,
    ) -> tuple[object | None, ValidationResult]:
        """Run one response attempt from the preceding response and validation state."""

    def _format_prompt_w_feedback(
        self,
        validation: ValidationResult,
        *,
        base_prompt: str | None = None,
    ) -> str:
        """Render one base prompt with only the latest response correction."""

        section = self._prompt_formatter_builder.generate
        retry_guidance = section(ResponseValidation.HEADER)(lambda: (
            ResponseValidation.RETRY
            + section("VALIDATION ERRORS", escape=True)(validation.stringify())
            + "\n\n"
            + section("REJECTED PREVIOUS RESPONSE")(self.escalation.last_response)
            + "\n\n"
            + ResponseValidation.RETURN_MOVE
        ))
        return "\n\n".join((base_prompt if base_prompt is not None else self.prompt, retry_guidance))

    def __call__(self, *args, **kwargs) -> bool:
        """Iterate child-owned response attempts until validation is clean."""

        # Localize each revived observation against this concrete decision before it decides.
        if not self._future_guidance_composed:
            localize_future_guidance(self)
            self._future_guidance_composed = True
        response: object | None = None
        validation: ValidationResult | None = None
        # Continue the reduction until one child attempt records an admitted decision.
        while validation is None or not validation.valid:
            self.escalation.validation_errors = validation.stringify() if validation else ""
            try:
                response, validation = self._decide(response, validation)
            except json.JSONDecodeError as error:
                response = error.doc
                validation = ValidationResult(errors=(ValidationError(ValidationLocation([]), str(error)),))
        return self.changes[-1][1]

    def last_result(self) -> dict[str, object]:
        """Return the admitted local result for an explicit boundary decision."""

        return {
            "decision": type(self).__name__,
            "outcome": self.outcome,
            "reason": self.reason,
            "additional": self.additional,
        }


class FutureGuidanceApplicability(StabilizationDecision):
    """Decide whether one future observation belongs to one decision input.

    Example:

        applies = FutureGuidanceApplicability(... )()
    """

    _VALID_RESPONSE_REFERENCES: list[object] = [{
        "applicable": True,
        "reason": "why this future observation changes this decision",
    }]

    @property
    def valid_response_references(self) -> list[object]:
        """Return the exact response shape for one applicability judgement."""

        return self._VALID_RESPONSE_REFERENCES

    def __init__(
        self,
        *,
        client: CompletionClient,
        file_manager: FileManager,
        request: AiderEditRequest,
        future_guidance: FutureGuidance,
        current_prompt: str,
        logger: logging.Logger,
        disambiguation: str | None = None,
    ) -> None:
        """Construct one future-guidance applicability decision."""

        super().__init__(
            client=client,
            file_manager=file_manager,
            request=request,
            logger=logger,
            disambiguation=disambiguation,
        )
        section = self._prompt_formatter_builder.generate
        self.prompt = section("FUTURE GUIDANCE APPLICABILITY")(lambda: "\n\n".join((
            "Decide whether the later observation applies to this exact current decision point. "
            "It applies only when changing this decision within its existing authority can avoid "
            "the described dead end. Do not use it to select a different phase, file, or owner.",
            section(
                "CURRENT DECISION PROMPT",
                escape=True,
                neutralize_json=True,
            )(current_prompt),
            section("FUTURE GUIDANCE")(lambda: "\n\n".join((
                section("MESSAGE", escape=True)(future_guidance.message),
                section("JUSTIFICATION", escape=True)(future_guidance.justification),
            ))),
            "Return strict JSON only. applicable must be true or false and reason must be a "
            "nonempty string:\n"
            + json.dumps(self.valid_response_references[0], separators=(",", ":")),
        )))

    def _decide(
        self,
        response: object | None = None,
        validation: ValidationResult | None = None,
    ) -> tuple[object | None, ValidationResult]:
        """Request and admit one applicability response."""

        prompt = self.prompt if validation is None else self._format_prompt_w_feedback(validation)
        root = ValidationLocation([])
        try:
            data = json.loads(self.escalation.complete(
                self.render_prompt_for_turn(prompt),
                disambiguation=self.disambiguation,
            ).strip())
        except json.JSONDecodeError as error:
            return error.doc, ValidationResult(errors=(ValidationError(
                location=root,
                error=str(error),
            ),))
        validation = validate_shape(data, self.valid_response_references, location=root)
        # Require a reason for both admission outcomes so the phase log remains diagnostic.
        if not validation.valid:
            return data, validation
        result = cast("dict[str, object]", data)
        reason = cast(str, result["reason"])
        # Reject a blank explanation because applicability must be auditable in the phase log.
        if not reason.strip():
            return data, ValidationResult(errors=(SemanticValidationError(
                location=root.child("reason"),
                error="must be a nonempty string.",
            ),))
        applicable = cast(bool, result["applicable"])
        self.outcome = "applicable" if applicable else "not_applicable"
        self.reason = reason
        self.additional = {"applicable": applicable}
        self.changes.append((self, applicable))
        return data, ValidationResult()


class FutureGuidanceRelocalization(StabilizationDecision):
    """Derive one current-decision correction from an applicable future observation.

    Example:

        FutureGuidanceRelocalization(... )()
    """

    _VALID_RESPONSE_REFERENCES: list[object] = [{
        "localized_prompt": "compact instruction for the current decision",
    }]

    @property
    def valid_response_references(self) -> list[object]:
        """Return the exact response shape for one localized instruction."""

        return self._VALID_RESPONSE_REFERENCES

    def __init__(
        self,
        *,
        client: CompletionClient,
        file_manager: FileManager,
        request: AiderEditRequest,
        future_guidance: FutureGuidance,
        current_prompt: str,
        logger: logging.Logger,
        disambiguation: str | None = None,
    ) -> None:
        """Construct one input-local future-guidance correction."""

        super().__init__(
            client=client,
            file_manager=file_manager,
            request=request,
            logger=logger,
            disambiguation=disambiguation,
        )
        section = self._prompt_formatter_builder.generate
        self.prompt = section("FUTURE GUIDANCE RELOCALIZATION")(lambda: "\n\n".join((
            "Rewrite one compact backlog instruction for this exact decision so it can avoid "
            "the described dead end while preserving the decision's current authority. "
            "Do not describe work for another phase, file, or owner.",
            section(
                "CURRENT DECISION PROMPT",
                escape=True,
                neutralize_json=True,
            )(current_prompt),
            section("APPLICABLE FUTURE GUIDANCE")(lambda: "\n\n".join((
                section("MESSAGE", escape=True)(future_guidance.message),
                section("JUSTIFICATION", escape=True)(future_guidance.justification),
            ))),
            "Return strict JSON only with this shape:\n"
            + json.dumps(self.valid_response_references[0], separators=(",", ":")),
        )))

    def _decide(
        self,
        response: object | None = None,
        validation: ValidationResult | None = None,
    ) -> tuple[object | None, ValidationResult]:
        """Request and admit one prompt-local correction."""

        prompt = self.prompt if validation is None else self._format_prompt_w_feedback(validation)
        root = ValidationLocation([])
        try:
            data = json.loads(self.escalation.complete(
                self.render_prompt_for_turn(prompt),
                disambiguation=self.disambiguation,
            ).strip())
        except json.JSONDecodeError as error:
            return error.doc, ValidationResult(errors=(ValidationError(
                location=root,
                error=str(error),
            ),))
        validation = validate_shape(data, self.valid_response_references, location=root)
        # Refuse an empty correction because it cannot alter the current backlog guidance.
        if not validation.valid:
            return data, validation
        result = cast("dict[str, object]", data)
        localized_prompt = cast(str, result["localized_prompt"])
        # Reject a blank instruction because localization must change the current decision request.
        if not localized_prompt.strip():
            return data, ValidationResult(errors=(SemanticValidationError(
                location=root.child("localized_prompt"),
                error="must be a nonempty string.",
            ),))
        self.outcome = "localized"
        self.reason = "Future guidance was localized to the current decision prompt."
        self.additional = {"localized_prompt": localized_prompt}
        self.changes.append((self, True))
        return data, ValidationResult()


def localize_future_guidance(decision: StabilizationDecision) -> None:
    """Apply only applicable future observations to one in-memory decision request."""

    logger = decision.logger or logging.getLogger(__name__)
    # Evaluate each later observation in provenance order against the prompt it can refine.
    for future_guidance in decision.request.future_guidance:
        helper_request = replace(
            decision.request,
            backlog_guidance=list(decision.request.backlog_guidance),
            future_guidance=(),
        )
        applicability = FutureGuidanceApplicability(
            client=decision._client,
            file_manager=decision._file_manager,
            request=helper_request,
            future_guidance=future_guidance,
            current_prompt=decision.prompt,
            logger=logger,
            disambiguation=decision.disambiguation,
        )
        applicable = applicability()
        logger.info(
            "Future guidance applicability decision=%s origin=%s addressed=%s applicable=%s reason=%s",
            type(decision).__name__,
            future_guidance.origin_message_id,
            future_guidance.addressed_message_id,
            applicable,
            applicability.reason,
        )
        # Change current guidance only after the model admits that this decision owns it.
        if applicable:
            relocalization = FutureGuidanceRelocalization(
                client=decision._client,
                file_manager=decision._file_manager,
                request=helper_request,
                future_guidance=future_guidance,
                current_prompt=decision.prompt,
                logger=logger,
                disambiguation=decision.disambiguation,
            )
            relocalization()
            localized_prompt = cast(str, relocalization.additional["localized_prompt"])
            decision.request.backlog_guidance.append(localized_prompt)
            decision.prompt = "\n\n".join((
                decision.prompt,
                decision._prompt_formatter_builder.generate(
                    "LOCALIZED BACKLOG GUIDANCE", escape=True, neutralize_json=True,
                )(localized_prompt),
            ))