batch_test_generated_states.py

# !/usr/bin/env python3
"""
Batch testing script for generated states

Usage:
python batch_test_generated_states.py --provider openrouter --workers 2
python batch_test_generated_states.py --provider xhub --workers 1 --reward-type heuristic
"""

import subprocess
import os
import json
import time
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed
import queue
import threading

def should_retest_config(result):
    """
    Determine whether a configuration should be retested

    Args:
        result: Test result dictionary

    Returns:
        tuple: (should_retest: bool, reason: str)
    """
    return_code = result.get('return_code', -1)
    stderr = result.get('stderr', '')
    # stdout = result.get('stdout', '')
    
    # According to the new return code rules:
    # 0: Normal completion. Cube solved successfully or reached max_steps
    # 100: time limit exceeded
    # 200: token limit exceeded
    # -1: other errors

    # Case 1: return_code == 0 - Normal completion (successful cube solving or exceeded max_steps)
    if return_code == 0:
        return False, "Normal completion"

    # Case 2: return_code == 100 - Runtime exceeded limit
    if return_code == 100:
        return False, "Runtime exceeded limit"

    # Case 3: return_code == 200 - Token limit exceeded
    if return_code == 200:
        return False, "Token limit exceeded"

    # Case 4: All other error cases need retesting
    # Including return_code == -1 and other undefined return codes
    return True, f"Other error (return_code={return_code})"

def run_single_test(config, states_data, output_base, test_total, states_file, reward_type='no_reward', timeout=1800, provider='openrouter', max_input_tokens=2000000, max_output_tokens=500000, step_observation_callback=False, max_images=200):
    """Execute a single test task

    Args:
        config: (agent_type, model, scramble_level, state_index, max_steps, test_id, observation_type)
        states_data: State data dictionary
        output_base: Output base directory
        test_total: Total number of tests
        states_file: State file path
        reward_type: Reward type ('no_reward', 'sticker', 'heuristic', 'face')
        timeout: Timeout time (seconds)
        provider: API provider

    Returns:
        dict: Test result
    """
    agent_type, model, scramble_level, state_index, max_steps, test_id, observation_type = config

    print(f"\n{'='*60}")
    print(f"Test {test_id}/{test_total}: {agent_type} agent, {model}, scramble_{scramble_level}, state {state_index}, {observation_type}")
    print(f"{'='*60}")
    
    # Verify if state exists
    scramble_key = f"scramble_{scramble_level}"
    if scramble_key not in states_data:
        print(f"❌ Scramble level {scramble_level} not available")
        return {
            'test_id': test_id,
            'agent_type': agent_type,
            'model': model,
            'scramble_level': scramble_level,
            'state_index': state_index,
            'max_steps': max_steps,
            'observation_type': observation_type,
            'duration': 0,
            'return_code': -1,
            # 'stdout': '',
            "reward_type": reward_type,
            'stderr': f'Scramble level {scramble_level} not available',
            'total_steps': 0
        }
        
    states = states_data[scramble_key]["states"]
    if state_index >= len(states):
        print(f"❌ State index {state_index} out of range (0-{len(states)-1})")
        return {
            'test_id': test_id,
            'agent_type': agent_type,
            'model': model,
            'scramble_level': scramble_level,
            'state_index': state_index,
            'max_steps': max_steps,
            'observation_type': observation_type,
            'duration': 0,
            'return_code': -1,
            # 'stdout': '',
            "reward_type": reward_type,
            'stderr': f'State index {state_index} out of range (0-{len(states)-1})',
            'total_steps': 0
        }
    
    # Build output directory path - Create independent subdirectory for each state
    model_dir = os.path.join(output_base, model.replace('/', '_'))
    agent_dir = os.path.join(model_dir, agent_type)
    observation_dir = os.path.join(agent_dir, observation_type)
    scramble_dir = os.path.join(observation_dir, f"scramble_{scramble_level}")
    state_dir = os.path.join(scramble_dir, f"state_{state_index}")

    # Ensure directory exists
    os.makedirs(state_dir, exist_ok=True)
    
    # Build command
    cmd = [
        'python', '-m', 'agent.agent',
        '--provider', provider,  # Add provider parameter
        '--agent-type', agent_type,
        '--use-generated-states',
        '--states-file', states_file,  # Add state file path parameter
        '--scramble-moves', str(scramble_level),
        '--state-index', str(state_index),
        '--max-steps', str(max_steps),
        '--model', model,
        '--output-dir', state_dir,  # Use independent state directory
        '--timeout', str(timeout),   # Add timeout parameter
        '--max-input-tokens', str(max_input_tokens),  # Add input token limit
        '--max-output-tokens', str(max_output_tokens),  # Add output token limit
        '--observation-type', observation_type,  # Add observation type parameter
        '--reward-type', reward_type,
    ]

    # Conditionally add --step-observation-callback parameter
    if step_observation_callback:
        cmd.append('--step-observation-callback')
    
    print(f"Running command: {' '.join(cmd)}")
    # exit(-1)
    
    # Record start time
    start_time = time.time()

    try:
        # Run test - No longer need subprocess timeout since agent.py handles it internally
        result = subprocess.run(cmd, capture_output=True, text=True)

        # Record results
        end_time = time.time()
        duration = end_time - start_time

        # Try to read session_summary.json to get actual execution steps
        total_steps = 0
        session_summary_path = os.path.join(state_dir, 'session_summary.json')
        if os.path.exists(session_summary_path):
            try:
                with open(session_summary_path, 'r', encoding='utf-8') as f:
                    session_data = json.load(f)
                total_steps = session_data.get('session_info', {}).get('dialogue_rounds', 0)
                print(f"📊 Extracted total_steps: {total_steps}")
            except Exception as e:
                print(f"⚠️  Failed to read session_summary.json: {e}")
                total_steps = 0
        else:
            print("⚠️  session_summary.json not found")

        test_result = {
            'test_id': test_id,
            'agent_type': agent_type,
            'model': model,
            'scramble_level': scramble_level,
            'state_index': state_index,
            'max_steps': max_steps,
            'observation_type': observation_type,
            'duration': duration,
            'return_code': result.returncode,
            # 'stdout': result.stdout,
            "reward_type": reward_type,
            'stderr': result.stderr,
            'total_steps': total_steps
        }
        
        if result.returncode == 0:
            print(f"✅ Test {test_id} completed successfully in {duration:.2f}s")
        elif result.returncode == 100:
            print(f"⏰ Test {test_id} terminated due to TIMEOUT in {duration:.2f}s")
        elif result.returncode == 200:
            print(f"🪙 Test {test_id} terminated due to TOKEN LIMIT in {duration:.2f}s")
        else:
            print(f"❌ Test {test_id} failed with return code {result.returncode}")
            print(f"Error: {result.stderr}")
        
        return test_result
        
    except Exception as e:
        print(f"❌ Test {test_id} failed with exception: {e}")
        return {
            'test_id': test_id,
            'agent_type': agent_type,
            'model': model,
            'scramble_level': scramble_level,
            'state_index': state_index,
            'max_steps': max_steps,
            'observation_type': observation_type,
            'duration': time.time() - start_time,
            'return_code': -1,
            # 'stdout': '',
            "reward_type": reward_type,
            'stderr': str(e),
            'total_steps': 0
        }


def run_batch_test():
    """Batch test pre-generated states"""
    
    # Load state data to verify availability
    import argparse

    parser = argparse.ArgumentParser(description="Batch test pre-generated states")
    parser.add_argument('--config', type=str, help='Configuration file path (JSON format)')
    parser.add_argument('--states-file', type=str, default='generated_states/initial_states.json', help='State file path')
    parser.add_argument('--resume', type=str, help='Resume experiment from specified result file')
    parser.add_argument('--checkpoint-interval', type=int, default=6, help='Save checkpoint every N tests')
    parser.add_argument('--workers', type=int, default=1, help='Number of parallel worker threads (1=serial execution)')
    parser.add_argument('--parallel-strategy', type=str, choices=['model', 'all'], default='all',
                       help='Parallel strategy: model=parallel by model, all=all tasks parallel')
    parser.add_argument('--reward-type', choices=['no_reward', 'sticker', 'heuristic', 'face'], default='no_reward', help='Reward type: no_reward(no reward), sticker(sticker solving metric), heuristic(heuristic solving metric), face(face solving metric)')
    parser.add_argument('--timeout', type=int, default=1800, help='Timeout time (seconds), default 30 minutes')
    parser.add_argument('--provider', type=str, default='openrouter',
                       help='Specify API provider (default: openrouter)')
    parser.add_argument('--max-input-tokens', type=int, default=2000000, help='Maximum input token limit, terminate conversation if exceeded')
    parser.add_argument('--max-output-tokens', type=int, default=500000, help='Maximum output token limit, terminate conversation if exceeded')
    parser.add_argument('--step-observation-callback', action='store_true', help='Enable step observation callback to provide current observation results to the agent')
    args = parser.parse_args()

    # Load configuration file (if provided)
    config_data = {}
    if args.config:
        print(f"Loading configuration from: {args.config}")
        try:
            with open(args.config, 'r', encoding='utf-8') as f:
                config_data = json.load(f)
            print("✅ Configuration loaded successfully")
        except Exception as e:
            print(f"❌ Failed to load configuration file: {e}")
            print("Using default values and command line arguments")
    
    # Load command line parameters from config file (if present)
    if args.config and 'experiment_parameters' in config_data:
        exp_params = config_data['experiment_parameters']
        # Override command line parameters (if present in config file)
        if 'states_file' in exp_params:
            args.states_file = exp_params['states_file']
        if 'reward_type' in exp_params:
            args.reward_type = exp_params['reward_type']
        if 'timeout' in exp_params:
            args.timeout = exp_params['timeout']
        if 'provider' in exp_params:
            args.provider = exp_params['provider']
        if 'max_input_tokens' in exp_params:
            args.max_input_tokens = exp_params['max_input_tokens']
        if 'max_output_tokens' in exp_params:
            args.max_output_tokens = exp_params['max_output_tokens']
        if 'workers' in exp_params:
            args.workers = exp_params['workers']
        if 'parallel_strategy' in exp_params:
            args.parallel_strategy = exp_params['parallel_strategy']
        if 'checkpoint_interval' in exp_params:
            args.checkpoint_interval = exp_params['checkpoint_interval']
        if 'step_observation_callback' in exp_params:
            args.step_observation_callback = exp_params['step_observation_callback']
        if 'max_images' in exp_params:
            args.max_images = exp_params['max_images']
        print("✅ Command line parameters updated from config file")
    
    # Load scramble configuration
    if args.config and 'scrambles' in config_data:
        # Convert string keys from JSON to integer keys
        scrambles_raw = config_data['scrambles']
        scrambles = {}
        for key, value in scrambles_raw.items():
            try:
                scrambles[int(key)] = value
            except ValueError:
                print(f"⚠️ Warning: Invalid scramble key '{key}', skipping...")
        print(f"Loaded scramble configs from file: {scrambles}")
    else:
        # Use hardcoded default values
        scrambles = {
            1: 20,   # scramble_moves=1 -> max_steps=30
            2: 20,   # scramble_moves=2 -> max_steps=30
            3: 20,   # scramble_moves=3 -> max_steps=30
            4: 20,
        }
        print(f"Using default scramble configs: {scrambles}")
    
    # Load agent types
    if args.config and 'agent_types' in config_data:
        agent_types = config_data['agent_types']
        print(f"Loaded agent types from file: {agent_types}")
    else:
        # Use hardcoded default values
        agent_types = ['basic', 'standard_solver', 'ideal_solver']
        print(f"Using default agent types: {agent_types}")

    # Load observation types - must be read from config file
    if args.config and 'observation_types' in config_data:
        observation_types = config_data['observation_types']
        print(f"Loaded observation types from file: {observation_types}")
    else:
        print("❌ Error: observation_types must be specified in the config file")
        print("Please provide a config file with 'observation_types' field")
        return
    
    # Load models
    if args.config and 'models' in config_data:
        models = config_data['models']
        print(f"Loaded models from file: {models}")
    else:
        # Use hardcoded default values
        models = ['qwen/qwen2.5-vl-72b-instruct']
        print(f"Using default models: {models}")

    with open(args.states_file, 'r', encoding='utf-8') as f:
        states_data = json.load(f)
    
    # Dynamically determine state count for each scramble level
    scramble_state_counts = {}
    for scramble_key, data in states_data.items():
        scramble_level = int(scramble_key.split('_')[1])
        scramble_state_counts[scramble_level] = data['num_states']

    print(f"Detected state counts per scramble level:")
    for level, count in sorted(scramble_state_counts.items()):
        print(f"  scramble_{level}: {count} states")
    
    # Generate test configurations: for each agent type, model, observation type, scramble_moves, state_index
    test_configs = []
    test_id = 1
    for agent_type in agent_types:
        for observation_type in observation_types:
            for scramble_moves, max_steps in scrambles.items():
                if scramble_moves in scramble_state_counts:
                    num_states = scramble_state_counts[scramble_moves]
                    for state_index in range(num_states):
                        for model in models:
                            test_configs.append((agent_type, model, scramble_moves, state_index, max_steps, test_id, observation_type))
                            test_id += 1
    
    print(f"Total test configurations: {len(test_configs)}")
    print(f"Agent types: {agent_types}")
    print(f"Models: {models}")
    print(f"Observation types: {observation_types}")
    print(f"Scramble moves: {list(scrambles.keys())}")
    print(f"States per scramble level: {scramble_state_counts}")
    print(f"Tests per agent-model-observation combination: {len(test_configs) // (len(agent_types) * len(models) * len(observation_types))}")
    
    # Handle resume functionality
    results = []
    completed_configs = set()
    start_index = 0
    test_id_to_index = {}  # Map test_id to index in results
    
    if args.resume:
        print(f"Resuming from checkpoint: {args.resume}")
        try:
            with open(args.resume, 'r', encoding='utf-8') as f:
                checkpoint_data = json.load(f)
            
            results = checkpoint_data.get('results', [])
            # Establish mapping from test_id to index
            test_id_to_index = {r['test_id']: i for i, r in enumerate(results)}
            # Record successfully completed configurations
            successful_configs = set()
            failed_configs = []

            for result in results:
                config_key = (result['agent_type'], result['model'],
                             result['scramble_level'], result['state_index'], result.get('observation_type', 'state_string'))

                # Check if should retest - according to new retest logic
                should_retest, reason = should_retest_config(result)

                if should_retest:
                    failed_configs.append((config_key, reason, result.get('stderr', '')[:100]))
                else:
                    # Cases that don't need retesting: normal completion, token limit exceeded, runtime timeout
                    successful_configs.add(config_key)
                
            # Replace completed_configs with successful_configs
            completed_configs = successful_configs

            print(f"Loaded {len(results)} previous results")
            print(f"Successfully completed: {len(successful_configs)}")
            if failed_configs:
                print(f"Found {len(failed_configs)} failed tests that will be retried:")
                for config_key, reason, detail in failed_configs:
                    print(f"  - {config_key}: {reason} - {detail}")
            print(f"Remaining tests: {len(test_configs) - len(successful_configs)}")
            
        except Exception as e:
            print(f"Failed to load checkpoint: {e}")
            print("Starting fresh...")
    
    print("Starting batch test...")
    print(f"Total test configurations: {len(test_configs)}")
    print(f"Already completed: {len(results)}")
    print(f"Remaining: {len(test_configs) - len(results)}")
    
    # Create output directory
    if args.resume:
        # Get output directory from checkpoint file path
        checkpoint_dir = os.path.dirname(args.resume)
        if os.path.basename(checkpoint_dir) == 'checkpoints':
            # If checkpoint is in checkpoints subdirectory, use its parent directory
            output_base = os.path.dirname(checkpoint_dir)
        else:
            # Otherwise create new output directory
            output_base = os.path.join("test_outputs", f"batch_tests_{datetime.now().strftime('%Y%m%d_%H%M%S')}")
    else:
        output_base = os.path.join("test_outputs", f"batch_tests_{datetime.now().strftime('%Y%m%d_%H%M%S')}")
    
    # Create checkpoints subdirectory under output directory
    checkpoints_dir = os.path.join(output_base, "checkpoints")
    os.makedirs(checkpoints_dir, exist_ok=True)
    
    # Save experiment configuration parameters
    def save_experiment_config():
        """Save experiment configuration parameters to file"""
        config_data = {
            'experiment_info': {
                'timestamp': datetime.now().isoformat(),
                'output_base': output_base,
                'checkpoints_dir': checkpoints_dir
            },
            'experiment_parameters': {
                'scrambles': scrambles,
                'agent_types': agent_types,
                'models': models,
                'reward_type': args.reward_type,
                'timeout': args.timeout,
                'provider': args.provider,
                'max_input_tokens': args.max_input_tokens,
                'max_output_tokens': args.max_output_tokens,
                'workers': args.workers,
                'parallel_strategy': args.parallel_strategy,
                'checkpoint_interval': args.checkpoint_interval,
                'states_file': args.states_file,
                'observation_types': observation_types,
                'step_observation_callback': args.step_observation_callback
            },
            'test_configuration': {
                'total_test_configs': len(test_configs),
                'scramble_state_counts': scramble_state_counts,
                'tests_per_agent_model_observation_combination': len(test_configs) // (len(agent_types) * len(models) * len(observation_types)) if len(agent_types) * len(models) * len(observation_types) > 0 else 0
            }
        }
        
        config_file = os.path.join(output_base, "experiment_config.json")
        with open(config_file, 'w', encoding='utf-8') as f:
            json.dump(config_data, f, indent=2, ensure_ascii=False)
        print(f"📋 Experiment configuration saved: {config_file}")

        # Also save a human-readable text version
        config_text_file = os.path.join(output_base, "experiment_config.txt")
        with open(config_text_file, 'w', encoding='utf-8') as f:
            f.write("Experiment Configuration Parameters\n")
            f.write("=" * 50 + "\n\n")
            
            f.write("Experiment Information:\n")
            f.write(f"  Timestamp: {config_data['experiment_info']['timestamp']}\n")
            f.write(f"  Output Directory: {config_data['experiment_info']['output_base']}\n")
            f.write(f"  Checkpoints Directory: {config_data['experiment_info']['checkpoints_dir']}\n\n")
            
            f.write("Experiment Parameters:\n")
            f.write(f"  Scramble Configs: {config_data['experiment_parameters']['scrambles']}\n")
            f.write(f"  Agent Types: {config_data['experiment_parameters']['agent_types']}\n")
            f.write(f"  Models: {config_data['experiment_parameters']['models']}\n")
            f.write(f"  Reward Type: {config_data['experiment_parameters']['reward_type']}\n")
            f.write(f"  Timeout: {config_data['experiment_parameters']['timeout']} seconds\n")
            f.write(f"  Provider: {config_data['experiment_parameters']['provider']}\n")
            f.write(f"  Max Input Tokens: {config_data['experiment_parameters']['max_input_tokens']:,}\n")
            f.write(f"  Max Output Tokens: {config_data['experiment_parameters']['max_output_tokens']:,}\n")
            f.write(f"  Workers: {config_data['experiment_parameters']['workers']}\n")
            f.write(f"  Parallel Strategy: {config_data['experiment_parameters']['parallel_strategy']}\n")
            f.write(f"  Checkpoint Interval: {config_data['experiment_parameters']['checkpoint_interval']}\n")
            f.write(f"  States File: {config_data['experiment_parameters']['states_file']}\n")
            f.write(f"  Observation Types: {config_data['experiment_parameters']['observation_types']}\n")
            f.write(f"  Step Observation Callback: {config_data['experiment_parameters']['step_observation_callback']}\n\n")
            
            f.write("Test Configuration:\n")
            f.write(f"  Total Test Configs: {config_data['test_configuration']['total_test_configs']}\n")
            f.write(f"  States per Scramble Level: {config_data['test_configuration']['scramble_state_counts']}\n")
            f.write(f"  Tests per Agent-Model-Observation Combination: {config_data['test_configuration']['tests_per_agent_model_observation_combination']}\n")
        
        print(f"📝 Human-readable config saved: {config_text_file}")
    
    # Save experiment configuration
    save_experiment_config()

    # Create checkpoint filename
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    checkpoint_file = os.path.join(checkpoints_dir, f"batch_test_checkpoint_{timestamp}.json")

    def save_checkpoint():
        """Save checkpoint"""
        checkpoint_data = {
            'timestamp': datetime.now().isoformat(),
            'total_configs': len(test_configs),
            'completed_configs': len(results),
            'results': results,
            'output_base': output_base
        }
        with open(checkpoint_file, 'w', encoding='utf-8') as f:
            json.dump(checkpoint_data, f, indent=2, ensure_ascii=False)
        print(f"💾 Checkpoint saved: {checkpoint_file}")
    
    # Filter out test configurations that need to run
    tasks_to_run = []
    for config in test_configs:
        agent_type, model, scramble_level, state_index, max_steps, test_id, observation_type = config
        config_key = (agent_type, model, scramble_level, state_index, observation_type)
        if config_key in completed_configs:
            print(f"⏭️  Skipping completed test {test_id}/{len(test_configs)}: {agent_type} + {model} + scramble_{scramble_level} + state_{state_index} + {observation_type}")
            continue
        tasks_to_run.append(config)

    print(f"Tasks to run: {len(tasks_to_run)}")

    if not tasks_to_run:
        print("No tasks to run!")
        return results
    
    # Thread-safe result collection
    results_lock = threading.Lock()

    def collect_result(future):
        """Callback function to collect individual test results"""
        try:
            result = future.result()
            with results_lock:
                test_id = result['test_id']
                
                # Check if need to replace old result
                if test_id in test_id_to_index:
                    # Replace old failed result
                    old_index = test_id_to_index[test_id]
                    old_result = results[old_index]

                    # Use new retest logic to judge old and new results
                    old_should_retest, old_reason = should_retest_config(old_result)
                    new_should_retest, new_reason = should_retest_config(result)

                    # If old result needs retesting but new result doesn't, replace
                    if old_should_retest and not new_should_retest:
                        print(f"✅ Replacing failed test {test_id} with successful result ({new_reason})")
                        results[old_index] = result
                    # If both need retesting, update failure info
                    elif old_should_retest and new_should_retest:
                        print(f"⚠️  Test {test_id} failed again, updating with new failure info ({new_reason})")
                        results[old_index] = result
                    # If old result doesn't need retesting, keep unchanged
                    else:
                        print(f"⏭️  Keeping previous result for test {test_id} ({old_reason})")
                else:
                    # New test result, add to list
                    results.append(result)
                    test_id_to_index[test_id] = len(results) - 1

                # Save checkpoint periodically
                if len(results) % args.checkpoint_interval == 0:
                    save_checkpoint()
        except Exception as e:
            print(f"❌ Failed to collect result: {e}")
    
    # Execute tests according to parallel strategy
    if args.parallel_strategy == 'model':
        # Execute in parallel by model grouping
        print(f"🚀 Using model-based parallelization with {args.workers} workers")

        # Group tasks by model
        model_tasks = {}
        for config in tasks_to_run:
            model = config[1]
            if model not in model_tasks:
                model_tasks[model] = []
            model_tasks[model].append(config)

        # Start a thread pool for each model
        for model, model_configs in model_tasks.items():
            print(f"\n🔄 Processing {len(model_configs)} tasks for model: {model}")

            with ThreadPoolExecutor(max_workers=min(args.workers, len(model_configs))) as executor:
                # Submit all tasks
                futures = []
                for config in model_configs:
                    future = executor.submit(run_single_test, config, states_data, output_base, len(test_configs), args.states_file, args.reward_type, args.timeout, args.provider, args.max_input_tokens, args.max_output_tokens, args.step_observation_callback, args.max_images)
                    future.add_done_callback(collect_result)
                    futures.append(future)

                # Wait for all tasks to complete
                for future in as_completed(futures):
                    pass  # Results are collected through callback function

    elif args.parallel_strategy == 'all':
        # Execute all tasks in parallel
        print(f"🚀 Using full parallelization with {args.workers} workers")
        
        with ThreadPoolExecutor(max_workers=min(args.workers, len(tasks_to_run))) as executor:
            # Submit all tasks
            futures = []
            for config in tasks_to_run:
                future = executor.submit(run_single_test, config, states_data, output_base, len(test_configs), args.states_file, args.reward_type, args.timeout, args.provider, args.max_input_tokens, args.max_output_tokens, args.step_observation_callback, args.max_images)
                future.add_done_callback(collect_result)
                futures.append(future)

            # Wait for all tasks to complete
            for future in as_completed(futures):
                pass  # Results are collected through callback function
    
    else:
        # Serial execution (workers=1 or other cases)
        print("🔄 Using serial execution")
        for config in tasks_to_run:
            result = run_single_test(config, states_data, output_base, len(test_configs), args.states_file, args.reward_type, args.timeout, args.provider, args.max_input_tokens, args.max_output_tokens, args.step_observation_callback, args.max_images)
            test_id = result['test_id']

            # Check if need to replace old result
            if test_id in test_id_to_index:
                # Replace old failed result
                old_index = test_id_to_index[test_id]
                old_result = results[old_index]

                # Use new retest logic to judge old and new results
                old_should_retest, old_reason = should_retest_config(old_result)
                new_should_retest, new_reason = should_retest_config(result)

                # If old result needs retesting but new result doesn't, replace
                if old_should_retest and not new_should_retest:
                    print(f"✅ Replacing failed test {test_id} with successful result ({new_reason})")
                    results[old_index] = result
                # If both need retesting, update failure info
                elif old_should_retest and new_should_retest:
                    print(f"⚠️  Test {test_id} failed again, updating with new failure info ({new_reason})")
                    results[old_index] = result
                # If old result doesn't need retesting, keep unchanged
                else:
                    print(f"⏭️  Keeping previous result for test {test_id} ({old_reason})")
            else:
                # New test result, add to list
                results.append(result)
                test_id_to_index[test_id] = len(results) - 1

            # Save checkpoint periodically
            if len(results) % args.checkpoint_interval == 0:
                save_checkpoint()
    
    # Save final checkpoint and results
    save_checkpoint()

    print(f"\n{'='*60}")
    print("BATCH TEST SUMMARY")
    print(f"{'='*60}")

    successful_tests = [r for r in results if r['return_code'] == 0]
    failed_tests = [r for r in results if r['return_code'] != 0]

    print(f"Total tests: {len(results)}")
    print(f"Successful: {len(successful_tests)}")
    print(f"Failed: {len(failed_tests)}")

    if successful_tests:
        avg_duration = sum(r['duration'] for r in successful_tests) / len(successful_tests)
        print(f"Average duration (successful): {avg_duration:.2f}s")

    # print(f"\nResults saved to: {results_file}")

    return results

if __name__ == "__main__":
    results = run_batch_test()