|
1
|
""" |
|
2
|
TaskQueue — work assignment for agent swarms via Redis. |
|
3
|
|
|
4
|
Agents enqueue tasks (e.g. "ingest file X") and other agents dequeue and |
|
5
|
claim them. Tasks move through PENDING -> IN_PROGRESS -> DONE | FAILED. |
|
6
|
|
|
7
|
Usage: |
|
8
|
queue = TaskQueue("redis://localhost:6379") |
|
9
|
|
|
10
|
task_id = queue.enqueue("ingest_file", {"path": "src/main.py"}) |
|
11
|
|
|
12
|
task = queue.dequeue("agent-1") # atomically claim next pending task |
|
13
|
if task: |
|
14
|
try: |
|
15
|
result = do_work(task) |
|
16
|
queue.complete(task.id, result) |
|
17
|
except Exception as e: |
|
18
|
queue.fail(task.id, str(e)) |
|
19
|
|
|
20
|
info = queue.status(task_id) |
|
21
|
n = queue.pending_count() |
|
22
|
""" |
|
23
|
|
|
24
|
from __future__ import annotations |
|
25
|
|
|
26
|
import json |
|
27
|
import logging |
|
28
|
import time |
|
29
|
import uuid |
|
30
|
from dataclasses import dataclass, field |
|
31
|
from enum import Enum |
|
32
|
from typing import Any |
|
33
|
|
|
34
|
logger = logging.getLogger(__name__) |
|
35
|
|
|
36
|
_QUEUE_KEY = "navegador:taskqueue:pending" # Redis list (RPUSH/BLPOP) |
|
37
|
_TASK_KEY_PREFIX = "navegador:task:" # Hash per task |
|
38
|
_INPROGRESS_KEY = "navegador:taskqueue:inprogress" # Set of in-progress task IDs |
|
39
|
|
|
40
|
|
|
41
|
class TaskStatus(str, Enum): |
|
42
|
PENDING = "pending" |
|
43
|
IN_PROGRESS = "in_progress" |
|
44
|
DONE = "done" |
|
45
|
FAILED = "failed" |
|
46
|
|
|
47
|
|
|
48
|
@dataclass |
|
49
|
class Task: |
|
50
|
"""A unit of work in the task queue.""" |
|
51
|
|
|
52
|
id: str |
|
53
|
type: str |
|
54
|
payload: dict[str, Any] |
|
55
|
status: TaskStatus = TaskStatus.PENDING |
|
56
|
agent_id: str | None = None |
|
57
|
result: Any = None |
|
58
|
error: str | None = None |
|
59
|
created_at: float = field(default_factory=time.time) |
|
60
|
updated_at: float = field(default_factory=time.time) |
|
61
|
|
|
62
|
def to_dict(self) -> dict[str, str]: |
|
63
|
"""Serialise to a flat string dict suitable for Redis HSET.""" |
|
64
|
return { |
|
65
|
"id": self.id, |
|
66
|
"type": self.type, |
|
67
|
"payload": json.dumps(self.payload), |
|
68
|
"status": self.status.value, |
|
69
|
"agent_id": self.agent_id or "", |
|
70
|
"result": json.dumps(self.result) if self.result is not None else "", |
|
71
|
"error": self.error or "", |
|
72
|
"created_at": str(self.created_at), |
|
73
|
"updated_at": str(self.updated_at), |
|
74
|
} |
|
75
|
|
|
76
|
@classmethod |
|
77
|
def from_dict(cls, d: dict[str, Any]) -> "Task": |
|
78
|
# Redis hgetall returns bytes; decode if necessary. |
|
79
|
decoded: dict[str, Any] = {} |
|
80
|
for k, v in d.items(): |
|
81
|
key = k.decode() if isinstance(k, bytes) else k |
|
82
|
val = v.decode() if isinstance(v, bytes) else v |
|
83
|
decoded[key] = val |
|
84
|
|
|
85
|
payload = json.loads(decoded.get("payload", "{}") or "{}") |
|
86
|
result_raw = decoded.get("result", "") |
|
87
|
result = json.loads(result_raw) if result_raw else None |
|
88
|
status_raw = decoded.get("status", TaskStatus.PENDING.value) |
|
89
|
status = TaskStatus(status_raw) |
|
90
|
|
|
91
|
return cls( |
|
92
|
id=decoded["id"], |
|
93
|
type=decoded["type"], |
|
94
|
payload=payload, |
|
95
|
status=status, |
|
96
|
agent_id=decoded.get("agent_id") or None, |
|
97
|
result=result, |
|
98
|
error=decoded.get("error") or None, |
|
99
|
created_at=float(decoded.get("created_at", 0)), |
|
100
|
updated_at=float(decoded.get("updated_at", 0)), |
|
101
|
) |
|
102
|
|
|
103
|
|
|
104
|
def _task_key(task_id: str) -> str: |
|
105
|
return f"{_TASK_KEY_PREFIX}{task_id}" |
|
106
|
|
|
107
|
|
|
108
|
class TaskQueue: |
|
109
|
""" |
|
110
|
Redis-backed task queue for coordinating work across agent swarms. |
|
111
|
|
|
112
|
Parameters |
|
113
|
---------- |
|
114
|
redis_url: |
|
115
|
URL of the Redis server. |
|
116
|
redis_client: |
|
117
|
Optional pre-built Redis client (for testing / DI). |
|
118
|
""" |
|
119
|
|
|
120
|
def __init__(self, redis_url: str, *, redis_client: Any = None) -> None: |
|
121
|
self.redis_url = redis_url |
|
122
|
self._redis = redis_client or self._connect_redis(redis_url) |
|
123
|
|
|
124
|
# ── Internal ────────────────────────────────────────────────────────────── |
|
125
|
|
|
126
|
@staticmethod |
|
127
|
def _connect_redis(url: str) -> Any: |
|
128
|
try: |
|
129
|
import redis # type: ignore[import] |
|
130
|
except ImportError as exc: |
|
131
|
raise ImportError("Install redis: pip install redis") from exc |
|
132
|
return redis.from_url(url) |
|
133
|
|
|
134
|
# ── Public API ──────────────────────────────────────────────────────────── |
|
135
|
|
|
136
|
def enqueue(self, task_type: str, payload: dict[str, Any]) -> str: |
|
137
|
""" |
|
138
|
Add a new task to the queue. |
|
139
|
|
|
140
|
Returns |
|
141
|
------- |
|
142
|
str |
|
143
|
The newly created task ID. |
|
144
|
""" |
|
145
|
task_id = str(uuid.uuid4()) |
|
146
|
task = Task(id=task_id, type=task_type, payload=payload) |
|
147
|
pipe = self._redis.pipeline() |
|
148
|
pipe.hset(_task_key(task_id), mapping=task.to_dict()) |
|
149
|
pipe.rpush(_QUEUE_KEY, task_id) |
|
150
|
pipe.execute() |
|
151
|
logger.debug("Enqueued task %s (type=%s)", task_id, task_type) |
|
152
|
return task_id |
|
153
|
|
|
154
|
def dequeue(self, agent_id: str) -> Task | None: |
|
155
|
""" |
|
156
|
Atomically claim the next pending task for *agent_id*. |
|
157
|
|
|
158
|
Returns ``None`` when the queue is empty. |
|
159
|
""" |
|
160
|
task_id_raw = self._redis.lpop(_QUEUE_KEY) |
|
161
|
if task_id_raw is None: |
|
162
|
return None |
|
163
|
|
|
164
|
task_id = task_id_raw.decode() if isinstance(task_id_raw, bytes) else task_id_raw |
|
165
|
now = time.time() |
|
166
|
pipe = self._redis.pipeline() |
|
167
|
pipe.hset( |
|
168
|
_task_key(task_id), |
|
169
|
mapping={ |
|
170
|
"status": TaskStatus.IN_PROGRESS.value, |
|
171
|
"agent_id": agent_id, |
|
172
|
"updated_at": now, |
|
173
|
}, |
|
174
|
) |
|
175
|
pipe.sadd(_INPROGRESS_KEY, task_id) |
|
176
|
pipe.execute() |
|
177
|
|
|
178
|
raw = self._redis.hgetall(_task_key(task_id)) |
|
179
|
task = Task.from_dict(raw) |
|
180
|
logger.debug("Agent %s claimed task %s", agent_id, task_id) |
|
181
|
return task |
|
182
|
|
|
183
|
def complete(self, task_id: str, result: Any = None) -> None: |
|
184
|
"""Mark a task as successfully completed.""" |
|
185
|
result_encoded = json.dumps(result) if result is not None else "" |
|
186
|
pipe = self._redis.pipeline() |
|
187
|
pipe.hset( |
|
188
|
_task_key(task_id), |
|
189
|
mapping={ |
|
190
|
"status": TaskStatus.DONE.value, |
|
191
|
"result": result_encoded, |
|
192
|
"updated_at": time.time(), |
|
193
|
}, |
|
194
|
) |
|
195
|
pipe.srem(_INPROGRESS_KEY, task_id) |
|
196
|
pipe.execute() |
|
197
|
logger.debug("Task %s completed", task_id) |
|
198
|
|
|
199
|
def fail(self, task_id: str, error: str) -> None: |
|
200
|
"""Mark a task as failed with an error message.""" |
|
201
|
pipe = self._redis.pipeline() |
|
202
|
pipe.hset( |
|
203
|
_task_key(task_id), |
|
204
|
mapping={ |
|
205
|
"status": TaskStatus.FAILED.value, |
|
206
|
"error": error, |
|
207
|
"updated_at": time.time(), |
|
208
|
}, |
|
209
|
) |
|
210
|
pipe.srem(_INPROGRESS_KEY, task_id) |
|
211
|
pipe.execute() |
|
212
|
logger.debug("Task %s failed: %s", task_id, error) |
|
213
|
|
|
214
|
def status(self, task_id: str) -> dict[str, Any]: |
|
215
|
""" |
|
216
|
Return a status dict for a task. |
|
217
|
|
|
218
|
Raises ``KeyError`` if the task does not exist. |
|
219
|
""" |
|
220
|
raw = self._redis.hgetall(_task_key(task_id)) |
|
221
|
if not raw: |
|
222
|
raise KeyError(f"Task not found: {task_id}") |
|
223
|
task = Task.from_dict(raw) |
|
224
|
return { |
|
225
|
"id": task.id, |
|
226
|
"type": task.type, |
|
227
|
"status": task.status.value, |
|
228
|
"agent_id": task.agent_id, |
|
229
|
"result": task.result, |
|
230
|
"error": task.error, |
|
231
|
"created_at": task.created_at, |
|
232
|
"updated_at": task.updated_at, |
|
233
|
} |
|
234
|
|
|
235
|
def pending_count(self) -> int: |
|
236
|
"""Return the number of tasks currently waiting in the queue.""" |
|
237
|
return self._redis.llen(_QUEUE_KEY) |
|
238
|
|