mirror of https://github.com/kortix-ai/suna.git
98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
import os
|
|
import asyncio
|
|
import subprocess
|
|
from agentpress.tool import Tool, ToolResult, tool_schema
|
|
from agentpress.state_manager import StateManager
|
|
|
|
class TerminalTool(Tool):
|
|
"""Terminal command execution tool for workspace operations.
|
|
|
|
Provides secure command execution within a workspace directory,
|
|
with command history tracking and output management.
|
|
|
|
Attributes:
|
|
workspace (str): Path to the workspace directory
|
|
"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.workspace = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'workspace')
|
|
os.makedirs(self.workspace, exist_ok=True)
|
|
self.state_manager = StateManager("state.json")
|
|
|
|
async def _update_command_history(self, command: str, output: str, success: bool):
|
|
"""Update command history in state"""
|
|
history = await self.state_manager.get("terminal_history") or []
|
|
history.append({
|
|
"command": command,
|
|
"output": output,
|
|
"success": success,
|
|
"cwd": os.path.relpath(os.getcwd(), self.workspace)
|
|
})
|
|
await self.state_manager.set("terminal_history", history)
|
|
|
|
@tool_schema({
|
|
"type": "function",
|
|
"function": {
|
|
"name": "execute_command",
|
|
"description": "Execute a shell command in the workspace directory. Commands are executed in an isolated environment.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"command": {
|
|
"type": "string",
|
|
"description": "The shell command to execute"
|
|
}
|
|
},
|
|
"required": ["command"]
|
|
}
|
|
}
|
|
})
|
|
async def execute_command(self, command: str) -> ToolResult:
|
|
original_dir = os.getcwd()
|
|
try:
|
|
# Always change to workspace directory before executing
|
|
os.chdir(self.workspace)
|
|
|
|
# Execute command
|
|
process = await asyncio.create_subprocess_shell(
|
|
command,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
cwd=self.workspace # Explicitly set working directory
|
|
)
|
|
stdout, stderr = await process.communicate()
|
|
|
|
# Prepare output
|
|
output = stdout.decode() if stdout else ""
|
|
error = stderr.decode() if stderr else ""
|
|
success = process.returncode == 0
|
|
|
|
# Update state with command history
|
|
await self._update_command_history(
|
|
command=command,
|
|
output=output + error,
|
|
success=success
|
|
)
|
|
|
|
if success:
|
|
return self.success_response({
|
|
"output": output,
|
|
"error": error,
|
|
"exit_code": process.returncode,
|
|
"cwd": self.workspace
|
|
})
|
|
else:
|
|
return self.fail_response(f"Command failed with exit code {process.returncode}: {error}")
|
|
|
|
except Exception as e:
|
|
await self._update_command_history(
|
|
command=command,
|
|
output=str(e),
|
|
success=False
|
|
)
|
|
return self.fail_response(f"Error executing command: {str(e)}")
|
|
finally:
|
|
# Always restore original directory
|
|
os.chdir(original_dir)
|