from __future__ import annotations from dataclasses import dataclass from datetime import datetime from typing import Optional from uuid import uuid4 from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session from app.db.base import DbOper from app.db.models.agenttask import ( AgentTask, _get_for_user_statement, _list_for_user_statement, ) from app.db.models.agenttaskrun import AgentTaskRun @dataclass(frozen=True, slots=True) class AgentTaskFinishRecord: """一次运行收口后由数据库事务确认的三个独立事实。""" run_finalized: bool task_projection_updated: bool date_task_disabled: bool class AgentTaskOper(DbOper): """ Agent 自主定时任务管理。 """ @staticmethod def _now() -> str: """生成当前数据库时间字符串。""" return datetime.now().strftime("%Y-%m-%d %H:%M:%S") def add(self, **kwargs: object) -> Optional[AgentTask]: """ 新增 Agent 定时任务。 """ now = self._now() task_id = self._execute_sync_write( lambda session: AgentTask.add_task( session, **kwargs, enabled=True, last_status="waiting", run_count=0, created_at=now, updated_at=now, ) ) return self.get(task_id) def get( self, task_id: int, user_id: Optional[str] = None, ) -> Optional[AgentTask]: """ 查询单个 Agent 定时任务。 """ def query(session: Session) -> Optional[AgentTask]: """在调用方会话中读取单个任务。""" return session.execute( _get_for_user_statement( AgentTask, task_id=task_id, user_id=user_id, ) ).scalars().first() return self._execute_sync_query(query) async def async_get( self, task_id: int, user_id: Optional[str] = None, ) -> Optional[AgentTask]: """通过异步会话查询单个 Agent 定时任务。""" async def query(session: AsyncSession) -> Optional[AgentTask]: """在调用方异步会话中执行与同步入口相同的查询语义。""" result = await session.execute( _get_for_user_statement( AgentTask, task_id=task_id, user_id=user_id, ) ) return result.scalars().first() return await self._execute_async_query(query) def list( self, user_id: Optional[str] = None, enabled: Optional[bool] = None, ) -> list[AgentTask]: """ 查询 Agent 定时任务列表。 """ def query(session: Session) -> list[AgentTask]: """在调用方会话中读取任务列表。""" return list(session.execute( _list_for_user_statement( AgentTask, user_id=user_id, enabled=enabled, ) ).scalars().all()) return self._execute_sync_query(query) def update( self, task_id: int, payload: dict, user_id: Optional[str] = None, ) -> bool: """ 更新 Agent 定时任务。 """ normalized_payload = { key: value for key, value in payload.items() if key in { "name", "content", "trigger_type", "cron_expression", "run_at", "enabled", "last_status", "last_result", } } if not normalized_payload: return False normalized_payload["updated_at"] = self._now() return self._execute_sync_write( lambda session: AgentTask.update_task( session, task_id=task_id, payload=normalized_payload, user_id=user_id, ) ) def delete(self, task_id: int, user_id: Optional[str] = None) -> bool: """ 删除非运行中的 Agent 定时任务及其运行历史。 """ return self._execute_sync_write( lambda session: AgentTaskRun.delete_task_and_runs( session, task_id=task_id, user_id=user_id, ) ) def begin_run( self, task_id: int, trigger_source: str = "scheduled", *, run_id: Optional[str] = None, started_at: Optional[str] = None, ) -> Optional[AgentTaskRun]: """ 原子创建一次运行并返回其任务快照。 可选运行 ID 和开始时间用于恢复/幂等验证;正常调度入口由本方法生成。 """ resolved_run_id = run_id or uuid4().hex resolved_started_at = started_at or self._now() created_run_id = self._execute_sync_write( lambda session: AgentTaskRun.begin_run( session, task_id=task_id, run_id=resolved_run_id, trigger_source=trigger_source, started_at=resolved_started_at, ) ) return self.get_run(created_run_id) if created_run_id else None def mark_running(self, task_id: int) -> bool: """兼容既有调用并为该次执行创建运行记录。""" return self.begin_run(task_id=task_id) is not None def mark_interrupted(self, task_id: int, result: str) -> bool: """ 将遗留的运行中任务标记为中断且结果未知。 """ finished_at = self._now() normalized_result = (result or "")[:20000] return self._execute_sync_write( lambda session: AgentTaskRun.interrupt_task( session, task_id=task_id, result=normalized_result, finished_at=finished_at, ) ) def get_run(self, run_id: str) -> Optional[AgentTaskRun]: """查询一次 Agent 任务运行。""" return self._execute_sync_query( lambda session: AgentTaskRun.get_by_run_id(session, run_id=run_id) ) def list_runs( self, task_id: int, user_id: Optional[str] = None, limit: int = 10, ) -> list[AgentTaskRun]: """查询任务最近的有界运行历史。""" return self._execute_sync_query( lambda session: AgentTaskRun.list_for_task( session, task_id=task_id, user_id=user_id, limit=limit, ) ) def finish_run( self, run_id: str, success: bool, result: str, disable_date_task: bool = False, ) -> bool: """收口精确运行并更新仍匹配的任务投影。""" finished_at = self._now() normalized_result = (result or "")[:20000] return self._execute_sync_write( lambda session: AgentTaskRun.finish_run( session, run_id=run_id, success=success, result=normalized_result, finished_at=finished_at, disable_date_task=disable_date_task, ) ) def finish_run_outcome( self, run_id: str, success: bool, result: str, ) -> AgentTaskFinishRecord: """收口运行,并返回当前事务实际更新的任务投影和停用事实。""" finished_at = self._now() normalized_result = (result or "")[:20000] expected_status = "success" if success else "failed" def finalize(session: Session) -> AgentTaskFinishRecord: """使用列查询绕过 ORM identity-map,读取刚写入的真实投影。""" finalized = AgentTaskRun.finish_run( session, run_id=run_id, success=success, result=normalized_result, finished_at=finished_at, disable_date_task=True, ) run = session.execute( select( AgentTaskRun.task_id, AgentTaskRun.trigger_type, ).where(AgentTaskRun.run_id == run_id) ).mappings().first() task = None if run: task = session.execute( select( AgentTask.last_run_id, AgentTask.last_status, AgentTask.enabled, ).where(AgentTask.id == run["task_id"]) ).mappings().first() projection_updated = bool( finalized and task and task["last_run_id"] == run_id and task["last_status"] == expected_status ) return AgentTaskFinishRecord( run_finalized=finalized, task_projection_updated=projection_updated, date_task_disabled=bool( projection_updated and run and run["trigger_type"] == "date" and task and not task["enabled"] ), ) return self._execute_sync_write(finalize) def finish( self, task_id: int, success: bool, result: str, disable: bool = False, ) -> bool: """ 记录 Agent 定时任务执行结果。 """ task = self.get(task_id) if not task or not task.last_run_id: return False return self.finish_run( run_id=task.last_run_id, success=success, result=result, disable_date_task=disable, ) @staticmethod def to_dict( task: AgentTask, next_run_at: Optional[str] = None, timezone: Optional[str] = None, ) -> dict: """ 将 Agent 定时任务转换为工具可返回的结构。 """ return { "id": task.id, "name": task.name, "content": task.content, "trigger_type": task.trigger_type, "cron_expression": task.cron_expression, "run_at": task.run_at, "timezone": timezone, "enabled": bool(task.enabled), "last_status": task.last_status, "last_run_at": task.last_run_at, "last_result": task.last_result, "last_run_id": task.last_run_id, "run_count": task.run_count or 0, "next_run_at": next_run_at, "created_at": task.created_at, "updated_at": task.updated_at, } @staticmethod def run_to_dict(run: AgentTaskRun) -> dict: """将一次 Agent 任务运行转换为工具返回结构。""" return { "run_id": run.run_id, "task_id": run.task_id, "trigger_source": run.trigger_source, "name": run.name, "content": run.content, "trigger_type": run.trigger_type, "cron_expression": run.cron_expression, "run_at": run.run_at, "status": run.status, "started_at": run.started_at, "finished_at": run.finished_at, "result": run.result, }