1mod attachment_manifest;
4mod lease_timings;
5pub mod queued_work;
6
7pub use attachment_manifest::{AttachmentIntent, AttachmentManifest, AttachmentManifestEntry};
8pub use lease_timings::{LeaseTimings, LeaseTimingsError};
9
10const PROC_BOOT_ID_PATH: &str = "/proc/sys/kernel/random/boot_id";
11
12fn default_root_session_id() -> String {
13 "root".to_string()
14}
15
16pub const SESSION_HEAD_META_SCHEMA_VERSION: u32 = 1;
17pub const SESSION_CHECKPOINT_SCHEMA_VERSION: u32 = 1;
18
19#[cfg(test)]
20mod persisted_state_tests {
21 use super::*;
22
23 #[test]
24 fn persisted_state_hydrates_provider_id_without_live_provider_rebinding() {
25 let state = persisted_session_state_from_head(
26 SessionHead {
27 session_id: "stored".to_string(),
28 head_revision: 7,
29 agent_frames: Vec::new(),
30 current_agent_frame_id: String::new(),
31 graph: crate::SessionGraph::default(),
32 config: crate::PersistedSessionConfig {
33 provider_id: "stored-provider".to_string(),
34 model: crate::ModelSpec::default(),
35 },
36 checkpoint_ref: None,
37 token_ledger: Vec::new(),
38 },
39 None,
40 );
41
42 assert_eq!(state.policy.recorded_provider_id(), "stored-provider");
43 assert!(
44 state
45 .agent_frames
46 .iter()
47 .all(|frame| frame.assignment.policy.recorded_provider_id() == "stored-provider")
48 );
49 assert_eq!(state.head_revision, Some(7));
50 }
51
52 #[test]
53 fn versioned_json_record_rejects_missing_schema_version() {
54 let err = decode_versioned_json_record::<SessionHeadMeta>(
55 "{}",
56 "SessionHeadMeta",
57 SESSION_HEAD_META_SCHEMA_VERSION,
58 )
59 .expect_err("pre-versioned session head should fail");
60
61 assert!(matches!(
62 err,
63 StoreError::MissingRecordSchemaVersion {
64 record_kind: "SessionHeadMeta",
65 expected: SESSION_HEAD_META_SCHEMA_VERSION
66 }
67 ));
68 }
69
70 #[test]
71 fn versioned_json_record_rejects_invalid_schema_version() {
72 let err = decode_versioned_json_record::<SessionHeadMeta>(
73 r#"{"schema_version":"1"}"#,
74 "SessionHeadMeta",
75 SESSION_HEAD_META_SCHEMA_VERSION,
76 )
77 .expect_err("invalid session head schema version should fail");
78
79 assert!(matches!(
80 err,
81 StoreError::InvalidRecordSchemaVersion {
82 record_kind: "SessionHeadMeta",
83 expected: SESSION_HEAD_META_SCHEMA_VERSION,
84 ..
85 }
86 ));
87 }
88
89 #[test]
90 fn versioned_json_record_rejects_unsupported_schema_version() {
91 let err = decode_versioned_json_record::<SessionHeadMeta>(
92 r#"{"schema_version":2}"#,
93 "SessionHeadMeta",
94 SESSION_HEAD_META_SCHEMA_VERSION,
95 )
96 .expect_err("unsupported session head schema version should fail");
97
98 assert!(matches!(
99 err,
100 StoreError::UnsupportedRecordSchemaVersion {
101 record_kind: "SessionHeadMeta",
102 actual: 2,
103 expected: SESSION_HEAD_META_SCHEMA_VERSION
104 }
105 ));
106 }
107}
108
109#[derive(Debug, thiserror::Error)]
110pub enum StoreError {
111 #[error(
112 "store is already bound to session `{bound_session_id}` and cannot be reused for `{attempted_session_id}`"
113 )]
114 SessionBindingMismatch {
115 bound_session_id: String,
116 attempted_session_id: String,
117 },
118 #[error("store does not support read scope {0:?}")]
119 UnsupportedReadScope(SessionReadScope),
120 #[error("store head revision conflict: expected {expected:?}, actual {actual}")]
121 HeadRevisionConflict { expected: Option<u64>, actual: u64 },
122 #[error(
123 "runtime turn `{turn_id}` for session `{session_id}` was already committed with a different commit hash"
124 )]
125 RuntimeTurnCommitConflict { session_id: String, turn_id: String },
126 #[error(
127 "queued work claim `{claim_id}` for session `{session_id}` is superseded by a newer session-lease generation"
128 )]
129 QueuedWorkClaimSuperseded {
130 session_id: String,
131 claim_id: String,
132 },
133 #[error(
134 "turn input claim `{claim_id}` for session `{session_id}` is superseded by a newer session-lease generation"
135 )]
136 TurnInputClaimSuperseded {
137 session_id: String,
138 claim_id: String,
139 },
140 #[error(
141 "runtime commit for session `{session_id}` includes queued-work-derived content without settling claim `{claim_id}`"
142 )]
143 UnsettledQueuedWorkClaim {
144 session_id: String,
145 claim_id: String,
146 },
147 #[error(
148 "runtime commit for session `{session_id}` includes turn-input-derived content without settling claim `{claim_id}`"
149 )]
150 UnsettledTurnInputClaim {
151 session_id: String,
152 claim_id: String,
153 },
154 #[error(
155 "pending turn input source_key `{source_key}` for session `{session_id}` is already bound to input `{existing_input_id}` with different submitted content"
156 )]
157 PendingTurnInputSourceKeyConflict {
158 session_id: String,
159 source_key: String,
160 existing_input_id: String,
161 },
162 #[error("session execution lease for session `{session_id}` is missing or expired")]
163 SessionExecutionLeaseExpired { session_id: String },
164 #[error(
165 "{record_kind} schema_version {actual} is not supported by this binary (expected {expected})"
166 )]
167 UnsupportedRecordSchemaVersion {
168 record_kind: &'static str,
169 actual: u32,
170 expected: u32,
171 },
172 #[error(
173 "{record_kind} is missing schema_version and was written by unsupported pre-versioned state (expected {expected})"
174 )]
175 MissingRecordSchemaVersion {
176 record_kind: &'static str,
177 expected: u32,
178 },
179 #[error("{record_kind} schema_version {actual} is invalid (expected integer {expected})")]
180 InvalidRecordSchemaVersion {
181 record_kind: &'static str,
182 actual: String,
183 expected: u32,
184 },
185 #[error("store backend error: {0}")]
186 Backend(String),
187}
188
189#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
190pub struct SessionMeta {
191 pub session_id: String,
192 pub session_name: String,
193 pub created_at: String,
194 pub model: String,
195 pub cwd: Option<String>,
196 pub relation: crate::SessionRelation,
197}
198
199impl SessionMeta {
200 pub fn parent_session_id(&self) -> Option<&str> {
203 self.relation.parent_session_id()
204 }
205}
206
207#[derive(Clone, Debug)]
209pub struct SessionPickerInfo {
210 pub session_id: String,
211 pub cwd: Option<String>,
212 pub relation: crate::SessionRelation,
213 pub first_user_message: String,
214 pub user_message_count: usize,
215}
216
217impl SessionPickerInfo {
218 pub fn parent_session_id(&self) -> Option<&str> {
219 self.relation.parent_session_id()
220 }
221}
222
223#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
224#[serde(transparent)]
225pub struct BlobRef(pub String);
226
227impl BlobRef {
228 pub fn as_str(&self) -> &str {
229 &self.0
230 }
231}
232
233impl std::fmt::Display for BlobRef {
234 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
235 f.write_str(&self.0)
236 }
237}
238
239impl From<String> for BlobRef {
240 fn from(value: String) -> Self {
241 Self(value)
242 }
243}
244
245#[derive(Clone, Debug, Default, PartialEq, Eq)]
246pub struct GcReport {
247 pub root_count: usize,
248 pub retained_blob_count: usize,
249 pub deleted_blob_count: usize,
250}
251
252#[derive(Clone, Debug, Default, PartialEq, Eq)]
258pub struct VacuumReport {
259 pub removed_node_count: usize,
260 pub removed_pending_turn_input_tombstone_count: usize,
261}
262
263#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
264pub struct SessionCheckpoint {
265 pub schema_version: u32,
266 pub turn_state: crate::PersistedTurnState,
267 #[serde(default, skip_serializing_if = "Option::is_none")]
268 pub tool_state_ref: Option<BlobRef>,
269 #[serde(default, skip_serializing_if = "Option::is_none")]
270 pub plugin_snapshot_ref: Option<BlobRef>,
271 #[serde(default, skip_serializing_if = "Option::is_none")]
272 pub plugin_snapshot_revision: Option<u64>,
273 #[serde(default, skip_serializing_if = "Option::is_none")]
274 pub execution_state_ref: Option<BlobRef>,
275}
276
277impl Default for SessionCheckpoint {
278 fn default() -> Self {
279 Self {
280 schema_version: SESSION_CHECKPOINT_SCHEMA_VERSION,
281 turn_state: crate::PersistedTurnState::default(),
282 tool_state_ref: None,
283 plugin_snapshot_ref: None,
284 plugin_snapshot_revision: None,
285 execution_state_ref: None,
286 }
287 }
288}
289
290impl SessionCheckpoint {
291 pub fn new(
292 turn_state: crate::PersistedTurnState,
293 tool_state_ref: Option<BlobRef>,
294 plugin_snapshot_ref: Option<BlobRef>,
295 plugin_snapshot_revision: Option<u64>,
296 execution_state_ref: Option<BlobRef>,
297 ) -> Self {
298 Self {
299 schema_version: SESSION_CHECKPOINT_SCHEMA_VERSION,
300 turn_state,
301 tool_state_ref,
302 plugin_snapshot_ref,
303 plugin_snapshot_revision,
304 execution_state_ref,
305 }
306 }
307}
308
309#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
310pub struct HydratedSessionCheckpoint {
311 pub turn_state: crate::PersistedTurnState,
312 pub tool_state_ref: Option<BlobRef>,
313 pub tool_state: Option<crate::ToolState>,
314 pub plugin_snapshot_ref: Option<BlobRef>,
315 pub plugin_snapshot: Option<crate::PluginSessionSnapshot>,
316 pub plugin_snapshot_revision: Option<u64>,
317 pub execution_state_ref: Option<BlobRef>,
318 pub execution_state: Option<Vec<u8>>,
319}
320
321#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
322pub struct SessionHead {
323 #[serde(default = "default_root_session_id")]
324 pub session_id: String,
325 #[serde(default)]
326 pub head_revision: u64,
327 #[serde(default)]
328 pub agent_frames: Vec<crate::AgentFrameRecord>,
329 #[serde(default, skip_serializing_if = "String::is_empty")]
330 pub current_agent_frame_id: crate::AgentFrameId,
331 pub graph: crate::SessionGraph,
332 pub config: crate::PersistedSessionConfig,
333 #[serde(default, skip_serializing_if = "Option::is_none")]
334 pub checkpoint_ref: Option<BlobRef>,
335 #[serde(default, skip_serializing_if = "Vec::is_empty")]
336 pub token_ledger: Vec<crate::TokenLedgerEntry>,
337}
338
339#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
340pub struct SessionHeadMeta {
341 pub schema_version: u32,
342 #[serde(default = "default_root_session_id")]
343 pub session_id: String,
344 #[serde(default)]
345 pub head_revision: u64,
346 pub config: crate::PersistedSessionConfig,
347 #[serde(default)]
348 pub agent_frames: Vec<crate::AgentFrameRecord>,
349 #[serde(default, skip_serializing_if = "String::is_empty")]
350 pub current_agent_frame_id: crate::AgentFrameId,
351 #[serde(default, skip_serializing_if = "Option::is_none")]
352 pub checkpoint_ref: Option<BlobRef>,
353 #[serde(default, skip_serializing_if = "Option::is_none")]
354 pub leaf_node_id: Option<String>,
355 #[serde(default)]
356 pub graph_node_count: usize,
357 #[serde(default, skip_serializing_if = "Vec::is_empty")]
358 pub token_ledger: Vec<crate::TokenLedgerEntry>,
359}
360
361fn persisted_session_config_from_state(
362 state: &crate::RuntimeSessionState,
363) -> crate::PersistedSessionConfig {
364 crate::PersistedSessionConfig {
365 provider_id: state.policy.recorded_provider_id().to_string(),
366 model: state.policy.model.clone(),
367 }
368}
369
370#[derive(Clone, Debug, PartialEq, Eq)]
371pub enum SessionReadScope {
372 FullGraph,
373 ActivePath { leaf_node_id: Option<String> },
374}
375
376#[derive(Clone, Debug)]
377pub struct PersistedSessionRead {
378 pub session_id: String,
379 pub head_revision: u64,
380 pub config: crate::PersistedSessionConfig,
381 pub agent_frames: Vec<crate::AgentFrameRecord>,
382 pub current_agent_frame_id: crate::AgentFrameId,
383 pub graph: crate::SessionGraph,
384 pub checkpoint_ref: Option<BlobRef>,
385 pub checkpoint: Option<HydratedSessionCheckpoint>,
386 pub token_ledger: Vec<crate::TokenLedgerEntry>,
387}
388
389#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
390pub enum GraphCommitDelta {
391 Unchanged {
392 leaf_node_id: Option<String>,
393 },
394 Append {
395 nodes: Vec<crate::SessionNodeRecord>,
396 leaf_node_id: Option<String>,
397 },
398 ReplaceFull(crate::SessionGraph),
399}
400
401impl GraphCommitDelta {
402 pub fn leaf_node_id(&self) -> Option<&String> {
403 match self {
404 Self::Unchanged { leaf_node_id } | Self::Append { leaf_node_id, .. } => {
405 leaf_node_id.as_ref()
406 }
407 Self::ReplaceFull(graph) => graph.leaf_node_id.as_ref(),
408 }
409 }
410}
411
412#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
413pub struct RuntimeCommit {
414 pub session_id: String,
415 pub expected_head_revision: Option<u64>,
416 #[serde(default, skip_serializing_if = "Option::is_none")]
417 pub session_execution_lease: Option<SessionExecutionLeaseFence>,
418 #[serde(default, skip_serializing_if = "Option::is_none")]
419 pub release_session_execution_lease: Option<SessionExecutionLeaseCompletion>,
420 pub config: crate::PersistedSessionConfig,
421 pub agent_frames: Vec<crate::AgentFrameRecord>,
422 pub current_agent_frame_id: crate::AgentFrameId,
423 pub graph: GraphCommitDelta,
424 pub checkpoint: HydratedSessionCheckpoint,
425 pub usage_deltas: Vec<crate::TokenLedgerEntry>,
426 pub turn_commit: Option<RuntimeTurnCommitStamp>,
427 pub completed_queue_claims: Vec<crate::QueuedWorkCompletion>,
428 pub completed_turn_input_claims: Vec<crate::TurnInputCompletion>,
429 #[serde(default, skip_serializing_if = "Vec::is_empty")]
430 pub enqueued_queue_batches: Vec<crate::QueuedWorkBatchDraft>,
431 #[serde(default, skip_serializing_if = "Option::is_none")]
432 pub interrupted_turn_input_turn_id: Option<String>,
433 pub committed_attachment_ids: Vec<crate::AttachmentId>,
441}
442
443#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
444pub struct RuntimeCommitResult {
445 pub head_revision: u64,
446 pub checkpoint_ref: BlobRef,
447 pub manifest: SessionCheckpoint,
448 #[serde(default, skip_serializing_if = "Vec::is_empty")]
449 pub enqueued_queue_batches: Vec<crate::QueuedWorkBatch>,
450}
451
452#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
459pub struct LeaseOwnerIdentity {
460 pub owner_id: String,
461 pub incarnation_id: String,
462 #[serde(default)]
463 pub liveness: LeaseOwnerLiveness,
464}
465
466impl LeaseOwnerIdentity {
467 pub fn opaque(
468 owner_id: impl Into<String>,
469 incarnation_id: impl Into<String>,
470 ) -> LeaseOwnerIdentity {
471 LeaseOwnerIdentity {
472 owner_id: owner_id.into(),
473 incarnation_id: incarnation_id.into(),
474 liveness: LeaseOwnerLiveness::Opaque,
475 }
476 }
477
478 pub fn local_process(
479 owner_id: impl Into<String>,
480 incarnation_id: impl Into<String>,
481 host_id: impl Into<String>,
482 ) -> LeaseOwnerIdentity {
483 let liveness = LeaseOwnerLiveness::current_local_process(host_id.into())
484 .unwrap_or(LeaseOwnerLiveness::Opaque);
485 LeaseOwnerIdentity {
486 owner_id: owner_id.into(),
487 incarnation_id: incarnation_id.into(),
488 liveness,
489 }
490 }
491
492 pub fn same_incarnation(&self, other: &LeaseOwnerIdentity) -> bool {
493 self.owner_id == other.owner_id && self.incarnation_id == other.incarnation_id
494 }
495
496 pub fn is_definitely_dead_for_claimant(&self, claimant: &LeaseOwnerIdentity) -> bool {
497 self.liveness
498 .is_definitely_dead_for_claimant(&claimant.liveness)
499 }
500}
501
502#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Default)]
503#[serde(tag = "kind", rename_all = "snake_case")]
504pub enum LeaseOwnerLiveness {
505 LocalProcess {
506 host_id: String,
507 boot_id: String,
508 pid: u32,
509 process_start: String,
510 },
511 #[default]
512 Opaque,
513}
514
515impl LeaseOwnerLiveness {
516 pub fn current_local_process(host_id: impl Into<String>) -> Option<LeaseOwnerLiveness> {
517 let boot_id = std::fs::read_to_string(PROC_BOOT_ID_PATH)
518 .ok()
519 .map(|value| value.trim().to_string())
520 .filter(|value| !value.is_empty())?;
521 let pid = std::process::id();
522 let process_start = read_linux_process_start(pid)?;
523 Some(LeaseOwnerLiveness::LocalProcess {
524 host_id: host_id.into(),
525 boot_id,
526 pid,
527 process_start,
528 })
529 }
530
531 pub fn local_process_for_test(
532 host_id: impl Into<String>,
533 boot_id: impl Into<String>,
534 pid: u32,
535 process_start: impl Into<String>,
536 ) -> LeaseOwnerLiveness {
537 LeaseOwnerLiveness::LocalProcess {
538 host_id: host_id.into(),
539 boot_id: boot_id.into(),
540 pid,
541 process_start: process_start.into(),
542 }
543 }
544
545 pub fn is_definitely_dead_for_claimant(&self, claimant: &LeaseOwnerLiveness) -> bool {
546 let (
547 LeaseOwnerLiveness::LocalProcess {
548 host_id,
549 boot_id,
550 pid,
551 process_start,
552 },
553 LeaseOwnerLiveness::LocalProcess {
554 host_id: claimant_host_id,
555 boot_id: claimant_boot_id,
556 ..
557 },
558 ) = (self, claimant)
559 else {
560 return false;
561 };
562 if host_id != claimant_host_id || boot_id != claimant_boot_id {
563 return false;
564 }
565 matches!(linux_process_is_live(*pid, process_start), Some(false))
566 }
567}
568
569fn read_linux_process_start(pid: u32) -> Option<String> {
570 let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;
571 parse_linux_process_start(&stat)
572}
573
574fn linux_process_is_live(pid: u32, expected_process_start: &str) -> Option<bool> {
575 match std::fs::read_to_string(format!("/proc/{pid}/stat")) {
576 Ok(stat) => parse_linux_process_start(&stat).map(|start| start == expected_process_start),
577 Err(err) if err.kind() == std::io::ErrorKind::NotFound => Some(false),
578 Err(_) => None,
579 }
580}
581
582fn parse_linux_process_start(stat: &str) -> Option<String> {
583 let after_comm = stat.rsplit_once(") ")?.1;
584 after_comm.split_whitespace().nth(19).map(ToOwned::to_owned)
585}
586
587#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
588pub struct SessionExecutionLease {
589 pub session_id: String,
590 pub owner: LeaseOwnerIdentity,
591 pub lease_token: String,
592 pub fencing_token: u64,
593 pub claimed_at_epoch_ms: u64,
594 pub expires_at_epoch_ms: u64,
595}
596
597#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
598pub struct SessionExecutionLeaseFence {
599 pub session_id: String,
600 pub owner: LeaseOwnerIdentity,
601 pub lease_token: String,
602 pub fencing_token: u64,
603}
604
605#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
606pub struct SessionExecutionLeaseCompletion {
607 pub session_id: String,
608 pub owner: LeaseOwnerIdentity,
609 pub lease_token: String,
610 pub fencing_token: u64,
611}
612
613impl SessionExecutionLease {
614 pub fn fence(&self) -> SessionExecutionLeaseFence {
615 SessionExecutionLeaseFence {
616 session_id: self.session_id.clone(),
617 owner: self.owner.clone(),
618 lease_token: self.lease_token.clone(),
619 fencing_token: self.fencing_token,
620 }
621 }
622
623 pub fn completion(&self) -> SessionExecutionLeaseCompletion {
624 SessionExecutionLeaseCompletion {
625 session_id: self.session_id.clone(),
626 owner: self.owner.clone(),
627 lease_token: self.lease_token.clone(),
628 fencing_token: self.fencing_token,
629 }
630 }
631}
632
633impl SessionExecutionLeaseCompletion {
634 pub fn from_lease(lease: &SessionExecutionLease) -> Self {
635 lease.completion()
636 }
637}
638
639#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
640pub enum SessionExecutionLeaseClaimOutcome {
641 Acquired(SessionExecutionLease),
642 Busy { holder: SessionExecutionLease },
643}
644
645impl SessionExecutionLeaseClaimOutcome {
646 pub fn acquired(self) -> Option<SessionExecutionLease> {
647 match self {
648 Self::Acquired(lease) => Some(lease),
649 Self::Busy { .. } => None,
650 }
651 }
652}
653
654pub fn ensure_supported_schema_version(
658 record_kind: &'static str,
659 actual: u32,
660 expected: u32,
661) -> Result<(), StoreError> {
662 if actual == expected {
663 Ok(())
664 } else {
665 Err(StoreError::UnsupportedRecordSchemaVersion {
666 record_kind,
667 actual,
668 expected,
669 })
670 }
671}
672
673pub fn ensure_supported_record_schema_version(
674 record_kind: &'static str,
675 value: &serde_json::Value,
676 expected: u32,
677) -> Result<(), StoreError> {
678 let Some(schema_version) = value.get("schema_version") else {
679 return Err(StoreError::MissingRecordSchemaVersion {
680 record_kind,
681 expected,
682 });
683 };
684 let Some(actual) = schema_version
685 .as_u64()
686 .and_then(|version| u32::try_from(version).ok())
687 else {
688 return Err(StoreError::InvalidRecordSchemaVersion {
689 record_kind,
690 actual: schema_version.to_string(),
691 expected,
692 });
693 };
694 ensure_supported_schema_version(record_kind, actual, expected)
695}
696
697pub fn decode_versioned_json_record<T>(
698 json: &str,
699 record_kind: &'static str,
700 expected: u32,
701) -> Result<T, StoreError>
702where
703 T: serde::de::DeserializeOwned,
704{
705 let value: serde_json::Value = serde_json::from_str(json)
706 .map_err(|err| StoreError::Backend(format!("failed to decode {record_kind}: {err}")))?;
707 ensure_supported_record_schema_version(record_kind, &value, expected)?;
708 serde_json::from_value(value)
709 .map_err(|err| StoreError::Backend(format!("failed to decode {record_kind}: {err}")))
710}
711
712#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
713pub struct RuntimeTurnCommitStamp {
714 pub session_id: String,
715 pub turn_id: String,
716 pub turn_commit_hash: String,
717}
718
719impl RuntimeTurnCommitStamp {
720 pub fn new(
721 session_id: impl Into<String>,
722 turn_id: impl Into<String>,
723 turn_commit_hash: impl Into<String>,
724 ) -> Self {
725 Self {
726 session_id: session_id.into(),
727 turn_id: turn_id.into(),
728 turn_commit_hash: turn_commit_hash.into(),
729 }
730 }
731}
732
733fn build_persisted_turn_state(state: &crate::RuntimeSessionState) -> crate::PersistedTurnState {
734 crate::PersistedTurnState {
735 turn_index: state.turn_index,
736 token_usage: state.token_usage.clone(),
737 last_prompt_usage: state.last_prompt_usage.clone(),
738 protocol_turn_options: state.protocol_turn_options.clone(),
739 }
740}
741
742fn build_checkpoint_from_persisted_state(
743 state: &crate::RuntimeSessionState,
744) -> HydratedSessionCheckpoint {
745 HydratedSessionCheckpoint {
746 turn_state: build_persisted_turn_state(state),
747 tool_state_ref: state.tool_state_ref.clone(),
748 tool_state: state.tool_state_snapshot.clone(),
749 plugin_snapshot_ref: state.plugin_snapshot_ref.clone(),
750 plugin_snapshot_revision: state.plugin_snapshot_revision,
751 plugin_snapshot: state.plugin_snapshot.clone(),
752 execution_state_ref: state.execution_state_ref.clone(),
753 execution_state: state.execution_state_snapshot.clone(),
754 }
755}
756
757impl RuntimeCommit {
758 pub(crate) fn validate_claim_settlement(
759 &self,
760 originating_queue_claims: &[crate::QueuedWorkCompletion],
761 originating_turn_input_claims: &[crate::TurnInputCompletion],
762 ) -> Result<(), StoreError> {
763 for originating in originating_queue_claims {
764 if !self.completed_queue_claims.iter().any(|completed| {
765 completed.session_id == originating.session_id
766 && completed.claim_id == originating.claim_id
767 }) {
768 return Err(StoreError::UnsettledQueuedWorkClaim {
769 session_id: originating.session_id.clone(),
770 claim_id: originating.claim_id.clone(),
771 });
772 }
773 }
774 for originating in originating_turn_input_claims {
775 if !self.completed_turn_input_claims.iter().any(|completed| {
776 completed.session_id == originating.session_id
777 && completed.claim_id == originating.claim_id
778 }) {
779 return Err(StoreError::UnsettledTurnInputClaim {
780 session_id: originating.session_id.clone(),
781 claim_id: originating.claim_id.clone(),
782 });
783 }
784 }
785 Ok(())
786 }
787
788 pub fn turn_commit_hash(&self) -> Result<String, StoreError> {
789 let mut semantic_commit = self.clone();
790 semantic_commit.expected_head_revision = None;
791 semantic_commit.session_execution_lease = None;
792 semantic_commit.release_session_execution_lease = None;
793 semantic_commit.turn_commit = None;
794 let mut semantic_commit = serde_json::to_value(&semantic_commit).map_err(|err| {
795 StoreError::Backend(format!("failed to serialize runtime turn commit: {err}"))
796 })?;
797 scrub_turn_commit_hash_value(&mut semantic_commit);
798 crate::stable_hash::stable_json_sha256_hex(&semantic_commit).map_err(|err| {
799 StoreError::Backend(format!(
800 "failed to serialize runtime turn commit hash: {err}"
801 ))
802 })
803 }
804
805 pub fn persisted_state(
806 state: &crate::RuntimeSessionState,
807 usage_deltas: &[crate::TokenLedgerEntry],
808 ) -> Self {
809 Self {
810 session_id: state.session_id.clone(),
811 expected_head_revision: state.head_revision,
812 session_execution_lease: None,
813 release_session_execution_lease: None,
814 config: persisted_session_config_from_state(state),
815 agent_frames: state.agent_frames.clone(),
816 current_agent_frame_id: state.current_agent_frame_id.clone(),
817 graph: if state.graph_replace_required || state.head_revision.is_none() {
818 GraphCommitDelta::ReplaceFull(state.session_graph.clone())
819 } else {
820 GraphCommitDelta::Unchanged {
821 leaf_node_id: state.session_graph.leaf_node_id.clone(),
822 }
823 },
824 checkpoint: build_checkpoint_from_persisted_state(state),
825 usage_deltas: usage_deltas.to_vec(),
826 turn_commit: None,
827 completed_queue_claims: Vec::new(),
828 completed_turn_input_claims: Vec::new(),
829 enqueued_queue_batches: Vec::new(),
830 interrupted_turn_input_turn_id: None,
831 committed_attachment_ids: Vec::new(),
832 }
833 }
834
835 pub(crate) fn persisted_state_with_graph_commit(
836 state: &crate::RuntimeSessionState,
837 graph: GraphCommitDelta,
838 usage_deltas: &[crate::TokenLedgerEntry],
839 ) -> Self {
840 Self {
841 session_id: state.session_id.clone(),
842 expected_head_revision: state.head_revision,
843 session_execution_lease: None,
844 release_session_execution_lease: None,
845 config: persisted_session_config_from_state(state),
846 agent_frames: state.agent_frames.clone(),
847 current_agent_frame_id: state.current_agent_frame_id.clone(),
848 graph,
849 checkpoint: build_checkpoint_from_persisted_state(state),
850 usage_deltas: usage_deltas.to_vec(),
851 turn_commit: None,
852 completed_queue_claims: Vec::new(),
853 completed_turn_input_claims: Vec::new(),
854 enqueued_queue_batches: Vec::new(),
855 interrupted_turn_input_turn_id: None,
856 committed_attachment_ids: Vec::new(),
857 }
858 }
859
860 pub fn with_turn_commit(mut self, turn_commit: RuntimeTurnCommitStamp) -> Self {
861 self.turn_commit = Some(turn_commit);
862 self
863 }
864
865 pub fn with_session_execution_lease(mut self, lease: SessionExecutionLeaseFence) -> Self {
866 self.session_execution_lease = Some(lease);
867 self
868 }
869
870 pub fn releasing_session_execution_lease(
871 mut self,
872 completion: SessionExecutionLeaseCompletion,
873 ) -> Self {
874 self.release_session_execution_lease = Some(completion);
875 self
876 }
877
878 pub fn completing_queue_claim(
879 mut self,
880 completed_queue_claim: crate::QueuedWorkCompletion,
881 ) -> Self {
882 self.completed_queue_claims.push(completed_queue_claim);
883 self
884 }
885
886 pub fn completing_queue_claims(
887 mut self,
888 completed_queue_claims: impl IntoIterator<Item = crate::QueuedWorkCompletion>,
889 ) -> Self {
890 self.completed_queue_claims.extend(completed_queue_claims);
891 self
892 }
893
894 pub fn completing_turn_input_claim(
895 mut self,
896 completed_turn_input_claim: crate::TurnInputCompletion,
897 ) -> Self {
898 self.completed_turn_input_claims
899 .push(completed_turn_input_claim);
900 self
901 }
902
903 pub fn completing_turn_input_claims(
904 mut self,
905 completed_turn_input_claims: impl IntoIterator<Item = crate::TurnInputCompletion>,
906 ) -> Self {
907 self.completed_turn_input_claims
908 .extend(completed_turn_input_claims);
909 self
910 }
911
912 pub fn deferring_interrupted_turn_inputs(mut self, turn_id: impl Into<String>) -> Self {
913 self.interrupted_turn_input_turn_id = Some(turn_id.into());
914 self
915 }
916
917 pub fn with_committed_attachments(
918 mut self,
919 attachment_ids: impl IntoIterator<Item = crate::AttachmentId>,
920 ) -> Self {
921 self.committed_attachment_ids = attachment_ids.into_iter().collect();
922 self
923 }
924}
925
926fn scrub_turn_commit_hash_value(value: &mut serde_json::Value) {
927 match value {
928 serde_json::Value::Object(map) => {
929 let is_message = map.contains_key("role") && map.contains_key("parts");
930 let is_message_part = map.contains_key("kind")
931 && map.contains_key("content")
932 && map.contains_key("prune_state");
933 if is_message || is_message_part {
934 map.remove("id");
935 }
936 for volatile_key in ["node_id", "parent_node_id", "leaf_node_id", "timestamp"] {
937 map.remove(volatile_key);
938 }
939 for child in map.values_mut() {
940 scrub_turn_commit_hash_value(child);
941 }
942 }
943 serde_json::Value::Array(items) => {
944 for item in items {
945 scrub_turn_commit_hash_value(item);
946 }
947 }
948 _ => {}
949 }
950}
951
952fn persisted_session_state_from_head(
953 head: SessionHead,
954 checkpoint: Option<HydratedSessionCheckpoint>,
955) -> crate::RuntimeSessionState {
956 let mut state = crate::RuntimeSessionState {
957 session_id: head.session_id,
958 policy: crate::SessionPolicy::default(),
959 agent_frames: head.agent_frames,
960 current_agent_frame_id: head.current_agent_frame_id,
961 session_graph: head.graph,
962 turn_index: 0,
963 token_usage: crate::TokenUsage::default(),
964 last_prompt_usage: None,
965 protocol_turn_options: crate::ProtocolTurnOptions::default(),
966 tool_state_ref: None,
967 tool_state_generation: None,
968 tool_state_snapshot: None,
969 plugin_snapshot_ref: None,
970 plugin_snapshot_revision: None,
971 plugin_snapshot: None,
972 execution_state_ref: None,
973 execution_state_snapshot: None,
974 token_ledger: head.token_ledger,
975 checkpoint_ref: head.checkpoint_ref.clone(),
976 head_revision: Some(head.head_revision),
977 graph_replace_required: false,
978 };
979 state.policy.model = head.config.model.clone();
980 state.policy.provider_id = head.config.provider_id.clone();
981 if let Some(checkpoint) = checkpoint {
982 state.turn_index = checkpoint.turn_state.turn_index;
983 state.token_usage = checkpoint.turn_state.token_usage;
984 state.last_prompt_usage = checkpoint.turn_state.last_prompt_usage;
985 state.protocol_turn_options = checkpoint.turn_state.protocol_turn_options;
986 state.tool_state_ref = checkpoint.tool_state_ref.clone();
987 state.tool_state_generation = checkpoint
988 .tool_state
989 .as_ref()
990 .map(|snapshot| snapshot.generation());
991 state.tool_state_snapshot = checkpoint.tool_state;
992 state.plugin_snapshot_ref = checkpoint.plugin_snapshot_ref.clone();
993 state.plugin_snapshot_revision = checkpoint.plugin_snapshot_revision;
994 state.plugin_snapshot = checkpoint.plugin_snapshot;
995 state.execution_state_ref = checkpoint.execution_state_ref.clone();
996 state.execution_state_snapshot = checkpoint.execution_state;
997 }
998 state.ensure_agent_frame_initialized();
999 state
1000}
1001
1002impl Default for SessionHead {
1003 fn default() -> Self {
1004 Self {
1005 session_id: default_root_session_id(),
1006 head_revision: 0,
1007 agent_frames: Vec::new(),
1008 current_agent_frame_id: String::new(),
1009 graph: crate::SessionGraph::default(),
1010 config: crate::PersistedSessionConfig::default(),
1011 checkpoint_ref: None,
1012 token_ledger: Vec::new(),
1013 }
1014 }
1015}
1016
1017impl Default for SessionHeadMeta {
1018 fn default() -> Self {
1019 Self {
1020 schema_version: SESSION_HEAD_META_SCHEMA_VERSION,
1021 session_id: default_root_session_id(),
1022 head_revision: 0,
1023 config: crate::PersistedSessionConfig::default(),
1024 agent_frames: Vec::new(),
1025 current_agent_frame_id: String::new(),
1026 checkpoint_ref: None,
1027 leaf_node_id: None,
1028 graph_node_count: 0,
1029 token_ledger: Vec::new(),
1030 }
1031 }
1032}
1033
1034#[async_trait::async_trait]
1051pub trait SessionCommitStore: AttachmentManifest + Send + Sync {
1052 fn durability_tier(&self) -> crate::DurabilityTier {
1055 crate::DurabilityTier::Inline
1056 }
1057
1058 async fn load_session(
1059 &self,
1060 scope: SessionReadScope,
1061 ) -> Result<Option<PersistedSessionRead>, StoreError>;
1062
1063 async fn load_node(
1064 &self,
1065 node_id: &str,
1066 ) -> Result<Option<crate::SessionNodeRecord>, StoreError>;
1067
1068 async fn commit_runtime_state(
1069 &self,
1070 commit: RuntimeCommit,
1071 ) -> Result<RuntimeCommitResult, StoreError>;
1072
1073 async fn save_session_meta(&self, meta: SessionMeta) -> Result<(), StoreError>;
1074 async fn load_session_meta(&self) -> Result<Option<SessionMeta>, StoreError>;
1075}
1076
1077#[async_trait::async_trait]
1085pub trait TurnInputStore: Send + Sync {
1086 async fn enqueue_pending_turn_input(
1088 &self,
1089 input: crate::PendingTurnInputDraft,
1090 ) -> Result<crate::PendingTurnInput, StoreError>;
1091
1092 async fn list_pending_turn_inputs(
1097 &self,
1098 session_id: &str,
1099 ) -> Result<Vec<crate::PendingTurnInput>, StoreError>;
1100
1101 async fn cancel_pending_turn_input(
1108 &self,
1109 session_id: &str,
1110 input_id: &str,
1111 ) -> Result<crate::PendingTurnInputCancelOutcome, StoreError> {
1112 let target = crate::PendingTurnInputCancelTarget::input_id(input_id);
1113 let targets = vec![target];
1114 let mut outcomes = self
1115 .cancel_pending_turn_inputs(session_id, &targets)
1116 .await?;
1117 Ok(outcomes
1118 .pop()
1119 .map(|result| result.outcome)
1120 .unwrap_or(crate::PendingTurnInputCancelOutcome::NotFound))
1121 }
1122
1123 async fn cancel_pending_turn_inputs(
1125 &self,
1126 session_id: &str,
1127 targets: &[crate::PendingTurnInputCancelTarget],
1128 ) -> Result<Vec<crate::PendingTurnInputCancelResult>, StoreError>;
1129
1130 async fn cancel_pending_turn_input_suffix(
1132 &self,
1133 session_id: &str,
1134 anchor: &crate::PendingTurnInputCancelTarget,
1135 ) -> Result<crate::PendingTurnInputSuffixCancelOutcome, StoreError>;
1136
1137 async fn claim_active_turn_inputs(
1143 &self,
1144 session_id: &str,
1145 session_execution_lease: &SessionExecutionLeaseFence,
1146 owner: &LeaseOwnerIdentity,
1147 turn_id: &str,
1148 checkpoint: crate::CheckpointKind,
1149 max_inputs: usize,
1150 ) -> Result<Option<crate::TurnInputClaim>, StoreError>;
1151
1152 async fn claim_next_turn_inputs(
1154 &self,
1155 session_id: &str,
1156 session_execution_lease: &SessionExecutionLeaseFence,
1157 owner: &LeaseOwnerIdentity,
1158 max_inputs: usize,
1159 ) -> Result<Option<crate::TurnInputClaim>, StoreError>;
1160
1161 async fn abandon_turn_input_claim(
1163 &self,
1164 claim: &crate::TurnInputClaim,
1165 ) -> Result<(), StoreError>;
1166}
1167
1168#[async_trait::async_trait]
1171pub trait SessionExecutionLeaseStore: Send + Sync {
1172 async fn try_claim_session_execution_lease(
1179 &self,
1180 session_id: &str,
1181 owner: &LeaseOwnerIdentity,
1182 lease_ttl_ms: u64,
1183 ) -> Result<SessionExecutionLeaseClaimOutcome, StoreError>;
1184
1185 async fn reclaim_session_execution_lease(
1191 &self,
1192 session_id: &str,
1193 owner: &LeaseOwnerIdentity,
1194 observed_holder: &SessionExecutionLeaseFence,
1195 lease_ttl_ms: u64,
1196 ) -> Result<SessionExecutionLeaseClaimOutcome, StoreError>;
1197
1198 async fn renew_session_execution_lease(
1203 &self,
1204 fence: &SessionExecutionLeaseFence,
1205 lease_ttl_ms: u64,
1206 ) -> Result<SessionExecutionLease, StoreError>;
1207
1208 async fn release_session_execution_lease(
1212 &self,
1213 completion: &SessionExecutionLeaseCompletion,
1214 ) -> Result<(), StoreError>;
1215}
1216
1217#[async_trait::async_trait]
1223pub trait QueuedWorkStore: Send + Sync {
1224 async fn enqueue_queued_work(
1226 &self,
1227 batch: crate::QueuedWorkBatchDraft,
1228 ) -> Result<crate::QueuedWorkBatch, StoreError>;
1229
1230 async fn claim_leading_ready_session_command(
1237 &self,
1238 session_id: &str,
1239 session_execution_lease: &SessionExecutionLeaseFence,
1240 owner: &LeaseOwnerIdentity,
1241 ) -> Result<Option<crate::QueuedWorkClaim>, StoreError>;
1242
1243 async fn claim_ready_queued_work(
1250 &self,
1251 session_id: &str,
1252 session_execution_lease: &SessionExecutionLeaseFence,
1253 owner: &LeaseOwnerIdentity,
1254 boundary: crate::QueuedWorkClaimBoundary,
1255 max_batches: usize,
1256 ) -> Result<Option<crate::QueuedWorkClaim>, StoreError>;
1257
1258 async fn claim_ready_queued_work_by_batch_ids(
1269 &self,
1270 session_id: &str,
1271 session_execution_lease: &SessionExecutionLeaseFence,
1272 owner: &LeaseOwnerIdentity,
1273 boundary: crate::QueuedWorkClaimBoundary,
1274 batch_ids: &[String],
1275 ) -> Result<Option<crate::QueuedWorkClaim>, StoreError>;
1276
1277 async fn abandon_queued_work_claim(
1279 &self,
1280 claim: &crate::QueuedWorkClaim,
1281 ) -> Result<(), StoreError>;
1282
1283 async fn cancel_queued_work_batch(
1290 &self,
1291 session_id: &str,
1292 batch_id: &str,
1293 ) -> Result<Option<crate::QueuedWorkBatch>, StoreError>;
1294
1295 async fn list_queued_work(
1298 &self,
1299 session_id: &str,
1300 ) -> Result<Vec<crate::QueuedWorkBatch>, StoreError>;
1301
1302 async fn list_pending_queued_work(
1315 &self,
1316 session_id: &str,
1317 ) -> Result<Vec<crate::QueuedWorkBatch>, StoreError>;
1318}
1319
1320#[async_trait::async_trait]
1323pub trait StoreMaintenance: Send + Sync {
1324 async fn tombstone_nodes(&self, ids: &[String]) -> Result<(), StoreError>;
1327
1328 async fn vacuum(&self) -> Result<VacuumReport, StoreError>;
1331
1332 async fn gc_unreachable(&self) -> Result<GcReport, StoreError>;
1334}
1335
1336pub trait RuntimePersistence:
1353 SessionCommitStore
1354 + TurnInputStore
1355 + SessionExecutionLeaseStore
1356 + QueuedWorkStore
1357 + StoreMaintenance
1358{
1359}
1360
1361impl<T> RuntimePersistence for T where
1362 T: SessionCommitStore
1363 + TurnInputStore
1364 + SessionExecutionLeaseStore
1365 + QueuedWorkStore
1366 + StoreMaintenance
1367 + ?Sized
1368{
1369}
1370
1371fn persisted_session_state_from_read(read: PersistedSessionRead) -> crate::RuntimeSessionState {
1372 persisted_session_state_from_head(
1373 SessionHead {
1374 session_id: read.session_id,
1375 head_revision: read.head_revision,
1376 agent_frames: read.agent_frames,
1377 current_agent_frame_id: read.current_agent_frame_id,
1378 graph: read.graph,
1379 config: read.config,
1380 checkpoint_ref: read.checkpoint_ref,
1381 token_ledger: read.token_ledger,
1382 },
1383 read.checkpoint,
1384 )
1385}
1386
1387pub async fn load_persisted_session_state(
1388 store: &(dyn RuntimePersistence + '_),
1389) -> Result<Option<crate::RuntimeSessionState>, StoreError> {
1390 Ok(store
1391 .load_session(SessionReadScope::FullGraph)
1392 .await?
1393 .map(persisted_session_state_from_read))
1394}
1395
1396pub async fn load_persisted_session_state_active_path(
1397 store: &(dyn RuntimePersistence + '_),
1398 leaf_node_id: Option<String>,
1399) -> Result<Option<crate::RuntimeSessionState>, StoreError> {
1400 Ok(store
1401 .load_session(SessionReadScope::ActivePath { leaf_node_id })
1402 .await?
1403 .map(persisted_session_state_from_read))
1404}
1405
1406pub async fn refresh_persisted_session_state(
1407 store: &(dyn RuntimePersistence + '_),
1408 state: &mut crate::RuntimeSessionState,
1409) -> Result<(), StoreError> {
1410 if let Some(mut fresh) = load_persisted_session_state(store).await? {
1411 fresh.policy.session_id = state.policy.session_id.clone();
1412 fresh.policy.max_turns = state.policy.max_turns;
1413 *state = fresh;
1414 }
1415 Ok(())
1416}
1417
1418#[cfg(test)]
1419mod tests {
1420 use super::{LeaseOwnerIdentity, LeaseOwnerLiveness};
1421
1422 fn local_liveness(
1423 host_id: &str,
1424 boot_id: &str,
1425 pid: u32,
1426 process_start: &str,
1427 ) -> LeaseOwnerLiveness {
1428 LeaseOwnerLiveness::local_process_for_test(host_id, boot_id, pid, process_start)
1429 }
1430
1431 #[test]
1432 fn lease_owner_identity_requires_same_incarnation() {
1433 let first = LeaseOwnerIdentity::opaque("owner", "incarnation-a");
1434 let same = LeaseOwnerIdentity::opaque("owner", "incarnation-a");
1435 let next = LeaseOwnerIdentity::opaque("owner", "incarnation-b");
1436
1437 assert!(first.same_incarnation(&same));
1438 assert!(!first.same_incarnation(&next));
1439 }
1440
1441 #[test]
1442 fn local_liveness_only_proves_same_host_boot_dead_processes() {
1443 let holder = local_liveness(
1444 "host-a",
1445 "boot-a",
1446 std::process::id(),
1447 "not-the-current-process-start",
1448 );
1449 let same_host_boot = local_liveness("host-a", "boot-a", std::process::id(), "claimant");
1450 let other_host = local_liveness("host-b", "boot-a", std::process::id(), "claimant");
1451 let other_boot = local_liveness("host-a", "boot-b", std::process::id(), "claimant");
1452
1453 assert!(holder.is_definitely_dead_for_claimant(&same_host_boot));
1454 assert!(!holder.is_definitely_dead_for_claimant(&other_host));
1455 assert!(!holder.is_definitely_dead_for_claimant(&other_boot));
1456 assert!(!holder.is_definitely_dead_for_claimant(&LeaseOwnerLiveness::Opaque));
1457 assert!(!LeaseOwnerLiveness::Opaque.is_definitely_dead_for_claimant(&same_host_boot));
1458 }
1459}