AlgoTune/algotune_agent

"""Reference-faithful AlgoTune command agent registered as a Galapagos scaffold."""
from __future__ import annotations

from ...config import GalapagosConfig
from ...models import GalapagosModel
from ..base_scaffold import GalapagosScaffold
from ..registry import register_scaffold
from .memory import AlgoTuneMemory
from .population import AlgoTunePopulation
from .prompt_builder import HISTORY_PER_ROLE, SPEND_LIMIT_USD, AlgoTunePromptBuilder
from .proposer import AlgoTuneProposer, format_evaluation
from .selection_policy import AlgoTuneSelectionPolicy

_CONTROL_COMMANDS = {
    "ls", "view_file", "reference", "eval_input", "eval", "profile", "profile_lines", "revert",
}


@register_scaffold("algotune_agent")
class AlgoTuneScaffold(GalapagosScaffold):
    """SWE-agent-style command loop shipped by the AlgoTune authors."""

    name = "algotune_agent"
    card_name = "algotune_agent"
    status_unit = "command"
    # Diagnostic tools and Cython must see the exact packages used by the task scorer.
    run_whole_loop_in_task_container = True

    @classmethod
    def build_components(cls, config: GalapagosConfig, model: GalapagosModel | None) -> dict:
        return {
            "population": AlgoTunePopulation(),
            "selection_policy": AlgoTuneSelectionPolicy(seed=int(config.seed)),
            "prompt_builder": AlgoTunePromptBuilder(),
            "proposer": AlgoTuneProposer(),
            "memory": AlgoTuneMemory(max_messages_per_role=HISTORY_PER_ROLE),
        }

    def setup(self, task) -> None:
        self._bind_task(task)
        super().setup(task)
        self.selection_policy.set_current(self.ctx.best)

    def _bind_task(self, task) -> None:
        if not str(getattr(task, "name", "")).startswith("algotune_"):
            raise ValueError(
                "the algotune_agent scaffold requires an `algotune_*` task because its reference, "
                "validator, and profiling commands use that task contract"
            )
        self.prompt_builder.bind_task(task)

    def _resume(self, task, path: str) -> None:
        # Base resume intentionally bypasses setup. Rebind immutable task prompt sources before the
        # first resumed command; component state restores only the mutable transcript/workspace.
        self._bind_task(task)
        super()._resume(task, path)

    def after_step(self, child, result) -> None:
        command = child.metadata.get("algotune_command", "command")
        detail = str(child.metadata.get("algotune_result") or "")
        no_candidate = bool((child.trace or {}).get("etif_no_candidate"))
        subject_id = child.parent_id if no_candidate else child.id
        terminal_event_id = (
            (child.trace or {}).get("etif_terminal_event_id")
            or (child.trace or {}).get("etif_tool_event_id")
            or self._candidate_last_event_ids.get(child.id)
        )
        if command == "revert" and child.metadata.get("algotune_revert_target_id"):
            target_id = child.metadata["algotune_revert_target_id"]
            self.selection_policy.set_current(target_id)
            self._emit_event(
                event_type="adaptation",
                action="restore",
                status="completed",
                component_role="selection_policy",
                evolution_track="solution",
                caused_by_event_ids=[terminal_event_id],
                inputs=[self._event_ref("candidate", subject_id, "workspace_before")],
                outputs=[self._event_ref("candidate", target_id, "restored_working_candidate")],
                details={
                    "target": "working_candidate",
                    "reason": "revert_command",
                    "state_changed": True,
                },
            )
        # The generic loop calls every unevaluated Genome a no-diff. These are intentional agent tool
        # turns, not failed mutations, so remove them from the no-diff run statistic after base.step
        # has recorded the command.
        if command in _CONTROL_COMMANDS and result is None:
            self.ctx.blackboard["no_diff"] = max(
                0, int(self.ctx.blackboard.get("no_diff", 0)) - 1
            )
        if child.metadata.get("algotune_needs_evaluation"):
            if result is not None:
                detail = f"{detail}\n{format_evaluation(result)}".strip()
            elif len(child.content) > self.general.max_solution_length:
                detail = (
                    f"{detail}\nCandidate exceeded max_solution_length; the edit was not admitted."
                ).strip()
            else:
                detail = f"{detail}\nThe workspace did not reach evaluation.".strip()

        # Upstream automatically rolls back Cython/Pythran edits whose compilation fails. Do not
        # conflate a compiled file that built successfully but returned a wrong answer with a build
        # failure: the former remains the mutable workspace, just as upstream does.
        feedback = ""
        if result is not None:
            feedback = str(
                result.text_feedback or (result.artifacts or {}).get("text_feedback") or ""
            ).lower()
        if result is not None and not result.valid and "compilation failed" in feedback:
            self.selection_policy.set_current(child.parent_id)
            detail += "\nCompiled workspace rejected; restored the preceding working version."
            self._emit_event(
                event_type="adaptation",
                action="restore",
                status="completed",
                component_role="selection_policy",
                evolution_track="solution",
                caused_by_event_ids=[self._candidate_last_event_ids.get(child.id)],
                inputs=[self._event_ref("candidate", child.id, "failed_working_candidate")],
                outputs=[self._event_ref(
                    "candidate", child.parent_id, "restored_working_candidate"
                )],
                details={
                    "target": "working_candidate",
                    "reason": "compilation_failure",
                    "state_changed": True,
                },
            )

        # The saved snapshot, not the mutable current workspace, is the search winner.
        best = self.population.best()
        if best is not None:
            self.ctx.best = best

        remaining = max(0.0, SPEND_LIMIT_USD - self.ctx.cost_usd)
        self.memory.add(
            "user",
            f"[{command}]\n{detail}\n\n"
            f"Budget used: ${self.ctx.cost_usd:.4f} / ${SPEND_LIMIT_USD:.4f}; "
            f"remaining: ${remaining:.4f}.",
        )
        self._emit_event(
            event_type="memory",
            action="write",
            status="completed",
            component_role="memory",
            evolution_track="solution",
            caused_by_event_ids=[terminal_event_id],
            inputs=[self._event_ref(
                "candidate", subject_id,
                "workspace" if no_candidate else "observed_candidate",
            )],
            details={
                "memory_name": "conversation_history",
                "entry_type": "evaluation_feedback",
                "command": command,
                "content": detail,
                "remaining_budget_usd": remaining,
            },
        )

    def _should_stop(self) -> bool:
        return super()._should_stop() or self.ctx.cost_usd >= SPEND_LIMIT_USD

    def _stop_reason(self) -> str:
        if self.ctx.cost_usd >= SPEND_LIMIT_USD:
            return "cost_budget_exhausted"
        return super()._stop_reason()

    def _record(self, genome, result, **kwargs) -> None:
        """Record intentional tool turns as applied controls, not failed mutations."""
        if (
            genome.metadata.get("algotune_command") in _CONTROL_COMMANDS
            and kwargs.get("reason") == "no_diff"
        ):
            kwargs["reason"] = "control_action"
        super()._record(genome, result, **kwargs)