1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
//! App-layer DAG persistence: debounced async writer backed by the core
//! [`SqliteDagStore`], wired into the engine via the [`DagPersistSink`]
//! contract.
//!
//! Saves are grouped per run session: each run is written to its owning
//! session's state file (`<cwd>/.pi/graph-engineering-state-<sessionId>.db`),
//! so resuming a different session never mixes another session's runs into
//! its state file. Session-less runs go to the default file.
//!
//! Lifecycle:
//! - [`DagPersistHandle::spawn`] creates the handle, wires it into the
//! engine, and starts the debounce task (500 ms merge window).
//! - Every engine state change fires [`DagPersistSink::notify_dirty`]; the
//! task coalesces notifications and writes once per window.
//! - Shutdown calls [`DagPersistSink::flush`] — a synchronous save of the
//! current running state — *before* `abort_all_runs`, so in-flight runs
//! survive a clean exit and are re-launched on the next `restore`.
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use parking_lot::Mutex;
use theway_contract::dag::{PersistedRun, state_path_for_project};
use theway_core::multiagent::graph::engine::DagEngine;
use theway_core::multiagent::graph::persist::{DagPersistSink, to_persisted};
use theway_core::multiagent::graph::types::DagStatus;
use theway_storage::sqlite_dag::SqliteDagStore;
use tokio::sync::Notify;
use tokio::task::JoinHandle;
use crate::session_execution::SessionExecutionRegistry;
/// Debounce window: state changes within this window coalesce into one write.
const DEBOUNCE: Duration = Duration::from_millis(500);
/// Shared handle driving debounced persistence for one engine.
pub struct DagPersistHandle {
engine: Arc<DagEngine>,
cwd: PathBuf,
sessions: SessionExecutionRegistry,
/// Per-owning-cwd/session stores (opened lazily, kept for the process lifetime).
stores: Mutex<HashMap<(PathBuf, Option<String>), SqliteDagStore>>,
/// Dirty signal from the engine.
dirty: Arc<Notify>,
/// Background write task.
task: Mutex<Option<JoinHandle<()>>>,
}
impl DagPersistHandle {
/// Create the handle, wire it into the engine as the persist sink, and
/// spawn the debounce task. Returns the handle (keep it alive; dropping
/// the JoinHandle detaches the task — the flush path is the caller's
/// shutdown responsibility regardless).
pub fn spawn(engine: Arc<DagEngine>, cwd: PathBuf) -> Arc<Self> {
Self::spawn_with_sessions(engine, cwd, SessionExecutionRegistry::new())
}
/// Create the handle with per-session cwd routing for registered contexts.
pub fn spawn_with_sessions(
engine: Arc<DagEngine>,
cwd: PathBuf,
sessions: SessionExecutionRegistry,
) -> Arc<Self> {
let dirty = Arc::new(Notify::new());
let handle = Arc::new(Self {
engine,
cwd,
sessions,
stores: Mutex::new(HashMap::new()),
dirty: dirty.clone(),
task: Mutex::new(None),
});
let task = tokio::spawn(handle.clone().run_loop());
*handle.task.lock() = Some(task);
handle.engine.set_persist_sink(Some(handle.clone()));
handle
}
/// Background loop: wait for a dirty signal, drain coalesced notifications
/// within the debounce window, then save the current engine state.
async fn run_loop(self: Arc<Self>) {
loop {
self.dirty.notified().await;
// Coalesce: keep draining notifications while they arrive within
// the window; stop when the window closes silently.
loop {
tokio::select! {
_ = self.dirty.notified() => {
// Another state change landed — reset the window.
tokio::time::sleep(DEBOUNCE).await;
}
_ = tokio::time::sleep(DEBOUNCE) => break,
}
}
if let Err(e) = self.save_all().await {
tracing::warn!("dag persist: {e}");
}
}
}
/// Save every running run to its owning session's store.
async fn save_all(&self) -> Result<(), String> {
let runs = self.engine.list_runs();
// Group by session id so each run lands in its own session's file.
let mut by_session: HashMap<Option<String>, Vec<PersistedRun>> = HashMap::new();
for run in runs {
let snapshots = by_session.entry(run.session_id.clone()).or_default();
if run.status == DagStatus::Running {
snapshots.push(to_persisted(&run));
}
}
for (session_id, session_runs) in by_session {
let store = self.store_for(session_id.as_deref()).await?;
store.save(&session_runs).await?;
}
Ok(())
}
/// Open (or reuse) the store for a session id, using the registered
/// context's cwd when one exists and the startup cwd otherwise.
async fn store_for(&self, session_id: Option<&str>) -> Result<SqliteDagStore, String> {
let owning_cwd = match session_id {
Some(session_id) => self
.sessions
.cwd_for(session_id)
.unwrap_or_else(|| self.cwd.clone()),
None => self.cwd.clone(),
};
let key = (owning_cwd.clone(), session_id.map(str::to_string));
if let Some(store) = self.stores.lock().get(&key) {
return Ok(store.clone());
}
let path = state_path_for_project(&owning_cwd.join(".pi"), session_id);
let store = SqliteDagStore::open(path).await?;
self.stores.lock().insert(key, store.clone());
Ok(store)
}
}
#[async_trait]
impl DagPersistSink for DagPersistHandle {
fn notify_dirty(&self) {
self.dirty.notify_one();
}
/// Synchronous save of the current running state (shutdown path). Skips
/// the debounce window; returns only after the write is durable.
async fn flush(&self) {
if let Err(e) = self.save_all().await {
tracing::warn!("dag persist flush: {e}");
}
}
}
/// Load persisted runs for a session (startup / session-factory restore path).
pub async fn load_session_runs(
cwd: &std::path::Path,
session_id: &str,
) -> Vec<theway_core::multiagent::graph::persist::PersistedRun> {
let path = state_path_for_project(&cwd.join(".pi"), Some(session_id));
match SqliteDagStore::open(path).await {
Ok(store) => store.load().await.unwrap_or_default(),
Err(e) => {
tracing::warn!("dag state open: {e}");
Vec::new()
}
}
}
#[cfg(test)]
tests_bridge_macro::tests_bridge!("dag_persist");