lib_run_single.py

# Copyright (c) 2026-present, the authors of StepJack: Benchmarking Computer-Use Agent Safety Against Multi-Step Indirect Prompt Injection.
# Copyright (c) 2025-present, the RedTeamCUA authors.
# Copyright (c) 2024-present, XLANG NLP Lab.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
# Code is based on RedTeamCUA's `lib_run_single.py` (Liao et al., 2025;
# arXiv:2505.21936) from https://github.com/OSU-NLP-Group/RedTeamCUA,
# licensed under Apache-2.0. RedTeamCUA itself derives this file from
# OSWorld (https://os-world.github.io/) at
# https://github.com/xlang-ai/OSWorld by XLANG NLP Lab, also licensed
# under Apache-2.0. Modified to add agent-specific multi-environment
# runner variants (run_single_example_kimi, _qwen, _gpt54) and a
# StepJack sub-step evaluator helper, extending the original
# single-runner scaffolding.
#
import inspect
import datetime
import json
import logging
import os

logger = logging.getLogger("desktopenv.experiment")


def save_sub_step_results(env, example_result_dir):
    """Run each alpha_i (sub-step) evaluator and write `alpha_{i}_result.txt`
    to `example_result_dir`. No-op if the task has no sub-step evaluators
    (1-step examples). Failures in a single alpha are recorded as 0 so the
    other alphas still get written."""
    n = env.num_sub_step_evaluators()
    for idx in range(n):
        try:
            alpha_result = env.sub_step_evaluate(idx)
        except Exception as e:
            logger.warning("alpha_%d evaluator failed: %s", idx + 1, e)
            alpha_result = 0
        logger.info("Alpha_%d Result: %.2f", idx + 1, float(alpha_result))
        with open(os.path.join(example_result_dir, f"alpha_{idx + 1}_result.txt"), "w", encoding="utf-8") as f:
            f.write(f"{alpha_result}\n")


def setup_logger(example, example_result_dir):
    runtime_logger = logging.getLogger(f"desktopenv.example.{example['id']}")
    runtime_logger.setLevel(logging.DEBUG)
    runtime_logger.addHandler(logging.FileHandler(os.path.join(example_result_dir, "runtime.log")))
    return runtime_logger


def run_single_example_kimi(agent, agent_type, env, example, max_steps, instruction, args, example_result_dir, scores):
    runtime_logger = setup_logger(example, example_result_dir)
    sig = inspect.signature(agent.reset)
    if len(sig.parameters) == 0:
        agent.reset()
    else:
        agent.reset(tags = example["tags"], _logger = runtime_logger)
    agent.example_result_dir = example_result_dir
    env.prepare_injection(task_config=example)
    obs = env._get_obs() # Get the initial observation
    action_timestamp = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
    with open(os.path.join(example_result_dir, f"step_0_{action_timestamp}.png"), "wb") as _f:
        _f.write(obs['screenshot'])
    done = False
    step_idx = 0
    env.controller.start_recording()

    while not done and step_idx < max_steps:
        agent.step_idx = step_idx + 1
        response, actions, info_dict = agent.predict(
            instruction,
            obs
        )

        if not actions or len(actions)==0 or actions[0]=="" or actions[0].lower().startswith("error"): 
            break
        # tool_call_is_triggered_l.append(tool_call_is_triggered)
        
        for action in actions:
            # Capture the timestamp before executing the action
            action_timestamp = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
            logger.info("Step %d: %s", step_idx + 1, action)
            obs, reward, done, info = env.step(action, args.sleep_after_execution)

            logger.info("Reward: %.2f", reward)
            logger.info("Done: %s", done)
            # Save screenshot and trajectory information
            with open(os.path.join(example_result_dir, f"step_{step_idx + 1}_{action_timestamp}.png"),
                      "wb") as _f:
                _f.write(obs['screenshot'])
            with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
                f.write(json.dumps({
                    "step_num": step_idx + 1,
                    "action_timestamp": action_timestamp,
                    "action": action,
                    "natural_language_action": info_dict.get("action"),
                    "reward": reward,
                    "done": done,
                    "info": info,
                    "screenshot_file": f"step_{step_idx + 1}_{action_timestamp}.png"
                }, ensure_ascii=False))
                f.write("\n")
            if done:
                logger.info("The episode is done.")
                break
            # Early exit: check both evaluators after each action
            if not getattr(args, 'disable_per_step_evaluator', False):
                try:
                    eval_result = env.evaluate(response, skip_postconfig=True)
                    adversary_eval_result = env.adversary_evaluate()
                    logger.info("Step %d eval: %.2f, adversary_eval: %.2f", step_idx + 1, eval_result, adversary_eval_result)
                    if eval_result and adversary_eval_result:
                        logger.info("Both evaluator and adversary_evaluator passed at step %d. Early exit.", step_idx + 1)
                        done = True
                        break
                except Exception as e:
                    logger.warning("Step %d early exit eval failed: %s", step_idx + 1, e)
        step_idx += 1


    result = env.evaluate(response)
    logger.info("Result: %.2f", result)
    scores.append(result)
    with open(os.path.join(example_result_dir, "result.txt"), "w", encoding="utf-8") as f:
        f.write(f"{result}\n")

    adversary_result = env.adversary_evaluate()
    logger.info("Adversary Result: %.2f", adversary_result)
    with open(os.path.join(example_result_dir, "adversary_result.txt"), "w", encoding="utf-8") as f:
        f.write(f"{adversary_result}\n")

    save_sub_step_results(env, example_result_dir)

    env.controller.end_recording(os.path.join(example_result_dir, "recording.mp4"))


def run_single_example_qwen(agent, agent_type, env, example, max_steps, instruction, args, example_result_dir, scores):
    runtime_logger = setup_logger(example, example_result_dir)
    sig = inspect.signature(agent.reset)
    if len(sig.parameters) == 0:
        agent.reset()
    else:
        agent.reset(tags = example["tags"], _logger = runtime_logger)
    agent.example_result_dir = example_result_dir
    env.prepare_injection(task_config=example)
    obs = env._get_obs() # Get the initial observation
    action_timestamp = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
    with open(os.path.join(example_result_dir, f"step_0_{action_timestamp}.png"), "wb") as _f:
        _f.write(obs['screenshot'])
    done = False
    step_idx = 0
    env.controller.start_recording()

    while not done and step_idx < max_steps:
        agent.step_idx = step_idx + 1
        response, actions = agent.predict(
            instruction,
            obs
        )

        if not actions or len(actions)==0 or actions[0]=="" or actions[0].lower().startswith("error"): 
            break
        # tool_call_is_triggered_l.append(tool_call_is_triggered)
        
        for action in actions:
            # Capture the timestamp before executing the action
            action_timestamp = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
            logger.info("Step %d: %s", step_idx + 1, action)
            obs, reward, done, info = env.step(action, args.sleep_after_execution)

            logger.info("Reward: %.2f", reward)
            logger.info("Done: %s", done)
            # Save screenshot and trajectory information
            with open(os.path.join(example_result_dir, f"step_{step_idx + 1}_{action_timestamp}.png"),
                      "wb") as _f:
                _f.write(obs['screenshot'])
            with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
                f.write(json.dumps({
                    "step_num": step_idx + 1,
                    "action_timestamp": action_timestamp,
                    "action": action,
                    "reward": reward,
                    "done": done,
                    "info": info,
                    "screenshot_file": f"step_{step_idx + 1}_{action_timestamp}.png"
                }, ensure_ascii=False))
                f.write("\n")
            if done:
                logger.info("The episode is done.")
                break
            # Early exit: check both evaluators after each action
            if not getattr(args, 'disable_per_step_evaluator', False):
                try:
                    eval_result = env.evaluate(response, skip_postconfig=True)
                    adversary_eval_result = env.adversary_evaluate()
                    logger.info("Step %d eval: %.2f, adversary_eval: %.2f", step_idx + 1, eval_result, adversary_eval_result)
                    if eval_result and adversary_eval_result:
                        logger.info("Both evaluator and adversary_evaluator passed at step %d. Early exit.", step_idx + 1)
                        done = True
                        break
                except Exception as e:
                    logger.warning("Step %d early exit eval failed: %s", step_idx + 1, e)
        step_idx += 1


    result = env.evaluate(response)
    logger.info("Result: %.2f", result)
    scores.append(result)
    with open(os.path.join(example_result_dir, "result.txt"), "w", encoding="utf-8") as f:
        f.write(f"{result}\n")

    adversary_result = env.adversary_evaluate()
    logger.info("Adversary Result: %.2f", adversary_result)
    with open(os.path.join(example_result_dir, "adversary_result.txt"), "w", encoding="utf-8") as f:
        f.write(f"{adversary_result}\n")

    save_sub_step_results(env, example_result_dir)

    env.controller.end_recording(os.path.join(example_result_dir, "recording.mp4"))

def run_single_example_gpt54(agent, agent_type, env, example, max_steps, instruction, args, example_result_dir, scores):
    runtime_logger = setup_logger(example, example_result_dir)
    agent.reset(tags=example["tags"], _logger=runtime_logger)
    env.prepare_injection(task_config=example)
    obs = env._get_obs()
    action_timestamp = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
    with open(os.path.join(example_result_dir, f"step_0_{action_timestamp}.png"), "wb") as _f:
        _f.write(obs['screenshot'])
    done = False
    step_idx = 0
    response_text = ""
    env.controller.start_recording()

    while not done and step_idx < max_steps:
        predict_info, actions = agent.predict(instruction, obs)

        response_text = predict_info.get("response", "")
        logger.info("Agent response: %s", response_text)
        logger.info("Agent state_correct: %s", predict_info.get("state_correct", False))
        logger.info("Agent model_usage: %s", predict_info.get("model_usage", {}))

        # No actions means model output [INFEASIBLE] or unsupported action — stop.
        if not actions:
            logger.warning("No actions returned (infeasible or unsupported). Breaking loop.")
            break

        for action in actions:
            action_timestamp = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
            logger.info("Step %d: %s", step_idx + 1, action.get("action", ""))
            obs, reward, done, info, step_info = agent.step(action)

            logger.info("Reward: %.2f", reward)
            logger.info("Done: %s", done)

            with open(os.path.join(example_result_dir, f"step_{step_idx + 1}_{action_timestamp}.png"),
                      "wb") as _f:
                _f.write(obs['screenshot'])

            action_for_log = {k: v for k, v in action.items() if k != "pending_checks"}
            with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
                f.write(json.dumps({
                    "step_num": step_idx + 1,
                    "action_timestamp": action_timestamp,
                    "action": action_for_log.get("action", ""),
                    "natural_language_action": response_text,
                    "thinking": predict_info.get("thinking"),
                    "state_correct": predict_info.get("state_correct", False),
                    "model_usage": predict_info.get("model_usage", {}),
                    "reward": reward,
                    "done": done,
                    "info": info,
                    "screenshot_file": f"step_{step_idx + 1}_{action_timestamp}.png"
                }, ensure_ascii=False))
                f.write("\n")
            if done:
                logger.info("The episode is done.")
                break
            # Early exit: check both evaluators after each action
            if not getattr(args, 'disable_per_step_evaluator', False):
                try:
                    eval_result = env.evaluate(response_text, skip_postconfig=True)
                    adversary_eval_result = env.adversary_evaluate()
                    logger.info("Step %d eval: %.2f, adversary_eval: %.2f", step_idx + 1, eval_result, adversary_eval_result)
                    if eval_result and adversary_eval_result:
                        logger.info("Both evaluator and adversary_evaluator passed at step %d. Early exit.", step_idx + 1)
                        done = True
                        break
                except Exception as e:
                    logger.warning("Step %d early exit eval failed: %s", step_idx + 1, e)
        step_idx += 1

    result = env.evaluate(response_text)
    logger.info("Result: %.2f", result)
    scores.append(result)
    with open(os.path.join(example_result_dir, "result.txt"), "w", encoding="utf-8") as f:
        f.write(f"{result}\n")

    adversary_result = env.adversary_evaluate()
    logger.info("Adversary Result: %.2f", adversary_result)
    with open(os.path.join(example_result_dir, "adversary_result.txt"), "w", encoding="utf-8") as f:
        f.write(f"{adversary_result}\n")

    save_sub_step_results(env, example_result_dir)

    env.controller.end_recording(os.path.join(example_result_dir, "recording.mp4"))

def run_single_example_claude(agent, agent_type, env, example, max_steps, instruction, args, example_result_dir, scores):
    """Run one example with the Claude (Anthropic) computer-use agent.

    Mirrors run_single_example_kimi but adapts to the Claude agent's predict
    contract (returns reasonings + a list of action dicts whose `command`
    field is the pyautogui code to execute, or DONE/FAIL/CALL_USER sentinel).
    """
    runtime_logger = setup_logger(example, example_result_dir)
    sig = inspect.signature(agent.reset)
    if len(sig.parameters) == 0:
        agent.reset()
    else:
        agent.reset(tags=example["tags"], _logger=runtime_logger)
    agent.example_result_dir = example_result_dir
    env.prepare_injection(task_config=example)
    obs = env._get_obs()
    action_timestamp = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
    with open(os.path.join(example_result_dir, f"step_0_{action_timestamp}.png"), "wb") as _f:
        _f.write(obs['screenshot'])

    done = False
    step_idx = 0
    response_text = ""
    env.controller.start_recording()

    while not done and step_idx < max_steps:
        agent.step_idx = step_idx + 1
        response_text, actions = agent.predict(instruction, obs)

        if not actions:
            logger.warning("No actions returned by Claude agent. Breaking loop.")
            break

        for action in actions:
            command = action.get("command")
            action_type = action.get("action_type")
            raw_response = action.get("raw_response", "")

            # Map Claude sentinel actions to env-level FAIL/DONE strings.
            if command in (None, "") and action_type in ("FAIL", "DONE", "CALL_USER"):
                command = "FAIL" if action_type in ("FAIL", "CALL_USER") else "DONE"
            if command in ("FAIL", "DONE"):
                env_action = command
            else:
                # Claude can emit multiple pyautogui statements per tool call;
                # we wrap them in `import pyautogui, time` so the executor has
                # access to both modules (mouse_down + sleep + mouse_up paths).
                env_action = "import pyautogui\nimport time\n" + (command or "")

            action_timestamp = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
            logger.info("Step %d: %s", step_idx + 1, env_action.strip().splitlines()[-1] if env_action else env_action)
            obs, reward, done, info = env.step(env_action, args.sleep_after_execution)

            logger.info("Reward: %.2f", reward)
            logger.info("Done: %s", done)

            with open(os.path.join(example_result_dir, f"step_{step_idx + 1}_{action_timestamp}.png"), "wb") as _f:
                _f.write(obs['screenshot'])
            with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
                f.write(json.dumps({
                    "step_num": step_idx + 1,
                    "action_timestamp": action_timestamp,
                    "action": env_action,
                    "tool_call": {"name": action.get("name"), "input": action.get("input")},
                    "action_type": action_type,
                    "raw_response": raw_response,
                    "reasonings": response_text,
                    "reward": reward,
                    "done": done,
                    "info": info,
                    "screenshot_file": f"step_{step_idx + 1}_{action_timestamp}.png",
                }, ensure_ascii=False))
                f.write("\n")

            if command in ("FAIL", "DONE"):
                done = True

            if done:
                logger.info("The episode is done.")
                break

            if not getattr(args, 'disable_per_step_evaluator', False):
                try:
                    eval_result = env.evaluate(response_text, skip_postconfig=True)
                    adversary_eval_result = env.adversary_evaluate()
                    logger.info("Step %d eval: %.2f, adversary_eval: %.2f", step_idx + 1, eval_result, adversary_eval_result)
                    if eval_result and adversary_eval_result:
                        logger.info("Both evaluator and adversary_evaluator passed at step %d. Early exit.", step_idx + 1)
                        done = True
                        break
                except Exception as e:
                    logger.warning("Step %d early exit eval failed: %s", step_idx + 1, e)

        step_idx += 1

    result = env.evaluate(response_text)
    logger.info("Result: %.2f", result)
    scores.append(result)
    with open(os.path.join(example_result_dir, "result.txt"), "w", encoding="utf-8") as f:
        f.write(f"{result}\n")

    adversary_result = env.adversary_evaluate()
    logger.info("Adversary Result: %.2f", adversary_result)
    with open(os.path.join(example_result_dir, "adversary_result.txt"), "w", encoding="utf-8") as f:
        f.write(f"{adversary_result}\n")

    save_sub_step_results(env, example_result_dir)

    env.controller.end_recording(os.path.join(example_result_dir, "recording.mp4"))