1"""

2 AgentOS v0.40 Task Queue — 异步任务调度与重试。

3 支持：内存队列（开发）/ Redis队列（生产）、优先级、重试、死信队列。

4"""

5

6 from __future__ import annotations

7

8 import asyncio

9 import heapq

10 import time

11 import uuid

12 from dataclasses import dataclass, field

13 from enum import Enum

14 from typing import Optional, Callable, Any

15

16

17 class TaskState(str, Enum):

18

19 """任务状态枚举。"""

20

21 PENDING = "pending"

22 RUNNING = "running"

23 SUCCESS = "success"

24 FAILED = "failed"

25 RETRYING = "retrying"

26 CANCELLED = "cancelled"

27 DEAD = "dead" # 死信

28

29

30 class TaskPriority(int, Enum):

31

32 """任务优先级枚举。"""

33

34 LOW = 0

35 NORMAL = 50

36 HIGH = 100

37 CRITICAL = 200

38

39

40@dataclass(order=True)

41 class QueuedTask:

42 """带优先级的任务节点（priority取负以实现最大堆）。"""

43 priority: int # -priority for max-heap

44 created_at: float = field(compare=False)

45 task: "ScheduledTask" = field(compare=False)

46

47

48@dataclass

49 class ScheduledTask:

50 """调度任务。"""

51 id: str = field(default_factory=lambda: uuid.uuid4().hex[:12])

52 name: str = ""

53 payload: dict = field(default_factory=dict)

54 priority: TaskPriority = TaskPriority.NORMAL

55 state: TaskState = TaskState.PENDING

56 max_retries: int = 3

57 retry_delay: float = 1.0 # 秒

58 timeout: float = 60.0

59 callback: Callable | None = field(default=None, repr=False)

60 result: Any = None

61 error: str = ""

62 retry_count: int = 0

63 created_at: float = field(default_factory=time.time)

64 started_at: float = 0

65 completed_at: float = 0

66 _tags: dict = field(default_factory=dict)

67

68

69 class MemoryQueue:

70 """基于堆内存的任务队列 — 开发环境默认。"""

71

72 def __init__(self, max_size: int = 10000):

73 self._heap: list[QueuedTask] = []

74 self._pending: dict[str, ScheduledTask] = {}

75 self._dead: list[ScheduledTask] = []

76 self.max_size = max_size

77 self._lock = asyncio.Lock()

78

79 async def enqueue(self, task: ScheduledTask) -> str:

80 async with self._lock:

81 if len(self._pending) >= self.max_size:

82 raise RuntimeError(f"Queue full ({self.max_size})")

83 self._pending[task.id] = task

84 heapq.heappush(self._heap, QueuedTask(priority=-task.priority.value, created_at=task.created_at, task=task))

85 return task.id

86

87 async def dequeue(self) -> ScheduledTask | None:

88 async with self._lock:

89 while self._heap:

90 qt = heapq.heappop(self._heap)

91 task = self._pending.pop(qt.task.id, None)

92 if task and task.state == TaskState.PENDING:

93 return task

94 return None

95

96 async def peek(self) -> ScheduledTask | None:

97 async with self._lock:

98 if self._heap:

99 return self._heap[0].task

100 return None

101

102 def pending_count(self) -> int:

103 return len(self._heap)

104

105 def dead_count(self) -> int:

106 return len(self._dead)

107

108 async def move_to_dead(self, task: ScheduledTask):

109 async with self._lock:

110 task.state = TaskState.DEAD

111 self._dead.append(task)

112 self._pending.pop(task.id, None)

113

114 def stats(self) -> dict:

115 return {"pending": len(self._heap), "dead": len(self._dead), "max_size": self.max_size}

116

117

118 class TaskQueue:

119 """任务队列管理器。"""

120

121 def __init__(self, queue: MemoryQueue | None = None, concurrency: int = 4):

122 self._queue = queue or MemoryQueue()

123 self._concurrency = concurrency

124 self._running: set[str] = set()

125 self._callbacks: dict[str, Callable] = {}

126 self._semaphore = asyncio.Semaphore(concurrency)

127 self._running_flag = False

128

129 def register_callback(self, task_name: str, handler: Callable):

130 """注册任务处理器。"""

131 self._callbacks[task_name] = handler

132

133 async def submit(self, task: ScheduledTask) -> str:

134 if task.name not in self._callbacks:

135 raise ValueError(f"No handler registered for task: {task.name}")

136 task_id = await self._queue.enqueue(task)

137 return task_id

138

139 async def start(self):

140 """启动Worker循环。"""

141 self._running_flag = True

142 while self._running_flag:

143 task = await self._queue.dequeue()

144 if not task:

145 await asyncio.sleep(0.1)

146 continue

147 asyncio.create_task(self._execute(task))

148

149 def stop(self):

150 self._running_flag = False

151

152 async def _execute(self, task: ScheduledTask):

153 async with self._semaphore:

154 self._running.add(task.id)

155 task.state = TaskState.RUNNING

156 task.started_at = time.time()

157

158 handler = self._callbacks.get(task.name)

159 if not handler:

160 task.state = TaskState.FAILED

161 task.error = f"No handler: {task.name}"

162 self._running.discard(task.id)

163 return

164

165 try:

166 result = handler(task.payload)

167 if asyncio.iscoroutine(result):

168 result = await asyncio.wait_for(result, timeout=task.timeout)

169 task.result = result

170 task.state = TaskState.SUCCESS

171 except asyncio.TimeoutError:

172 task.error = f"Timeout after {task.timeout}s"

173 await self._handle_failure(task)

174 except Exception as e:

175 task.error = str(e)

176 await self._handle_failure(task)

177 finally:

178 task.completed_at = time.time()

179 self._running.discard(task.id)

180

181 async def _handle_failure(self, task: ScheduledTask):

182 if task.retry_count < task.max_retries:

183 task.retry_count += 1

184 task.state = TaskState.RETRYING

185 await asyncio.sleep(task.retry_delay * (2 ** (task.retry_count - 1))) # 指数退避

186 task.state = TaskState.PENDING

187 task.priority = TaskPriority(task.priority.value + 10) # 提升优先级

188 await self._queue.enqueue(task)

189 else:

190 task.state = TaskState.FAILED

191 await self._queue.move_to_dead(task)

192

193 def cancel(self, task_id: str):

194 """取消任务。"""

195 task = self._queue._pending.get(task_id)

196 if task and task.state in (TaskState.PENDING, TaskState.RETRYING):

197 task.state = TaskState.CANCELLED

198

199 def stats(self) -> dict:

200 return {

201 "running": len(self._running),

202 "concurrency": self._concurrency,

203 "queue": self._queue.stats(),

204 "handlers": list(self._callbacks.keys()),

205 }