run_kimi_multienv.py

# Copyright (c) 2026-present, the authors of StepJack: Benchmarking Computer-Use Agent Safety Against Multi-Step Indirect Prompt Injection.
# 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 OSWorld's `scripts/python/run_multienv_kimi_k25.py`
# from https://github.com/xlang-ai/OSWorld by XLANG NLP Lab, licensed
# under Apache-2.0. 
#
"""
Script to run end-to-end evaluation on the benchmark (multi-environment parallel version).
"""

from __future__ import annotations
import argparse
import datetime
import json
import logging
import os
import sys
import signal
import threading
import time
import traceback
from typing import List
from multiprocessing import Process, Manager, current_process

import lib_run_single
from desktop_env.desktop_env import DesktopEnv
from mm_agents.kimi import KimiAgent

# Global variables for signal handling
active_environments = []
processes = []
is_terminating = False

# load the environment variables from .env file
if os.path.exists(".env"):
    from dotenv import load_dotenv
    load_dotenv()

#  Logger Configs {{{ #
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)

datetime_str: str = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
if not os.path.exists("logs"):
    os.makedirs("logs")
file_handler = logging.FileHandler(
    os.path.join("logs", "normal-{:}.log".format(datetime_str)), encoding="utf-8"
)
debug_handler = logging.FileHandler(
    os.path.join("logs", "debug-{:}.log".format(datetime_str)), encoding="utf-8"
)
stdout_handler = logging.StreamHandler(sys.stdout)
sdebug_handler = logging.FileHandler(
    os.path.join("logs", "sdebug-{:}.log".format(datetime_str)), encoding="utf-8"
)

file_handler.setLevel(logging.INFO)
debug_handler.setLevel(logging.DEBUG)
stdout_handler.setLevel(logging.INFO)
sdebug_handler.setLevel(logging.DEBUG)

formatter = logging.Formatter(
    fmt="\x1b[1;33m[%(asctime)s \x1b[31m%(levelname)s \x1b[32m%(module)s/%(lineno)d-%(processName)s\x1b[1;33m] \x1b[0m%(message)s"
)
file_handler.setFormatter(formatter)
debug_handler.setFormatter(formatter)
stdout_handler.setFormatter(formatter)
sdebug_handler.setFormatter(formatter)

stdout_handler.addFilter(logging.Filter("desktopenv"))
sdebug_handler.addFilter(logging.Filter("desktopenv"))

logger.addHandler(file_handler)
logger.addHandler(debug_handler)
logger.addHandler(stdout_handler)
logger.addHandler(sdebug_handler)
#  }}} Logger Configs #

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


def config() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run end-to-end evaluation on the benchmark"
    )

    # environment config
    parser.add_argument("--provider_name", type=str, default="aws")
    parser.add_argument("--region", type=str, default="us-east-2")
    parser.add_argument("--path_to_vm", type=str, default=None)
    parser.add_argument(
        "--headless", action="store_true", help="Run in headless machine"
    )
    parser.add_argument(
        "--action_space", type=str, default="pyautogui", help="Action type"
    )
    parser.add_argument(
        "--observation_type",
        choices=["screenshot", "a11y_tree", "screenshot_a11y_tree", "som"],
        default="screenshot",
        help="Observation type",
    )
    parser.add_argument("--screen_width", type=int, default=1920)
    parser.add_argument("--screen_height", type=int, default=1080)
    parser.add_argument("--sleep_after_execution", type=float, default=2.0)
    parser.add_argument("--max_steps", type=int, default=50)
    parser.add_argument("--aws_ami", type=str, default=None)

    # agent config
    parser.add_argument("--max_trajectory_length", type=int, default=15)
    parser.add_argument(
        "--test_config_base_dir", type=str, default="forum"
    )

    # lm config
    parser.add_argument("--model", type=str, default="gpt-4o")
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top_p", type=float, default=0.9)
    parser.add_argument("--max_tokens", type=int, default=4096)
    parser.add_argument("--stop_token", type=str, default=None)

    # example config
    parser.add_argument("--domain", type=str, default="all")
    parser.add_argument(
        "--test_all_meta_path", type=str, default="forum/test_all_reddit.json"
    )

    # logging related
    parser.add_argument("--result_dir", type=str, default="./results")


    # new args
    parser.add_argument("--snapshot_name", type=str, default="init_state")
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--disable-per-step-evaluator", action="store_true", help="Disable per-step evaluator to prevent Chrome focus stealing")

    # agent
    parser.add_argument("--agent_type", choices=["KimiAgent"], default="KimiAgent", help="Type of agent to use")

    # parallelism
    parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to run in parallel")

    # watchdog: per-task wall-clock timeout. If a single task hangs longer than this
    # (e.g. end_recording retry loop), a watchdog thread force-terminates the AWS VM
    # and hard-exits this env process so it stops incurring cloud cost.
    parser.add_argument("--task_timeout_sec", type=int, default=40 * 60,
                        help="Per-task timeout in seconds (default 2400 = 40 min)")

    args = parser.parse_args()

    if args.debug:
        logger.info("Debug mode")
        args.result_dir = os.path.join(args.result_dir, "debug_{uuid}".format(uuid=datetime_str))
        logger.info("result_dir: %s", args.result_dir)

    args.result_model_name = args.model

    # fix platform
    if args.aws_ami is not None:
        # fake patch, it's a bit hacky..
        import platform
        platform.machine = lambda: "x86_64"

    return args


def distribute_tasks(test_all_meta: dict) -> List[tuple]:
    """Flatten domain->example_ids dict into a list of (domain, example_id) tuples."""
    all_tasks = []
    for domain, examples in test_all_meta.items():
        for example_id in examples:
            all_tasks.append((domain, example_id))
    return all_tasks


def run_env_tasks(task_queue, args: argparse.Namespace, shared_scores: list):
    """Worker function: each process creates its own DesktopEnv + KimiAgent and pulls tasks from queue."""
    env = None
    try:
        env = DesktopEnv(
            provider_name=args.provider_name,
            region=args.region,
            path_to_vm=args.path_to_vm,
            action_space=args.action_space,
            snapshot_name=args.snapshot_name,
            screen_size=(args.screen_width, args.screen_height),
            headless=args.headless,
            os_type="Ubuntu",
            require_a11y_tree=args.observation_type in ["a11y_tree", "screenshot_a11y_tree", "som"],
            aws_ami=args.aws_ami,
        )

        logger.info(f"Process {current_process().name} started with env.")

        while True:
            try:
                item = task_queue.get(timeout=5)
            except Exception:
                # Queue empty or timeout
                break

            domain, example_id = item
            try:
                config_file = os.path.join(
                    args.test_config_base_dir, f"examples/{domain}/{example_id}.json"
                )
                with open(config_file, "r", encoding="utf-8") as f:
                    example = json.load(f)

                logger.info(f"[{current_process().name}][Domain]: {domain}")
                logger.info(f"[{current_process().name}][Example ID]: {example_id}")
                logger.info(f"[{current_process().name}][Instruction]: {example['instruction']}")

                example_result_dir = os.path.join(
                    args.result_dir,
                    args.action_space,
                    args.observation_type,
                    args.result_model_name,
                    domain,
                    example_id,
                )
                os.makedirs(example_result_dir, exist_ok=True)

                # Determine API provider
                if args.model.split(' | ')[0] == "Azure":
                    agent_api_provider = "Azure"
                else:
                    agent_api_provider = "Kimi"

                agent = KimiAgent(
                    model="kimi-k2.5",
                    max_tokens=4096,
                    top_p=0.95,
                    temperature=1,
                    action_space="pyautogui",
                    password="password",
                    observation_type="screenshot",
                    max_image_history_length=3,
                    max_steps=args.max_steps,
                    provider=agent_api_provider,
                )
                agent_type = args.agent_type

                def _watchdog_force_terminate():
                    logger.error(
                        f"[WATCHDOG] {current_process().name} task {domain}/{example_id} "
                        f"exceeded {args.task_timeout_sec}s. Force-terminating AWS VM and exiting env process."
                    )
                    try:
                        if env is not None and env.provider_name == "aws":
                            env.terminate()
                    except Exception as te:
                        logger.error(f"[WATCHDOG] terminate failed: {te}")
                    os._exit(137)

                watchdog = threading.Timer(args.task_timeout_sec, _watchdog_force_terminate)
                watchdog.daemon = True
                watchdog.start()

                try:
                    lib_run_single.run_single_example_kimi(
                        agent,
                        agent_type,
                        env,
                        example,
                        args.max_steps,
                        example["instruction"],
                        args,
                        example_result_dir,
                        shared_scores,
                    )
                except Exception as e:
                    logger.error(f"Exception in {current_process().name} {domain}/{example_id}: {e}")
                    logger.error(traceback.format_exc())

                    try:
                        env.controller.end_recording(
                            os.path.join(example_result_dir, "recording.mp4")
                        )
                    except Exception as rec_e:
                        logger.error(f"Failed to end recording: {rec_e}")

                    with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
                        f.write(
                            json.dumps(
                                {"Error": f"{e} --- {domain}/{example_id}"}
                            )
                        )
                        f.write("\n")
                finally:
                    watchdog.cancel()
                    if env.provider_name == "aws":
                        try:
                            env.terminate()
                        except Exception as te:
                            logger.error(f"Failed to terminate VM for {domain}/{example_id}: {te}")

            except Exception as e:
                logger.error(f"Task-level error in {current_process().name}: {e}")
                logger.error(traceback.format_exc())

    except Exception as e:
        logger.error(f"Process-level error in {current_process().name}: {e}")
        logger.error(traceback.format_exc())
    finally:
        logger.info(f"{current_process().name} cleaning up environment...")
        try:
            if env:
                if env.provider_name == "aws":
                    env.terminate()
                else:
                    env.close()
                logger.info(f"{current_process().name} environment closed successfully")
        except Exception as e:
            logger.error(f"{current_process().name} error during environment cleanup: {e}")


def signal_handler(signum, frame):
    """Handle termination signals (SIGINT, SIGTERM) to gracefully shutdown environments."""
    global is_terminating, active_environments, processes

    if is_terminating:
        return

    is_terminating = True
    logger.info(f"Received signal {signum}. Gracefully shutting down...")

    for env in active_environments:
        try:
            logger.info("Closing environment...")
            env.close()
            logger.info("Environment closed successfully")
        except Exception as e:
            logger.error(f"Error closing environment: {e}")

    for p in processes:
        if p.is_alive():
            try:
                logger.info(f"Sending termination signal to process {p.name}...")
                p.terminate()
            except Exception as e:
                logger.error(f"Error sending termination signal to process: {e}")

    time.sleep(1)

    for p in processes:
        if p.is_alive():
            try:
                logger.info(f"Forcefully terminating process {p.name}...")
                os.kill(p.pid, signal.SIGKILL)
            except Exception as e:
                logger.error(f"Error forcefully terminating process: {e}")

    logger.info("Shutdown complete. Exiting.")
    sys.exit(0)


def test(args: argparse.Namespace, test_all_meta: dict) -> None:
    global processes

    logger.info("Args: %s", args)
    assert args.action_space in ["pyautogui"], "we only support pyautogui now since we will do the second call to convert the response into pyautogui actions"

    all_tasks = distribute_tasks(test_all_meta)
    logger.info(f"Total tasks: {len(all_tasks)}")

    num_envs = args.num_envs

    with Manager() as manager:
        shared_scores = manager.list()
        task_queue = manager.Queue()
        for item in all_tasks:
            task_queue.put(item)

        processes = []
        for i in range(num_envs):
            p = Process(
                target=run_env_tasks,
                args=(task_queue, args, shared_scores),
                name=f"EnvProcess-{i+1}"
            )
            p.daemon = True
            p.start()
            processes.append(p)
            logger.info(f"Started process {p.name} with PID {p.pid}")

        try:
            while True:
                alive_count = 0
                for idx, p in enumerate(processes):
                    if not p.is_alive():
                        if not task_queue.empty():
                            logger.warning(f"Process {p.name} died, restarting...")
                            new_p = Process(
                                target=run_env_tasks,
                                args=(task_queue, args, shared_scores),
                                name=f"EnvProcess-Restart-{idx+1}"
                            )
                            new_p.daemon = True
                            new_p.start()
                            processes[idx] = new_p
                            logger.info(f"Restarted process {new_p.name} with PID {new_p.pid}")
                    else:
                        alive_count += 1

                if task_queue.empty():
                    logger.info("All tasks dispatched. Waiting for processes to finish...")
                    break
                if alive_count == 0:
                    logger.error("All processes died, exiting.")
                    break
                time.sleep(5)

            for p in processes:
                p.join()

        except KeyboardInterrupt:
            logger.info("Main process received KeyboardInterrupt. Initiating graceful shutdown...")
            raise
        except Exception as e:
            logger.error(f"Unexpected error while waiting for processes: {e}", exc_info=True)
            for p in processes:
                if p.is_alive():
                    try:
                        logger.info(f"Terminating process {p.name} due to error...")
                        p.terminate()
                    except Exception as term_e:
                        logger.error(f"Error terminating process {p.name}: {term_e}")
            raise

        scores = list(shared_scores)

    if scores:
        logger.info(f"Average score: {sum(scores) / len(scores)}")
    else:
        logger.info("No scores collected.")


def get_unfinished(
    action_space, use_model, observation_type, result_dir, total_file_json
):
    target_dir = os.path.join(result_dir, action_space, observation_type, use_model)

    if not os.path.exists(target_dir):
        return total_file_json

    finished = {}
    for domain in os.listdir(target_dir):
        finished[domain] = []
        domain_path = os.path.join(target_dir, domain)
        if os.path.isdir(domain_path):
            for example_id in os.listdir(domain_path):
                if example_id == "onboard":
                    continue
                example_path = os.path.join(domain_path, example_id)
                if os.path.isdir(example_path):
                    if "result.txt" not in os.listdir(example_path) or "adversary_result.txt" not in os.listdir(example_path):
                        # empty all files under example_id
                        for file in os.listdir(example_path):
                            os.remove(os.path.join(example_path, file))
                    else:
                        finished[domain].append(example_id)

    if not finished:
        return total_file_json

    for domain, examples in finished.items():
        if domain in total_file_json:
            total_file_json[domain] = [
                x for x in total_file_json[domain] if x not in examples
            ]

    return total_file_json


def get_result(action_space, use_model, observation_type, result_dir, total_file_json):
    target_dir = os.path.join(result_dir, action_space, observation_type, use_model)
    if not os.path.exists(target_dir):
        print("New experiment, no result yet.")
        return None

    all_result = []

    for domain in os.listdir(target_dir):
        domain_path = os.path.join(target_dir, domain)
        if os.path.isdir(domain_path):
            for example_id in os.listdir(domain_path):
                example_path = os.path.join(domain_path, example_id)
                if os.path.isdir(example_path):
                    if "result.txt" in os.listdir(example_path):
                        try:
                            all_result.append(
                                float(
                                    open(
                                        os.path.join(example_path, "result.txt"), "r"
                                    ).read()
                                )
                            )
                        except:
                            all_result.append(0.0)

    if not all_result:
        print("New experiment, no result yet.")
        return None
    else:
        print("Current Success Rate:", sum(all_result) / len(all_result) * 100, "%")
        return all_result


if __name__ == "__main__":
    ####### The complete version of the list of examples #######
    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    # Register signal handlers for graceful termination
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    try:
        args = config()

        # save args to json
        path_to_args = os.path.join(
            args.result_dir,
            args.action_space,
            args.observation_type,
            args.result_model_name,
            "args.json",
        )
        os.makedirs(os.path.dirname(path_to_args), exist_ok=True)
        with open(path_to_args, "w", encoding="utf-8") as f:
            json.dump(vars(args), f, indent=4)

        with open(args.test_all_meta_path, "r", encoding="utf-8") as f:
            test_all_meta = json.load(f)

        if args.domain != "all":
            test_all_meta = {args.domain: test_all_meta[args.domain]}

        test_file_list = get_unfinished(
            args.action_space,
            args.result_model_name,
            args.observation_type,
            args.result_dir,
            test_all_meta,
        )
        left_info = ""
        total_number = 0
        for domain in test_file_list:
            left_info += f"{domain}: {len(test_file_list[domain])}\n"
            total_number += len(test_file_list[domain])
        logger.info(f"Left tasks:\n{left_info}")

        if total_number == 0:
            sys.exit(0)

        get_result(
            args.action_space,
            args.result_model_name,
            args.observation_type,
            args.result_dir,
            test_all_meta,
        )
        test(args, test_file_list)

    except KeyboardInterrupt:
        logger.info("Main process received KeyboardInterrupt.")
    except Exception as e:
        logger.error(f"Unexpected error in main process: {e}", exc_info=True)
        signal_handler(signal.SIGTERM, None)
    finally:
        logger.info("Main process final cleanup...")
        for env in active_environments:
            if env is not None:
                try:
                    logger.info("Closing environment in final cleanup...")
                    env.close()
                    logger.info("Environment closed successfully in final cleanup")
                except Exception as e:
                    logger.error(f"Error during final environment cleanup: {e}")

        for p in processes:
            if p is not None and p.is_alive():
                try:
                    logger.info(f"Terminating process {p.name}...")
                    p.terminate()
                except Exception as e:
                    logger.error(f"Error terminating process: {e}")

        time.sleep(1)

        for p in processes:
            if p is not None and p.is_alive():
                try:
                    logger.info(f"Force killing process {p.name}...")
                    os.kill(p.pid, signal.SIGKILL)
                    logger.info(f"Process {p.name} force killed")
                except Exception as e:
                    logger.error(f"Error force killing process: {e}")