Agentswarm/orchestrator/queen.py

"""Queen — the swarm's terminal arbitration layer (agent_swarm#8/#12, Queen role).

The Queen does NOT execute or dispatch tasks. She only judges the FINAL result of a run:
  - aggregate the candidate artifacts produced by the fan-out agents,
  - score each candidate by running its tests in the sandbox,
  - SELECT the single best candidate (best-of-N — the emergence lever a single model lacks),
  - (M2) promote the winner to the artifact repo's `main` as one coherent deliverable,
  - (M3) if no candidate meets the quality bar, send work BACK for another round (reopen_task),
  - (P0, in main.py) arbitrate competition and reject cross-run grabs.

This module is intentionally import-light and does NOT import orchestrator.main (avoids an import
cycle): `aggregate_run` returns a plain summary dict; the caller (refresh_swarm_run_status) folds it
into the run's deliverable. `select_best` is a pure function so it can be unit-tested without I/O.
"""
from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


@dataclass
class Candidate:
    """One agent's contribution to the run, with its sandbox score."""
    task_id: str
    agent_id: Optional[str]
    impl_files: List[Any] = field(default_factory=list)   # SandboxFile (impl)
    score: Optional[float] = None                         # pass_rate 0-100, None = not scored
    total: int = 0
    passed: int = 0


def select_best(candidates: List[Candidate]) -> Optional[Candidate]:
    """Pure selection: the candidate with the highest test pass_rate wins. Unscored (None)
    candidates rank below any scored one; ties and all-unscored fall back to the first candidate
    that actually carries impl files (deterministic — preserves input order). None if no candidate
    has impl files."""
    with_impl = [c for c in candidates if c.impl_files]
    if not with_impl:
        return None
    scored = [c for c in with_impl if c.score is not None]
    if scored:
        # max by score; stable on ties (first in input order wins)
        return max(scored, key=lambda c: (c.score, c.passed, -with_impl.index(c)))
    return with_impl[0]


def should_bounce(summary: Dict[str, Any], threshold: Optional[float],
                  cycles: int, max_cycles: int) -> bool:
    """Pure decision (M3/SC-9): should the run be sent BACK for another round?

    True only when the best candidate WAS scored, fell BELOW `threshold`, and the review-cycle cap
    isn't hit yet. False (= accept / converge) when: no threshold (gate disabled), not scored
    (honest — don't bounce on a score we couldn't compute, rule #9), already meets the bar, or the
    cap is reached (convergence then marks MAX_ROUNDS_REACHED on the best-so-far)."""
    if threshold is None:
        return False
    winner = summary.get("winner")
    if not winner or winner.get("score") is None:
        return False
    if winner["score"] >= threshold:
        return False
    return cycles < max_cycles


def _result_of(task) -> Dict[str, Any]:
    """task.result → dict (JSON string or dict), {} on failure. Mirrors main.parse_task_result
    without importing main."""
    import json
    raw = getattr(task, "result", None)
    if isinstance(raw, dict):
        return raw
    if isinstance(raw, str):
        try:
            return json.loads(raw)
        except Exception:
            return {}
    return {}


def collect_candidates(tasks) -> List[Candidate]:
    """One Candidate per completed task that produced impl files. Reuses quality.collect_generated_files
    (per task) to split impl vs test_*.py, so the Queen scores each agent's implementation."""
    from .quality import collect_generated_files
    out: List[Candidate] = []
    for t in tasks:
        files = collect_generated_files([t])
        impl = files.get("impl") or []
        if not impl:
            continue
        out.append(Candidate(
            task_id=getattr(t, "task_id", "?"),
            agent_id=getattr(t, "assigned_agent_id", None) or getattr(t, "agent_role", None),
            impl_files=impl,
        ))
    return out


def _shared_tests(tasks) -> List[Any]:
    """All test_*.py the swarm produced this run — the shared yardstick the Queen scores impls against."""
    from .quality import collect_generated_files
    seen: Dict[str, Any] = {}
    for t in tasks:
        for tf in (collect_generated_files([t]).get("agent_tests") or []):
            seen[tf.path] = tf  # de-dup by path, last writer wins
    return list(seen.values())


async def score_candidates(candidates: List[Candidate], tests: List[Any]) -> None:
    """Score each candidate in-place: run its impl against the shared test set in the sandbox.

    FAIL-CLOSED + honest: if isolation isn't confirmed (sandbox.assert_isolated would refuse) or
    there are no tests, scores stay None (not 0 — 'not scored' != 'scored zero', org rule #9)."""
    from .sandbox import run_tests, isolation_confirmed
    if not tests or not isolation_confirmed():
        return
    for c in candidates:
        try:
            res = await asyncio.to_thread(run_tests, c.impl_files, tests)
            c.score = res.pass_rate
            c.total = res.total
            c.passed = res.passed
        except Exception as exc:  # a scoring error leaves this candidate unscored, never crashes
            logger.warning("queen: scoring candidate %s failed: %s", c.task_id, exc)


async def aggregate_run(run, tasks) -> Dict[str, Any]:
    """Queen entry point at run terminal. Collect → score → select best. Returns a summary dict the
    caller folds into the deliverable. Never raises (best-effort; a failure leaves winner=None and
    the caller keeps the legacy per-task deliverable)."""
    try:
        candidates = collect_candidates(tasks)
        if not candidates:
            return {"winner": None, "candidate_count": 0, "reason": "no_impl_artifacts"}
        tests = _shared_tests(tasks)
        await score_candidates(candidates, tests)
        best = select_best(candidates)
        return {
            "winner": (
                {"task_id": best.task_id, "agent_id": best.agent_id,
                 "score": best.score, "passed": best.passed, "total": best.total}
                if best else None
            ),
            "candidate_count": len(candidates),
            "scored": sum(1 for c in candidates if c.score is not None),
            "candidates": [
                {"task_id": c.task_id, "agent_id": c.agent_id, "score": c.score}
                for c in candidates
            ],
        }
    except Exception as exc:  # Queen never breaks the run's terminal path
        logger.warning("queen: aggregate_run failed for %s: %s", getattr(run, "swarm_id", "?"), exc)
        return {"winner": None, "candidate_count": 0, "reason": f"error:{exc!r}"}


def _auth_url(repo_url: str, user: Optional[str], pw: Optional[str]) -> str:
    """Embed credentials into an http(s) clone URL. URL-encodes the password (handles '@','/', etc).
    The result is NEVER logged."""
    import urllib.parse
    if not user or not pw or "://" not in repo_url:
        return repo_url
    scheme, rest = repo_url.split("://", 1)
    return f"{scheme}://{urllib.parse.quote(user, safe='')}:{urllib.parse.quote(pw, safe='')}@{rest}"


def _git_promote(env: Dict[str, str], base_branch: str, files: List[Any], swarm_id: str) -> Dict[str, Any]:
    """Sync git: clone base_branch → overwrite with winner files → commit → push base_branch.
    Blocking (run via asyncio.to_thread). Credentials live only in the clone URL, never logged."""
    import os
    import shutil
    import subprocess
    import tempfile
    repo_url = env["GIT_REPO_URL"]
    auth_url = _auth_url(repo_url, env.get("GIT_USERNAME"), env.get("GIT_PASSWORD"))
    workdir = tempfile.mkdtemp(prefix="queen-promote-")
    repo_dir = os.path.join(workdir, "repo")

    def git(*args, cwd=None):
        return subprocess.run(["git", *args], cwd=cwd or workdir,
                              capture_output=True, text=True, timeout=180)

    try:
        if git("clone", "--depth", "1", "--branch", base_branch, auth_url, "repo").returncode != 0:
            return {"promoted": False, "reason": "clone_failed"}
        for f in files:
            dest = os.path.join(repo_dir, f.path)
            os.makedirs(os.path.dirname(dest) or repo_dir, exist_ok=True)
            with open(dest, "w", encoding="utf-8") as fh:
                fh.write(f.content)
        git("config", "user.email", "queen@heicode.swarm", cwd=repo_dir)
        git("config", "user.name", "HeiCode Queen", cwd=repo_dir)
        git("add", "-A", cwd=repo_dir)
        c = git("commit", "-m", f"queen: promote best swarm artifact ({swarm_id})", cwd=repo_dir)
        if c.returncode != 0:
            return {"promoted": False, "reason": "no_changes"}
        if git("push", "origin", base_branch, cwd=repo_dir).returncode != 0:
            return {"promoted": False, "reason": "push_failed"}
        sha = git("rev-parse", "HEAD", cwd=repo_dir).stdout.strip()
        return {"promoted": True, "commit_sha": sha, "branch": base_branch}
    except Exception as exc:
        return {"promoted": False, "reason": f"error:{exc!r}"}
    finally:
        shutil.rmtree(workdir, ignore_errors=True)


async def promote_to_main(run, tasks, winner_task_id: str) -> Dict[str, Any]:
    """SC-7: push the winning candidate's files to the artifact repo's base branch (main), so the
    run delivers ONE coherent artifact, not N scattered agent branches. Credentials are resolved
    from the grant's secret_ref at call time (never stored/logged). Best-effort: never raises."""
    try:
        grant = (run.metadata or {}).get("git_grant") or {}
        repo_url = grant.get("repo_url")
        if not repo_url:
            return {"promoted": False, "reason": "no_git_grant"}
        wtask = next((t for t in tasks if getattr(t, "task_id", None) == winner_task_id), None)
        if wtask is None:
            return {"promoted": False, "reason": "winner_not_found"}
        from .quality import collect_generated_files
        cf = collect_generated_files([wtask])
        files = (cf.get("impl") or []) + (cf.get("agent_tests") or [])
        if not files:
            return {"promoted": False, "reason": "winner_no_files"}
        from .agent_launcher import resolve_git_grant
        env = resolve_git_grant({"resource_grants": [{
            "resource_type": "git", "secret_ref": grant.get("secret_ref"),
            "metadata": {"repo_url": repo_url, "base_branch": grant.get("base_branch", "main")},
        }]})
        if not env or "GIT_REPO_URL" not in env:
            return {"promoted": False, "reason": "grant_unresolved"}
        return await asyncio.to_thread(_git_promote, env, grant.get("base_branch", "main"),
                                       files, getattr(run, "swarm_id", "?"))
    except Exception as exc:
        logger.warning("queen: promote_to_main failed for %s: %s", getattr(run, "swarm_id", "?"), exc)
        return {"promoted": False, "reason": f"error:{exc!r}"}