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 "pending turn input source_key `{source_key}` for session `{session_id}` is already bound to input `{existing_input_id}` with different submitted content"
142 )]
143 PendingTurnInputSourceKeyConflict {
144 session_id: String,
145 source_key: String,
146 existing_input_id: String,
147 },
148 #[error("session execution lease for session `{session_id}` is missing or expired")]
149 SessionExecutionLeaseExpired { session_id: String },
150 #[error(
151 "{record_kind} schema_version {actual} is not supported by this binary (expected {expected})"
152 )]
153 UnsupportedRecordSchemaVersion {
154 record_kind: &'static str,
155 actual: u32,
156 expected: u32,
157 },
158 #[error(
159 "{record_kind} is missing schema_version and was written by unsupported pre-versioned state (expected {expected})"
160 )]
161 MissingRecordSchemaVersion {
162 record_kind: &'static str,
163 expected: u32,
164 },
165 #[error("{record_kind} schema_version {actual} is invalid (expected integer {expected})")]
166 InvalidRecordSchemaVersion {
167 record_kind: &'static str,
168 actual: String,
169 expected: u32,
170 },
171 #[error("store backend error: {0}")]
172 Backend(String),
173}
174
175#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
176pub struct SessionMeta {
177 pub session_id: String,
178 pub session_name: String,
179 pub created_at: String,
180 pub model: String,
181 pub cwd: Option<String>,
182 pub relation: crate::SessionRelation,
183}
184
185impl SessionMeta {
186 pub fn parent_session_id(&self) -> Option<&str> {
189 self.relation.parent_session_id()
190 }
191}
192
193#[derive(Clone, Debug)]
195pub struct SessionPickerInfo {
196 pub session_id: String,
197 pub cwd: Option<String>,
198 pub relation: crate::SessionRelation,
199 pub first_user_message: String,
200 pub user_message_count: usize,
201}
202
203impl SessionPickerInfo {
204 pub fn parent_session_id(&self) -> Option<&str> {
205 self.relation.parent_session_id()
206 }
207}
208
209#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
210#[serde(transparent)]
211pub struct BlobRef(pub String);
212
213impl BlobRef {
214 pub fn as_str(&self) -> &str {
215 &self.0
216 }
217}
218
219impl std::fmt::Display for BlobRef {
220 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
221 f.write_str(&self.0)
222 }
223}
224
225impl From<String> for BlobRef {
226 fn from(value: String) -> Self {
227 Self(value)
228 }
229}
230
231#[derive(Clone, Debug, Default, PartialEq, Eq)]
232pub struct GcReport {
233 pub root_count: usize,
234 pub retained_blob_count: usize,
235 pub deleted_blob_count: usize,
236}
237
238#[derive(Clone, Debug, Default, PartialEq, Eq)]
244pub struct VacuumReport {
245 pub removed_node_count: usize,
246 pub removed_pending_turn_input_tombstone_count: usize,
247}
248
249#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
250pub struct SessionCheckpoint {
251 pub schema_version: u32,
252 pub turn_state: crate::PersistedTurnState,
253 #[serde(default, skip_serializing_if = "Option::is_none")]
254 pub tool_state_ref: Option<BlobRef>,
255 #[serde(default, skip_serializing_if = "Option::is_none")]
256 pub plugin_snapshot_ref: Option<BlobRef>,
257 #[serde(default, skip_serializing_if = "Option::is_none")]
258 pub plugin_snapshot_revision: Option<u64>,
259 #[serde(default, skip_serializing_if = "Option::is_none")]
260 pub execution_state_ref: Option<BlobRef>,
261}
262
263impl Default for SessionCheckpoint {
264 fn default() -> Self {
265 Self {
266 schema_version: SESSION_CHECKPOINT_SCHEMA_VERSION,
267 turn_state: crate::PersistedTurnState::default(),
268 tool_state_ref: None,
269 plugin_snapshot_ref: None,
270 plugin_snapshot_revision: None,
271 execution_state_ref: None,
272 }
273 }
274}
275
276impl SessionCheckpoint {
277 pub fn new(
278 turn_state: crate::PersistedTurnState,
279 tool_state_ref: Option<BlobRef>,
280 plugin_snapshot_ref: Option<BlobRef>,
281 plugin_snapshot_revision: Option<u64>,
282 execution_state_ref: Option<BlobRef>,
283 ) -> Self {
284 Self {
285 schema_version: SESSION_CHECKPOINT_SCHEMA_VERSION,
286 turn_state,
287 tool_state_ref,
288 plugin_snapshot_ref,
289 plugin_snapshot_revision,
290 execution_state_ref,
291 }
292 }
293}
294
295#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
296pub struct HydratedSessionCheckpoint {
297 pub turn_state: crate::PersistedTurnState,
298 pub tool_state_ref: Option<BlobRef>,
299 pub tool_state: Option<crate::ToolState>,
300 pub plugin_snapshot_ref: Option<BlobRef>,
301 pub plugin_snapshot: Option<crate::PluginSessionSnapshot>,
302 pub plugin_snapshot_revision: Option<u64>,
303 pub execution_state_ref: Option<BlobRef>,
304 pub execution_state: Option<Vec<u8>>,
305}
306
307#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
308pub struct SessionHead {
309 #[serde(default = "default_root_session_id")]
310 pub session_id: String,
311 #[serde(default)]
312 pub head_revision: u64,
313 #[serde(default)]
314 pub agent_frames: Vec<crate::AgentFrameRecord>,
315 #[serde(default, skip_serializing_if = "String::is_empty")]
316 pub current_agent_frame_id: crate::AgentFrameId,
317 pub graph: crate::SessionGraph,
318 pub config: crate::PersistedSessionConfig,
319 #[serde(default, skip_serializing_if = "Option::is_none")]
320 pub checkpoint_ref: Option<BlobRef>,
321 #[serde(default, skip_serializing_if = "Vec::is_empty")]
322 pub token_ledger: Vec<crate::TokenLedgerEntry>,
323}
324
325#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
326pub struct SessionHeadMeta {
327 pub schema_version: u32,
328 #[serde(default = "default_root_session_id")]
329 pub session_id: String,
330 #[serde(default)]
331 pub head_revision: u64,
332 pub config: crate::PersistedSessionConfig,
333 #[serde(default)]
334 pub agent_frames: Vec<crate::AgentFrameRecord>,
335 #[serde(default, skip_serializing_if = "String::is_empty")]
336 pub current_agent_frame_id: crate::AgentFrameId,
337 #[serde(default, skip_serializing_if = "Option::is_none")]
338 pub checkpoint_ref: Option<BlobRef>,
339 #[serde(default, skip_serializing_if = "Option::is_none")]
340 pub leaf_node_id: Option<String>,
341 #[serde(default)]
342 pub graph_node_count: usize,
343 #[serde(default, skip_serializing_if = "Vec::is_empty")]
344 pub token_ledger: Vec<crate::TokenLedgerEntry>,
345}
346
347fn persisted_session_config_from_state(
348 state: &crate::RuntimeSessionState,
349) -> crate::PersistedSessionConfig {
350 crate::PersistedSessionConfig {
351 provider_id: state.policy.recorded_provider_id().to_string(),
352 model: state.policy.model.clone(),
353 }
354}
355
356#[derive(Clone, Debug, PartialEq, Eq)]
357pub enum SessionReadScope {
358 FullGraph,
359 ActivePath { leaf_node_id: Option<String> },
360}
361
362#[derive(Clone, Debug)]
363pub struct PersistedSessionRead {
364 pub session_id: String,
365 pub head_revision: u64,
366 pub config: crate::PersistedSessionConfig,
367 pub agent_frames: Vec<crate::AgentFrameRecord>,
368 pub current_agent_frame_id: crate::AgentFrameId,
369 pub graph: crate::SessionGraph,
370 pub checkpoint_ref: Option<BlobRef>,
371 pub checkpoint: Option<HydratedSessionCheckpoint>,
372 pub token_ledger: Vec<crate::TokenLedgerEntry>,
373}
374
375#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
376pub enum GraphCommitDelta {
377 Unchanged {
378 leaf_node_id: Option<String>,
379 },
380 Append {
381 nodes: Vec<crate::SessionNodeRecord>,
382 leaf_node_id: Option<String>,
383 },
384 ReplaceFull(crate::SessionGraph),
385}
386
387impl GraphCommitDelta {
388 pub fn leaf_node_id(&self) -> Option<&String> {
389 match self {
390 Self::Unchanged { leaf_node_id } | Self::Append { leaf_node_id, .. } => {
391 leaf_node_id.as_ref()
392 }
393 Self::ReplaceFull(graph) => graph.leaf_node_id.as_ref(),
394 }
395 }
396}
397
398#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
399pub struct RuntimeCommit {
400 pub session_id: String,
401 pub expected_head_revision: Option<u64>,
402 #[serde(default, skip_serializing_if = "Option::is_none")]
403 pub session_execution_lease: Option<SessionExecutionLeaseFence>,
404 #[serde(default, skip_serializing_if = "Option::is_none")]
405 pub release_session_execution_lease: Option<SessionExecutionLeaseCompletion>,
406 pub config: crate::PersistedSessionConfig,
407 pub agent_frames: Vec<crate::AgentFrameRecord>,
408 pub current_agent_frame_id: crate::AgentFrameId,
409 pub graph: GraphCommitDelta,
410 pub checkpoint: HydratedSessionCheckpoint,
411 pub usage_deltas: Vec<crate::TokenLedgerEntry>,
412 pub turn_commit: Option<RuntimeTurnCommitStamp>,
413 pub completed_queue_claims: Vec<crate::QueuedWorkCompletion>,
414 pub completed_turn_input_claims: Vec<crate::TurnInputCompletion>,
415 #[serde(default, skip_serializing_if = "Option::is_none")]
416 pub interrupted_turn_input_turn_id: Option<String>,
417 pub committed_attachment_ids: Vec<crate::AttachmentId>,
425}
426
427#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
428pub struct RuntimeCommitResult {
429 pub head_revision: u64,
430 pub checkpoint_ref: BlobRef,
431 pub manifest: SessionCheckpoint,
432}
433
434#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
435pub struct LeaseOwnerIdentity {
436 pub owner_id: String,
437 pub incarnation_id: String,
438 #[serde(default)]
439 pub liveness: LeaseOwnerLiveness,
440}
441
442impl LeaseOwnerIdentity {
443 pub fn opaque(
444 owner_id: impl Into<String>,
445 incarnation_id: impl Into<String>,
446 ) -> LeaseOwnerIdentity {
447 LeaseOwnerIdentity {
448 owner_id: owner_id.into(),
449 incarnation_id: incarnation_id.into(),
450 liveness: LeaseOwnerLiveness::Opaque,
451 }
452 }
453
454 pub fn local_process(
455 owner_id: impl Into<String>,
456 incarnation_id: impl Into<String>,
457 host_id: impl Into<String>,
458 ) -> LeaseOwnerIdentity {
459 let liveness = LeaseOwnerLiveness::current_local_process(host_id.into())
460 .unwrap_or(LeaseOwnerLiveness::Opaque);
461 LeaseOwnerIdentity {
462 owner_id: owner_id.into(),
463 incarnation_id: incarnation_id.into(),
464 liveness,
465 }
466 }
467
468 pub fn same_incarnation(&self, other: &LeaseOwnerIdentity) -> bool {
469 self.owner_id == other.owner_id && self.incarnation_id == other.incarnation_id
470 }
471
472 pub fn is_definitely_dead_for_claimant(&self, claimant: &LeaseOwnerIdentity) -> bool {
473 self.liveness
474 .is_definitely_dead_for_claimant(&claimant.liveness)
475 }
476}
477
478#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Default)]
479#[serde(tag = "kind", rename_all = "snake_case")]
480pub enum LeaseOwnerLiveness {
481 LocalProcess {
482 host_id: String,
483 boot_id: String,
484 pid: u32,
485 process_start: String,
486 },
487 #[default]
488 Opaque,
489}
490
491impl LeaseOwnerLiveness {
492 pub fn current_local_process(host_id: impl Into<String>) -> Option<LeaseOwnerLiveness> {
493 let boot_id = std::fs::read_to_string(PROC_BOOT_ID_PATH)
494 .ok()
495 .map(|value| value.trim().to_string())
496 .filter(|value| !value.is_empty())?;
497 let pid = std::process::id();
498 let process_start = read_linux_process_start(pid)?;
499 Some(LeaseOwnerLiveness::LocalProcess {
500 host_id: host_id.into(),
501 boot_id,
502 pid,
503 process_start,
504 })
505 }
506
507 pub fn local_process_for_test(
508 host_id: impl Into<String>,
509 boot_id: impl Into<String>,
510 pid: u32,
511 process_start: impl Into<String>,
512 ) -> LeaseOwnerLiveness {
513 LeaseOwnerLiveness::LocalProcess {
514 host_id: host_id.into(),
515 boot_id: boot_id.into(),
516 pid,
517 process_start: process_start.into(),
518 }
519 }
520
521 pub fn is_definitely_dead_for_claimant(&self, claimant: &LeaseOwnerLiveness) -> bool {
522 let (
523 LeaseOwnerLiveness::LocalProcess {
524 host_id,
525 boot_id,
526 pid,
527 process_start,
528 },
529 LeaseOwnerLiveness::LocalProcess {
530 host_id: claimant_host_id,
531 boot_id: claimant_boot_id,
532 ..
533 },
534 ) = (self, claimant)
535 else {
536 return false;
537 };
538 if host_id != claimant_host_id || boot_id != claimant_boot_id {
539 return false;
540 }
541 matches!(linux_process_is_live(*pid, process_start), Some(false))
542 }
543}
544
545fn read_linux_process_start(pid: u32) -> Option<String> {
546 let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")).ok()?;
547 parse_linux_process_start(&stat)
548}
549
550fn linux_process_is_live(pid: u32, expected_process_start: &str) -> Option<bool> {
551 match std::fs::read_to_string(format!("/proc/{pid}/stat")) {
552 Ok(stat) => parse_linux_process_start(&stat).map(|start| start == expected_process_start),
553 Err(err) if err.kind() == std::io::ErrorKind::NotFound => Some(false),
554 Err(_) => None,
555 }
556}
557
558fn parse_linux_process_start(stat: &str) -> Option<String> {
559 let after_comm = stat.rsplit_once(") ")?.1;
560 after_comm.split_whitespace().nth(19).map(ToOwned::to_owned)
561}
562
563#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
564pub struct SessionExecutionLease {
565 pub session_id: String,
566 pub owner: LeaseOwnerIdentity,
567 pub lease_token: String,
568 pub fencing_token: u64,
569 pub claimed_at_epoch_ms: u64,
570 pub expires_at_epoch_ms: u64,
571}
572
573#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
574pub struct SessionExecutionLeaseFence {
575 pub session_id: String,
576 pub owner: LeaseOwnerIdentity,
577 pub lease_token: String,
578 pub fencing_token: u64,
579}
580
581#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
582pub struct SessionExecutionLeaseCompletion {
583 pub session_id: String,
584 pub owner: LeaseOwnerIdentity,
585 pub lease_token: String,
586 pub fencing_token: u64,
587}
588
589impl SessionExecutionLease {
590 pub fn fence(&self) -> SessionExecutionLeaseFence {
591 SessionExecutionLeaseFence {
592 session_id: self.session_id.clone(),
593 owner: self.owner.clone(),
594 lease_token: self.lease_token.clone(),
595 fencing_token: self.fencing_token,
596 }
597 }
598
599 pub fn completion(&self) -> SessionExecutionLeaseCompletion {
600 SessionExecutionLeaseCompletion {
601 session_id: self.session_id.clone(),
602 owner: self.owner.clone(),
603 lease_token: self.lease_token.clone(),
604 fencing_token: self.fencing_token,
605 }
606 }
607}
608
609impl SessionExecutionLeaseCompletion {
610 pub fn from_lease(lease: &SessionExecutionLease) -> Self {
611 lease.completion()
612 }
613}
614
615#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
616pub enum SessionExecutionLeaseClaimOutcome {
617 Acquired(SessionExecutionLease),
618 Busy { holder: SessionExecutionLease },
619}
620
621impl SessionExecutionLeaseClaimOutcome {
622 pub fn acquired(self) -> Option<SessionExecutionLease> {
623 match self {
624 Self::Acquired(lease) => Some(lease),
625 Self::Busy { .. } => None,
626 }
627 }
628}
629
630pub fn ensure_supported_schema_version(
634 record_kind: &'static str,
635 actual: u32,
636 expected: u32,
637) -> Result<(), StoreError> {
638 if actual == expected {
639 Ok(())
640 } else {
641 Err(StoreError::UnsupportedRecordSchemaVersion {
642 record_kind,
643 actual,
644 expected,
645 })
646 }
647}
648
649pub fn ensure_supported_record_schema_version(
650 record_kind: &'static str,
651 value: &serde_json::Value,
652 expected: u32,
653) -> Result<(), StoreError> {
654 let Some(schema_version) = value.get("schema_version") else {
655 return Err(StoreError::MissingRecordSchemaVersion {
656 record_kind,
657 expected,
658 });
659 };
660 let Some(actual) = schema_version
661 .as_u64()
662 .and_then(|version| u32::try_from(version).ok())
663 else {
664 return Err(StoreError::InvalidRecordSchemaVersion {
665 record_kind,
666 actual: schema_version.to_string(),
667 expected,
668 });
669 };
670 ensure_supported_schema_version(record_kind, actual, expected)
671}
672
673pub fn decode_versioned_json_record<T>(
674 json: &str,
675 record_kind: &'static str,
676 expected: u32,
677) -> Result<T, StoreError>
678where
679 T: serde::de::DeserializeOwned,
680{
681 let value: serde_json::Value = serde_json::from_str(json)
682 .map_err(|err| StoreError::Backend(format!("failed to decode {record_kind}: {err}")))?;
683 ensure_supported_record_schema_version(record_kind, &value, expected)?;
684 serde_json::from_value(value)
685 .map_err(|err| StoreError::Backend(format!("failed to decode {record_kind}: {err}")))
686}
687
688#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
689pub struct RuntimeTurnCommitStamp {
690 pub session_id: String,
691 pub turn_id: String,
692 pub turn_commit_hash: String,
693}
694
695impl RuntimeTurnCommitStamp {
696 pub fn new(
697 session_id: impl Into<String>,
698 turn_id: impl Into<String>,
699 turn_commit_hash: impl Into<String>,
700 ) -> Self {
701 Self {
702 session_id: session_id.into(),
703 turn_id: turn_id.into(),
704 turn_commit_hash: turn_commit_hash.into(),
705 }
706 }
707}
708
709fn build_persisted_turn_state(state: &crate::RuntimeSessionState) -> crate::PersistedTurnState {
710 crate::PersistedTurnState {
711 turn_index: state.turn_index,
712 token_usage: state.token_usage.clone(),
713 last_prompt_usage: state.last_prompt_usage.clone(),
714 protocol_turn_options: state.protocol_turn_options.clone(),
715 }
716}
717
718fn build_checkpoint_from_persisted_state(
719 state: &crate::RuntimeSessionState,
720) -> HydratedSessionCheckpoint {
721 HydratedSessionCheckpoint {
722 turn_state: build_persisted_turn_state(state),
723 tool_state_ref: state.tool_state_ref.clone(),
724 tool_state: state.tool_state_snapshot.clone(),
725 plugin_snapshot_ref: state.plugin_snapshot_ref.clone(),
726 plugin_snapshot_revision: state.plugin_snapshot_revision,
727 plugin_snapshot: state.plugin_snapshot.clone(),
728 execution_state_ref: state.execution_state_ref.clone(),
729 execution_state: state.execution_state_snapshot.clone(),
730 }
731}
732
733impl RuntimeCommit {
734 pub fn turn_commit_hash(&self) -> Result<String, StoreError> {
735 let mut semantic_commit = self.clone();
736 semantic_commit.expected_head_revision = None;
737 semantic_commit.session_execution_lease = None;
738 semantic_commit.release_session_execution_lease = None;
739 semantic_commit.turn_commit = None;
740 let mut semantic_commit = serde_json::to_value(&semantic_commit).map_err(|err| {
741 StoreError::Backend(format!("failed to serialize runtime turn commit: {err}"))
742 })?;
743 scrub_turn_commit_hash_value(&mut semantic_commit);
744 crate::stable_hash::stable_json_sha256_hex(&semantic_commit).map_err(|err| {
745 StoreError::Backend(format!(
746 "failed to serialize runtime turn commit hash: {err}"
747 ))
748 })
749 }
750
751 pub fn persisted_state(
752 state: &crate::RuntimeSessionState,
753 usage_deltas: &[crate::TokenLedgerEntry],
754 ) -> Self {
755 Self {
756 session_id: state.session_id.clone(),
757 expected_head_revision: state.head_revision,
758 session_execution_lease: None,
759 release_session_execution_lease: None,
760 config: persisted_session_config_from_state(state),
761 agent_frames: state.agent_frames.clone(),
762 current_agent_frame_id: state.current_agent_frame_id.clone(),
763 graph: if state.graph_replace_required || state.head_revision.is_none() {
764 GraphCommitDelta::ReplaceFull(state.session_graph.clone())
765 } else {
766 GraphCommitDelta::Unchanged {
767 leaf_node_id: state.session_graph.leaf_node_id.clone(),
768 }
769 },
770 checkpoint: build_checkpoint_from_persisted_state(state),
771 usage_deltas: usage_deltas.to_vec(),
772 turn_commit: None,
773 completed_queue_claims: Vec::new(),
774 completed_turn_input_claims: Vec::new(),
775 interrupted_turn_input_turn_id: None,
776 committed_attachment_ids: Vec::new(),
777 }
778 }
779
780 pub(crate) fn persisted_state_with_graph_commit(
781 state: &crate::RuntimeSessionState,
782 graph: GraphCommitDelta,
783 usage_deltas: &[crate::TokenLedgerEntry],
784 ) -> Self {
785 Self {
786 session_id: state.session_id.clone(),
787 expected_head_revision: state.head_revision,
788 session_execution_lease: None,
789 release_session_execution_lease: None,
790 config: persisted_session_config_from_state(state),
791 agent_frames: state.agent_frames.clone(),
792 current_agent_frame_id: state.current_agent_frame_id.clone(),
793 graph,
794 checkpoint: build_checkpoint_from_persisted_state(state),
795 usage_deltas: usage_deltas.to_vec(),
796 turn_commit: None,
797 completed_queue_claims: Vec::new(),
798 completed_turn_input_claims: Vec::new(),
799 interrupted_turn_input_turn_id: None,
800 committed_attachment_ids: Vec::new(),
801 }
802 }
803
804 pub fn with_turn_commit(mut self, turn_commit: RuntimeTurnCommitStamp) -> Self {
805 self.turn_commit = Some(turn_commit);
806 self
807 }
808
809 pub fn with_session_execution_lease(mut self, lease: SessionExecutionLeaseFence) -> Self {
810 self.session_execution_lease = Some(lease);
811 self
812 }
813
814 pub fn releasing_session_execution_lease(
815 mut self,
816 completion: SessionExecutionLeaseCompletion,
817 ) -> Self {
818 self.release_session_execution_lease = Some(completion);
819 self
820 }
821
822 pub fn completing_queue_claim(
823 mut self,
824 completed_queue_claim: crate::QueuedWorkCompletion,
825 ) -> Self {
826 self.completed_queue_claims.push(completed_queue_claim);
827 self
828 }
829
830 pub fn completing_queue_claims(
831 mut self,
832 completed_queue_claims: impl IntoIterator<Item = crate::QueuedWorkCompletion>,
833 ) -> Self {
834 self.completed_queue_claims.extend(completed_queue_claims);
835 self
836 }
837
838 pub fn completing_turn_input_claim(
839 mut self,
840 completed_turn_input_claim: crate::TurnInputCompletion,
841 ) -> Self {
842 self.completed_turn_input_claims
843 .push(completed_turn_input_claim);
844 self
845 }
846
847 pub fn completing_turn_input_claims(
848 mut self,
849 completed_turn_input_claims: impl IntoIterator<Item = crate::TurnInputCompletion>,
850 ) -> Self {
851 self.completed_turn_input_claims
852 .extend(completed_turn_input_claims);
853 self
854 }
855
856 pub fn deferring_interrupted_turn_inputs(mut self, turn_id: impl Into<String>) -> Self {
857 self.interrupted_turn_input_turn_id = Some(turn_id.into());
858 self
859 }
860
861 pub fn with_committed_attachments(
862 mut self,
863 attachment_ids: impl IntoIterator<Item = crate::AttachmentId>,
864 ) -> Self {
865 self.committed_attachment_ids = attachment_ids.into_iter().collect();
866 self
867 }
868}
869
870fn scrub_turn_commit_hash_value(value: &mut serde_json::Value) {
871 match value {
872 serde_json::Value::Object(map) => {
873 let is_message = map.contains_key("role") && map.contains_key("parts");
874 let is_message_part = map.contains_key("kind")
875 && map.contains_key("content")
876 && map.contains_key("prune_state");
877 if is_message || is_message_part {
878 map.remove("id");
879 }
880 for volatile_key in ["node_id", "parent_node_id", "leaf_node_id", "timestamp"] {
881 map.remove(volatile_key);
882 }
883 for child in map.values_mut() {
884 scrub_turn_commit_hash_value(child);
885 }
886 }
887 serde_json::Value::Array(items) => {
888 for item in items {
889 scrub_turn_commit_hash_value(item);
890 }
891 }
892 _ => {}
893 }
894}
895
896fn persisted_session_state_from_head(
897 head: SessionHead,
898 checkpoint: Option<HydratedSessionCheckpoint>,
899) -> crate::RuntimeSessionState {
900 let mut state = crate::RuntimeSessionState {
901 session_id: head.session_id,
902 policy: crate::SessionPolicy::default(),
903 agent_frames: head.agent_frames,
904 current_agent_frame_id: head.current_agent_frame_id,
905 session_graph: head.graph,
906 turn_index: 0,
907 token_usage: crate::TokenUsage::default(),
908 last_prompt_usage: None,
909 protocol_turn_options: crate::ProtocolTurnOptions::default(),
910 tool_state_ref: None,
911 tool_state_generation: None,
912 tool_state_snapshot: None,
913 plugin_snapshot_ref: None,
914 plugin_snapshot_revision: None,
915 plugin_snapshot: None,
916 execution_state_ref: None,
917 execution_state_snapshot: None,
918 token_ledger: head.token_ledger,
919 checkpoint_ref: head.checkpoint_ref.clone(),
920 head_revision: Some(head.head_revision),
921 graph_replace_required: false,
922 };
923 state.policy.model = head.config.model.clone();
924 state.policy.provider_id = head.config.provider_id.clone();
925 if let Some(checkpoint) = checkpoint {
926 state.turn_index = checkpoint.turn_state.turn_index;
927 state.token_usage = checkpoint.turn_state.token_usage;
928 state.last_prompt_usage = checkpoint.turn_state.last_prompt_usage;
929 state.protocol_turn_options = checkpoint.turn_state.protocol_turn_options;
930 state.tool_state_ref = checkpoint.tool_state_ref.clone();
931 state.tool_state_generation = checkpoint
932 .tool_state
933 .as_ref()
934 .map(|snapshot| snapshot.generation());
935 state.tool_state_snapshot = checkpoint.tool_state;
936 state.plugin_snapshot_ref = checkpoint.plugin_snapshot_ref.clone();
937 state.plugin_snapshot_revision = checkpoint.plugin_snapshot_revision;
938 state.plugin_snapshot = checkpoint.plugin_snapshot;
939 state.execution_state_ref = checkpoint.execution_state_ref.clone();
940 state.execution_state_snapshot = checkpoint.execution_state;
941 }
942 state.ensure_agent_frame_initialized();
943 state
944}
945
946impl Default for SessionHead {
947 fn default() -> Self {
948 Self {
949 session_id: default_root_session_id(),
950 head_revision: 0,
951 agent_frames: Vec::new(),
952 current_agent_frame_id: String::new(),
953 graph: crate::SessionGraph::default(),
954 config: crate::PersistedSessionConfig::default(),
955 checkpoint_ref: None,
956 token_ledger: Vec::new(),
957 }
958 }
959}
960
961impl Default for SessionHeadMeta {
962 fn default() -> Self {
963 Self {
964 schema_version: SESSION_HEAD_META_SCHEMA_VERSION,
965 session_id: default_root_session_id(),
966 head_revision: 0,
967 config: crate::PersistedSessionConfig::default(),
968 agent_frames: Vec::new(),
969 current_agent_frame_id: String::new(),
970 checkpoint_ref: None,
971 leaf_node_id: None,
972 graph_node_count: 0,
973 token_ledger: Vec::new(),
974 }
975 }
976}
977
978#[async_trait::async_trait]
995pub trait SessionCommitStore: AttachmentManifest + Send + Sync {
996 fn durability_tier(&self) -> crate::DurabilityTier {
999 crate::DurabilityTier::Inline
1000 }
1001
1002 async fn load_session(
1003 &self,
1004 scope: SessionReadScope,
1005 ) -> Result<Option<PersistedSessionRead>, StoreError>;
1006
1007 async fn load_node(
1008 &self,
1009 node_id: &str,
1010 ) -> Result<Option<crate::SessionNodeRecord>, StoreError>;
1011
1012 async fn commit_runtime_state(
1013 &self,
1014 commit: RuntimeCommit,
1015 ) -> Result<RuntimeCommitResult, StoreError>;
1016
1017 async fn save_session_meta(&self, meta: SessionMeta) -> Result<(), StoreError>;
1018 async fn load_session_meta(&self) -> Result<Option<SessionMeta>, StoreError>;
1019}
1020
1021#[async_trait::async_trait]
1029pub trait TurnInputStore: Send + Sync {
1030 async fn enqueue_pending_turn_input(
1032 &self,
1033 input: crate::PendingTurnInputDraft,
1034 ) -> Result<crate::PendingTurnInput, StoreError>;
1035
1036 async fn list_pending_turn_inputs(
1041 &self,
1042 session_id: &str,
1043 ) -> Result<Vec<crate::PendingTurnInput>, StoreError>;
1044
1045 async fn cancel_pending_turn_input(
1052 &self,
1053 session_id: &str,
1054 input_id: &str,
1055 ) -> Result<crate::PendingTurnInputCancelOutcome, StoreError> {
1056 let target = crate::PendingTurnInputCancelTarget::input_id(input_id);
1057 let targets = vec![target];
1058 let mut outcomes = self
1059 .cancel_pending_turn_inputs(session_id, &targets)
1060 .await?;
1061 Ok(outcomes
1062 .pop()
1063 .map(|result| result.outcome)
1064 .unwrap_or(crate::PendingTurnInputCancelOutcome::NotFound))
1065 }
1066
1067 async fn cancel_pending_turn_inputs(
1069 &self,
1070 session_id: &str,
1071 targets: &[crate::PendingTurnInputCancelTarget],
1072 ) -> Result<Vec<crate::PendingTurnInputCancelResult>, StoreError>;
1073
1074 async fn cancel_pending_turn_input_suffix(
1076 &self,
1077 session_id: &str,
1078 anchor: &crate::PendingTurnInputCancelTarget,
1079 ) -> Result<crate::PendingTurnInputSuffixCancelOutcome, StoreError>;
1080
1081 async fn claim_active_turn_inputs(
1087 &self,
1088 session_id: &str,
1089 session_execution_lease: &SessionExecutionLeaseFence,
1090 owner: &LeaseOwnerIdentity,
1091 turn_id: &str,
1092 checkpoint: crate::CheckpointKind,
1093 max_inputs: usize,
1094 ) -> Result<Option<crate::TurnInputClaim>, StoreError>;
1095
1096 async fn claim_next_turn_inputs(
1098 &self,
1099 session_id: &str,
1100 session_execution_lease: &SessionExecutionLeaseFence,
1101 owner: &LeaseOwnerIdentity,
1102 max_inputs: usize,
1103 ) -> Result<Option<crate::TurnInputClaim>, StoreError>;
1104
1105 async fn abandon_turn_input_claim(
1107 &self,
1108 claim: &crate::TurnInputClaim,
1109 ) -> Result<(), StoreError>;
1110}
1111
1112#[async_trait::async_trait]
1115pub trait SessionExecutionLeaseStore: Send + Sync {
1116 async fn try_claim_session_execution_lease(
1123 &self,
1124 session_id: &str,
1125 owner: &LeaseOwnerIdentity,
1126 lease_ttl_ms: u64,
1127 ) -> Result<SessionExecutionLeaseClaimOutcome, StoreError>;
1128
1129 async fn reclaim_session_execution_lease(
1135 &self,
1136 session_id: &str,
1137 owner: &LeaseOwnerIdentity,
1138 observed_holder: &SessionExecutionLeaseFence,
1139 lease_ttl_ms: u64,
1140 ) -> Result<SessionExecutionLeaseClaimOutcome, StoreError>;
1141
1142 async fn renew_session_execution_lease(
1147 &self,
1148 fence: &SessionExecutionLeaseFence,
1149 lease_ttl_ms: u64,
1150 ) -> Result<SessionExecutionLease, StoreError>;
1151
1152 async fn release_session_execution_lease(
1156 &self,
1157 completion: &SessionExecutionLeaseCompletion,
1158 ) -> Result<(), StoreError>;
1159}
1160
1161#[async_trait::async_trait]
1167pub trait QueuedWorkStore: Send + Sync {
1168 async fn enqueue_queued_work(
1170 &self,
1171 batch: crate::QueuedWorkBatchDraft,
1172 ) -> Result<crate::QueuedWorkBatch, StoreError>;
1173
1174 async fn claim_leading_ready_session_command(
1181 &self,
1182 session_id: &str,
1183 session_execution_lease: &SessionExecutionLeaseFence,
1184 owner: &LeaseOwnerIdentity,
1185 ) -> Result<Option<crate::QueuedWorkClaim>, StoreError>;
1186
1187 async fn claim_ready_queued_work(
1194 &self,
1195 session_id: &str,
1196 session_execution_lease: &SessionExecutionLeaseFence,
1197 owner: &LeaseOwnerIdentity,
1198 boundary: crate::QueuedWorkClaimBoundary,
1199 max_batches: usize,
1200 ) -> Result<Option<crate::QueuedWorkClaim>, StoreError>;
1201
1202 async fn claim_ready_queued_work_by_batch_ids(
1214 &self,
1215 session_id: &str,
1216 session_execution_lease: &SessionExecutionLeaseFence,
1217 owner: &LeaseOwnerIdentity,
1218 boundary: crate::QueuedWorkClaimBoundary,
1219 batch_ids: &[String],
1220 ) -> Result<Option<crate::QueuedWorkClaim>, StoreError> {
1221 if batch_ids.is_empty() {
1222 return Ok(None);
1223 }
1224 let Some(claim) = self
1225 .claim_ready_queued_work(
1226 session_id,
1227 session_execution_lease,
1228 owner,
1229 boundary,
1230 batch_ids.len(),
1231 )
1232 .await?
1233 else {
1234 return Ok(None);
1235 };
1236 let claimed_ids = claim
1237 .batches
1238 .iter()
1239 .map(|batch| batch.batch_id.as_str())
1240 .collect::<Vec<_>>();
1241 if claimed_ids == batch_ids.iter().map(String::as_str).collect::<Vec<_>>() {
1242 return Ok(Some(claim));
1243 }
1244 self.abandon_queued_work_claim(&claim).await?;
1245 Ok(None)
1246 }
1247
1248 async fn abandon_queued_work_claim(
1250 &self,
1251 claim: &crate::QueuedWorkClaim,
1252 ) -> Result<(), StoreError>;
1253
1254 async fn cancel_queued_work_batch(
1261 &self,
1262 session_id: &str,
1263 batch_id: &str,
1264 ) -> Result<Option<crate::QueuedWorkBatch>, StoreError>;
1265
1266 async fn list_queued_work(
1269 &self,
1270 session_id: &str,
1271 ) -> Result<Vec<crate::QueuedWorkBatch>, StoreError>;
1272
1273 async fn list_pending_queued_work(
1286 &self,
1287 session_id: &str,
1288 ) -> Result<Vec<crate::QueuedWorkBatch>, StoreError>;
1289}
1290
1291#[async_trait::async_trait]
1294pub trait StoreMaintenance: Send + Sync {
1295 async fn tombstone_nodes(&self, ids: &[String]) -> Result<(), StoreError>;
1298
1299 async fn vacuum(&self) -> Result<VacuumReport, StoreError>;
1302
1303 async fn gc_unreachable(&self) -> Result<GcReport, StoreError>;
1305}
1306
1307pub trait RuntimePersistence:
1324 SessionCommitStore
1325 + TurnInputStore
1326 + SessionExecutionLeaseStore
1327 + QueuedWorkStore
1328 + StoreMaintenance
1329{
1330}
1331
1332impl<T> RuntimePersistence for T where
1333 T: SessionCommitStore
1334 + TurnInputStore
1335 + SessionExecutionLeaseStore
1336 + QueuedWorkStore
1337 + StoreMaintenance
1338 + ?Sized
1339{
1340}
1341
1342fn persisted_session_state_from_read(read: PersistedSessionRead) -> crate::RuntimeSessionState {
1343 persisted_session_state_from_head(
1344 SessionHead {
1345 session_id: read.session_id,
1346 head_revision: read.head_revision,
1347 agent_frames: read.agent_frames,
1348 current_agent_frame_id: read.current_agent_frame_id,
1349 graph: read.graph,
1350 config: read.config,
1351 checkpoint_ref: read.checkpoint_ref,
1352 token_ledger: read.token_ledger,
1353 },
1354 read.checkpoint,
1355 )
1356}
1357
1358pub async fn load_persisted_session_state(
1359 store: &(dyn RuntimePersistence + '_),
1360) -> Result<Option<crate::RuntimeSessionState>, StoreError> {
1361 Ok(store
1362 .load_session(SessionReadScope::FullGraph)
1363 .await?
1364 .map(persisted_session_state_from_read))
1365}
1366
1367pub async fn load_persisted_session_state_active_path(
1368 store: &(dyn RuntimePersistence + '_),
1369 leaf_node_id: Option<String>,
1370) -> Result<Option<crate::RuntimeSessionState>, StoreError> {
1371 Ok(store
1372 .load_session(SessionReadScope::ActivePath { leaf_node_id })
1373 .await?
1374 .map(persisted_session_state_from_read))
1375}
1376
1377pub async fn refresh_persisted_session_state(
1378 store: &(dyn RuntimePersistence + '_),
1379 state: &mut crate::RuntimeSessionState,
1380) -> Result<(), StoreError> {
1381 if let Some(mut fresh) = load_persisted_session_state(store).await? {
1382 fresh.policy.session_id = state.policy.session_id.clone();
1383 fresh.policy.max_turns = state.policy.max_turns;
1384 *state = fresh;
1385 }
1386 Ok(())
1387}
1388
1389#[cfg(test)]
1390mod tests {
1391 use super::{LeaseOwnerIdentity, LeaseOwnerLiveness};
1392
1393 fn local_liveness(
1394 host_id: &str,
1395 boot_id: &str,
1396 pid: u32,
1397 process_start: &str,
1398 ) -> LeaseOwnerLiveness {
1399 LeaseOwnerLiveness::local_process_for_test(host_id, boot_id, pid, process_start)
1400 }
1401
1402 #[test]
1403 fn lease_owner_identity_requires_same_incarnation() {
1404 let first = LeaseOwnerIdentity::opaque("owner", "incarnation-a");
1405 let same = LeaseOwnerIdentity::opaque("owner", "incarnation-a");
1406 let next = LeaseOwnerIdentity::opaque("owner", "incarnation-b");
1407
1408 assert!(first.same_incarnation(&same));
1409 assert!(!first.same_incarnation(&next));
1410 }
1411
1412 #[test]
1413 fn local_liveness_only_proves_same_host_boot_dead_processes() {
1414 let holder = local_liveness(
1415 "host-a",
1416 "boot-a",
1417 std::process::id(),
1418 "not-the-current-process-start",
1419 );
1420 let same_host_boot = local_liveness("host-a", "boot-a", std::process::id(), "claimant");
1421 let other_host = local_liveness("host-b", "boot-a", std::process::id(), "claimant");
1422 let other_boot = local_liveness("host-a", "boot-b", std::process::id(), "claimant");
1423
1424 assert!(holder.is_definitely_dead_for_claimant(&same_host_boot));
1425 assert!(!holder.is_definitely_dead_for_claimant(&other_host));
1426 assert!(!holder.is_definitely_dead_for_claimant(&other_boot));
1427 assert!(!holder.is_definitely_dead_for_claimant(&LeaseOwnerLiveness::Opaque));
1428 assert!(!LeaseOwnerLiveness::Opaque.is_definitely_dead_for_claimant(&same_host_boot));
1429 }
1430}