1use super::*;
25
26#[async_trait::async_trait]
27impl RuntimePersistence for Store {
28 fn durability_tier(&self) -> DurabilityTier {
29 DurabilityTier::Durable
30 }
31
32 async fn load_session(
33 &self,
34 scope: SessionReadScope,
35 ) -> Result<Option<PersistedSessionRead>, StoreError> {
36 self.conn
37 .call(move |conn| {
38 let outcome: Result<Option<PersistedSessionRead>, StoreError> = (|| {
39 let Some(meta) = try_load_session_head_meta_from_conn(conn)? else {
40 return Ok(None);
41 };
42 let leaf_node_id = match &scope {
43 SessionReadScope::FullGraph => meta.leaf_node_id.clone(),
44 SessionReadScope::ActivePath { leaf_node_id } => {
45 leaf_node_id.clone().or_else(|| meta.leaf_node_id.clone())
46 }
47 };
48 let mut graph = match scope {
49 SessionReadScope::FullGraph => {
50 Self::load_session_graph_from_conn(conn, meta.leaf_node_id.clone())
51 }
52 SessionReadScope::ActivePath { .. } => {
53 Self::load_active_path_session_graph_from_conn(
54 conn,
55 leaf_node_id.clone(),
56 )
57 .map_err(sqlite_error)?
58 }
59 };
60 graph.set_leaf_node_id(leaf_node_id);
61 let checkpoint = meta
62 .checkpoint_ref
63 .as_ref()
64 .map(|blob_ref| Self::get_checkpoint_conn(conn, blob_ref))
65 .transpose()?
66 .flatten();
67 Ok(Some(PersistedSessionRead {
68 session_id: meta.session_id,
69 head_revision: meta.head_revision,
70 config: meta.config,
71 agent_frames: meta.agent_frames,
72 current_agent_frame_id: meta.current_agent_frame_id,
73 graph,
74 checkpoint_ref: meta.checkpoint_ref,
75 checkpoint,
76 token_ledger: merge_token_ledger_entries(Self::load_usage_deltas_conn(
77 conn,
78 )),
79 }))
80 })(
81 );
82 Ok(outcome)
83 })
84 .await
85 .map_err(sqlite_error)?
86 }
87
88 async fn load_node(
89 &self,
90 node_id: &str,
91 ) -> Result<Option<lash_core::SessionNodeRecord>, StoreError> {
92 let node_id = node_id.to_string();
93 let row: Option<String> = self
94 .conn
95 .call(move |conn| {
96 conn.query_row(
97 "SELECT node_json FROM graph_nodes WHERE node_id = ?1 AND tombstoned = 0",
98 params![node_id],
99 |row| row.get(0),
100 )
101 .optional()
102 })
103 .await
104 .map_err(sqlite_error)?;
105 Ok(row.and_then(|json| serde_json::from_str(&json).ok()))
106 }
107
108 async fn commit_runtime_state(
109 &self,
110 commit: RuntimeCommit,
111 ) -> Result<RuntimeCommitResult, StoreError> {
112 let blob_profile = self.options.blob_profile;
113 let result = self
114 .conn
115 .write_flow(move |tx| {
116 let outcome: Result<RuntimeCommitResult, StoreError> = (|| {
117 let existing = try_load_session_head_meta_from_conn(tx)?;
118 if let Some(bound_session_id) =
119 existing.as_ref().map(|meta| meta.session_id.as_str())
120 && bound_session_id != commit.session_id
121 {
122 return Err(StoreError::SessionBindingMismatch {
123 bound_session_id: bound_session_id.to_string(),
124 attempted_session_id: commit.session_id.clone(),
125 });
126 }
127 if let Some(completed) = &commit.turn_commit {
128 if completed.session_id != commit.session_id {
129 return Err(StoreError::RuntimeTurnCommitConflict {
130 session_id: completed.session_id.clone(),
131 turn_id: completed.turn_id.clone(),
132 });
133 }
134 let prior: Option<(String, String)> = tx
135 .query_row(
136 "SELECT turn_commit_hash, result_json FROM runtime_turn_commits
137 WHERE session_id = ?1 AND turn_id = ?2",
138 params![completed.session_id, completed.turn_id],
139 |row| Ok((row.get(0)?, row.get(1)?)),
140 )
141 .optional()
142 .map_err(sqlite_error)?;
143 if let Some((turn_commit_hash, result_json)) = prior {
144 if turn_commit_hash == completed.turn_commit_hash {
145 let result: RuntimeCommitResult =
146 serde_json::from_str(&result_json).map_err(|err| {
147 StoreError::Backend(format!(
148 "failed to decode runtime turn commit result: {err}"
149 ))
150 })?;
151 if let Some(completion) =
152 commit.release_session_execution_lease.as_ref()
153 {
154 release_session_execution_lease_conn(tx, completion)?;
155 }
156 return Ok(result);
157 }
158 return Err(StoreError::RuntimeTurnCommitConflict {
159 session_id: completed.session_id.clone(),
160 turn_id: completed.turn_id.clone(),
161 });
162 }
163 }
164 let Some(session_execution_lease) = commit.session_execution_lease.as_ref()
165 else {
166 return Err(StoreError::SessionExecutionLeaseExpired {
167 session_id: commit.session_id.clone(),
168 });
169 };
170 ensure_session_execution_lease_conn(
171 tx,
172 &commit.session_id,
173 session_execution_lease,
174 )?;
175 let actual_revision = existing.as_ref().map_or(0, |meta| meta.head_revision);
176 if commit.expected_head_revision.is_some()
177 && commit.expected_head_revision != Some(actual_revision)
178 {
179 return Err(StoreError::HeadRevisionConflict {
180 expected: commit.expected_head_revision,
181 actual: actual_revision,
182 });
183 }
184 for completed in &commit.completed_queue_claims {
185 if completed.session_id != commit.session_id {
186 return Err(StoreError::QueuedWorkClaimExpired {
187 session_id: completed.session_id.clone(),
188 claim_id: completed.claim_id.clone(),
189 });
190 }
191 ensure_queued_work_completion_conn(tx, completed)?;
192 }
193 for completed in &commit.completed_turn_input_claims {
194 if completed.session_id != commit.session_id {
195 return Err(StoreError::TurnInputClaimExpired {
196 session_id: completed.session_id.clone(),
197 claim_id: completed.claim_id.clone(),
198 });
199 }
200 let owned_rows: usize = tx
201 .query_row(
202 "SELECT COUNT(*)
203 FROM pending_turn_inputs
204 WHERE session_id = ?1
205 AND claim_id = ?2
206 AND claim_token = ?3",
207 params![
208 completed.session_id,
209 completed.claim_id,
210 completed.lease_token
211 ],
212 |row| row.get::<_, i64>(0),
213 )
214 .map_err(sqlite_error)? as usize;
215 ensure_turn_input_completion_owns_all_inputs(completed, owned_rows)?;
216 }
217
218 let stored_checkpoint =
219 Self::put_checkpoint_conn(tx, &commit.checkpoint, blob_profile)
220 .map_err(sqlite_error)?;
221
222 if !commit.usage_deltas.is_empty() {
223 let mut stmt = tx
224 .prepare(
225 "INSERT INTO usage_deltas (
226 source, model, input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens, reasoning_output_tokens
227 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
228 )
229 .map_err(sqlite_error)?;
230 for entry in &commit.usage_deltas {
231 stmt.execute(params![
232 entry.source,
233 entry.model,
234 entry.usage.input_tokens,
235 entry.usage.output_tokens,
236 entry.usage.cache_read_input_tokens,
237 entry.usage.cache_write_input_tokens,
238 entry.usage.reasoning_output_tokens,
239 ])
240 .map_err(sqlite_error)?;
241 }
242 }
243
244 let leaf_node_id = match &commit.graph {
245 GraphCommitDelta::Unchanged { leaf_node_id } => leaf_node_id.clone(),
246 GraphCommitDelta::Append {
247 nodes,
248 leaf_node_id,
249 } => {
250 for node in nodes {
251 let node_json = encode_json(node);
252 tx.execute(
253 "INSERT INTO graph_nodes (node_id, node_json) VALUES (?1, ?2)",
254 params![node.node_id, node_json],
255 )
256 .map_err(sqlite_error)?;
257 }
258 leaf_node_id.clone()
259 }
260 GraphCommitDelta::ReplaceFull(graph) => {
261 tx.execute("DELETE FROM graph_nodes", [])
262 .map_err(sqlite_error)?;
263 for node in &graph.nodes {
264 let node_json = encode_json(node);
265 tx.execute(
266 "INSERT INTO graph_nodes (node_id, node_json) VALUES (?1, ?2)",
267 params![node.node_id, node_json],
268 )
269 .map_err(sqlite_error)?;
270 }
271 graph.leaf_node_id.clone()
272 }
273 };
274 let graph_node_count: usize = tx
275 .query_row(
276 "SELECT COUNT(*) FROM graph_nodes WHERE tombstoned = 0",
277 [],
278 |row| row.get::<_, i64>(0),
279 )
280 .map_err(sqlite_error)? as usize;
281 let next_revision = actual_revision + 1;
282 let meta = SessionHeadMeta {
283 schema_version: lash_core::store::SESSION_HEAD_META_SCHEMA_VERSION,
284 session_id: commit.session_id.clone(),
285 head_revision: next_revision,
286 config: commit.config.clone(),
287 agent_frames: commit.agent_frames.clone(),
288 current_agent_frame_id: commit.current_agent_frame_id.clone(),
289 checkpoint_ref: Some(stored_checkpoint.checkpoint_ref.clone()),
290 leaf_node_id,
291 graph_node_count,
292 token_ledger: Vec::new(),
293 };
294 tx.execute(
295 "INSERT OR REPLACE INTO session_head (singleton, session_id, head_json, head_revision)
296 VALUES (1, ?1, ?2, ?3)",
297 params![
298 meta.session_id,
299 encode_json(&meta),
300 meta.head_revision as i64
301 ],
302 )
303 .map_err(sqlite_error)?;
304 for completed in &commit.completed_queue_claims {
305 for batch_id in &completed.batch_ids {
306 tx.execute(
307 "DELETE FROM queued_work_batches
308 WHERE session_id = ?1
309 AND batch_id = ?2
310 AND claim_id = ?3
311 AND claim_token = ?4",
312 params![
313 completed.session_id,
314 batch_id,
315 completed.claim_id,
316 completed.lease_token
317 ],
318 )
319 .map_err(sqlite_error)?;
320 }
321 }
322 for completed in &commit.completed_turn_input_claims {
323 for input_id in &completed.input_ids {
324 tx.execute(
325 "UPDATE pending_turn_inputs
326 SET state = ?5,
327 claim_id = NULL,
328 claim_owner_id = NULL,
329 claim_owner_incarnation_id = NULL,
330 claim_owner_liveness_json = NULL,
331 claim_token = NULL,
332 claim_claimed_at_ms = 0,
333 claim_expires_at_ms = 0
334 WHERE session_id = ?1
335 AND input_id = ?2
336 AND claim_id = ?3
337 AND claim_token = ?4",
338 params![
339 completed.session_id,
340 input_id,
341 completed.claim_id,
342 completed.lease_token,
343 lash_core::TurnInputState::Completed.as_str(),
344 ],
345 )
346 .map_err(sqlite_error)?;
347 }
348 }
349 if let Some(turn_id) = commit.interrupted_turn_input_turn_id.as_deref() {
350 let input_ids = {
351 let mut stmt = tx
352 .prepare(
353 "SELECT input_id, ingress_json
354 FROM pending_turn_inputs
355 WHERE session_id = ?1 AND state = ?2",
356 )
357 .map_err(sqlite_error)?;
358 let rows = stmt
359 .query_map(
360 params![
361 commit.session_id,
362 lash_core::TurnInputState::PendingActive.as_str()
363 ],
364 |row| {
365 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
366 },
367 )
368 .map_err(sqlite_error)?;
369 let mut input_ids = Vec::new();
370 for row in rows {
371 let (input_id, ingress_json) = row.map_err(sqlite_error)?;
372 let ingress = decode_turn_input_ingress(ingress_json)?;
373 if ingress
374 .active_turn_id()
375 .is_some_and(|active| active == turn_id)
376 {
377 input_ids.push(input_id);
378 }
379 }
380 input_ids
381 };
382 let next_turn_ingress = encode_json(&lash_core::TurnInputIngress::NextTurn);
383 let mut stmt = tx
384 .prepare(
385 "UPDATE pending_turn_inputs
386 SET state = ?3,
387 ingress_json = ?4,
388 claim_id = NULL,
389 claim_owner_id = NULL,
390 claim_owner_incarnation_id = NULL,
391 claim_owner_liveness_json = NULL,
392 claim_token = NULL,
393 claim_claimed_at_ms = 0,
394 claim_expires_at_ms = 0
395 WHERE session_id = ?1 AND input_id = ?2",
396 )
397 .map_err(sqlite_error)?;
398 for input_id in input_ids {
399 stmt.execute(params![
400 commit.session_id,
401 input_id,
402 lash_core::TurnInputState::DeferredNextTurn.as_str(),
403 next_turn_ingress
404 ])
405 .map_err(sqlite_error)?;
406 }
407 }
408 if !commit.committed_attachment_ids.is_empty() {
409 let now = current_epoch_ms() as i64;
410 let mut stmt = tx
411 .prepare(
412 "UPDATE attachment_manifest
413 SET committed_at_ms = COALESCE(committed_at_ms, ?1)
414 WHERE attachment_id = ?2 AND session_id = ?3",
415 )
416 .map_err(sqlite_error)?;
417 for id in &commit.committed_attachment_ids {
418 stmt.execute(params![now, id.as_str(), commit.session_id])
419 .map_err(sqlite_error)?;
420 }
421 }
422 let result = RuntimeCommitResult {
423 head_revision: next_revision,
424 checkpoint_ref: stored_checkpoint.checkpoint_ref,
425 manifest: stored_checkpoint.manifest,
426 };
427 if let Some(completed) = &commit.turn_commit {
428 tx.execute(
429 "INSERT INTO runtime_turn_commits (
430 session_id, turn_id, turn_commit_hash, result_json, committed_at_ms
431 )
432 VALUES (?1, ?2, ?3, ?4, ?5)",
433 params![
434 completed.session_id,
435 completed.turn_id,
436 completed.turn_commit_hash,
437 encode_json(&result),
438 current_epoch_ms() as i64
439 ],
440 )
441 .map_err(sqlite_error)?;
442 }
443 if let Some(completion) = commit.release_session_execution_lease.as_ref() {
444 release_session_execution_lease_conn(tx, completion)?;
445 }
446 Ok(result)
447 })();
448 match outcome {
453 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
454 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
455 }
456 })
457 .await
458 .map_err(sqlite_error)??;
459 self.maybe_auto_gc().await;
460 Ok(result)
461 }
462
463 async fn try_claim_session_execution_lease(
464 &self,
465 session_id: &str,
466 owner: &LeaseOwnerIdentity,
467 lease_ttl_ms: u64,
468 ) -> Result<SessionExecutionLeaseClaimOutcome, StoreError> {
469 let session_id = session_id.to_string();
470 let owner = owner.clone();
471 self.conn
472 .write_flow(move |tx| {
473 let outcome: Result<SessionExecutionLeaseClaimOutcome, StoreError> = (|| {
474 let now = current_epoch_ms();
475 let current = load_session_execution_lease_row_conn(tx, &session_id)?;
476 if current.as_ref().is_some_and(|lease| {
477 lease.lease_token.is_some() && lease.expires_at_ms > now
478 }) {
479 let current = current.expect("checked current lease is present");
480 if current
481 .owner
482 .as_ref()
483 .is_some_and(|current_owner| current_owner.same_incarnation(&owner))
484 {
485 let expires_at = now.saturating_add(lease_ttl_ms);
486 tx.execute(
487 "UPDATE session_execution_leases
488 SET lease_expires_at_ms = ?2
489 WHERE session_id = ?1",
490 params![session_id, expires_at as i64],
491 )
492 .map_err(sqlite_error)?;
493 return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
494 SessionExecutionLease {
495 session_id,
496 owner,
497 lease_token: current.lease_token.expect("live lease token set"),
498 fencing_token: current.fencing_token,
499 claimed_at_epoch_ms: current.claimed_at_ms,
500 expires_at_epoch_ms: expires_at,
501 },
502 ));
503 }
504 return Ok(SessionExecutionLeaseClaimOutcome::Busy {
505 holder: row_to_session_execution_lease(&session_id, current)?,
506 });
507 }
508 Ok(SessionExecutionLeaseClaimOutcome::Acquired(
509 acquire_session_execution_lease_conn(
510 tx,
511 &session_id,
512 &owner,
513 current.as_ref().map_or(0, |lease| lease.fencing_token),
514 now,
515 lease_ttl_ms,
516 )?,
517 ))
518 })(
519 );
520 match outcome {
521 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
522 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
523 }
524 })
525 .await
526 .map_err(sqlite_error)?
527 }
528
529 async fn reclaim_session_execution_lease(
530 &self,
531 session_id: &str,
532 owner: &LeaseOwnerIdentity,
533 observed_holder: &SessionExecutionLeaseFence,
534 lease_ttl_ms: u64,
535 ) -> Result<SessionExecutionLeaseClaimOutcome, StoreError> {
536 let session_id = session_id.to_string();
537 let owner = owner.clone();
538 let observed_holder = observed_holder.clone();
539 self.conn
540 .write_flow(move |tx| {
541 let outcome: Result<SessionExecutionLeaseClaimOutcome, StoreError> = (|| {
542 let now = current_epoch_ms();
543 let current = load_session_execution_lease_row_conn(tx, &session_id)?;
544 let Some(current) = current else {
545 return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
546 acquire_session_execution_lease_conn(
547 tx,
548 &session_id,
549 &owner,
550 0,
551 now,
552 lease_ttl_ms,
553 )?,
554 ));
555 };
556 if current.lease_token.is_none() || current.expires_at_ms <= now {
557 return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
558 acquire_session_execution_lease_conn(
559 tx,
560 &session_id,
561 &owner,
562 current.fencing_token,
563 now,
564 lease_ttl_ms,
565 )?,
566 ));
567 }
568 let holder = row_to_session_execution_lease(&session_id, current)?;
569 if observed_holder.session_id == session_id
570 && holder.owner.same_incarnation(&observed_holder.owner)
571 && holder.lease_token == observed_holder.lease_token
572 && holder.fencing_token == observed_holder.fencing_token
573 && holder.owner.is_definitely_dead_for_claimant(&owner)
574 {
575 let fencing_token = holder.fencing_token.saturating_add(1);
576 let lease_token = format!(
577 "{}:{}:{}:{now}:{fencing_token}",
578 session_id, owner.owner_id, owner.incarnation_id
579 );
580 let expires_at = now.saturating_add(lease_ttl_ms);
581 let liveness_json = encode_liveness(&owner.liveness)?;
582 let changed = tx
583 .execute(
584 "UPDATE session_execution_leases
585 SET lease_owner_id = ?1,
586 lease_owner_incarnation_id = ?2,
587 lease_owner_liveness_json = ?3,
588 lease_token = ?4,
589 lease_fencing_token = ?5,
590 lease_claimed_at_ms = ?6,
591 lease_expires_at_ms = ?7
592 WHERE session_id = ?8
593 AND lease_owner_id = ?9
594 AND lease_owner_incarnation_id = ?10
595 AND lease_token = ?11
596 AND lease_fencing_token = ?12",
597 params![
598 owner.owner_id,
599 owner.incarnation_id,
600 liveness_json,
601 lease_token,
602 fencing_token as i64,
603 now as i64,
604 expires_at as i64,
605 session_id,
606 observed_holder.owner.owner_id,
607 observed_holder.owner.incarnation_id,
608 observed_holder.lease_token,
609 observed_holder.fencing_token as i64,
610 ],
611 )
612 .map_err(sqlite_error)?;
613 if changed == 1 {
614 return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
615 SessionExecutionLease {
616 session_id,
617 owner,
618 lease_token,
619 fencing_token,
620 claimed_at_epoch_ms: now,
621 expires_at_epoch_ms: expires_at,
622 },
623 ));
624 }
625 let current = load_session_execution_lease_row_conn(tx, &session_id)?;
626 if current.as_ref().is_some_and(|lease| {
627 lease.lease_token.is_some() && lease.expires_at_ms > now
628 }) {
629 let current = current.expect("checked current lease is present");
630 return Ok(SessionExecutionLeaseClaimOutcome::Busy {
631 holder: row_to_session_execution_lease(&session_id, current)?,
632 });
633 }
634 let previous_fencing_token =
635 current.as_ref().map_or(0, |lease| lease.fencing_token);
636 return Ok(SessionExecutionLeaseClaimOutcome::Acquired(
637 acquire_session_execution_lease_conn(
638 tx,
639 &session_id,
640 &owner,
641 previous_fencing_token,
642 now,
643 lease_ttl_ms,
644 )?,
645 ));
646 }
647 Ok(SessionExecutionLeaseClaimOutcome::Busy { holder })
648 })(
649 );
650 match outcome {
651 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
652 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
653 }
654 })
655 .await
656 .map_err(sqlite_error)?
657 }
658
659 async fn renew_session_execution_lease(
660 &self,
661 fence: &SessionExecutionLeaseFence,
662 lease_ttl_ms: u64,
663 ) -> Result<SessionExecutionLease, StoreError> {
664 let fence = fence.clone();
665 self.conn
666 .write_flow(move |tx| {
667 let outcome: Result<SessionExecutionLease, StoreError> = (|| {
668 let now = current_epoch_ms();
669 let current = load_session_execution_lease_row_conn(tx, &fence.session_id)?;
670 let Some(current) = current else {
671 return Err(StoreError::SessionExecutionLeaseExpired {
672 session_id: fence.session_id.clone(),
673 });
674 };
675 if !current
676 .owner
677 .as_ref()
678 .is_some_and(|owner| owner.same_incarnation(&fence.owner))
679 || current.lease_token.as_deref() != Some(fence.lease_token.as_str())
680 || current.fencing_token != fence.fencing_token
681 || current.expires_at_ms <= now
682 {
683 return Err(StoreError::SessionExecutionLeaseExpired {
684 session_id: fence.session_id.clone(),
685 });
686 }
687 let expires_at = now.saturating_add(lease_ttl_ms);
688 tx.execute(
689 "UPDATE session_execution_leases
690 SET lease_expires_at_ms = ?5
691 WHERE session_id = ?1
692 AND lease_owner_id = ?2
693 AND lease_owner_incarnation_id = ?3
694 AND lease_token = ?4
695 AND lease_fencing_token = ?6",
696 params![
697 fence.session_id,
698 fence.owner.owner_id,
699 fence.owner.incarnation_id,
700 fence.lease_token,
701 expires_at as i64,
702 fence.fencing_token as i64
703 ],
704 )
705 .map_err(sqlite_error)?;
706 Ok(SessionExecutionLease {
707 session_id: fence.session_id,
708 owner: fence.owner,
709 lease_token: fence.lease_token,
710 fencing_token: fence.fencing_token,
711 claimed_at_epoch_ms: current.claimed_at_ms,
712 expires_at_epoch_ms: expires_at,
713 })
714 })();
715 match outcome {
716 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
717 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
718 }
719 })
720 .await
721 .map_err(sqlite_error)?
722 }
723
724 async fn release_session_execution_lease(
725 &self,
726 completion: &SessionExecutionLeaseCompletion,
727 ) -> Result<(), StoreError> {
728 let completion = completion.clone();
729 self.conn
730 .write_flow(move |tx| {
731 let outcome = release_session_execution_lease_conn(tx, &completion);
732 match outcome {
733 Ok(()) => Ok(TxOutcome::Commit(Ok(()))),
734 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
735 }
736 })
737 .await
738 .map_err(sqlite_error)?
739 }
740
741 async fn enqueue_queued_work(
742 &self,
743 batch: QueuedWorkBatchDraft,
744 ) -> Result<QueuedWorkBatch, StoreError> {
745 let nonce = self.commit_count.fetch_add(1, AtomicOrdering::Relaxed);
746 self.conn
747 .write_flow(move |tx| {
748 let outcome: Result<QueuedWorkBatch, StoreError> = (|| {
749 if let Some(source_key) = batch.source_key.as_deref() {
750 let existing_id: Option<String> = tx
751 .query_row(
752 "SELECT batch_id
753 FROM queued_work_batches
754 WHERE session_id = ?1 AND source_key = ?2",
755 params![batch.session_id, source_key],
756 |row| row.get(0),
757 )
758 .optional()
759 .map_err(sqlite_error)?;
760 if let Some(batch_id) = existing_id {
761 let existing = load_queued_batch_by_id_conn(tx, &batch_id)?
762 .ok_or_else(|| {
763 StoreError::Backend(
764 "queued work source row disappeared".to_string(),
765 )
766 })?;
767 return Ok(existing);
768 }
769 }
770 let now = current_epoch_ms();
771 let batch_id =
772 derive_batch_id(&batch.session_id, batch.source_key.as_deref(), now, Some(nonce));
773 tx.execute(
774 "INSERT INTO queued_work_batches (
775 batch_id, session_id, source_key, delivery_policy, slot_policy,
776 merge_key_json, available_at_ms, enqueued_at_ms
777 )
778 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
779 params![
780 batch_id,
781 batch.session_id,
782 batch.source_key.as_deref(),
783 batch.delivery_policy.as_str(),
784 batch.slot_policy.as_str(),
785 encode_json(&batch.merge_key),
786 batch.available_at_ms as i64,
787 now as i64,
788 ],
789 )
790 .map_err(sqlite_error)?;
791 for (index, payload) in batch.payloads.iter().enumerate() {
792 let item_id = format!("{batch_id}:item:{index}");
793 tx.execute(
794 "INSERT INTO queued_work_items (batch_id, item_index, item_id, payload_json)
795 VALUES (?1, ?2, ?3, ?4)",
796 params![batch_id, index as i64, item_id, encode_json(payload)],
797 )
798 .map_err(sqlite_error)?;
799 }
800 load_queued_batch_by_id_conn(tx, &batch_id)?.ok_or_else(|| {
801 StoreError::Backend("queued work insert disappeared".to_string())
802 })
803 })();
804 match outcome {
807 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
808 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
809 }
810 })
811 .await
812 .map_err(sqlite_error)?
813 }
814
815 async fn claim_leading_ready_session_command(
816 &self,
817 session_id: &str,
818 session_execution_lease: &SessionExecutionLeaseFence,
819 owner: &LeaseOwnerIdentity,
820 lease_ttl_ms: u64,
821 ) -> Result<Option<QueuedWorkClaim>, StoreError> {
822 let session_id = session_id.to_string();
823 let session_execution_lease = session_execution_lease.clone();
824 let owner = owner.clone();
825 self.conn
826 .write_flow(move |tx| {
827 let outcome: Result<TxOutcome<Option<QueuedWorkClaim>>, StoreError> = (|| {
828 ensure_session_execution_lease_conn(
829 tx,
830 &session_id,
831 &session_execution_lease,
832 )?;
833 let now = current_epoch_ms();
834 let candidate_rows = {
835 let mut stmt = tx
836 .prepare(
837 "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
838 slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
839 claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
840 claim_owner_liveness_json, claim_token, claim_expires_at_ms
841 FROM queued_work_batches
842 WHERE session_id = ?1
843 AND available_at_ms <= ?2
844 ORDER BY enqueue_seq ASC
845 LIMIT ?3",
846 )
847 .map_err(sqlite_error)?;
848 let rows = stmt
849 .query_map(
850 params![session_id, now as i64, claim_scan_limit(1)],
851 queued_batch_row_from_sql,
852 )
853 .map_err(sqlite_error)?;
854 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
855 };
856 let candidate_rows = candidate_rows
857 .into_iter()
858 .filter(|row| {
859 row.claim_token.is_none()
860 || row.claim_expires_at_ms <= now
861 || row
862 .claim_owner
863 .as_ref()
864 .is_some_and(|holder| holder.is_definitely_dead_for_claimant(&owner))
865 })
866 .collect::<Vec<_>>();
867 let candidate_batches = candidate_rows
868 .iter()
869 .map(|row| queued_work_batch_from_conn(tx, row.clone()))
870 .collect::<Result<Vec<_>, StoreError>>()?;
871 let candidates = candidate_rows
872 .iter()
873 .zip(candidate_batches.iter())
874 .map(|(row, batch)| {
875 Ok(ClaimCandidate {
876 enqueue_seq: row.enqueue_seq,
877 claim_fencing_token: row.claim_fencing_token,
878 work_class: batch.work_class().ok_or_else(|| {
879 StoreError::Backend(format!(
880 "queued-work batch `{}` has mixed or empty payload classes",
881 batch.batch_id
882 ))
883 })?,
884 delivery_policy: decode_delivery_policy(
885 row.delivery_policy.clone(),
886 )?,
887 slot_policy: decode_slot_policy(row.slot_policy.clone())?,
888 merge_key: decode_merge_key(row.merge_key_json.clone())?,
889 })
890 })
891 .collect::<Result<Vec<_>, StoreError>>()?;
892 let selected_len = select_leading_session_command(&candidates);
893 if selected_len == 0 {
894 return Ok(TxOutcome::Commit(None));
895 }
896 let mut selected = candidate_rows;
897 selected.truncate(selected_len);
898 let mut selected_batches = candidate_batches;
899 selected_batches.truncate(selected_len);
900 let lease = QueuedWorkClaimLease::derive(
901 &candidates[0],
902 &session_id,
903 &owner,
904 now,
905 lease_ttl_ms,
906 );
907 let liveness_json = encode_liveness(&owner.liveness)?;
908 for row in &selected {
909 let claimed = tx
910 .execute(
911 "UPDATE queued_work_batches
912 SET claim_id = ?3,
913 claim_owner_id = ?4,
914 claim_owner_incarnation_id = ?5,
915 claim_owner_liveness_json = ?6,
916 claim_token = ?7,
917 claim_fencing_token = claim_fencing_token + 1,
918 claim_claimed_at_ms = ?8,
919 claim_expires_at_ms = ?9
920 WHERE session_id = ?1
921 AND batch_id = ?2
922 AND (
923 claim_token IS NULL
924 OR claim_expires_at_ms <= ?8
925 OR (
926 claim_token = ?10
927 AND claim_owner_id = ?11
928 AND claim_owner_incarnation_id = ?12
929 )
930 )",
931 params![
932 session_id,
933 row.batch_id,
934 lease.claim_id,
935 owner.owner_id.as_str(),
936 owner.incarnation_id.as_str(),
937 liveness_json.as_str(),
938 lease.lease_token,
939 now as i64,
940 lease.expires_at_epoch_ms as i64,
941 row.claim_token,
942 row.claim_owner.as_ref().map(|owner| owner.owner_id.as_str()),
943 row.claim_owner
944 .as_ref()
945 .map(|owner| owner.incarnation_id.as_str())
946 ],
947 )
948 .map_err(sqlite_error)?;
949 if claimed == 0 {
950 return Ok(TxOutcome::Rollback(None));
951 }
952 }
953 Ok(TxOutcome::Commit(Some(QueuedWorkClaim {
954 session_id: session_id.clone(),
955 claim_id: lease.claim_id,
956 owner: owner.clone(),
957 lease_token: lease.lease_token,
958 fencing_token: lease.fencing_token,
959 claimed_at_epoch_ms: lease.claimed_at_epoch_ms,
960 expires_at_epoch_ms: lease.expires_at_epoch_ms,
961 batches: selected_batches,
962 })))
963 })();
964 match outcome {
965 Ok(TxOutcome::Commit(value)) => Ok(TxOutcome::Commit(Ok(value))),
966 Ok(TxOutcome::Rollback(value)) => Ok(TxOutcome::Rollback(Ok(value))),
967 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
968 }
969 })
970 .await
971 .map_err(sqlite_error)?
972 }
973
974 async fn claim_ready_queued_work(
975 &self,
976 session_id: &str,
977 session_execution_lease: &SessionExecutionLeaseFence,
978 owner: &LeaseOwnerIdentity,
979 boundary: QueuedWorkClaimBoundary,
980 lease_ttl_ms: u64,
981 max_batches: usize,
982 ) -> Result<Option<QueuedWorkClaim>, StoreError> {
983 if max_batches == 0 {
984 return Ok(None);
985 }
986 let session_id = session_id.to_string();
987 let session_execution_lease = session_execution_lease.clone();
988 let owner = owner.clone();
989 self.conn
990 .write_flow(move |tx| {
991 let outcome: Result<TxOutcome<Option<QueuedWorkClaim>>, StoreError> = (|| {
992 ensure_session_execution_lease_conn(
993 tx,
994 &session_id,
995 &session_execution_lease,
996 )?;
997 let now = current_epoch_ms();
998 let candidate_rows = {
999 let mut stmt = tx
1000 .prepare(
1001 "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
1002 slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
1003 claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
1004 claim_owner_liveness_json, claim_token, claim_expires_at_ms
1005 FROM queued_work_batches
1006 WHERE session_id = ?1
1007 AND available_at_ms <= ?2
1008 ORDER BY enqueue_seq ASC
1009 LIMIT ?3",
1010 )
1011 .map_err(sqlite_error)?;
1012 let rows = stmt
1013 .query_map(
1014 params![session_id, now as i64, claim_scan_limit(max_batches)],
1015 queued_batch_row_from_sql,
1016 )
1017 .map_err(sqlite_error)?;
1018 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1019 };
1020 let candidate_rows = candidate_rows
1021 .into_iter()
1022 .filter(|row| {
1023 row.claim_token.is_none()
1024 || row.claim_expires_at_ms <= now
1025 || row
1026 .claim_owner
1027 .as_ref()
1028 .is_some_and(|holder| holder.is_definitely_dead_for_claimant(&owner))
1029 })
1030 .collect::<Vec<_>>();
1031 let candidate_batches = candidate_rows
1032 .iter()
1033 .map(|row| queued_work_batch_from_conn(tx, row.clone()))
1034 .collect::<Result<Vec<_>, StoreError>>()?;
1035 let candidates = candidate_rows
1036 .iter()
1037 .zip(candidate_batches.iter())
1038 .map(|(row, batch)| {
1039 Ok(ClaimCandidate {
1040 enqueue_seq: row.enqueue_seq,
1041 claim_fencing_token: row.claim_fencing_token,
1042 work_class: batch.work_class().ok_or_else(|| {
1043 StoreError::Backend(format!(
1044 "queued-work batch `{}` has mixed or empty payload classes",
1045 batch.batch_id
1046 ))
1047 })?,
1048 delivery_policy: decode_delivery_policy(
1049 row.delivery_policy.clone(),
1050 )?,
1051 slot_policy: decode_slot_policy(row.slot_policy.clone())?,
1052 merge_key: decode_merge_key(row.merge_key_json.clone())?,
1053 })
1054 })
1055 .collect::<Result<Vec<_>, StoreError>>()?;
1056 let selected_len =
1057 select_turn_work_claim_prefix(&candidates, boundary, max_batches);
1058 if selected_len == 0 {
1059 return Ok(TxOutcome::Commit(None));
1060 }
1061 let mut selected = candidate_rows;
1062 selected.truncate(selected_len);
1063 let mut selected_batches = candidate_batches;
1064 selected_batches.truncate(selected_len);
1065 let lease = QueuedWorkClaimLease::derive(
1066 &candidates[0],
1067 &session_id,
1068 &owner,
1069 now,
1070 lease_ttl_ms,
1071 );
1072 let liveness_json = encode_liveness(&owner.liveness)?;
1073 for row in &selected {
1074 let claimed = tx
1083 .execute(
1084 "UPDATE queued_work_batches
1085 SET claim_id = ?3,
1086 claim_owner_id = ?4,
1087 claim_owner_incarnation_id = ?5,
1088 claim_owner_liveness_json = ?6,
1089 claim_token = ?7,
1090 claim_fencing_token = claim_fencing_token + 1,
1091 claim_claimed_at_ms = ?8,
1092 claim_expires_at_ms = ?9
1093 WHERE session_id = ?1
1094 AND batch_id = ?2
1095 AND (
1096 claim_token IS NULL
1097 OR claim_expires_at_ms <= ?8
1098 OR (
1099 claim_token = ?10
1100 AND claim_owner_id = ?11
1101 AND claim_owner_incarnation_id = ?12
1102 )
1103 )",
1104 params![
1105 session_id,
1106 row.batch_id,
1107 lease.claim_id,
1108 owner.owner_id.as_str(),
1109 owner.incarnation_id.as_str(),
1110 liveness_json.as_str(),
1111 lease.lease_token,
1112 now as i64,
1113 lease.expires_at_epoch_ms as i64,
1114 row.claim_token,
1115 row.claim_owner.as_ref().map(|owner| owner.owner_id.as_str()),
1116 row.claim_owner
1117 .as_ref()
1118 .map(|owner| owner.incarnation_id.as_str())
1119 ],
1120 )
1121 .map_err(sqlite_error)?;
1122 if claimed == 0 {
1123 return Ok(TxOutcome::Rollback(None));
1127 }
1128 }
1129 Ok(TxOutcome::Commit(Some(QueuedWorkClaim {
1130 session_id: session_id.clone(),
1131 claim_id: lease.claim_id,
1132 owner: owner.clone(),
1133 lease_token: lease.lease_token,
1134 fencing_token: lease.fencing_token,
1135 claimed_at_epoch_ms: lease.claimed_at_epoch_ms,
1136 expires_at_epoch_ms: lease.expires_at_epoch_ms,
1137 batches: selected_batches,
1138 })))
1139 })();
1140 match outcome {
1144 Ok(TxOutcome::Commit(value)) => Ok(TxOutcome::Commit(Ok(value))),
1145 Ok(TxOutcome::Rollback(value)) => Ok(TxOutcome::Rollback(Ok(value))),
1146 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1147 }
1148 })
1149 .await
1150 .map_err(sqlite_error)?
1151 }
1152
1153 async fn renew_queued_work_claim(
1154 &self,
1155 claim: &QueuedWorkClaim,
1156 lease_ttl_ms: u64,
1157 ) -> Result<QueuedWorkClaim, StoreError> {
1158 let now = current_epoch_ms();
1159 let expires_at = now.saturating_add(lease_ttl_ms);
1160 let session_id = claim.session_id.clone();
1161 let claim_id = claim.claim_id.clone();
1162 let lease_token = claim.lease_token.clone();
1163 let changed = self
1164 .conn
1165 .write(move |tx| {
1166 tx.execute(
1167 "UPDATE queued_work_batches
1168 SET claim_expires_at_ms = ?4
1169 WHERE session_id = ?1 AND claim_id = ?2 AND claim_token = ?3",
1170 params![session_id, claim_id, lease_token, expires_at as i64],
1171 )
1172 })
1173 .await
1174 .map_err(sqlite_error)?;
1175 renewed_claim(claim, changed, expires_at)
1176 }
1177
1178 async fn abandon_queued_work_claim(&self, claim: &QueuedWorkClaim) -> Result<(), StoreError> {
1179 let session_id = claim.session_id.clone();
1180 let claim_id = claim.claim_id.clone();
1181 let lease_token = claim.lease_token.clone();
1182 self.conn
1183 .write(move |tx| {
1184 tx.execute(
1185 "UPDATE queued_work_batches
1186 SET claim_id = NULL,
1187 claim_owner_id = NULL,
1188 claim_owner_incarnation_id = NULL,
1189 claim_owner_liveness_json = NULL,
1190 claim_token = NULL,
1191 claim_claimed_at_ms = 0,
1192 claim_expires_at_ms = 0
1193 WHERE session_id = ?1 AND claim_id = ?2 AND claim_token = ?3",
1194 params![session_id, claim_id, lease_token],
1195 )
1196 })
1197 .await
1198 .map_err(sqlite_error)?;
1199 Ok(())
1200 }
1201
1202 async fn cancel_queued_work_batch(
1203 &self,
1204 session_id: &str,
1205 batch_id: &str,
1206 ) -> Result<Option<QueuedWorkBatch>, StoreError> {
1207 let session_id = session_id.to_string();
1208 let batch_id = batch_id.to_string();
1209 self.conn
1210 .write_flow(move |tx| {
1211 let outcome: Result<Option<QueuedWorkBatch>, StoreError> = (|| {
1212 let now = current_epoch_ms() as i64;
1213 let row = tx
1214 .query_row(
1215 "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
1216 slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
1217 claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
1218 claim_owner_liveness_json, claim_token, claim_expires_at_ms
1219 FROM queued_work_batches
1220 WHERE session_id = ?1
1221 AND batch_id = ?2
1222 AND (claim_token IS NULL OR claim_expires_at_ms <= ?3)",
1223 params![session_id, batch_id, now],
1224 queued_batch_row_from_sql,
1225 )
1226 .optional()
1227 .map_err(sqlite_error)?;
1228 let Some(row) = row else {
1229 return Ok(None);
1230 };
1231 let batch = queued_work_batch_from_conn(tx, row)?;
1232 tx.execute(
1233 "DELETE FROM queued_work_batches
1234 WHERE session_id = ?1
1235 AND batch_id = ?2
1236 AND (claim_token IS NULL OR claim_expires_at_ms <= ?3)",
1237 params![session_id, batch_id, now],
1238 )
1239 .map_err(sqlite_error)?;
1240 Ok(Some(batch))
1241 })();
1242 match outcome {
1243 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1244 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1245 }
1246 })
1247 .await
1248 .map_err(sqlite_error)?
1249 }
1250
1251 async fn list_queued_work(&self, session_id: &str) -> Result<Vec<QueuedWorkBatch>, StoreError> {
1252 let session_id = session_id.to_string();
1253 self.conn
1254 .call(move |conn| {
1255 let outcome: Result<Vec<QueuedWorkBatch>, StoreError> = (|| {
1256 let rows = {
1257 let mut stmt = conn
1258 .prepare(
1259 "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
1260 slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
1261 claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
1262 claim_owner_liveness_json, claim_token, claim_expires_at_ms
1263 FROM queued_work_batches
1264 WHERE session_id = ?1
1265 ORDER BY enqueue_seq ASC",
1266 )
1267 .map_err(sqlite_error)?;
1268 let rows = stmt
1269 .query_map(params![session_id], queued_batch_row_from_sql)
1270 .map_err(sqlite_error)?;
1271 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1272 };
1273 rows.into_iter()
1274 .map(|row| queued_work_batch_from_conn(conn, row))
1275 .collect()
1276 })();
1277 Ok(outcome)
1278 })
1279 .await
1280 .map_err(sqlite_error)?
1281 }
1282
1283 async fn list_pending_queued_work(
1284 &self,
1285 session_id: &str,
1286 ) -> Result<Vec<QueuedWorkBatch>, StoreError> {
1287 let session_id = session_id.to_string();
1288 self.conn
1289 .call(move |conn| {
1290 let outcome: Result<Vec<QueuedWorkBatch>, StoreError> = (|| {
1291 let now = current_epoch_ms();
1292 let rows = {
1293 let mut stmt = conn
1294 .prepare(
1295 "SELECT enqueue_seq, batch_id, session_id, source_key, delivery_policy,
1296 slot_policy, merge_key_json, available_at_ms, enqueued_at_ms,
1297 claim_fencing_token, claim_owner_id, claim_owner_incarnation_id,
1298 claim_owner_liveness_json, claim_token, claim_expires_at_ms
1299 FROM queued_work_batches
1300 WHERE session_id = ?1
1301 AND (claim_token IS NULL OR claim_expires_at_ms <= ?2)
1302 ORDER BY enqueue_seq ASC",
1303 )
1304 .map_err(sqlite_error)?;
1305 let rows = stmt
1306 .query_map(
1307 params![session_id, now as i64],
1308 queued_batch_row_from_sql,
1309 )
1310 .map_err(sqlite_error)?;
1311 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1312 };
1313 rows.into_iter()
1314 .map(|row| queued_work_batch_from_conn(conn, row))
1315 .collect()
1316 })();
1317 Ok(outcome)
1318 })
1319 .await
1320 .map_err(sqlite_error)?
1321 }
1322
1323 async fn enqueue_pending_turn_input(
1324 &self,
1325 draft: lash_core::PendingTurnInputDraft,
1326 ) -> Result<lash_core::PendingTurnInput, StoreError> {
1327 let nonce = self.commit_count.fetch_add(1, AtomicOrdering::Relaxed);
1328 self.conn
1329 .write_flow(move |tx| {
1330 let outcome: Result<lash_core::PendingTurnInput, StoreError> = (|| {
1331 if let Some(source_key) = draft.source_key.as_deref() {
1332 let existing_id: Option<String> = tx
1333 .query_row(
1334 "SELECT input_id
1335 FROM pending_turn_inputs
1336 WHERE session_id = ?1 AND source_key = ?2",
1337 params![draft.session_id, source_key],
1338 |row| row.get(0),
1339 )
1340 .optional()
1341 .map_err(sqlite_error)?;
1342 if let Some(input_id) = existing_id {
1343 let existing = load_pending_turn_input_by_id_conn(
1344 tx,
1345 &draft.session_id,
1346 &input_id,
1347 )?
1348 .ok_or_else(|| {
1349 StoreError::Backend(
1350 "pending turn input source row disappeared".to_string(),
1351 )
1352 })?;
1353 if !draft.submitted_content_matches(&existing).map_err(|err| {
1354 StoreError::Backend(format!(
1355 "failed to compare pending turn input submission: {err}"
1356 ))
1357 })? {
1358 return Err(StoreError::PendingTurnInputSourceKeyConflict {
1359 session_id: draft.session_id.clone(),
1360 source_key: source_key.to_string(),
1361 existing_input_id: existing.input_id.clone(),
1362 });
1363 }
1364 return Ok(existing);
1365 }
1366 }
1367 let now = current_epoch_ms();
1368 let input_id = draft.input_id.clone().unwrap_or_else(|| {
1369 derive_pending_turn_input_id(
1370 &draft.session_id,
1371 draft.source_key.as_deref(),
1372 now,
1373 nonce,
1374 )
1375 });
1376 let state = match draft.ingress {
1377 lash_core::TurnInputIngress::ActiveTurn { .. } => {
1378 lash_core::TurnInputState::PendingActive
1379 }
1380 lash_core::TurnInputIngress::NextTurn => {
1381 lash_core::TurnInputState::DeferredNextTurn
1382 }
1383 };
1384 tx.execute(
1385 "INSERT INTO pending_turn_inputs (
1386 input_id, session_id, source_key, ingress_json, state,
1387 input_json, enqueued_at_ms
1388 )
1389 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
1390 params![
1391 input_id,
1392 draft.session_id,
1393 draft.source_key.as_deref(),
1394 encode_json(&draft.ingress),
1395 state.as_str(),
1396 encode_json(&draft.input),
1397 now as i64,
1398 ],
1399 )
1400 .map_err(sqlite_error)?;
1401 load_pending_turn_input_by_id_conn(tx, &draft.session_id, &input_id)?
1402 .ok_or_else(|| {
1403 StoreError::Backend("pending turn input insert disappeared".to_string())
1404 })
1405 })();
1406 match outcome {
1407 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1408 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1409 }
1410 })
1411 .await
1412 .map_err(sqlite_error)?
1413 }
1414
1415 async fn list_pending_turn_inputs(
1416 &self,
1417 session_id: &str,
1418 ) -> Result<Vec<lash_core::PendingTurnInput>, StoreError> {
1419 let session_id = session_id.to_string();
1420 self.conn
1421 .call(move |conn| {
1422 let outcome: Result<Vec<lash_core::PendingTurnInput>, StoreError> = (|| {
1423 let now = current_epoch_ms();
1424 let rows = {
1425 let mut stmt = conn
1426 .prepare(
1427 "SELECT enqueue_seq, input_id, session_id, source_key, ingress_json,
1428 state, input_json, enqueued_at_ms, claim_id, claim_fencing_token,
1429 claim_owner_id, claim_owner_incarnation_id,
1430 claim_owner_liveness_json, claim_token, claim_expires_at_ms
1431 FROM pending_turn_inputs
1432 WHERE session_id = ?1
1433 AND state IN (?2, ?3)
1434 AND (claim_token IS NULL OR claim_expires_at_ms <= ?4)
1435 ORDER BY enqueue_seq ASC",
1436 )
1437 .map_err(sqlite_error)?;
1438 let rows = stmt
1439 .query_map(
1440 params![
1441 session_id,
1442 lash_core::TurnInputState::PendingActive.as_str(),
1443 lash_core::TurnInputState::DeferredNextTurn.as_str(),
1444 now as i64
1445 ],
1446 pending_turn_input_row_from_sql,
1447 )
1448 .map_err(sqlite_error)?;
1449 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1450 };
1451 rows.into_iter().map(pending_turn_input_from_row).collect()
1452 })(
1453 );
1454 Ok(outcome)
1455 })
1456 .await
1457 .map_err(sqlite_error)?
1458 }
1459
1460 async fn cancel_pending_turn_input(
1461 &self,
1462 session_id: &str,
1463 input_id: &str,
1464 ) -> Result<lash_core::PendingTurnInputCancelOutcome, StoreError> {
1465 let target = lash_core::PendingTurnInputCancelTarget::input_id(input_id);
1466 let targets = vec![target];
1467 let mut outcomes = self
1468 .cancel_pending_turn_inputs(session_id, &targets)
1469 .await?;
1470 Ok(outcomes
1471 .pop()
1472 .map(|result| result.outcome)
1473 .unwrap_or(lash_core::PendingTurnInputCancelOutcome::NotFound))
1474 }
1475
1476 async fn cancel_pending_turn_inputs(
1477 &self,
1478 session_id: &str,
1479 targets: &[lash_core::PendingTurnInputCancelTarget],
1480 ) -> Result<Vec<lash_core::PendingTurnInputCancelResult>, StoreError> {
1481 let session_id = session_id.to_string();
1482 let targets = targets.to_vec();
1483 self.conn
1484 .write_flow(move |tx| {
1485 let outcome: Result<Vec<lash_core::PendingTurnInputCancelResult>, StoreError> =
1486 (|| {
1487 let now = current_epoch_ms();
1488 let mut results = Vec::with_capacity(targets.len());
1489 for target in targets {
1490 let outcome = match load_pending_turn_input_row_by_target_conn(
1491 tx,
1492 &session_id,
1493 &target,
1494 )? {
1495 Some(row) => cancel_pending_turn_input_row_conn(tx, row, now)?,
1496 None => lash_core::PendingTurnInputCancelOutcome::NotFound,
1497 };
1498 results
1499 .push(lash_core::PendingTurnInputCancelResult { target, outcome });
1500 }
1501 Ok(results)
1502 })();
1503 match outcome {
1504 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1505 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1506 }
1507 })
1508 .await
1509 .map_err(sqlite_error)?
1510 }
1511
1512 async fn cancel_pending_turn_input_suffix(
1513 &self,
1514 session_id: &str,
1515 anchor: &lash_core::PendingTurnInputCancelTarget,
1516 ) -> Result<lash_core::PendingTurnInputSuffixCancelOutcome, StoreError> {
1517 let session_id = session_id.to_string();
1518 let anchor = anchor.clone();
1519 self.conn
1520 .write_flow(move |tx| {
1521 let outcome: Result<lash_core::PendingTurnInputSuffixCancelOutcome, StoreError> =
1522 (|| {
1523 let now = current_epoch_ms();
1524 let Some(anchor_row) =
1525 load_pending_turn_input_row_by_target_conn(tx, &session_id, &anchor)?
1526 else {
1527 return Ok(
1528 lash_core::PendingTurnInputSuffixCancelOutcome::AnchorNotFound {
1529 anchor,
1530 },
1531 );
1532 };
1533 let rows = {
1534 let mut stmt = tx
1535 .prepare(
1536 "SELECT enqueue_seq, input_id, session_id, source_key, ingress_json,
1537 state, input_json, enqueued_at_ms, claim_id, claim_fencing_token,
1538 claim_owner_id, claim_owner_incarnation_id,
1539 claim_owner_liveness_json, claim_token, claim_expires_at_ms
1540 FROM pending_turn_inputs
1541 WHERE session_id = ?1 AND enqueue_seq >= ?2
1542 ORDER BY enqueue_seq ASC",
1543 )
1544 .map_err(sqlite_error)?;
1545 let rows = stmt
1546 .query_map(
1547 params![session_id, anchor_row.enqueue_seq as i64],
1548 pending_turn_input_row_from_sql,
1549 )
1550 .map_err(sqlite_error)?;
1551 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1552 };
1553 let mut outcomes = Vec::with_capacity(rows.len());
1554 for row in rows {
1555 outcomes.push(cancel_pending_turn_input_row_conn(tx, row, now)?);
1556 }
1557 Ok(lash_core::PendingTurnInputSuffixCancelOutcome::Outcomes {
1558 anchor,
1559 outcomes,
1560 })
1561 })();
1562 match outcome {
1563 Ok(value) => Ok(TxOutcome::Commit(Ok(value))),
1564 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1565 }
1566 })
1567 .await
1568 .map_err(sqlite_error)?
1569 }
1570
1571 async fn claim_active_turn_inputs(
1572 &self,
1573 session_id: &str,
1574 session_execution_lease: &SessionExecutionLeaseFence,
1575 owner: &LeaseOwnerIdentity,
1576 turn_id: &str,
1577 checkpoint: lash_core::CheckpointKind,
1578 lease_ttl_ms: u64,
1579 max_inputs: usize,
1580 ) -> Result<Option<lash_core::TurnInputClaim>, StoreError> {
1581 claim_pending_turn_inputs_sqlite(
1582 &self.conn,
1583 session_id,
1584 session_execution_lease,
1585 owner,
1586 lease_ttl_ms,
1587 max_inputs,
1588 lash_core::TurnInputClaimMode::ActiveTurn {
1589 turn_id: turn_id.to_string(),
1590 checkpoint,
1591 },
1592 )
1593 .await
1594 }
1595
1596 async fn claim_next_turn_inputs(
1597 &self,
1598 session_id: &str,
1599 session_execution_lease: &SessionExecutionLeaseFence,
1600 owner: &LeaseOwnerIdentity,
1601 lease_ttl_ms: u64,
1602 max_inputs: usize,
1603 ) -> Result<Option<lash_core::TurnInputClaim>, StoreError> {
1604 claim_pending_turn_inputs_sqlite(
1605 &self.conn,
1606 session_id,
1607 session_execution_lease,
1608 owner,
1609 lease_ttl_ms,
1610 max_inputs,
1611 lash_core::TurnInputClaimMode::NextTurn,
1612 )
1613 .await
1614 }
1615
1616 async fn abandon_turn_input_claim(
1617 &self,
1618 claim: &lash_core::TurnInputClaim,
1619 ) -> Result<(), StoreError> {
1620 let session_id = claim.session_id.clone();
1621 let claim_id = claim.claim_id.clone();
1622 let lease_token = claim.lease_token.clone();
1623 let restored_state = match claim.mode {
1624 lash_core::TurnInputClaimMode::ActiveTurn { .. } => {
1625 lash_core::TurnInputState::PendingActive
1626 }
1627 lash_core::TurnInputClaimMode::NextTurn => lash_core::TurnInputState::DeferredNextTurn,
1628 };
1629 self.conn
1630 .write(move |tx| {
1631 tx.execute(
1632 "UPDATE pending_turn_inputs
1633 SET state = CASE
1634 WHEN state = ?4 THEN ?5
1635 ELSE state
1636 END,
1637 claim_id = NULL,
1638 claim_owner_id = NULL,
1639 claim_owner_incarnation_id = NULL,
1640 claim_owner_liveness_json = NULL,
1641 claim_token = NULL,
1642 claim_claimed_at_ms = 0,
1643 claim_expires_at_ms = 0
1644 WHERE session_id = ?1 AND claim_id = ?2 AND claim_token = ?3",
1645 params![
1646 session_id,
1647 claim_id,
1648 lease_token,
1649 lash_core::TurnInputState::Accepted.as_str(),
1650 restored_state.as_str(),
1651 ],
1652 )
1653 })
1654 .await
1655 .map_err(sqlite_error)?;
1656 Ok(())
1657 }
1658
1659 async fn save_session_meta(&self, meta: SessionMeta) -> Result<(), StoreError> {
1660 Store::save_session_meta(self, meta).await;
1661 Ok(())
1662 }
1663
1664 async fn load_session_meta(&self) -> Result<Option<SessionMeta>, StoreError> {
1665 Ok(Store::load_session_meta(self).await)
1666 }
1667
1668 async fn tombstone_nodes(&self, ids: &[String]) -> Result<(), StoreError> {
1669 if ids.is_empty() {
1670 return Ok(());
1671 }
1672 let ids = ids.to_vec();
1673 self.conn
1674 .write(move |tx| {
1675 let mut stmt =
1676 tx.prepare("UPDATE graph_nodes SET tombstoned = 1 WHERE node_id = ?1")?;
1677 for id in &ids {
1678 stmt.execute(params![id])?;
1679 }
1680 Ok(())
1681 })
1682 .await
1683 .map_err(sqlite_error)
1684 }
1685
1686 async fn vacuum(&self) -> Result<VacuumReport, StoreError> {
1687 let (removed_node_count, removed_pending_turn_input_tombstone_count) = self
1688 .conn
1689 .write(move |tx| {
1690 let removed_node_count =
1691 tx.execute("DELETE FROM graph_nodes WHERE tombstoned = 1", [])?;
1692 let removed_pending_turn_input_tombstone_count = tx.execute(
1693 "DELETE FROM pending_turn_inputs
1694 WHERE state IN (?1, ?2)",
1695 params![
1696 lash_core::TurnInputState::Cancelled.as_str(),
1697 lash_core::TurnInputState::Completed.as_str()
1698 ],
1699 )?;
1700 Ok((
1701 removed_node_count,
1702 removed_pending_turn_input_tombstone_count,
1703 ))
1704 })
1705 .await
1706 .map_err(sqlite_error)?;
1707 Ok(VacuumReport {
1708 removed_node_count,
1709 removed_pending_turn_input_tombstone_count,
1710 })
1711 }
1712
1713 async fn gc_unreachable(&self) -> Result<GcReport, StoreError> {
1714 Ok(Store::gc_unreachable(self).await)
1715 }
1716}
1717
1718fn derive_pending_turn_input_id(
1719 session_id: &str,
1720 source_key: Option<&str>,
1721 now_epoch_ms: u64,
1722 nonce: u64,
1723) -> String {
1724 format!(
1725 "ti:{:x}",
1726 Sha256::digest(format!("{session_id}:{source_key:?}:{now_epoch_ms}:{nonce}").as_bytes())
1727 )
1728}
1729
1730fn cancel_pending_turn_input_row_conn(
1731 conn: &Connection,
1732 row: PendingTurnInputRow,
1733 now_epoch_ms: u64,
1734) -> Result<lash_core::PendingTurnInputCancelOutcome, StoreError> {
1735 let mut input = pending_turn_input_from_row(row.clone())?;
1736 match input.state {
1737 lash_core::TurnInputState::Cancelled => Ok(
1738 lash_core::PendingTurnInputCancelOutcome::AlreadyCancelled(input),
1739 ),
1740 lash_core::TurnInputState::Completed => Ok(
1741 lash_core::PendingTurnInputCancelOutcome::AlreadyCompleted(input),
1742 ),
1743 lash_core::TurnInputState::Accepted => {
1744 Ok(lash_core::PendingTurnInputCancelOutcome::AlreadyClaimed {
1745 claim: pending_turn_input_claim_diagnostics_from_row(&row, input.state),
1746 input,
1747 })
1748 }
1749 lash_core::TurnInputState::PendingActive | lash_core::TurnInputState::DeferredNextTurn => {
1750 let live_claim = row.claim_token.is_some() && row.claim_expires_at_ms > now_epoch_ms;
1751 if live_claim {
1752 return Ok(lash_core::PendingTurnInputCancelOutcome::AlreadyClaimed {
1753 claim: pending_turn_input_claim_diagnostics_from_row(&row, input.state),
1754 input,
1755 });
1756 }
1757 conn.execute(
1758 "UPDATE pending_turn_inputs
1759 SET state = ?3,
1760 claim_id = NULL,
1761 claim_owner_id = NULL,
1762 claim_owner_incarnation_id = NULL,
1763 claim_owner_liveness_json = NULL,
1764 claim_token = NULL,
1765 claim_claimed_at_ms = 0,
1766 claim_expires_at_ms = 0
1767 WHERE session_id = ?1 AND input_id = ?2",
1768 params![
1769 row.session_id,
1770 row.input_id,
1771 lash_core::TurnInputState::Cancelled.as_str(),
1772 ],
1773 )
1774 .map_err(sqlite_error)?;
1775 input.state = lash_core::TurnInputState::Cancelled;
1776 Ok(lash_core::PendingTurnInputCancelOutcome::Cancelled(input))
1777 }
1778 }
1779}
1780
1781async fn claim_pending_turn_inputs_sqlite(
1782 conn: &SqliteConnection,
1783 session_id: &str,
1784 session_execution_lease: &SessionExecutionLeaseFence,
1785 owner: &LeaseOwnerIdentity,
1786 lease_ttl_ms: u64,
1787 max_inputs: usize,
1788 mode: lash_core::TurnInputClaimMode,
1789) -> Result<Option<lash_core::TurnInputClaim>, StoreError> {
1790 if max_inputs == 0 {
1791 return Ok(None);
1792 }
1793 let session_id = session_id.to_string();
1794 let session_execution_lease = session_execution_lease.clone();
1795 let owner = owner.clone();
1796 conn.write_flow(move |tx| {
1797 let outcome: Result<TxOutcome<Option<lash_core::TurnInputClaim>>, StoreError> = (|| {
1798 ensure_session_execution_lease_conn(tx, &session_id, &session_execution_lease)?;
1799 let now = current_epoch_ms();
1800 let wanted_state = match &mode {
1801 lash_core::TurnInputClaimMode::ActiveTurn { .. } => {
1802 lash_core::TurnInputState::PendingActive
1803 }
1804 lash_core::TurnInputClaimMode::NextTurn => {
1805 lash_core::TurnInputState::DeferredNextTurn
1806 }
1807 };
1808 let candidate_rows = {
1809 let mut stmt = tx
1810 .prepare(
1811 "SELECT enqueue_seq, input_id, session_id, source_key, ingress_json,
1812 state, input_json, enqueued_at_ms, claim_id, claim_fencing_token,
1813 claim_owner_id, claim_owner_incarnation_id,
1814 claim_owner_liveness_json, claim_token, claim_expires_at_ms
1815 FROM pending_turn_inputs
1816 WHERE session_id = ?1 AND state = ?2
1817 ORDER BY enqueue_seq ASC
1818 LIMIT ?3",
1819 )
1820 .map_err(sqlite_error)?;
1821 let rows = stmt
1822 .query_map(
1823 params![
1824 session_id,
1825 wanted_state.as_str(),
1826 (max_inputs as i64).saturating_add(32)
1827 ],
1828 pending_turn_input_row_from_sql,
1829 )
1830 .map_err(sqlite_error)?;
1831 rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)?
1832 };
1833 let mut selected = Vec::new();
1834 for row in candidate_rows {
1835 let claim_available = row.claim_token.is_none()
1836 || row.claim_expires_at_ms <= now
1837 || row
1838 .claim_owner
1839 .as_ref()
1840 .is_some_and(|holder| holder.is_definitely_dead_for_claimant(&owner));
1841 if !claim_available {
1842 continue;
1843 }
1844 let input = pending_turn_input_from_row(row.clone())?;
1845 let matches_mode = match &mode {
1846 lash_core::TurnInputClaimMode::ActiveTurn {
1847 turn_id,
1848 checkpoint,
1849 } => {
1850 input
1851 .ingress
1852 .active_turn_id()
1853 .is_some_and(|active| active == turn_id)
1854 && input.ingress.admits_checkpoint(*checkpoint)
1855 }
1856 lash_core::TurnInputClaimMode::NextTurn => input.state.is_next_turn_pending(),
1857 };
1858 if matches_mode {
1859 selected.push((row, input));
1860 if selected.len() >= max_inputs {
1861 break;
1862 }
1863 }
1864 }
1865 let Some((head, _)) = selected.first() else {
1866 return Ok(TxOutcome::Commit(None));
1867 };
1868 let lease = TurnInputClaimLease::derive(head, &session_id, &owner, now, lease_ttl_ms);
1869 let liveness_json = encode_liveness(&owner.liveness)?;
1870 let state_after_claim = match &mode {
1871 lash_core::TurnInputClaimMode::ActiveTurn { .. } => {
1872 lash_core::TurnInputState::Accepted
1873 }
1874 lash_core::TurnInputClaimMode::NextTurn => {
1875 lash_core::TurnInputState::DeferredNextTurn
1876 }
1877 };
1878 let mut inputs = Vec::new();
1879 for (row, mut input) in selected {
1880 let claimed = tx
1881 .execute(
1882 "UPDATE pending_turn_inputs
1883 SET state = ?3,
1884 claim_id = ?4,
1885 claim_owner_id = ?5,
1886 claim_owner_incarnation_id = ?6,
1887 claim_owner_liveness_json = ?7,
1888 claim_token = ?8,
1889 claim_fencing_token = claim_fencing_token + 1,
1890 claim_claimed_at_ms = ?9,
1891 claim_expires_at_ms = ?10
1892 WHERE session_id = ?1
1893 AND input_id = ?2
1894 AND (
1895 claim_token IS NULL
1896 OR claim_expires_at_ms <= ?9
1897 OR (
1898 claim_token = ?11
1899 AND claim_owner_id = ?12
1900 AND claim_owner_incarnation_id = ?13
1901 )
1902 )",
1903 params![
1904 session_id,
1905 row.input_id,
1906 state_after_claim.as_str(),
1907 lease.claim_id,
1908 owner.owner_id.as_str(),
1909 owner.incarnation_id.as_str(),
1910 liveness_json.as_str(),
1911 lease.lease_token,
1912 now as i64,
1913 lease.expires_at_epoch_ms as i64,
1914 row.claim_token,
1915 row.claim_owner
1916 .as_ref()
1917 .map(|owner| owner.owner_id.as_str()),
1918 row.claim_owner
1919 .as_ref()
1920 .map(|owner| owner.incarnation_id.as_str())
1921 ],
1922 )
1923 .map_err(sqlite_error)?;
1924 if claimed == 0 {
1925 return Ok(TxOutcome::Rollback(None));
1926 }
1927 input.state = state_after_claim;
1928 inputs.push(input);
1929 }
1930 Ok(TxOutcome::Commit(Some(lash_core::TurnInputClaim {
1931 session_id: session_id.clone(),
1932 claim_id: lease.claim_id,
1933 owner: owner.clone(),
1934 lease_token: lease.lease_token,
1935 fencing_token: lease.fencing_token,
1936 claimed_at_epoch_ms: lease.claimed_at_epoch_ms,
1937 expires_at_epoch_ms: lease.expires_at_epoch_ms,
1938 mode,
1939 inputs,
1940 })))
1941 })(
1942 );
1943 match outcome {
1944 Ok(TxOutcome::Commit(value)) => Ok(TxOutcome::Commit(Ok(value))),
1945 Ok(TxOutcome::Rollback(value)) => Ok(TxOutcome::Rollback(Ok(value))),
1946 Err(err) => Ok(TxOutcome::Rollback(Err(err))),
1947 }
1948 })
1949 .await
1950 .map_err(sqlite_error)?
1951}
1952
1953struct SessionExecutionLeaseRow {
1954 owner: Option<LeaseOwnerIdentity>,
1955 lease_token: Option<String>,
1956 fencing_token: u64,
1957 claimed_at_ms: u64,
1958 expires_at_ms: u64,
1959}
1960
1961fn load_session_execution_lease_row_conn(
1962 conn: &Connection,
1963 session_id: &str,
1964) -> Result<Option<SessionExecutionLeaseRow>, StoreError> {
1965 let row = conn
1966 .query_row(
1967 "SELECT lease_owner_id, lease_token, lease_fencing_token,
1968 lease_claimed_at_ms, lease_expires_at_ms,
1969 lease_owner_incarnation_id, lease_owner_liveness_json
1970 FROM session_execution_leases
1971 WHERE session_id = ?1",
1972 params![session_id],
1973 |row| {
1974 let owner_id: Option<String> = row.get(0)?;
1975 let incarnation_id: Option<String> = row.get(5)?;
1976 let liveness_json: Option<String> = row.get(6)?;
1977 Ok(SessionExecutionLeaseRow {
1978 owner: lease_owner_from_columns(owner_id, incarnation_id, liveness_json),
1979 lease_token: row.get(1)?,
1980 fencing_token: row.get::<_, i64>(2)? as u64,
1981 claimed_at_ms: row.get::<_, i64>(3)? as u64,
1982 expires_at_ms: row.get::<_, i64>(4)? as u64,
1983 })
1984 },
1985 )
1986 .optional()
1987 .map_err(sqlite_error)?;
1988 Ok(row)
1989}
1990
1991fn lease_owner_from_columns(
1992 owner_id: Option<String>,
1993 incarnation_id: Option<String>,
1994 liveness_json: Option<String>,
1995) -> Option<LeaseOwnerIdentity> {
1996 owner_id.map(|owner_id| LeaseOwnerIdentity {
1997 incarnation_id: incarnation_id.unwrap_or_else(|| owner_id.clone()),
1998 owner_id,
1999 liveness: liveness_json
2000 .as_deref()
2001 .and_then(|json| serde_json::from_str(json).ok())
2002 .unwrap_or(LeaseOwnerLiveness::Opaque),
2003 })
2004}
2005
2006fn encode_liveness(liveness: &LeaseOwnerLiveness) -> Result<String, StoreError> {
2007 serde_json::to_string(liveness)
2008 .map_err(|err| StoreError::Backend(format!("failed to encode lease liveness: {err}")))
2009}
2010
2011fn row_to_session_execution_lease(
2012 session_id: &str,
2013 row: SessionExecutionLeaseRow,
2014) -> Result<SessionExecutionLease, StoreError> {
2015 Ok(SessionExecutionLease {
2016 session_id: session_id.to_string(),
2017 owner: row
2018 .owner
2019 .ok_or_else(|| StoreError::Backend("live session lease missing owner".to_string()))?,
2020 lease_token: row.lease_token.ok_or_else(|| {
2021 StoreError::Backend("live session lease missing lease token".to_string())
2022 })?,
2023 fencing_token: row.fencing_token,
2024 claimed_at_epoch_ms: row.claimed_at_ms,
2025 expires_at_epoch_ms: row.expires_at_ms,
2026 })
2027}
2028
2029fn acquire_session_execution_lease_conn(
2030 conn: &Connection,
2031 session_id: &str,
2032 owner: &LeaseOwnerIdentity,
2033 previous_fencing_token: u64,
2034 now: u64,
2035 lease_ttl_ms: u64,
2036) -> Result<SessionExecutionLease, StoreError> {
2037 let fencing_token = previous_fencing_token.saturating_add(1);
2038 let lease_token = format!(
2039 "{}:{}:{}:{now}:{fencing_token}",
2040 session_id, owner.owner_id, owner.incarnation_id
2041 );
2042 let expires_at = now.saturating_add(lease_ttl_ms);
2043 let liveness_json = encode_liveness(&owner.liveness)?;
2044 conn.execute(
2045 "INSERT INTO session_execution_leases (
2046 session_id, lease_owner_id, lease_owner_incarnation_id, lease_owner_liveness_json,
2047 lease_token, lease_fencing_token, lease_claimed_at_ms, lease_expires_at_ms
2048 )
2049 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)
2050 ON CONFLICT(session_id) DO UPDATE SET
2051 lease_owner_id = excluded.lease_owner_id,
2052 lease_owner_incarnation_id = excluded.lease_owner_incarnation_id,
2053 lease_owner_liveness_json = excluded.lease_owner_liveness_json,
2054 lease_token = excluded.lease_token,
2055 lease_fencing_token = excluded.lease_fencing_token,
2056 lease_claimed_at_ms = excluded.lease_claimed_at_ms,
2057 lease_expires_at_ms = excluded.lease_expires_at_ms",
2058 params![
2059 session_id,
2060 owner.owner_id,
2061 owner.incarnation_id,
2062 liveness_json,
2063 lease_token,
2064 fencing_token as i64,
2065 now as i64,
2066 expires_at as i64
2067 ],
2068 )
2069 .map_err(sqlite_error)?;
2070 Ok(SessionExecutionLease {
2071 session_id: session_id.to_string(),
2072 owner: owner.clone(),
2073 lease_token,
2074 fencing_token,
2075 claimed_at_epoch_ms: now,
2076 expires_at_epoch_ms: expires_at,
2077 })
2078}
2079
2080fn ensure_session_execution_lease_conn(
2081 conn: &Connection,
2082 session_id: &str,
2083 fence: &SessionExecutionLeaseFence,
2084) -> Result<(), StoreError> {
2085 if fence.session_id != session_id {
2086 return Err(StoreError::SessionExecutionLeaseExpired {
2087 session_id: session_id.to_string(),
2088 });
2089 }
2090 let now = current_epoch_ms();
2091 let current = load_session_execution_lease_row_conn(conn, session_id)?;
2092 let Some(current) = current else {
2093 return Err(StoreError::SessionExecutionLeaseExpired {
2094 session_id: session_id.to_string(),
2095 });
2096 };
2097 if current
2098 .owner
2099 .as_ref()
2100 .is_some_and(|owner| owner.same_incarnation(&fence.owner))
2101 && current.lease_token.as_deref() == Some(fence.lease_token.as_str())
2102 && current.fencing_token == fence.fencing_token
2103 && current.expires_at_ms > now
2104 {
2105 Ok(())
2106 } else {
2107 Err(StoreError::SessionExecutionLeaseExpired {
2108 session_id: session_id.to_string(),
2109 })
2110 }
2111}
2112
2113fn release_session_execution_lease_conn(
2114 conn: &Connection,
2115 completion: &SessionExecutionLeaseCompletion,
2116) -> Result<(), StoreError> {
2117 conn.execute(
2118 "UPDATE session_execution_leases
2119 SET lease_owner_id = NULL,
2120 lease_owner_incarnation_id = NULL,
2121 lease_owner_liveness_json = NULL,
2122 lease_token = NULL,
2123 lease_claimed_at_ms = 0,
2124 lease_expires_at_ms = 0
2125 WHERE session_id = ?1
2126 AND lease_owner_id = ?2
2127 AND lease_owner_incarnation_id = ?3
2128 AND lease_token = ?4
2129 AND lease_fencing_token = ?5",
2130 params![
2131 completion.session_id,
2132 completion.owner.owner_id,
2133 completion.owner.incarnation_id,
2134 completion.lease_token,
2135 completion.fencing_token as i64
2136 ],
2137 )
2138 .map_err(sqlite_error)?;
2139 Ok(())
2140}