refactor: simplify tool server to asyncio tasks with per-agent isolation
- Replace multiprocessing/threading with single asyncio task per agent - Add task cancellation: new request cancels previous for same agent - Add per-agent state isolation via ContextVar for Terminal, Browser, Python managers - Add posthog telemetry for tool execution errors (timeout, http, sandbox) - Fix proxy manager singleton pattern - Increase client timeout buffer over server timeout - Add context.py to Dockerfile
This commit is contained in:
0xallam
+45
-165
@@ -2,16 +2,10 @@ from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
import queue as stdlib_queue
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
from multiprocessing import Process, Queue
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
import uvicorn
|
||||
from fastapi import Depends, FastAPI, HTTPException, status
|
||||
@@ -40,13 +34,9 @@ REQUEST_TIMEOUT = args.timeout
|
||||
|
||||
app = FastAPI()
|
||||
security = HTTPBearer()
|
||||
|
||||
security_dependency = Depends(security)
|
||||
|
||||
agent_processes: dict[str, dict[str, Any]] = {}
|
||||
agent_queues: dict[str, dict[str, Queue[Any]]] = {}
|
||||
pending_responses: dict[str, dict[str, asyncio.Future[Any]]] = {}
|
||||
agent_listeners: dict[str, dict[str, Any]] = {}
|
||||
agent_tasks: dict[str, asyncio.Task[Any]] = {}
|
||||
|
||||
|
||||
def verify_token(credentials: HTTPAuthorizationCredentials) -> str:
|
||||
@@ -78,107 +68,19 @@ class ToolExecutionResponse(BaseModel):
|
||||
error: str | None = None
|
||||
|
||||
|
||||
def agent_worker(_agent_id: str, request_queue: Queue[Any], response_queue: Queue[Any]) -> None:
|
||||
null_handler = logging.NullHandler()
|
||||
|
||||
root_logger = logging.getLogger()
|
||||
root_logger.handlers = [null_handler]
|
||||
root_logger.setLevel(logging.CRITICAL)
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from strix.tools.argument_parser import ArgumentConversionError, convert_arguments
|
||||
async def _run_tool(agent_id: str, tool_name: str, kwargs: dict[str, Any]) -> Any:
|
||||
from strix.tools.argument_parser import convert_arguments
|
||||
from strix.tools.context import set_current_agent_id
|
||||
from strix.tools.registry import get_tool_by_name
|
||||
|
||||
def _execute_request(request: dict[str, Any]) -> None:
|
||||
request_id = request.get("request_id", "")
|
||||
tool_name = request["tool_name"]
|
||||
kwargs = request["kwargs"]
|
||||
set_current_agent_id(agent_id)
|
||||
|
||||
try:
|
||||
tool_func = get_tool_by_name(tool_name)
|
||||
if not tool_func:
|
||||
response_queue.put(
|
||||
{"request_id": request_id, "error": f"Tool '{tool_name}' not found"}
|
||||
)
|
||||
return
|
||||
tool_func = get_tool_by_name(tool_name)
|
||||
if not tool_func:
|
||||
raise ValueError(f"Tool '{tool_name}' not found")
|
||||
|
||||
converted_kwargs = convert_arguments(tool_func, kwargs)
|
||||
result = tool_func(**converted_kwargs)
|
||||
|
||||
response_queue.put({"request_id": request_id, "result": result})
|
||||
|
||||
except (ArgumentConversionError, ValidationError) as e:
|
||||
response_queue.put({"request_id": request_id, "error": f"Invalid arguments: {e}"})
|
||||
except (RuntimeError, ValueError, ImportError) as e:
|
||||
response_queue.put({"request_id": request_id, "error": f"Tool execution error: {e}"})
|
||||
except Exception as e: # noqa: BLE001
|
||||
response_queue.put({"request_id": request_id, "error": f"Unexpected error: {e}"})
|
||||
|
||||
with ThreadPoolExecutor() as executor:
|
||||
while True:
|
||||
request = None
|
||||
try:
|
||||
request = request_queue.get()
|
||||
|
||||
if request is None:
|
||||
break
|
||||
|
||||
executor.submit(_execute_request, request)
|
||||
|
||||
except (RuntimeError, ValueError, ImportError) as e:
|
||||
req_id = request.get("request_id", "") if request else ""
|
||||
response_queue.put({"request_id": req_id, "error": f"Worker error: {e}"})
|
||||
|
||||
|
||||
def _ensure_response_listener(agent_id: str, response_queue: Queue[Any]) -> None:
|
||||
if agent_id in agent_listeners:
|
||||
return
|
||||
|
||||
stop_event = threading.Event()
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def _listener() -> None:
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
item = response_queue.get(timeout=0.5)
|
||||
except stdlib_queue.Empty:
|
||||
continue
|
||||
except (BrokenPipeError, EOFError):
|
||||
break
|
||||
|
||||
request_id = item.get("request_id")
|
||||
if not request_id or agent_id not in pending_responses:
|
||||
continue
|
||||
|
||||
future = pending_responses[agent_id].pop(request_id, None)
|
||||
if future and not future.done():
|
||||
with contextlib.suppress(RuntimeError):
|
||||
loop.call_soon_threadsafe(future.set_result, item)
|
||||
|
||||
listener_thread = threading.Thread(target=_listener, daemon=True)
|
||||
listener_thread.start()
|
||||
|
||||
agent_listeners[agent_id] = {"thread": listener_thread, "stop_event": stop_event}
|
||||
|
||||
|
||||
def ensure_agent_process(agent_id: str) -> tuple[Queue[Any], Queue[Any]]:
|
||||
if agent_id not in agent_processes:
|
||||
request_queue: Queue[Any] = Queue()
|
||||
response_queue: Queue[Any] = Queue()
|
||||
|
||||
process = Process(
|
||||
target=agent_worker, args=(agent_id, request_queue, response_queue), daemon=True
|
||||
)
|
||||
process.start()
|
||||
|
||||
agent_processes[agent_id] = {"process": process, "pid": process.pid}
|
||||
agent_queues[agent_id] = {"request": request_queue, "response": response_queue}
|
||||
pending_responses[agent_id] = {}
|
||||
|
||||
_ensure_response_listener(agent_id, response_queue)
|
||||
|
||||
return agent_queues[agent_id]["request"], agent_queues[agent_id]["response"]
|
||||
converted_kwargs = convert_arguments(tool_func, kwargs)
|
||||
return await asyncio.to_thread(tool_func, **converted_kwargs)
|
||||
|
||||
|
||||
@app.post("/execute", response_model=ToolExecutionResponse)
|
||||
@@ -187,33 +89,42 @@ async def execute_tool(
|
||||
) -> ToolExecutionResponse:
|
||||
verify_token(credentials)
|
||||
|
||||
request_queue, _response_queue = ensure_agent_process(request.agent_id)
|
||||
agent_id = request.agent_id
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
req_id = uuid4().hex
|
||||
future: asyncio.Future[Any] = loop.create_future()
|
||||
pending_responses[request.agent_id][req_id] = future
|
||||
if agent_id in agent_tasks:
|
||||
old_task = agent_tasks[agent_id]
|
||||
if not old_task.done():
|
||||
old_task.cancel()
|
||||
|
||||
request_queue.put(
|
||||
{
|
||||
"request_id": req_id,
|
||||
"tool_name": request.tool_name,
|
||||
"kwargs": request.kwargs,
|
||||
}
|
||||
task = asyncio.create_task(
|
||||
asyncio.wait_for(
|
||||
_run_tool(agent_id, request.tool_name, request.kwargs), timeout=REQUEST_TIMEOUT
|
||||
)
|
||||
)
|
||||
agent_tasks[agent_id] = task
|
||||
|
||||
try:
|
||||
response = await asyncio.wait_for(future, timeout=REQUEST_TIMEOUT)
|
||||
result = await task
|
||||
return ToolExecutionResponse(result=result)
|
||||
|
||||
if "error" in response:
|
||||
return ToolExecutionResponse(error=response["error"])
|
||||
return ToolExecutionResponse(result=response.get("result"))
|
||||
except asyncio.CancelledError:
|
||||
return ToolExecutionResponse(error="Cancelled by newer request")
|
||||
|
||||
except TimeoutError:
|
||||
pending_responses[request.agent_id].pop(req_id, None)
|
||||
return ToolExecutionResponse(error=f"Request timed out after {REQUEST_TIMEOUT} seconds")
|
||||
except (RuntimeError, ValueError, OSError) as e:
|
||||
return ToolExecutionResponse(error=f"Worker error: {e}")
|
||||
return ToolExecutionResponse(error=f"Tool timed out after {REQUEST_TIMEOUT}s")
|
||||
|
||||
except ValidationError as e:
|
||||
return ToolExecutionResponse(error=f"Invalid arguments: {e}")
|
||||
|
||||
except (ValueError, RuntimeError, ImportError) as e:
|
||||
return ToolExecutionResponse(error=f"Tool execution error: {e}")
|
||||
|
||||
except Exception as e: # noqa: BLE001
|
||||
return ToolExecutionResponse(error=f"Unexpected error: {e}")
|
||||
|
||||
finally:
|
||||
if agent_tasks.get(agent_id) is task:
|
||||
del agent_tasks[agent_id]
|
||||
|
||||
|
||||
@app.post("/register_agent")
|
||||
@@ -221,8 +132,6 @@ async def register_agent(
|
||||
agent_id: str, credentials: HTTPAuthorizationCredentials = security_dependency
|
||||
) -> dict[str, str]:
|
||||
verify_token(credentials)
|
||||
|
||||
ensure_agent_process(agent_id)
|
||||
return {"status": "registered", "agent_id": agent_id}
|
||||
|
||||
|
||||
@@ -233,42 +142,16 @@ async def health_check() -> dict[str, Any]:
|
||||
"sandbox_mode": str(SANDBOX_MODE),
|
||||
"environment": "sandbox" if SANDBOX_MODE else "main",
|
||||
"auth_configured": "true" if EXPECTED_TOKEN else "false",
|
||||
"active_agents": len(agent_processes),
|
||||
"agents": list(agent_processes.keys()),
|
||||
"active_agents": len(agent_tasks),
|
||||
"agents": list(agent_tasks.keys()),
|
||||
}
|
||||
|
||||
|
||||
def cleanup_all_agents() -> None:
|
||||
for agent_id in list(agent_processes.keys()):
|
||||
try:
|
||||
if agent_id in agent_listeners:
|
||||
agent_listeners[agent_id]["stop_event"].set()
|
||||
|
||||
agent_queues[agent_id]["request"].put(None)
|
||||
process = agent_processes[agent_id]["process"]
|
||||
|
||||
process.join(timeout=1)
|
||||
|
||||
if process.is_alive():
|
||||
process.terminate()
|
||||
process.join(timeout=1)
|
||||
|
||||
if process.is_alive():
|
||||
process.kill()
|
||||
|
||||
if agent_id in agent_listeners:
|
||||
listener_thread = agent_listeners[agent_id]["thread"]
|
||||
listener_thread.join(timeout=0.5)
|
||||
|
||||
except (BrokenPipeError, EOFError, OSError):
|
||||
pass
|
||||
except (RuntimeError, ValueError) as e:
|
||||
logging.getLogger(__name__).debug(f"Error during agent cleanup: {e}")
|
||||
|
||||
|
||||
def signal_handler(_signum: int, _frame: Any) -> None:
|
||||
signal.signal(signal.SIGPIPE, signal.SIG_IGN) if hasattr(signal, "SIGPIPE") else None
|
||||
cleanup_all_agents()
|
||||
if hasattr(signal, "SIGPIPE"):
|
||||
signal.signal(signal.SIGPIPE, signal.SIG_IGN)
|
||||
for task in agent_tasks.values():
|
||||
task.cancel()
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
@@ -279,7 +162,4 @@ signal.signal(signal.SIGTERM, signal_handler)
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
uvicorn.run(app, host=args.host, port=args.port, log_level="info")
|
||||
finally:
|
||||
cleanup_all_agents()
|
||||
uvicorn.run(app, host=args.host, port=args.port, log_level="info")
|
||||
|
||||
Reference in New Issue
Block a user