1use std::future::Future;
32use std::pin::Pin;
33
34use a2a_protocol_types::error::{A2aError, A2aResult};
35use a2a_protocol_types::params::ListTasksParams;
36use a2a_protocol_types::responses::TaskListResponse;
37use a2a_protocol_types::task::{Task, TaskId};
38use sqlx::sqlite::SqlitePool;
39
40use super::event_log_sql::{
41 decode_text_row, encode_text, limit_to_i64, report_no_op_append, seq_from_i64, seq_to_i64,
42};
43use super::task_store::{RecordedEvent, TaskStore};
44use super::tenant::TenantContext;
45use super::tenant_event_log as evlog;
46use super::tenant_idempotency as idem;
47use crate::metrics::MetricsHandle;
48
49#[derive(Debug, Clone)]
71pub struct TenantAwareSqliteTaskStore {
72 pool: SqlitePool,
73 max_page_size: u32,
76 metrics: MetricsHandle,
79}
80
81impl TenantAwareSqliteTaskStore {
82 #[must_use]
90 pub const fn with_max_page_size(mut self, max: u32) -> Self {
91 self.max_page_size = max;
92 self
93 }
94
95 #[must_use]
102 pub fn with_metrics(mut self, metrics: MetricsHandle) -> Self {
103 self.metrics = metrics;
104 self
105 }
106
107 pub async fn new(url: &str) -> Result<Self, sqlx::Error> {
113 let pool = sqlite_pool(url).await?;
114 Self::from_pool(pool).await
115 }
116
117 pub async fn from_pool(pool: SqlitePool) -> Result<Self, sqlx::Error> {
123 sqlx::query(
124 "CREATE TABLE IF NOT EXISTS tenant_tasks (
125 tenant_id TEXT NOT NULL DEFAULT '',
126 id TEXT NOT NULL,
127 context_id TEXT NOT NULL,
128 state TEXT NOT NULL,
129 data TEXT NOT NULL,
130 updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%d %H:%M:%f','now')),
131 created_at TEXT NOT NULL DEFAULT (datetime('now')),
132 PRIMARY KEY (tenant_id, id)
133 )",
134 )
135 .execute(&pool)
136 .await?;
137
138 sqlx::query(idem::SQLITE_CREATE_TABLE)
148 .execute(&pool)
149 .await?;
150
151 sqlx::query(evlog::SQLITE_CREATE_TABLE)
157 .execute(&pool)
158 .await?;
159
160 sqlx::query(
161 "CREATE INDEX IF NOT EXISTS idx_tenant_tasks_ctx ON tenant_tasks(tenant_id, context_id)",
162 )
163 .execute(&pool)
164 .await?;
165
166 sqlx::query(
167 "CREATE INDEX IF NOT EXISTS idx_tenant_tasks_state ON tenant_tasks(tenant_id, state)",
168 )
169 .execute(&pool)
170 .await?;
171
172 sqlx::query(
173 "CREATE INDEX IF NOT EXISTS idx_tenant_tasks_ctx_state ON tenant_tasks(tenant_id, context_id, state)",
174 )
175 .execute(&pool)
176 .await?;
177
178 sqlx::query(
181 "CREATE INDEX IF NOT EXISTS idx_tenant_tasks_updated_at ON tenant_tasks(tenant_id, updated_at DESC, id DESC)",
182 )
183 .execute(&pool)
184 .await?;
185
186 Ok(Self {
187 pool,
188 max_page_size: crate::store::DEFAULT_MAX_PAGE_SIZE,
189 metrics: MetricsHandle::default(),
190 })
191 }
192
193 pub async fn purge_expired(
217 &self,
218 policy: &super::retention::RetentionPolicy,
219 ) -> A2aResult<super::retention::PurgeReport> {
220 super::retention::sqlite::purge(
221 &self.pool,
222 "tenant_tasks",
223 &[evlog::SQLITE_DELETE_ORPHANS],
224 super::tenant_idempotency::SQLITE_EXPIRE,
225 policy,
226 )
227 .await
228 .map_err(|e| to_a2a_error(&e))
229 }
230}
231
232use crate::sqlite_pool::sqlite_pool;
233
234fn to_a2a_error(e: &sqlx::Error) -> A2aError {
235 A2aError::internal(format!("sqlite error: {e}"))
236}
237
238#[allow(clippy::manual_async_fn)]
239impl TaskStore for TenantAwareSqliteTaskStore {
240 fn supports_idempotency(&self) -> bool {
241 true
242 }
243
244 fn claim_idempotency_key<'a>(
245 &'a self,
246 key: &'a str,
247 message_id: &'a a2a_protocol_types::message::MessageId,
248 task_id: &'a TaskId,
249 ) -> Pin<
250 Box<dyn Future<Output = A2aResult<crate::store::task_store::IdempotencyClaim>> + Send + 'a>,
251 > {
252 use crate::store::task_store::IdempotencyClaim;
253 use sqlx::Row as _;
254
255 Box::pin(async move {
256 let tenant = TenantContext::current();
257 let mut tx = self.pool.begin().await.map_err(|e| to_a2a_error(&e))?;
261
262 let inserted = sqlx::query(idem::SQLITE_CLAIM)
263 .bind(&tenant)
264 .bind(key)
265 .bind(message_id.0.as_str())
266 .bind(task_id.0.as_str())
267 .execute(&mut *tx)
268 .await
269 .map_err(|e| to_a2a_error(&e))?
270 .rows_affected();
271
272 if inserted == 1 {
273 tx.commit().await.map_err(|e| to_a2a_error(&e))?;
274 return Ok(IdempotencyClaim::Claimed);
275 }
276
277 let holder = sqlx::query(idem::SQLITE_HOLDER)
278 .bind(&tenant)
279 .bind(key)
280 .fetch_optional(&mut *tx)
281 .await
282 .map_err(|e| to_a2a_error(&e))?;
283 tx.commit().await.map_err(|e| to_a2a_error(&e))?;
284
285 let Some(row) = holder else {
286 return Err(A2aError::internal(
287 "idempotency key was neither claimed nor held; the row vanished \
288 inside the claiming transaction",
289 ));
290 };
291
292 let held_by: String = row.try_get("message_id").map_err(|e| to_a2a_error(&e))?;
293 let held_task: String = row.try_get("task_id").map_err(|e| to_a2a_error(&e))?;
294
295 Ok(idem::outcome(held_by, held_task, message_id))
296 })
297 }
298
299 fn release_idempotency_key<'a>(
300 &'a self,
301 key: &'a str,
302 ) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
303 Box::pin(async move {
304 let tenant = TenantContext::current();
305 sqlx::query(idem::SQLITE_RELEASE)
306 .bind(&tenant)
307 .bind(key)
308 .execute(&self.pool)
309 .await
310 .map_err(|e| to_a2a_error(&e))?;
311 Ok(())
312 })
313 }
314
315 fn save<'a>(
316 &'a self,
317 task: &'a Task,
318 ) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
319 Box::pin(async move {
320 let tenant = TenantContext::current();
321 let id = task.id.0.as_str();
322 let context_id = task.context_id.0.as_str();
323 let state = task.status.state.to_string();
324 let data = serde_json::to_string(task)
325 .map_err(|e| A2aError::internal(format!("failed to serialize task: {e}")))?;
326 let status_ts = super::status_timestamp_sqlite(task.status.timestamp.as_deref());
330
331 let written = sqlx::query(crate::store::terminal::sql_write_allowed!(
336 "INSERT INTO tenant_tasks (tenant_id, id, context_id, state, data, updated_at)
337 VALUES (?1, ?2, ?3, ?4, ?5, COALESCE(?6, strftime('%Y-%m-%d %H:%M:%f','now')))
338 ON CONFLICT(tenant_id, id) DO UPDATE SET
339 context_id = excluded.context_id,
340 state = excluded.state,
341 data = excluded.data,
342 updated_at = excluded.updated_at
343 WHERE ",
344 "tenant_tasks.state",
345 "excluded.state"
346 ))
347 .bind(&tenant)
348 .bind(id)
349 .bind(context_id)
350 .bind(&state)
351 .bind(&data)
352 .bind(&status_ts)
353 .execute(&self.pool)
354 .await
355 .map_err(|e| to_a2a_error(&e))?
356 .rows_affected();
357
358 if written == 0 {
359 let stored: Option<(String,)> = sqlx::query_as(
362 "SELECT state FROM tenant_tasks WHERE tenant_id = ?1 AND id = ?2",
363 )
364 .bind(&tenant)
365 .bind(id)
366 .fetch_optional(&self.pool)
367 .await
368 .map_err(|e| to_a2a_error(&e))?;
369 return Err(crate::store::terminal::refusal(
370 task,
371 stored.map(|(s,)| s).as_deref(),
372 ));
373 }
374
375 Ok(())
376 })
377 }
378
379 fn get<'a>(
380 &'a self,
381 id: &'a TaskId,
382 ) -> Pin<Box<dyn Future<Output = A2aResult<Option<Task>>> + Send + 'a>> {
383 Box::pin(async move {
384 let tenant = TenantContext::current();
385 let row: Option<(String,)> =
386 sqlx::query_as("SELECT data FROM tenant_tasks WHERE tenant_id = ?1 AND id = ?2")
387 .bind(&tenant)
388 .bind(id.0.as_str())
389 .fetch_optional(&self.pool)
390 .await
391 .map_err(|e| to_a2a_error(&e))?;
392
393 match row {
394 Some((data,)) => {
395 let task: Task = serde_json::from_str(&data)
396 .map_err(|e| A2aError::internal(format!("deserialize: {e}")))?;
397 Ok(Some(task))
398 }
399 None => Ok(None),
400 }
401 })
402 }
403
404 #[allow(clippy::too_many_lines)]
405 fn list<'a>(
406 &'a self,
407 params: &'a ListTasksParams,
408 ) -> Pin<Box<dyn Future<Output = A2aResult<TaskListResponse>> + Send + 'a>> {
409 Box::pin(async move {
410 let tenant = TenantContext::current();
411 let mut conditions = vec!["tenant_id = ?1".to_string()];
412 let mut bind_values: Vec<String> = vec![tenant];
413
414 if let Some(ref ctx) = params.context_id {
415 conditions.push(format!("context_id = ?{}", bind_values.len() + 1));
416 bind_values.push(ctx.clone());
417 }
418 if let Some(ref status) = params.status {
419 conditions.push(format!("state = ?{}", bind_values.len() + 1));
420 bind_values.push(status.to_string());
421 }
422 if let Some(ref after) = params.status_timestamp_after {
427 let Some(after_dt) = super::status_timestamp_sqlite(Some(after)) else {
428 return Ok(TaskListResponse::new(Vec::new()));
429 };
430 conditions.push(format!("updated_at > ?{}", bind_values.len() + 1));
431 bind_values.push(after_dt);
432 }
433 if let Some(ref token) = params.page_token {
437 let Some((cursor_ua, cursor_id)) = super::cursor::decode(token) else {
438 return Ok(TaskListResponse::new(Vec::new()));
439 };
440 let p = bind_values.len();
441 conditions.push(format!("(updated_at, id) < (?{}, ?{})", p + 1, p + 2));
442 bind_values.push(cursor_ua.to_string());
443 bind_values.push(cursor_id.to_string());
444 }
445
446 let where_clause = format!("WHERE {}", conditions.join(" AND "));
447
448 let page_size = match params.page_size {
449 Some(0) | None => 50_u32,
450 Some(n) => n.min(self.max_page_size),
451 };
452
453 let limit = super::pagination::fetch_limit(page_size);
454 let sql = format!(
455 "SELECT updated_at, data FROM tenant_tasks {where_clause} \
456 ORDER BY updated_at DESC, id DESC LIMIT {limit}"
457 );
458
459 let mut query = sqlx::query_as::<_, (String, String)>(&sql);
460 for val in &bind_values {
461 query = query.bind(val);
462 }
463
464 let rows: Vec<(String, String)> = query
465 .fetch_all(&self.pool)
466 .await
467 .map_err(|e| to_a2a_error(&e))?;
468
469 let mut rows: Vec<(String, Task)> = rows
470 .into_iter()
471 .map(|(updated_at, data)| {
472 serde_json::from_str::<Task>(&data)
473 .map(|task| (updated_at, task))
474 .map_err(|e| A2aError::internal(format!("deserialize: {e}")))
475 })
476 .collect::<A2aResult<Vec<_>>>()?;
477
478 let next_page_token =
479 if super::pagination::has_next_page(rows.len(), page_size as usize) {
480 rows.truncate(page_size as usize);
481 rows.last()
482 .map(|(ua, task)| super::cursor::encode(ua, task.id.0.as_str()))
483 .unwrap_or_default()
484 } else {
485 String::new()
486 };
487
488 #[allow(clippy::cast_possible_truncation)]
489 let page_len = rows.len() as u32;
490 let tasks: Vec<Task> = rows.into_iter().map(|(_, task)| task).collect();
491 let mut response = TaskListResponse::new(tasks);
492 response.next_page_token = next_page_token;
493 response.page_size = page_len;
494 Ok(response)
495 })
496 }
497
498 fn insert_if_absent<'a>(
499 &'a self,
500 task: &'a Task,
501 ) -> Pin<Box<dyn Future<Output = A2aResult<bool>> + Send + 'a>> {
502 Box::pin(async move {
503 let tenant = TenantContext::current();
504 let id = task.id.0.as_str();
505 let context_id = task.context_id.0.as_str();
506 let state = task.status.state.to_string();
507 let data = serde_json::to_string(task)
508 .map_err(|e| A2aError::internal(format!("serialize: {e}")))?;
509 let status_ts = super::status_timestamp_sqlite(task.status.timestamp.as_deref());
510
511 let result = sqlx::query(
512 "INSERT OR IGNORE INTO tenant_tasks (tenant_id, id, context_id, state, data, updated_at)
513 VALUES (?1, ?2, ?3, ?4, ?5, COALESCE(?6, strftime('%Y-%m-%d %H:%M:%f','now')))",
514 )
515 .bind(&tenant)
516 .bind(id)
517 .bind(context_id)
518 .bind(&state)
519 .bind(&data)
520 .bind(&status_ts)
521 .execute(&self.pool)
522 .await
523 .map_err(|e| to_a2a_error(&e))?;
524
525 Ok(result.rows_affected() > 0)
526 })
527 }
528
529 fn delete<'a>(
530 &'a self,
531 id: &'a TaskId,
532 ) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
533 Box::pin(async move {
534 let tenant = TenantContext::current();
535 sqlx::query(evlog::SQLITE_DELETE_FOR_TASK)
540 .bind(&tenant)
541 .bind(id.0.as_str())
542 .execute(&self.pool)
543 .await
544 .map_err(|e| to_a2a_error(&e))?;
545
546 sqlx::query("DELETE FROM tenant_tasks WHERE tenant_id = ?1 AND id = ?2")
547 .bind(&tenant)
548 .bind(id.0.as_str())
549 .execute(&self.pool)
550 .await
551 .map_err(|e| to_a2a_error(&e))?;
552 Ok(())
553 })
554 }
555
556 fn supports_event_log(&self) -> bool {
557 true
558 }
559
560 fn append_event<'a>(
561 &'a self,
562 task_id: &'a TaskId,
563 seq: u64,
564 event: &'a a2a_protocol_types::events::StreamResponse,
565 ) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
566 Box::pin(async move {
567 let tenant = TenantContext::current();
568 let position = seq_to_i64(seq)?;
569 let payload = encode_text(event)?;
570 let changed_nothing = sqlx::query(evlog::SQLITE_APPEND)
573 .bind(&tenant)
574 .bind(task_id.0.as_str())
575 .bind(position)
576 .bind(&payload)
577 .execute(&self.pool)
578 .await
579 .map_err(|e| to_a2a_error(&e))?
580 .rows_affected()
581 == 0;
582 if changed_nothing {
583 let stored: Option<(String,)> = sqlx::query_as(evlog::SQLITE_SELECT_PAYLOAD)
584 .bind(&tenant)
585 .bind(task_id.0.as_str())
586 .bind(position)
587 .fetch_optional(&self.pool)
588 .await
589 .unwrap_or(None);
590 let same = stored.is_some_and(|(held,)| held == payload);
591 report_no_op_append(&*self.metrics, task_id, seq, same);
592 }
593 Ok(())
594 })
595 }
596
597 fn last_event_seq<'a>(
598 &'a self,
599 task_id: &'a TaskId,
600 ) -> Pin<Box<dyn Future<Output = A2aResult<u64>> + Send + 'a>> {
601 Box::pin(async move {
602 let tenant = TenantContext::current();
603 let (max,): (i64,) = sqlx::query_as(evlog::SQLITE_LAST_SEQ)
604 .bind(&tenant)
605 .bind(task_id.0.as_str())
606 .fetch_one(&self.pool)
607 .await
608 .map_err(|e| to_a2a_error(&e))?;
609 seq_from_i64(max)
610 })
611 }
612
613 fn earliest_event_seq<'a>(
614 &'a self,
615 task_id: &'a TaskId,
616 ) -> Pin<Box<dyn Future<Output = A2aResult<Option<u64>>> + Send + 'a>> {
617 Box::pin(async move {
618 let tenant = TenantContext::current();
619 let (min,): (Option<i64>,) = sqlx::query_as(evlog::SQLITE_EARLIEST_SEQ)
620 .bind(&tenant)
621 .bind(task_id.0.as_str())
622 .fetch_one(&self.pool)
623 .await
624 .map_err(|e| to_a2a_error(&e))?;
625 min.map(seq_from_i64).transpose()
626 })
627 }
628
629 fn read_events<'a>(
630 &'a self,
631 task_id: &'a TaskId,
632 after_seq: u64,
633 limit: usize,
634 ) -> Pin<Box<dyn Future<Output = A2aResult<Vec<RecordedEvent>>> + Send + 'a>> {
635 Box::pin(async move {
636 let tenant = TenantContext::current();
637 let rows: Vec<(i64, String)> = sqlx::query_as(evlog::SQLITE_SELECT_AFTER)
638 .bind(&tenant)
639 .bind(task_id.0.as_str())
640 .bind(seq_to_i64(after_seq)?)
641 .bind(limit_to_i64(limit))
642 .fetch_all(&self.pool)
643 .await
644 .map_err(|e| to_a2a_error(&e))?;
645 rows.into_iter().map(decode_text_row).collect()
646 })
647 }
648
649 fn count<'a>(&'a self) -> Pin<Box<dyn Future<Output = A2aResult<u64>> + Send + 'a>> {
650 Box::pin(async move {
651 let tenant = TenantContext::current();
652 let row: (i64,) =
653 sqlx::query_as("SELECT COUNT(*) FROM tenant_tasks WHERE tenant_id = ?1")
654 .bind(&tenant)
655 .fetch_one(&self.pool)
656 .await
657 .map_err(|e| to_a2a_error(&e))?;
658 #[allow(clippy::cast_sign_loss)]
659 Ok(row.0 as u64)
660 })
661 }
662}
663
664#[cfg(test)]
665mod tests {
666 use super::*;
667 use a2a_protocol_types::task::{ContextId, Task, TaskId, TaskState, TaskStatus};
668
669 async fn make_store() -> TenantAwareSqliteTaskStore {
670 TenantAwareSqliteTaskStore::new("sqlite::memory:")
671 .await
672 .expect("failed to create in-memory tenant store")
673 }
674
675 fn make_task(id: &str, ctx: &str, state: TaskState) -> Task {
676 Task {
677 id: TaskId::new(id),
678 context_id: ContextId::new(ctx),
679 status: TaskStatus::new(state),
680 history: None,
681 artifacts: None,
682 metadata: None,
683 }
684 }
685
686 #[tokio::test]
690 async fn sqlite_tenant_store_advertises_idempotency_support() {
691 let store = make_store().await;
692 assert!(
693 store.supports_idempotency(),
694 "the tenant-aware SQLite store implements claim_idempotency_key; \
695 it must advertise support"
696 );
697 }
698
699 #[tokio::test]
708 async fn purging_one_tenant_leaves_the_same_id_under_another() {
709 use crate::store::retention::RetentionPolicy;
710 use std::time::Duration;
711
712 let store = make_store().await;
713 for tenant in ["acme", "globex"] {
714 TenantContext::scope(tenant, async {
715 store
716 .save(&make_task("shared-id", "ctx", TaskState::Completed))
717 .await
718 .unwrap();
719 })
720 .await;
721 }
722
723 sqlx::query(
725 "UPDATE tenant_tasks \
726 SET updated_at = strftime('%Y-%m-%d %H:%M:%f','now','-7200 seconds') \
727 WHERE tenant_id = 'acme'",
728 )
729 .execute(&store.pool)
730 .await
731 .expect("backdate");
732
733 let report = store
734 .purge_expired(&RetentionPolicy::new(Duration::from_secs(3_600)))
735 .await
736 .expect("purge");
737 assert_eq!(report.tasks_deleted, 1, "only acme's copy was old enough");
738
739 TenantContext::scope("globex", async {
740 assert!(
741 store
742 .get(&TaskId::new("shared-id"))
743 .await
744 .unwrap()
745 .is_some(),
746 "globex must still have its own task of the same id"
747 );
748 })
749 .await;
750 TenantContext::scope("acme", async {
751 assert!(
752 store
753 .get(&TaskId::new("shared-id"))
754 .await
755 .unwrap()
756 .is_none(),
757 "acme's expired copy should be gone"
758 );
759 })
760 .await;
761 }
762
763 #[tokio::test]
764 async fn save_and_get_within_tenant() {
765 let store = make_store().await;
766 TenantContext::scope("acme", async {
767 store
768 .save(&make_task("t1", "ctx1", TaskState::Submitted))
769 .await
770 .unwrap();
771 let task = store.get(&TaskId::new("t1")).await.unwrap();
772 assert!(
773 task.is_some(),
774 "task should be retrievable within its tenant"
775 );
776 assert_eq!(task.unwrap().id, TaskId::new("t1"));
777 })
778 .await;
779 }
780
781 #[tokio::test]
782 async fn tenant_isolation_get() {
783 let store = make_store().await;
784 TenantContext::scope("tenant-a", async {
785 store
786 .save(&make_task("t1", "ctx1", TaskState::Submitted))
787 .await
788 .unwrap();
789 })
790 .await;
791
792 TenantContext::scope("tenant-b", async {
793 let result = store.get(&TaskId::new("t1")).await.unwrap();
794 assert!(result.is_none(), "tenant-b should not see tenant-a's task");
795 })
796 .await;
797 }
798
799 #[tokio::test]
800 async fn tenant_isolation_list() {
801 let store = make_store().await;
802 TenantContext::scope("tenant-a", async {
803 store
804 .save(&make_task("t1", "ctx1", TaskState::Submitted))
805 .await
806 .unwrap();
807 store
808 .save(&make_task("t2", "ctx1", TaskState::Working))
809 .await
810 .unwrap();
811 })
812 .await;
813
814 TenantContext::scope("tenant-b", async {
815 store
816 .save(&make_task("t3", "ctx1", TaskState::Submitted))
817 .await
818 .unwrap();
819 })
820 .await;
821
822 TenantContext::scope("tenant-a", async {
823 let response = store.list(&ListTasksParams::default()).await.unwrap();
824 assert_eq!(
825 response.tasks.len(),
826 2,
827 "tenant-a should see only its 2 tasks"
828 );
829 })
830 .await;
831
832 TenantContext::scope("tenant-b", async {
833 let response = store.list(&ListTasksParams::default()).await.unwrap();
834 assert_eq!(
835 response.tasks.len(),
836 1,
837 "tenant-b should see only its 1 task"
838 );
839 })
840 .await;
841 }
842
843 #[tokio::test]
844 async fn tenant_isolation_count() {
845 let store = make_store().await;
846 TenantContext::scope("tenant-a", async {
847 store
848 .save(&make_task("t1", "ctx1", TaskState::Submitted))
849 .await
850 .unwrap();
851 store
852 .save(&make_task("t2", "ctx1", TaskState::Working))
853 .await
854 .unwrap();
855 })
856 .await;
857
858 TenantContext::scope("tenant-b", async {
859 let count = store.count().await.unwrap();
860 assert_eq!(count, 0, "tenant-b should have zero tasks");
861 })
862 .await;
863
864 TenantContext::scope("tenant-a", async {
865 let count = store.count().await.unwrap();
866 assert_eq!(count, 2, "tenant-a should have 2 tasks");
867 })
868 .await;
869 }
870
871 #[tokio::test]
872 async fn tenant_isolation_delete() {
873 let store = make_store().await;
874 TenantContext::scope("tenant-a", async {
875 store
876 .save(&make_task("t1", "ctx1", TaskState::Submitted))
877 .await
878 .unwrap();
879 })
880 .await;
881
882 TenantContext::scope("tenant-b", async {
884 store.delete(&TaskId::new("t1")).await.unwrap();
885 })
886 .await;
887
888 TenantContext::scope("tenant-a", async {
889 let task = store.get(&TaskId::new("t1")).await.unwrap();
890 assert!(
891 task.is_some(),
892 "tenant-a's task should still exist after tenant-b's delete"
893 );
894 })
895 .await;
896 }
897
898 #[tokio::test]
899 async fn same_task_id_different_tenants() {
900 let store = make_store().await;
901 TenantContext::scope("tenant-a", async {
902 store
903 .save(&make_task("t1", "ctx-a", TaskState::Submitted))
904 .await
905 .unwrap();
906 })
907 .await;
908
909 TenantContext::scope("tenant-b", async {
910 store
911 .save(&make_task("t1", "ctx-b", TaskState::Working))
912 .await
913 .unwrap();
914 })
915 .await;
916
917 TenantContext::scope("tenant-a", async {
918 let task = store.get(&TaskId::new("t1")).await.unwrap().unwrap();
919 assert_eq!(
920 task.context_id,
921 ContextId::new("ctx-a"),
922 "tenant-a should get its own version of t1"
923 );
924 assert_eq!(task.status.state, TaskState::Submitted);
925 })
926 .await;
927
928 TenantContext::scope("tenant-b", async {
929 let task = store.get(&TaskId::new("t1")).await.unwrap().unwrap();
930 assert_eq!(
931 task.context_id,
932 ContextId::new("ctx-b"),
933 "tenant-b should get its own version of t1"
934 );
935 assert_eq!(task.status.state, TaskState::Working);
936 })
937 .await;
938 }
939
940 #[tokio::test]
941 async fn insert_if_absent_respects_tenant_scope() {
942 let store = make_store().await;
943 TenantContext::scope("tenant-a", async {
944 let inserted = store
945 .insert_if_absent(&make_task("t1", "ctx1", TaskState::Submitted))
946 .await
947 .unwrap();
948 assert!(inserted, "first insert should succeed");
949
950 let inserted = store
951 .insert_if_absent(&make_task("t1", "ctx1", TaskState::Working))
952 .await
953 .unwrap();
954 assert!(!inserted, "duplicate insert in same tenant should fail");
955 })
956 .await;
957
958 TenantContext::scope("tenant-b", async {
960 let inserted = store
961 .insert_if_absent(&make_task("t1", "ctx1", TaskState::Working))
962 .await
963 .unwrap();
964 assert!(
965 inserted,
966 "insert of same task id in different tenant should succeed"
967 );
968 })
969 .await;
970 }
971
972 #[tokio::test]
973 async fn list_pagination_within_tenant() {
974 let store = make_store().await;
975 TenantContext::scope("tenant-a", async {
976 for i in 0..5 {
977 store
978 .save(&make_task(
979 &format!("task-{i:03}"),
980 "ctx1",
981 TaskState::Submitted,
982 ))
983 .await
984 .unwrap();
985 }
986
987 let params = ListTasksParams {
988 page_size: Some(2),
989 ..Default::default()
990 };
991 let response = store.list(¶ms).await.unwrap();
992 assert_eq!(response.tasks.len(), 2, "first page should have 2 tasks");
993 assert!(
994 !response.next_page_token.is_empty(),
995 "should have a next page token"
996 );
997
998 let params2 = ListTasksParams {
999 page_size: Some(2),
1000 page_token: Some(response.next_page_token),
1001 ..Default::default()
1002 };
1003 let response2 = store.list(¶ms2).await.unwrap();
1004 assert_eq!(response2.tasks.len(), 2, "second page should have 2 tasks");
1005 })
1006 .await;
1007 }
1008
1009 #[test]
1011 fn to_a2a_error_formats_message() {
1012 let sqlite_err = sqlx::Error::RowNotFound;
1013 let a2a_err = to_a2a_error(&sqlite_err);
1014 let msg = format!("{a2a_err}");
1015 assert!(
1016 msg.contains("sqlite error"),
1017 "error message should contain 'sqlite error': {msg}"
1018 );
1019 }
1020
1021 #[tokio::test]
1022 async fn default_tenant_context_uses_empty_string() {
1023 let store = make_store().await;
1024 store
1026 .save(&make_task("t1", "ctx1", TaskState::Submitted))
1027 .await
1028 .unwrap();
1029 let task = store.get(&TaskId::new("t1")).await.unwrap();
1030 assert!(task.is_some(), "default (empty) tenant should work");
1031 }
1032}
1033
1034#[cfg(test)]
1035mod idempotency_tests {
1036 use super::*;
1037 use crate::store::task_store::IdempotencyClaim;
1038 use a2a_protocol_types::message::MessageId;
1039
1040 const KEY: &str = "8f14e45fceea167a5a36dedd4bea2543";
1041
1042 async fn store() -> TenantAwareSqliteTaskStore {
1043 TenantAwareSqliteTaskStore::new("sqlite::memory:")
1044 .await
1045 .expect("in-memory tenant store")
1046 }
1047
1048 #[tokio::test]
1053 async fn release_frees_the_key_in_its_own_tenant_only() {
1054 let store = store().await;
1055 for tenant in ["tenant-a", "tenant-b"] {
1056 TenantContext::scope(tenant, async {
1057 let claim = store
1058 .claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("t1"))
1059 .await
1060 .unwrap();
1061 assert_eq!(claim, IdempotencyClaim::Claimed);
1062 })
1063 .await;
1064 }
1065
1066 TenantContext::scope("tenant-a", async {
1067 store.release_idempotency_key(KEY).await.unwrap();
1068 let reclaim = store
1069 .claim_idempotency_key(KEY, &MessageId::new("m2"), &TaskId::new("t2"))
1070 .await
1071 .unwrap();
1072 assert_eq!(reclaim, IdempotencyClaim::Claimed, "released, so free");
1073 })
1074 .await;
1075
1076 TenantContext::scope("tenant-b", async {
1077 let held = store
1078 .claim_idempotency_key(KEY, &MessageId::new("m2"), &TaskId::new("t2"))
1079 .await
1080 .unwrap();
1081 assert_ne!(
1082 held,
1083 IdempotencyClaim::Claimed,
1084 "tenant-a's release must not free tenant-b's key"
1085 );
1086 })
1087 .await;
1088 }
1089
1090 #[tokio::test]
1091 async fn one_tenants_key_never_names_another_tenants_task() {
1092 let store = store().await;
1096
1097 let a = TenantContext::scope("tenant-a", async {
1098 store
1099 .claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("task-a"))
1100 .await
1101 .unwrap()
1102 })
1103 .await;
1104 assert_eq!(a, IdempotencyClaim::Claimed);
1105
1106 let b = TenantContext::scope("tenant-b", async {
1107 store
1108 .claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("task-b"))
1109 .await
1110 .unwrap()
1111 })
1112 .await;
1113 assert_eq!(
1114 b,
1115 IdempotencyClaim::Claimed,
1116 "tenant-b must claim its own key, not replay tenant-a's task"
1117 );
1118 }
1119
1120 #[tokio::test]
1121 async fn a_retry_within_one_tenant_replays() {
1122 let store = store().await;
1123 TenantContext::scope("tenant-a", async {
1124 assert_eq!(
1125 store
1126 .claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("task-a"))
1127 .await
1128 .unwrap(),
1129 IdempotencyClaim::Claimed
1130 );
1131 assert_eq!(
1132 store
1133 .claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("other"))
1134 .await
1135 .unwrap(),
1136 IdempotencyClaim::Replay(TaskId::new("task-a"))
1137 );
1138 })
1139 .await;
1140 }
1141
1142 #[tokio::test]
1143 async fn a_reused_key_conflicts_within_its_own_tenant() {
1144 let store = store().await;
1145 TenantContext::scope("tenant-a", async {
1146 store
1147 .claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("task-a"))
1148 .await
1149 .unwrap();
1150 assert_eq!(
1151 store
1152 .claim_idempotency_key(KEY, &MessageId::new("m2"), &TaskId::new("task-b"))
1153 .await
1154 .unwrap(),
1155 IdempotencyClaim::Conflict {
1156 held_by: MessageId::new("m1")
1157 }
1158 );
1159 })
1160 .await;
1161 }
1162
1163 #[tokio::test]
1164 async fn releasing_in_one_tenant_leaves_another_tenants_key_held() {
1165 let store = store().await;
1166 TenantContext::scope("tenant-a", async {
1167 store
1168 .claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("task-a"))
1169 .await
1170 .unwrap();
1171 })
1172 .await;
1173 TenantContext::scope("tenant-b", async {
1174 store
1175 .claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("task-b"))
1176 .await
1177 .unwrap();
1178 store.release_idempotency_key(KEY).await.unwrap();
1179 })
1180 .await;
1181
1182 let still_held = TenantContext::scope("tenant-a", async {
1183 store
1184 .claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("ignored"))
1185 .await
1186 .unwrap()
1187 })
1188 .await;
1189 assert_eq!(
1190 still_held,
1191 IdempotencyClaim::Replay(TaskId::new("task-a")),
1192 "a release must not reach across the tenant boundary"
1193 );
1194 }
1195}
1196
1197#[cfg(test)]
1198mod event_log_tests {
1199 use super::*;
1200 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
1201 use a2a_protocol_types::task::{ContextId, TaskState, TaskStatus};
1202
1203 async fn store() -> TenantAwareSqliteTaskStore {
1204 TenantAwareSqliteTaskStore::new("sqlite::memory:")
1205 .await
1206 .expect("in-memory tenant store")
1207 }
1208
1209 fn task(id: &str) -> Task {
1210 Task {
1211 id: TaskId::new(id),
1212 context_id: ContextId::new("c-1"),
1213 status: TaskStatus::new(TaskState::Working),
1214 history: None,
1215 artifacts: None,
1216 metadata: None,
1217 }
1218 }
1219
1220 fn event(id: &str, state: TaskState) -> StreamResponse {
1221 StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
1222 task_id: TaskId::new(id),
1223 context_id: ContextId::new("c-1"),
1224 status: TaskStatus::new(state),
1225 metadata: None,
1226 })
1227 }
1228
1229 fn states(events: &[RecordedEvent]) -> Vec<TaskState> {
1230 events
1231 .iter()
1232 .map(|r| match &r.event {
1233 StreamResponse::StatusUpdate(u) => u.status.state,
1234 other => panic!("only status events are written here, got {other:?}"),
1235 })
1236 .collect()
1237 }
1238
1239 #[tokio::test]
1244 async fn the_log_reports_its_earliest_position_per_tenant() {
1245 let store = store().await;
1246 let id = TaskId::new("shared-id");
1247
1248 TenantContext::scope("tenant-a", async {
1249 store.save(&task("shared-id")).await.expect("save a");
1250 assert_eq!(store.earliest_event_seq(&id).await.expect("empty"), None);
1251 for seq in [5, 6, 7] {
1252 store
1253 .append_event(&id, seq, &event("shared-id", TaskState::Working))
1254 .await
1255 .expect("append a");
1256 }
1257 assert_eq!(store.earliest_event_seq(&id).await.expect("a"), Some(5));
1258 })
1259 .await;
1260
1261 TenantContext::scope("tenant-b", async {
1262 store.save(&task("shared-id")).await.expect("save b");
1263 assert_eq!(
1264 store.earliest_event_seq(&id).await.expect("b empty"),
1265 None,
1266 "tenant-a's log is not tenant-b's"
1267 );
1268 store
1269 .append_event(&id, 3, &event("shared-id", TaskState::Working))
1270 .await
1271 .expect("append b");
1272 assert_eq!(store.earliest_event_seq(&id).await.expect("b"), Some(3));
1273 })
1274 .await;
1275
1276 TenantContext::scope("tenant-a", async {
1278 assert_eq!(store.earliest_event_seq(&id).await.expect("a"), Some(5));
1279 })
1280 .await;
1281 }
1282
1283 #[tokio::test]
1284 async fn the_tenant_store_reports_that_it_keeps_a_log() {
1285 assert!(store().await.supports_event_log());
1286 }
1287
1288 #[tokio::test]
1294 async fn one_tenants_task_id_never_reaches_another_tenants_log() {
1295 let store = store().await;
1296
1297 TenantContext::scope("tenant-a", async {
1298 store.save(&task("shared-id")).await.expect("save a");
1299 store
1300 .append_event(
1301 &TaskId::new("shared-id"),
1302 1,
1303 &event("shared-id", TaskState::Working),
1304 )
1305 .await
1306 .expect("append a");
1307 })
1308 .await;
1309
1310 TenantContext::scope("tenant-b", async {
1311 store.save(&task("shared-id")).await.expect("save b");
1312 assert_eq!(
1313 store
1314 .last_event_seq(&TaskId::new("shared-id"))
1315 .await
1316 .expect("last b"),
1317 0,
1318 "tenant-b's log for this id is its own, and it is empty"
1319 );
1320 assert!(
1321 store
1322 .read_events(&TaskId::new("shared-id"), 0, 10)
1323 .await
1324 .expect("read b")
1325 .is_empty(),
1326 "tenant-b must not be handed tenant-a's events"
1327 );
1328
1329 store
1330 .append_event(
1331 &TaskId::new("shared-id"),
1332 1,
1333 &event("shared-id", TaskState::Completed),
1334 )
1335 .await
1336 .expect("append b");
1337 })
1338 .await;
1339
1340 TenantContext::scope("tenant-a", async {
1343 assert_eq!(
1344 states(
1345 &store
1346 .read_events(&TaskId::new("shared-id"), 0, 10)
1347 .await
1348 .expect("read a")
1349 ),
1350 vec![TaskState::Working],
1351 );
1352 })
1353 .await;
1354 }
1355
1356 #[tokio::test]
1357 async fn events_round_trip_in_order_with_their_positions() {
1358 let store = store().await;
1359 TenantContext::scope("tenant-a", async {
1360 store.save(&task("t-1")).await.expect("save");
1361 for (seq, state) in [
1362 (1, TaskState::Submitted),
1363 (2, TaskState::Working),
1364 (3, TaskState::Completed),
1365 ] {
1366 store
1367 .append_event(&TaskId::new("t-1"), seq, &event("t-1", state))
1368 .await
1369 .expect("append");
1370 }
1371
1372 let all = store
1373 .read_events(&TaskId::new("t-1"), 0, 100)
1374 .await
1375 .expect("read");
1376 assert_eq!(all.iter().map(|r| r.seq).collect::<Vec<_>>(), vec![1, 2, 3]);
1377 assert_eq!(
1378 states(&all),
1379 vec![
1380 TaskState::Submitted,
1381 TaskState::Working,
1382 TaskState::Completed
1383 ],
1384 );
1385 assert_eq!(
1386 store
1387 .last_event_seq(&TaskId::new("t-1"))
1388 .await
1389 .expect("last"),
1390 3
1391 );
1392 })
1393 .await;
1394 }
1395
1396 #[tokio::test]
1397 async fn appending_the_same_position_twice_leaves_one_row() {
1398 let store = store().await;
1399 TenantContext::scope("tenant-a", async {
1400 store.save(&task("t-1")).await.expect("save");
1401 store
1402 .append_event(&TaskId::new("t-1"), 1, &event("t-1", TaskState::Working))
1403 .await
1404 .expect("append");
1405 store
1406 .append_event(&TaskId::new("t-1"), 1, &event("t-1", TaskState::Completed))
1407 .await
1408 .expect("replay must not error");
1409
1410 let all = store
1411 .read_events(&TaskId::new("t-1"), 0, 10)
1412 .await
1413 .expect("read");
1414 assert_eq!(all.len(), 1, "one position, one row");
1415 assert_eq!(
1416 states(&all),
1417 vec![TaskState::Working],
1418 "the first write wins; a replay must not rewrite history"
1419 );
1420 })
1421 .await;
1422 }
1423
1424 #[tokio::test]
1425 async fn reading_after_an_offset_is_exclusive_and_honours_the_limit() {
1426 let store = store().await;
1427 TenantContext::scope("tenant-a", async {
1428 store.save(&task("t-1")).await.expect("save");
1429 for seq in 1..=5 {
1430 store
1431 .append_event(&TaskId::new("t-1"), seq, &event("t-1", TaskState::Working))
1432 .await
1433 .expect("append");
1434 }
1435
1436 let after_two = store
1437 .read_events(&TaskId::new("t-1"), 2, 100)
1438 .await
1439 .expect("read");
1440 assert_eq!(after_two.first().map(|r| r.seq), Some(3), "exclusive");
1441 assert_eq!(after_two.len(), 3);
1442 assert_eq!(
1443 store
1444 .read_events(&TaskId::new("t-1"), 0, 2)
1445 .await
1446 .expect("read")
1447 .len(),
1448 2
1449 );
1450 })
1451 .await;
1452 }
1453
1454 #[tokio::test]
1455 async fn deleting_a_task_removes_its_log_and_leaves_other_tenants_alone() {
1456 let store = store().await;
1457 for tenant in ["tenant-a", "tenant-b"] {
1458 TenantContext::scope(tenant, async {
1459 store.save(&task("t-1")).await.expect("save");
1460 store
1461 .append_event(&TaskId::new("t-1"), 1, &event("t-1", TaskState::Working))
1462 .await
1463 .expect("append");
1464 })
1465 .await;
1466 }
1467
1468 TenantContext::scope("tenant-a", async {
1469 store.delete(&TaskId::new("t-1")).await.expect("delete");
1470 assert_eq!(
1471 store
1472 .last_event_seq(&TaskId::new("t-1"))
1473 .await
1474 .expect("last"),
1475 0
1476 );
1477 })
1478 .await;
1479
1480 TenantContext::scope("tenant-b", async {
1481 assert_eq!(
1482 store
1483 .last_event_seq(&TaskId::new("t-1"))
1484 .await
1485 .expect("last"),
1486 1,
1487 "a delete in one tenant must not reach across the boundary"
1488 );
1489 })
1490 .await;
1491 }
1492
1493 #[tokio::test]
1498 async fn a_retention_sweep_reclaims_orphaned_events() {
1499 use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
1504 use std::str::FromStr as _;
1505
1506 let opts = SqliteConnectOptions::from_str("sqlite::memory:")
1507 .expect("options")
1508 .pragma("foreign_keys", "OFF")
1509 .create_if_missing(true);
1510 let pool = SqlitePoolOptions::new()
1511 .max_connections(1)
1512 .connect_with(opts)
1513 .await
1514 .expect("pool");
1515 let store = TenantAwareSqliteTaskStore::from_pool(pool.clone())
1516 .await
1517 .expect("store from a pool without the pragma");
1518
1519 TenantContext::scope("tenant-a", async {
1520 let mut done = task("t-1");
1521 done.status = TaskStatus::new(TaskState::Completed);
1522 store.save(&done).await.expect("save");
1523 store
1524 .append_event(&TaskId::new("t-1"), 1, &event("t-1", TaskState::Completed))
1525 .await
1526 .expect("append");
1527 })
1528 .await;
1529
1530 sqlx::query("UPDATE tenant_tasks SET updated_at = '2000-01-01 00:00:00.000'")
1532 .execute(&pool)
1533 .await
1534 .expect("backdate");
1535
1536 let report = store
1537 .purge_expired(&crate::store::RetentionPolicy::new(
1538 std::time::Duration::from_secs(60),
1539 ))
1540 .await
1541 .expect("purge");
1542 assert_eq!(report.tasks_deleted, 1);
1543 assert_eq!(
1544 report.orphan_rows_deleted, 1,
1545 "without the pragma the cascade does not fire, so the sweep must \
1546 reclaim the event row itself"
1547 );
1548
1549 TenantContext::scope("tenant-a", async {
1550 assert_eq!(
1551 store
1552 .last_event_seq(&TaskId::new("t-1"))
1553 .await
1554 .expect("last"),
1555 0
1556 );
1557 })
1558 .await;
1559 }
1560}