163 lines
5.1 KiB
Python
163 lines
5.1 KiB
Python
"""
|
||
停等账仓储壳：suspend_wait（waitVersion）。
|
||
|
||
本文件职责：登记/关闭停点；与 Graph interrupt、H5 token 对齐的版本号。
|
||
禁止：用本表冒充主账六态；禁止用户文字伪造 Worker 完成。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import threading
|
||
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 SuspendRecord:
|
||
"""一条停等记录。"""
|
||
|
||
thread_id: str
|
||
wait_version: int
|
||
pause_name: str = ""
|
||
inquiry_no: str = ""
|
||
status: str = "open"
|
||
payload: dict[str, Any] = field(default_factory=dict)
|
||
|
||
|
||
class MemorySuspendStore:
|
||
"""内存停等壳。"""
|
||
|
||
def __init__(self) -> None:
|
||
self._rows: dict[tuple[str, int], SuspendRecord] = {}
|
||
self._lock = threading.Lock()
|
||
|
||
def open_wait(self, record: SuspendRecord) -> None:
|
||
key = (record.thread_id, record.wait_version)
|
||
with self._lock:
|
||
self._rows[key] = SuspendRecord(
|
||
thread_id=record.thread_id,
|
||
wait_version=record.wait_version,
|
||
pause_name=record.pause_name,
|
||
inquiry_no=record.inquiry_no,
|
||
status="open",
|
||
payload=dict(record.payload),
|
||
)
|
||
logger.info(
|
||
"suspend.memory.open thread=%s wait=%s pause=%s",
|
||
record.thread_id,
|
||
record.wait_version,
|
||
record.pause_name,
|
||
)
|
||
|
||
def get(self, thread_id: str, wait_version: int) -> Optional[SuspendRecord]:
|
||
with self._lock:
|
||
rec = self._rows.get((thread_id, wait_version))
|
||
if rec is None:
|
||
return None
|
||
return SuspendRecord(
|
||
thread_id=rec.thread_id,
|
||
wait_version=rec.wait_version,
|
||
pause_name=rec.pause_name,
|
||
inquiry_no=rec.inquiry_no,
|
||
status=rec.status,
|
||
payload=dict(rec.payload),
|
||
)
|
||
|
||
def close_wait(self, thread_id: str, wait_version: int) -> bool:
|
||
with self._lock:
|
||
rec = self._rows.get((thread_id, wait_version))
|
||
if rec is None:
|
||
return False
|
||
rec.status = "closed"
|
||
return True
|
||
|
||
|
||
class PostgresSuspendStore:
|
||
"""PG 停等壳（表 suspend_wait）。"""
|
||
|
||
def __init__(self, dsn: str) -> None:
|
||
assert_safe_database(dsn)
|
||
self._dsn = dsn
|
||
|
||
def open_wait(self, record: SuspendRecord) -> None:
|
||
import json
|
||
|
||
with connect_psycopg(self._dsn) as conn:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO suspend_wait (
|
||
thread_id, wait_version, pause_name, inquiry_no, status, payload
|
||
) VALUES (%s, %s, %s, %s, 'open', %s::jsonb)
|
||
ON CONFLICT (thread_id, wait_version) DO UPDATE SET
|
||
pause_name = EXCLUDED.pause_name,
|
||
inquiry_no = EXCLUDED.inquiry_no,
|
||
status = 'open',
|
||
payload = EXCLUDED.payload,
|
||
closed_at = NULL
|
||
""",
|
||
(
|
||
record.thread_id,
|
||
record.wait_version,
|
||
record.pause_name,
|
||
record.inquiry_no,
|
||
json.dumps(record.payload, ensure_ascii=False),
|
||
),
|
||
)
|
||
conn.commit()
|
||
logger.info(
|
||
"suspend.pg.open thread=%s wait=%s",
|
||
record.thread_id,
|
||
record.wait_version,
|
||
)
|
||
|
||
def get(self, thread_id: str, wait_version: int) -> Optional[SuspendRecord]:
|
||
with connect_psycopg(self._dsn) as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT thread_id, wait_version, pause_name, inquiry_no, status, payload
|
||
FROM suspend_wait
|
||
WHERE thread_id = %s AND wait_version = %s
|
||
""",
|
||
(thread_id, wait_version),
|
||
).fetchone()
|
||
if not row:
|
||
return None
|
||
return SuspendRecord(
|
||
thread_id=str(row[0]),
|
||
wait_version=int(row[1]),
|
||
pause_name=str(row[2] or ""),
|
||
inquiry_no=str(row[3] or ""),
|
||
status=str(row[4] or ""),
|
||
payload=dict(row[5] or {}),
|
||
)
|
||
|
||
def close_wait(self, thread_id: str, wait_version: int) -> bool:
|
||
with connect_psycopg(self._dsn) as conn:
|
||
cur = conn.execute(
|
||
"""
|
||
UPDATE suspend_wait
|
||
SET status = 'closed', closed_at = %s
|
||
WHERE thread_id = %s AND wait_version = %s AND status = 'open'
|
||
""",
|
||
(_utcnow(), thread_id, wait_version),
|
||
)
|
||
conn.commit()
|
||
return cur.rowcount > 0
|
||
|
||
|
||
def create_suspend_store(*, dsn: str = "", backend: str = "auto"):
|
||
"""工厂：Memory 或 PG。"""
|
||
if backend == "memory" or not dsn:
|
||
return MemorySuspendStore()
|
||
return PostgresSuspendStore(dsn)
|