inference.py

from __future__ import annotations

import argparse
import json
import os
import re
import shlex
import sys
import traceback
from typing import Any

import requests
from openai import OpenAI


MAX_STEPS = 40
DEFAULT_ENV_URL = "http://localhost:7860"
ENV_NAME = "secrets_audit"
PRIMARY_FILE_BY_TASK_ID = {
    1: "config.py",
    2: "db.py",
    3: "settings.js",
    4: "logger.py",
    5: ".env",
    6: "utils.py",
    7: "deploy.yml",
    8: "app.toml",
    9: "migrate.sql",
    10: "deploy.sh",
    11: "service_a.py",
    12: "crypto.py",
    13: "config.py",
}


def normalize_action(raw: str) -> str:
    text = raw.strip()
    if not text:
        return "true"
    if _looks_like_provider_error(text):
        return "true"

    # Try extracting from ```bash ... ``` fences first
    fence_match = re.search(r"```(?:bash|sh)?\s*(.*?)```", text, re.DOTALL)
    if fence_match:
        text = fence_match.group(1).strip()
    if _looks_like_provider_error(text):
        return "true"

    # Try extracting from XML tool_call (Minimax, etc.)
    xml_match = re.search(r'<parameter\s+name="command">(.*?)</parameter>', text, re.DOTALL)
    if xml_match:
        text = xml_match.group(1).strip()

    lines = [
        line.strip()
        for line in text.splitlines()
        if line.strip() and not line.strip().startswith(("#", "-", "*"))
    ]
    if not lines:
        return "true"

    # If first line looks like a shell command, use it
    first_line = lines[0]
    if _looks_like_shell(first_line):
        return first_line

    # If first line is prose, scan remaining lines for a shell command
    for line in lines[1:]:
        if _looks_like_shell(line):
            return line

    return "true"


def _looks_like_shell(line: str) -> bool:
    if not line:
        return False
    prose_prefixes = (
        "here",
        "this",
        "i ",
        "i'",
        "you ",
        "the ",
        "to ",
        "we ",
        "run ",
        "use ",
        "first ",
        "let ",
        "let's",
        "looking",
        "now ",
        "next ",
        "since ",
        "note",
        "okay",
        "sure",
        "great",
        "step ",
        "<",
    )
    lowered = line.lower()
    if lowered.startswith(prose_prefixes):
        return False
    if _looks_like_provider_error(line):
        return False
    return True


def _looks_like_provider_error(text: str) -> bool:
    lowered = text.strip().lower()
    error_markers = (
        "internal server error",
        "server error",
        "bad gateway",
        "gateway timeout",
        "service unavailable",
        "rate limit",
        "too many requests",
        "upstream error",
        "provider error",
        "api error",
        "error code:",
        "request failed",
    )
    return any(lowered.startswith(marker) for marker in error_markers)


def stderr_log(message: str) -> None:
    print(message, file=sys.stderr, flush=True)


def format_field(value: Any) -> str:
    text = "none" if value is None else re.sub(r"\s+", " ", str(value).strip())
    if not text:
        text = "none"
    if " " in text:
        return json.dumps(text)
    return text


def bool_text(value: bool) -> str:
    return "true" if value else "false"


def stdout_tag(tag: str, **fields: Any) -> None:
    parts = [f"[{tag}]"]
    for key, value in fields.items():
        parts.append(f"{key}={format_field(value)}")
    print(" ".join(parts), flush=True)


def build_prompt(
    state: dict[str, Any],
    recent_actions: list[str] | None = None,
    stuck_warning: str | None = None,
) -> str:
    sanitized_state = sanitize_state_for_prompt(state)
    session = sanitized_state["session"]
    last_result = session.get("last_result") or {}
    recent_actions = recent_actions or []
    recent_summary = ", ".join(recent_actions[-5:]) if recent_actions else "none"
    health_stdout = sanitize_observation_text(session.get("health_stdout", ""))
    health_stderr = sanitize_observation_text(session.get("health_stderr", ""))
    observation = sanitize_observation_text(session.get("observation", ""))
    stuck_section = f"\n[CRITICAL WARNING: YOU ARE STUCK] {stuck_warning}\n" if stuck_warning else ""
    return (
        "You are remediating leaked secrets in a single benchmark workspace.\n"
        f"Task {session['task_id']}: {session['title']}\n"
        f"Description: {session['description']}\n"
        f"Workspace: {session['workspace']}\n"
        f"Current reward: {session['reward']}\n"
        f"Current leaks: {session['current_leaks']}\n"
        f"Health score: {session['health_score']}\n"
        f"Health stdout:\n{health_stdout}\n"
        f"Health stderr:\n{health_stderr}\n"
        f"Observation:\n{observation}\n"
        f"Recent actions: {recent_summary}\n"
        f"Previous action: {last_result.get('action', '')}\n"
        f"Previous stdout:\n{last_result.get('stdout', '')}\n"
        f"Previous stderr:\n{last_result.get('stderr', '')}\n"
        "Return exactly one executable bash command and nothing else.\n"
        "Senior Dev SOP:\n"
        "Rule A: Atomic Only. Run exactly ONE command per step. Chaining with && or {} is strictly forbidden.\n"
        "Rule B: You MUST add import os to the top of any file where you use os.getenv or os.environ.get.\n"
        "Rule C: Reward Math. Total = Security x Health. A 0.5 reward means Health is 1.0 but Security is 0.5. A 0.0 reward means your code is broken.\n"
        "Rule D: If reward is 0.5, the live file is fixed but git history still leaks the secret. You MUST use git filter-repo or git filter-branch.\n"
        "Use relative paths only. You are already running inside the task root. Never use absolute filesystem paths.\n"
        "Start by inspecting or editing the relevant project files. Prefer precise file reads over repeating repository-wide status commands."
        f"{stuck_section}"
    )


def call_model(client: OpenAI, model_name: str, prompt: str) -> str:
    response = client.chat.completions.create(
        model=model_name,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1,
    )
    content = extract_response_text(response)
    if content:
        return content

    stderr_log(f"model_response_without_text={safe_model_dump(response)}")
    return ""


def extract_response_text(response: Any) -> str:
    choices = getattr(response, "choices", None)
    if choices:
        first_choice = choices[0]
        message = getattr(first_choice, "message", None)
        if message is not None:
            content = getattr(message, "content", None)
            extracted = coerce_content_to_text(content)
            if extracted:
                return extracted

    output_text = getattr(response, "output_text", None)
    if output_text:
        return str(output_text)

    output = getattr(response, "output", None)
    if output:
        extracted = coerce_content_to_text(output)
        if extracted:
            return extracted

    dumped = safe_model_dump(response)
    return extract_text_from_dump(dumped)


def coerce_content_to_text(content: Any) -> str:
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: list[str] = []
        for item in content:
            if isinstance(item, str):
                parts.append(item)
                continue
            if isinstance(item, dict):
                text = item.get("text")
                if isinstance(text, str):
                    parts.append(text)
                    continue
                inner = item.get("content")
                if isinstance(inner, str):
                    parts.append(inner)
            else:
                text = getattr(item, "text", None)
                if isinstance(text, str):
                    parts.append(text)
        return "\n".join(part for part in parts if part).strip()
    if isinstance(content, dict):
        for key in ("text", "content", "output_text"):
            value = content.get(key)
            if isinstance(value, str) and value.strip():
                return value
        return ""
    return str(content).strip()


def safe_model_dump(response: Any) -> dict[str, Any]:
    if hasattr(response, "model_dump"):
        try:
            dumped = response.model_dump()
            if isinstance(dumped, dict):
                return dumped
        except Exception:
            return {"repr": repr(response)}
    if isinstance(response, dict):
        return response
    return {"repr": repr(response)}


def extract_text_from_dump(payload: Any) -> str:
    if isinstance(payload, str):
        return payload.strip()
    if isinstance(payload, list):
        for item in payload:
            extracted = extract_text_from_dump(item)
            if extracted:
                return extracted
        return ""
    if isinstance(payload, dict):
        for key in ("content", "text", "output_text"):
            value = payload.get(key)
            if isinstance(value, str) and value.strip():
                return value.strip()
            extracted = extract_text_from_dump(value)
            if extracted:
                return extracted
        for value in payload.values():
            extracted = extract_text_from_dump(value)
            if extracted:
                return extracted
    return ""


def post_json(base_url: str, path: str, payload: dict[str, Any], timeout: int) -> dict[str, Any]:
    response = requests.post(f"{base_url}{path}", json=payload, timeout=timeout)
    response.raise_for_status()
    return response.json()


def is_done(state: dict[str, Any]) -> bool:
    session = state["session"]
    return float(session["reward"]) >= 1.0


def extract_error(session: dict[str, Any]) -> str:
    session_error = (session.get("error") or "").strip()
    if session_error and session_error != "none":
        return session_error
    last_result = session.get("last_result") or {}
    if last_result.get("timed_out"):
        return f"timeout:{last_result.get('stderr') or 'command timed out'}"
    if int(last_result.get("exit_code", 0)) != 0:
        return (last_result.get("stderr") or last_result.get("stdout") or "command failed").strip()
    return "none"


def parse_task_id(task_value: str) -> int:
    text = str(task_value).strip()
    match = re.fullmatch(r"task_(\d+)", text)
    if match:
        return int(match.group(1))
    return int(text)


def detect_repeated_action(actions: list[str], rewards: list[float]) -> tuple[str, str] | None:
    if len(actions) < 3 or len(rewards) < 3:
        return None
    if actions[-1] == actions[-2] == actions[-3] and rewards[-1] == rewards[-2] == rewards[-3]:
        warning = (
            f"The last three steps repeated {actions[-1]!r} and reward stayed at {rewards[-1]}. "
            "You are stuck. Choose a different single command that inspects or patches the primary source file."
        )
        return actions[-1], warning
    return None


def choose_fallback_action(task_id: int, repeated_action: str) -> str:
    candidate = PRIMARY_FILE_BY_TASK_ID.get(task_id)
    if "git status" in repeated_action:
        if candidate:
            return f"sed -n '1,200p' {shlex.quote(candidate)}"
        return "find . -maxdepth 2 -type f"
    if repeated_action.startswith("cat ") or repeated_action.startswith("sed -n"):
        return "pytest -q"
    if candidate:
        return f"sed -n '1,200p' {shlex.quote(candidate)}"
    return "find . -maxdepth 2 -type f"


def sanitize_state_for_prompt(state: dict[str, Any]) -> dict[str, Any]:
    return sanitize_observation_value(state)


def sanitize_observation_value(value: Any) -> Any:
    if isinstance(value, dict):
        return {key: sanitize_observation_value(item) for key, item in value.items()}
    if isinstance(value, list):
        return [sanitize_observation_value(item) for item in value]
    if isinstance(value, str):
        return sanitize_observation_text(value)
    return value


def sanitize_observation_text(text: str) -> str:
    if not text:
        return text
    sanitized = text
    repo_root = os.getcwd()
    sanitized = sanitized.replace(f"{repo_root}/", "")
    sanitized = sanitized.replace(repo_root, ".")
    sanitized = re.sub(r"/[^/\s]*/runtime/session_[^/\s]+/", "", sanitized)
    sanitized = re.sub(r"/home/[^/\s]+/", "", sanitized)
    sanitized = re.sub(r"\.{2,}", ".", sanitized)
    return sanitized


def enforce_atomic_action(action: str) -> str:
    if not action:
        return "true"
    if re.search(r"&&|\|\||[;{}]", action):
        return "true"
    return action


def main() -> None:
    parser = argparse.ArgumentParser(description="Meta OpenEnv Round 1 inference loop.")
    parser.add_argument("--task-id", default=os.environ.get("TASK_ID", "task_1"))
    args = parser.parse_args()

    api_base_url = os.environ["API_BASE_URL"].rstrip("/")
    hf_token = os.environ["HF_TOKEN"]
    model_name = os.environ["MODEL_NAME"]
    env_url = os.environ.get("ENV_URL", DEFAULT_ENV_URL).rstrip("/")
    task_id_value = str(args.task_id)
    task_id = parse_task_id(task_id_value)

    client = OpenAI(base_url=api_base_url, api_key=hf_token, timeout=60)

    rewards: list[float] = []
    actions: list[str] = []
    success = False
    steps_run = 0
    final_score = 0.0

    stdout_tag("START", task=task_id_value, env=ENV_NAME, model=model_name)

    try:
        state = post_json(env_url, "/reset", {"task_id": task_id}, timeout=30)
        final_score = float(state["session"]["reward"])

        for step_num in range(1, MAX_STEPS + 1):
            stderr_log(f"step={step_num} building prompt")
            repeated = detect_repeated_action(actions, rewards)
            forbidden_action = None
            stuck_warning = None
            if repeated:
                forbidden_action, stuck_warning = repeated
                stderr_log(f"step={step_num} repeated_action_detected={forbidden_action!r}")
            prompt = build_prompt(state, recent_actions=actions, stuck_warning=stuck_warning)
            raw = call_model(client, model_name, prompt)
            stderr_log(f"step={step_num} raw_response={raw!r}")

            action = normalize_action(raw)
            atomic_action = enforce_atomic_action(action)
            if atomic_action != action:
                stderr_log(f"step={step_num} rejected_non_atomic_action={action!r}")
                action = atomic_action
            if forbidden_action and action == forbidden_action:
                fallback_action = choose_fallback_action(task_id, forbidden_action)
                stderr_log(
                    f"step={step_num} overriding_repeated_action={forbidden_action!r} fallback={fallback_action!r}"
                )
                action = fallback_action
            state = post_json(env_url, "/step", {"action": action}, timeout=90)

            session = state["session"]
            final_score = float(session["reward"])
            actions.append(action)
            rewards.append(final_score)
            done = is_done(state)
            error_text = extract_error(session)

            stdout_tag(
                "STEP",
                step=step_num,
                action=action,
                reward=final_score,
                done=bool_text(done),
                error=error_text,
            )

            steps_run = step_num
            if done:
                success = True
                break

    except Exception as exc:
        stderr_log(f"fatal_error={exc!r}")
        traceback.print_exc(file=sys.stderr)

    stdout_tag(
        "END",
        success=bool_text(success),
        steps=steps_run,
        score=final_score,
        rewards=json.dumps(rewards),
    )


if __name__ == "__main__":
    main()