1use crate::state::types::{
4 DefinitionSnapshot, Manifest, ManifestPaths, RunState, SessionBinding, SessionCapture,
5 SessionEntryRecord, SessionEventRecord, TraceEvent, DEFINITION_SNAPSHOT_SCHEMA,
6 RUN_STATE_SCHEMA,
7};
8use anyhow::{bail, Context, Result};
9use chrono::{TimeZone, Utc};
10use rusqlite::{Connection, OpenFlags, OptionalExtension};
11use serde_json::{json, Value};
12use std::path::Path;
13
14const APPLICATION_ID: i64 = 0x5049_5746;
15const USER_VERSION: i64 = 1;
16
17type LoadedSession = (
18 Option<SessionBinding>,
19 Vec<SessionEntryRecord>,
20 Vec<SessionEventRecord>,
21 Option<SessionCapture>,
22);
23
24#[derive(Debug, Clone)]
25pub struct LoadedRun {
26 pub manifest: Manifest,
27 pub state: RunState,
28 pub snapshot: Option<DefinitionSnapshot>,
29 pub trace: Vec<TraceEvent>,
30 pub session_binding: Option<SessionBinding>,
31 pub session_entries: Vec<SessionEntryRecord>,
32 pub session_events: Vec<SessionEventRecord>,
33 pub session_capture: Option<SessionCapture>,
34 pub possibly_interrupted: bool,
35}
36
37pub fn read_run(database_path: &Path, run_id: &str) -> Result<LoadedRun> {
38 let connection = open(database_path)?;
39 let row = connection
40 .query_row(
41 "SELECT r.output_hash, d.definition_hash, l.owner_id, l.expires_at
42 FROM runs r
43 JOIN workflow_definitions d ON d.definition_digest = r.definition_digest
44 JOIN leases l ON l.resource_id = r.resource_id
45 WHERE r.run_id = ?1",
46 [run_id],
47 |row| {
48 Ok((
49 row.get::<_, Vec<u8>>(0)?,
50 row.get::<_, Vec<u8>>(1)?,
51 row.get::<_, Option<String>>(2)?,
52 row.get::<_, Option<i64>>(3)?,
53 ))
54 },
55 )
56 .optional()?;
57 let Some((state_hash, definition_hash, owner_id, lease_expires_at)) = row else {
58 bail!("workflow run not found: {run_id}");
59 };
60 let state_value = read_json_blob(&connection, &state_hash)?;
61 let state: RunState = serde_json::from_value(state_value.clone())?;
62 if state.schema != RUN_STATE_SCHEMA {
63 bail!("unsupported workflow state schema: {}", state.schema);
64 }
65 let definition_value = read_json_blob(&connection, &definition_hash)?;
66 let snapshot: DefinitionSnapshot = serde_json::from_value(definition_value)?;
67 if snapshot.schema != DEFINITION_SNAPSHOT_SCHEMA {
68 bail!(
69 "unsupported workflow definition schema: {}",
70 snapshot.schema
71 );
72 }
73 let trace = read_trace(&connection, run_id)?;
74 let (session_binding, session_entries, session_events, session_capture) =
75 read_session(&connection, run_id)?;
76 let manifest = manifest_from_state(&state);
77 let possibly_interrupted = state.status.label() == "running"
78 && (owner_id.is_none()
79 || lease_expires_at
80 .is_none_or(|expires_at| expires_at <= Utc::now().timestamp_millis()));
81 Ok(LoadedRun {
82 manifest,
83 state,
84 snapshot: Some(snapshot),
85 trace,
86 session_binding,
87 session_entries,
88 session_events,
89 session_capture,
90 possibly_interrupted,
91 })
92}
93
94pub fn list_runs(database_path: &Path) -> Vec<(String, Manifest)> {
95 let Ok(connection) = open(database_path) else {
96 return Vec::new();
97 };
98 let Ok(mut statement) = connection.prepare("SELECT run_id FROM runs ORDER BY created_at DESC")
99 else {
100 return Vec::new();
101 };
102 let Ok(ids) = statement.query_map([], |row| row.get::<_, String>(0)) else {
103 return Vec::new();
104 };
105 ids.filter_map(|id| {
106 let id = id.ok()?;
107 let run = read_run(database_path, &id).ok()?;
108 Some((id, run.manifest))
109 })
110 .collect()
111}
112
113pub fn with_artifact_placeholders(value: &Value) -> Value {
114 value.clone()
115}
116
117pub fn resolve_artifacts(value: &Value, _database_path: &Path, _max_bytes: u64) -> Value {
118 value.clone()
119}
120
121fn open(path: &Path) -> Result<Connection> {
122 let connection = Connection::open_with_flags(
123 path,
124 OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
125 )
126 .with_context(|| format!("could not open {}", path.display()))?;
127 connection.pragma_update(None, "query_only", true)?;
128 connection.pragma_update(None, "foreign_keys", true)?;
129 let application_id: i64 =
130 connection.pragma_query_value(None, "application_id", |row| row.get(0))?;
131 let user_version: i64 =
132 connection.pragma_query_value(None, "user_version", |row| row.get(0))?;
133 if application_id != APPLICATION_ID || user_version != USER_VERSION {
134 bail!("incompatible Pi Workflows SQLite schema");
135 }
136 Ok(connection)
137}
138
139fn read_json_blob(connection: &Connection, hash: &[u8]) -> Result<Value> {
140 let content: Vec<u8> = connection.query_row(
141 "SELECT content FROM blobs WHERE blob_hash = ?1 AND media_type = 'application/json'",
142 [hash],
143 |row| row.get(0),
144 )?;
145 Ok(serde_json::from_slice(&content)?)
146}
147
148fn read_trace(connection: &Connection, run_id: &str) -> Result<Vec<TraceEvent>> {
149 let resource_id: String = connection.query_row(
150 "SELECT resource_id FROM runs WHERE run_id = ?1",
151 [run_id],
152 |row| row.get(0),
153 )?;
154 let mut statement = connection.prepare(
155 "SELECT resource_revision, event_type, payload_hash, recorded_at
156 FROM events WHERE resource_id = ?1 ORDER BY resource_revision",
157 )?;
158 let rows = statement.query_map([resource_id], |row| {
159 Ok((
160 row.get::<_, u64>(0)?,
161 row.get::<_, String>(1)?,
162 row.get::<_, Option<Vec<u8>>>(2)?,
163 row.get::<_, i64>(3)?,
164 ))
165 })?;
166 let mut events = Vec::new();
167 for row in rows {
168 let (seq, event_type, payload_hash, recorded_at) = row?;
169 let envelope = match payload_hash {
170 Some(hash) => read_json_blob(connection, &hash)?,
171 None => json!({}),
172 };
173 let payload = envelope
174 .get("payload")
175 .cloned()
176 .unwrap_or_else(|| json!({}));
177 let mut event = json!({
178 "seq": seq,
179 "at": timestamp(recorded_at),
180 "runId": run_id,
181 "scope": envelope.get("scope").and_then(Value::as_str).unwrap_or("run"),
182 "type": event_type,
183 "payload": payload,
184 });
185 if let Some(node_id) = envelope.get("nodeId") {
186 event["nodeId"] = node_id.clone();
187 }
188 if let Some(attempt_id) = envelope.get("attemptId") {
189 event["attemptId"] = attempt_id.clone();
190 }
191 events.push(serde_json::from_value(event)?);
192 }
193 Ok(events)
194}
195
196fn read_session(connection: &Connection, run_id: &str) -> Result<LoadedSession> {
197 let mut segments_statement = connection.prepare(
198 "SELECT segment_id, binding_hash, status, entry_count, event_count,
199 failure_hash
200 FROM session_segments
201 WHERE run_id = ?1
202 ORDER BY created_at, segment_id",
203 )?;
204 let segment_rows = segments_statement
205 .query_map([run_id], |row| {
206 Ok((
207 row.get::<_, String>(0)?,
208 row.get::<_, Option<Vec<u8>>>(1)?,
209 row.get::<_, String>(2)?,
210 row.get::<_, u64>(3)?,
211 row.get::<_, u64>(4)?,
212 row.get::<_, Option<Vec<u8>>>(5)?,
213 ))
214 })?
215 .collect::<Result<Vec<_>, _>>()?;
216 if segment_rows.is_empty() {
217 return Ok((None, Vec::new(), Vec::new(), None));
218 }
219
220 let mut binding = None;
221 let mut entries = Vec::new();
222 let mut events = Vec::new();
223 let mut status = "complete".to_string();
224 let mut failure = None;
225
226 for (segment_id, binding_hash, segment_status, _entry_count, _event_count, failure_hash) in
227 segment_rows
228 {
229 if binding.is_none() {
230 if let Some(hash) = binding_hash {
231 binding = Some(serde_json::from_value(read_json_blob(connection, &hash)?)?);
232 }
233 }
234 if segment_status == "failed" {
235 status = "failed".to_string();
236 if failure.is_none() {
237 if let Some(hash) = failure_hash {
238 failure = Some(read_json_blob(connection, &hash)?);
239 }
240 }
241 } else if segment_status == "recording" && status != "failed" {
242 status = "recording".to_string();
243 }
244
245 let mut entries_statement = connection.prepare(
246 "SELECT entry_hash, recorded_at
247 FROM session_entries WHERE segment_id = ?1 ORDER BY entry_seq",
248 )?;
249 let segment_entries = entries_statement
250 .query_map([&segment_id], |row| {
251 Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, i64>(1)?))
252 })?
253 .collect::<Result<Vec<_>, _>>()?;
254 for (hash, at) in segment_entries {
255 entries.push(serde_json::from_value(json!({
256 "seq": entries.len() + 1,
257 "at": timestamp(at),
258 "entry": read_json_blob(connection, &hash)?,
259 }))?);
260 }
261
262 let mut events_statement = connection.prepare(
263 "SELECT event_type, node_id, attempt_id, turn_id,
264 message_id, tool_call_id, payload_hash, recorded_at
265 FROM session_events WHERE segment_id = ?1 ORDER BY event_seq",
266 )?;
267 let segment_events = events_statement
268 .query_map([&segment_id], |row| {
269 Ok((
270 row.get::<_, String>(0)?,
271 row.get::<_, String>(1)?,
272 row.get::<_, String>(2)?,
273 row.get::<_, Option<String>>(3)?,
274 row.get::<_, Option<String>>(4)?,
275 row.get::<_, Option<String>>(5)?,
276 row.get::<_, Vec<u8>>(6)?,
277 row.get::<_, i64>(7)?,
278 ))
279 })?
280 .collect::<Result<Vec<_>, _>>()?;
281 for (event_type, node_id, attempt_id, turn_id, message_id, tool_call_id, hash, at) in
282 segment_events
283 {
284 let mut value = json!({
285 "seq": events.len() + 1,
286 "at": timestamp(at),
287 "nodeId": node_id,
288 "attemptId": attempt_id,
289 "type": event_type,
290 "payload": read_json_blob(connection, &hash)?,
291 });
292 if let Some(turn_id) = turn_id {
293 value["turnId"] = json!(turn_id);
294 }
295 if let Some(message_id) = message_id {
296 value["messageId"] = json!(message_id);
297 }
298 if let Some(tool_call_id) = tool_call_id {
299 value["toolCallId"] = json!(tool_call_id);
300 }
301 events.push(serde_json::from_value(value)?);
302 }
303 }
304
305 let mut capture = json!({
306 "schema": "pi-workflows.session-capture.v1",
307 "eventSchema": "pi-workflows.session-event.v1",
308 "status": status,
309 "eventCount": events.len(),
310 "entryCount": entries.len(),
311 "lastEventSeq": events.len(),
312 });
313 if let Some(failure) = failure {
314 capture["failure"] = failure;
315 }
316 Ok((
317 binding,
318 entries,
319 events,
320 Some(serde_json::from_value(capture)?),
321 ))
322}
323
324fn manifest_from_state(state: &RunState) -> Manifest {
325 Manifest {
326 schema: "pi-workflows.sqlite-view.v1".to_string(),
327 run_id: state.run_id.clone(),
328 workflow_name: state.workflow_name.clone(),
329 run_title: state.run_title.clone(),
330 workflow_source: state.workflow_source.clone(),
331 started_at: state.started_at.clone(),
332 finished_at: state.finished_at.clone(),
333 status: state.status,
334 trace_schema: "pi-workflows.event.v1".to_string(),
335 paths: ManifestPaths {
336 workflow: String::new(),
337 state: String::new(),
338 trace: String::new(),
339 session: None,
340 artifacts: None,
341 },
342 }
343}
344
345fn timestamp(milliseconds: i64) -> String {
346 Utc.timestamp_millis_opt(milliseconds)
347 .single()
348 .map(|value| value.to_rfc3339_opts(chrono::SecondsFormat::Millis, true))
349 .unwrap_or_else(|| "1970-01-01T00:00:00.000Z".to_string())
350}