diff --git a/backend/app/api/admin.py b/backend/app/api/admin.py index bc823d8..d0a3601 100644 --- a/backend/app/api/admin.py +++ b/backend/app/api/admin.py @@ -1,17 +1,24 @@ """后台管理接口：任务/Agent/日志/手动管控。""" from typing import Optional -from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi import APIRouter, Body, Depends, Header, HTTPException, Request from app.constants import TaskStatus from app.models.schemas import AgentInfo, LogEntry, TaskInfo from app.repository.log_repo import LogRepo +from app.constants import AgentStatus from app.services.agent_service import AgentService +from app.services.security import require_admin from app.services.task_service import TaskService router = APIRouter() +def require_admin_dep(x_admin_auth: str = Header(..., alias="X-Admin-Auth")) -> None: + """后台管理接口鉴权：校验 X-Admin-Auth 头与 ADMIN_AUTH 一致。""" + require_admin(x_admin_auth) + + def get_task_service(request: Request) -> TaskService: return TaskService(request.app.state.redis) @@ -30,12 +37,17 @@ async def list_tasks( status: Optional[TaskStatus] = None, limit: int = 100, svc: TaskService = Depends(get_task_service), + _admin: None = Depends(require_admin_dep), ): return await svc.list(status=status, limit=limit) @router.get("/tasks/{request_id}", response_model=TaskInfo, summary="任务详情") -async def task_detail(request_id: str, svc: TaskService = Depends(get_task_service)): +async def task_detail( + request_id: str, + svc: TaskService = Depends(get_task_service), + _admin: None = Depends(require_admin_dep), +): task = await svc.get(request_id) if not task: raise HTTPException(status_code=404, detail="task not found") @@ -43,14 +55,22 @@ async def task_detail(request_id: str, svc: TaskService = Depends(get_task_servi @router.post("/tasks/{request_id}/cancel", summary="取消任务") -async def admin_cancel(request_id: str, svc: TaskService = Depends(get_task_service)): +async def admin_cancel( + request_id: str, + svc: TaskService = Depends(get_task_service), + _admin: None = Depends(require_admin_dep), +): if not await svc.cancel(request_id): raise HTTPException(status_code=400, detail="cannot cancel task") return {"ok": True} @router.post("/tasks/{request_id}/reset", summary="重置任务状态") -async def admin_reset(request_id: str, svc: TaskService = Depends(get_task_service)): +async def admin_reset( + request_id: str, + svc: TaskService = Depends(get_task_service), + _admin: None = Depends(require_admin_dep), +): if not await svc.reset(request_id): raise HTTPException(status_code=400, detail="cannot reset task") return {"ok": True} @@ -58,17 +78,48 @@ async def admin_reset(request_id: str, svc: TaskService = Depends(get_task_servi # ---------- Agent 管理 ---------- @router.get("/agents", response_model=list[AgentInfo], summary="Agent 列表") -async def list_agents(svc: AgentService = Depends(get_agent_service)): +async def list_agents( + svc: AgentService = Depends(get_agent_service), + _admin: None = Depends(require_admin_dep), +): return await svc.all() -@router.post("/agents/{agent_id}/offline", summary="下线 Agent") -async def agent_offline(agent_id: str, svc: AgentService = Depends(get_agent_service)): - if not await svc.take_offline(agent_id): +@router.post("/agents/{agent_id}/unavailable", summary="置 Agent 不可用") +async def agent_unavailable( + agent_id: str, + svc: AgentService = Depends(get_agent_service), + _admin: None = Depends(require_admin_dep), +): + if not await svc.set_unavailable(agent_id): raise HTTPException(status_code=404, detail="agent not found") return {"ok": True} +@router.post("/agents/{agent_id}/available", summary="置 Agent 可用") +async def agent_available( + agent_id: str, + svc: AgentService = Depends(get_agent_service), + _admin: None = Depends(require_admin_dep), +): + if not await svc.set_available(agent_id): + raise HTTPException(status_code=404, detail="agent not found") + return {"ok": True} + + +@router.post("/agents/{agent_id}/priority", response_model=AgentInfo, summary="设置 Agent 优先级") +async def agent_priority( + agent_id: str, + priority: int = Body(..., ge=1, le=5, description="优先级 1-5，越小越先分配"), + svc: AgentService = Depends(get_agent_service), + _admin: None = Depends(require_admin_dep), +): + agent = await svc.set_priority(agent_id, priority) + if not agent: + raise HTTPException(status_code=404, detail="agent not found") + return agent + + # ---------- 日志审计 ---------- @router.get("/logs", response_model=list[LogEntry], summary="日志审计") async def logs( @@ -76,6 +127,7 @@ async def logs( request_id: Optional[str] = None, agent_id: Optional[str] = None, repo: LogRepo = Depends(get_log_repo), + _admin: None = Depends(require_admin_dep), ): return await repo.list(limit=limit, request_id=request_id, agent_id=agent_id) @@ -85,6 +137,7 @@ async def logs( async def overview( task_svc: TaskService = Depends(get_task_service), agent_svc: AgentService = Depends(get_agent_service), + _admin: None = Depends(require_admin_dep), ): tasks = await task_svc.list(limit=1000) agents = await agent_svc.all() @@ -98,6 +151,9 @@ async def overview( "task_success": status_count.get(TaskStatus.SUCCESS.value, 0), "task_failed": status_count.get(TaskStatus.FAILED.value, 0), "agent_total": len(agents), - "agent_online": sum(1 for a in agents if a.status.value == "online"), - "agent_offline": sum(1 for a in agents if a.status.value == "offline"), + "agent_ready": sum(1 for a in agents if a.status == AgentStatus.READY), + "agent_processing": sum(1 for a in agents if a.status == AgentStatus.PROCESSING), + "agent_stopping": sum(1 for a in agents if a.status == AgentStatus.STOPPING), + "agent_offline": sum(1 for a in agents if a.status == AgentStatus.OFFLINE), + "agent_unavailable": sum(1 for a in agents if a.status == AgentStatus.UNAVAILABLE), } \ No newline at end of file diff --git a/backend/app/api/agent.py b/backend/app/api/agent.py index b0d6be7..368a04a 100644 --- a/backend/app/api/agent.py +++ b/backend/app/api/agent.py @@ -27,6 +27,7 @@ async def register(body: AgentRegister, svc: AgentService = Depends(get_agent_se async def unregister(body: dict, svc: AgentService = Depends(get_agent_service)): from fastapi import HTTPException + require_auth(body.get("auth", "")) agent_id = body.get("agent_id") if not agent_id: raise HTTPException(status_code=400, detail="agent_id required") @@ -39,6 +40,7 @@ async def unregister(body: dict, svc: AgentService = Depends(get_agent_service)) async def heartbeat(body: AgentHeartbeat, svc: AgentService = Depends(get_agent_service)): from fastapi import HTTPException + require_auth(body.auth) if not await svc.heartbeat(body): raise HTTPException(status_code=404, detail="agent not registered") return {"ok": True} @@ -48,6 +50,7 @@ async def heartbeat(body: AgentHeartbeat, svc: AgentService = Depends(get_agent_ async def result(body: TaskResult, relay: RelayService = Depends(get_relay_service)): from fastapi import HTTPException + require_auth(body.auth) if not await relay.on_result(body): raise HTTPException(status_code=400, detail="result rejected") return {"ok": True} \ No newline at end of file diff --git a/backend/app/api/cli.py b/backend/app/api/cli.py index 969fa5c..28fa658 100644 --- a/backend/app/api/cli.py +++ b/backend/app/api/cli.py @@ -70,9 +70,14 @@ async def submit_task(body: TaskSubmit, svc: TaskService = Depends(get_task_serv @router.get("/tasks/{request_id}", response_model=TaskInfo, summary="查询任务") -async def get_task(request_id: str, svc: TaskService = Depends(get_task_service)): +async def get_task( + request_id: str, + auth: str = "", + svc: TaskService = Depends(get_task_service), +): from fastapi import HTTPException + require_auth(auth) task = await svc.get(request_id) if not task: raise HTTPException(status_code=404, detail="task not found") @@ -80,9 +85,14 @@ async def get_task(request_id: str, svc: TaskService = Depends(get_task_service) @router.post("/tasks/{request_id}/cancel", summary="取消任务") -async def cancel_task(request_id: str, svc: TaskService = Depends(get_task_service)): +async def cancel_task( + request_id: str, + auth: str = "", + svc: TaskService = Depends(get_task_service), +): from fastapi import HTTPException + require_auth(auth) if not await svc.cancel(request_id): raise HTTPException(status_code=400, detail="cannot cancel task") return {"ok": True, "request_id": request_id} \ No newline at end of file diff --git a/backend/app/config.py b/backend/app/config.py index 033ce07..2a6d052 100644 --- a/backend/app/config.py +++ b/backend/app/config.py @@ -42,6 +42,9 @@ class Settings(BaseSettings): # 认证：CLI 提交任务 / Agent 注册时必须携带的密码 gateway_auth: str = "dev-gateway-auth" + # 认证：Admin 后台管理接口必须携带的密码（与协议层 gateway_auth 分离） + admin_auth: str = "dev-admin-auth" + # 分布式锁相关 lock_timeout: int = 30 diff --git a/backend/app/constants.py b/backend/app/constants.py index 2a3fc8c..255c04b 100644 --- a/backend/app/constants.py +++ b/backend/app/constants.py @@ -10,8 +10,11 @@ class TaskStatus(str, Enum): class AgentStatus(str, Enum): - ONLINE = "online" - OFFLINE = "offline" + OFFLINE = "offline" # 离线：心跳失联（超时） + UNAVAILABLE = "unavailable" # 不可用：后台手动禁用（心跳不复活） + READY = "ready" # 就绪（注册后、空闲时），调度只分配此状态 + PROCESSING = "processing" # 处理中（调度选中并下发后） + STOPPING = "stopping" # 正在停止（收到取消指令后） # ---- 任务池键 ---- diff --git a/backend/app/models/schemas.py b/backend/app/models/schemas.py index d41ee31..7269deb 100644 --- a/backend/app/models/schemas.py +++ b/backend/app/models/schemas.py @@ -53,9 +53,11 @@ class AgentRegister(BaseModel): agent_tags: list[str] = Field(default_factory=list) max_concurrent: int = Field(default=1, ge=1) current_load: int = Field(default=0, ge=0) + priority: int = Field(default=3, ge=1, le=5, description="调度优先级，越小越先分配") class AgentHeartbeat(BaseModel): + auth: str = Field(..., description="网关认证密码，需与 GATEWAY_AUTH 一致") agent_id: str current_load: int = Field(default=0, ge=0) @@ -66,13 +68,15 @@ class AgentInfo(BaseModel): agent_tags: list[str] = Field(default_factory=list) max_concurrent: int = 1 current_load: int = 0 + priority: int = Field(default=3, ge=1, le=5, description="调度优先级，越小越先分配") last_heartbeat: float = Field(default_factory=time.time) - status: AgentStatus = AgentStatus.ONLINE + status: AgentStatus = AgentStatus.READY create_time: float = Field(default_factory=time.time) # ---------- 结果回传 ---------- class TaskResult(BaseModel): + auth: str = Field(..., description="网关认证密码，需与 GATEWAY_AUTH 一致") request_id: str agent_id: str status: TaskStatus = TaskStatus.SUCCESS diff --git a/backend/app/repository/agent_repo.py b/backend/app/repository/agent_repo.py index 4df38e4..a74376b 100644 --- a/backend/app/repository/agent_repo.py +++ b/backend/app/repository/agent_repo.py @@ -37,6 +37,7 @@ class AgentRepo: "agent_tags": ",".join(agent.agent_tags), "max_concurrent": str(agent.max_concurrent), "current_load": str(agent.current_load), + "priority": str(agent.priority), "last_heartbeat": str(agent.last_heartbeat), "status": agent.status.value, "create_time": str(agent.create_time), @@ -51,23 +52,86 @@ class AgentRepo: if not await self.redis.exists(key): return False await self.redis.hset(key, "status", status.value) - if status == AgentStatus.OFFLINE: + # 离线（心跳失联）与不可用均退出心跳 ZSet，避免被误判/干扰 + if status in (AgentStatus.OFFLINE, AgentStatus.UNAVAILABLE): await self.redis.zrem(C.AGENT_HEARTBEAT_ZSET, agent_id) return True + async def set_unavailable(self, agent_id: str) -> bool: + """后台手动置为不可用（unavailable），心跳不会复活。""" + key = self._info_key(agent_id) + if not await self.redis.exists(key): + return False + await self.redis.hset(key, "status", AgentStatus.UNAVAILABLE.value) + await self.redis.zrem(C.AGENT_HEARTBEAT_ZSET, agent_id) + return True + + async def set_available(self, agent_id: str) -> bool: + """后台手动恢复可用（ready），清除不可用状态。""" + key = self._info_key(agent_id) + if not await self.redis.exists(key): + return False + await self.redis.hset(key, "status", AgentStatus.READY.value) + # 恢复就绪时补齐心跳时间戳，避免被误判为超时 + await self.redis.zadd(C.AGENT_HEARTBEAT_ZSET, {agent_id: time.time()}) + return True + + async def mark_idle(self, agent_id: str) -> bool: + """任务结束/取消后标记 Agent 空闲。 + + 若 agent 已被手动置为不可用（unavailable），保持不可用、不复活； + 否则置回 ready。 + """ + key = self._info_key(agent_id) + if not await self.redis.exists(key): + return False + cur = await self.get(agent_id) + if cur is None: + return False + if cur.status == AgentStatus.UNAVAILABLE: + return True + await self.redis.hset(key, "status", AgentStatus.READY.value) + return True + async def beat(self, agent_id: str, current_load: int) -> bool: - """更新心跳时间与负载。""" + """更新心跳时间与负载，按状态机规则流转状态。 + + - 不可用（unavailable）：保持不可用，不因心跳复活。 + - 离线（offline，心跳失联）：恢复 ready。 + - ready/processing/stopping：保持原状态。 + """ key = self._info_key(agent_id) if not await self.redis.exists(key): return False now = time.time() + cur = await self.get(agent_id) + if cur is None: + return False + if cur.status == AgentStatus.UNAVAILABLE: + status = AgentStatus.UNAVAILABLE + elif cur.status == AgentStatus.OFFLINE: + status = AgentStatus.READY + else: + status = cur.status await self.redis.hset( key, - mapping={"last_heartbeat": str(now), "current_load": str(current_load), "status": AgentStatus.ONLINE.value}, + mapping={ + "last_heartbeat": str(now), + "current_load": str(current_load), + "status": status.value, + }, ) await self.redis.zadd(C.AGENT_HEARTBEAT_ZSET, {agent_id: now}) return True + async def update_priority(self, agent_id: str, priority: int) -> bool: + """更新 Agent 调度优先级。""" + key = self._info_key(agent_id) + if not await self.redis.exists(key): + return False + await self.redis.hset(key, "priority", str(priority)) + return True + async def adjust_load(self, agent_id: str, delta: int) -> bool: """原子调整 Agent 负载（delta 可为正/负），负载下限为 0。""" key = self._info_key(agent_id) @@ -90,13 +154,16 @@ class AgentRepo: agent_tags=[t for t in raw.get("agent_tags", "").split(",") if t], max_concurrent=int(raw.get("max_concurrent", 1) or 1), current_load=int(raw.get("current_load", 0) or 0), + priority=int(raw.get("priority", 3) or 3), last_heartbeat=float(raw.get("last_heartbeat", 0) or 0), - status=AgentStatus(raw.get("status", AgentStatus.ONLINE.value)), + status=AgentStatus(raw.get("status", AgentStatus.READY.value)), create_time=float(raw.get("create_time", 0) or 0), ) async def online(self) -> list[AgentInfo]: - return [a for a in await self.all() if a.status == AgentStatus.ONLINE] + # 可参与调度的状态：ready（空闲）与 processing（忙碌但仍有并发容量）。 + # offline（心跳失联）、unavailable（手动禁用）与 stopping（正在停止）不参与调度。 + return [a for a in await self.all() if a.status in (AgentStatus.READY, AgentStatus.PROCESSING)] async def all(self) -> list[AgentInfo]: ids = list(await self.redis.smembers(C.AGENT_ALL_SET)) @@ -108,7 +175,10 @@ class AgentRepo: return out async def by_tags(self, tags: list[str]) -> list[AgentInfo]: - """按标签索引取 Agent 池（取所有标签的交集，无标签则返回全部在线）。""" + """按标签索引取候选 Agent 池（取所有标签的交集，无标签则返回全部可调度）。 + + 只返回可调度的 agent（ready/processing），offline/unavailable/stopping 被排除。 + """ if not tags: return await self.online() keys = [self._tag_key(t) for t in tags] @@ -121,7 +191,7 @@ class AgentRepo: out = [] for aid in ids: a = await self.get(aid) - if a and a.status == AgentStatus.ONLINE: + if a and a.status in (AgentStatus.READY, AgentStatus.PROCESSING): out.append(a) return out diff --git a/backend/app/services/agent_service.py b/backend/app/services/agent_service.py index 81d4783..ea4ac6f 100644 --- a/backend/app/services/agent_service.py +++ b/backend/app/services/agent_service.py @@ -23,7 +23,8 @@ class AgentService: agent_tags=reg.agent_tags, max_concurrent=reg.max_concurrent, current_load=reg.current_load, - status=AgentStatus.ONLINE, + priority=reg.priority, + status=AgentStatus.READY, create_time=time.time(), last_heartbeat=time.time(), ) @@ -54,19 +55,46 @@ class AgentService: async def get(self, agent_id: str) -> AgentInfo | None: return await self.repo.get(agent_id) - async def take_offline(self, agent_id: str) -> bool: - """手动下线 Agent。""" + async def set_unavailable(self, agent_id: str) -> bool: + """后台手动置为「不可用」（unavailable），心跳不会复活。""" a = await self.repo.get(agent_id) if not a: return False - await self.repo.update_status(agent_id, AgentStatus.OFFLINE) - logger.info("agent taken offline agent_id=%s", agent_id) + await self.repo.set_unavailable(agent_id) + logger.info("agent set unavailable agent_id=%s", agent_id) return True + async def set_available(self, agent_id: str) -> bool: + """后台手动恢复「可用」（ready），清除不可用状态。""" + a = await self.repo.get(agent_id) + if not a: + return False + await self.repo.set_available(agent_id) + logger.info("agent set available agent_id=%s", agent_id) + return True + + async def set_priority(self, agent_id: str, priority: int) -> AgentInfo | None: + """设置 Agent 调度优先级（1-5，越小越先分配）。""" + if priority < 1 or priority > 5: + raise ValueError("priority must be in [1, 5]") + ok = await self.repo.update_priority(agent_id, priority) + if not ok: + return None + logger.info("agent priority updated agent_id=%s priority=%s", agent_id, priority) + return await self.repo.get(agent_id) + async def purge_stale(self, timeout: float) -> list[str]: - """剔除超时未心跳的 Agent。""" + """剔除超时未心跳的 Agent：非 unavailable 的置为 offline（心跳失联）。 + + 手动置为不可用的 agent 不参与心跳计时（已从心跳 ZSet 移除），保持 unavailable。 + """ stale = await self.repo.stale_agents(timeout) + purged = [] for aid in stale: + a = await self.repo.get(aid) + if a and a.status == AgentStatus.UNAVAILABLE: + continue await self.repo.update_status(aid, AgentStatus.OFFLINE) logger.info("agent purged (stale) agent_id=%s", aid) - return stale \ No newline at end of file + purged.append(aid) + return purged \ No newline at end of file diff --git a/backend/app/services/notifier.py b/backend/app/services/notifier.py index b676e38..b1456e9 100644 --- a/backend/app/services/notifier.py +++ b/backend/app/services/notifier.py @@ -60,6 +60,24 @@ class TaskNotifier: logger.info("notify task done session=%s request=%s status=%s", task.cli_session_id, task.request_id, task.status.value) + async def publish_task_stop(self, cli_session_id: str, request_id: str) -> None: + """取消任务时向 agent 下发停止指令（复用同一 CLI 会话通道，event=task_stop）。 + + 允许未携带 cli_session_id 时跳过（此时无法通过 SSE 通知，只能等待 agent 超时）。 + """ + if not cli_session_id: + return + payload = json.dumps( + { + "event": "task_stop", + "request_id": request_id, + }, + ensure_ascii=False, + ) + channel = self.channel(cli_session_id) + await self.redis.publish(channel, payload) + logger.info("notify task stop session=%s request=%s", cli_session_id, request_id) + # ---------- 消费 ---------- async def replay(self, cli_session_id: str) -> list[str]: """返回该会话的历史完成事件（新 → 旧），供 SSE 连接回放。""" diff --git a/backend/app/services/relay.py b/backend/app/services/relay.py index 23bcf0a..c052f5b 100644 --- a/backend/app/services/relay.py +++ b/backend/app/services/relay.py @@ -5,7 +5,7 @@ import httpx import redis.asyncio as aioredis from app.config import settings -from app.constants import TaskStatus +from app.constants import AgentStatus, TaskStatus from app.models.schemas import TaskResult from app.repository.agent_repo import AgentRepo from app.repository.task_repo import TaskRepo @@ -35,11 +35,16 @@ class RelayService: if not agent.endpoint: logger.warning("relay dispatch failed: endpoint empty agent=%s request=%s", agent_id, request_id) return False + task = await self.task_repo.get(request_id) url = f"{agent.endpoint.rstrip('/')}/tasks/{request_id}" body = { "auth": settings.gateway_auth, "request_id": request_id, "payload": payload, + "cli_session_id": getattr(task, "cli_session_id", None) if task else None, + "task_type": getattr(task, "task_type", None) if task else None, + "task_tags": getattr(task, "task_tags", None) if task else None, + "timeout": getattr(task, "timeout", None) if task else None, } try: async with httpx.AsyncClient(timeout=5.0) as client: @@ -57,8 +62,14 @@ class RelayService: logger.info("relay dispatch ok agent=%s request=%s url=%s", agent_id, request_id, url) return True + _TERMINAL = {TaskStatus.SUCCESS, TaskStatus.FAILED} + async def on_result(self, result: TaskResult) -> bool: - """反向：Agent 回传结果，更新任务状态并记录日志。""" + """反向：Agent 回传结果，更新任务状态并记录日志。 + + 终态保护：任务已处于终态（如已取消置 failed）时，忽略本次任务状态覆盖， + 但匹配的 Agent 仍会释放负载并置回 ready。 + """ task = await self.task_repo.get(result.request_id) if not task: logger.warning("result for unknown task request=%s", result.request_id) @@ -67,20 +78,25 @@ class RelayService: logger.warning("result agent mismatch request=%s expected=%s got=%s", result.request_id, task.agent_id, result.agent_id) return False - await self.task_repo.update( - result.request_id, - status=result.status, - progress=result.progress, - result=result.result, - error_info=result.error_info, - ) - await self.task_repo.mark_finished(result.request_id) + already_terminal = task.status in self._TERMINAL + if not already_terminal: + await self.task_repo.update( + result.request_id, + status=result.status, + progress=result.progress, + result=result.result, + error_info=result.error_info, + ) + await self.task_repo.mark_finished(result.request_id) # 释放 Agent 算力（任务确实绑定在该 Agent 时才释放） if task.agent_id == result.agent_id: await self.agent_repo.adjust_load(result.agent_id, -1) + # Agent 空闲，置回就绪（若已被手动置为不可用则保持不可用，不会复活） + if task.agent_id: + await self.agent_repo.mark_idle(task.agent_id) await self._log(result.agent_id, result.request_id, "result", f"result received status={result.status.value}") # 任务到达终态，通知对应 CLI 会话（若携带了 cli_session_id） - if task.cli_session_id: + if not already_terminal and task.cli_session_id: finished = await self.task_repo.get(result.request_id) if finished: await TaskNotifier(self.redis).publish_task_done(finished) diff --git a/backend/app/services/scheduler.py b/backend/app/services/scheduler.py index c365d3a..9f71bb5 100644 --- a/backend/app/services/scheduler.py +++ b/backend/app/services/scheduler.py @@ -3,7 +3,7 @@ import logging import redis.asyncio as aioredis -from app.constants import TaskStatus +from app.constants import AgentStatus, TaskStatus from app.models.schemas import AgentInfo from app.repository.agent_repo import AgentRepo from app.repository.task_repo import TaskRepo @@ -26,14 +26,14 @@ class Scheduler: self.relay = relay or RelayService(redis) async def pick_candidate(self, task_tags: list[str]) -> AgentInfo | None: - """按规则挑选最优 Agent：标签匹配→负载过滤→最低负载。""" + """按规则挑选最优 Agent：标签匹配→负载过滤→优先级最小且空闲。""" candidates = await self.agent_repo.by_tags(task_tags) # 负载过滤：剔除满载 candidates = [a for a in candidates if a.current_load < a.max_concurrent] if not candidates: return None - # 最优：负载最低、在线时间最长（稳定） - candidates.sort(key=lambda a: (a.current_load, -a.create_time)) + # 最优：优先级最小（越先分配），同优先级取负载最低、在线时间最长（稳定） + candidates.sort(key=lambda a: (a.priority, a.current_load, -a.create_time)) return candidates[0] async def dispatch(self, request_id: str) -> bool: @@ -52,15 +52,18 @@ class Scheduler: progress=0, ) await self.task_repo.mark_running(request_id) + # 标记 Agent 为处理中 + await self.agent_repo.update_status(agent.agent_id, AgentStatus.PROCESSING) # 更新 Agent 负载 await self.agent_repo.adjust_load(agent.agent_id, 1) # 真实推送 payload 到 Agent 端点 pushed = await self.relay.dispatch_command(agent.agent_id, request_id, task.payload or {}) if not pushed: - # 推送失败：回退 pending 并释放负载，等待下次调度（避免重复扣负载） + # 推送失败：回退 pending 并恢复 agent 就绪、释放负载，等待下次调度（避免重复扣负载） await self.task_repo.update(request_id, status=TaskStatus.PENDING, agent_id="", progress=0) await self.task_repo.mark_pending(request_id) await self.agent_repo.adjust_load(agent.agent_id, -1) + await self.agent_repo.update_status(agent.agent_id, AgentStatus.READY) logger.warning("task dispatch push failed, reverted request_id=%s agent=%s", request_id, agent.agent_id) return False logger.info("task dispatched request_id=%s -> agent=%s", request_id, agent.agent_id) diff --git a/backend/app/services/security.py b/backend/app/services/security.py index 816299b..08279ee 100644 --- a/backend/app/services/security.py +++ b/backend/app/services/security.py @@ -1,10 +1,23 @@ -"""认证校验：CLI / Agent 请求必须携带与 GATEWAY_AUTH 一致的密码。""" +"""认证校验：CLI / Agent 请求携带 GATEWAY_AUTH，后台管理请求携带 ADMIN_AUTH。""" +import secrets + from fastapi import HTTPException from app.config import settings +def _constant_time_eq(a: str, b: str) -> bool: + """常量时间字符串比较，避免时序侧信道。""" + return secrets.compare_digest((a or "").encode("utf-8"), (b or "").encode("utf-8")) + + def require_auth(auth: str) -> None: - """auth 不匹配时抛出 401。""" - if auth != settings.gateway_auth: + """协议层鉴权：auth 必须与 GATEWAY_AUTH 一致，否则 401。""" + if not _constant_time_eq(auth, settings.gateway_auth): raise HTTPException(status_code=401, detail="invalid auth") + + +def require_admin(auth: str) -> None: + """后台管理鉴权：auth 必须与 ADMIN_AUTH 一致，否则 401。""" + if not _constant_time_eq(auth, settings.admin_auth): + raise HTTPException(status_code=401, detail="invalid admin auth") \ No newline at end of file diff --git a/backend/app/services/task_service.py b/backend/app/services/task_service.py index 12de5f6..425d32a 100644 --- a/backend/app/services/task_service.py +++ b/backend/app/services/task_service.py @@ -5,7 +5,7 @@ import time import redis.asyncio as aioredis from app.config import settings -from app.constants import TaskStatus +from app.constants import AgentStatus, TaskStatus from app.models.schemas import TaskInfo, TaskSubmit from app.repository.agent_repo import AgentRepo from app.repository.task_repo import TaskRepo @@ -53,15 +53,26 @@ class TaskService: return await self.repo.list(status=status, limit=limit) async def cancel(self, request_id: str) -> bool: - """取消任务（仅 pending/running 可取消）。""" + """取消任务（仅 pending/running 可取消）。 + + - running 且绑定 agent：置 agent 为 stopping，并通过同一 cli_session_id 的 + SSE 通道下发 task_stop 停止指令；agent 停止后回传结果，由 on_result 终态 + 保护将其置回 ready。 + - pending：无需下发，直接标记失败。 + """ task = await self.repo.get(request_id) if not task: return False if task.status in (TaskStatus.SUCCESS, TaskStatus.FAILED): return False + was_running = task.status == TaskStatus.RUNNING await self._release_if_running(task) await self.repo.update(request_id, status=TaskStatus.FAILED, error_info="cancelled by user") await self.repo.mark_finished(request_id) + if was_running and task.agent_id: + # 通知 agent 停止当前任务 + await self.agent_repo.update_status(task.agent_id, AgentStatus.STOPPING) + await TaskNotifier(self.redis).publish_task_stop(task.cli_session_id, request_id) if task.cli_session_id: finished = await self.repo.get(request_id) if finished: @@ -88,6 +99,8 @@ class TaskService: if not task: continue if task.timeout and (now - task.create_time) > task.timeout: + if task.status == TaskStatus.RUNNING and task.agent_id: + await self.agent_repo.mark_idle(task.agent_id) await self._release_if_running(task) await self.repo.update(rid, status=TaskStatus.FAILED, error_info="timeout") await self.repo.mark_finished(rid) diff --git a/backend/tests/test_agent_repo.py b/backend/tests/test_agent_repo.py index 4a30af9..cd9f0f2 100644 --- a/backend/tests/test_agent_repo.py +++ b/backend/tests/test_agent_repo.py @@ -16,7 +16,20 @@ async def test_upsert_and_get(agent_repo): got = await agent_repo.get("a1") assert got is not None assert got.agent_tags == ["compile", "test"] - assert got.status == AgentStatus.ONLINE + assert got.status == AgentStatus.READY + assert got.priority == 3 # 默认优先级 + + +@pytest.mark.asyncio +async def test_priority_roundtrip_and_update(agent_repo): + a = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], priority=1) + await agent_repo.upsert(a) + assert (await agent_repo.get("a1")).priority == 1 + + assert await agent_repo.update_priority("a1", 5) is True + assert (await agent_repo.get("a1")).priority == 5 + + assert await agent_repo.update_priority("missing", 2) is False @pytest.mark.asyncio @@ -46,6 +59,99 @@ async def test_by_tags(agent_repo): assert ids == {"a2"} +@pytest.mark.asyncio +async def test_by_tags_excludes_offline_unavailable_stopping(agent_repo): + """调度候选排除 offline/unavailable/stopping，允许 ready/processing（processing 仍有并发容量）。""" + a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2) + a2 = AgentInfo(agent_id="a2", endpoint="e2", agent_tags=["compile"], max_concurrent=2) + a3 = AgentInfo(agent_id="a3", endpoint="e3", agent_tags=["compile"], max_concurrent=2) + a4 = AgentInfo(agent_id="a4", endpoint="e4", agent_tags=["compile"], max_concurrent=2) + a5 = AgentInfo(agent_id="a5", endpoint="e5", agent_tags=["compile"], max_concurrent=2) + await agent_repo.upsert(a1) + await agent_repo.upsert(a2) + await agent_repo.upsert(a3) + await agent_repo.upsert(a4) + await agent_repo.upsert(a5) + await agent_repo.update_status("a2", AgentStatus.OFFLINE) # 离线（心跳失联）不可调度 + await agent_repo.set_unavailable("a3") # 手动不可用不可调度 + await agent_repo.update_status("a4", AgentStatus.STOPPING) # 停止中不可调度 + await agent_repo.update_status("a5", AgentStatus.PROCESSING) # 处理中但还有容量，可调度 + + ids = {a.agent_id for a in await agent_repo.by_tags(["compile"])} + assert ids == {"a1", "a5"} + + +@pytest.mark.asyncio +async def test_unavailable_heartbeat_not_revive(agent_repo): + """手动置不可用（unavailable）后，即使发心跳也不复活。""" + a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2) + await agent_repo.upsert(a1) + await agent_repo.set_unavailable("a1") + assert (await agent_repo.get("a1")).status == AgentStatus.UNAVAILABLE + # 心跳不应复活 + await agent_repo.beat("a1", 0) + assert (await agent_repo.get("a1")).status == AgentStatus.UNAVAILABLE + # 恢复可用 + await agent_repo.set_available("a1") + assert (await agent_repo.get("a1")).status == AgentStatus.READY + + +@pytest.mark.asyncio +async def test_unavailable_distinct_from_offline(agent_repo): + """离线（心跳失联）与不可用（手动禁用）是两个独立状态。""" + a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2) + a2 = AgentInfo(agent_id="a2", endpoint="e2", agent_tags=["compile"], max_concurrent=2) + await agent_repo.upsert(a1) + await agent_repo.upsert(a2) + # a1 心跳失联→离线；a2 手动禁用→不可用 + await agent_repo.update_status("a1", AgentStatus.OFFLINE) + await agent_repo.set_unavailable("a2") + assert (await agent_repo.get("a1")).status == AgentStatus.OFFLINE + assert (await agent_repo.get("a2")).status == AgentStatus.UNAVAILABLE + # 发心跳：离线的 a1 恢复 ready，不可用的 a2 保持 unavailable + await agent_repo.beat("a1", 0) + await agent_repo.beat("a2", 0) + assert (await agent_repo.get("a1")).status == AgentStatus.READY + assert (await agent_repo.get("a2")).status == AgentStatus.UNAVAILABLE + + +@pytest.mark.asyncio +async def test_mark_idle_keeps_unavailable(agent_repo): + """任务结束后 mark_idle：unavailable 的 agent 保持不可用，其余置回 ready。""" + a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2) + a2 = AgentInfo(agent_id="a2", endpoint="e2", agent_tags=["compile"], max_concurrent=2) + await agent_repo.upsert(a1) + await agent_repo.upsert(a2) + await agent_repo.set_unavailable("a1") + await agent_repo.update_status("a2", AgentStatus.PROCESSING) + await agent_repo.mark_idle("a1") + await agent_repo.mark_idle("a2") + assert (await agent_repo.get("a1")).status == AgentStatus.UNAVAILABLE + assert (await agent_repo.get("a2")).status == AgentStatus.READY + + +@pytest.mark.asyncio +async def test_heartbeat_recovers_stale_offline(agent_repo): + """超时离线（非手动）的 agent 在发心跳后恢复为 ready。""" + a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2) + await agent_repo.upsert(a1) + await agent_repo.update_status("a1", AgentStatus.OFFLINE) + assert (await agent_repo.get("a1")).status == AgentStatus.OFFLINE + await agent_repo.beat("a1", 0) + assert (await agent_repo.get("a1")).status == AgentStatus.READY + + +@pytest.mark.asyncio +async def test_heartbeat_keeps_busy_states(agent_repo): + """processing/stopping 状态不因心跳改变。""" + for st in (AgentStatus.PROCESSING, AgentStatus.STOPPING): + a = AgentInfo(agent_id=f"a-{st.value}", endpoint="e1", agent_tags=["compile"], max_concurrent=2) + await agent_repo.upsert(a) + await agent_repo.update_status(a.agent_id, st) + await agent_repo.beat(a.agent_id, 1) + assert (await agent_repo.get(a.agent_id)).status == st + + @pytest.mark.asyncio async def test_remove(agent_repo): a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"]) diff --git a/backend/tests/test_api.py b/backend/tests/test_api.py index a085d98..8796bc8 100644 --- a/backend/tests/test_api.py +++ b/backend/tests/test_api.py @@ -8,6 +8,7 @@ from app.config import settings from app.main import app AUTH = settings.gateway_auth +ADMIN_AUTH = settings.admin_auth @pytest_asyncio.fixture @@ -36,7 +37,7 @@ async def test_submit_and_query_task(client): assert data["status"] == "pending" rid = data["request_id"] - r2 = await client.get(f"/api/cli/tasks/{rid}") + r2 = await client.get(f"/api/cli/tasks/{rid}", params={"auth": AUTH}) assert r2.status_code == 200 assert r2.json()["request_id"] == rid @@ -48,7 +49,7 @@ async def test_agent_register_heartbeat_result(client): assert r.status_code == 200 # 心跳 - r = await client.post("/api/agent/heartbeat", json={"agent_id": "a1", "current_load": 0}) + r = await client.post("/api/agent/heartbeat", json={"auth": AUTH, "agent_id": "a1", "current_load": 0}) assert r.status_code == 200 # 提交任务并等待调度（手动触发一次调度） @@ -58,20 +59,44 @@ async def test_agent_register_heartbeat_result(client): await Scheduler(app.state.redis).dispatch_pending(10) # 回传结果 - r = await client.post("/api/agent/result", json={"request_id": rid, "agent_id": "a1", "status": "success", "progress": 100, "result": {"ok": True}}) + r = await client.post("/api/agent/result", json={"auth": AUTH, "request_id": rid, "agent_id": "a1", "status": "success", "progress": 100, "result": {"ok": True}}) assert r.status_code == 200 - detail = await client.get(f"/api/cli/tasks/{rid}") + detail = await client.get(f"/api/cli/tasks/{rid}", params={"auth": AUTH}) assert detail.json()["status"] == "success" @pytest.mark.asyncio async def test_admin_overview(client): - r = await client.get("/api/admin/overview") + r = await client.get("/api/admin/overview", headers={"X-Admin-Auth": ADMIN_AUTH}) assert r.status_code == 200 assert "task_total" in r.json() +@pytest.mark.asyncio +async def test_admin_rejects_bad_auth(client): + r = await client.get("/api/admin/overview", headers={"X-Admin-Auth": "wrong-admin"}) + assert r.status_code == 401 + + +@pytest.mark.asyncio +async def test_admin_requires_auth(client): + r = await client.get("/api/admin/overview") + assert r.status_code == 422 # 缺 X-Admin-Auth header + + +@pytest.mark.asyncio +async def test_agent_heartbeat_rejects_bad_auth(client): + r = await client.post("/api/agent/heartbeat", json={"auth": "wrong", "agent_id": "x1", "current_load": 0}) + assert r.status_code == 401 + + +@pytest.mark.asyncio +async def test_cli_query_rejects_bad_auth(client): + r = await client.get("/api/cli/tasks/whatever", params={"auth": "wrong"}) + assert r.status_code == 401 + + @pytest.mark.asyncio async def test_submit_task_rejects_bad_auth(client): r = await client.post("/api/cli/tasks", json={"auth": "wrong-password", "task_type": "compile"}) diff --git a/backend/tests/test_scheduler.py b/backend/tests/test_scheduler.py index 182eb79..f65a74d 100644 --- a/backend/tests/test_scheduler.py +++ b/backend/tests/test_scheduler.py @@ -1,13 +1,27 @@ """调度服务单元测试。""" import pytest -from app.constants import TaskStatus +from app.constants import AgentStatus, TaskStatus from app.models.schemas import AgentInfo, TaskInfo from app.repository.agent_repo import AgentRepo from app.repository.task_repo import TaskRepo from app.services.scheduler import Scheduler +class _StubRelay: + """模拟 relay：不发起真实网络请求，直接返回推送成功。""" + + def __init__(self, ok: bool = True): + self.ok = ok + + async def dispatch_command(self, agent_id: str, request_id: str, payload: dict) -> bool: + return self.ok + + +def _sched(redis, task_repo, agent_repo): + return Scheduler(redis, task_repo, agent_repo, relay=_StubRelay()) + + @pytest.mark.asyncio async def test_dispatch_binds_min_load_agent(redis): task_repo = TaskRepo(redis) @@ -19,7 +33,7 @@ async def test_dispatch_binds_min_load_agent(redis): await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"])) - sched = Scheduler(redis, task_repo, agent_repo) + sched = _sched(redis, task_repo, agent_repo) ok = await sched.dispatch("t1") assert ok is True @@ -29,6 +43,35 @@ async def test_dispatch_binds_min_load_agent(redis): assert "t1" in await task_repo.running_tasks() +@pytest.mark.asyncio +async def test_pick_candidate_lowest_priority(redis): + agent_repo = AgentRepo(redis) + + # 三个 compile Agent：b1 负载最低但优先级高(最不被优先)，b2 负载高但优先级低(最优先) + await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2, current_load=0, priority=5)) + await agent_repo.upsert(AgentInfo(agent_id="b2", endpoint="e2", agent_tags=["compile"], max_concurrent=2, current_load=1, priority=1)) + await agent_repo.upsert(AgentInfo(agent_id="b3", endpoint="e3", agent_tags=["compile"], max_concurrent=2, current_load=0, priority=3)) + + sched = Scheduler(redis) + picked = await sched.pick_candidate(["compile"]) + assert picked is not None + assert picked.agent_id == "b2" # 优先级最小优先分配 + + +@pytest.mark.asyncio +async def test_pick_candidate_same_priority_min_load(redis): + agent_repo = AgentRepo(redis) + + # 同优先级(默认3)时，负载更低者优先 + await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2, current_load=1)) + await agent_repo.upsert(AgentInfo(agent_id="b2", endpoint="e2", agent_tags=["compile"], max_concurrent=2, current_load=0)) + + sched = Scheduler(redis) + picked = await sched.pick_candidate(["compile"]) + assert picked is not None + assert picked.agent_id == "b2" + + @pytest.mark.asyncio async def test_dispatch_no_agent_stays_pending(redis): task_repo = TaskRepo(redis) @@ -62,6 +105,70 @@ async def test_dispatch_pending(redis): await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"])) await task_repo.create(TaskInfo(request_id="t2", task_type="compile", task_tags=["compile"])) - sched = Scheduler(redis, task_repo, agent_repo) + sched = _sched(redis, task_repo, agent_repo) n = await sched.dispatch_pending(10) - assert n == 2 \ No newline at end of file + assert n == 2 + + +@pytest.mark.asyncio +async def test_dispatch_excludes_offline_unavailable_stopping(redis): + """调度排除 offline（心跳失联）、unavailable（手动不可用）与 stopping 的 agent。""" + task_repo = TaskRepo(redis) + agent_repo = AgentRepo(redis) + # b1 手动不可用，b2 正在停止，b3 离线，b4 就绪 + await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)) + await agent_repo.upsert(AgentInfo(agent_id="b2", endpoint="e2", agent_tags=["compile"], max_concurrent=2)) + await agent_repo.upsert(AgentInfo(agent_id="b3", endpoint="e3", agent_tags=["compile"], max_concurrent=2)) + await agent_repo.upsert(AgentInfo(agent_id="b4", endpoint="e4", agent_tags=["compile"], max_concurrent=2)) + await agent_repo.set_unavailable("b1") + await agent_repo.update_status("b2", AgentStatus.STOPPING) + await agent_repo.update_status("b3", AgentStatus.OFFLINE) + + await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"])) + sched = _sched(redis, task_repo, agent_repo) + ok = await sched.dispatch("t1") + assert ok is True + task = await task_repo.get("t1") + assert task.agent_id == "b4" + + +@pytest.mark.asyncio +async def test_dispatch_processing_agent_with_capacity(redis): + """processing 的 agent 若还有并发容量，仍可接收新任务。""" + task_repo = TaskRepo(redis) + agent_repo = AgentRepo(redis) + await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)) + await agent_repo.update_status("b1", AgentStatus.PROCESSING) + await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"])) + sched = _sched(redis, task_repo, agent_repo) + assert await sched.dispatch("t1") is True + assert (await task_repo.get("t1")).agent_id == "b1" + + +@pytest.mark.asyncio +async def test_dispatch_marks_agent_processing(redis): + """调度下发成功后，agent 状态流转为 processing。""" + task_repo = TaskRepo(redis) + agent_repo = AgentRepo(redis) + await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)) + await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"])) + sched = _sched(redis, task_repo, agent_repo) + assert await sched.dispatch("t1") is True + assert (await agent_repo.get("b1")).status == AgentStatus.PROCESSING + + +@pytest.mark.asyncio +async def test_dispatch_push_fail_reverts_agent_to_ready(redis): + """推送失败时任务回退 pending，agent 释放负载并回到 ready。""" + task_repo = TaskRepo(redis) + agent_repo = AgentRepo(redis) + await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)) + await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"])) + sched = Scheduler(redis, task_repo, agent_repo, relay=_StubRelay(ok=False)) + assert await sched.dispatch("t1") is False + task = await task_repo.get("t1") + assert task.status == TaskStatus.PENDING + assert not task.agent_id + a = await agent_repo.get("b1") + assert a.status == AgentStatus.READY + assert a.current_load == 0 \ No newline at end of file diff --git a/frontend/src/App.vue b/frontend/src/App.vue index 358e74b..252752d 100644 --- a/frontend/src/App.vue +++ b/frontend/src/App.vue @@ -1,14 +1,43 @@