nearxos
0375b20bb4
- Added configuration options for requiring human approval before applying LLM-generated MCP patches. - Updated Docker setup to include skills directory. - Integrated skills management into the backend, allowing for procedural guides and skill matching. - Refactored database initialization to apply Alembic migrations. - Enhanced task approval process to handle MCP patch applications with optional approval. - Introduced new schemas for skills and updated existing APIs to support skills functionality. This commit lays the groundwork for improved agent capabilities and better management of MCP development processes.
138 lines
4.6 KiB
Python
138 lines
4.6 KiB
Python
"""Background worker: owns MCP sessions and runs the troubleshooting agent per task.
|
|
|
|
Consumes task ids from the Redis queue, executes the LangGraph agent, and keeps the
|
|
MCP status snapshot fresh for the API/UI.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
|
|
from sqlalchemy import select
|
|
|
|
from app.agent.graph import run_task_agent
|
|
from app.config import settings
|
|
from app.database import SessionLocal, init_db
|
|
from app.mcp_manager import get_manager
|
|
from app.models.enums import TaskStatus
|
|
from app.models.task import Task
|
|
from app.services.events import TASK_QUEUE, get_redis
|
|
from app.services.task_runtime import is_task_cancel_requested
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("worker")
|
|
|
|
|
|
async def _refresh_status_loop() -> None:
|
|
manager = get_manager()
|
|
while True:
|
|
await asyncio.sleep(5)
|
|
try:
|
|
while await manager.process_command_queue_once():
|
|
await asyncio.sleep(0) # yield so task queue blpop can run
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.warning("MCP command processing failed: %s", exc)
|
|
try:
|
|
await manager.publish_status()
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.warning("status publish failed: %s", exc)
|
|
|
|
|
|
async def _task_should_run(task_id: int) -> bool:
|
|
async with SessionLocal() as db:
|
|
res = await db.execute(select(Task).where(Task.id == task_id))
|
|
task = res.scalar_one_or_none()
|
|
if not task:
|
|
return False
|
|
if task.status == TaskStatus.cancelled:
|
|
return False
|
|
if await is_task_cancel_requested(task_id):
|
|
return False
|
|
return True
|
|
|
|
|
|
async def _mark_task_running(task_id: int) -> None:
|
|
async with SessionLocal() as db:
|
|
res = await db.execute(select(Task).where(Task.id == task_id))
|
|
task = res.scalar_one_or_none()
|
|
if task and task.status == TaskStatus.queued:
|
|
task.status = TaskStatus.running
|
|
await db.commit()
|
|
|
|
|
|
async def _run_task(task_id: int, sem: asyncio.Semaphore) -> None:
|
|
async with sem:
|
|
if not await _task_should_run(task_id):
|
|
logger.info("Skipping task %s (cancelled or missing)", task_id)
|
|
return
|
|
logger.info("Running task %s", task_id)
|
|
try:
|
|
await run_task_agent(task_id)
|
|
except asyncio.CancelledError:
|
|
logger.info("task %s asyncio task cancelled", task_id)
|
|
except Exception: # noqa: BLE001
|
|
logger.exception("task %s crashed", task_id)
|
|
|
|
|
|
async def _cancel_watch_loop(in_flight: dict[int, asyncio.Task]) -> None:
|
|
"""Cancel in-flight asyncio tasks when the user requests stop via the API."""
|
|
while True:
|
|
await asyncio.sleep(1)
|
|
for task_id, handle in list(in_flight.items()):
|
|
if handle.done():
|
|
continue
|
|
try:
|
|
if await is_task_cancel_requested(task_id):
|
|
logger.info("Cancelling in-flight asyncio task %s", task_id)
|
|
handle.cancel()
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.debug("cancel watch for task %s: %s", task_id, exc)
|
|
|
|
|
|
async def main() -> None:
|
|
await init_db()
|
|
manager = get_manager()
|
|
logger.info("Starting MCP manager (cloning + launching MCP servers)...")
|
|
await manager.startup()
|
|
logger.info("MCP manager ready: %s", [s["name"] for s in manager.status_snapshot()])
|
|
|
|
asyncio.create_task(_refresh_status_loop())
|
|
|
|
max_concurrent = max(1, settings.worker_max_concurrent_tasks)
|
|
sem = asyncio.Semaphore(max_concurrent)
|
|
in_flight: dict[int, asyncio.Task] = {}
|
|
asyncio.create_task(_cancel_watch_loop(in_flight))
|
|
|
|
redis = get_redis()
|
|
logger.info(
|
|
"Worker waiting for tasks on %s (max concurrent: %s)",
|
|
TASK_QUEUE,
|
|
max_concurrent,
|
|
)
|
|
while True:
|
|
item = await redis.blpop(TASK_QUEUE, timeout=5)
|
|
if not item:
|
|
in_flight = {tid: t for tid, t in in_flight.items() if not t.done()}
|
|
continue
|
|
_, task_id_str = item
|
|
try:
|
|
task_id = int(task_id_str)
|
|
except ValueError:
|
|
continue
|
|
if not await _task_should_run(task_id):
|
|
logger.info("Skipped dequeued task %s (already cancelled)", task_id)
|
|
continue
|
|
logger.info("Picked up task %s", task_id)
|
|
await _mark_task_running(task_id)
|
|
handle = asyncio.create_task(_run_task(task_id, sem))
|
|
in_flight[task_id] = handle
|
|
|
|
def _done(t: asyncio.Task, tid: int = task_id) -> None:
|
|
in_flight.pop(tid, None)
|
|
|
|
handle.add_done_callback(_done)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|