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,