diff --git a/examples/deepswe/swe_agent.py b/examples/deepswe/swe_agent.py index 8a64272df1..6e1c3fe6cf 100644 --- a/examples/deepswe/swe_agent.py +++ b/examples/deepswe/swe_agent.py @@ -7,315 +7,22 @@ from absl import logging -SWE_SYSTEM_PROMPT_FN_CALL = """You are a programming agent who is provided a github issue and repository bash environment and is tasked to solve certain tasks (e.g., file localization, testcase generation, code repair and editing etc) to resolve the issue. -""" - -SWE_SYSTEM_PROMPT = """You are a programming agent who is provided a github issue and repository bash environment and is tasked to solve certain tasks (e.g., file localization, testcase generation, code repair and editing etc) to resolve the issue. - -We have access to the following functions: - -–– BEGIN FUNCTION #1: file_editor –– -Description: -Custom editing tool for viewing, creating and editing files - • State is persistent across command calls and discussions with the user - • If path is a file, view displays the result of applying cat -n. If path is a directory, view lists non-hidden files and directories up to 2 levels deep - • The create command cannot be used if the specified path already exists as a file - • If a command generates a long output, it will be truncated and marked with - • The undo_edit command will revert the last edit made to the file at path - -Notes for using the str_replace command: - • The old_str parameter should match EXACTLY one or more consecutive lines from the original file. Be mindful of whitespaces! - • If the old_str parameter is not unique in the file, the replacement will not be performed. Make sure to include enough context in old_str to make it unique - • The new_str parameter should contain the edited lines that should replace the old_str - -Parameters: - 1. command (string, required) -Allowed values: [view, create, str_replace, insert, undo_edit] -The command to run. - 2. path (string, required) -Absolute path to file or directory, e.g. /testbed/file.py or /testbed. - 3. file_text (string, optional) -Required for the create command. Contains the content of the file to be created. - 4. old_str (string, optional) -Required for the str_replace command. The exact string in path to replace. - 5. new_str (string, optional) - • Optional for the str_replace command to specify the replacement string. - • Required for the insert command to specify the string to insert. - 6. insert_line (integer, optional) -Required for the insert command. The new_str will be inserted after the line number specified here. - 7. view_range (array, optional) - • Optional for the view command (when path is a file). - • If provided, specifies the line range to view, e.g. [11, 12] shows lines 11 and 12. - • [start_line, -1] will show all lines from start_line to the end of file. - 8. concise (boolean, optional) - • Optional for the view command. - • Defaults to True; displays a concise skeletal view of the file. If set to False, displays the full content in the specified view_range. - -–– END FUNCTION #1 –– - -–– BEGIN FUNCTION #2: execute_bash –– -Description: -Execute a bash command in the terminal. - -Behavior notes: - • If a command may run indefinitely (long-running), consider running it in the background and redirecting output, e.g. python3 app.py > server.log 2>&1 &. - • If the bash command returns exit code -1, it means the process is still running. The assistant may: - • Call this function again with command as an empty string ("") to retrieve additional logs. - • Send more input to STDIN of the running process by calling this function again with command set to the text input. - • Send command="ctrl+c" to interrupt the currently running process. - • If the command times out, it will be interrupted (SIGINT). The assistant may then retry or do further steps if needed. - -Parameters: - 1. cmd (string, required) -The bash command (and optional arguments) to execute. - • Can be empty ("") to retrieve more logs if the process is still running. - • Can be "ctrl+c" to interrupt the running process. - -–– END FUNCTION #2 –– - -–– BEGIN FUNCTION #3: search –– -Description: -Search for a term in a directory or a single file. - • If path is a directory (or unspecified, default is .), it recursively searches all non-hidden files and directories for the search term. - • If path points to a file, it runs a grep -n in that file to show line numbers matching the search term. - • If more than 100 files match in a directory search, results are truncated and the tool will inform you to narrow your search. - • If no matches are found, it will inform you as well. - -Parameters: - 1. search_term (string, required) -The term or string to search for in files. - 2. path (string, optional) -The file or directory to search in. Defaults to . if not specified. - -–– END FUNCTION #3 –– - -–– BEGIN FUNCTION #4: finish –– -Description: -Finish the interaction once the task is complete or if no further progress can be made. - -Behavior notes: - • The submit command finalizes your output. - -Parameters: - 1. command (string, required) -Currently allowed value: [submit] - 2. result (string, optional) -The result text or final message to submit. Defaults to an empty string if not provided. - -–– END FUNCTION #4 –– - -If you choose to call a function ONLY reply in the following format with NO suffix: - - -value_1 - -This is the value for the second parameter -that can span -multiple lines - - - - -Reminder: -- Function calls MUST follow the specified format, start with -- Required parameters MUST be specified -- Only call one function at a time -- VERY IMPORTANT: Each response must include both reasoning (as natural text) and function call (in above format) to solve the task. - -""" - -SWEAGENT_SYSTEM_PROMPT = """You are a programming agent who is provided a github issue and repository bash environment and is tasked to solve certain tasks (e.g., file localization, testcase generation, code repair and editing etc) to resolve the issue. - -We have access to the following functions: - ----- BEGIN FUNCTION #1: execute_bash ---- -Description: Execute a bash command in the terminal. -Parameters: - (1) command (string, required): The bash command to execute. For example: `python my_script.py`. If not provided, will show help. ----- END FUNCTION #1 ---- - - ----- BEGIN FUNCTION #2: submit ---- -Description: Finish the interaction when the task is complete OR if the assistant cannot proceed further with the task. -No parameters are required for this function. ----- END FUNCTION #2 ---- - - ----- BEGIN FUNCTION #3: str_replace_editor ---- -Description: Custom editing tool for viewing, creating and editing files -* State is persistent across command calls and discussions with the user -* If `path` is a file, `view` displays the result of applying `cat -n`. If `path` is a directory, `view` lists non-hidden files and directories up to 2 levels deep -* The `create` command cannot be used if the specified `path` already exists as a file -* If a `command` generates a long output, it will be truncated and marked with `` -Notes for using the `str_replace` command: -* The `old_str` parameter should match EXACTLY one or more consecutive lines from the original file. Be mindful of whitespaces! -* If the `old_str` parameter is not unique in the file, the replacement will not be performed. Make sure to include enough context in `old_str` to make it unique -* The `new_str` parameter should contain the edited lines that should replace the `old_str` -Parameters: - (1) command (string, required): The commands to run. Allowed options are: `view`, `create`, `str_replace`, `insert`. -Allowed values: [`view`, `create`, `str_replace`, `insert`] - (2) path (string, required): Absolute path to file or directory, e.g. `/repo/file.py` or `/repo`. - (3) file_text (string, optional): Required parameter of `create` command, with the content of the file to be created. - (4) old_str (string, optional): Required parameter of `str_replace` command containing the string in `path` to replace. - (5) new_str (string, optional): Optional parameter of `str_replace` command containing the new string (if not given, no string will be added). Required parameter of `insert` command containing the string to insert. - (6) insert_line (integer, optional): Required parameter of `insert` command. The `new_str` will be inserted AFTER the line `insert_line` of `path`. - (7) view_range (array, optional): Optional parameter of `view` command when `path` points to a file. If none is given, the full file is shown. If provided, the file will be shown in the indicated line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to start. Setting `[start_line, -1]` shows all lines from `start_line` to the end of the file. ----- END FUNCTION #3 ---- - - -If you choose to call a function ONLY reply in the following format with NO suffix: - -Provide any reasoning for the function call here. - -value_1 - -This is the value for the second parameter -that can span -multiple lines - - - - -Reminder: -- Function calls MUST follow the specified format, start with -- Required parameters MUST be specified -- Only call one function at a time -- Always provide reasoning for your function call in natural language BEFORE the function call (not after) -""" - - -SWE_USER_PROMPT_FN_CALL = """Consider the following github issue: - -{problem_statement} - - -Can you help me implement the necessary changes to the repository to fix the ? -I've already taken care of all changes to any of the test files described in the . This means you DON'T have to modify the testing logic or any of the tests in any way! -Your task is to make the minimal changes to non-tests files in the /testbed directory to ensure the is satisfied. - -IMPORTANT TIP: -Follow these steps to resolve the issue: -1. As a first step, it might be a good idea to explore the repo to familiarize yourself with its structure. -2. Create a script ('reproduce_issue.py') to reproduce the error and execute it to confirm the error - 2.1 reproduce_issue.py script finishes quickly after checking the error, fix etc. There no long running background servers for django for instance etc. It should be a quick script which checks the error and fix to provide a visible response. - 2.2 SUPER IMPORTANT: to ensure this reproduce_script.py must have a timeout logic of 20 seconds. If the script runs for more than 30 seconds, it should output a timeout message and you can interpret accordingly. -3. Edit the sourcecode of the repo to resolve the issue -4. Rerun your reproduce script and confirm that the error is fixed! -5. Think about edgecases and make sure your fix handles them as well - -VERY IMPORTANT: each response must include both reasoning and function call to solve the task. -You are being told a million times, each response must include a function call. Must inlcude a function call at all costs. - -You can take multiple turns to solve the task. So please only finish / submit when you are confident in your response. Dont rush. Be comprehensive. -You are being told a million times, please dont just submit without proper reasoning. Try to fully analyse the problem statement, explore the repository, reproduce the issue, fix it, check edge cases and then submit. - -Your thinking should be thorough and so it's fine if it's very long. -VERY IMPORTANT: file_editor old_str and new_str must be w/o the line numbers. line numbers are only shown in the view for clarity. - -Also if a file_editor edit fails, its a good idea to view the file near the edit location before trying to edit again. Dont keep trying the same edit over and over again. It will keep leading to the same failure. -Again do not get stuck trying to do the same thing over and over again. Please be efficient. -""" - -SWE_USER_PROMPT = """Consider the following github issue: - -{problem_statement} - - -Can you help me implement the necessary changes to the repository to fix the ? -I've already taken care of all changes to any of the test files described in the . This means you DON'T have to modify the testing logic or any of the tests in any way! -Your task is to make the minimal changes to non-tests files in the /testbed directory to ensure the is satisfied. - -IMPORTANT TIP: -Follow these steps to resolve the issue: -1. As a first step, it might be a good idea to explore the repo to familiarize yourself with its structure. -2. Create a script ('reproduce_issue.py') to reproduce the error and execute it to confirm the error -3. Edit the sourcecode of the repo to resolve the issue -4. Rerun your reproduce script and confirm that the error is fixed! -5. Think about edgecases and make sure your fix handles them as well -6. When viewing large files, use specific line-ranges, usually within 50 to 100 lines) as required -7. NOTE: The repository is at '/testbed' and the current working directory is already '/testbed', so DO NOT include 'testbed/' or 'testbed.' in relative paths in bash commands or reproduction python files. -""" - -SWEAGENT_USER_PROMPT = """I have uploaded a python code repository in the /testbed directory. - -Now consider the following Github issue: - - -{problem_statement} - - -Can you help me implement the necessary changes to the repository to fix the ? -I have already taken care of all changes to any of the test files described in the . This means you DON'T have to modify the testing logic or any of the tests in any way! Your task is to make changes to non-test files in the /testbed directory to ensure the is resolved. - -Follow these steps to resolve the issue: -1. First, explore the codebase to locate and understand the code relevant to the . - - Use efficient search commands to identify key files and functions. - - You should err on the side of caution and look at various relevant files and build your understanding of - - how the code works - - what are the expected behaviors and edge cases - - what are the potential root causes for the given issue - -2. Assess whether you can reproduce the issue: - - Create a script at '/testbed/reproduce_issue.py' that demonstrates the error. - - Execute this script to confirm the error behavior. - - You should reproduce the issue before fixing it. - - Your reproduction script should also assert the expected behavior for the fixed code. - -3. Analyze the root cause: - - Identify the underlying problem based on your code exploration and reproduction results. - - Critically analyze different potential approaches to fix the issue. - - You NEED to explicitly reason about multiple approaches to fix the issue. Next, find the most elegant and effective solution among them considering the tradeoffs (correctness, generality, side effects, etc.). - - You would need to reason about execution paths, edge cases, and other potential issues. You should look at the unit tests to understand the expected behavior of the relevant code. - -4. Implement your solution: - - Make targeted changes to the necessary files following idiomatic code patterns once you determine the root cause. - - You should be thorough and methodical. - -5. Verify your solution: - - Rerun your reproduction script to confirm the error is fixed. - - If verification fails, iterate on your solution until successful. If you identify the reproduction script is buggy, adjust it as needed. - -6. Run unit tests: - - Find and run the relevant unit tests relevant to the performed fix. - - You should run the unit tests to ensure your solution is correct and does not cause any regressions. - - In cases where the unit tests are do not pass, you should consider whether the unit tests does not reflect the *new* expected behavior of the code. If so, you can test it by writing additional edge test cases. - - Use the existing test runner to run the unit tests you identify as relevant to the changes you made. For example: - - `python -m pytest -xvs sympy/physics/units/tests/test_dimensions_transcendental.py` - - `python -m pytest tests/test_domain_py.py::test_pymethod_options` - - `./tests/runtests.py constraints.tests.CheckConstraintTests -v 2` - - RUN ALL relevant unit tests to ensure your solution is correct and does not cause any regressions. - -7. Test edge cases: - - Identify potential edge cases that might challenge your solution. - - Create additional test cases in a separate file '/testbed/edge_case_tests.py'. - - Execute these tests to verify your solution's robustness. - - You should run multiple rounds of edge cases. When creating edge cases: - - Consider complex scenarios beyond the original issue description - - Test for regressions to ensure existing functionality remains intact - -8. Refine if necessary: - - If edge case testing reveals issues, refine your solution accordingly. - - Ensure your final implementation handles all identified scenarios correctly. - - Document any assumptions or limitations of your solution. - -9. Submit your solution: - - Once you have verified your solution, submit your solution using the `submit` tool. - -A successful resolution means: -- The specific error/issue described no longer occurs -- Your changes maintain compatibility with existing functionality -- Edge cases are properly handled - - -Additional recommendations: -- You should be thorough, methodical, and prioritize quality over speed. Be comprehensive. -- You should think carefully before making the tool call about what should be done. However, each step should only use one tool call. YOU SHOULD NOT USE TOOLS INSIDE YOUR THOUGHT PROCESS. YOU SHOULD PRIMARILY USE THINKING FOR IDENTIFYING THE ROOT CAUSE OF THE ISSUE, MAKING THE CHANGES, AND CREATING TEST CASES (REPRODUCTION OR EDGE CASES). -- Each action you take is somewhat expensive. Wherever possible, combine multiple actions into a single action (e.g., combine multiple bash commands, use sed/grep for bulk operations). - - Your grep commands should identify both relevant files and line numbers so you can use the file_editor tool. - - Use grep with `-A -B -C` flags to quickly identify the relevant code blocks during your exploration. -- When exploring the codebase, use targeted search patterns to minimize unnecessary operations. -- When creating edge cases, you should look at the relevant existing tests to understand existing "regression" test cases. Ensure the fix doesn't break existing functionality. -""" +try: + from examples.deepswe import template +except ImportError: + import template # pytype: disable=import-error + +OPENHANDS_SYSTEM_PROMPT = template.OPENHANDS_SYSTEM_PROMPT +SWE_SYSTEM_PROMPT = template.SWE_SYSTEM_PROMPT +SWE_SYSTEM_PROMPT_FN_CALL = template.SWE_SYSTEM_PROMPT_FN_CALL +SWE_USER_PROMPT = template.SWE_USER_PROMPT +SWE_USER_PROMPT_FN_CALL = template.SWE_USER_PROMPT_FN_CALL +SWEAGENT_SYSTEM_PROMPT = template.SWEAGENT_SYSTEM_PROMPT +SWEAGENT_USER_PROMPT = template.SWEAGENT_USER_PROMPT +get_system_prompt = template.get_system_prompt +get_user_prompt_template = template.get_user_prompt_template + + from tunix.rl.agentic.agents.agent_types import Action from tunix.rl.agentic.agents.agent_types import Step from tunix.rl.agentic.agents.agent_types import Trajectory @@ -393,17 +100,15 @@ def __init__( assert scaffold in [ "r2egym", "sweagent", - ], f"Invalid scaffold: {scaffold}, must be one of ['r2egym', 'sweagent']" - system_prompt = ( - SWE_SYSTEM_PROMPT_FN_CALL if use_fn_calling else SWE_SYSTEM_PROMPT - ) - if scaffold == "sweagent": - system_prompt = SWEAGENT_SYSTEM_PROMPT - self.user_prompt_template = ( - SWE_USER_PROMPT_FN_CALL if use_fn_calling else SWE_USER_PROMPT + "openhands", + ], f"Invalid scaffold: {scaffold}, must be one of ['r2egym', 'sweagent', 'openhands']" + if system_prompt is None: + system_prompt = get_system_prompt( + scaffold=scaffold, use_fn_calling=use_fn_calling + ) + self.user_prompt_template = get_user_prompt_template( + scaffold=scaffold, use_fn_calling=use_fn_calling ) - if scaffold == "sweagent": - self.user_prompt_template = SWEAGENT_USER_PROMPT super().__init__(system_prompt) diff --git a/examples/deepswe/swe_env.py b/examples/deepswe/swe_env.py index 1827ad6744..3ca69b66a2 100644 --- a/examples/deepswe/swe_env.py +++ b/examples/deepswe/swe_env.py @@ -2,10 +2,17 @@ import json import logging import os +import re import threading +import time from typing import Any, Optional, cast import numpy as np +try: + from examples.deepswe import template as template_mod +except ImportError: + import template as template_mod # pytype: disable=import-error + _GLOBAL_FLEET = None _FLEET_LOCK = threading.Lock() _PATCH_LOCK = threading.Lock() @@ -79,26 +86,57 @@ def _patched_start_container( logging.debug("[SandboxFleet] r2egym in-memory patch note: %s", e) -def _normalize_tasks_for_fleet(tasks: Any) -> list[Any]: +def _get_image_rewrite_fn(image_rewrite: Any | None = None) -> Any | None: + """Retrieve or construct the image rewrite function from prefix if configured.""" + if image_rewrite is not None: + return image_rewrite + if os.getenv("IMAGE_REWRITE_PREFIX"): + prefix = os.environ["IMAGE_REWRITE_PREFIX"].rstrip("/") + return lambda img: f"{prefix}/{img.split('/')[-1]}" + return None + + +def _normalize_tasks_for_fleet( + tasks: Any, scaffold: str = "r2egym" +) -> list[Any]: """Normalize heterogeneous dataset entries into Task objects for SandboxFleet.""" + TaskCls = None try: from agent_sandbox_rl import Task # pytype: disable=import-error + if isinstance(Task, type): + TaskCls = Task except ImportError: - return list(tasks) + pass + + if TaskCls is None: + from dataclasses import dataclass + + @dataclass + class TaskCls: # pytype: disable=reimported + id: str + image: str + metadata: dict[str, Any] + agent_server_override = ( + os.getenv("AGENT_SERVER_IMAGE") if scaffold == "openhands" else None + ) normalized = [] for item in tasks: - if isinstance(item, Task): + if hasattr(item, "image") and hasattr(item, "id"): normalized.append(item) elif isinstance(item, dict): - img = item.get("docker_image") or item.get("image", "default") + img = ( + agent_server_override + or item.get("docker_image") + or item.get("image", "default") + ) if isinstance(img, (list, np.ndarray)): img = img[0] if len(img) > 0 else "default" t_id = item.get("instance_id") or item.get("id") or img if isinstance(t_id, (list, np.ndarray)): t_id = t_id[0] if len(t_id) > 0 else "default" normalized.append( - Task(id=str(t_id), image=str(img), metadata={"ds": item}) + TaskCls(id=str(t_id), image=str(img), metadata={"ds": item}) ) else: normalized.append(item) @@ -111,6 +149,8 @@ def _init_global_fleet( num_generations: int = 8, batch_size: int = 8, max_warmpool_replicas: int | None = None, + scaffold: str = "r2egym", + image_rewrite: Any | None = None, ) -> Any: """Initialize the process-wide SandboxFleet instance once upfront.""" global _GLOBAL_FLEET @@ -126,7 +166,6 @@ def _init_global_fleet( FleetConfig, SandboxFleet, Task, - TemplateSpec, ) except ImportError as e: raise ImportError( @@ -144,8 +183,17 @@ def _init_global_fleet( effective_max_concurrent = max( max_concurrency, batch_size * num_generations * 2 ) - fleet_cfg = FleetConfig( - clusters=[ + + template = template_mod.get_template(scaffold, node_sel) + + is_shared_image = scaffold == "openhands" and bool( + os.getenv("AGENT_SERVER_IMAGE") + ) + default_warmpool_size = ( + effective_max_concurrent if is_shared_image else num_generations + ) + fleet_kwargs = { + "clusters": [ ClusterConfig( name="default", namespace=fleet_ns, @@ -153,21 +201,38 @@ def _init_global_fleet( in_cluster=in_cluster, ) ], - max_concurrent=effective_max_concurrent, - window_size=batch_size, - max_warmpool_size=max_warmpool_replicas - if max_warmpool_replicas is not None - else num_generations, - warm_per_task=True, - ) + "max_concurrent": effective_max_concurrent, + "window_size": batch_size, + "max_warmpool_size": ( + max_warmpool_replicas + if max_warmpool_replicas is not None + else default_warmpool_size + ), + "warm_per_task": True, + } + if template is not None: + fleet_kwargs["template"] = template + if scaffold == "openhands": + fleet_kwargs["template_name_prefix"] = "oh-img-" + fleet_cfg = FleetConfig(**fleet_kwargs) fleet_inst = SandboxFleet(fleet_cfg) + image_rewrite_fn = _get_image_rewrite_fn(image_rewrite) + fleet_inst._image_rewrite_fn = image_rewrite_fn if tasks is not None: - fleet_inst.load_tasks(_normalize_tasks_for_fleet(tasks)) + normalized_tasks = _normalize_tasks_for_fleet(tasks, scaffold=scaffold) + should_rewrite = ( + image_rewrite_fn is not None + and (scaffold != "openhands" or not os.getenv("AGENT_SERVER_IMAGE")) + ) + if should_rewrite: + fleet_inst.load_tasks(normalized_tasks, image_rewrite=image_rewrite_fn) + else: + fleet_inst.load_tasks(normalized_tasks) msg = ( f"[SandboxFleet] Initializing pipelined fleet in namespace={fleet_ns}" f" (max_concurrent={effective_max_concurrent}," f" window_size={batch_size}," - f" max_warmpool_replicas={max_warmpool_replicas if max_warmpool_replicas is not None else num_generations}," + f" max_warmpool_replicas={fleet_kwargs['max_warmpool_size']}," " warm_per_task=True)..." ) logging.info(msg) @@ -201,12 +266,18 @@ def __init__( num_generations: int = 8, batch_size: int = 8, max_warmpool_replicas: int | None = None, + scaffold: str = "r2egym", + image_rewrite: Any | None = None, ): self.dataset_iter = iter(dataset) self.num_generations = num_generations self.batch_size = batch_size self.max_warmpool_replicas = max_warmpool_replicas + self.scaffold = scaffold self.fleet = fleet or _get_global_fleet() + self.image_rewrite = _get_image_rewrite_fn( + image_rewrite or getattr(self.fleet, "_image_rewrite_fn", None) + ) self.current_batch = None self.next_batch = None self.prev_batch_images: list[str] = [] @@ -230,6 +301,14 @@ def __init__( pass def _extract_images(self, batch: Any) -> list[str]: + agent_server_override = ( + os.getenv("AGENT_SERVER_IMAGE") + if self.scaffold == "openhands" + else None + ) + if agent_server_override: + return [agent_server_override] + raw_images = [] if isinstance(batch, dict) and "docker_image" in batch: raw = batch["docker_image"] @@ -246,19 +325,29 @@ def _extract_images(self, batch: Any) -> list[str]: if isinstance(item, dict) and item.get("docker_image") ] - # Safely decode/stringify all elements + # Safely decode/stringify all elements and apply rewrite if configured + rewrite_fn = getattr(self, "image_rewrite", None) or _get_image_rewrite_fn() str_images = [] for img in raw_images: - str_images.append( - img.decode("utf-8") if hasattr(img, "decode") else str(img) - ) + s = img.decode("utf-8") if hasattr(img, "decode") else str(img) + if rewrite_fn is not None: + s = rewrite_fn(s) + str_images.append(s) return list(dict.fromkeys(str_images)) def _warm_batch(self, batch: Any, wait: bool = False): images = self._extract_images(batch) if images and self.fleet: - target_replicas = self.max_warmpool_replicas or self.num_generations + is_shared_image = self.scaffold == "openhands" and bool( + os.getenv("AGENT_SERVER_IMAGE") + ) + default_replicas = ( + getattr(self.fleet.config, "max_concurrent", self.num_generations) + if is_shared_image + else self.num_generations + ) + target_replicas = self.max_warmpool_replicas or default_replicas try: self.fleet.warm_images( images, replicas_override=target_replicas, wait=wait @@ -276,6 +365,8 @@ def _warm_batch(self, batch: Any, wait: bool = False): def _unwarm_batch(self, images: list[str]): if images and self.fleet: for img in images: + if self.scaffold == "openhands" and os.getenv("AGENT_SERVER_IMAGE"): + continue try: self.fleet.unwarm_image(img) logging.info( @@ -341,7 +432,25 @@ def _teardown_global_fleet() -> None: r2egym = cast(Any, None) EnvArgs = cast(Any, None) RepoEnv = cast(Any, None) - Action = cast(Any, None) + Action = None + + +class _ActionFallback: + """Minimal Action parser fallback when r2egym is not installed.""" + + def __init__(self, function_name: str, parameters: dict[str, str]): + self.function_name = function_name + self.parameters = parameters + + @classmethod + def from_string(cls, action_str: str) -> "_ActionFallback": + fn_match = re.search(r"]+)>", action_str) + function_name = fn_match.group(1).strip() if fn_match else "" + pattern = r"]+)>(.*?)" + param_matches = re.findall(pattern, action_str, flags=re.DOTALL) + params = {k.strip(): v.strip() for k, v in param_matches} + return cls(function_name, params) + from tunix.rl.agentic.environments.base_environment import BaseTaskEnv, EnvStepResult @@ -410,7 +519,7 @@ def __init__( backend: Backend to use for the environment. delete_image: Whether to delete the Docker image after closing. verbose: Verbose output toggle. - scaffold: Scaffold tool set ('r2egym' or 'sweagent'). + scaffold: Scaffold tool set ('r2egym', 'sweagent', or 'openhands'). max_steps: Maximum interaction steps. use_agent_sandbox: If True, strictly forces SandboxFleet and AgentSandboxRuntime. @@ -423,16 +532,19 @@ def __init__( self.delete_image = delete_image self.backend = backend self.env = None + self.workspace = None self.handle = None self.verbose = verbose self.scaffold = scaffold self.use_agent_sandbox = use_agent_sandbox self.fleet = fleet + self._cached_reward = None assert scaffold in [ "r2egym", "sweagent", - ], f"Invalid scaffold: {scaffold}, must be one of ['r2egym', 'sweagent']" + "openhands", + ], f"Invalid scaffold: {scaffold}, must be one of ['r2egym', 'sweagent', 'openhands']" super().__init__(max_steps=max_steps) if not hasattr(self, "extra_kwargs"): @@ -442,34 +554,86 @@ def __init__( self.extra_kwargs["pair_index"] = pair_index def _initial_observation(self) -> Any: - if not self.env: + if not self.env and not self.workspace: if self.use_agent_sandbox: _patch_r2egym_for_agent_sandbox() from agent_sandbox_rl import Task # pytype: disable=import-error - from agent_sandbox_rl.adapters.r2egym import ( # pytype: disable=import-error - make_fleet_repo_env, - r2egym_command_files, - ) fleet = self.fleet or _get_global_fleet() msg = ( - "[SWEEnv] Acquiring SandboxHandle from SandboxFleet and" - " constructing FleetRepoEnv!" + "[SWEEnv] Acquiring SandboxHandle from SandboxFleet!" ) logging.info(msg) - task = Task( - id=str( - self.entry.get( - "instance_id", self.entry.get("docker_image", "default") - ) - ), - image=self.entry.get("docker_image", "default"), - metadata={"ds": self.entry}, + task_id = str( + self.entry.get( + "instance_id", self.entry.get("docker_image", "default") + ) ) - self.handle = fleet.acquire(task) - # TODO(wuhao): Revisit command_files once other harnesses (such as OpenHands) are supported. - cmd_files = r2egym_command_files() - self.env = make_fleet_repo_env(self.handle, command_files=cmd_files) + task = None + if hasattr(fleet, "tasks") and fleet.tasks: + for t in fleet.tasks: + if t.id == task_id: + task = t + break + if task is None: + task_img = ( + os.getenv("AGENT_SERVER_IMAGE") + if self.scaffold == "openhands" and os.getenv("AGENT_SERVER_IMAGE") + else self.entry.get("docker_image", "default") + ) + if isinstance(task_img, (list, np.ndarray)): + task_img = task_img[0] if len(task_img) > 0 else "default" + task_img_str = str(task_img) + rewrite_fn = _get_image_rewrite_fn( + getattr(fleet, "_image_rewrite_fn", None) + ) + if rewrite_fn and not ( + self.scaffold == "openhands" and os.getenv("AGENT_SERVER_IMAGE") + ): + task_img_str = rewrite_fn(task_img_str) + task = Task( + id=task_id, + image=task_img_str, + metadata={"ds": self.entry}, + ) + max_acquire_retries = 5 + for attempt in range(max_acquire_retries): + try: + self.handle = fleet.acquire(task) + break + except Exception as e: + if attempt < max_acquire_retries - 1: + logging.warning( + "[SWEEnv] fleet.acquire failed (attempt %d/%d): %s; retrying in %ds...", + attempt + 1, + max_acquire_retries, + e, + 5 * (attempt + 1), + ) + time.sleep(5 * (attempt + 1)) + else: + raise + if self.scaffold == "openhands": + from agent_sandbox_rl.adapters.openhands import make_handle_workspace # pytype: disable=import-error + ws_kwargs = {} + if os.getenv("SANDBOX_SESSION_KEY"): + ws_kwargs["api_key"] = os.getenv("SANDBOX_SESSION_KEY") + if os.getenv("ROUTER_URL"): + ws_kwargs["router_url"] = os.getenv("ROUTER_URL") + if os.getenv("ROUTER_AUTH_TOKEN"): + ws_kwargs["router_auth_token"] = os.getenv("ROUTER_AUTH_TOKEN") + ws_kwargs["working_dir"] = os.getenv( + "OPENHANDS_WORKING_DIR", "/testbed" + ) + self.workspace = make_handle_workspace(self.handle, **ws_kwargs) + self._setup_openhands_workspace() + else: + from agent_sandbox_rl.adapters.r2egym import ( # pytype: disable=import-error + make_fleet_repo_env, + r2egym_command_files, + ) + cmd_files = r2egym_command_files() + self.env = make_fleet_repo_env(self.handle, command_files=cmd_files) else: # Initialize standard local Docker RepoEnv global EnvArgs, RepoEnv, Action @@ -486,28 +650,285 @@ def _initial_observation(self) -> Any: ) if self.scaffold == "r2egym": self.env.add_commands(R2EGYM_COMMAND_FILES) - else: + elif self.scaffold == "sweagent": self.env.add_commands(SWEAGENT_COMMAND_FILES) else: - self.env.reset() + if self.env is not None: + self.env.reset() - self.final_reward_fn = self.env.compute_reward # pytype: disable=attribute-error + self.final_reward_fn = self.compute_reward + self._cached_reward = None self.total_steps = 0 + if self.workspace is not None: + return str( + self.entry.get("problem_statement") + or self.entry.get("instruction") + or "" + ) + # Polls docker runtime to get task instruction. return self.env.get_task_instruction() # pytype: disable=attribute-error + def _setup_openhands_workspace(self) -> None: + """Configure repository environment in the OpenHands workspace.""" + if self.workspace is None: + return + + entry = getattr(self, "entry", None) or {} + repo = ( + entry.get("repo_name") + or entry.get("repo") + or "" + ) + commit = ( + entry.get("commit_hash") + or entry.get("base_commit") + or "" + ) + if isinstance(repo, (list, np.ndarray)): + repo = repo[0] if len(repo) > 0 else "" + if isinstance(commit, (list, np.ndarray)): + commit = commit[0] if len(commit) > 0 else "" + + repo_map = { + "pandas": "https://github.com/pandas-dev/pandas.git", + "numpy": "https://github.com/numpy/numpy.git", + "pillow": "https://github.com/python-pillow/Pillow.git", + "tornado": "https://github.com/tornadoweb/tornado.git", + "orange3": "https://github.com/biolab/orange3.git", + "datalad": "https://github.com/datalad/datalad.git", + "aiohttp": "https://github.com/aio-libs/aiohttp.git", + "pyramid": "https://github.com/Pylons/pyramid.git", + "scrapy": "https://github.com/scrapy/scrapy.git", + "coveragepy": "https://github.com/nedbat/coveragepy.git", + } + if repo in repo_map: + repo_url = repo_map[repo] + elif "/" in repo: + repo_url = ( + f"https://github.com/{repo}.git" + if not repo.startswith("http") + else repo + ) + elif repo: + repo_url = f"https://github.com/{repo}/{repo}.git" + else: + repo_url = "" + + setup_cmds = [ + "git config --global user.email 'openhands@agent.sandbox'", + "git config --global user.name 'OpenHands Agent'", + "git config --global --add safe.directory /testbed 2>/dev/null || true", + "git config --global --add safe.directory /workspace 2>/dev/null || true", + ] + if repo_url: + setup_cmds.append( + f"if [ ! -d /testbed/.git ] && [ ! -d /workspace/.git ]; then " + f" git clone {repo_url} /tmp/repo_clone && " + f" cp -a /tmp/repo_clone/. /workspace/ && " + f" rm -rf /tmp/repo_clone && " + f" (([ ! -d /testbed ] || rmdir /testbed 2>/dev/null || true) && ln -s /workspace /testbed 2>/dev/null || true); " + f"elif [ ! -e /testbed ] && [ -d /workspace ]; then " + f" ln -s /workspace /testbed 2>/dev/null || true; " + f"fi" + ) + if commit: + setup_cmds.append( + f"if [ ! -d /testbed/.git ] || [ -L /testbed ]; then " + f" if [ -d /workspace/.git ]; then " + f" cd /workspace && git fetch origin 2>/dev/null || true && " + f" git checkout -f {commit} 2>/dev/null || true; " + f" fi; " + f"fi" + ) + else: + setup_cmds.append( + f"if [ ! -e /testbed ] && [ -d /workspace ]; then " + f" ln -s /workspace /testbed 2>/dev/null || true; " + f"fi" + ) + + full_setup_cmd = " && ".join(setup_cmds) + try: + logging.info( + "[SWEEnv] Configuring repository environment for %s...", + repo, + ) + res = self.workspace.execute_command(full_setup_cmd, timeout=180.0) + if res.exit_code != 0: + logging.warning( + "[SWEEnv] Repository setup exit code %s: %s", + res.exit_code, + res.stderr or res.stdout, + ) + else: + logging.info( + "[SWEEnv] Successfully configured repository environment for %s", repo + ) + except Exception as e: + logging.warning( + "[SWEEnv] Failed to set up repository in workspace: %s", e + ) + + def compute_reward(self, verbose: bool = False) -> float: + """Compute task reward for the current environment state.""" + if self._cached_reward is not None: + return self._cached_reward + + if self.env is not None and hasattr(self.env, "compute_reward"): + try: + reward = float(self.env.compute_reward(verbose=verbose)) + except TypeError: + reward = float(self.env.compute_reward()) + self._cached_reward = reward + return reward + + if self.workspace is not None: + test_patch = self.entry.get("test_patch") + if test_patch: + if hasattr(test_patch, "decode"): + test_patch = test_patch.decode("utf-8") + apply_cmd = ( + f"cat << '__EOF_PATCH__' > /tmp/test_patch.diff\n{test_patch}\n__EOF_PATCH__\n" + "(cd /testbed 2>/dev/null || cd /workspace) && " + "(git apply --whitespace=nowarn /tmp/test_patch.diff 2>/dev/null || patch -p1 < /tmp/test_patch.diff)" + ) + try: + self.workspace.execute_command(apply_cmd, timeout=60.0) + except Exception as e: + logging.warning("[SWEEnv] Failed to apply test_patch: %s", e) + + eval_script = ( + self.entry.get("eval_script") + or self.entry.get("run_tests_regression") + ) + if eval_script: + if hasattr(eval_script, "decode"): + eval_script = eval_script.decode("utf-8") + eval_cmd = ( + f"cat << '__EOF_EVAL__' > /tmp/eval.sh\n{eval_script}\n__EOF_EVAL__\n" + "chmod +x /tmp/eval.sh && (cd /testbed 2>/dev/null || cd /workspace) && /tmp/eval.sh" + ) + try: + res = self.workspace.execute_command( + eval_cmd, timeout=float(self.reward_timeout) + ) + reward = 1.0 if res.exit_code == 0 else 0.0 + self._cached_reward = reward + return reward + except Exception as e: + logging.warning("[SWEEnv] Reward computation error: %s", e) + self._cached_reward = 0.0 + return 0.0 + + test_cmd = self.entry.get("test_command") + if test_cmd: + try: + res = self.workspace.execute_command( + f"(cd /testbed 2>/dev/null || cd /workspace) && {test_cmd}", + timeout=float(self.reward_timeout), + ) + reward = 1.0 if res.exit_code == 0 else 0.0 + self._cached_reward = reward + return reward + except Exception as e: + logging.warning("[SWEEnv] Reward computation error: %s", e) + self._cached_reward = 0.0 + return 0.0 + + try: + run_cmd = ( + "if [ -f /testbed/run_tests.sh ]; then " + "cd /testbed && bash /testbed/run_tests.sh; " + "elif [ -f /run_tests.sh ]; then " + "(cd /testbed 2>/dev/null || cd /workspace) && bash /run_tests.sh; " + "else exit 1; fi" + ) + res = self.workspace.execute_command( + run_cmd, + timeout=float(self.reward_timeout), + ) + reward = 1.0 if res.exit_code == 0 else 0.0 + self._cached_reward = reward + return reward + except Exception as e: + logging.warning("[SWEEnv] Reward computation error: %s", e) + self._cached_reward = 0.0 + return 0.0 + + return 0.0 + def _step_impl(self, action: Any) -> EnvStepResult: global Action if Action is None: - from r2egym.agenthub.action import Action # pytype: disable=import-error + try: + from r2egym.agenthub.action import Action # pytype: disable=import-error + except ImportError: + Action = _ActionFallback if isinstance(action, str): action_obj = Action.from_string(action) else: action_obj = action if not action_obj.function_name: - return EnvStepResult(observation="", reward=0, done=False, info={}) + return EnvStepResult( + observation="", + reward=0, + done=False, + info={"max_steps": self.max_steps}, + ) + + if self.scaffold == "openhands" and self.workspace is not None: + if action_obj.function_name in ("finish", "submit"): + return EnvStepResult( + observation="Task submitted.", + reward=0, + done=True, + info={"max_steps": self.max_steps}, + ) + + if action_obj.function_name != "execute_bash": + return EnvStepResult( + observation=( + f"ERROR: Tool '{action_obj.function_name}' is not recognized. " + "Only 'execute_bash' and 'submit' are available." + ), + reward=0, + done=False, + info={"max_steps": self.max_steps}, + ) + + cmd = action_obj.parameters.get("command") or action_obj.parameters.get( + "cmd" + ) + if not cmd: + return EnvStepResult( + observation="ERROR: No command specified for execute_bash.", + reward=0, + done=False, + info={"max_steps": self.max_steps}, + ) + + try: + wrapped_cmd = f"(cd /testbed 2>/dev/null || cd /workspace) && {cmd}" + result = self.workspace.execute_command( + wrapped_cmd, timeout=float(self.step_timeout) + ) + obs = ( + result.stdout + if result.exit_code == 0 + else f"{result.stdout}\n{result.stderr}" + ) + except Exception as e: + obs = f"Command execution failed: {e}" + self.total_steps += 1 + return EnvStepResult( + observation=obs, + reward=0, + done=False, + info={"max_steps": self.max_steps}, + ) # RepoEnv always returns 0 reward, must be evaluated by DockerRuntime. if not self.env: @@ -522,9 +943,17 @@ def _step_impl(self, action: Any) -> EnvStepResult: def close(self) -> None: """Close the environment and clean up resources.""" + self._cached_reward = None if self.env is not None: self.env.close() + if getattr(self, "workspace", None) is not None: + try: + self.workspace.cleanup() + except Exception as e: + logging.warning("[SWEEnv] Workspace cleanup note: %s", e) + self.workspace = None + fleet = self.fleet or _GLOBAL_FLEET if ( hasattr(self, "handle") diff --git a/examples/deepswe/template.py b/examples/deepswe/template.py new file mode 100644 index 0000000000..54ba2612ab --- /dev/null +++ b/examples/deepswe/template.py @@ -0,0 +1,481 @@ +# Copyright 2026 The Google Research Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Templates and specifications for DeepSWE agents and environments.""" + +import json +import os +from typing import Any, Optional + +# ============================================================================== +# Agent System Prompts +# ============================================================================== + +SWE_SYSTEM_PROMPT_FN_CALL = """You are a programming agent who is provided a github issue and repository bash environment and is tasked to solve certain tasks (e.g., file localization, testcase generation, code repair and editing etc) to resolve the issue. +""" + +SWE_SYSTEM_PROMPT = """You are a programming agent who is provided a github issue and repository bash environment and is tasked to solve certain tasks (e.g., file localization, testcase generation, code repair and editing etc) to resolve the issue. + +We have access to the following functions: + +–– BEGIN FUNCTION #1: file_editor –– +Description: +Custom editing tool for viewing, creating and editing files + • State is persistent across command calls and discussions with the user + • If path is a file, view displays the result of applying cat -n. If path is a directory, view lists non-hidden files and directories up to 2 levels deep + • The create command cannot be used if the specified path already exists as a file + • If a command generates a long output, it will be truncated and marked with + • The undo_edit command will revert the last edit made to the file at path + +Notes for using the str_replace command: + • The old_str parameter should match EXACTLY one or more consecutive lines from the original file. Be mindful of whitespaces! + • If the old_str parameter is not unique in the file, the replacement will not be performed. Make sure to include enough context in old_str to make it unique + • The new_str parameter should contain the edited lines that should replace the old_str + +Parameters: + 1. command (string, required) +Allowed values: [view, create, str_replace, insert, undo_edit] +The command to run. + 2. path (string, required) +Absolute path to file or directory, e.g. /testbed/file.py or /testbed. + 3. file_text (string, optional) +Required for the create command. Contains the content of the file to be created. + 4. old_str (string, optional) +Required for the str_replace command. The exact string in path to replace. + 5. new_str (string, optional) + • Optional for the str_replace command to specify the replacement string. + • Required for the insert command to specify the string to insert. + 6. insert_line (integer, optional) +Required for the insert command. The new_str will be inserted after the line number specified here. + 7. view_range (array, optional) + • Optional for the view command (when path is a file). + • If provided, specifies the line range to view, e.g. [11, 12] shows lines 11 and 12. + • [start_line, -1] will show all lines from start_line to the end of file. + 8. concise (boolean, optional) + • Optional for the view command. + • Defaults to True; displays a concise skeletal view of the file. If set to False, displays the full content in the specified view_range. + +–– END FUNCTION #1 –– + +–– BEGIN FUNCTION #2: execute_bash –– +Description: +Execute a bash command in the terminal. + +Behavior notes: + • If a command may run indefinitely (long-running), consider running it in the background and redirecting output, e.g. python3 app.py > server.log 2>&1 &. + • If the bash command returns exit code -1, it means the process is still running. The assistant may: + • Call this function again with command as an empty string ("") to retrieve additional logs. + • Send more input to STDIN of the running process by calling this function again with command set to the text input. + • Send command="ctrl+c" to interrupt the currently running process. + • If the command times out, it will be interrupted (SIGINT). The assistant may then retry or do further steps if needed. + +Parameters: + 1. cmd (string, required) +The bash command (and optional arguments) to execute. + • Can be empty ("") to retrieve more logs if the process is still running. + • Can be "ctrl+c" to interrupt the running process. + +–– END FUNCTION #2 –– + +–– BEGIN FUNCTION #3: search –– +Description: +Search for a term in a directory or a single file. + • If path is a directory (or unspecified, default is .), it recursively searches all non-hidden files and directories for the search term. + • If path points to a file, it runs a grep -n in that file to show line numbers matching the search term. + • If more than 100 files match in a directory search, results are truncated and the tool will inform you to narrow your search. + • If no matches are found, it will inform you as well. + +Parameters: + 1. search_term (string, required) +The term or string to search for in files. + 2. path (string, optional) +The file or directory to search in. Defaults to . if not specified. + +–– END FUNCTION #3 –– + +–– BEGIN FUNCTION #4: finish –– +Description: +Finish the interaction once the task is complete or if no further progress can be made. + +Behavior notes: + • The submit command finalizes your output. + +Parameters: + 1. command (string, required) +Currently allowed value: [submit] + 2. result (string, optional) +The result text or final message to submit. Defaults to an empty string if not provided. + +–– END FUNCTION #4 –– + +If you choose to call a function ONLY reply in the following format with NO suffix: + + +value_1 + +This is the value for the second parameter +that can span +multiple lines + + + + +Reminder: +- Function calls MUST follow the specified format, start with +- Required parameters MUST be specified +- Only call one function at a time +- VERY IMPORTANT: Each response must include both reasoning (as natural text) and function call (in above format) to solve the task. + +""" + +SWEAGENT_SYSTEM_PROMPT = """You are a programming agent who is provided a github issue and repository bash environment and is tasked to solve certain tasks (e.g., file localization, testcase generation, code repair and editing etc) to resolve the issue. + +We have access to the following functions: + +---- BEGIN FUNCTION #1: execute_bash ---- +Description: Execute a bash command in the terminal. +Parameters: + (1) command (string, required): The bash command to execute. For example: `python my_script.py`. If not provided, will show help. +---- END FUNCTION #1 ---- + + +---- BEGIN FUNCTION #2: submit ---- +Description: Finish the interaction when the task is complete OR if the assistant cannot proceed further with the task. +No parameters are required for this function. +---- END FUNCTION #2 ---- + + +---- BEGIN FUNCTION #3: str_replace_editor ---- +Description: Custom editing tool for viewing, creating and editing files +* State is persistent across command calls and discussions with the user +* If `path` is a file, `view` displays the result of applying `cat -n`. If `path` is a directory, `view` lists non-hidden files and directories up to 2 levels deep +* The `create` command cannot be used if the specified `path` already exists as a file +* If a `command` generates a long output, it will be truncated and marked with `` +Notes for using the `str_replace` command: +* The `old_str` parameter should match EXACTLY one or more consecutive lines from the original file. Be mindful of whitespaces! +* If the `old_str` parameter is not unique in the file, the replacement will not be performed. Make sure to include enough context in `old_str` to make it unique +* The `new_str` parameter should contain the edited lines that should replace the `old_str` +Parameters: + (1) command (string, required): The commands to run. Allowed options are: `view`, `create`, `str_replace`, `insert`. +Allowed values: [`view`, `create`, `str_replace`, `insert`] + (2) path (string, required): Absolute path to file or directory, e.g. `/repo/file.py` or `/repo`. + (3) file_text (string, optional): Required parameter of `create` command, with the content of the file to be created. + (4) old_str (string, optional): Required parameter of `str_replace` command containing the string in `path` to replace. + (5) new_str (string, optional): Optional parameter of `str_replace` command containing the new string (if not given, no string will be added). Required parameter of `insert` command containing the string to insert. + (6) insert_line (integer, optional): Required parameter of `insert` command. The `new_str` will be inserted AFTER the line `insert_line` of `path`. + (7) view_range (array, optional): Optional parameter of `view` command when `path` points to a file. If none is given, the full file is shown. If provided, the file will be shown in the indicated line number range, e.g. [11, 12] will show lines 11 and 12. Indexing at 1 to start. Setting `[start_line, -1]` shows all lines from `start_line` to the end of the file. +---- END FUNCTION #3 ---- + + +If you choose to call a function ONLY reply in the following format with NO suffix: + +Provide any reasoning for the function call here. + +value_1 + +This is the value for the second parameter +that can span +multiple lines + + + + +Reminder: +- Function calls MUST follow the specified format, start with +- Required parameters MUST be specified +- Only call one function at a time +- Always provide reasoning for your function call in natural language BEFORE the function call (not after) +""" + + +OPENHANDS_SYSTEM_PROMPT = """You are OpenHands agent, a helpful AI assistant that can interact with a computer to solve tasks. +You are provided a github issue and repository bash environment and is tasked to solve certain tasks (e.g., file localization, testcase generation, code repair and editing etc) to resolve the issue. + +We have access to the following functions: + +–– BEGIN FUNCTION #1: execute_bash –– +Description: +Execute a bash command in the terminal. +Parameters: + 1. command (string, required) +The bash command to execute. For example: `python my_script.py`. If not provided, will show help. +–– END FUNCTION #1 –– + +–– BEGIN FUNCTION #2: submit –– +Description: +Finish the interaction when the task is complete OR if the assistant cannot proceed further with the task. +No parameters are required for this function. +–– END FUNCTION #2 –– + +If you choose to call a function ONLY reply in the following format with NO suffix: + + +value_1 + +This is the value for the second parameter +that can span +multiple lines + + + + +Reminder: +- Function calls MUST follow the specified format, start with +- Required parameters MUST be specified +- Only call one function at a time +- VERY IMPORTANT: Each response must include both reasoning (as natural text) and function call (in above format) to solve the task. + +""" + +# ============================================================================== +# Agent User Prompts +# ============================================================================== + +SWE_USER_PROMPT_FN_CALL = """Consider the following github issue: + +{problem_statement} + + +Can you help me implement the necessary changes to the repository to fix the ? +I've already taken care of all changes to any of the test files described in the . This means you DON'T have to modify the testing logic or any of the tests in any way! +Your task is to make the minimal changes to non-tests files in the /testbed directory to ensure the is satisfied. + +IMPORTANT TIP: +Follow these steps to resolve the issue: +1. As a first step, it might be a good idea to explore the repo to familiarize yourself with its structure. +2. Create a script ('reproduce_issue.py') to reproduce the error and execute it to confirm the error + 2.1 reproduce_issue.py script finishes quickly after checking the error, fix etc. There no long running background servers for django for instance etc. It should be a quick script which checks the error and fix to provide a visible response. + 2.2 SUPER IMPORTANT: to ensure this reproduce_script.py must have a timeout logic of 20 seconds. If the script runs for more than 30 seconds, it should output a timeout message and you can interpret accordingly. +3. Edit the sourcecode of the repo to resolve the issue +4. Rerun your reproduce script and confirm that the error is fixed! +5. Think about edgecases and make sure your fix handles them as well + +VERY IMPORTANT: each response must include both reasoning and function call to solve the task. +You are being told a million times, each response must include a function call. Must inlcude a function call at all costs. + +You can take multiple turns to solve the task. So please only finish / submit when you are confident in your response. Dont rush. Be comprehensive. +You are being told a million times, please dont just submit without proper reasoning. Try to fully analyse the problem statement, explore the repository, reproduce the issue, fix it, check edge cases and then submit. + +Your thinking should be thorough and so it's fine if it's very long. +VERY IMPORTANT: file_editor old_str and new_str must be w/o the line numbers. line numbers are only shown in the view for clarity. + +Also if a file_editor edit fails, its a good idea to view the file near the edit location before trying to edit again. Dont keep trying the same edit over and over again. It will keep leading to the same failure. +Again do not get stuck trying to do the same thing over and over again. Please be efficient. +""" + +SWE_USER_PROMPT = """Consider the following github issue: + +{problem_statement} + + +Can you help me implement the necessary changes to the repository to fix the ? +I've already taken care of all changes to any of the test files described in the . This means you DON'T have to modify the testing logic or any of the tests in any way! +Your task is to make the minimal changes to non-tests files in the /testbed directory to ensure the is satisfied. + +IMPORTANT TIP: +Follow these steps to resolve the issue: +1. As a first step, it might be a good idea to explore the repo to familiarize yourself with its structure. +2. Create a script ('reproduce_issue.py') to reproduce the error and execute it to confirm the error +3. Edit the sourcecode of the repo to resolve the issue +4. Rerun your reproduce script and confirm that the error is fixed! +5. Think about edgecases and make sure your fix handles them as well +6. When viewing large files, use specific line-ranges, usually within 50 to 100 lines) as required +7. NOTE: The repository is at '/testbed' and the current working directory is already '/testbed', so DO NOT include 'testbed/' or 'testbed.' in relative paths in bash commands or reproduction python files. +""" + +SWEAGENT_USER_PROMPT = """I have uploaded a python code repository in the /testbed directory. + +Now consider the following Github issue: + + +{problem_statement} + + +Can you help me implement the necessary changes to the repository to fix the ? +I have already taken care of all changes to any of the test files described in the . This means you DON'T have to modify the testing logic or any of the tests in any way! Your task is to make changes to non-test files in the /testbed directory to ensure the is resolved. + +Follow these steps to resolve the issue: +1. First, explore the codebase to locate and understand the code relevant to the . + - Use efficient search commands to identify key files and functions. + - You should err on the side of caution and look at various relevant files and build your understanding of + - how the code works + - what are the expected behaviors and edge cases + - what are the potential root causes for the given issue + +2. Assess whether you can reproduce the issue: + - Create a script at '/testbed/reproduce_issue.py' that demonstrates the error. + - Execute this script to confirm the error behavior. + - You should reproduce the issue before fixing it. + - Your reproduction script should also assert the expected behavior for the fixed code. + +3. Analyze the root cause: + - Identify the underlying problem based on your code exploration and reproduction results. + - Critically analyze different potential approaches to fix the issue. + - You NEED to explicitly reason about multiple approaches to fix the issue. Next, find the most elegant and effective solution among them considering the tradeoffs (correctness, generality, side effects, etc.). + - You would need to reason about execution paths, edge cases, and other potential issues. You should look at the unit tests to understand the expected behavior of the relevant code. + +4. Implement your solution: + - Make targeted changes to the necessary files following idiomatic code patterns once you determine the root cause. + - You should be thorough and methodical. + +5. Verify your solution: + - Rerun your reproduction script to confirm the error is fixed. + - If verification fails, iterate on your solution until successful. If you identify the reproduction script is buggy, adjust it as needed. + +6. Run unit tests: + - Find and run the relevant unit tests relevant to the performed fix. + - You should run the unit tests to ensure your solution is correct and does not cause any regressions. + - In cases where the unit tests are do not pass, you should consider whether the unit tests does not reflect the *new* expected behavior of the code. If so, you can test it by writing additional edge test cases. + - Use the existing test runner to run the unit tests you identify as relevant to the changes you made. For example: + - `python -m pytest -xvs sympy/physics/units/tests/test_dimensions_transcendental.py` + - `python -m pytest tests/test_domain_py.py::test_pymethod_options` + - `./tests/runtests.py constraints.tests.CheckConstraintTests -v 2` + - RUN ALL relevant unit tests to ensure your solution is correct and does not cause any regressions. + +7. Test edge cases: + - Identify potential edge cases that might challenge your solution. + - Create additional test cases in a separate file '/testbed/edge_case_tests.py'. + - Execute these tests to verify your solution's robustness. + - You should run multiple rounds of edge cases. When creating edge cases: + - Consider complex scenarios beyond the original issue description + - Test for regressions to ensure existing functionality remains intact + +8. Refine if necessary: + - If edge case testing reveals issues, refine your solution accordingly. + - Ensure your final implementation handles all identified scenarios correctly. + - Document any assumptions or limitations of your solution. + +9. Submit your solution: + - Once you have verified your solution, submit your solution using the `submit` tool. + +A successful resolution means: +- The specific error/issue described no longer occurs +- Your changes maintain compatibility with existing functionality +- Edge cases are properly handled + + +Additional recommendations: +- You should be thorough, methodical, and prioritize quality over speed. Be comprehensive. +- You should think carefully before making the tool call about what should be done. However, each step should only use one tool call. YOU SHOULD NOT USE TOOLS INSIDE YOUR THOUGHT PROCESS. YOU SHOULD PRIMARILY USE THINKING FOR IDENTIFYING THE ROOT CAUSE OF THE ISSUE, MAKING THE CHANGES, AND CREATING TEST CASES (REPRODUCTION OR EDGE CASES). +- Each action you take is somewhat expensive. Wherever possible, combine multiple actions into a single action (e.g., combine multiple bash commands, use sed/grep for bulk operations). + - Your grep commands should identify both relevant files and line numbers so you can use the file_editor tool. + - Use grep with `-A -B -C` flags to quickly identify the relevant code blocks during your exploration. +- When exploring the codebase, use targeted search patterns to minimize unnecessary operations. +- When creating edge cases, you should look at the relevant existing tests to understand existing "regression" test cases. Ensure the fix doesn't break existing functionality. +""" + + +def get_system_prompt( + scaffold: str = "r2egym", + use_fn_calling: bool = False, +) -> str: + """Get system prompt for the given scaffold and function calling mode.""" + if scaffold == "sweagent": + return SWEAGENT_SYSTEM_PROMPT + elif scaffold == "openhands": + return OPENHANDS_SYSTEM_PROMPT + return SWE_SYSTEM_PROMPT_FN_CALL if use_fn_calling else SWE_SYSTEM_PROMPT + + +def get_user_prompt_template( + scaffold: str = "r2egym", + use_fn_calling: bool = False, +) -> str: + """Get user prompt template for the given scaffold and function calling mode.""" + if scaffold == "sweagent": + return SWEAGENT_USER_PROMPT + return SWE_USER_PROMPT_FN_CALL if use_fn_calling else SWE_USER_PROMPT + + +# ============================================================================== +# Environment Sandbox Fleet Pod Templates +# ============================================================================== + +DEFAULT_OPENHANDS_KEEPALIVE_CMD = [ + "sh", + "-c", + ( + "if [ -x /usr/local/bin/openhands-agent-server ]; then " + "exec tini -- /usr/local/bin/openhands-agent-server --host 0.0.0.0 --port 8000; " + "else " + "exec tini -- /agent-server/.venv/bin/python -m openhands.agent_server --host 0.0.0.0 --port 8000; " + "fi" + ), +] + + +def get_openhands_pod_template( + node_selector: Optional[dict[str, str]] = None, +) -> Any: + """Builds and returns the TemplateSpec for OpenHands agent sandbox.""" + try: + from agent_sandbox_rl import ( # pytype: disable=import-error + ResourceSpec, + TemplateSpec, + ) + except ImportError as e: + raise ImportError( + "use_agent_sandbox=True strictly requires the 'agent_sandbox_rl'" + " package. Install via: pip install" + " git+https://github.com/kubernetes-sigs/agent-sandbox.git#subdirectory=examples/agent-sandbox-rl" + ) from e + + session_key = os.getenv("SANDBOX_SESSION_KEY", "") + if os.getenv("AGENT_SERVER_COMMAND"): + try: + keepalive_cmd = json.loads(os.environ["AGENT_SERVER_COMMAND"]) + except Exception: + keepalive_cmd = os.environ["AGENT_SERVER_COMMAND"].split() + else: + keepalive_cmd = list(DEFAULT_OPENHANDS_KEEPALIVE_CMD) + + return TemplateSpec( + keepalive_command=keepalive_cmd, + resources=ResourceSpec( + cpu=os.getenv("SANDBOX_CPU", "500m"), + memory=os.getenv("SANDBOX_MEM", "1Gi"), + ), + extra_pod_spec={ + "containers": [{ + "ports": [{"containerPort": 8000}], + "readinessProbe": { + "httpGet": {"path": "/health", "port": 8000}, + "periodSeconds": 2, + "failureThreshold": 150, + }, + "resources": { + "limits": { + "cpu": os.getenv("SANDBOX_CPU_LIMIT", "2"), + "memory": os.getenv("SANDBOX_MEM_LIMIT", "4Gi"), + } + }, + "env": ( + [{"name": "OH_SESSION_API_KEYS_0", "value": session_key}] + if session_key + else [] + ), + }] + }, + node_selector=node_selector, + ) + + +def get_template( + scaffold: str, + node_selector: Optional[dict[str, str]] = None, +) -> Any: + """Returns the fleet TemplateSpec for the given scaffold, or None.""" + if scaffold == "openhands": + return get_openhands_pod_template(node_selector=node_selector) + return None diff --git a/examples/deepswe/train_maxtext_nb.py b/examples/deepswe/train_maxtext_nb.py index 0b2e5864aa..16447e03d8 100644 --- a/examples/deepswe/train_maxtext_nb.py +++ b/examples/deepswe/train_maxtext_nb.py @@ -349,6 +349,13 @@ def str2bool(v): choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], help="Logging level for the script and relevant libraries.", ) +parser.add_argument( + "--scaffold", + type=str, + default="r2egym", + choices=["r2egym", "sweagent", "openhands"], + help="Agent scaffold/sandbox toolset to use ('r2egym', 'sweagent', or 'openhands').", +) args, _ = parser.parse_known_args() @@ -762,6 +769,7 @@ def mixed_type_batch_fn(elements): num_generations=NUM_GENERATIONS, batch_size=MINI_BATCH_SIZE, max_warmpool_replicas=args.max_warmpool_replicas, + scaffold=args.scaffold, ) train_dataset = swe_env.PrewarmDatasetIterator( train_dataset, @@ -769,6 +777,7 @@ def mixed_type_batch_fn(elements): num_generations=NUM_GENERATIONS, batch_size=MINI_BATCH_SIZE, max_warmpool_replicas=args.max_warmpool_replicas, + scaffold=args.scaffold, ) @@ -1181,13 +1190,14 @@ def _generate_maxtext_config_with_no_remat(vllm_config_param): rl_engine=rl_engine, reward_fns=None, agent_class=swe_agent.SWEAgent, - agent_kwargs={}, + agent_kwargs={"scaffold": args.scaffold}, env_class=swe_env.SWEEnv, env_kwargs={ "max_steps": MAX_TURNS, "step_timeout": STEP_TIMEOUT_SECS, "reward_timeout": REWARD_TIMEOUT_SECS, "verbose": True, + "scaffold": args.scaffold, "use_agent_sandbox": USE_AGENT_SANDBOX, "fleet": fleet, }, @@ -1248,4 +1258,13 @@ def _generate_maxtext_config_with_no_remat(vllm_config_param): ) print("Starting training...", flush=True) -agentic_grpo_learner.train(train_dataset=train_dataset) +try: + agentic_grpo_learner.train(train_dataset=train_dataset) +finally: + if USE_AGENT_SANDBOX: + print("Tearing down agent sandbox fleet...", flush=True) + try: + swe_env._teardown_global_fleet() + except Exception as teardown_e: + print(f"Failed to teardown agent sandbox fleet: {teardown_e}", flush=True) + diff --git a/tests/examples/swe_env_test.py b/tests/examples/swe_env_test.py new file mode 100644 index 0000000000..b16f309a08 --- /dev/null +++ b/tests/examples/swe_env_test.py @@ -0,0 +1,296 @@ +# Copyright 2026 The Google Research Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for SWEEnv and PrewarmDatasetIterator.""" + +import os +import sys +from unittest import mock +import pytest + +from examples.deepswe import swe_env + + +def test_extract_images_gated_by_scaffold(): + """Test that AGENT_SERVER_IMAGE is gated on scaffold == 'openhands'.""" + batch = [{"docker_image": "docker.io/swebench/test-image:v1"}] + + with mock.patch.dict(os.environ, {"AGENT_SERVER_IMAGE": "ghcr.io/openhands/agent-server:custom"}): + # For r2egym scaffold, AGENT_SERVER_IMAGE must be ignored + iter_r2e = object.__new__(swe_env.PrewarmDatasetIterator) + iter_r2e.scaffold = "r2egym" + images_r2e = iter_r2e._extract_images(batch) + assert images_r2e == ["docker.io/swebench/test-image:v1"] + + # For openhands scaffold, AGENT_SERVER_IMAGE should be used + iter_oh = object.__new__(swe_env.PrewarmDatasetIterator) + iter_oh.scaffold = "openhands" + images_oh = iter_oh._extract_images(batch) + assert images_oh == ["ghcr.io/openhands/agent-server:custom"] + + +def test_extract_images_without_override(): + """Test image extraction when AGENT_SERVER_IMAGE is unset.""" + batch = [{"docker_image": "docker.io/swebench/test-image:v1"}] + + with mock.patch.dict(os.environ, {}, clear=True): + iter_oh = object.__new__(swe_env.PrewarmDatasetIterator) + iter_oh.scaffold = "openhands" + assert iter_oh._extract_images(batch) == ["docker.io/swebench/test-image:v1"] + + +def test_warm_batch_shared_image_sizing(): + """Test warm pool sizing for shared image vs per-task images.""" + mock_fleet = mock.MagicMock() + mock_fleet.config.max_concurrent = 64 + + iter_oh_shared = object.__new__(swe_env.PrewarmDatasetIterator) + iter_oh_shared.fleet = mock_fleet + iter_oh_shared.scaffold = "openhands" + iter_oh_shared.num_generations = 8 + iter_oh_shared.max_warmpool_replicas = None + + with mock.patch.dict(os.environ, {"AGENT_SERVER_IMAGE": "ghcr.io/openhands/agent-server:custom"}): + iter_oh_shared._warm_batch([{"docker_image": "dummy"}]) + mock_fleet.warm_images.assert_called_with( + ["ghcr.io/openhands/agent-server:custom"], replicas_override=64, wait=False + ) + + mock_fleet.reset_mock() + with mock.patch.dict(os.environ, {}, clear=True): + iter_oh_shared._warm_batch([{"docker_image": "img1"}]) + mock_fleet.warm_images.assert_called_with( + ["img1"], replicas_override=8, wait=False + ) + + +def test_compute_reward_fallback_uses_tmp_eval_sh(): + """Verify fallback uses /tmp/eval.sh for non-root user compatibility.""" + env = object.__new__(swe_env.SWEEnv) + env._cached_reward = None + env.env = None + env.reward_timeout = 30 + env.entry = {"eval_script": "pytest"} + + mock_workspace = mock.MagicMock() + mock_res = mock.MagicMock(exit_code=0) + mock_workspace.execute_command.return_value = mock_res + env.workspace = mock_workspace + + reward = env.compute_reward() + assert reward == 1.0 + + called_cmd = mock_workspace.execute_command.call_args[0][0] + assert "/tmp/eval.sh" in called_cmd + assert "cat << '__EOF_EVAL__' > /eval.sh" not in called_cmd + + +def test_normalize_tasks_scaffold_gating(): + """Test that _normalize_tasks_for_fleet gates AGENT_SERVER_IMAGE on scaffold.""" + tasks = [{"docker_image": "r2e_img:latest", "instance_id": "task-1"}] + + with mock.patch.dict(os.environ, {"AGENT_SERVER_IMAGE": "agent_server:latest"}): + # r2egym scaffold keeps task image + norm_r2e = swe_env._normalize_tasks_for_fleet(tasks, scaffold="r2egym") + assert getattr(norm_r2e[0], "image", None) == "r2e_img:latest" + + # openhands scaffold uses AGENT_SERVER_IMAGE + norm_oh = swe_env._normalize_tasks_for_fleet(tasks, scaffold="openhands") + assert getattr(norm_oh[0], "image", None) == "agent_server:latest" + + +def _setup_mock_agent_sandbox(): + mock_as = mock.MagicMock() + mock_as.FleetConfig.side_effect = lambda **kw: mock.MagicMock(**kw) + mock_as.TemplateSpec.side_effect = lambda **kw: mock.MagicMock(**kw) + mock_as.ResourceSpec.side_effect = lambda **kw: mock.MagicMock(**kw) + return mock_as + + +def test_init_global_fleet_openhands_template_and_sizing(): + """Test TemplateSpec resources/limits and warmpool sizing in _init_global_fleet.""" + mock_as_rl = _setup_mock_agent_sandbox() + with mock.patch.dict(sys.modules, {"agent_sandbox_rl": mock_as_rl}), \ + mock.patch.dict(os.environ, {"AGENT_SERVER_IMAGE": "ghcr.io/openhands/agent-server:1.44.1"}): + swe_env._GLOBAL_FLEET = None + fleet = swe_env._init_global_fleet( + tasks=[{"docker_image": "r2e_img", "instance_id": "task-1"}], + max_concurrency=32, + num_generations=8, + batch_size=4, + scaffold="openhands", + ) + assert mock_as_rl.SandboxFleet.called + fleet_cfg = mock_as_rl.SandboxFleet.call_args[0][0] + + # Shared image mode pool size must equal effective_max_concurrent = max(32, 4*8*2) = 64 + assert fleet_cfg.max_warmpool_size == 64 + + # Template resources and limits + template = fleet_cfg.template + assert template.resources.cpu == "500m" + assert template.resources.memory == "1Gi" + container = template.extra_pod_spec["containers"][0] + assert container["resources"]["limits"]["cpu"] == "2" + assert container["resources"]["limits"]["memory"] == "4Gi" + assert container["readinessProbe"]["httpGet"]["path"] == "/health" + + swe_env._GLOBAL_FLEET = None + + +def test_extract_images_with_image_rewrite(): + """Test that image_rewrite is applied in PrewarmDatasetIterator._extract_images.""" + batch = [{"docker_image": "docker.io/swebench/test-image:v1"}] + iter_oh = object.__new__(swe_env.PrewarmDatasetIterator) + iter_oh.scaffold = "openhands" + iter_oh.image_rewrite = lambda img: f"gcr.io/custom/{img.split('/')[-1]}" + + with mock.patch.dict(os.environ, {}, clear=True): + assert iter_oh._extract_images(batch) == ["gcr.io/custom/test-image:v1"] + + +def test_extract_images_with_image_rewrite_prefix_env(): + """Test that IMAGE_REWRITE_PREFIX env var is picked up in _extract_images.""" + batch = [{"docker_image": "docker.io/swebench/test-image:v1"}] + iter_oh = object.__new__(swe_env.PrewarmDatasetIterator) + iter_oh.scaffold = "openhands" + + with mock.patch.dict(os.environ, {"IMAGE_REWRITE_PREFIX": "gcr.io/prefix"}): + assert iter_oh._extract_images(batch) == ["gcr.io/prefix/test-image:v1"] + + +def test_keepalive_command_default_and_override(): + """Test adaptive keepalive_command and custom AGENT_SERVER_COMMAND override.""" + mock_as_rl = _setup_mock_agent_sandbox() + with mock.patch.dict(sys.modules, {"agent_sandbox_rl": mock_as_rl}), \ + mock.patch.dict(os.environ, {}, clear=True): + swe_env._GLOBAL_FLEET = None + swe_env._init_global_fleet([], scaffold="openhands") + fleet_cfg = mock_as_rl.SandboxFleet.call_args[0][0] + cmd = fleet_cfg.template.keepalive_command + assert "openhands-agent-server" in cmd[2] + assert "openhands.agent_server" in cmd[2] + + with mock.patch.dict(sys.modules, {"agent_sandbox_rl": mock_as_rl}), \ + mock.patch.dict(os.environ, {"AGENT_SERVER_COMMAND": '["custom", "server"]'}): + swe_env._GLOBAL_FLEET = None + swe_env._init_global_fleet([], scaffold="openhands") + fleet_cfg = mock_as_rl.SandboxFleet.call_args[0][0] + assert fleet_cfg.template.keepalive_command == ["custom", "server"] + swe_env._GLOBAL_FLEET = None + + +def test_setup_openhands_workspace_conditional_clone_and_symlink(): + """Verify setup includes conditional git clone and symlink creation.""" + env = object.__new__(swe_env.SWEEnv) + mock_workspace = mock.MagicMock() + mock_res = mock.MagicMock(exit_code=0) + mock_workspace.execute_command.return_value = mock_res + env.workspace = mock_workspace + env.entry = {"repo": "pandas", "commit_hash": "12345"} + + env._setup_openhands_workspace() + assert mock_workspace.execute_command.called + cmd = mock_workspace.execute_command.call_args[0][0] + assert "git clone https://github.com/pandas-dev/pandas.git" in cmd + assert "ln -s /workspace /testbed" in cmd + assert "safe.directory /testbed" in cmd + assert "safe.directory /workspace" in cmd + + +def test_load_tasks_image_rewrite_scaffold_and_agent_server_gating(): + """Test that load_tasks only rewrites images when scaffold != 'openhands' or AGENT_SERVER_IMAGE is unset.""" + mock_as_rl = _setup_mock_agent_sandbox() + with mock.patch.dict(sys.modules, {"agent_sandbox_rl": mock_as_rl}): + # Case 1: openhands scaffold with AGENT_SERVER_IMAGE set -> no rewrite in load_tasks + with mock.patch.dict(os.environ, { + "AGENT_SERVER_IMAGE": "ghcr.io/openhands/agent-server:1.44.1", + "IMAGE_REWRITE_PREFIX": "gcr.io/custom", + }): + swe_env._GLOBAL_FLEET = None + fleet = swe_env._init_global_fleet( + tasks=[{"docker_image": "r2e_img", "instance_id": "task-1"}], + scaffold="openhands", + ) + assert fleet.load_tasks.called + assert "image_rewrite" not in fleet.load_tasks.call_args[1] + + # Case 2: openhands scaffold with AGENT_SERVER_IMAGE unset (derived images) -> rewrite in load_tasks + with mock.patch.dict(os.environ, { + "IMAGE_REWRITE_PREFIX": "gcr.io/custom", + }, clear=True): + swe_env._GLOBAL_FLEET = None + fleet = swe_env._init_global_fleet( + tasks=[{"docker_image": "r2e_img", "instance_id": "task-1"}], + scaffold="openhands", + ) + assert fleet.load_tasks.called + assert "image_rewrite" in fleet.load_tasks.call_args[1] + + # Case 3: r2egym scaffold -> rewrite in load_tasks + with mock.patch.dict(os.environ, { + "AGENT_SERVER_IMAGE": "ghcr.io/openhands/agent-server:1.44.1", + "IMAGE_REWRITE_PREFIX": "gcr.io/custom", + }): + swe_env._GLOBAL_FLEET = None + fleet = swe_env._init_global_fleet( + tasks=[{"docker_image": "r2e_img", "instance_id": "task-1"}], + scaffold="r2egym", + ) + assert fleet.load_tasks.called + assert "image_rewrite" in fleet.load_tasks.call_args[1] + + swe_env._GLOBAL_FLEET = None + + +def test_acquire_retry_on_transient_error(): + """Verify fleet.acquire retries on transient errors.""" + import types + env = object.__new__(swe_env.SWEEnv) + env.entry = {"instance_id": "test_inst", "docker_image": "test_image:latest"} + env.scaffold = "r2egym" + env.env = None + env.workspace = None + env.use_agent_sandbox = True + env.step_timeout = 60 + env.reward_timeout = 180 + env.verbose = False + mock_fleet = mock.MagicMock() + mock_handle = mock.MagicMock() + mock_fleet.acquire.side_effect = [ + RuntimeError("transient protocol error"), + mock_handle, + ] + env.fleet = mock_fleet + swe_env._fleet = mock_fleet + + mock_r2e = types.ModuleType("agent_sandbox_rl.adapters.r2egym") + mock_make = mock.MagicMock() + mock_repo_env = mock.MagicMock() + mock_repo_env.reset.return_value = ("obs", {}) + mock_make.return_value = mock_repo_env + mock_r2e.make_fleet_repo_env = mock_make + mock_r2e.r2egym_command_files = mock.MagicMock(return_value=[]) + + mock_as_rl = _setup_mock_agent_sandbox() + with mock.patch.dict(sys.modules, { + "agent_sandbox_rl": mock_as_rl, + "agent_sandbox_rl.adapters.r2egym": mock_r2e, + }): + with mock.patch("time.sleep") as mock_sleep: + env._initial_observation() + + assert mock_fleet.acquire.call_count == 2 + assert env.handle == mock_handle + mock_sleep.assert_called_once_with(5) diff --git a/tests/examples/template_test.py b/tests/examples/template_test.py new file mode 100644 index 0000000000..8f365a4260 --- /dev/null +++ b/tests/examples/template_test.py @@ -0,0 +1,137 @@ +# Copyright 2026 The Google Research Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for examples/deepswe/template.py.""" + +import os +import sys +from unittest import mock +import pytest + +from examples.deepswe import template + + +def _setup_mock_agent_sandbox(): + mock_as = mock.MagicMock() + mock_as.TemplateSpec.side_effect = lambda **kw: mock.MagicMock(**kw) + mock_as.ResourceSpec.side_effect = lambda **kw: mock.MagicMock(**kw) + return mock_as + + +def test_prompt_constants_exist(): + """Verify that all DeepSWE prompt constants are defined and non-empty.""" + assert template.SWE_SYSTEM_PROMPT_FN_CALL + assert template.SWE_SYSTEM_PROMPT + assert template.SWEAGENT_SYSTEM_PROMPT + assert template.OPENHANDS_SYSTEM_PROMPT + assert template.SWE_USER_PROMPT_FN_CALL + assert template.SWE_USER_PROMPT + assert template.SWEAGENT_USER_PROMPT + + +def test_get_system_prompt(): + """Verify get_system_prompt returns correct prompt for all scaffolds and modes.""" + assert template.get_system_prompt("r2egym", use_fn_calling=False) == template.SWE_SYSTEM_PROMPT + assert template.get_system_prompt("r2egym", use_fn_calling=True) == template.SWE_SYSTEM_PROMPT_FN_CALL + assert template.get_system_prompt("sweagent", use_fn_calling=False) == template.SWEAGENT_SYSTEM_PROMPT + assert template.get_system_prompt("sweagent", use_fn_calling=True) == template.SWEAGENT_SYSTEM_PROMPT + assert template.get_system_prompt("openhands", use_fn_calling=False) == template.OPENHANDS_SYSTEM_PROMPT + assert template.get_system_prompt("openhands", use_fn_calling=True) == template.OPENHANDS_SYSTEM_PROMPT + + +def test_get_user_prompt_template(): + """Verify get_user_prompt_template returns correct prompt for all scaffolds.""" + assert template.get_user_prompt_template("r2egym", use_fn_calling=False) == template.SWE_USER_PROMPT + assert template.get_user_prompt_template("r2egym", use_fn_calling=True) == template.SWE_USER_PROMPT_FN_CALL + assert template.get_user_prompt_template("sweagent", use_fn_calling=False) == template.SWEAGENT_USER_PROMPT + assert template.get_user_prompt_template("openhands", use_fn_calling=False) == template.SWE_USER_PROMPT + + +def test_get_openhands_pod_template_default(): + """Verify default openhands TemplateSpec construction.""" + mock_as = _setup_mock_agent_sandbox() + with mock.patch.dict(sys.modules, {"agent_sandbox_rl": mock_as}), \ + mock.patch.dict(os.environ, {}, clear=True): + pod_template = template.get_openhands_pod_template(node_selector={"node": "worker"}) + assert pod_template is not None + assert pod_template.node_selector == {"node": "worker"} + assert "openhands-agent-server" in pod_template.keepalive_command[2] + assert pod_template.resources.cpu == "500m" + assert pod_template.resources.memory == "1Gi" + container = pod_template.extra_pod_spec["containers"][0] + assert container["resources"]["limits"]["cpu"] == "2" + assert container["resources"]["limits"]["memory"] == "4Gi" + assert container["readinessProbe"]["httpGet"]["path"] == "/health" + assert container["ports"] == [{"containerPort": 8000}] + assert container["env"] == [] + + +def test_get_openhands_pod_template_with_overrides(): + """Verify openhands TemplateSpec honors env var overrides.""" + mock_as = _setup_mock_agent_sandbox() + with mock.patch.dict(sys.modules, {"agent_sandbox_rl": mock_as}), \ + mock.patch.dict(os.environ, { + "SANDBOX_SESSION_KEY": "secret_key_123", + "AGENT_SERVER_COMMAND": '["custom", "entrypoint"]', + "SANDBOX_CPU": "1", + "SANDBOX_MEM": "2Gi", + "SANDBOX_CPU_LIMIT": "4", + "SANDBOX_MEM_LIMIT": "8Gi", + }): + pod_template = template.get_openhands_pod_template() + assert pod_template.keepalive_command == ["custom", "entrypoint"] + assert pod_template.resources.cpu == "1" + assert pod_template.resources.memory == "2Gi" + container = pod_template.extra_pod_spec["containers"][0] + assert container["resources"]["limits"]["cpu"] == "4" + assert container["resources"]["limits"]["memory"] == "8Gi" + assert container["env"] == [{"name": "OH_SESSION_API_KEYS_0", "value": "secret_key_123"}] + + +def test_get_template_scaffolds(): + """Verify get_template delegates to openhands or returns None for other scaffolds.""" + mock_as = _setup_mock_agent_sandbox() + with mock.patch.dict(sys.modules, {"agent_sandbox_rl": mock_as}), \ + mock.patch.dict(os.environ, {}, clear=True): + assert template.get_template("openhands") is not None + assert template.get_template("r2egym") is None + assert template.get_template("sweagent") is None + + +def test_swe_agent_reexports(): + """Verify swe_agent re-exports all prompt constants and helpers for backward compatibility.""" + mock_r2e = mock.MagicMock() + with mock.patch.dict(sys.modules, {"r2egym": mock_r2e, "r2egym.agenthub.action": mock_r2e}): + from examples.deepswe import swe_agent + assert swe_agent.SWE_SYSTEM_PROMPT == template.SWE_SYSTEM_PROMPT + assert swe_agent.SWE_SYSTEM_PROMPT_FN_CALL == template.SWE_SYSTEM_PROMPT_FN_CALL + assert swe_agent.SWEAGENT_SYSTEM_PROMPT == template.SWEAGENT_SYSTEM_PROMPT + assert swe_agent.OPENHANDS_SYSTEM_PROMPT == template.OPENHANDS_SYSTEM_PROMPT + assert swe_agent.SWE_USER_PROMPT == template.SWE_USER_PROMPT + assert swe_agent.SWE_USER_PROMPT_FN_CALL == template.SWE_USER_PROMPT_FN_CALL + assert swe_agent.SWEAGENT_USER_PROMPT == template.SWEAGENT_USER_PROMPT + assert swe_agent.get_system_prompt == template.get_system_prompt + assert swe_agent.get_user_prompt_template == template.get_user_prompt_template + + agent_r2e = swe_agent.SWEAgent(scaffold="r2egym") + assert agent_r2e.system_prompt == template.SWE_SYSTEM_PROMPT + assert agent_r2e.user_prompt_template == template.SWE_USER_PROMPT + + agent_swe = swe_agent.SWEAgent(scaffold="sweagent") + assert agent_swe.system_prompt == template.SWEAGENT_SYSTEM_PROMPT + assert agent_swe.user_prompt_template == template.SWEAGENT_USER_PROMPT + + agent_oh = swe_agent.SWEAgent(scaffold="openhands") + assert agent_oh.system_prompt == template.OPENHANDS_SYSTEM_PROMPT + assert agent_oh.user_prompt_template == template.SWE_USER_PROMPT