client.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799
  1. """
  2. This is the main file for the runtime client.
  3. It is responsible for executing actions received from OpenHands backend and producing observations.
  4. NOTE: this will be executed inside the docker sandbox.
  5. """
  6. import argparse
  7. import asyncio
  8. import os
  9. import re
  10. import shutil
  11. import subprocess
  12. import time
  13. from contextlib import asynccontextmanager
  14. from pathlib import Path
  15. import pexpect
  16. from fastapi import Depends, FastAPI, HTTPException, Request, UploadFile
  17. from fastapi.exceptions import RequestValidationError
  18. from fastapi.responses import JSONResponse
  19. from fastapi.security import APIKeyHeader
  20. from pydantic import BaseModel
  21. from starlette.exceptions import HTTPException as StarletteHTTPException
  22. from uvicorn import run
  23. from openhands.core.logger import openhands_logger as logger
  24. from openhands.events.action import (
  25. Action,
  26. BrowseInteractiveAction,
  27. BrowseURLAction,
  28. CmdRunAction,
  29. FileReadAction,
  30. FileWriteAction,
  31. IPythonRunCellAction,
  32. )
  33. from openhands.events.observation import (
  34. CmdOutputObservation,
  35. ErrorObservation,
  36. FileReadObservation,
  37. FileWriteObservation,
  38. IPythonRunCellObservation,
  39. Observation,
  40. )
  41. from openhands.events.serialization import event_from_dict, event_to_dict
  42. from openhands.runtime.browser import browse
  43. from openhands.runtime.browser.browser_env import BrowserEnv
  44. from openhands.runtime.plugins import (
  45. ALL_PLUGINS,
  46. JupyterPlugin,
  47. Plugin,
  48. )
  49. from openhands.runtime.utils import split_bash_commands
  50. from openhands.runtime.utils.files import insert_lines, read_lines
  51. class ActionRequest(BaseModel):
  52. action: dict
  53. ROOT_GID = 0
  54. INIT_COMMANDS = [
  55. 'git config --global user.name "openhands" && git config --global user.email "openhands@all-hands.dev" && alias git="git --no-pager"',
  56. ]
  57. SOFT_TIMEOUT_SECONDS = 5
  58. SESSION_API_KEY = os.environ.get('SESSION_API_KEY')
  59. api_key_header = APIKeyHeader(name='X-Session-API-Key', auto_error=False)
  60. def verify_api_key(api_key: str = Depends(api_key_header)):
  61. if SESSION_API_KEY and api_key != SESSION_API_KEY:
  62. raise HTTPException(status_code=403, detail='Invalid API Key')
  63. return api_key
  64. class RuntimeClient:
  65. """RuntimeClient is running inside docker sandbox.
  66. It is responsible for executing actions received from OpenHands backend and producing observations.
  67. """
  68. def __init__(
  69. self,
  70. plugins_to_load: list[Plugin],
  71. work_dir: str,
  72. username: str,
  73. user_id: int,
  74. browsergym_eval_env: str | None,
  75. ) -> None:
  76. self.plugins_to_load = plugins_to_load
  77. self.username = username
  78. self.user_id = user_id
  79. self.pwd = work_dir # current PWD
  80. self._initial_pwd = work_dir
  81. self._init_user(self.username, self.user_id)
  82. self._init_bash_shell(self.pwd, self.username)
  83. self.lock = asyncio.Lock()
  84. self.plugins: dict[str, Plugin] = {}
  85. self.browser = BrowserEnv(browsergym_eval_env)
  86. self.start_time = time.time()
  87. self.last_execution_time = self.start_time
  88. @property
  89. def initial_pwd(self):
  90. return self._initial_pwd
  91. async def ainit(self):
  92. for plugin in self.plugins_to_load:
  93. await plugin.initialize(self.username)
  94. self.plugins[plugin.name] = plugin
  95. logger.info(f'Initializing plugin: {plugin.name}')
  96. if isinstance(plugin, JupyterPlugin):
  97. await self.run_ipython(
  98. IPythonRunCellAction(code=f'import os; os.chdir("{self.pwd}")')
  99. )
  100. # This is a temporary workaround
  101. # TODO: refactor AgentSkills to be part of JupyterPlugin
  102. # AFTER ServerRuntime is deprecated
  103. if 'agent_skills' in self.plugins and 'jupyter' in self.plugins:
  104. obs = await self.run_ipython(
  105. IPythonRunCellAction(
  106. code='from openhands.runtime.plugins.agent_skills.agentskills import *\n'
  107. )
  108. )
  109. logger.info(f'AgentSkills initialized: {obs}')
  110. await self._init_bash_commands()
  111. logger.info('Runtime client initialized.')
  112. def _init_user(self, username: str, user_id: int) -> None:
  113. """Create working directory and user if not exists.
  114. It performs the following steps effectively:
  115. * Creates the Working Directory:
  116. - Uses mkdir -p to create the directory.
  117. - Sets ownership to username:root.
  118. - Adjusts permissions to be readable and writable by group and others.
  119. * User Verification and Creation:
  120. - Checks if the user exists using id -u.
  121. - If the user exists with the correct UID, it skips creation.
  122. - If the UID differs, it logs a warning and updates self.user_id.
  123. - If the user doesn't exist, it proceeds to create the user.
  124. * Sudo Configuration:
  125. - Appends %sudo ALL=(ALL) NOPASSWD:ALL to /etc/sudoers to grant
  126. passwordless sudo access to the sudo group.
  127. - Adds the user to the sudo group with the useradd command, handling
  128. UID conflicts by incrementing the UID if necessary.
  129. """
  130. # First create the working directory, independent of the user
  131. logger.info(f'Client working directory: {self.initial_pwd}')
  132. command = f'umask 002; mkdir -p {self.initial_pwd}'
  133. output = subprocess.run(command, shell=True, capture_output=True)
  134. out_str = output.stdout.decode()
  135. command = f'chown -R {username}:root {self.initial_pwd}'
  136. output = subprocess.run(command, shell=True, capture_output=True)
  137. out_str += output.stdout.decode()
  138. command = f'chmod g+rw {self.initial_pwd}'
  139. output = subprocess.run(command, shell=True, capture_output=True)
  140. out_str += output.stdout.decode()
  141. logger.debug(f'Created working directory. Output: [{out_str}]')
  142. # Skip root since it is already created
  143. if username == 'root':
  144. return
  145. # Check if the username already exists
  146. existing_user_id = -1
  147. try:
  148. result = subprocess.run(
  149. f'id -u {username}', shell=True, check=True, capture_output=True
  150. )
  151. existing_user_id = int(result.stdout.decode().strip())
  152. # The user ID already exists, skip setup
  153. if existing_user_id == user_id:
  154. logger.debug(
  155. f'User `{username}` already has the provided UID {user_id}. Skipping user setup.'
  156. )
  157. else:
  158. logger.warning(
  159. f'User `{username}` already exists with UID {existing_user_id}. Skipping user setup.'
  160. )
  161. self.user_id = existing_user_id
  162. return
  163. except subprocess.CalledProcessError as e:
  164. # Returncode 1 indicates, that the user does not exist yet
  165. if e.returncode == 1:
  166. logger.debug(
  167. f'User `{username}` does not exist. Proceeding with user creation.'
  168. )
  169. else:
  170. logger.error(
  171. f'Error checking user `{username}`, skipping setup:\n{e}\n'
  172. )
  173. raise
  174. # Add sudoer
  175. sudoer_line = r"echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers"
  176. output = subprocess.run(sudoer_line, shell=True, capture_output=True)
  177. if output.returncode != 0:
  178. raise RuntimeError(f'Failed to add sudoer: {output.stderr.decode()}')
  179. logger.debug(f'Added sudoer successfully. Output: [{output.stdout.decode()}]')
  180. command = (
  181. f'useradd -rm -d /home/{username} -s /bin/bash '
  182. f'-g root -G sudo -u {user_id} {username}'
  183. )
  184. output = subprocess.run(command, shell=True, capture_output=True)
  185. if output.returncode == 0:
  186. logger.debug(
  187. f'Added user `{username}` successfully with UID {user_id}. Output: [{output.stdout.decode()}]'
  188. )
  189. else:
  190. raise RuntimeError(
  191. f'Failed to create user `{username}` with UID {user_id}. Output: [{output.stderr.decode()}]'
  192. )
  193. def _init_bash_shell(self, work_dir: str, username: str) -> None:
  194. self.shell = pexpect.spawn(
  195. f'su {username}',
  196. encoding='utf-8',
  197. echo=False,
  198. )
  199. self.__bash_PS1 = (
  200. r'[PEXPECT_BEGIN]\n'
  201. r'$(which python >/dev/null 2>&1 && echo "[Python Interpreter: $(which python)]\n")'
  202. r'\u@\h:\w\n'
  203. r'[PEXPECT_END]'
  204. )
  205. # This should NOT match "PS1=\u@\h:\w [PEXPECT]$" when `env` is executed
  206. self.__bash_expect_regex = r'\[PEXPECT_BEGIN\]\s*(.*?)\s*([a-z0-9_-]*)@([a-zA-Z0-9.-]*):(.+)\s*\[PEXPECT_END\]'
  207. # Set umask to allow group write permissions
  208. self.shell.sendline(f'umask 002; export PS1="{self.__bash_PS1}"; export PS2=""')
  209. self.shell.expect(self.__bash_expect_regex)
  210. self.shell.sendline(
  211. f'if [ ! -d "{work_dir}" ]; then mkdir -p "{work_dir}"; fi && cd "{work_dir}"'
  212. )
  213. self.shell.expect(self.__bash_expect_regex)
  214. logger.debug(
  215. f'Bash initialized. Working directory: {work_dir}. Output: [{self.shell.before}]'
  216. )
  217. # Ensure the group has write permissions on the working directory
  218. self.shell.sendline(f'chmod g+rw "{work_dir}"')
  219. self.shell.expect(self.__bash_expect_regex)
  220. async def _init_bash_commands(self):
  221. logger.info(f'Initializing by running {len(INIT_COMMANDS)} bash commands...')
  222. for command in INIT_COMMANDS:
  223. action = CmdRunAction(command=command)
  224. action.timeout = 300
  225. logger.debug(f'Executing init command: {command}')
  226. obs: CmdOutputObservation = await self.run(action)
  227. logger.debug(
  228. f'Init command outputs (exit code: {obs.exit_code}): {obs.content}'
  229. )
  230. assert obs.exit_code == 0
  231. logger.info('Bash init commands completed')
  232. def _get_bash_prompt_and_update_pwd(self):
  233. ps1 = self.shell.after
  234. if ps1 == pexpect.EOF:
  235. logger.error(f'Bash shell EOF! {self.shell.after=}, {self.shell.before=}')
  236. raise RuntimeError('Bash shell EOF')
  237. if ps1 == pexpect.TIMEOUT:
  238. logger.warning('Bash shell timeout')
  239. return ''
  240. # begin at the last occurrence of '[PEXPECT_BEGIN]'.
  241. # In multi-line bash commands, the prompt will be repeated
  242. # and the matched regex captures all of them
  243. # - we only want the last one (newest prompt)
  244. _begin_pos = ps1.rfind('[PEXPECT_BEGIN]')
  245. if _begin_pos != -1:
  246. ps1 = ps1[_begin_pos:]
  247. # parse the ps1 to get username, hostname, and working directory
  248. matched = re.match(self.__bash_expect_regex, ps1)
  249. assert (
  250. matched is not None
  251. ), f'Failed to parse bash prompt: {ps1}. This should not happen.'
  252. other_info, username, hostname, working_dir = matched.groups()
  253. working_dir = working_dir.rstrip()
  254. self.pwd = os.path.expanduser(working_dir)
  255. # re-assemble the prompt
  256. prompt = f'{other_info.strip()}\n{username}@{hostname}:{working_dir} '
  257. if username == 'root':
  258. prompt += '#'
  259. else:
  260. prompt += '$'
  261. return prompt + ' '
  262. def _execute_bash(
  263. self,
  264. command: str,
  265. timeout: int | None,
  266. keep_prompt: bool = True,
  267. kill_on_timeout: bool = True,
  268. ) -> tuple[str, int]:
  269. logger.debug(f'Executing command: {command}')
  270. self.shell.sendline(command)
  271. return self._continue_bash(
  272. timeout=timeout, keep_prompt=keep_prompt, kill_on_timeout=kill_on_timeout
  273. )
  274. def _interrupt_bash(self, timeout: int | None = None) -> tuple[str, int]:
  275. self.shell.sendintr() # send SIGINT to the shell
  276. self.shell.expect(self.__bash_expect_regex, timeout=timeout)
  277. output = self.shell.before
  278. exit_code = 130 # SIGINT
  279. return output, exit_code
  280. def _continue_bash(
  281. self,
  282. timeout: int | None,
  283. keep_prompt: bool = True,
  284. kill_on_timeout: bool = True,
  285. ) -> tuple[str, int]:
  286. try:
  287. self.shell.expect(self.__bash_expect_regex, timeout=timeout)
  288. output = self.shell.before
  289. # Get exit code
  290. self.shell.sendline('echo $?')
  291. logger.debug('Requesting exit code...')
  292. self.shell.expect(self.__bash_expect_regex, timeout=timeout)
  293. _exit_code_output = self.shell.before
  294. try:
  295. exit_code = int(_exit_code_output.strip().split()[0])
  296. except Exception:
  297. logger.error('Error getting exit code from bash script')
  298. # If we try to run an invalid shell script the output sometimes includes error text
  299. # rather than the error code - we assume this is an error
  300. exit_code = 2
  301. except pexpect.TIMEOUT as e:
  302. if kill_on_timeout:
  303. output, exit_code = self._interrupt_bash()
  304. output += (
  305. '\r\n\r\n'
  306. + f'[Command timed out after {timeout} seconds. SIGINT was sent to interrupt it.]'
  307. )
  308. logger.error(f'Failed to execute command. Error: {e}')
  309. else:
  310. output = self.shell.before or ''
  311. exit_code = -1
  312. finally:
  313. bash_prompt = self._get_bash_prompt_and_update_pwd()
  314. if keep_prompt:
  315. output += '\r\n' + bash_prompt
  316. # logger.debug(f'Command output:\n{output}')
  317. return output, exit_code
  318. async def run_action(self, action) -> Observation:
  319. action_type = action.action
  320. logger.debug(f'Running action:\n{action}')
  321. observation = await getattr(self, action_type)(action)
  322. logger.debug(f'Action output:\n{observation}')
  323. return observation
  324. async def run(self, action: CmdRunAction) -> CmdOutputObservation:
  325. try:
  326. assert (
  327. action.timeout is not None
  328. ), f'Timeout argument is required for CmdRunAction: {action}'
  329. commands = split_bash_commands(action.command)
  330. all_output = ''
  331. for command in commands:
  332. if command == '':
  333. output, exit_code = self._continue_bash(
  334. timeout=SOFT_TIMEOUT_SECONDS,
  335. keep_prompt=action.keep_prompt,
  336. kill_on_timeout=False,
  337. )
  338. elif command.lower() == 'ctrl+c':
  339. output, exit_code = self._interrupt_bash(
  340. timeout=SOFT_TIMEOUT_SECONDS
  341. )
  342. else:
  343. output, exit_code = self._execute_bash(
  344. command,
  345. timeout=SOFT_TIMEOUT_SECONDS
  346. if not action.blocking
  347. else action.timeout,
  348. keep_prompt=action.keep_prompt,
  349. kill_on_timeout=False if not action.blocking else True,
  350. )
  351. if all_output:
  352. # previous output already exists with prompt "user@hostname:working_dir #""
  353. # we need to add the command to the previous output,
  354. # so model knows the following is the output of another action)
  355. all_output = all_output.rstrip() + ' ' + command + '\r\n'
  356. all_output += str(output) + '\r\n'
  357. if exit_code != 0:
  358. break
  359. return CmdOutputObservation(
  360. command_id=-1,
  361. content=all_output.rstrip('\r\n'),
  362. command=action.command,
  363. exit_code=exit_code,
  364. )
  365. except UnicodeDecodeError:
  366. raise RuntimeError('Command output could not be decoded as utf-8')
  367. async def run_ipython(self, action: IPythonRunCellAction) -> Observation:
  368. if 'jupyter' in self.plugins:
  369. _jupyter_plugin: JupyterPlugin = self.plugins['jupyter'] # type: ignore
  370. # This is used to make AgentSkills in Jupyter aware of the
  371. # current working directory in Bash
  372. jupyter_pwd = getattr(self, '_jupyter_pwd', None)
  373. if self.pwd != jupyter_pwd:
  374. logger.debug(f'{self.pwd} != {jupyter_pwd} -> reset Jupyter PWD')
  375. reset_jupyter_pwd_code = f'import os; os.chdir("{self.pwd}")'
  376. _aux_action = IPythonRunCellAction(code=reset_jupyter_pwd_code)
  377. _reset_obs = await _jupyter_plugin.run(_aux_action)
  378. logger.debug(
  379. f'Changed working directory in IPython to: {self.pwd}. Output: {_reset_obs}'
  380. )
  381. self._jupyter_pwd = self.pwd
  382. obs: IPythonRunCellObservation = await _jupyter_plugin.run(action)
  383. obs.content = obs.content.rstrip()
  384. obs.content += f'\n[Jupyter current working directory: {self.pwd}]'
  385. obs.content += f'\n[Jupyter Python interpreter: {_jupyter_plugin.python_interpreter_path}]'
  386. return obs
  387. else:
  388. raise RuntimeError(
  389. 'JupyterRequirement not found. Unable to run IPython action.'
  390. )
  391. def _get_working_directory(self):
  392. # NOTE: this is part of initialization, so we hard code the timeout
  393. result, exit_code = self._execute_bash('pwd', timeout=60, keep_prompt=False)
  394. if exit_code != 0:
  395. raise RuntimeError('Failed to get working directory')
  396. return result.strip()
  397. def _resolve_path(self, path: str, working_dir: str) -> str:
  398. filepath = Path(path)
  399. if not filepath.is_absolute():
  400. return str(Path(working_dir) / filepath)
  401. return str(filepath)
  402. async def read(self, action: FileReadAction) -> Observation:
  403. # NOTE: the client code is running inside the sandbox,
  404. # so there's no need to check permission
  405. working_dir = self._get_working_directory()
  406. filepath = self._resolve_path(action.path, working_dir)
  407. try:
  408. with open(filepath, 'r', encoding='utf-8') as file:
  409. lines = read_lines(file.readlines(), action.start, action.end)
  410. except FileNotFoundError:
  411. return ErrorObservation(
  412. f'File not found: {filepath}. Your current working directory is {working_dir}.'
  413. )
  414. except UnicodeDecodeError:
  415. return ErrorObservation(f'File could not be decoded as utf-8: {filepath}.')
  416. except IsADirectoryError:
  417. return ErrorObservation(
  418. f'Path is a directory: {filepath}. You can only read files'
  419. )
  420. code_view = ''.join(lines)
  421. return FileReadObservation(path=filepath, content=code_view)
  422. async def write(self, action: FileWriteAction) -> Observation:
  423. working_dir = self._get_working_directory()
  424. filepath = self._resolve_path(action.path, working_dir)
  425. insert = action.content.split('\n')
  426. try:
  427. if not os.path.exists(os.path.dirname(filepath)):
  428. os.makedirs(os.path.dirname(filepath))
  429. file_exists = os.path.exists(filepath)
  430. if file_exists:
  431. file_stat = os.stat(filepath)
  432. else:
  433. file_stat = None
  434. mode = 'w' if not file_exists else 'r+'
  435. try:
  436. with open(filepath, mode, encoding='utf-8') as file:
  437. if mode != 'w':
  438. all_lines = file.readlines()
  439. new_file = insert_lines(
  440. insert, all_lines, action.start, action.end
  441. )
  442. else:
  443. new_file = [i + '\n' for i in insert]
  444. file.seek(0)
  445. file.writelines(new_file)
  446. file.truncate()
  447. # Handle file permissions
  448. if file_exists:
  449. assert file_stat is not None
  450. # restore the original file permissions if the file already exists
  451. os.chmod(filepath, file_stat.st_mode)
  452. os.chown(filepath, file_stat.st_uid, file_stat.st_gid)
  453. else:
  454. # set the new file permissions if the file is new
  455. os.chmod(filepath, 0o664)
  456. os.chown(filepath, self.user_id, self.user_id)
  457. except FileNotFoundError:
  458. return ErrorObservation(f'File not found: {filepath}')
  459. except IsADirectoryError:
  460. return ErrorObservation(
  461. f'Path is a directory: {filepath}. You can only write to files'
  462. )
  463. except UnicodeDecodeError:
  464. return ErrorObservation(
  465. f'File could not be decoded as utf-8: {filepath}'
  466. )
  467. except PermissionError:
  468. return ErrorObservation(f'Malformed paths not permitted: {filepath}')
  469. return FileWriteObservation(content='', path=filepath)
  470. async def browse(self, action: BrowseURLAction) -> Observation:
  471. return await browse(action, self.browser)
  472. async def browse_interactive(self, action: BrowseInteractiveAction) -> Observation:
  473. return await browse(action, self.browser)
  474. def close(self):
  475. self.shell.close()
  476. self.browser.close()
  477. if __name__ == '__main__':
  478. parser = argparse.ArgumentParser()
  479. parser.add_argument('port', type=int, help='Port to listen on')
  480. parser.add_argument('--working-dir', type=str, help='Working directory')
  481. parser.add_argument('--plugins', type=str, help='Plugins to initialize', nargs='+')
  482. parser.add_argument(
  483. '--username', type=str, help='User to run as', default='openhands'
  484. )
  485. parser.add_argument('--user-id', type=int, help='User ID to run as', default=1000)
  486. parser.add_argument(
  487. '--browsergym-eval-env',
  488. type=str,
  489. help='BrowserGym environment used for browser evaluation',
  490. default=None,
  491. )
  492. # example: python client.py 8000 --working-dir /workspace --plugins JupyterRequirement
  493. args = parser.parse_args()
  494. plugins_to_load: list[Plugin] = []
  495. if args.plugins:
  496. for plugin in args.plugins:
  497. if plugin not in ALL_PLUGINS:
  498. raise ValueError(f'Plugin {plugin} not found')
  499. plugins_to_load.append(ALL_PLUGINS[plugin]()) # type: ignore
  500. client: RuntimeClient | None = None
  501. @asynccontextmanager
  502. async def lifespan(app: FastAPI):
  503. global client
  504. client = RuntimeClient(
  505. plugins_to_load,
  506. work_dir=args.working_dir,
  507. username=args.username,
  508. user_id=args.user_id,
  509. browsergym_eval_env=args.browsergym_eval_env,
  510. )
  511. await client.ainit()
  512. yield
  513. # Clean up & release the resources
  514. client.close()
  515. app = FastAPI(lifespan=lifespan)
  516. # TODO below 3 exception handlers were recommended by Sonnet.
  517. # Are these something we should keep?
  518. @app.exception_handler(Exception)
  519. async def global_exception_handler(request: Request, exc: Exception):
  520. logger.exception('Unhandled exception occurred:')
  521. return JSONResponse(
  522. status_code=500,
  523. content={
  524. 'message': 'An unexpected error occurred. Please try again later.'
  525. },
  526. )
  527. @app.exception_handler(StarletteHTTPException)
  528. async def http_exception_handler(request: Request, exc: StarletteHTTPException):
  529. logger.error(f'HTTP exception occurred: {exc.detail}')
  530. return JSONResponse(
  531. status_code=exc.status_code, content={'message': exc.detail}
  532. )
  533. @app.exception_handler(RequestValidationError)
  534. async def validation_exception_handler(
  535. request: Request, exc: RequestValidationError
  536. ):
  537. logger.error(f'Validation error occurred: {exc}')
  538. return JSONResponse(
  539. status_code=422,
  540. content={'message': 'Invalid request parameters', 'details': exc.errors()},
  541. )
  542. @app.middleware('http')
  543. async def one_request_at_a_time(request: Request, call_next):
  544. assert client is not None
  545. async with client.lock:
  546. response = await call_next(request)
  547. return response
  548. @app.middleware('http')
  549. async def authenticate_requests(request: Request, call_next):
  550. if request.url.path != '/alive' and request.url.path != '/server_info':
  551. try:
  552. verify_api_key(request.headers.get('X-Session-API-Key'))
  553. except HTTPException as e:
  554. return e
  555. response = await call_next(request)
  556. return response
  557. @app.get('/server_info')
  558. async def get_server_info():
  559. assert client is not None
  560. current_time = time.time()
  561. uptime = current_time - client.start_time
  562. idle_time = current_time - client.last_execution_time
  563. return {'uptime': uptime, 'idle_time': idle_time}
  564. @app.post('/execute_action')
  565. async def execute_action(action_request: ActionRequest):
  566. assert client is not None
  567. try:
  568. action = event_from_dict(action_request.action)
  569. if not isinstance(action, Action):
  570. raise HTTPException(status_code=400, detail='Invalid action type')
  571. client.last_execution_time = time.time()
  572. observation = await client.run_action(action)
  573. return event_to_dict(observation)
  574. except Exception as e:
  575. logger.error(
  576. f'Error processing command: {str(e)}', exc_info=True, stack_info=True
  577. )
  578. raise HTTPException(status_code=500, detail=str(e))
  579. @app.post('/upload_file')
  580. async def upload_file(
  581. file: UploadFile, destination: str = '/', recursive: bool = False
  582. ):
  583. assert client is not None
  584. try:
  585. # Ensure the destination directory exists
  586. if not os.path.isabs(destination):
  587. raise HTTPException(
  588. status_code=400, detail='Destination must be an absolute path'
  589. )
  590. full_dest_path = destination
  591. if not os.path.exists(full_dest_path):
  592. os.makedirs(full_dest_path, exist_ok=True)
  593. if recursive:
  594. # For recursive uploads, we expect a zip file
  595. if not file.filename.endswith('.zip'):
  596. raise HTTPException(
  597. status_code=400, detail='Recursive uploads must be zip files'
  598. )
  599. zip_path = os.path.join(full_dest_path, file.filename)
  600. with open(zip_path, 'wb') as buffer:
  601. shutil.copyfileobj(file.file, buffer)
  602. # Extract the zip file
  603. shutil.unpack_archive(zip_path, full_dest_path)
  604. os.remove(zip_path) # Remove the zip file after extraction
  605. logger.info(
  606. f'Uploaded file {file.filename} and extracted to {destination}'
  607. )
  608. else:
  609. # For single file uploads
  610. file_path = os.path.join(full_dest_path, file.filename)
  611. with open(file_path, 'wb') as buffer:
  612. shutil.copyfileobj(file.file, buffer)
  613. logger.info(f'Uploaded file {file.filename} to {destination}')
  614. return JSONResponse(
  615. content={
  616. 'filename': file.filename,
  617. 'destination': destination,
  618. 'recursive': recursive,
  619. },
  620. status_code=200,
  621. )
  622. except Exception as e:
  623. raise HTTPException(status_code=500, detail=str(e))
  624. @app.get('/alive')
  625. async def alive():
  626. return {'status': 'ok'}
  627. # ================================
  628. # File-specific operations for UI
  629. # ================================
  630. @app.post('/list_files')
  631. async def list_files(request: Request):
  632. """List files in the specified path.
  633. This function retrieves a list of files from the agent's runtime file store,
  634. excluding certain system and hidden files/directories.
  635. To list files:
  636. ```sh
  637. curl http://localhost:3000/api/list-files
  638. ```
  639. Args:
  640. request (Request): The incoming request object.
  641. path (str, optional): The path to list files from. Defaults to '/'.
  642. Returns:
  643. list: A list of file names in the specified path.
  644. Raises:
  645. HTTPException: If there's an error listing the files.
  646. """
  647. assert client is not None
  648. # get request as dict
  649. request_dict = await request.json()
  650. path = request_dict.get('path', None)
  651. # Get the full path of the requested directory
  652. if path is None:
  653. full_path = client.initial_pwd
  654. elif os.path.isabs(path):
  655. full_path = path
  656. else:
  657. full_path = os.path.join(client.initial_pwd, path)
  658. if not os.path.exists(full_path):
  659. # if user just removed a folder, prevent server error 500 in UI
  660. return []
  661. try:
  662. # Check if the directory exists
  663. if not os.path.exists(full_path) or not os.path.isdir(full_path):
  664. return []
  665. entries = os.listdir(full_path)
  666. # Separate directories and files
  667. directories = []
  668. files = []
  669. for entry in entries:
  670. # Remove leading slash and any parent directory components
  671. entry_relative = entry.lstrip('/').split('/')[-1]
  672. # Construct the full path by joining the base path with the relative entry path
  673. full_entry_path = os.path.join(full_path, entry_relative)
  674. if os.path.exists(full_entry_path):
  675. is_dir = os.path.isdir(full_entry_path)
  676. if is_dir:
  677. # add trailing slash to directories
  678. # required by FE to differentiate directories and files
  679. entry = entry.rstrip('/') + '/'
  680. directories.append(entry)
  681. else:
  682. files.append(entry)
  683. # Sort directories and files separately
  684. directories.sort(key=lambda s: s.lower())
  685. files.sort(key=lambda s: s.lower())
  686. # Combine sorted directories and files
  687. sorted_entries = directories + files
  688. return sorted_entries
  689. except Exception as e:
  690. logger.error(f'Error listing files: {e}', exc_info=True)
  691. return []
  692. logger.info('Runtime client initialized.')
  693. logger.info(f'Starting action execution API on port {args.port}')
  694. run(app, host='0.0.0.0', port=args.port)