command_manager.py 3.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. import subprocess
  2. import select
  3. from typing import List
  4. from opendevin.lib.event import Event
  5. class BackgroundCommand:
  6. def __init__(self, id: int, command: str, process: subprocess.Popen):
  7. self.command = command
  8. self.id = id
  9. self.process = process
  10. def _get_log_from_stream(self, stream):
  11. logs = ""
  12. while True:
  13. readable, _, _ = select.select([stream], [], [], .1)
  14. if not readable:
  15. break
  16. next = stream.readline()
  17. if next == '':
  18. break
  19. logs += next
  20. if logs == "": return
  21. return logs
  22. def get_logs(self):
  23. stdout = self._get_log_from_stream(self.process.stdout)
  24. stderr = self._get_log_from_stream(self.process.stderr)
  25. exit_code = self.process.poll()
  26. return stdout, stderr, exit_code
  27. class CommandManager:
  28. def __init__(self):
  29. self.cur_id = 0
  30. self.background_commands = {}
  31. def run_command(self, command: str, background=False) -> str:
  32. if background:
  33. return self.run_background(command)
  34. else:
  35. return self.run_immediately(command)
  36. def run_immediately(self, command: str) -> str:
  37. result = subprocess.run(["/bin/bash", "-c", command], capture_output=True, text=True)
  38. output = result.stdout + result.stderr
  39. exit_code = result.returncode
  40. if exit_code != 0:
  41. raise ValueError('Command failed with exit code ' + str(exit_code) + ': ' + output)
  42. return output
  43. def run_background(self, command: str) -> str:
  44. process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=True)
  45. bg_cmd = BackgroundCommand(self.cur_id, command, process)
  46. self.cur_id += 1
  47. self.background_commands[bg_cmd.id] = bg_cmd
  48. return "Background command started. To stop it, send a `kill` action with id " + str(bg_cmd.id)
  49. def kill_command(self, id: int) -> str:
  50. # TODO: get log events before killing
  51. self.background_commands[id].processs.kill()
  52. del self.background_commands[id]
  53. def get_background_events(self) -> List[Event]:
  54. events = []
  55. for id, cmd in self.background_commands.items():
  56. stdout, stderr, exit_code = cmd.get_logs()
  57. if stdout is not None:
  58. events.append(Event('output', {
  59. 'output': stdout,
  60. 'stream': 'stdout',
  61. 'id': id,
  62. 'command': cmd.command,
  63. }))
  64. if stderr is not None:
  65. events.append(Event('output', {
  66. 'output': stderr,
  67. 'stream': 'stderr',
  68. 'id': id,
  69. 'command': cmd.command,
  70. }))
  71. if exit_code is not None:
  72. events.append(Event('output', {
  73. 'exit_code': exit_code,
  74. 'output': 'Background command %d exited with code %d' % (idx, exit_code),
  75. 'id': id,
  76. 'command': cmd.command,
  77. }))
  78. del self.background_commands[id]
  79. return events