Persist active run state for replay
This commit is contained in:
wuyang6
@@ -0,0 +1,298 @@
|
||||
"""持久化保存会话运行状态。
|
||||
|
||||
这个模块只保存“运行事实”：当前 run 的状态、开始/结束时间、阶段说明和
|
||||
活动事件。它不负责取消进程；取消仍由内存里的 RunManager 持有进程句柄来做。
|
||||
前端刷新、切换会话或服务短暂重启后，可以通过这里恢复可靠的展示状态。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
import time
|
||||
from pathlib import Path
|
||||
from threading import RLock
|
||||
from typing import Any
|
||||
|
||||
|
||||
ACTIVE_RUN_STATUSES = {'queued', 'running'}
|
||||
TERMINAL_RUN_STATUSES = {'completed', 'failed', 'cancelled', 'interrupted'}
|
||||
|
||||
|
||||
class RunStateStore:
|
||||
"""SQLite backed run-state store.
|
||||
|
||||
数据库放在 `.port_sessions/run_state.db`，用 account_key/session_id 做隔离。
|
||||
SQLite 足够支撑这个场景：写入量小，读多写少，并且可以避免前端回放依赖
|
||||
后端进程内存。
|
||||
"""
|
||||
|
||||
def __init__(self, db_path: Path, *, max_events_per_run: int = 240) -> None:
|
||||
self.db_path = db_path
|
||||
self.max_events_per_run = max_events_per_run
|
||||
self._lock = RLock()
|
||||
|
||||
def start(
|
||||
self,
|
||||
*,
|
||||
run_id: str,
|
||||
account_key: str,
|
||||
session_id: str,
|
||||
pending_prompt: str,
|
||||
started_at: float,
|
||||
status: str = 'queued',
|
||||
) -> None:
|
||||
now = time.time()
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
insert into run_states (
|
||||
run_id, account_key, session_id, status, started_at,
|
||||
updated_at, pending_prompt, cancellable
|
||||
)
|
||||
values (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
on conflict(run_id) do update set
|
||||
account_key=excluded.account_key,
|
||||
session_id=excluded.session_id,
|
||||
status=excluded.status,
|
||||
started_at=excluded.started_at,
|
||||
updated_at=excluded.updated_at,
|
||||
pending_prompt=excluded.pending_prompt,
|
||||
cancellable=excluded.cancellable
|
||||
""",
|
||||
(
|
||||
run_id,
|
||||
account_key,
|
||||
session_id,
|
||||
status,
|
||||
started_at,
|
||||
now,
|
||||
pending_prompt,
|
||||
1 if status in ACTIVE_RUN_STATUSES else 0,
|
||||
),
|
||||
)
|
||||
|
||||
def update(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
status: str | None = None,
|
||||
stage: str | None = None,
|
||||
error: str | None = None,
|
||||
cancellable: bool | None = None,
|
||||
elapsed_ms: int | None = None,
|
||||
finished_at: float | None = None,
|
||||
) -> None:
|
||||
assignments: list[str] = ['updated_at = ?']
|
||||
values: list[Any] = [time.time()]
|
||||
if status is not None:
|
||||
assignments.append('status = ?')
|
||||
values.append(status)
|
||||
if cancellable is None:
|
||||
cancellable = status in ACTIVE_RUN_STATUSES
|
||||
if stage is not None:
|
||||
assignments.append('current_stage = ?')
|
||||
values.append(stage)
|
||||
if error is not None:
|
||||
assignments.append('error = ?')
|
||||
values.append(error)
|
||||
if cancellable is not None:
|
||||
assignments.append('cancellable = ?')
|
||||
values.append(1 if cancellable else 0)
|
||||
if elapsed_ms is not None:
|
||||
assignments.append('elapsed_ms = ?')
|
||||
values.append(max(0, int(elapsed_ms)))
|
||||
if finished_at is not None:
|
||||
assignments.append('finished_at = ?')
|
||||
values.append(finished_at)
|
||||
values.append(run_id)
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
f"update run_states set {', '.join(assignments)} where run_id = ?",
|
||||
values,
|
||||
)
|
||||
|
||||
def finish(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
status: str,
|
||||
elapsed_ms: int | None = None,
|
||||
stage: str | None = None,
|
||||
error: str | None = None,
|
||||
) -> None:
|
||||
self.update(
|
||||
run_id,
|
||||
status=status,
|
||||
stage=stage,
|
||||
error=error,
|
||||
cancellable=False,
|
||||
elapsed_ms=elapsed_ms,
|
||||
finished_at=time.time(),
|
||||
)
|
||||
|
||||
def record_event(self, run_id: str, event: dict[str, Any]) -> None:
|
||||
recorded_at = event.get('recorded_at')
|
||||
if not isinstance(recorded_at, (int, float)):
|
||||
recorded_at = time.time()
|
||||
event = {**event, 'recorded_at': recorded_at}
|
||||
event_json = json.dumps(event, ensure_ascii=False, sort_keys=True)
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
insert into run_events (run_id, recorded_at, event_json)
|
||||
values (?, ?, ?)
|
||||
""",
|
||||
(run_id, float(recorded_at), event_json),
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
delete from run_events
|
||||
where run_id = ?
|
||||
and id not in (
|
||||
select id
|
||||
from run_events
|
||||
where run_id = ?
|
||||
order by id desc
|
||||
limit ?
|
||||
)
|
||||
""",
|
||||
(run_id, run_id, self.max_events_per_run),
|
||||
)
|
||||
|
||||
def snapshot_latest(self, account_key: str, session_id: str) -> dict[str, Any] | None:
|
||||
with self._connect() as conn:
|
||||
row = conn.execute(
|
||||
"""
|
||||
select *
|
||||
from run_states
|
||||
where account_key = ? and session_id = ?
|
||||
order by updated_at desc, started_at desc
|
||||
limit 1
|
||||
""",
|
||||
(account_key, session_id),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return self._snapshot_from_row(conn, row)
|
||||
|
||||
def snapshot_run(self, run_id: str) -> dict[str, Any] | None:
|
||||
with self._connect() as conn:
|
||||
row = conn.execute(
|
||||
"select * from run_states where run_id = ?",
|
||||
(run_id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return self._snapshot_from_row(conn, row)
|
||||
|
||||
def mark_interrupted_if_active(
|
||||
self,
|
||||
run_id: str,
|
||||
*,
|
||||
stage: str = '上次运行已中断，后台没有正在执行的进程',
|
||||
) -> dict[str, Any] | None:
|
||||
snapshot = self.snapshot_run(run_id)
|
||||
if snapshot is None:
|
||||
return None
|
||||
if snapshot.get('status') not in ACTIVE_RUN_STATUSES:
|
||||
return snapshot
|
||||
started_at = snapshot.get('started_at')
|
||||
elapsed_ms = None
|
||||
if isinstance(started_at, (int, float)):
|
||||
elapsed_ms = max(0, int((time.time() - float(started_at)) * 1000))
|
||||
self.finish(
|
||||
run_id,
|
||||
status='interrupted',
|
||||
elapsed_ms=elapsed_ms,
|
||||
stage=stage,
|
||||
)
|
||||
return self.snapshot_run(run_id)
|
||||
|
||||
def _snapshot_from_row(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
row: sqlite3.Row,
|
||||
) -> dict[str, Any]:
|
||||
event_rows = conn.execute(
|
||||
"""
|
||||
select event_json
|
||||
from run_events
|
||||
where run_id = ?
|
||||
order by id asc
|
||||
""",
|
||||
(row['run_id'],),
|
||||
).fetchall()
|
||||
events: list[dict[str, Any]] = []
|
||||
for event_row in event_rows:
|
||||
try:
|
||||
event = json.loads(event_row['event_json'])
|
||||
except (TypeError, json.JSONDecodeError):
|
||||
continue
|
||||
if isinstance(event, dict):
|
||||
events.append(event)
|
||||
return {
|
||||
'run_id': row['run_id'],
|
||||
'session_id': row['session_id'],
|
||||
'status': row['status'],
|
||||
'current_stage': row['current_stage'] or '',
|
||||
'started_at': row['started_at'],
|
||||
'updated_at': row['updated_at'],
|
||||
'finished_at': row['finished_at'],
|
||||
'elapsed_ms': row['elapsed_ms'],
|
||||
'pending_prompt': row['pending_prompt'] or '',
|
||||
'error': row['error'] or '',
|
||||
'cancellable': bool(row['cancellable']),
|
||||
'events': events,
|
||||
}
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
with self._lock:
|
||||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
conn = sqlite3.connect(self.db_path, timeout=10)
|
||||
conn.row_factory = sqlite3.Row
|
||||
self._init_schema(conn)
|
||||
return conn
|
||||
|
||||
def _init_schema(self, conn: sqlite3.Connection) -> None:
|
||||
conn.execute('pragma journal_mode=wal')
|
||||
conn.execute(
|
||||
"""
|
||||
create table if not exists run_states (
|
||||
run_id text primary key,
|
||||
account_key text not null,
|
||||
session_id text not null,
|
||||
status text not null,
|
||||
started_at real not null,
|
||||
updated_at real not null,
|
||||
finished_at real,
|
||||
elapsed_ms integer,
|
||||
pending_prompt text default '',
|
||||
current_stage text default '',
|
||||
error text default '',
|
||||
cancellable integer not null default 0
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
create index if not exists idx_run_states_session
|
||||
on run_states(account_key, session_id, updated_at)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
create table if not exists run_events (
|
||||
id integer primary key autoincrement,
|
||||
run_id text not null,
|
||||
recorded_at real not null,
|
||||
event_json text not null
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
create index if not exists idx_run_events_run_id
|
||||
on run_events(run_id, id)
|
||||
"""
|
||||
)
|
||||
Reference in New Issue
Block a user