218 lines
8.1 KiB
Python
218 lines
8.1 KiB
Python
"""A2A 网关接入客户端：Agent 服务启动后自动注册 + 心跳保活 + 注销。
|
||
|
||
用法（在 api_server.py 的 main() 中）:
|
||
from gateway_client import register_agent, start_heartbeat, unregister_agent
|
||
|
||
if register_agent(dev_app.name, f"http://127.0.0.1:{PORT}"):
|
||
start_heartbeat(dev_app.name)
|
||
|
||
环境变量:
|
||
GATEWAY_URL 网关地址，默认 http://127.0.0.1:8000
|
||
GATEWAY_AUTH 网关认证密码，需与网关 GATEWAY_AUTH 一致
|
||
AGENT_TAGS 能力标签，逗号分隔，默认 code
|
||
AGENT_MAX_CONCURRENT 最大并发数，默认 1
|
||
HEARTBEAT_INTERVAL 心跳间隔秒，默认 10
|
||
"""
|
||
import logging
|
||
import os
|
||
import threading
|
||
import time
|
||
|
||
import httpx
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
GATEWAY_URL = os.getenv("GATEWAY_URL", "http://127.0.0.1:8000").rstrip("/")
|
||
GATEWAY_AUTH = os.getenv("GATEWAY_AUTH", "dev-gateway-auth")
|
||
HEARTBEAT_INTERVAL = int(os.getenv("HEARTBEAT_INTERVAL", "10"))
|
||
|
||
_client = httpx.Client(timeout=5.0)
|
||
|
||
# 取消/停止指令下发：网关复用 CLI 的 /api/cli/events 通道（event=task_stop）
|
||
# Agent 通过同一个 cli_session_id 订阅，收到匹配自身正在执行 request_id 的 task_stop 时置停止标志。
|
||
_stop_registry: dict[str, threading.Event] = {}
|
||
|
||
|
||
def mark_stop_requested(request_id: str) -> None:
|
||
"""标记某任务需要停止（收到 task_stop 后调用）。"""
|
||
ev = _stop_registry.get(request_id)
|
||
if ev is None:
|
||
ev = threading.Event()
|
||
_stop_registry[request_id] = ev
|
||
ev.set()
|
||
|
||
|
||
def clear_stop_requested(request_id: str) -> None:
|
||
"""任务开始执行前清除停止标志。"""
|
||
ev = _stop_registry.get(request_id)
|
||
if ev is not None:
|
||
ev.clear()
|
||
|
||
|
||
def is_stop_requested(request_id: str) -> bool:
|
||
"""判断某任务是否已被要求停止（供执行循环轮询检查）。"""
|
||
ev = _stop_registry.get(request_id)
|
||
return ev is not None and ev.is_set()
|
||
|
||
|
||
def wait_stop(request_id: str, timeout: float = 0.5) -> bool:
|
||
"""等待停止标志，返回 True 表示已收到停止请求。执行循环可用它做可中断 sleep。"""
|
||
ev = _stop_registry.get(request_id)
|
||
if ev is None:
|
||
ev = threading.Event()
|
||
_stop_registry[request_id] = ev
|
||
return ev.wait(timeout)
|
||
|
||
|
||
def _sse_subscribe_poll(cli_session_id: str) -> None:
|
||
"""后台线程：订阅网关 /api/cli/events 通道，识别 task_stop 指令并标记停止。"""
|
||
url = f"{GATEWAY_URL}/api/cli/events?cli_session_id={cli_session_id}&auth={GATEWAY_AUTH}"
|
||
while True:
|
||
try:
|
||
with _client.stream("GET", url, timeout=None) as resp:
|
||
if resp.status_code != 200:
|
||
logger.warning("sse subscribe failed status=%s", resp.status_code)
|
||
time.sleep(HEARTBEAT_INTERVAL)
|
||
continue
|
||
for line in resp.iter_lines():
|
||
if not line or not line.startswith("data:"):
|
||
continue
|
||
try:
|
||
import json
|
||
evt = json.loads(line[len("data:"):].strip())
|
||
except Exception:
|
||
continue
|
||
if evt.get("event") == "task_stop":
|
||
rid = evt.get("request_id")
|
||
if rid:
|
||
mark_stop_requested(rid)
|
||
logger.info("stop requested received request=%s", rid)
|
||
except Exception as e:
|
||
logger.warning("sse subscribe loop error err=%s", e)
|
||
time.sleep(HEARTBEAT_INTERVAL)
|
||
|
||
|
||
def start_stop_listener(cli_session_id: str) -> threading.Thread:
|
||
"""启动 SSE 停止指令订阅线程（daemon）。"""
|
||
t = threading.Thread(
|
||
target=_sse_subscribe_poll,
|
||
args=(cli_session_id,),
|
||
name=f"gateway-sse-{cli_session_id[:8]}",
|
||
daemon=True,
|
||
)
|
||
t.start()
|
||
logger.info("stop listener started session=%s", cli_session_id)
|
||
return t
|
||
|
||
|
||
def _parse_tags() -> list[str]:
|
||
raw = os.getenv("AGENT_TAGS", "code")
|
||
return [t.strip() for t in raw.split(",") if t.strip()]
|
||
|
||
|
||
def register_agent(agent_id: str, endpoint: str, agent_tags: list[str] | None = None,
|
||
max_concurrent: int | None = None) -> bool:
|
||
"""向网关注册本 Agent。成功返回 True，失败（网关未启动/认证失败）返回 False。"""
|
||
tags = agent_tags if agent_tags is not None else _parse_tags()
|
||
max_conc = max_concurrent or int(os.getenv("AGENT_MAX_CONCURRENT", "1"))
|
||
body = {
|
||
"auth": GATEWAY_AUTH,
|
||
"agent_id": agent_id,
|
||
"endpoint": endpoint,
|
||
"agent_tags": tags,
|
||
"max_concurrent": max_conc,
|
||
"current_load": 0,
|
||
}
|
||
try:
|
||
resp = _client.post(f"{GATEWAY_URL}/api/agent/register", json=body)
|
||
if resp.status_code == 200:
|
||
logger.info("registered to gateway agent=%s endpoint=%s tags=%s", agent_id, endpoint, tags)
|
||
return True
|
||
logger.error("register failed agent=%s status=%s body=%s", agent_id, resp.status_code, resp.text)
|
||
except httpx.HTTPError as e:
|
||
logger.error("register network error agent=%s err=%s", agent_id, e)
|
||
return False
|
||
|
||
|
||
def _heartbeat_once(agent_id: str, current_load: int) -> bool:
|
||
try:
|
||
resp = _client.post(
|
||
f"{GATEWAY_URL}/api/agent/heartbeat",
|
||
json={"auth": GATEWAY_AUTH, "agent_id": agent_id, "current_load": current_load},
|
||
)
|
||
if resp.status_code == 200:
|
||
return True
|
||
logger.warning("heartbeat failed status=%s body=%s", resp.status_code, resp.text)
|
||
except httpx.HTTPError as e:
|
||
logger.warning("heartbeat network error err=%s", e)
|
||
return False
|
||
|
||
|
||
def _heartbeat_loop(agent_id: str) -> None:
|
||
while True:
|
||
try:
|
||
load = get_current_load()
|
||
except Exception:
|
||
load = 0
|
||
_heartbeat_once(agent_id, load)
|
||
time.sleep(HEARTBEAT_INTERVAL)
|
||
|
||
|
||
def start_heartbeat(agent_id: str) -> threading.Thread:
|
||
"""启动后台心跳线程（daemon，随进程退出）。"""
|
||
t = threading.Thread(target=_heartbeat_loop, args=(agent_id,), name=f"gateway-hb-{agent_id}", daemon=True)
|
||
t.start()
|
||
logger.info("heartbeat started agent=%s interval=%ss", agent_id, HEARTBEAT_INTERVAL)
|
||
return t
|
||
|
||
|
||
def unregister_agent(agent_id: str) -> bool:
|
||
"""向网关注销本 Agent。"""
|
||
try:
|
||
resp = _client.post(f"{GATEWAY_URL}/api/agent/unregister", json={"auth": GATEWAY_AUTH, "agent_id": agent_id})
|
||
if resp.status_code in (200, 404):
|
||
logger.info("unregistered from gateway agent=%s", agent_id)
|
||
return True
|
||
logger.warning("unregister failed status=%s body=%s", resp.status_code, resp.text)
|
||
except httpx.HTTPError as e:
|
||
logger.error("unregister network error agent=%s err=%s", agent_id, e)
|
||
return False
|
||
|
||
|
||
def report_result(request_id: str, agent_id: str, status: str = "success",
|
||
progress: int = 100, result: dict | None = None,
|
||
error_info: str | None = None) -> bool:
|
||
"""任务执行完成后，将结果回传给网关（POST /api/agent/result）。
|
||
|
||
Args:
|
||
request_id: 网关分配的任务 ID
|
||
agent_id: 本 Agent ID
|
||
status: success / failed
|
||
progress: 进度百分比
|
||
result: 结果字典（可选）
|
||
error_info: 错误信息（失败时必填）
|
||
"""
|
||
body = {
|
||
"auth": GATEWAY_AUTH,
|
||
"request_id": request_id,
|
||
"agent_id": agent_id,
|
||
"status": status,
|
||
"progress": progress,
|
||
"result": result,
|
||
"error_info": error_info,
|
||
}
|
||
try:
|
||
resp = _client.post(f"{GATEWAY_URL}/api/agent/result", json=body)
|
||
if resp.status_code == 200:
|
||
logger.info("result reported request=%s status=%s", request_id, status)
|
||
return True
|
||
logger.error("report result failed request=%s status=%s body=%s", request_id, resp.status_code, resp.text)
|
||
except httpx.HTTPError as e:
|
||
logger.error("report result network error request=%s err=%s", request_id, e)
|
||
return False
|
||
|
||
|
||
def get_current_load() -> int:
|
||
"""Agent 当前负载。子类/调用方可覆写以报告真实并发数。"""
|
||
return 0
|