落地并存测服全端：智能体 Graph/路由、后台切 8180、主账与部署脚本齐备，便于线上重发与验收。
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
linmianxun
@@ -0,0 +1,267 @@
|
||||
"""
|
||||
PG agent_job 认领壳（与 Redis Stream 二选一真相时用）。
|
||||
|
||||
本文件职责：enqueue / claim / complete / requeue_expired；不含业务处理体。
|
||||
禁止：一把大锁串行所有工单；禁止 Redis wecom:worker:singleton。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Optional
|
||||
|
||||
from agent.runtime_db import assert_safe_database, connect_psycopg
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _utcnow() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobRecord:
|
||||
"""一条任务行。"""
|
||||
|
||||
id: str
|
||||
job_type: str
|
||||
payload: dict[str, Any] = field(default_factory=dict)
|
||||
state: str = "pending"
|
||||
idempotency_key: str = ""
|
||||
attempts: int = 0
|
||||
last_error: str = ""
|
||||
|
||||
|
||||
class MemoryJobStore:
|
||||
"""内存任务认领壳。"""
|
||||
|
||||
def __init__(self, *, claim_timeout_sec: float = 180.0) -> None:
|
||||
self._rows: dict[str, JobRecord] = {}
|
||||
self._claimed_at: dict[str, float] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._claim_timeout = claim_timeout_sec
|
||||
|
||||
def enqueue(
|
||||
self,
|
||||
*,
|
||||
job_type: str,
|
||||
payload: dict[str, Any],
|
||||
idempotency_key: str = "",
|
||||
) -> tuple[bool, str]:
|
||||
with self._lock:
|
||||
if idempotency_key:
|
||||
for j in self._rows.values():
|
||||
if j.job_type == job_type and j.idempotency_key == idempotency_key:
|
||||
return False, "duplicate_idempotency_key"
|
||||
job_id = str(uuid.uuid4())
|
||||
self._rows[job_id] = JobRecord(
|
||||
id=job_id,
|
||||
job_type=job_type,
|
||||
payload=dict(payload),
|
||||
idempotency_key=idempotency_key,
|
||||
)
|
||||
return True, job_id
|
||||
|
||||
def claim(self, *, worker_id: str, job_type: str = "") -> Optional[JobRecord]:
|
||||
import time
|
||||
|
||||
now = time.time()
|
||||
with self._lock:
|
||||
for jid, job in self._rows.items():
|
||||
if job_type and job.job_type != job_type:
|
||||
continue
|
||||
if job.state == "pending" or (
|
||||
job.state == "claimed"
|
||||
and now - self._claimed_at.get(jid, 0) > self._claim_timeout
|
||||
):
|
||||
job.state = "claimed"
|
||||
job.attempts += 1
|
||||
self._claimed_at[jid] = now
|
||||
return JobRecord(
|
||||
id=job.id,
|
||||
job_type=job.job_type,
|
||||
payload=dict(job.payload),
|
||||
state=job.state,
|
||||
idempotency_key=job.idempotency_key,
|
||||
attempts=job.attempts,
|
||||
last_error=job.last_error,
|
||||
)
|
||||
return None
|
||||
|
||||
def complete(self, job_id: str, *, error: str = "", result: Optional[dict] = None) -> None:
|
||||
with self._lock:
|
||||
job = self._rows.get(job_id)
|
||||
if not job:
|
||||
return
|
||||
job.state = "failed" if error else "done"
|
||||
job.last_error = error
|
||||
logger.info("job.memory.complete id=%s state=%s", job_id, job.state)
|
||||
|
||||
|
||||
class PostgresJobStore:
|
||||
"""PG agent_job 认领壳。"""
|
||||
|
||||
def __init__(self, dsn: str, *, claim_timeout_sec: float = 180.0) -> None:
|
||||
assert_safe_database(dsn)
|
||||
self._dsn = dsn
|
||||
self._claim_timeout = int(claim_timeout_sec)
|
||||
|
||||
def enqueue(
|
||||
self,
|
||||
*,
|
||||
job_type: str,
|
||||
payload: dict[str, Any],
|
||||
idempotency_key: str = "",
|
||||
) -> tuple[bool, str]:
|
||||
job_id = str(uuid.uuid4())
|
||||
try:
|
||||
with connect_psycopg(self._dsn) as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO agent_job (id, job_type, idempotency_key, payload, state)
|
||||
VALUES (%s, %s, %s, %s::jsonb, 'pending')
|
||||
""",
|
||||
(
|
||||
job_id,
|
||||
job_type,
|
||||
idempotency_key,
|
||||
json.dumps(payload, ensure_ascii=False),
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
return True, job_id
|
||||
except Exception as exc: # noqa: BLE001
|
||||
if "unique" in str(exc).lower() or "duplicate" in str(exc).lower():
|
||||
return False, "duplicate_idempotency_key"
|
||||
raise
|
||||
|
||||
def claim(self, *, worker_id: str, job_type: str = "") -> Optional[JobRecord]:
|
||||
type_filter = "AND job_type = %s" if job_type else ""
|
||||
params: list[Any] = [self._claim_timeout, worker_id]
|
||||
if job_type:
|
||||
params.append(job_type)
|
||||
sql = f"""
|
||||
WITH cte AS (
|
||||
SELECT id FROM agent_job
|
||||
WHERE (
|
||||
state = 'pending'
|
||||
OR (state = 'claimed' AND claimed_at < NOW() - (%s || ' seconds')::interval)
|
||||
)
|
||||
{type_filter}
|
||||
ORDER BY created_at
|
||||
FOR UPDATE SKIP LOCKED
|
||||
LIMIT 1
|
||||
)
|
||||
UPDATE agent_job j SET
|
||||
state = 'claimed',
|
||||
claimed_at = NOW(),
|
||||
claimed_by = %s,
|
||||
attempts = j.attempts + 1,
|
||||
updated_at = NOW()
|
||||
FROM cte WHERE j.id = cte.id
|
||||
RETURNING j.id, j.job_type, j.payload, j.state, j.idempotency_key, j.attempts, j.last_error
|
||||
"""
|
||||
# 参数顺序：timeout, [job_type], worker_id — 上面拼法易错，改用固定两段
|
||||
with connect_psycopg(self._dsn) as conn:
|
||||
if job_type:
|
||||
row = conn.execute(
|
||||
"""
|
||||
WITH cte AS (
|
||||
SELECT id FROM agent_job
|
||||
WHERE job_type = %s AND (
|
||||
state = 'pending'
|
||||
OR (state = 'claimed'
|
||||
AND claimed_at < NOW() - (%s || ' seconds')::interval)
|
||||
)
|
||||
ORDER BY created_at
|
||||
FOR UPDATE SKIP LOCKED
|
||||
LIMIT 1
|
||||
)
|
||||
UPDATE agent_job j SET
|
||||
state = 'claimed',
|
||||
claimed_at = NOW(),
|
||||
claimed_by = %s,
|
||||
attempts = j.attempts + 1,
|
||||
updated_at = NOW()
|
||||
FROM cte WHERE j.id = cte.id
|
||||
RETURNING j.id, j.job_type, j.payload, j.state,
|
||||
j.idempotency_key, j.attempts, j.last_error
|
||||
""",
|
||||
(job_type, str(self._claim_timeout), worker_id),
|
||||
).fetchone()
|
||||
else:
|
||||
row = conn.execute(
|
||||
"""
|
||||
WITH cte AS (
|
||||
SELECT id FROM agent_job
|
||||
WHERE state = 'pending'
|
||||
OR (state = 'claimed'
|
||||
AND claimed_at < NOW() - (%s || ' seconds')::interval)
|
||||
ORDER BY created_at
|
||||
FOR UPDATE SKIP LOCKED
|
||||
LIMIT 1
|
||||
)
|
||||
UPDATE agent_job j SET
|
||||
state = 'claimed',
|
||||
claimed_at = NOW(),
|
||||
claimed_by = %s,
|
||||
attempts = j.attempts + 1,
|
||||
updated_at = NOW()
|
||||
FROM cte WHERE j.id = cte.id
|
||||
RETURNING j.id, j.job_type, j.payload, j.state,
|
||||
j.idempotency_key, j.attempts, j.last_error
|
||||
""",
|
||||
(str(self._claim_timeout), worker_id),
|
||||
).fetchone()
|
||||
conn.commit()
|
||||
if not row:
|
||||
return None
|
||||
return JobRecord(
|
||||
id=str(row[0]),
|
||||
job_type=str(row[1]),
|
||||
payload=dict(row[2] or {}),
|
||||
state=str(row[3]),
|
||||
idempotency_key=str(row[4] or ""),
|
||||
attempts=int(row[5] or 0),
|
||||
last_error=str(row[6] or ""),
|
||||
)
|
||||
|
||||
def complete(
|
||||
self,
|
||||
job_id: str,
|
||||
*,
|
||||
error: str = "",
|
||||
result: Optional[dict[str, Any]] = None,
|
||||
) -> None:
|
||||
state = "failed" if error else "done"
|
||||
with connect_psycopg(self._dsn) as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE agent_job SET
|
||||
state = %s,
|
||||
last_error = %s,
|
||||
result = %s::jsonb,
|
||||
updated_at = NOW()
|
||||
WHERE id = %s
|
||||
""",
|
||||
(
|
||||
state,
|
||||
error,
|
||||
json.dumps(result or {}, ensure_ascii=False),
|
||||
job_id,
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
logger.info("job.pg.complete id=%s state=%s", job_id, state)
|
||||
|
||||
|
||||
def create_job_store(*, dsn: str = "", backend: str = "auto", claim_timeout_sec: float = 180.0):
|
||||
"""工厂：Memory 或 PG。"""
|
||||
if backend == "memory" or not dsn:
|
||||
return MemoryJobStore(claim_timeout_sec=claim_timeout_sec)
|
||||
return PostgresJobStore(dsn, claim_timeout_sec=claim_timeout_sec)
|
||||
Reference in New Issue
Block a user