Skip to content

Fix async bottlenecks in bash_server file operations #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kcoopermiller wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix async bottlenecks in bash_server file operations #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 39 additions & 10 deletions universal/bash_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,12 @@
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn
import logging
import time

# Configure logging for performance monitoring
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
from pathlib import Path


Expand Down Expand Up @@ -678,6 +684,8 @@ class FileTool(BaseAnthropicTool):
"""

name: ClassVar[Literal["file"]] = "file"
# Maximum file size (in bytes) to store in history to prevent memory issues
MAX_HISTORY_FILE_SIZE = 1024 * 1024 # 1MB
_file_history: Dict[Path, List[str]] # Undo history for text edits

def __init__(self, base_path: Path | None = None):
Expand All @@ -691,6 +699,19 @@ async def _ensure_base_path_exists(self):
if not await aiofiles.os.path.exists(str(self.base_path)):
await asyncio.to_thread(self.base_path.mkdir, parents=True, exist_ok=True)


def _add_to_history(self, full_path: Path, content: str) -> None:
"""Add content to file history with size limits to prevent memory issues"""
# Skip history for very large files to prevent memory problems
if len(content.encode("utf-8")) > self.MAX_HISTORY_FILE_SIZE:
# Keep only the most recent version for large files
self._file_history[full_path] = []
return

self._file_history[full_path].append(content)
if len(self._file_history[full_path]) > 5:
self._file_history[full_path].pop(0)

async def _validate_path(self, path: str) -> Path:
try:
path_obj = Path(path)
Expand Down Expand Up @@ -998,9 +1019,7 @@ async def replace(self, path: str, old_str: str, new_str: str, all_occurrences:
else:
new_content = norm_new_content

self._file_history[full_path].append(content)
if len(self._file_history[full_path]) > 5:
self._file_history[full_path].pop(0)
self._add_to_history(full_path, content)
async with aiofiles.open(str(full_path), 'w') as f:
await f.write(new_content)
return ToolResult(output=f"Replaced \"{_shorten(old_str)}\" with \"{_shorten(new_str)}\"")
Expand All @@ -1020,9 +1039,7 @@ async def insert(self, path: str, line: int, text: str) -> ToolResult:
raise ToolError(f"Line number {line} is out of range")
lines.insert(line - 1, text)
new_content = "\n".join(lines)
self._file_history[full_path].append(content)
if len(self._file_history[full_path]) > 5:
self._file_history[full_path].pop(0)
self._add_to_history(full_path, content)
async with aiofiles.open(str(full_path), 'w') as f:
await f.write(new_content)
return ToolResult(output=f"Inserted \"{_shorten(text)}\" at line {line}")
Expand All @@ -1041,9 +1058,7 @@ async def delete_lines(self, path: str, lines: List[int]) -> ToolResult:
lines_to_delete = set(lines)
new_lines = [line for i, line in enumerate(file_lines, 1) if i not in lines_to_delete]
new_content = "\n".join(new_lines)
self._file_history[full_path].append(content)
if len(self._file_history[full_path]) > 5:
self._file_history[full_path].pop(0)
self._add_to_history(full_path, content)
async with aiofiles.open(str(full_path), 'w') as f:
await f.write(new_content)
return ToolResult(output=f"Deleted lines {lines}")
Expand Down Expand Up @@ -1218,14 +1233,22 @@ async def bash_action(request: BashRequest):
@app.post("/file", response_model=ToolResponse)
async def file_action(request: FileRequest):
"""Execute file operations"""
start_time = time.time()
try:
# Convert request to kwargs, excluding None values
kwargs = request.model_dump(exclude_none=True)
logger.info(f"Processing file command: {request.command} for path: {getattr(request, 'path', 'unknown')}")
result = await file_tool(**kwargs)
elapsed = time.time() - start_time
logger.info(f"File operation completed in {elapsed:.2f}s")
return _tool_result_to_response(result)
except ToolError as e:
elapsed = time.time() - start_time
logger.error(f"File operation failed after {elapsed:.2f}s: {str(e)}")
raise HTTPException(status_code=400, detail=str(e))
except Exception as e:
elapsed = time.time() - start_time
logger.error(f"Unexpected error after {elapsed:.2f}s: {str(e)}")
raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")


Expand Down Expand Up @@ -1384,5 +1407,11 @@ async def root():
app,
host="0.0.0.0",
port=8000,
log_level="info"
log_level="info",
timeout_keep_alive=30,
timeout_graceful_shutdown=10,
# Prevent hanging connections from blocking the server
limit_concurrency=100,
# Add request timeout to prevent indefinite hangs
timeout_notify=30
)
118 changes: 118 additions & 0 deletions universal/fix_bash_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
#!/usr/bin/env python3
"""
Script to fix async bottlenecks in bash_server.py
"""

def fix_bash_server():
with open('bash_server.py', 'r') as f:
content = f.read()

# 1. Add logging imports
if 'import logging' not in content:
content = content.replace(
'import uvicorn',
'import uvicorn\nimport logging\nimport time\n\n# Configure logging for performance monitoring\nlogging.basicConfig(level=logging.INFO)\nlogger = logging.getLogger(__name__)'
)

# 2. Add file size constant to FileTool class
content = content.replace(
'class FileTool(BaseAnthropicTool):\n """\n A filesystem editor tool that allows the agent to view, create, and edit files.\n """\n\n name: ClassVar[Literal["file"]] = "file"',
'class FileTool(BaseAnthropicTool):\n """\n A filesystem editor tool that allows the agent to view, create, and edit files.\n """\n\n name: ClassVar[Literal["file"]] = "file"\n # Maximum file size (in bytes) to store in history to prevent memory issues\n MAX_HISTORY_FILE_SIZE = 1024 * 1024 # 1MB'
)

# 3. Add history helper method after _ensure_base_path_exists
history_method = '''
def _add_to_history(self, full_path: Path, content: str) -> None:
"""Add content to file history with size limits to prevent memory issues"""
# Skip history for very large files to prevent memory problems
if len(content.encode("utf-8")) > self.MAX_HISTORY_FILE_SIZE:
# Keep only the most recent version for large files
self._file_history[full_path] = []
return

self._file_history[full_path].append(content)
if len(self._file_history[full_path]) > 5:
self._file_history[full_path].pop(0)
'''

content = content.replace(
'async def _validate_path(self, path: str) -> Path:',
history_method + '\n async def _validate_path(self, path: str) -> Path:'
)

# 4. Replace direct history usage with helper method
content = content.replace(
'self._file_history[full_path].append(content)\n if len(self._file_history[full_path]) > 5:\n self._file_history[full_path].pop(0)',
'self._add_to_history(full_path, content)'
)

# 5. Add timing to file_action endpoint
old_file_action = '''@app.post("/file", response_model=ToolResponse)
async def file_action(request: FileRequest):
"""Execute file operations"""
try:
# Convert request to kwargs, excluding None values
kwargs = request.model_dump(exclude_none=True)
result = await file_tool(**kwargs)
return _tool_result_to_response(result)
except ToolError as e:
raise HTTPException(status_code=400, detail=str(e))
except Exception as e:
raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")'''

new_file_action = '''@app.post("/file", response_model=ToolResponse)
async def file_action(request: FileRequest):
"""Execute file operations"""
start_time = time.time()
try:
# Convert request to kwargs, excluding None values
kwargs = request.model_dump(exclude_none=True)
logger.info(f"Processing file command: {request.command} for path: {getattr(request, 'path', 'unknown')}")
result = await file_tool(**kwargs)
elapsed = time.time() - start_time
logger.info(f"File operation completed in {elapsed:.2f}s")
return _tool_result_to_response(result)
except ToolError as e:
elapsed = time.time() - start_time
logger.error(f"File operation failed after {elapsed:.2f}s: {str(e)}")
raise HTTPException(status_code=400, detail=str(e))
except Exception as e:
elapsed = time.time() - start_time
logger.error(f"Unexpected error after {elapsed:.2f}s: {str(e)}")
raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")'''

content = content.replace(old_file_action, new_file_action)

# 6. Update uvicorn configuration
old_uvicorn = '''if __name__ == "__main__":
uvicorn.run(
app,
host="0.0.0.0",
port=8000,
log_level="info"
)'''

new_uvicorn = '''if __name__ == "__main__":
uvicorn.run(
app,
host="0.0.0.0",
port=8000,
log_level="info",
timeout_keep_alive=30,
timeout_graceful_shutdown=10,
# Prevent hanging connections from blocking the server
limit_concurrency=100,
# Add request timeout to prevent indefinite hangs
timeout_notify=30
)'''

content = content.replace(old_uvicorn, new_uvicorn)

# Write the fixed content
with open('bash_server.py', 'w') as f:
f.write(content)

print("Successfully applied async bottleneck fixes to bash_server.py")

if __name__ == '__main__':
fix_bash_server()