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