wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Persist active run state for replay

Browse Source
This commit is contained in:
wuyang6
2026-05-14 11:47:37 +08:00
parent e3569b6901
commit 105d287ebd
4 changed files with 423 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/api/server.py
+104 -11
View File
@@ -55,6 +55,7 @@ from src.personal_memory import (
USER_MEMORY_FILENAME,
PersonalMemoryManager,
)
from src.run_state_store import ACTIVE_RUN_STATUSES, RunStateStore
from src.session_store import (
DEFAULT_AGENT_SESSION_DIR,
StoredAgentSession,
@@ -414,19 +415,20 @@ class RunManager:
record.error = error
record.updated_at = time.time()
def record_event(self, run_id: str, event: dict[str, object]) -> None:
def record_event(self, run_id: str, event: dict[str, object]) -> dict[str, Any] | None:
normalized = _normalize_run_event(event)
if normalized is None:
return
return None
normalized['recorded_at'] = time.time()
with self._lock:
record = self._runs.get(run_id)
if record is None:
return
return normalized
record.events.append(normalized)
if len(record.events) > 160:
del record.events[: len(record.events) - 160]
record.updated_at = time.time()
return normalized
def finish(self, run_id: str, status: str) -> None:
self.update(run_id, status=status)
@@ -515,6 +517,7 @@ class AgentState:
self._run_locks: dict[str, Lock] = {}
self._account_configs: dict[str, AgentInstanceConfig] = {}
self.run_manager = RunManager()
self.run_state_store = RunStateStore(self.session_directory.parent / 'run_state.db')
self.jupyter_runtime_manager = JupyterRuntimeManager()
self._default_config = AgentInstanceConfig(
cwd=cwd.resolve(),
@@ -1573,6 +1576,27 @@ def create_app(state: AgentState) -> FastAPI:
)
return _serialize_token_budget(snapshot)
def _merge_run_snapshot(
memory_snapshot: dict[str, Any] | None,
stored_snapshot: dict[str, Any] | None,
) -> dict[str, Any] | None:
if memory_snapshot is None:
return stored_snapshot
if stored_snapshot is None:
merged = dict(memory_snapshot)
else:
merged = {**stored_snapshot, **memory_snapshot}
if not memory_snapshot.get('events') and stored_snapshot.get('events'):
merged['events'] = stored_snapshot.get('events')
if memory_snapshot.get('current_stage'):
merged['current_stage'] = memory_snapshot.get('current_stage')
if stored_snapshot.get('elapsed_ms') is not None:
merged['elapsed_ms'] = stored_snapshot.get('elapsed_ms')
if stored_snapshot.get('finished_at') is not None:
merged['finished_at'] = stored_snapshot.get('finished_at')
merged['cancellable'] = merged.get('status') in ACTIVE_RUN_STATUSES
return merged
@app.get('/api/runs/latest')
async def latest_run(
session_id: str,
@@ -1581,12 +1605,18 @@ def create_app(state: AgentState) -> FastAPI:
safe_id = _safe_session_id(session_id)
if safe_id is None:
raise HTTPException(status_code=400, detail='Invalid session id')
snapshot = state.run_manager.snapshot_latest(
state._account_key(account_id),
safe_id,
)
account_key = state._account_key(account_id)
snapshot = state.run_manager.snapshot_latest(account_key, safe_id)
stored_snapshot = state.run_state_store.snapshot_latest(account_key, safe_id)
if snapshot is not None:
return snapshot
return _merge_run_snapshot(snapshot, stored_snapshot) or snapshot
if stored_snapshot is not None:
if stored_snapshot.get('status') in ACTIVE_RUN_STATUSES:
interrupted = state.run_state_store.mark_interrupted_if_active(
str(stored_snapshot.get('run_id') or ''),
)
return interrupted or stored_snapshot
return stored_snapshot
directory = state.account_paths(account_id)['sessions']
try:
stored = load_agent_session(safe_id, directory=directory)
@@ -1614,6 +1644,25 @@ def create_app(state: AgentState) -> FastAPI:
)
)
if cancelled:
run_id = payload.run_id
if not run_id:
snapshot = state.run_state_store.snapshot_latest(
state._account_key(payload.account_id),
safe_id,
)
run_id = str(snapshot.get('run_id') or '') if snapshot else ''
if run_id:
stored = state.run_state_store.snapshot_run(run_id)
elapsed_ms = None
started_at = stored.get('started_at') if stored else None
if isinstance(started_at, (int, float)):
elapsed_ms = max(0, int((time.time() - float(started_at)) * 1000))
state.run_state_store.finish(
run_id,
status='cancelled',
elapsed_ms=elapsed_ms,
stage='用户已取消',
)
_mark_session_interrupted(
state.account_paths(payload.account_id)['sessions'],
safe_id,
@@ -1631,6 +1680,13 @@ def create_app(state: AgentState) -> FastAPI:
account_key = state._account_key(request.account_id)
prompt = request.prompt.strip()
run_record = state.run_manager.start(account_key, requested_session_id, prompt)
state.run_state_store.start(
run_id=run_record.run_id,
account_key=account_key,
session_id=requested_session_id,
pending_prompt=prompt,
started_at=run_record.started_at,
)
run_lock = state.run_lock_for(request.account_id, requested_session_id)
agent = state.agent_for(request.account_id, requested_session_id)
config = state.config_for(request.account_id)
@@ -1657,7 +1713,10 @@ def create_app(state: AgentState) -> FastAPI:
stage = _runtime_event_stage(event)
if stage:
state.run_manager.update(run_record.run_id, stage=stage)
state.run_manager.record_event(run_record.run_id, event)
state.run_state_store.update(run_record.run_id, stage=stage)
stored_event = state.run_manager.record_event(run_record.run_id, event)
if stored_event is not None:
state.run_state_store.record_event(run_record.run_id, stored_event)
_emit_runtime_event(event_sink, event)
previous_title = _read_session_title(
@@ -1679,7 +1738,9 @@ def create_app(state: AgentState) -> FastAPI:
'session_id': requested_session_id or '',
'account_id': request.account_id or '',
}
state.run_manager.record_event(run_record.run_id, queued_event)
stored_event = state.run_manager.record_event(run_record.run_id, queued_event)
if stored_event is not None:
state.run_state_store.record_event(run_record.run_id, stored_event)
_emit_runtime_event(
event_sink,
queued_event,
@@ -1687,6 +1748,12 @@ def create_app(state: AgentState) -> FastAPI:
with run_lock:
if run_record.cancel_event.is_set():
state.run_manager.finish(run_record.run_id, 'cancelled')
state.run_state_store.finish(
run_record.run_id,
status='cancelled',
elapsed_ms=max(0, int((time.time() - run_record.started_at) * 1000)),
stage='已取消排队中的请求',
)
return {
'final_output': '已取消排队中的请求。',
'turns': 0,
@@ -1698,6 +1765,12 @@ def create_app(state: AgentState) -> FastAPI:
'stop_reason': 'cancelled',
}
state.run_manager.update(run_record.run_id, status='running', stage='后端开始执行')
state.run_state_store.update(
run_record.run_id,
status='running',
stage='后端开始执行',
cancellable=True,
)
started_event = {
'type': 'run_started',
'run_id': run_record.run_id,
@@ -1705,7 +1778,9 @@ def create_app(state: AgentState) -> FastAPI:
'session_id': requested_session_id or '',
'account_id': request.account_id or '',
}
state.run_manager.record_event(run_record.run_id, started_event)
stored_event = state.run_manager.record_event(run_record.run_id, started_event)
if stored_event is not None:
state.run_state_store.record_event(run_record.run_id, stored_event)
_emit_runtime_event(event_sink, started_event)
agent.tool_context = replace(
agent.tool_context,
@@ -1761,6 +1836,19 @@ def create_app(state: AgentState) -> FastAPI:
),
error=str(exc),
)
state.run_state_store.finish(
run_record.run_id,
status=(
'cancelled'
if run_record.cancel_event.is_set()
else 'failed'
),
elapsed_ms=max(
0,
int((time.time() - run_record.started_at) * 1000),
),
error=str(exc),
)
raise
finally:
agent.tool_context = replace(
@@ -1806,6 +1894,11 @@ def create_app(state: AgentState) -> FastAPI:
run_record.run_id,
'cancelled' if run_record.cancel_event.is_set() else 'completed',
)
state.run_state_store.finish(
run_record.run_id,
status='cancelled' if run_record.cancel_event.is_set() else 'completed',
elapsed_ms=max(0, int((time.time() - run_record.started_at) * 1000)),
)
if run_record.cancel_event.is_set() and result.session_id:
_mark_session_interrupted(
session_directory,
frontend/app/components/assistant-ui/thread-list.tsx
+12 -1
View File
@@ -69,6 +69,10 @@ export type ClawRunStatus = {
| "interrupted";
current_stage?: string;
started_at?: number;
updated_at?: number;
finished_at?: number;
elapsed_ms?: number;
cancellable?: boolean;
pending_prompt?: string;
events?: ClawRunEvent[];
};
@@ -590,8 +594,9 @@ function buildActiveRunParts(runStatus: ClawActiveRunStatus) {
: runStatus.current_stage || "后端正在执行这个会话。";
const reasoningText = eventLines.length ? eventLines.join("\n") : fallback;
const runStartedAtMs = resolveRunStartedAtMs(runStatus);
const elapsedMs = normalizeElapsedMs(runStatus.elapsed_ms);
return [
{ type: "reasoning", text: reasoningText, runStartedAtMs },
{ type: "reasoning", text: reasoningText, runStartedAtMs, elapsedMs },
...buildRunToolParts(runStatus.events ?? []),
];
}
@@ -795,6 +800,12 @@ function resolveRunStartedAtMs(runStatus: ClawRunStatus) {
const direct = normalizeRunTimestampMs(runStatus.started_at);
if (direct !== undefined) return direct;
const updatedAt = normalizeRunTimestampMs(runStatus.updated_at);
const elapsedMs = normalizeElapsedMs(runStatus.elapsed_ms);
if (updatedAt !== undefined && elapsedMs !== undefined) {
return Math.max(0, updatedAt - elapsedMs);
}
for (const event of runStatus.events ?? []) {
if (event.type !== "run_started") continue;
const eventStarted = normalizeRunTimestampMs(event.started_at);
frontend/app/components/assistant-ui/thread.tsx
+9 -7
View File
@@ -1884,7 +1884,6 @@ const ActivityChainSummary: FC<{
partIndex: number | undefined;
}> = ({ messageId, partIndex }) => {
const { openItem } = useActivityPanel();
const startedAtRef = useRef(Date.now());
const metadataRunStartedAtMs = useAuiState((s) => {
const value = readMetadataValue(s.message.metadata, "runStartedAtMs");
return typeof value === "number" && Number.isFinite(value) ? value : null;
@@ -1933,9 +1932,11 @@ const ActivityChainSummary: FC<{
const startedAt =
metadataRunStartedAtMs ??
contentRunStartedAtMs ??
(storedElapsedMs === null
? startedAtRef.current
: Date.now() - storedElapsedMs);
(storedElapsedMs === null ? null : Date.now() - storedElapsedMs);
if (startedAt === null) {
setElapsedMs(storedElapsedMs);
return;
}
setElapsedMs(Math.max(0, Date.now() - startedAt));
const timer = window.setInterval(() => {
setElapsedMs(Math.max(0, Date.now() - startedAt));
@@ -2021,11 +2022,12 @@ function summarizeActivityChainLabel(
(isGenericActivityLine(lastReasoningLine) ? "" : lastReasoningLine);
if (status === "running") {
const elapsed = `用时 ${formatActivityDuration(elapsedMs ?? 0)}`;
const elapsed =
elapsedMs === null ? "" : `，用时 ${formatActivityDuration(elapsedMs)}`;
if (activeDetail) {
return `${compactActivityLabel(activeDetail, 15)}，${elapsed}`;
return `${compactActivityLabel(activeDetail, 15)}${elapsed}`;
}
return `思考中，${elapsed}`;
return elapsed ? `思考中${elapsed}` : "思考中";
}
const label = compactActivityLabel(
src/run_state_store.py
+298
View File
@@ -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)
"""
)