Skip to content

Commit 3f93a5a

Browse files
authored
rock agentic swe agent impl (#143)
* feat: add rock agentic demo for swe agent * feat: update logger print for swe agent * fix: correct SweAgent import path * feat: add base deps to sdk * feat: update deps in pyproject.toml * feat: update lock * feat: add retry logic for swe-agent installation
1 parent 81876fb commit 3f93a5a

4 files changed

Lines changed: 308 additions & 6 deletions

File tree

pyproject.toml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,10 +69,15 @@ builder = [
6969
"docker"
7070
]
7171

72+
examples = [
73+
"pyyaml"
74+
]
75+
7276
all = [
7377
"rl-rock[admin]",
7478
"rl-rock[rocklet]",
7579
"rl-rock[builder]",
80+
"rl-rock[examples]",
7681
]
7782

7883
[project.scripts]

rock/env_vars.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import json
12
import os
23
import sys
34
from collections.abc import Callable
@@ -39,6 +40,10 @@
3940
# Model Service Config
4041
ROCK_MODEL_SERVICE_DATA_DIR: str
4142

43+
# Agentic
44+
ROCK_AGENT_PRE_STARTUP_BASH_CMD_LIST: list[str] = []
45+
ROCK_AGENT_PYTHON_INSTALL_CMD: str
46+
4247

4348
environment_variables: dict[str, Callable[[], Any]] = {
4449
"ROCK_LOGGING_PATH": lambda: os.getenv("ROCK_LOGGING_PATH"),
@@ -74,6 +79,11 @@
7479
"ROCK_CLI_DEFAULT_CONFIG_PATH", Path.home() / ".rock" / "config.ini"
7580
),
7681
"ROCK_MODEL_SERVICE_DATA_DIR": lambda: os.getenv("ROCK_MODEL_SERVICE_DATA_DIR", "/data/logs"),
82+
"ROCK_AGENT_PYTHON_INSTALL_CMD": lambda: os.getenv(
83+
"ROCK_AGENT_PYTHON_INSTALL_CMD",
84+
"[ -f cpython31114.tar.gz ] && rm cpython31114.tar.gz; [ -d python ] && rm -rf python; wget -q -O cpython31114.tar.gz https://github.com/astral-sh/python-build-standalone/releases/download/20251120/cpython-3.11.14+20251120-x86_64-unknown-linux-gnu-install_only.tar.gz && tar -xzf cpython31114.tar.gz",
85+
),
86+
"ROCK_AGENT_PRE_STARTUP_BASH_CMD_LIST": lambda: json.loads(os.getenv("ROCK_AGENT_PRE_STARTUP_BASH_CMD_LIST", "[]")),
7787
}
7888

7989

Lines changed: 286 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,304 @@
1+
"""
2+
SWE-agent Integration Module
3+
4+
This module provides integration with SWE-agent (Software Engineering Agent) for automated
5+
software engineering tasks within a sandboxed environment. It handles the complete lifecycle
6+
of SWE-agent including environment initialization, dependency installation, and execution.
7+
8+
Key Components:
9+
- SweAgentConfig: Configuration dataclass for SWE-agent setup parameters
10+
- SweAgent: Main agent implementation managing initialization and execution
11+
12+
Usage Example:
13+
```python
14+
from rock.sdk.sandbox.client import Sandbox
15+
16+
swe_agent_config = SweAgentConfig(
17+
agent_type="swe-agent",
18+
version="unknown",
19+
swe_agent_workdir=self.sweagent_dir,
20+
swe_session=self.swe_session,
21+
)
22+
23+
sandbox = Sandbox(...)
24+
sandbox.agent = SweAgent(sandbox, config)
25+
26+
await sandbox.agent.init()
27+
await sandbox.agent.run("path/to/config.yaml")
28+
```
29+
30+
Note:
31+
Currently supports LocalDeployment and RunSingleConfig modes only.
32+
Requires a Sandbox instance (not AbstractSandbox) for execution.
33+
"""
34+
35+
import os
36+
import shlex
37+
from pathlib import Path
138
from typing import Literal
239

40+
from rock import env_vars
341
from rock.actions.sandbox.base import AbstractSandbox
42+
from rock.actions.sandbox.request import CreateBashSessionRequest, UploadRequest
43+
from rock.logger import init_logger
444
from rock.sdk.sandbox.agent.base import Agent
545
from rock.sdk.sandbox.agent.config import AgentConfig
46+
from rock.sdk.sandbox.client import Sandbox
47+
from rock.utils import retry_async
48+
49+
logger = init_logger(__name__)
650

751

852
class SweAgentConfig(AgentConfig):
53+
"""
54+
Configuration dataclass for SWE-agent initialization and execution.
55+
56+
This class defines all configurable parameters for setting up and running
57+
SWE-agent in a sandboxed environment, including installation commands,
58+
working directories, and execution timeouts.
59+
60+
Attributes:
61+
agent_type: Fixed identifier for this agent type ("swe-agent")
62+
swe_session: Name of the bash session used for SWE-agent execution
63+
pre_startup_bash_cmd_list: Commands executed before agent initialization
64+
post_startup_bash_cmd_list: Commands executed after agent initialization
65+
swe_agent_workdir: Working directory for agent installation and execution
66+
python_install_cmd: Command to install Python environment
67+
swe_agent_install_cmd: Command to clone and install SWE-agent repository
68+
python_install_timeout: Maximum seconds to wait for Python installation
69+
swe_agent_install_timeout: Maximum seconds to wait for SWE-agent installation
70+
agent_run_timeout: Maximum seconds to wait for agent execution completion
71+
agent_run_check_interval: Seconds between status checks during execution
72+
"""
73+
974
agent_type: Literal["swe-agent"] = "swe-agent"
1075

76+
swe_session: str = "swe-agent-session"
77+
78+
# Commands to execute before agent initialization (e.g., bashrc setup, hosts config)
79+
pre_startup_bash_cmd_list: list[str] = env_vars.ROCK_AGENT_PRE_STARTUP_BASH_CMD_LIST
80+
81+
# Commands to execute after agent initialization
82+
post_startup_bash_cmd_list: list[str] = []
83+
84+
# Working directory where SWE-agent will be installed and executed
85+
swe_agent_workdir: str = "/tmp_sweagent"
86+
87+
# Command to download and set up Python environment
88+
python_install_cmd: str = env_vars.ROCK_AGENT_PYTHON_INSTALL_CMD
89+
90+
# Command to clone SWE-agent repository and install dependencies
91+
swe_agent_install_cmd: str = "[ -d SWE-agent ] && rm -rf SWE-agent; git clone https://github.com/SWE-agent/SWE-agent.git && cd SWE-agent && pip install -e . -i https://mirrors.aliyun.com/pypi/simple/"
92+
93+
python_install_timeout: int = 300
94+
95+
swe_agent_install_timeout: int = 600
96+
97+
agent_run_timeout: int = 1800
98+
99+
agent_run_check_interval: int = 30
100+
11101

12102
class SweAgent(Agent):
103+
"""
104+
SWE-agent implementation for automated software engineering tasks.
105+
106+
This class manages the complete lifecycle of SWE-agent including environment
107+
initialization, dependency installation, and task execution within a sandboxed
108+
environment. It provides an asynchronous interface for agent operations.
109+
110+
Attributes:
111+
config: Configuration parameters for agent setup and execution
112+
swe_session: Name of the bash session used for agent operations
113+
114+
Note:
115+
Currently requires a Sandbox instance (not AbstractSandbox).
116+
Only supports LocalDeployment and RunSingleConfig modes.
117+
"""
118+
13119
def __init__(self, sandbox: AbstractSandbox, config: SweAgentConfig):
120+
"""
121+
Initialize SWE-agent with sandbox environment and configuration.
122+
123+
Args:
124+
sandbox: Sandbox instance for isolated agent execution
125+
config: Configuration parameters for agent setup
126+
127+
Raises:
128+
AssertionError: If sandbox is not an instance of Sandbox class
129+
"""
14130
super().__init__(sandbox)
15131
self.config = config
132+
self.swe_session = self.config.swe_session
16133

17134
async def init(self):
18-
# Initialization logic for SWE agent
19-
pass
135+
"""
136+
Initialize the SWE-agent environment within the sandbox.
137+
138+
Performs the following initialization steps in sequence:
139+
1. Creates a dedicated bash session for agent execution
140+
2. Executes pre-startup configuration commands
141+
3. Creates working directory for agent installation
142+
4. Installs Python environment
143+
5. Clones and installs SWE-agent from GitHub repository
144+
145+
The initialization process is asynchronous and uses the configured
146+
timeouts for long-running operations like dependency installation.
147+
148+
Raises:
149+
Exception: If any initialization step fails
150+
"""
151+
assert isinstance(self._sandbox, Sandbox), "Sandbox must be an instance of Sandbox class"
152+
153+
sandbox_id = self._sandbox.sandbox_id
154+
155+
logger.info(f"[{sandbox_id}] Starting SWE-agent initialization")
156+
157+
# Step 1: Create dedicated bash session for agent operations
158+
logger.info(f"[{sandbox_id}] Creating bash session: {self.swe_session}")
159+
await self._sandbox.create_session(
160+
CreateBashSessionRequest(
161+
session=self.swe_session,
162+
env_enable=True,
163+
)
164+
)
165+
166+
# Step 2: Execute pre-startup configuration commands
167+
logger.info(f"[{sandbox_id}] Executing {len(self.config.pre_startup_bash_cmd_list)} pre-startup commands")
168+
for idx, cmd in enumerate(self.config.pre_startup_bash_cmd_list, 1):
169+
logger.debug(f"→ Pre-startup command {idx}/{len(self.config.pre_startup_bash_cmd_list)}: {cmd[:100]}...")
170+
await self._sandbox.arun(
171+
cmd=cmd,
172+
session=self.swe_session,
173+
)
174+
175+
# Step 3: Create working directory structure
176+
logger.info(f"[{sandbox_id}] Creating working directory: {self.config.swe_agent_workdir}")
177+
await self._sandbox.arun(
178+
cmd=f"mkdir -p {self.config.swe_agent_workdir}",
179+
session=self.swe_session,
180+
)
181+
182+
# Step 4: Install Python environment with retry
183+
logger.info(f"[{sandbox_id}] Installing Python environment")
184+
185+
python_install_cmd = f"cd {self.config.swe_agent_workdir} && {self.config.python_install_cmd}"
186+
await self._arun_with_retry(
187+
cmd=f"bash -c {shlex.quote(python_install_cmd)}",
188+
session=self.swe_session,
189+
mode="nohup",
190+
wait_timeout=self.config.python_install_timeout,
191+
error_msg="Python installation failed",
192+
)
193+
logger.info(f"[{sandbox_id}] Python installation completed")
194+
195+
# Step 5: Install SWE-agent repository with retry
196+
# Note: Temporarily using standalone pip from installed Python
197+
logger.info(f"[{sandbox_id}] Installing SWE-agent from repository")
198+
199+
swe_agent_install_cmd = f"export PATH={self.config.swe_agent_workdir}/python/bin:$PATH && cd {self.config.swe_agent_workdir} && {self.config.swe_agent_install_cmd}"
200+
await self._arun_with_retry(
201+
cmd=f"bash -c {shlex.quote(swe_agent_install_cmd)}",
202+
session=self.swe_session,
203+
mode="nohup",
204+
wait_timeout=self.config.swe_agent_install_timeout,
205+
error_msg="SWE-agent installation failed",
206+
)
207+
logger.info(f"[{sandbox_id}] SWE-agent installation completed successfully")
208+
209+
@retry_async(max_attempts=3, delay_seconds=5.0, backoff=2.0)
210+
async def _arun_with_retry(
211+
self,
212+
cmd: str,
213+
session: str,
214+
mode: str = "nohup",
215+
wait_timeout: int = 300,
216+
wait_interval: int = 10,
217+
error_msg: str = "Command failed",
218+
):
219+
"""
220+
Execute a command with retry logic based on exit code.
221+
222+
Args:
223+
cmd: Command to execute
224+
session: Session name to execute command in
225+
mode: Execution mode (normal, nohup, etc.)
226+
wait_timeout: Timeout for command execution
227+
wait_interval: Check interval for nohup commands
228+
error_msg: Error message to use when raising exception
229+
230+
Returns:
231+
Command result upon success
232+
"""
233+
result = await self._sandbox.arun(
234+
cmd=cmd, session=session, mode=mode, wait_timeout=wait_timeout, wait_interval=wait_interval
235+
)
236+
# If exit_code is not 0, raise an exception to trigger retry
237+
if result.exit_code != 0:
238+
raise Exception(f"{error_msg} with exit code: {result.exit_code}, output: {result.output}")
239+
return result
240+
241+
async def run(self, swe_agent_config_path: str | Path):
242+
"""
243+
Execute SWE-agent with the specified configuration file.
244+
245+
This method uploads the configuration file to the sandbox and executes
246+
SWE-agent with monitoring for completion. The execution runs in nohup
247+
mode with periodic status checks based on the configured interval.
248+
249+
Args:
250+
swe_agent_config_path: Local path to the SWE-agent configuration file
251+
(YAML format). The file will be uploaded to the
252+
sandbox before execution.
253+
254+
Returns:
255+
CommandResult: Execution result containing exit code, stdout, and stderr
256+
257+
Raises:
258+
AssertionError: If sandbox is not an instance of Sandbox class
259+
Exception: If file upload or command execution fails
260+
261+
Example:
262+
```python
263+
result = await agent.run("configs/swe_task.yaml")
264+
if result.exit_code == 0:
265+
print("Agent completed successfully")
266+
```
267+
"""
268+
assert isinstance(self._sandbox, Sandbox), "Sandbox must be an instance of Sandbox class"
269+
270+
logger.info(f"→ Starting SWE-agent execution with config: {swe_agent_config_path}")
271+
272+
config_filename = Path(swe_agent_config_path).name
273+
274+
# Upload configuration file to sandbox working directory
275+
logger.info(f"↑ Uploading configuration file: {config_filename}")
276+
await self._sandbox.upload(
277+
UploadRequest(
278+
source_path=os.path.abspath(swe_agent_config_path),
279+
target_path=f"{self.config.swe_agent_workdir}/{config_filename}",
280+
)
281+
)
282+
logger.debug(f"✓ Configuration file uploaded to: {self.config.swe_agent_workdir}/{config_filename}")
283+
284+
# Construct and execute SWE-agent run command
285+
swe_agent_run_cmd = f"cd {self.config.swe_agent_workdir} && {self.config.swe_agent_workdir}/python/bin/sweagent run --config {config_filename}"
286+
logger.info(
287+
f"▶ Executing SWE-agent (timeout: {self.config.agent_run_timeout}s, check interval: {self.config.agent_run_check_interval}s)"
288+
)
289+
290+
result = await self._sandbox.arun(
291+
cmd=f"bash -c {shlex.quote(swe_agent_run_cmd)}",
292+
session=self.swe_session,
293+
mode="nohup",
294+
wait_timeout=self.config.agent_run_timeout,
295+
wait_interval=self.config.agent_run_check_interval,
296+
)
297+
298+
# Log execution outcome
299+
if result.exit_code == 0:
300+
logger.info(f"✓ SWE-agent completed successfully (exit_code: {result.exit_code})")
301+
else:
302+
logger.error(f"✗ SWE-agent failed with exit_code: {result.exit_code}")
20303

21-
async def run(self, **kwargs):
22-
# Execution logic for SWE agent
23-
pass
304+
return result

uv.lock

Lines changed: 7 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)