use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use a2a_protocol_types::error::A2aResult;
use a2a_protocol_types::params::ListTasksParams;
use a2a_protocol_types::responses::TaskListResponse;
use a2a_protocol_types::task::{Task, TaskId};
use tokio::sync::RwLock;
use super::super::task_store::{InMemoryTaskStore, RecordedEvent, TaskStore, TaskStoreConfig};
use super::context::TenantContext;
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct TenantStoreConfig {
pub per_tenant: TaskStoreConfig,
pub max_tenants: usize,
}
impl Default for TenantStoreConfig {
fn default() -> Self {
Self {
per_tenant: TaskStoreConfig::default(),
max_tenants: 1000,
}
}
}
impl TenantStoreConfig {
#[must_use]
pub const fn with_per_tenant(mut self, config: TaskStoreConfig) -> Self {
self.per_tenant = config;
self
}
#[must_use]
pub const fn with_max_tenants(mut self, max: usize) -> Self {
self.max_tenants = max;
self
}
}
#[derive(Debug)]
pub struct TenantAwareInMemoryTaskStore {
stores: RwLock<HashMap<String, Arc<InMemoryTaskStore>>>,
config: TenantStoreConfig,
overrides: HashMap<String, TaskStoreConfig>,
metrics: crate::metrics::MetricsHandle,
}
impl Default for TenantAwareInMemoryTaskStore {
fn default() -> Self {
Self::new()
}
}
impl TenantAwareInMemoryTaskStore {
#[must_use]
pub fn new() -> Self {
Self {
stores: RwLock::new(HashMap::new()),
config: TenantStoreConfig::default(),
overrides: HashMap::new(),
metrics: crate::metrics::MetricsHandle::default(),
}
}
#[must_use]
pub fn with_config(config: TenantStoreConfig) -> Self {
Self {
stores: RwLock::new(HashMap::new()),
config,
overrides: HashMap::new(),
metrics: crate::metrics::MetricsHandle::default(),
}
}
#[must_use]
pub fn with_metrics(mut self, metrics: crate::metrics::MetricsHandle) -> Self {
self.metrics = metrics;
self
}
#[must_use]
pub fn with_tenant_override(
mut self,
tenant: impl Into<String>,
config: TaskStoreConfig,
) -> Self {
self.overrides.insert(tenant.into(), config);
self
}
async fn get_store(&self) -> A2aResult<Arc<InMemoryTaskStore>> {
let tenant = TenantContext::current();
{
let stores = self.stores.read().await;
if let Some(store) = stores.get(&tenant) {
return Ok(Arc::clone(store));
}
}
let mut stores = self.stores.write().await;
if let Some(store) = stores.get(&tenant) {
return Ok(Arc::clone(store));
}
if stores.len() >= self.config.max_tenants {
return Err(a2a_protocol_types::error::A2aError::internal(format!(
"tenant limit exceeded: max {} tenants",
self.config.max_tenants
)));
}
let store = Arc::new(
InMemoryTaskStore::with_config(
self.overrides
.get(&tenant)
.unwrap_or(&self.config.per_tenant)
.clone(),
)
.with_metrics(self.metrics.clone()),
);
stores.insert(tenant, Arc::clone(&store));
drop(stores);
Ok(store)
}
async fn get_existing_store(&self) -> Option<Arc<InMemoryTaskStore>> {
let tenant = TenantContext::current();
let stores = self.stores.read().await;
stores.get(&tenant).map(Arc::clone)
}
pub async fn tenant_count(&self) -> usize {
self.stores.read().await.len()
}
pub async fn run_eviction_all(&self) {
let stores = self.stores.read().await;
for store in stores.values() {
store.run_eviction().await;
}
}
pub async fn prune_empty_tenants(&self) {
let mut stores = self.stores.write().await;
let mut empty_tenants = Vec::new();
for (tenant, store) in stores.iter() {
if store.is_prunable().await {
empty_tenants.push(tenant.clone());
}
}
for tenant in empty_tenants {
stores.remove(&tenant);
}
}
}
#[allow(clippy::manual_async_fn)]
impl TaskStore for TenantAwareInMemoryTaskStore {
fn supports_idempotency(&self) -> bool {
true
}
fn claim_idempotency_key<'a>(
&'a self,
key: &'a str,
message_id: &'a a2a_protocol_types::message::MessageId,
task_id: &'a TaskId,
) -> Pin<
Box<dyn Future<Output = A2aResult<super::super::task_store::IdempotencyClaim>> + Send + 'a>,
> {
Box::pin(async move {
let store = self.get_store().await?;
store.claim_idempotency_key(key, message_id, task_id).await
})
}
fn release_idempotency_key<'a>(
&'a self,
key: &'a str,
) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
Box::pin(async move {
match self.get_existing_store().await {
Some(store) => store.release_idempotency_key(key).await,
None => Ok(()),
}
})
}
fn save<'a>(
&'a self,
task: &'a Task,
) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
Box::pin(async move {
let store = self.get_store().await?;
store.save(task).await
})
}
fn get<'a>(
&'a self,
id: &'a TaskId,
) -> Pin<Box<dyn Future<Output = A2aResult<Option<Task>>> + Send + 'a>> {
Box::pin(async move {
match self.get_existing_store().await {
Some(store) => store.get(id).await,
None => Ok(None),
}
})
}
fn list<'a>(
&'a self,
params: &'a ListTasksParams,
) -> Pin<Box<dyn Future<Output = A2aResult<TaskListResponse>> + Send + 'a>> {
Box::pin(async move {
match self.get_existing_store().await {
Some(store) => store.list(params).await,
None => Ok(TaskListResponse::new(Vec::new())),
}
})
}
fn insert_if_absent<'a>(
&'a self,
task: &'a Task,
) -> Pin<Box<dyn Future<Output = A2aResult<bool>> + Send + 'a>> {
Box::pin(async move {
let store = self.get_store().await?;
store.insert_if_absent(task).await
})
}
fn delete<'a>(
&'a self,
id: &'a TaskId,
) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
Box::pin(async move {
match self.get_existing_store().await {
Some(store) => store.delete(id).await,
None => Ok(()),
}
})
}
fn supports_event_log(&self) -> bool {
true
}
fn append_event<'a>(
&'a self,
task_id: &'a TaskId,
seq: u64,
event: &'a a2a_protocol_types::events::StreamResponse,
) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
Box::pin(async move {
let store = self.get_store().await?;
store.append_event(task_id, seq, event).await
})
}
fn last_event_seq<'a>(
&'a self,
task_id: &'a TaskId,
) -> Pin<Box<dyn Future<Output = A2aResult<u64>> + Send + 'a>> {
Box::pin(async move {
match self.get_existing_store().await {
Some(store) => store.last_event_seq(task_id).await,
None => Ok(0),
}
})
}
fn earliest_event_seq<'a>(
&'a self,
task_id: &'a TaskId,
) -> Pin<Box<dyn Future<Output = A2aResult<Option<u64>>> + Send + 'a>> {
Box::pin(async move {
match self.get_existing_store().await {
Some(store) => store.earliest_event_seq(task_id).await,
None => Ok(None),
}
})
}
fn read_events<'a>(
&'a self,
task_id: &'a TaskId,
after_seq: u64,
limit: usize,
) -> Pin<Box<dyn Future<Output = A2aResult<Vec<RecordedEvent>>> + Send + 'a>> {
Box::pin(async move {
match self.get_existing_store().await {
Some(store) => store.read_events(task_id, after_seq, limit).await,
None => Ok(Vec::new()),
}
})
}
fn count<'a>(&'a self) -> Pin<Box<dyn Future<Output = A2aResult<u64>> + Send + 'a>> {
Box::pin(async move {
match self.get_existing_store().await {
Some(store) => store.count().await,
None => Ok(0),
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use a2a_protocol_types::task::{ContextId, TaskState, TaskStatus};
fn make_task(id: &str, state: TaskState) -> Task {
Task {
id: TaskId::new(id),
context_id: ContextId::new("ctx-default"),
status: TaskStatus::new(state),
history: None,
artifacts: None,
metadata: None,
}
}
#[tokio::test]
async fn with_metrics_reaches_the_per_tenant_stores() {
use a2a_protocol_types::events::StreamResponse;
use a2a_protocol_types::events::TaskStatusUpdateEvent;
#[derive(Default)]
struct Recorder(std::sync::Mutex<Vec<(String, String)>>);
impl crate::metrics::Metrics for Recorder {
fn on_persistence_error(&self, operation: &str, error_kind: &str) {
self.0
.lock()
.expect("not poisoned")
.push((operation.to_owned(), error_kind.to_owned()));
}
}
let recorder = Arc::new(Recorder::default());
let store = TenantAwareInMemoryTaskStore::new()
.with_metrics(crate::metrics::MetricsHandle::from_arc(recorder.clone()));
store
.append_event(
&TaskId::new("gone"),
1,
&StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
task_id: TaskId::new("gone"),
context_id: ContextId::new("ctx-default"),
status: TaskStatus::new(TaskState::Working),
metadata: None,
}),
)
.await
.expect("a missing task must not fail the agent");
assert_eq!(
*recorder.0.lock().expect("not poisoned"),
vec![(
crate::metrics::persistence_operation::EVENT_APPEND.to_owned(),
crate::metrics::event_append_error::TASK_ABSENT.to_owned(),
)],
"the partition must report to the handle the store was given"
);
}
#[tokio::test]
async fn per_tenant_page_size_cap_reaches_the_delegate() {
let store = TenantAwareInMemoryTaskStore::with_config(TenantStoreConfig {
per_tenant: TaskStoreConfig {
max_page_size: 2,
..TaskStoreConfig::default()
},
max_tenants: 10,
});
TenantContext::scope("capped", async {
for i in 0..5 {
store
.save(&make_task(&format!("t{i}"), TaskState::Submitted))
.await
.expect("save");
}
let listed = store
.list(&ListTasksParams {
page_size: Some(100),
..Default::default()
})
.await
.expect("list");
assert_eq!(
listed.tasks.len(),
2,
"the caller asked for 100; per_tenant.max_page_size is 2. \
A delegate built with the default config would return 5"
);
})
.await;
}
#[tokio::test]
async fn an_override_gives_that_tenant_its_own_store_config() {
async fn saved_then_listed(
store: &TenantAwareInMemoryTaskStore,
tenant: &'static str,
) -> usize {
TenantContext::scope(tenant, async {
for i in 0..5 {
store
.save(&make_task(&format!("{tenant}-{i}"), TaskState::Submitted))
.await
.expect("save");
}
store
.list(&ListTasksParams {
page_size: Some(100),
..Default::default()
})
.await
.expect("list")
.tasks
.len()
})
.await
}
let store = TenantAwareInMemoryTaskStore::with_config(TenantStoreConfig {
per_tenant: TaskStoreConfig {
max_page_size: 50,
..TaskStoreConfig::default()
},
max_tenants: 10,
})
.with_tenant_override(
"small",
TaskStoreConfig {
max_page_size: 1,
..TaskStoreConfig::default()
},
);
assert_eq!(
saved_then_listed(&store, "small").await,
1,
"the override caps this tenant's page size at 1"
);
assert_eq!(
saved_then_listed(&store, "ordinary").await,
5,
"a tenant with no override keeps per_tenant's cap of 50"
);
}
#[tokio::test]
async fn tenant_context_default_is_empty_string() {
let tenant = TenantContext::current();
assert_eq!(tenant, "", "default tenant should be empty string");
}
#[tokio::test]
async fn tenant_context_scope_sets_and_restores() {
let before = TenantContext::current();
assert_eq!(before, "");
let inside = TenantContext::scope("acme", async { TenantContext::current() }).await;
assert_eq!(inside, "acme", "scope should set the tenant");
let after = TenantContext::current();
assert_eq!(after, "", "tenant should revert after scope exits");
}
#[tokio::test]
async fn tenant_context_nested_scopes() {
TenantContext::scope("outer", async {
assert_eq!(TenantContext::current(), "outer");
TenantContext::scope("inner", async {
assert_eq!(TenantContext::current(), "inner");
})
.await;
assert_eq!(
TenantContext::current(),
"outer",
"should restore outer tenant after inner scope"
);
})
.await;
}
#[tokio::test]
async fn tenant_isolation_save_and_get() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("tenant-a", async {
store
.save(&make_task("t1", TaskState::Submitted))
.await
.unwrap();
})
.await;
let found = TenantContext::scope("tenant-a", async {
store.get(&TaskId::new("t1")).await.unwrap()
})
.await;
assert!(found.is_some(), "tenant-a should see its own task");
let not_found = TenantContext::scope("tenant-b", async {
store.get(&TaskId::new("t1")).await.unwrap()
})
.await;
assert!(
not_found.is_none(),
"tenant-b should not see tenant-a's task"
);
}
#[tokio::test]
async fn tenant_isolation_list() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("alpha", async {
store
.save(&make_task("a1", TaskState::Submitted))
.await
.unwrap();
store
.save(&make_task("a2", TaskState::Working))
.await
.unwrap();
})
.await;
TenantContext::scope("beta", async {
store
.save(&make_task("b1", TaskState::Submitted))
.await
.unwrap();
})
.await;
let alpha_list = TenantContext::scope("alpha", async {
let params = ListTasksParams::default();
store.list(¶ms).await.unwrap()
})
.await;
assert_eq!(
alpha_list.tasks.len(),
2,
"alpha should see only its 2 tasks"
);
let beta_list = TenantContext::scope("beta", async {
let params = ListTasksParams::default();
store.list(¶ms).await.unwrap()
})
.await;
assert_eq!(beta_list.tasks.len(), 1, "beta should see only its 1 task");
}
#[tokio::test]
async fn tenant_isolation_delete() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("tenant-a", async {
store
.save(&make_task("t1", TaskState::Submitted))
.await
.unwrap();
})
.await;
TenantContext::scope("tenant-b", async {
store.delete(&TaskId::new("t1")).await.unwrap();
})
.await;
let still_exists = TenantContext::scope("tenant-a", async {
store.get(&TaskId::new("t1")).await.unwrap()
})
.await;
assert!(
still_exists.is_some(),
"tenant-a's task should survive tenant-b's delete"
);
}
#[tokio::test]
async fn tenant_isolation_insert_if_absent() {
let store = TenantAwareInMemoryTaskStore::new();
let inserted_a = TenantContext::scope("tenant-a", async {
store
.insert_if_absent(&make_task("shared-id", TaskState::Submitted))
.await
.unwrap()
})
.await;
assert!(inserted_a, "tenant-a insert should succeed");
let inserted_b = TenantContext::scope("tenant-b", async {
store
.insert_if_absent(&make_task("shared-id", TaskState::Working))
.await
.unwrap()
})
.await;
assert!(
inserted_b,
"tenant-b insert of same ID should also succeed (different partition)"
);
}
#[tokio::test]
async fn tenant_isolation_count() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("x", async {
store
.save(&make_task("t1", TaskState::Submitted))
.await
.unwrap();
store
.save(&make_task("t2", TaskState::Submitted))
.await
.unwrap();
})
.await;
TenantContext::scope("y", async {
store
.save(&make_task("t3", TaskState::Submitted))
.await
.unwrap();
})
.await;
let count_x = TenantContext::scope("x", async { store.count().await.unwrap() }).await;
assert_eq!(count_x, 2, "tenant x should have 2 tasks");
let count_y = TenantContext::scope("y", async { store.count().await.unwrap() }).await;
assert_eq!(count_y, 1, "tenant y should have 1 task");
}
#[tokio::test]
async fn a_full_tenant_table_refuses_new_tenants_and_pruning_does_not_help() {
let store = TenantAwareInMemoryTaskStore::with_config(TenantStoreConfig {
per_tenant: TaskStoreConfig::default(),
max_tenants: 2,
});
for junk in ["junk-1", "junk-2"] {
TenantContext::scope(junk, async {
store
.save(&make_task("t", TaskState::Working))
.await
.expect("a fresh tenant partition is created on demand");
})
.await;
}
assert_eq!(store.tenant_count().await, 2, "the table is full");
let refused = TenantContext::scope("legitimate", async {
store.save(&make_task("t", TaskState::Working)).await
})
.await;
assert!(
refused.is_err(),
"a new tenant must be refused once the cap is reached"
);
store.prune_empty_tenants().await;
assert_eq!(
store.tenant_count().await,
2,
"pruning reclaims nothing while the junk partitions hold live tasks"
);
let still_refused = TenantContext::scope("legitimate", async {
store.save(&make_task("t", TaskState::Working)).await
})
.await;
assert!(
still_refused.is_err(),
"so the lockout outlives the pruning that is supposed to end it"
);
for junk in ["junk-1", "junk-2"] {
TenantContext::scope(junk, async {
store.delete(&TaskId::new("t")).await.expect("delete");
})
.await;
}
store.prune_empty_tenants().await;
assert_eq!(store.tenant_count().await, 0, "now the slots come back");
let admitted = TenantContext::scope("legitimate", async {
store.save(&make_task("t", TaskState::Working)).await
})
.await;
assert!(admitted.is_ok(), "and the legitimate tenant gets in");
}
#[tokio::test]
async fn tenant_count_reflects_active_tenants() {
let store = TenantAwareInMemoryTaskStore::new();
assert_eq!(store.tenant_count().await, 0);
TenantContext::scope("a", async {
store
.save(&make_task("t1", TaskState::Submitted))
.await
.unwrap();
})
.await;
assert_eq!(store.tenant_count().await, 1);
TenantContext::scope("b", async {
store
.save(&make_task("t2", TaskState::Submitted))
.await
.unwrap();
})
.await;
assert_eq!(store.tenant_count().await, 2);
}
#[tokio::test]
async fn max_tenants_limit_enforced() {
let config = TenantStoreConfig {
per_tenant: TaskStoreConfig::default(),
max_tenants: 2,
};
let store = TenantAwareInMemoryTaskStore::with_config(config);
TenantContext::scope("t1", async {
store
.save(&make_task("task-a", TaskState::Submitted))
.await
.unwrap();
})
.await;
TenantContext::scope("t2", async {
store
.save(&make_task("task-b", TaskState::Submitted))
.await
.unwrap();
})
.await;
let result = TenantContext::scope("t3", async {
store.save(&make_task("task-c", TaskState::Submitted)).await
})
.await;
assert!(
result.is_err(),
"exceeding max_tenants should return an error"
);
}
#[tokio::test]
async fn existing_tenant_does_not_count_against_limit() {
let config = TenantStoreConfig {
per_tenant: TaskStoreConfig::default(),
max_tenants: 1,
};
let store = TenantAwareInMemoryTaskStore::with_config(config);
TenantContext::scope("only", async {
store
.save(&make_task("t1", TaskState::Submitted))
.await
.unwrap();
store
.save(&make_task("t2", TaskState::Working))
.await
.unwrap();
})
.await;
let count = TenantContext::scope("only", async { store.count().await.unwrap() }).await;
assert_eq!(count, 2, "existing tenant can add more tasks");
}
#[tokio::test]
async fn no_tenant_context_uses_default_partition() {
let store = TenantAwareInMemoryTaskStore::new();
store
.save(&make_task("default-task", TaskState::Submitted))
.await
.unwrap();
let fetched = store.get(&TaskId::new("default-task")).await.unwrap();
assert!(
fetched.is_some(),
"task saved without tenant context should be retrievable without context"
);
let not_found = TenantContext::scope("other", async {
store.get(&TaskId::new("default-task")).await.unwrap()
})
.await;
assert!(
not_found.is_none(),
"default partition task should not leak to named tenants"
);
}
#[tokio::test]
async fn prune_empty_tenants_removes_empty_partitions() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("keep", async {
store
.save(&make_task("t1", TaskState::Submitted))
.await
.unwrap();
})
.await;
TenantContext::scope("remove", async {
store
.save(&make_task("t2", TaskState::Submitted))
.await
.unwrap();
})
.await;
assert_eq!(store.tenant_count().await, 2);
TenantContext::scope("remove", async {
store.delete(&TaskId::new("t2")).await.unwrap();
})
.await;
store.prune_empty_tenants().await;
assert_eq!(
store.tenant_count().await,
1,
"empty tenant partition should be pruned"
);
}
#[tokio::test]
async fn a_partition_still_holding_idempotency_keys_is_not_pruned() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("keyed", async {
store
.save(&make_task("t1", TaskState::Submitted))
.await
.unwrap();
store
.claim_idempotency_key(
"8f14e45fceea167a5a36dedd4bea2543",
&a2a_protocol_types::message::MessageId::new("msg-1"),
&TaskId::new("t1"),
)
.await
.unwrap();
store.delete(&TaskId::new("t1")).await.unwrap();
assert_eq!(
store.count().await.unwrap(),
0,
"the partition holds no task, which is what used to decide this"
);
})
.await;
assert_eq!(store.tenant_count().await, 1);
store.prune_empty_tenants().await;
assert_eq!(
store.tenant_count().await,
1,
"a partition holding a live idempotency key must survive the prune"
);
let claim = TenantContext::scope("keyed", async {
store
.claim_idempotency_key(
"8f14e45fceea167a5a36dedd4bea2543",
&a2a_protocol_types::message::MessageId::new("msg-1"),
&TaskId::new("t2"),
)
.await
.unwrap()
})
.await;
assert!(
matches!(
claim,
crate::store::task_store::IdempotencyClaim::Replay(ref id) if id.0 == "t1"
),
"the surviving key must still name the first task, got {claim:?}"
);
}
#[tokio::test]
async fn a_partition_holding_neither_tasks_nor_keys_is_still_pruned() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("transient", async {
store
.save(&make_task("t1", TaskState::Submitted))
.await
.unwrap();
store.delete(&TaskId::new("t1")).await.unwrap();
})
.await;
assert_eq!(store.tenant_count().await, 1);
store.prune_empty_tenants().await;
assert_eq!(
store.tenant_count().await,
0,
"a genuinely empty partition must still give its slot back"
);
}
#[test]
fn default_creates_new_tenant_store() {
let store = TenantAwareInMemoryTaskStore::default();
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
let count = rt.block_on(store.tenant_count());
assert_eq!(count, 0, "default store should have no tenants");
}
#[tokio::test]
async fn run_eviction_all_runs_without_error() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("t1", async {
store
.save(&make_task("task-a", TaskState::Completed))
.await
.unwrap();
})
.await;
TenantContext::scope("t2", async {
store
.save(&make_task("task-b", TaskState::Working))
.await
.unwrap();
})
.await;
store.run_eviction_all().await;
}
#[tokio::test]
async fn get_store_double_check_path() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("racer", async {
store
.save(&make_task("t1", TaskState::Submitted))
.await
.unwrap();
store
.save(&make_task("t2", TaskState::Working))
.await
.unwrap();
let count = store.count().await.unwrap();
assert_eq!(count, 2, "both tasks should be in same tenant store");
})
.await;
assert_eq!(
store.tenant_count().await,
1,
"should have exactly 1 tenant"
);
}
#[test]
fn default_tenant_store_config() {
let cfg = TenantStoreConfig::default();
assert_eq!(cfg.max_tenants, 1000);
}
#[tokio::test]
async fn run_eviction_all_evicts_in_every_tenant() {
let store = TenantAwareInMemoryTaskStore::with_config(TenantStoreConfig {
per_tenant: TaskStoreConfig {
task_ttl: Some(std::time::Duration::from_millis(1)),
..TaskStoreConfig::default()
},
..TenantStoreConfig::default()
});
for tenant in ["tenant-a", "tenant-b"] {
TenantContext::scope(tenant, async {
store
.save(&make_task("t1", TaskState::Completed))
.await
.expect("save");
})
.await;
}
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
store.run_eviction_all().await;
for tenant in ["tenant-a", "tenant-b"] {
let still_there = TenantContext::scope(tenant, async {
store.get(&TaskId::new("t1")).await.expect("get")
})
.await;
assert!(
still_there.is_none(),
"a terminal task past its TTL must be evicted in {tenant}; \
surviving means run_eviction_all did not reach this partition"
);
}
}
#[test]
fn in_memory_tenant_store_advertises_idempotency_support() {
let store = TenantAwareInMemoryTaskStore::new();
assert!(
store.supports_idempotency(),
"the tenant-aware in-memory store partitions by tenant, which is \
exactly the scope a key needs; it must advertise support"
);
}
#[test]
fn every_tenant_store_config_setter_sets_its_field() {
let d = TenantStoreConfig::default();
let per_tenant =
TaskStoreConfig::default().with_max_page_size(d.per_tenant.max_page_size + 1);
let cfg = TenantStoreConfig::default()
.with_per_tenant(per_tenant)
.with_max_tenants(d.max_tenants + 1);
assert_eq!(cfg.per_tenant.max_page_size, d.per_tenant.max_page_size + 1);
assert_eq!(cfg.max_tenants, d.max_tenants + 1);
}
}
#[cfg(test)]
mod idempotency_isolation_tests {
use super::*;
use crate::store::task_store::IdempotencyClaim;
use a2a_protocol_types::message::MessageId;
const KEY: &str = "8f14e45fceea167a5a36dedd4bea2543";
#[tokio::test]
async fn one_tenants_key_never_names_another_tenants_task() {
let store = TenantAwareInMemoryTaskStore::new();
let a = TenantContext::scope("tenant-a", async {
store
.claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("task-a"))
.await
.unwrap()
})
.await;
assert_eq!(a, IdempotencyClaim::Claimed);
let b = TenantContext::scope("tenant-b", async {
store
.claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("task-b"))
.await
.unwrap()
})
.await;
assert_eq!(
b,
IdempotencyClaim::Claimed,
"tenant-b must claim its own key, not replay tenant-a's task"
);
}
#[tokio::test]
async fn a_retry_within_one_tenant_still_replays() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("tenant-a", async {
let first = store
.claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("task-a"))
.await
.unwrap();
assert_eq!(first, IdempotencyClaim::Claimed);
let retry = store
.claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("task-other"))
.await
.unwrap();
assert_eq!(retry, IdempotencyClaim::Replay(TaskId::new("task-a")));
})
.await;
}
#[tokio::test]
async fn release_frees_the_key_in_its_own_tenant_only() {
let store = TenantAwareInMemoryTaskStore::new();
for tenant in ["tenant-a", "tenant-b"] {
TenantContext::scope(tenant, async {
let claim = store
.claim_idempotency_key(KEY, &MessageId::new("m1"), &TaskId::new("t1"))
.await
.unwrap();
assert_eq!(claim, IdempotencyClaim::Claimed);
})
.await;
}
TenantContext::scope("tenant-a", async {
store.release_idempotency_key(KEY).await.unwrap();
let reclaim = store
.claim_idempotency_key(KEY, &MessageId::new("m2"), &TaskId::new("t2"))
.await
.unwrap();
assert_eq!(reclaim, IdempotencyClaim::Claimed, "released, so free");
})
.await;
TenantContext::scope("tenant-b", async {
let held = store
.claim_idempotency_key(KEY, &MessageId::new("m2"), &TaskId::new("t2"))
.await
.unwrap();
assert_ne!(
held,
IdempotencyClaim::Claimed,
"tenant-a's release must not free tenant-b's key"
);
})
.await;
}
#[tokio::test]
async fn releasing_under_an_unknown_tenant_creates_no_partition() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("never-seen", async {
store.release_idempotency_key(KEY).await.unwrap();
})
.await;
assert_eq!(store.tenant_count().await, 0);
}
}
#[cfg(test)]
mod event_log_tests {
use super::*;
use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
use a2a_protocol_types::task::{ContextId, TaskState, TaskStatus};
fn make_task(id: &str, state: TaskState) -> Task {
Task {
id: TaskId::new(id),
context_id: ContextId::new("ctx-default"),
status: TaskStatus::new(state),
history: None,
artifacts: None,
metadata: None,
}
}
fn event(id: &str, state: TaskState) -> StreamResponse {
StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
task_id: TaskId::new(id),
context_id: ContextId::new("ctx-default"),
status: TaskStatus::new(state),
metadata: None,
})
}
#[tokio::test]
async fn the_log_reports_its_earliest_position_per_tenant() {
let store = TenantAwareInMemoryTaskStore::new();
let id = TaskId::new("shared");
TenantContext::scope("tenant-a", async {
store
.save(&make_task("shared", TaskState::Working))
.await
.expect("save a");
assert_eq!(store.earliest_event_seq(&id).await.expect("empty"), None);
for seq in [5, 6, 7] {
store
.append_event(&id, seq, &event("shared", TaskState::Working))
.await
.expect("append a");
}
assert_eq!(store.earliest_event_seq(&id).await.expect("a"), Some(5));
})
.await;
TenantContext::scope("tenant-b", async {
assert_eq!(
store.earliest_event_seq(&id).await.expect("unknown"),
None,
"tenant-a's log is not tenant-b's"
);
})
.await;
assert_eq!(
store.tenant_count().await,
1,
"asking for tenant-b's log must not create its partition"
);
}
#[tokio::test]
async fn the_dispatcher_reports_that_it_keeps_a_log() {
assert!(TenantAwareInMemoryTaskStore::new().supports_event_log());
}
#[tokio::test]
async fn one_tenants_task_id_never_reaches_another_tenants_log() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("tenant-a", async {
store
.save(&make_task("shared", TaskState::Working))
.await
.unwrap();
store
.append_event(
&TaskId::new("shared"),
1,
&event("shared", TaskState::Working),
)
.await
.unwrap();
})
.await;
TenantContext::scope("tenant-b", async {
store
.save(&make_task("shared", TaskState::Working))
.await
.unwrap();
assert_eq!(
store.last_event_seq(&TaskId::new("shared")).await.unwrap(),
0
);
assert!(
store
.read_events(&TaskId::new("shared"), 0, 10)
.await
.unwrap()
.is_empty(),
"tenant-b must not be handed tenant-a's events"
);
})
.await;
TenantContext::scope("tenant-a", async {
assert_eq!(
store.last_event_seq(&TaskId::new("shared")).await.unwrap(),
1
);
})
.await;
}
#[tokio::test]
async fn reading_under_an_unknown_tenant_creates_no_partition() {
let store = TenantAwareInMemoryTaskStore::new();
TenantContext::scope("never-seen", async {
assert_eq!(store.last_event_seq(&TaskId::new("t1")).await.unwrap(), 0);
assert!(
store
.read_events(&TaskId::new("t1"), 0, 10)
.await
.unwrap()
.is_empty()
);
})
.await;
assert_eq!(store.tenant_count().await, 0);
}
}