heaven_base.tools.bash_tool

  1import asyncio
  2import os
  3from typing import Optional, Dict, Any, ClassVar, Type
  4from langchain.tools import Tool, BaseTool
  5from langchain.callbacks.manager import AsyncCallbackManagerForToolRun
  6from langchain.schema.runnable import RunnableConfig
  7from collections.abc import Callable
  8
  9from ..baseheaventool import BaseHeavenTool, ToolResult, CLIResult, ToolError, ToolArgsSchema
 10
 11class _BashSession:
 12    """A session of a bash shell."""
 13
 14    _started: bool
 15    _process: asyncio.subprocess.Process
 16
 17    command: str = "/bin/bash"
 18    _output_delay: float = 0.2  # seconds
 19    _timeout: float = 120.0  # seconds
 20    _sentinel: str = "<<exit>>"
 21
 22    def __init__(self):
 23        self._started = False
 24        self._timed_out = False
 25
 26    async def start(self):
 27        if self._started:
 28            return
 29
 30        self._process = await asyncio.create_subprocess_shell(
 31            self.command,
 32            preexec_fn=os.setsid,
 33            shell=True,
 34            bufsize=0,
 35            stdin=asyncio.subprocess.PIPE,
 36            stdout=asyncio.subprocess.PIPE,
 37            stderr=asyncio.subprocess.PIPE,
 38        )
 39
 40        self._started = True
 41
 42    def stop(self):
 43        """Terminate the bash shell."""
 44        if not self._started:
 45            raise ToolError("ERROR: Session has not started.")
 46        if self._process.returncode is not None:
 47            return
 48        self._process.terminate()
 49
 50    async def run(self, command: str):
 51        """Execute a command in the bash shell."""
 52        if not self._started:
 53            raise ToolError("ERROR: Session has not started.")
 54        if self._process.returncode is not None:
 55            return ToolResult(
 56                system="tool must be restarted",
 57                error=f"bash has exited with returncode {self._process.returncode}",
 58            )
 59        if self._timed_out:
 60            raise ToolError(
 61                f"ERROR: timed out: bash has not returned in {self._timeout} seconds and must be restarted",
 62            )
 63
 64        # we know these are not None because we created the process with PIPEs
 65        assert self._process.stdin
 66        assert self._process.stdout
 67        assert self._process.stderr
 68
 69        # send command to the process
 70        self._process.stdin.write(
 71            command.encode() + f"; echo '{self._sentinel}'\n".encode()
 72        )
 73        await self._process.stdin.drain()
 74
 75        # read output from the process, until the sentinel is found
 76        try:
 77            async with asyncio.timeout(self._timeout):
 78                while True:
 79                    await asyncio.sleep(self._output_delay)
 80                    # if we read directly from stdout/stderr, it will wait forever for
 81                    # EOF. use the StreamReader buffer directly instead.
 82                    output = self._process.stdout._buffer.decode()  # pyright: ignore[reportAttributeAccessIssue]
 83                    if self._sentinel in output:
 84                        # strip the sentinel and break
 85                        output = output[: output.index(self._sentinel)]
 86                        break
 87        except asyncio.TimeoutError:
 88            self._timed_out = True
 89            raise ToolError(
 90                f"ERROR: timed out: bash has not returned in {self._timeout} seconds and must be restarted",
 91            ) from None
 92
 93        if output.endswith("\n"):
 94            output = output[:-1]
 95
 96        error = self._process.stderr._buffer.decode()  # pyright: ignore[reportAttributeAccessIssue]
 97        if error.endswith("\n"):
 98            error = error[:-1]
 99
100        # clear the buffers so that the next output can be read correctly
101        self._process.stdout._buffer.clear()  # pyright: ignore[reportAttributeAccessIssue]
102        self._process.stderr._buffer.clear()  # pyright: ignore[reportAttributeAccessIssue]
103
104        return CLIResult(output=output, error=error)
105
106        # # read output from the process, until the sentinel is found
107        # try:
108        #     async with asyncio.timeout(self._timeout):
109        #         while True:
110        #             await asyncio.sleep(self._output_delay)
111        #             # if we read directly from stdout/stderr, it will wait forever for
112        #             # EOF. use the StreamReader buffer directly instead.
113        #             output = self._process.stdout._buffer.decode()  # pyright: ignore[reportAttributeAccessIssue]
114        #             if self._sentinel in output:
115        #                 # strip the sentinel and break
116        #                 output = output[: output.index(self._sentinel)]
117        #                 break
118        # except asyncio.TimeoutError:
119        #     self._timed_out = True
120        #     raise ToolError(
121        #         f"timed out: bash has not returned in {self._timeout} seconds and must be restarted",
122        #     ) from None
123
124        # if output.endswith("\n"):
125        #     output = output[:-1]
126
127        # error = self._process.stderr._buffer.decode()  # pyright: ignore[reportAttributeAccessIssue]
128        # if error.endswith("\n"):
129        #     error = error[:-1]
130
131        # # clear the buffers so that the next output can be read correctly
132        # self._process.stdout._buffer.clear()  # pyright: ignore[reportAttributeAccessIssue]
133        # self._process.stderr._buffer.clear()  # pyright: ignore[reportAttributeAccessIssue]
134
135        # return CLIResult(output=output, error=error)
136
137
138# OLD
139# class BashToolArgsSchema(ToolArgsSchema):
140#     arguments: Dict[str, Dict[str, Any]] = {
141#         'command': {
142#             'name': 'command',
143#             'type': 'str',
144#             'description': 'The bash command to run. Required unless the tool is being restarted. For running files inside /core/, you may need to set PYTHONPATH=/home/GOD/core'
145#         },
146#         'restart': {
147#             'name': 'restart',
148#             'type': 'bool',
149#             'description': 'Specifying true will restart this tool. Otherwise, leave this unspecified. If you receive a timeout error, BashTool must be restarted without any command before it can be used again.'
150#         }
151#     }
152
153class BashToolArgsSchema(ToolArgsSchema):
154    arguments: Dict[str, Dict[str, Any]] = {
155        'command': {
156            'name': 'command',
157            'type': 'str',
158            'description': 'The bash command to run. Required unless restart is set to true. For running files inside /core/, you may need to set PYTHONPATH=/home/GOD/core',
159            'required': False  # Not required when restarting
160        },
161        'restart': {
162            'name': 'restart',
163            'type': 'bool',
164            'description': 'Specifying true will restart this tool. Otherwise, leave this unspecified or false. If you receive a timeout error, BashTool must be restarted without any command before it can be used again.',
165            'required': False  # Defaults to false
166        }
167    }
168
169# class BashTool(BaseHeavenTool):
170#     name = "BashTool"
171#     description = "Run commands in a bash shell"
172#     args_schema = BashToolArgsSchema
173#     is_async = True
174
175#     def __init__(self, base_tool: BaseTool, args_schema: Type[ToolArgsSchema], is_async: bool = False):
176#         super().__init__(base_tool=base_tool, args_schema=args_schema, is_async=is_async)
177#         self._bash_session = _BashSession()  # Each instance gets its own session
178
179#     @classmethod
180#     def create(cls, adk: bool = False):
181#         async def wrapped_func(command: str | None = None, restart: bool = False):
182#             # Get the instance from the bound function
183#             self = wrapped_func.__self__
184
185#             if restart:
186#                 if self._bash_session and self._bash_session._started:
187#                     self._bash_session.stop()
188#                 self._bash_session = _BashSession()
189#                 await self._bash_session.start()
190#                 return "tool has been restarted."
191
192#             if not self._bash_session._started:
193#                 await self._bash_session.start()
194
195#             if command is not None:
196#                 return await self._bash_session.run(command)
197
198#             raise ToolError("ERROR: no command provided.")
199
200#         cls.func = wrapped_func
201#         instance = super().create()
202#         wrapped_func.__self__ = instance
203#         return instance
204
205# With ADK
206class BashTool(BaseHeavenTool):
207    name = "BashTool"
208    description = "Run commands in a bash shell"
209    args_schema = BashToolArgsSchema
210    is_async = True
211
212    def __init__(self, base_tool: BaseTool, args_schema: Type[ToolArgsSchema], is_async: bool = False):
213        super().__init__(base_tool=base_tool, args_schema=args_schema, is_async=is_async)
214        self._bash_session = _BashSession()
215
216    @classmethod
217    def create(cls, adk: bool = False):
218        session = _BashSession()
219    
220        async def wrapped_func(command: Optional[str] = None,
221                               restart: Optional[bool] = None):
222            nonlocal session
223            if restart:
224                if session._started:
225                    session.stop()
226                session = _BashSession()
227                await session.start()
228                return "tool has been restarted."
229            if not session._started:
230                await session.start()
231            if command is not None:
232                return await session.run(command)
233            raise ToolError("ERROR: no command provided.")
234    
235        cls.func = wrapped_func
236        instance = super().create(adk=adk)
237        wrapped_func.__self__ = instance
238        return instance
239
240
241
242
243
244      
245      
246# #### Has timeout problem
247# class BashTool(BaseHeavenTool):
248#     name = "BashTool"
249#     description = "Run commands in a bash shell"
250#     args_schema = BashToolArgsSchema
251#     is_async = True
252
253#     _bash_session = None  # Class-level singleton
254
255#     @classmethod
256#     def get_session(cls):
257#         if cls._bash_session is None:
258#             cls._bash_session = _BashSession()
259#         return cls._bash_session
260
261#     @classmethod
262#     def create(cls):
263#         # Get/create the persistent session
264#         session = cls.get_session()
265
266#         async def wrapped_func(command: str | None = None, restart: bool = False):
267#             if restart:
268#                 if cls._bash_session and cls._bash_session._started:
269#                     cls._bash_session.stop()
270#                 cls._bash_session = None
271#                 cls._bash_session = cls.get_session()
272#                 await cls._bash_session.start()
273#                 return "tool has been restarted."
274
275#             if not cls._bash_session._started:
276#                 await cls._bash_session.start()
277
278#             if command is not None:
279#                 return await cls._bash_session.run(command)
280
281#             raise ToolError("no command provided.")
282
283#         cls.func = wrapped_func
284#         return super().create()
285
286# class BashTool(BaseHeavenTool):
287#     name = "BashTool"
288#     description = "Run commands in a bash shell"
289#     args_schema = BashToolArgsSchema
290#     is_async = True
291
292#     # Dictionary to store sessions per container
293#     _bash_sessions: ClassVar[Dict[str, _BashSession]] = {}
294
295#     @classmethod
296#     def get_session(cls, container_id: str = "default"):
297#         if container_id not in cls._bash_sessions:
298#             cls._bash_sessions[container_id] = _BashSession()
299#         return cls._bash_sessions[container_id]
300
301#     @classmethod
302#     def create(cls):
303#         async def wrapped_func(command: str | None = None, restart: bool = False, container_id: str = "default"):
304#             session = cls.get_session(container_id)
305
306#             if restart:
307#                 if session and session._started:
308#                     session.stop()
309#                 cls._bash_sessions[container_id] = _BashSession()
310#                 session = cls._bash_sessions[container_id]
311#                 await session.start()
312#                 return "tool has been restarted."
313
314#             if not session._started:
315#                 await session.start()
316
317#             if command is not None:
318#                 return await session.run(command)
319
320#             raise ToolError("no command provided.")
321
322#         cls.func = wrapped_func
323#         return super().create()
324
325
326
327  
328# # Global session management
329# _bash_session = None
330
331# # The main function that our tool will wrap
332# async def execute_bash_command(command: str | None = None, restart: bool = False) -> str:
333#     """Main function that handles all bash execution logic"""
334#     global _bash_session
335
336#     try:
337#                 # Add safety checks
338#         if command is not None:
339#             # Check for rm -rf
340#             if "rm -rf" in command:
341#                 raise ToolError("SAFETY ERROR: 'rm -rf' commands are forbidden to prevent accidental data loss. If you need to delete something large, ask the user to help you by doing it themselves.")
342
343#             # Could also add other dangerous commands to check for
344#             dangerous_commands = [
345#                 "rm -rf",
346#                 "rm -r",
347#                 "rmdir",  # Maybe allow this one with specific checks
348#                 "> /dev/",  # Prevent direct device writes
349#                 "mkfs",    # Prevent filesystem formatting
350#                 "dd",      # Prevent direct disk operations
351#             ]
352
353#             for dangerous_cmd in dangerous_commands:
354#                 if dangerous_cmd in command:
355#                     raise ToolError(f"SAFETY ERROR: '{dangerous_cmd}' is a protected command that could cause data loss. Ask the user to help you if you are tangled.")
356
357#         if _bash_session is None:
358#             _bash_session = _BashSession()
359#             await _bash_session.start()
360
361#         if restart:
362#             if _bash_session:
363#                 _bash_session.stop()
364#             _bash_session = _BashSession()
365#             await _bash_session.start()
366#             return "tool has been restarted."
367
368#         if command is not None:
369#             result = await _bash_session.run(command)
370#             # Convert CLIResult to string format for the tool
371#             return f"Output: {result.output}\nError: {result.error}" if result.error else result.output
372
373#         raise ToolError("no command provided.")
374#     except Exception as e:
375#         raise ToolError(f"Error in bash command: {e}")
376
377# class BashTool(BaseHeavenTool):
378#     """A tool that allows the agent to run bash commands."""
379#     name = "BashTool"
380#     description ="Run commands in a bash shell"
381#     func = execute_bash_command
382#     args_schema = BashToolArgsSchema
383#     is_async = True