416 lines
15 KiB
Python
416 lines
15 KiB
Python
"""任务分发器 - 通过 Redis Queue 将任务分发给 Agent
|
|
|
|
与业务系统解耦：通过依赖注入获取 Redis 连接和数据库会话。
|
|
"""
|
|
|
|
import ipaddress
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from typing import Any, Callable, Awaitable
|
|
from urllib.parse import urlparse
|
|
|
|
from agentkit.core.exceptions import (
|
|
NoAvailableAgentError,
|
|
TaskDispatchError,
|
|
TaskNotFoundError,
|
|
)
|
|
from agentkit.core.protocol import (
|
|
AgentStatus,
|
|
TaskMessage,
|
|
TaskProgress,
|
|
TaskResult,
|
|
TaskStatus,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_PRIVATE_NETWORKS = [
|
|
ipaddress.ip_network("127.0.0.0/8"),
|
|
ipaddress.ip_network("10.0.0.0/8"),
|
|
ipaddress.ip_network("172.16.0.0/12"),
|
|
ipaddress.ip_network("192.168.0.0/16"),
|
|
ipaddress.ip_network("169.254.0.0/16"),
|
|
ipaddress.ip_network("::1/128"),
|
|
ipaddress.ip_network("fc00::/7"),
|
|
ipaddress.ip_network("fe80::/10"),
|
|
]
|
|
|
|
|
|
def _validate_callback_url(url: str) -> bool:
|
|
"""Validate callback URL to prevent SSRF attacks.
|
|
|
|
Rules:
|
|
- Only http/https protocols allowed
|
|
- No localhost or loopback addresses
|
|
- No private/internal IP ranges
|
|
- No link-local addresses
|
|
|
|
Returns True if valid, False if should be blocked.
|
|
"""
|
|
try:
|
|
parsed = urlparse(url)
|
|
except Exception:
|
|
return False
|
|
|
|
if parsed.scheme not in ("http", "https"):
|
|
return False
|
|
|
|
hostname = parsed.hostname
|
|
if not hostname:
|
|
return False
|
|
|
|
if hostname.lower() in ("localhost", "127.0.0.1", "::1"):
|
|
return False
|
|
|
|
try:
|
|
ip = ipaddress.ip_address(hostname)
|
|
for network in _PRIVATE_NETWORKS:
|
|
if ip in network:
|
|
return False
|
|
except ValueError:
|
|
pass
|
|
|
|
return True
|
|
|
|
|
|
class TaskDispatcher:
|
|
"""任务分发器，通过 Redis Queue 将任务分发给 Agent"""
|
|
|
|
def __init__(
|
|
self,
|
|
redis_factory: Callable[[], Awaitable[Any]],
|
|
session_factory: Callable[[], Any],
|
|
agent_model: Any,
|
|
task_model: Any,
|
|
task_log_model: Any,
|
|
):
|
|
"""
|
|
Args:
|
|
redis_factory: 返回 Redis 连接的异步工厂
|
|
session_factory: 返回 async context manager 的工厂，用于获取数据库会话
|
|
agent_model: Agent ORM 模型类
|
|
task_model: AgentTask ORM 模型类
|
|
task_log_model: AgentTaskLog ORM 模型类
|
|
"""
|
|
self._redis_factory = redis_factory
|
|
self._session_factory = session_factory
|
|
self._agent_model = agent_model
|
|
self._task_model = task_model
|
|
self._task_log_model = task_log_model
|
|
|
|
async def _get_redis(self):
|
|
return await self._redis_factory()
|
|
|
|
async def dispatch(
|
|
self,
|
|
task: TaskMessage,
|
|
organization_id: str | None = None,
|
|
created_by: str | None = None,
|
|
) -> str:
|
|
"""分发任务到对应 Agent 的队列，返回 task_id"""
|
|
async with self._session_factory() as db:
|
|
try:
|
|
from sqlalchemy import select
|
|
|
|
AgentModel = self._agent_model
|
|
TaskModel = self._task_model
|
|
|
|
# 查找目标 Agent
|
|
stmt = select(AgentModel).where(AgentModel.name == task.agent_name)
|
|
result = await db.execute(stmt)
|
|
agent = result.scalar_one_or_none()
|
|
|
|
if not agent:
|
|
raise TaskDispatchError(task.task_id, f"Agent '{task.agent_name}' not found")
|
|
|
|
if agent.status != AgentStatus.ONLINE:
|
|
raise TaskDispatchError(
|
|
task.task_id,
|
|
f"Agent '{task.agent_name}' is not online (status={agent.status})",
|
|
)
|
|
|
|
# 写入 agent_tasks 表
|
|
task_id_uuid = uuid.UUID(task.task_id)
|
|
agent_task = TaskModel(
|
|
id=task_id_uuid,
|
|
agent_id=agent.id,
|
|
task_type=task.task_type,
|
|
status=TaskStatus.PENDING,
|
|
priority=task.priority,
|
|
input_data=task.input_data,
|
|
organization_id=uuid.UUID(organization_id) if organization_id else agent.id,
|
|
created_by=uuid.UUID(created_by) if created_by else None,
|
|
)
|
|
db.add(agent_task)
|
|
await db.commit()
|
|
|
|
# 推送到 Redis Queue
|
|
redis = await self._get_redis()
|
|
queue_key = f"agent:{task.agent_name}:tasks"
|
|
await redis.lpush(queue_key, json.dumps(task.to_dict()))
|
|
|
|
logger.info(
|
|
f"Task {task.task_id} dispatched to agent '{task.agent_name}' "
|
|
f"(type={task.task_type}, priority={task.priority})"
|
|
)
|
|
return task.task_id
|
|
|
|
except TaskDispatchError:
|
|
raise
|
|
except Exception as e:
|
|
await db.rollback()
|
|
logger.error(f"Failed to dispatch task {task.task_id}: {e}")
|
|
raise TaskDispatchError(task.task_id, str(e))
|
|
|
|
async def cancel_task(self, task_id: str):
|
|
"""取消任务"""
|
|
async with self._session_factory() as db:
|
|
try:
|
|
from sqlalchemy import select
|
|
|
|
TaskModel = self._task_model
|
|
task_uuid = uuid.UUID(task_id)
|
|
stmt = select(TaskModel).where(TaskModel.id == task_uuid)
|
|
result = await db.execute(stmt)
|
|
task = result.scalar_one_or_none()
|
|
|
|
if not task:
|
|
raise TaskNotFoundError(task_id)
|
|
|
|
if task.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED):
|
|
logger.warning(f"Cannot cancel task {task_id} with status '{task.status}'")
|
|
return
|
|
|
|
task.status = TaskStatus.CANCELLED
|
|
task.completed_at = datetime.now(timezone.utc)
|
|
await db.commit()
|
|
|
|
await self._write_log(
|
|
db, task_id=task_id, agent_id=str(task.agent_id),
|
|
log_level="info", message="Task cancelled by user",
|
|
)
|
|
await db.commit()
|
|
|
|
logger.info(f"Task {task_id} cancelled")
|
|
|
|
except TaskNotFoundError:
|
|
raise
|
|
except Exception as e:
|
|
await db.rollback()
|
|
logger.error(f"Failed to cancel task {task_id}: {e}")
|
|
raise
|
|
|
|
async def get_task_status(self, task_id: str) -> dict:
|
|
"""获取任务状态"""
|
|
async with self._session_factory() as db:
|
|
from sqlalchemy import select
|
|
|
|
TaskModel = self._task_model
|
|
task_uuid = uuid.UUID(task_id)
|
|
stmt = select(TaskModel).where(TaskModel.id == task_uuid)
|
|
result = await db.execute(stmt)
|
|
task = result.scalar_one_or_none()
|
|
|
|
if not task:
|
|
raise TaskNotFoundError(task_id)
|
|
|
|
return self._task_to_dict(task)
|
|
|
|
async def handle_result(self, result: TaskResult):
|
|
"""处理 Agent 返回的结果"""
|
|
async with self._session_factory() as db:
|
|
try:
|
|
from sqlalchemy import select
|
|
|
|
TaskModel = self._task_model
|
|
task_uuid = uuid.UUID(result.task_id)
|
|
stmt = select(TaskModel).where(TaskModel.id == task_uuid)
|
|
db_result = await db.execute(stmt)
|
|
task = db_result.scalar_one_or_none()
|
|
|
|
if not task:
|
|
logger.error(f"Task {result.task_id} not found when handling result")
|
|
return
|
|
|
|
task.status = result.status
|
|
task.output_data = result.output_data
|
|
task.error_message = result.error_message
|
|
task.started_at = result.started_at
|
|
task.completed_at = result.completed_at
|
|
await db.commit()
|
|
|
|
log_level = "info" if result.status == TaskStatus.COMPLETED else "error"
|
|
log_message = (
|
|
f"Task {result.status}"
|
|
if result.status == TaskStatus.COMPLETED
|
|
else f"Task failed: {result.error_message}"
|
|
)
|
|
await self._write_log(
|
|
db,
|
|
task_id=result.task_id,
|
|
agent_id=str(task.agent_id),
|
|
log_level=log_level,
|
|
message=log_message,
|
|
extra_metadata=result.metrics,
|
|
)
|
|
await db.commit()
|
|
|
|
# 触发回调
|
|
if result.output_data and result.output_data.get("callback_url"):
|
|
await self._trigger_callback(result.output_data["callback_url"], result)
|
|
|
|
logger.info(f"Task {result.task_id} result handled (status={result.status})")
|
|
|
|
except Exception as e:
|
|
await db.rollback()
|
|
logger.error(f"Failed to handle result for task {result.task_id}: {e}")
|
|
|
|
async def handle_progress(self, progress: TaskProgress):
|
|
"""处理进度上报"""
|
|
async with self._session_factory() as db:
|
|
try:
|
|
from sqlalchemy import select
|
|
|
|
AgentModel = self._agent_model
|
|
stmt = select(AgentModel).where(AgentModel.name == progress.agent_name)
|
|
result = await db.execute(stmt)
|
|
agent = result.scalar_one_or_none()
|
|
|
|
if not agent:
|
|
logger.warning(f"Agent '{progress.agent_name}' not found for progress report")
|
|
return
|
|
|
|
await self._write_log(
|
|
db,
|
|
task_id=progress.task_id,
|
|
agent_id=str(agent.id),
|
|
log_level="info",
|
|
message=f"Progress: {progress.progress:.0%} - {progress.message}",
|
|
extra_metadata={
|
|
"progress": progress.progress,
|
|
"updated_at": progress.updated_at.isoformat(),
|
|
},
|
|
)
|
|
await db.commit()
|
|
|
|
except Exception as e:
|
|
await db.rollback()
|
|
logger.error(f"Failed to handle progress for task {progress.task_id}: {e}")
|
|
|
|
async def retry_failed_tasks(self, max_retries: int = 3):
|
|
"""重试失败的任务"""
|
|
async with self._session_factory() as db:
|
|
try:
|
|
from sqlalchemy import select
|
|
|
|
TaskModel = self._task_model
|
|
AgentModel = self._agent_model
|
|
LogModel = self._task_log_model
|
|
|
|
stmt = select(TaskModel).where(TaskModel.status == TaskStatus.FAILED)
|
|
result = await db.execute(stmt)
|
|
failed_tasks = result.scalars().all()
|
|
|
|
retried = 0
|
|
for task in failed_tasks:
|
|
log_stmt = select(LogModel).where(
|
|
LogModel.task_id == task.id,
|
|
LogModel.message.like("%retry%"),
|
|
)
|
|
log_result = await db.execute(log_stmt)
|
|
retry_count = len(log_result.scalars().all())
|
|
|
|
if retry_count < max_retries:
|
|
task.status = TaskStatus.PENDING
|
|
task.error_message = None
|
|
task.started_at = None
|
|
task.completed_at = None
|
|
|
|
agent_stmt = select(AgentModel).where(AgentModel.id == task.agent_id)
|
|
agent_result = await db.execute(agent_stmt)
|
|
agent = agent_result.scalar_one_or_none()
|
|
|
|
if agent and agent.status == AgentStatus.ONLINE:
|
|
task_msg = TaskMessage(
|
|
task_id=str(task.id),
|
|
agent_name=agent.name,
|
|
task_type=task.task_type,
|
|
priority=task.priority,
|
|
input_data=task.input_data or {},
|
|
callback_url=None,
|
|
created_at=datetime.now(timezone.utc),
|
|
)
|
|
redis = await self._get_redis()
|
|
queue_key = f"agent:{agent.name}:tasks"
|
|
await redis.lpush(queue_key, json.dumps(task_msg.to_dict()))
|
|
|
|
await self._write_log(
|
|
db,
|
|
task_id=str(task.id),
|
|
agent_id=str(agent.id),
|
|
log_level="info",
|
|
message=f"Task retry attempt {retry_count + 1}/{max_retries}",
|
|
)
|
|
retried += 1
|
|
|
|
await db.commit()
|
|
if retried > 0:
|
|
logger.info(f"Retried {retried} failed tasks")
|
|
|
|
except Exception as e:
|
|
await db.rollback()
|
|
logger.error(f"Failed to retry failed tasks: {e}")
|
|
|
|
async def _write_log(
|
|
self,
|
|
db: Any,
|
|
task_id: str,
|
|
agent_id: str,
|
|
log_level: str,
|
|
message: str,
|
|
extra_metadata: dict | None = None,
|
|
):
|
|
LogModel = self._task_log_model
|
|
log_entry = LogModel(
|
|
task_id=uuid.UUID(task_id),
|
|
agent_id=uuid.UUID(agent_id),
|
|
log_level=log_level,
|
|
message=message,
|
|
extra_metadata=extra_metadata,
|
|
)
|
|
db.add(log_entry)
|
|
|
|
async def _trigger_callback(self, callback_url: str, result: TaskResult):
|
|
if not _validate_callback_url(callback_url):
|
|
logger.warning(f"Callback URL rejected (SSRF protection): {callback_url}")
|
|
return
|
|
|
|
try:
|
|
import httpx
|
|
async with httpx.AsyncClient(timeout=10) as client:
|
|
await client.post(callback_url, json=result.to_dict())
|
|
logger.info(f"Callback triggered for task {result.task_id}")
|
|
except Exception as e:
|
|
logger.warning(f"Callback failed for task {result.task_id}: {e}")
|
|
|
|
def _task_to_dict(self, task: Any) -> dict:
|
|
return {
|
|
"id": str(task.id),
|
|
"agent_id": str(task.agent_id),
|
|
"task_type": task.task_type,
|
|
"status": task.status,
|
|
"priority": task.priority,
|
|
"input_data": task.input_data,
|
|
"output_data": task.output_data,
|
|
"error_message": task.error_message,
|
|
"created_by": str(task.created_by) if task.created_by else None,
|
|
"organization_id": str(task.organization_id),
|
|
"project_id": str(task.project_id) if task.project_id else None,
|
|
"scheduled_at": task.scheduled_at.isoformat() if task.scheduled_at else None,
|
|
"started_at": task.started_at.isoformat() if task.started_at else None,
|
|
"completed_at": task.completed_at.isoformat() if task.completed_at else None,
|
|
"created_at": task.created_at.isoformat() if task.created_at else None,
|
|
}
|