lambdagent/agentpaas/engine/worker.py

"""
engine.worker — Background worker that consumes jobs from the queue.

Usage: python -m agentpaas worker --concurrency 4
"""
from __future__ import annotations
import asyncio
import json
import time
from agentpaas.engine.queue import create_queue, Job
from agentpaas.engine.sandbox import Sandbox
from agentpaas.db.session import get_db
from agentpaas.db.models import gen_id, now_utc
from agentpaas.config import settings


async def run_worker(concurrency: int = 4):
    """Main worker loop: dequeue jobs and execute them."""
    queue = create_queue(settings.redis_url)
    sandbox = Sandbox(level=settings.sandbox_level)
    worker_id = gen_id("wk_")
    print(f"Worker {worker_id} started, concurrency={concurrency}")

    semaphore = asyncio.Semaphore(concurrency)

    async def process_job(job: Job):
        async with semaphore:
            db = get_db()
            run_id = gen_id("run_")
            now = now_utc()

            db.execute(
                "INSERT INTO runs (id, agent_id, agent_version, tenant_id, job_id, input, status, created_at) "
                "VALUES (?, ?, ?, ?, ?, ?, 'running', ?)",
                (run_id, job.agent_id, job.agent_version, job.tenant_id, job.id, job.input, now)
            )
            db.commit()

            result = await sandbox.execute(job.config, job.input, settings.default_timeout)

            db.execute(
                "UPDATE runs SET status=?, output=?, duration_ms=?, steps=?, "
                "input_tokens=?, completed_at=? WHERE id=?",
                (result.status, result.output, result.duration_ms, result.steps,
                 result.input_tokens, now_utc(), run_id)
            )
            queue.update_job(job.id, status=result.status, result=result.output)
            db.commit()
            print(f"  Job {job.id} → {result.status} ({result.duration_ms}ms)")

    while True:
        job = await queue.dequeue(timeout=5.0)
        if job:
            asyncio.create_task(process_job(job))
        await asyncio.sleep(0.1)