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