394 lines
15 KiB
Python
394 lines
15 KiB
Python
"""子智能体状态管理 Mixin。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from utils.logger import setup_logger
|
||
from modules.multi_agent.debug_logger import ma_debug
|
||
|
||
logger = setup_logger(__name__)
|
||
TERMINAL_STATUSES = {"completed", "failed", "timeout"}
|
||
|
||
|
||
class SubAgentStateMixin:
|
||
"""提供子智能体任务状态加载、保存、刷新与结果落盘能力。"""
|
||
|
||
state_file: Path
|
||
tasks: Dict[str, Dict[str, Any]]
|
||
conversation_agents: Dict[str, List[int]]
|
||
_running_tasks: Dict[str, Any]
|
||
_state_lock: Any
|
||
|
||
def _load_state(self):
|
||
with self._state_lock:
|
||
if not self.state_file.exists():
|
||
self.tasks = {}
|
||
self.conversation_agents = {}
|
||
return
|
||
try:
|
||
data = json.loads(self.state_file.read_text(encoding="utf-8"))
|
||
loaded_tasks = data.get("tasks", {})
|
||
loaded_agents = data.get("conversation_agents", {})
|
||
except json.JSONDecodeError:
|
||
logger.warning("子智能体状态文件损坏，已忽略。")
|
||
self.tasks = {}
|
||
self.conversation_agents = {}
|
||
return
|
||
|
||
runtime_only_keys = {"_stdout_lines"}
|
||
merged_tasks: Dict[str, Dict] = {}
|
||
for task_id, task in loaded_tasks.items():
|
||
existing = self.tasks.get(task_id)
|
||
if existing:
|
||
for key in runtime_only_keys:
|
||
if key in existing:
|
||
task[key] = existing[key]
|
||
merged_tasks[task_id] = task
|
||
self.tasks = merged_tasks
|
||
self.conversation_agents = loaded_agents
|
||
|
||
if self.tasks:
|
||
migrated = False
|
||
for task in self.tasks.values():
|
||
if task.get("parent_conversation_id"):
|
||
continue
|
||
candidate = task.get("conversation_id") or (task.get("service_payload") or {}).get("parent_conversation_id")
|
||
if candidate:
|
||
task["parent_conversation_id"] = candidate
|
||
migrated = True
|
||
if migrated:
|
||
self._save_state_unsafe()
|
||
|
||
def _save_state(self):
|
||
with self._state_lock:
|
||
self._save_state_unsafe()
|
||
|
||
def _save_state_unsafe(self):
|
||
payload = {
|
||
"tasks": self.tasks,
|
||
"conversation_agents": self.conversation_agents,
|
||
}
|
||
try:
|
||
self.state_file.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
except Exception as exc:
|
||
logger.warning(f"保存子智能体状态失败: {exc}")
|
||
|
||
def _check_task_status(self, task: Dict) -> Dict:
|
||
"""检查任务状态，读取输出文件。"""
|
||
task_id = task["task_id"]
|
||
output_file = Path(task.get("output_file", ""))
|
||
if not output_file.exists():
|
||
running_task = self._running_tasks.get(task_id)
|
||
if running_task and running_task.done():
|
||
try:
|
||
running_task.result(timeout=0)
|
||
except Exception:
|
||
pass
|
||
task["status"] = "failed"
|
||
task["updated_at"] = time.time()
|
||
return {"status": task.get("status", "running"), "task_id": task_id}
|
||
|
||
try:
|
||
output = json.loads(output_file.read_text(encoding="utf-8"))
|
||
except Exception as exc:
|
||
logger.warning(f"[check_task_status] output 文件解析失败: {exc}")
|
||
task["status"] = "failed"
|
||
task["updated_at"] = time.time()
|
||
return {"success": False, "status": "failed", "task_id": task_id, "message": f"输出文件解析失败: {exc}"}
|
||
|
||
success = output.get("success", False)
|
||
summary = output.get("summary", "")
|
||
stats = output.get("stats", {})
|
||
elapsed_seconds = self._compute_elapsed_seconds(task)
|
||
|
||
ma_debug(
|
||
"check_task_status",
|
||
task_id=task_id,
|
||
agent_id=task.get("agent_id"),
|
||
multi_agent_mode=task.get("multi_agent_mode"),
|
||
output_status=output.get("status"),
|
||
output_success=success,
|
||
)
|
||
|
||
# 多智能体模式：output 中 status 为 running/idle 表示子智能体仍在运行或
|
||
# 本轮结束但上下文保留、可继续接收消息，都不是失败/完成。
|
||
# 此时不生成 final_result，只保持原状态，避免 reconcile 把运行中误判为失败。
|
||
if task.get("multi_agent_mode") and output.get("status") in {"running", "idle"}:
|
||
task["status"] = output["status"]
|
||
task["updated_at"] = time.time()
|
||
ma_debug("check_task_status_keep_alive", task_id=task_id, status=output["status"])
|
||
return {"status": output["status"], "task_id": task_id}
|
||
|
||
if output.get("timeout"):
|
||
status = "timeout"
|
||
elif output.get("max_turns_exceeded"):
|
||
status = "failed"
|
||
summary = f"任务执行超过最大轮次限制。{summary}"
|
||
elif success:
|
||
status = "completed"
|
||
else:
|
||
status = "failed"
|
||
|
||
ma_debug("check_task_status_result", task_id=task_id, agent_id=task.get("agent_id"), status=status)
|
||
task["status"] = status
|
||
task["updated_at"] = time.time()
|
||
if status == "completed" and elapsed_seconds is not None:
|
||
task["elapsed_seconds"] = elapsed_seconds
|
||
task["runtime_seconds"] = elapsed_seconds
|
||
|
||
agent_id = task.get("agent_id")
|
||
task_summary = task.get("summary")
|
||
deliverables_dir = task.get("deliverables_dir")
|
||
stats_summary = self._build_stats_summary(stats)
|
||
|
||
if status == "completed":
|
||
system_message = self._compose_sub_agent_message(
|
||
prefix=f"✅ 子智能体{agent_id} 任务摘要：{task_summary} 已完成。",
|
||
stats_summary=stats_summary,
|
||
summary=summary,
|
||
deliverables_dir=deliverables_dir,
|
||
duration_seconds=elapsed_seconds,
|
||
)
|
||
elif status == "timeout":
|
||
system_message = self._compose_sub_agent_message(
|
||
prefix=f"⏱️ 子智能体{agent_id} 任务摘要：{task_summary} 超时未完成。",
|
||
stats_summary=stats_summary,
|
||
summary=summary,
|
||
)
|
||
else:
|
||
system_message = self._compose_sub_agent_message(
|
||
prefix=f"❌ 子智能体{agent_id} 任务摘要：{task_summary} 执行失败。",
|
||
stats_summary=stats_summary,
|
||
summary=summary,
|
||
)
|
||
|
||
result = {
|
||
"success": success,
|
||
"status": status,
|
||
"task_id": task_id,
|
||
"agent_id": agent_id,
|
||
"message": summary,
|
||
"deliverables_dir": deliverables_dir,
|
||
"stats": stats,
|
||
"stats_summary": stats_summary,
|
||
"system_message": system_message,
|
||
}
|
||
if status == "completed" and elapsed_seconds is not None:
|
||
result["elapsed_seconds"] = elapsed_seconds
|
||
result["runtime_seconds"] = elapsed_seconds
|
||
task["final_result"] = result
|
||
return result
|
||
|
||
def _handle_timeout(self, task: Dict) -> Dict:
|
||
"""处理任务超时。"""
|
||
task_id = task["task_id"]
|
||
running_task = self._running_tasks.pop(task_id, None)
|
||
if running_task and not running_task.done():
|
||
running_task.cancel()
|
||
deadline = time.time() + 5
|
||
while not running_task.done() and time.time() < deadline:
|
||
time.sleep(0.05)
|
||
|
||
task["status"] = "timeout"
|
||
task["updated_at"] = time.time()
|
||
|
||
stats = {}
|
||
stats_file = Path(task.get("stats_file", ""))
|
||
if stats_file.exists():
|
||
try:
|
||
stats = json.loads(stats_file.read_text(encoding="utf-8"))
|
||
except Exception:
|
||
stats = {}
|
||
stats_summary = self._build_stats_summary(stats)
|
||
system_message = self._compose_sub_agent_message(
|
||
prefix=f"⏱️ 子智能体{task.get('agent_id')} 任务摘要：{task.get('summary')} 超时未完成。",
|
||
stats_summary=stats_summary,
|
||
summary="等待超时，子智能体已被终止。",
|
||
)
|
||
|
||
result = {
|
||
"success": False,
|
||
"status": "timeout",
|
||
"task_id": task_id,
|
||
"agent_id": task.get("agent_id"),
|
||
"message": "等待超时，子智能体已被终止。",
|
||
"stats": stats,
|
||
"stats_summary": stats_summary,
|
||
"system_message": system_message,
|
||
}
|
||
task["final_result"] = result
|
||
self._save_state()
|
||
return result
|
||
|
||
def _mark_task_terminated(
|
||
self,
|
||
task: Dict[str, Any],
|
||
*,
|
||
message: str,
|
||
system_message: Optional[str] = None,
|
||
notified: bool = False,
|
||
) -> Dict[str, Any]:
|
||
task["status"] = "terminated"
|
||
task["updated_at"] = time.time()
|
||
if notified:
|
||
task["notified"] = True
|
||
result = {
|
||
"success": False,
|
||
"status": "terminated",
|
||
"task_id": task.get("task_id"),
|
||
"agent_id": task.get("agent_id"),
|
||
"message": message,
|
||
"system_message": system_message or message,
|
||
}
|
||
task["final_result"] = result
|
||
return result
|
||
|
||
def _mark_task_done(
|
||
self,
|
||
task_id: str,
|
||
success: bool,
|
||
summary: str,
|
||
runtime_seconds: int,
|
||
) -> None:
|
||
task = self.tasks.get(task_id)
|
||
if not task:
|
||
return
|
||
status = "completed" if success else "failed"
|
||
ma_debug(
|
||
"mark_task_done",
|
||
task_id=task_id,
|
||
agent_id=task.get("agent_id"),
|
||
success=success,
|
||
new_status=status,
|
||
)
|
||
task["status"] = status
|
||
task["updated_at"] = time.time()
|
||
task["runtime_seconds"] = runtime_seconds
|
||
self._check_task_status(task)
|
||
self._save_state()
|
||
|
||
def _refresh_task_runtime_state(self, task: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""刷新单个任务运行态。"""
|
||
status = task.get("status")
|
||
task_id = task.get("task_id")
|
||
agent_id = task.get("agent_id")
|
||
multi_agent_flag = task.get("multi_agent_mode")
|
||
ma_debug(
|
||
"refresh_task_runtime_state_start",
|
||
task_id=task_id,
|
||
agent_id=agent_id,
|
||
before_status=status,
|
||
multi_agent_mode=multi_agent_flag,
|
||
)
|
||
if status in TERMINAL_STATUSES.union({"terminated"}):
|
||
return {"status": status, "task_id": task_id}
|
||
|
||
running_task = self._running_tasks.get(task_id) if task_id else None
|
||
if running_task:
|
||
if running_task.done():
|
||
try:
|
||
running_task.result(timeout=0)
|
||
except Exception:
|
||
pass
|
||
# 多智能体模式：任务自然进入 idle 时不写输出文件，不应标记为失败
|
||
if task.get("multi_agent_mode") and not Path(task.get("output_file", "")).exists():
|
||
task["status"] = "idle"
|
||
task["updated_at"] = time.time()
|
||
ma_debug("refresh_task_runtime_state_idle_no_output", task_id=task_id)
|
||
return {"status": "idle", "task_id": task_id}
|
||
result = self._check_task_status(task)
|
||
ma_debug("refresh_task_runtime_state_done_result", task_id=task_id, result=result)
|
||
return result
|
||
ma_debug("refresh_task_runtime_state_still_running", task_id=task_id)
|
||
return {"status": "running", "task_id": task_id}
|
||
|
||
output_file = Path(task.get("output_file", ""))
|
||
if output_file.exists():
|
||
result = self._check_task_status(task)
|
||
ma_debug("refresh_task_runtime_state_output_result", task_id=task_id, result=result)
|
||
return result
|
||
|
||
# 多智能体模式：没有输出文件表示 idle，不强制清理
|
||
if task.get("multi_agent_mode"):
|
||
task["status"] = "idle"
|
||
task["updated_at"] = time.time()
|
||
ma_debug("refresh_task_runtime_state_idle_no_output_file", task_id=task_id)
|
||
return {"status": "idle", "task_id": task_id}
|
||
|
||
if self._should_force_cleanup_stale_task(task):
|
||
return self._mark_task_terminated(
|
||
task,
|
||
message="子智能体疑似僵尸任务，已超时自动清理运行状态。",
|
||
system_message="⚠️ 子智能体长时间未结束，系统已自动清理运行状态。",
|
||
notified=True,
|
||
)
|
||
|
||
return self._mark_task_terminated(
|
||
task,
|
||
message="检测到子智能体任务已退出，已自动清理运行状态。",
|
||
system_message="⚠️ 子智能体任务异常退出，系统已自动清理运行状态。",
|
||
notified=True,
|
||
)
|
||
|
||
def reconcile_task_states(self, conversation_id: Optional[str] = None) -> int:
|
||
"""修正运行态任务状态，返回修正条目数。"""
|
||
changed = 0
|
||
ma_debug("reconcile_task_states_start", conversation_id=conversation_id, task_count=len(self.tasks))
|
||
for task in self.tasks.values():
|
||
if not isinstance(task, dict):
|
||
continue
|
||
if conversation_id and task.get("conversation_id") != conversation_id:
|
||
continue
|
||
before_status = task.get("status")
|
||
before_notified = task.get("notified")
|
||
self._refresh_task_runtime_state(task)
|
||
after_status = task.get("status")
|
||
if after_status != before_status or task.get("notified") != before_notified:
|
||
ma_debug(
|
||
"reconcile_task_states_changed",
|
||
task_id=task.get("task_id"),
|
||
agent_id=task.get("agent_id"),
|
||
before_status=before_status,
|
||
after_status=after_status,
|
||
)
|
||
changed += 1
|
||
if changed:
|
||
self._save_state()
|
||
ma_debug("reconcile_task_states_end", conversation_id=conversation_id, changed=changed)
|
||
return changed
|
||
|
||
def _should_force_cleanup_stale_task(self, task: Dict[str, Any]) -> bool:
|
||
try:
|
||
created_at = float(task.get("created_at") or 0)
|
||
except (TypeError, ValueError):
|
||
created_at = 0
|
||
if created_at <= 0:
|
||
return False
|
||
# timeout_seconds 为 None 表示永久子智能体，不强制清理
|
||
raw_timeout = task.get("timeout_seconds")
|
||
if raw_timeout is None:
|
||
return False
|
||
from config import SUB_AGENT_DEFAULT_TIMEOUT
|
||
timeout_seconds = int(raw_timeout or SUB_AGENT_DEFAULT_TIMEOUT or 0)
|
||
timeout_seconds = max(timeout_seconds, 1)
|
||
grace_seconds = 120
|
||
elapsed = time.time() - created_at
|
||
if elapsed <= (timeout_seconds + grace_seconds):
|
||
return False
|
||
output_file = Path(task.get("output_file", ""))
|
||
if output_file.exists():
|
||
return False
|
||
progress_file = Path(task.get("progress_file", ""))
|
||
if progress_file.exists():
|
||
try:
|
||
stale_span = time.time() - progress_file.stat().st_mtime
|
||
if stale_span <= grace_seconds:
|
||
return False
|
||
except Exception:
|
||
pass
|
||
return True
|