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