use crate::auth::RequestContext;
use crate::envelope::{
codes, A2aHeaders, A2aMethod, JsonRpcError, JsonRpcRequest, JsonRpcResponse,
};
use crate::error::{A2aBuilderError, A2aError};
use crate::handler::A2aHandler;
use crate::types::{
CancelTaskParams, DeletePushNotificationConfigParams, GetExtendedAgentCardParams,
GetPushNotificationConfigParams, GetTaskParams, ListPushNotificationConfigsParams,
ListTasksParams, Message, PushNotificationConfigParams, SendMessageParams, SendMessageResult,
SubscribeToTaskParams, Task, TaskStatus,
};
use bytes::Bytes;
use klieo_auth_common::Authenticator;
use klieo_core::{DurableName, Headers, Msg, Pubsub};
const A2A_CANCEL_SUBJECT_PREFIX: &str = "klieo.a2a.cancel.";
const A2A_CANCEL_SUBJECT_PATTERN: &str = "klieo.a2a.cancel.>";
const A2A_CANCEL_LOG_TARGET: &str = "a2a.cancel";
use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json::Value;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::Semaphore;
use tokio_stream::StreamExt;
use tracing::{error, instrument, warn};
use tracing_opentelemetry::OpenTelemetrySpanExt as _;
pub(crate) const DEFAULT_PUBLISH_PERMITS: usize = 64;
pub const LEADER_TTL: std::time::Duration = std::time::Duration::from_secs(5);
const A2A_LEADER_KEY_PREFIX: &str = "a2a.";
const A2A_OWNERSHIP_KEY_PREFIX: &str = "a2a.";
fn ownership_key(task_id: &str) -> String {
format!("{A2A_OWNERSHIP_KEY_PREFIX}{task_id}")
}
pub type TaskEventStream = Pin<Box<dyn futures::Stream<Item = TaskEvent> + Send>>;
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
#[non_exhaustive]
pub struct TaskEvent {
pub task_id: String,
pub status: TaskStatus,
pub message: Option<Message>,
pub final_event: bool,
#[serde(default)]
pub event_id: u64,
}
impl TaskEvent {
pub fn new(
task_id: impl Into<String>,
status: TaskStatus,
message: Option<Message>,
final_event: bool,
) -> Self {
Self {
task_id: task_id.into(),
status,
message,
final_event,
event_id: 0,
}
}
#[must_use]
pub fn with_event_id(mut self, event_id: u64) -> Self {
self.event_id = event_id;
self
}
#[allow(dead_code)]
pub(crate) fn synthetic_corrupt(task_id: &str) -> Self {
Self {
task_id: task_id.into(),
status: TaskStatus::Failed,
message: None,
final_event: true,
event_id: 0,
}
}
}
struct LeaderHoldStream<S> {
inner: S,
_leader: Option<klieo_core::LeaderHandle>,
_ownership: Option<klieo_core::OwnershipHandle>,
}
impl<S: futures::Stream + Unpin> futures::Stream for LeaderHoldStream<S> {
type Item = S::Item;
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<S::Item>> {
std::pin::Pin::new(&mut self.inner).poll_next(cx)
}
}
enum LeaderProbe {
NoRegistry,
Alive,
Dead,
}
#[derive(Clone)]
pub struct TaskEventSink {
pubsub: Arc<dyn klieo_core::Pubsub>,
permits: Arc<Semaphore>,
}
impl TaskEventSink {
pub fn new(pubsub: Arc<dyn klieo_core::Pubsub>) -> Self {
Self::with_permits(pubsub, Arc::new(Semaphore::new(DEFAULT_PUBLISH_PERMITS)))
}
pub fn with_permits(pubsub: Arc<dyn klieo_core::Pubsub>, permits: Arc<Semaphore>) -> Self {
Self { pubsub, permits }
}
#[instrument(
skip_all,
fields(
messaging.system = "klieo-bus",
messaging.destination = tracing::field::Empty,
messaging.operation = "publish",
messaging.message.payload_size_bytes = tracing::field::Empty,
klieo.stream_id = %event.task_id,
),
err,
)]
pub async fn send(&self, event: TaskEvent) -> Result<(), A2aError> {
klieo_core::validate_subject_token(&event.task_id)?;
let subject = format!("klieo.a2a.task.{}", event.task_id);
tracing::Span::current().record("messaging.destination", subject.as_str());
let bytes = serde_json::to_vec(&event).map_err(|e| A2aError::Internal {
source: Box::new(e),
})?;
tracing::Span::current().record("messaging.message.payload_size_bytes", bytes.len());
let mut headers = klieo_core::Headers::default();
let cx = tracing::Span::current().context();
klieo_core::inject_traceparent(&mut headers, &cx);
let permit = match self.permits.clone().try_acquire_owned() {
Ok(p) => p,
Err(_) => {
tracing::warn!(
target: "a2a.fanout",
subject = %subject,
available_permits = self.permits.available_permits(),
"task event publish dropped: concurrency cap reached",
);
return Ok(()); }
};
let result = self
.pubsub
.publish(&subject, bytes::Bytes::from(bytes), headers)
.await;
drop(permit);
result?;
Ok(())
}
}
#[derive(Default)]
pub struct A2aDispatcherBuilder {
handler: Option<Arc<dyn A2aHandler>>,
authenticator: Option<Arc<dyn Authenticator>>,
pubsub: Option<Arc<dyn klieo_core::Pubsub>>,
publish_concurrency: Option<usize>,
subscribe_cancels: bool,
leader_kv: Option<Arc<dyn klieo_core::KvStore>>,
tenant_kv: Option<Arc<dyn klieo_core::KvStore>>,
tenant_strict: bool,
leader_ttl: Option<std::time::Duration>,
leader_heartbeat_interval: Option<std::time::Duration>,
max_failover_attempts: Option<u32>,
kv_reaper_interval: Option<std::time::Duration>,
profile: klieo_core::DeploymentProfile,
}
impl A2aDispatcherBuilder {
pub fn handler(mut self, handler: Arc<dyn A2aHandler>) -> Self {
self.handler = Some(handler);
self
}
pub fn authenticator(mut self, authenticator: Arc<dyn Authenticator>) -> Self {
self.authenticator = Some(authenticator);
self
}
pub fn profile(mut self, profile: klieo_core::DeploymentProfile) -> Self {
self.profile = profile;
self
}
pub fn pubsub(mut self, pubsub: Arc<dyn klieo_core::Pubsub>) -> Self {
self.pubsub = Some(pubsub);
self
}
#[must_use]
pub fn with_in_process_pubsub(mut self) -> Self {
self.pubsub = Some(klieo_bus_memory::MemoryBus::new().pubsub.clone());
self
}
#[must_use]
pub fn with_publish_concurrency(mut self, permits: usize) -> Self {
self.publish_concurrency = Some(permits);
self
}
#[must_use]
pub fn with_cancel_subscription(mut self) -> Self {
self.subscribe_cancels = true;
self
}
#[must_use]
pub fn with_leader_election(mut self, kv: Arc<dyn klieo_core::KvStore>) -> Self {
self.leader_kv = Some(kv);
self
}
#[must_use]
pub fn with_tenant_binding(mut self, kv: Arc<dyn klieo_core::KvStore>) -> Self {
self.tenant_kv = Some(kv);
self.tenant_strict = false;
self
}
#[must_use]
pub fn with_tenant_binding_strict(mut self, kv: Arc<dyn klieo_core::KvStore>) -> Self {
self.tenant_kv = Some(kv);
self.tenant_strict = true;
self
}
#[must_use]
pub fn with_leader_ttl(mut self, ttl: std::time::Duration) -> Self {
self.leader_ttl = Some(ttl);
self
}
#[must_use]
pub fn with_leader_heartbeat_interval(mut self, interval: std::time::Duration) -> Self {
self.leader_heartbeat_interval = Some(interval);
self
}
#[must_use]
pub fn with_max_failover_attempts(mut self, cap: u32) -> Self {
self.max_failover_attempts = Some(cap);
self
}
#[must_use]
pub fn with_kv_reaper(mut self, interval: std::time::Duration) -> Self {
self.kv_reaper_interval = Some(interval);
self
}
pub fn build(self) -> Result<A2aDispatcher, A2aBuilderError> {
if self.subscribe_cancels {
return Err(A2aBuilderError::CancelRequiresArc);
}
self.build_inner()
}
fn build_inner(self) -> Result<A2aDispatcher, A2aBuilderError> {
let handler = self.handler.ok_or(A2aBuilderError::MissingHandler)?;
let authenticator = self
.authenticator
.ok_or(A2aBuilderError::MissingAuthenticator)?;
let pubsub = self.pubsub.ok_or(A2aBuilderError::MissingPubsub)?;
let permits = self.publish_concurrency.unwrap_or(DEFAULT_PUBLISH_PERMITS);
let leader_registry = self.leader_kv.map(|kv| {
klieo_core::LeaderRegistry::new(
kv,
"klieo-leaders".into(),
uuid::Uuid::new_v4().to_string(),
)
});
let profile = self.profile;
profile.validate(
self.tenant_kv.is_some(),
Some(authenticator.allows_anonymous()),
)?;
let tenant_strict = self.tenant_strict || profile.requires_strict_binding();
let ownership_registry = self.tenant_kv.map(|kv| {
let bucket = "klieo-tenants".into();
if tenant_strict {
klieo_core::OwnershipRegistry::new_strict(kv, bucket)
} else {
klieo_core::OwnershipRegistry::new(kv, bucket)
}
});
if profile.requires_strict_binding() || profile.requires_named_principal() {
tracing::warn!(
target: "klieo.security",
cwe = 639,
"regulated multi-tenant profile active on this replica; \
cross-replica tenant isolation assumes ALL replicas run the \
same profile — a lenient peer reintroduces CWE-639. Fleet \
homogeneity is NOT verified by this replica."
);
}
let leader_ttl = self.leader_ttl.unwrap_or(LEADER_TTL);
let leader_heartbeat_interval = self.leader_heartbeat_interval.unwrap_or(leader_ttl / 2);
let max_failover_attempts = self
.max_failover_attempts
.unwrap_or(klieo_core::FAILOVER_ATTEMPT_CAP);
Ok(A2aDispatcher {
handler,
authenticator,
pubsub,
cancel_registry: klieo_core::CancelRegistry::new(),
publish_permits: Arc::new(Semaphore::new(permits)),
leader_registry,
ownership_registry,
leader_ttl,
leader_heartbeat_interval,
max_failover_attempts,
kv_reaper_interval: self.kv_reaper_interval,
})
}
pub fn build_arc(self) -> Result<Arc<A2aDispatcher>, A2aBuilderError> {
let spawn_subscriber = self.subscribe_cancels;
let dispatcher = Arc::new(self.build_inner()?);
if spawn_subscriber {
dispatcher.with_cancel_subscription();
}
Ok(dispatcher)
}
}
pub struct A2aDispatcher {
handler: Arc<dyn A2aHandler>,
authenticator: Arc<dyn Authenticator>,
pubsub: Arc<dyn klieo_core::Pubsub>,
cancel_registry: klieo_core::CancelRegistry<String>,
publish_permits: Arc<Semaphore>,
leader_registry: Option<klieo_core::LeaderRegistry>,
ownership_registry: Option<klieo_core::OwnershipRegistry>,
leader_ttl: std::time::Duration,
leader_heartbeat_interval: std::time::Duration,
max_failover_attempts: u32,
kv_reaper_interval: Option<std::time::Duration>,
}
impl A2aDispatcher {
pub fn builder() -> A2aDispatcherBuilder {
A2aDispatcherBuilder::default()
}
pub fn new(
handler: Arc<dyn A2aHandler>,
authenticator: Arc<dyn Authenticator>,
pubsub: Arc<dyn klieo_core::Pubsub>,
) -> Self {
Self {
handler,
authenticator,
pubsub,
cancel_registry: klieo_core::CancelRegistry::new(),
publish_permits: Arc::new(Semaphore::new(DEFAULT_PUBLISH_PERMITS)),
leader_registry: None,
ownership_registry: None,
leader_ttl: LEADER_TTL,
leader_heartbeat_interval: LEADER_TTL / 2,
max_failover_attempts: klieo_core::FAILOVER_ATTEMPT_CAP,
kv_reaper_interval: None,
}
}
pub fn leader_ttl(&self) -> std::time::Duration {
self.leader_ttl
}
pub fn leader_heartbeat_interval(&self) -> std::time::Duration {
self.leader_heartbeat_interval
}
pub fn max_failover_attempts(&self) -> u32 {
self.max_failover_attempts
}
pub fn kv_reaper_interval(&self) -> Option<std::time::Duration> {
self.kv_reaper_interval
}
pub fn leader_registry(&self) -> Option<&klieo_core::LeaderRegistry> {
self.leader_registry.as_ref()
}
pub fn ownership_registry(&self) -> Option<&klieo_core::OwnershipRegistry> {
self.ownership_registry.as_ref()
}
#[must_use]
pub fn with_publish_concurrency(mut self, permits: usize) -> Self {
self.publish_permits = Arc::new(Semaphore::new(permits));
self
}
pub fn with_in_process_pubsub(
handler: Arc<dyn A2aHandler>,
authenticator: Arc<dyn Authenticator>,
) -> Self {
let bus = klieo_bus_memory::MemoryBus::new();
Self::new(handler, authenticator, bus.pubsub.clone())
}
#[cfg(feature = "test-fixtures")]
pub fn local(handler: Arc<dyn A2aHandler>) -> Self {
let bus = klieo_bus_memory::MemoryBus::new();
Self::new(
handler,
Arc::new(klieo_auth_common::AllowAnonymous),
bus.pubsub.clone(),
)
}
pub fn authenticator(&self) -> &Arc<dyn Authenticator> {
&self.authenticator
}
pub fn handler(&self) -> &Arc<dyn A2aHandler> {
&self.handler
}
pub fn event_sink(&self) -> TaskEventSink {
TaskEventSink::with_permits(self.pubsub.clone(), self.publish_permits.clone())
}
pub fn pubsub(&self) -> &Arc<dyn klieo_core::Pubsub> {
&self.pubsub
}
pub fn publish_permits(&self) -> &Arc<Semaphore> {
&self.publish_permits
}
pub fn cancel_registry(&self) -> &klieo_core::CancelRegistry<String> {
&self.cancel_registry
}
pub async fn publish_cancel(&self, task_id: &str) -> Result<(), A2aError> {
klieo_core::cancel::publish_cancel_signal(&self.pubsub, A2A_CANCEL_SUBJECT_PREFIX, task_id)
.await?;
Ok(())
}
pub fn with_cancel_subscription(self: &Arc<Self>) {
klieo_core::cancel::spawn_wildcard_cancel_subscriber(
self.pubsub.clone(),
A2A_CANCEL_SUBJECT_PATTERN.to_string(),
A2A_CANCEL_SUBJECT_PREFIX.to_string(),
self.cancel_registry.clone(),
A2A_CANCEL_LOG_TARGET,
);
}
#[allow(clippy::too_many_lines)]
#[instrument(
skip_all,
fields(rpc.system = "klieo-a2a", rpc.method = tracing::field::Empty),
)]
pub async fn dispatch(&self, ctx: &RequestContext, payload: &[u8]) -> JsonRpcResponse {
let req: JsonRpcRequest = match serde_json::from_slice(payload) {
Ok(r) => r,
Err(e) => return A2aError::Json(e).to_json_rpc_error(Value::Null),
};
tracing::Span::current().record("rpc.method", req.method.as_str());
let method = match A2aMethod::from_str(&req.method) {
Ok(m) => m,
Err(e) => return e.to_json_rpc_error(req.id.clone()),
};
let id = req.id.clone();
if let Some(identity) = ctx.caller.as_ref() {
if let Err(e) = self
.authenticator
.authorize_method(identity, method.wire_name())
.await
{
warn!(
target: "security",
error = %e,
method = %req.method,
"a2a authorize_method rejected",
);
return A2aError::Unauthorized(e.to_string()).to_json_rpc_error(id);
}
} else if !self.authenticator.allows_anonymous() {
warn!(
target: "security",
method = %req.method,
"a2a anonymous caller rejected (authenticator requires an identity)",
);
return A2aError::Unauthorized("anonymous caller not permitted".into())
.to_json_rpc_error(id);
}
let params_raw = req.params;
let result: Result<Value, A2aError> = match method {
A2aMethod::SendMessage => {
run_method::<SendMessageParams, _, _>(params_raw, |p| async move {
let result = self.handler.send_message(ctx, p).await?;
if let SendMessageResult::Task(task) = &result {
self.claim_owner_persistent(ctx, &task.id).await?;
}
Ok(result)
})
.await
}
A2aMethod::SendStreamingMessage => {
Err(A2aError::MethodNotFound("SendStreamingMessage".into()))
}
A2aMethod::GetTask => {
run_method::<GetTaskParams, _, _>(params_raw, |p| async move {
self.enforce_owner(ctx, &p.id).await?;
self.handler.get_task(ctx, p).await
})
.await
}
A2aMethod::ListTasks => {
run_method::<ListTasksParams, _, _>(params_raw, |p| async move {
let mut result = self.handler.list_tasks(ctx, p).await?;
self.retain_owned_tasks(ctx, &mut result.tasks).await?;
Ok(result)
})
.await
}
A2aMethod::CancelTask => {
run_method::<CancelTaskParams, _, _>(params_raw, |p| async move {
self.enforce_owner(ctx, &p.id).await?;
self.handler.cancel_task(ctx, p).await
})
.await
}
A2aMethod::SubscribeToTask => Err(A2aError::MethodNotFound("SubscribeToTask".into())),
A2aMethod::CreateTaskPushNotificationConfig => {
run_method::<PushNotificationConfigParams, _, _>(params_raw, |p| async move {
crate::ssrf::validate_callback_url(&p.url)?;
self.enforce_owner(ctx, &p.taskId).await?;
self.handler
.create_task_push_notification_config(ctx, p)
.await
})
.await
}
A2aMethod::GetTaskPushNotificationConfig => {
run_method::<GetPushNotificationConfigParams, _, _>(params_raw, |p| async move {
self.enforce_owner(ctx, &p.taskId).await?;
self.handler.get_task_push_notification_config(ctx, p).await
})
.await
}
A2aMethod::ListTaskPushNotificationConfigs => {
run_method::<ListPushNotificationConfigsParams, _, _>(params_raw, |p| async move {
self.enforce_owner(ctx, &p.taskId).await?;
self.handler
.list_task_push_notification_configs(ctx, p)
.await
})
.await
}
A2aMethod::DeleteTaskPushNotificationConfig => {
run_method::<DeletePushNotificationConfigParams, _, _>(params_raw, |p| async move {
self.enforce_owner(ctx, &p.taskId).await?;
self.handler
.delete_task_push_notification_config(ctx, p)
.await
})
.await
}
A2aMethod::GetExtendedAgentCard => {
run_method::<GetExtendedAgentCardParams, _, _>(params_raw, |p| async move {
self.handler.get_extended_agent_card(ctx, p).await
})
.await
}
};
match result {
Ok(value) => JsonRpcResponse {
jsonrpc: "2.0".into(),
id,
result: Some(value),
error: None,
},
Err(e) => {
log_internal_before_wire_seam(&e, &req.method);
e.to_json_rpc_error(id)
}
}
}
pub async fn handle_request(&self, headers: A2aHeaders, payload: &[u8]) -> JsonRpcResponse {
match self.authenticator.authenticate(&headers, payload).await {
Ok(identity) => {
let ctx = RequestContext::new(headers, Some(identity));
self.dispatch(&ctx, payload).await
}
Err(e) => {
warn!(target: "security", error = %e, "a2a auth rejected");
let req_id = serde_json::from_slice::<JsonRpcRequest>(payload)
.map(|r| r.id)
.unwrap_or(Value::Null);
JsonRpcResponse {
jsonrpc: "2.0".into(),
id: req_id,
result: None,
error: Some(JsonRpcError {
code: codes::UNAUTHENTICATED,
message: "Unauthenticated".into(),
data: None,
}),
}
}
}
}
pub async fn handle_streaming(
&self,
headers: A2aHeaders,
payload: &[u8],
task_store: &crate::task_store::A2aTaskStore,
cancel: tokio_util::sync::CancellationToken,
last_event_id: Option<u64>,
resume_buffer: Arc<dyn klieo_core::resume::ResumeBuffer>,
) -> Result<TaskEventStream, A2aError> {
match self.authenticator.authenticate(&headers, payload).await {
Ok(identity) => {
let ctx = RequestContext::new(headers, Some(identity)).with_cancel(cancel);
self.dispatch_streaming(&ctx, payload, task_store, last_event_id, resume_buffer)
.await
}
Err(e) => {
warn!(target: "security", error = %e, "a2a auth rejected (streaming path)");
Err(A2aError::Unauthorized(e.to_string()))
}
}
}
#[instrument(
skip_all,
fields(rpc.system = "klieo-a2a", rpc.method = tracing::field::Empty),
err,
)]
pub async fn dispatch_streaming(
&self,
ctx: &RequestContext,
payload: &[u8],
task_store: &crate::task_store::A2aTaskStore,
last_event_id: Option<u64>,
resume_buffer: Arc<dyn klieo_core::resume::ResumeBuffer>,
) -> Result<TaskEventStream, A2aError> {
let req: JsonRpcRequest =
serde_json::from_slice(payload).map_err(|e| A2aError::InvalidParams(e.to_string()))?;
tracing::Span::current().record("rpc.method", req.method.as_str());
let method = A2aMethod::from_str(&req.method)
.map_err(|_| A2aError::MethodNotFound(req.method.clone()))?;
if let Some(identity) = ctx.caller.as_ref() {
if let Err(e) = self
.authenticator
.authorize_method(identity, method.wire_name())
.await
{
warn!(
target: "security",
error = %e,
method = %req.method,
"a2a authorize_method rejected (streaming path)",
);
return Err(A2aError::Unauthorized(e.to_string()));
}
} else if !self.authenticator.allows_anonymous() {
warn!(
target: "security",
method = %req.method,
"a2a anonymous caller rejected (streaming path; authenticator requires an identity)",
);
return Err(A2aError::Unauthorized(
"anonymous caller not permitted".into(),
));
}
let params_raw = req.params;
match method {
A2aMethod::SendStreamingMessage => {
let params: SendMessageParams = serde_json::from_value(params_raw)
.map_err(|e| A2aError::InvalidParams(e.to_string()))?;
self.dispatch_send_streaming(ctx, params, task_store).await
}
A2aMethod::SubscribeToTask => {
let params: SubscribeToTaskParams = serde_json::from_value(params_raw)
.map_err(|e| A2aError::InvalidParams(e.to_string()))?;
self.dispatch_subscribe_to_task(
ctx,
params,
task_store,
last_event_id,
resume_buffer,
)
.await
}
_ => Err(A2aError::MethodNotFound(req.method)),
}
}
#[instrument(
skip_all,
fields(
rpc.system = "klieo-a2a",
rpc.method = "SendStreamingMessage",
klieo.stream_id = tracing::field::Empty,
),
err,
)]
async fn dispatch_send_streaming(
&self,
ctx: &RequestContext,
params: SendMessageParams,
task_store: &crate::task_store::A2aTaskStore,
) -> Result<TaskEventStream, A2aError> {
self.run_streaming_invoke(ctx, params, task_store, None)
.await
}
async fn run_streaming_invoke(
&self,
ctx: &RequestContext,
params: SendMessageParams,
task_store: &crate::task_store::A2aTaskStore,
existing_leader: Option<klieo_core::LeaderHandle>,
) -> Result<TaskEventStream, A2aError> {
match self
.handler
.send_streaming_message(ctx, params.clone())
.await
{
Ok(stream) => return Ok(stream),
Err(A2aError::MethodNotFound(_)) => {}
Err(other) => return Err(other),
}
let payload_bytes_for_failover = match existing_leader {
Some(_) => None,
None => match serde_json::to_vec(¶ms) {
Ok(v) => Some(Bytes::from(v)),
Err(e) => {
warn!(
target: "a2a.failover",
error = %e,
"send-streaming params serialise for failover failed; \
proceeding without cached payload",
);
None
}
},
};
let principal_for_failover = match existing_leader {
Some(_) => None,
None => ctx
.caller
.as_ref()
.filter(|id| !id.is_anonymous())
.map(|id| id.as_str().to_string()),
};
let result = self.handler.send_message(ctx, params).await?;
let task = match result {
SendMessageResult::Task(t) => t,
SendMessageResult::Message(_) => {
return Ok(Box::pin(futures::stream::empty()));
}
};
let task_id = task.id.clone();
tracing::Span::current().record("klieo.stream_id", task_id.as_str());
let terminal = task.status.is_terminal();
let event_id = task_store.next_event_id(&task_id);
let initial = TaskEvent::new(
task_id.clone(),
task.status,
task.history.last().cloned(),
terminal,
)
.with_event_id(event_id);
if terminal {
return Ok(Box::pin(futures::stream::once(futures::future::ready(
initial,
))));
}
let initial_stream = futures::stream::once(futures::future::ready(initial));
let tail = self.task_event_stream(task_id.clone()).await?;
let chained = futures::StreamExt::chain(initial_stream, tail);
let leader_handle = match existing_leader {
Some(handle) => Some(handle),
None => {
self.try_claim_leader(&task_id, payload_bytes_for_failover, principal_for_failover)
.await
}
};
let ownership_handle = self.try_claim_ownership(ctx, &task_id).await?;
Ok(Box::pin(LeaderHoldStream {
inner: chained,
_leader: leader_handle,
_ownership: ownership_handle,
}))
}
async fn try_claim_ownership(
&self,
ctx: &RequestContext,
task_id: &str,
) -> Result<Option<klieo_core::OwnershipHandle>, A2aError> {
let Some(registry) = self.ownership_registry.as_ref() else {
return Ok(None);
};
let Some(caller) = ctx.caller.as_ref() else {
return Ok(None);
};
if caller.is_anonymous() {
return Ok(None);
}
let key = ownership_key(task_id);
let principal = caller.as_str().to_string();
match registry.claim_guarded(key, principal).await {
klieo_core::OwnershipClaim::Claimed(handle) => Ok(Some(handle)),
klieo_core::OwnershipClaim::Proceed => Ok(None),
klieo_core::OwnershipClaim::Unavailable => Err(A2aError::Server(
"ownership store unavailable; invoke denied (strict tenant binding)".into(),
)),
_ => Err(A2aError::Server(
"ownership claim inconclusive; invoke denied".into(),
)),
}
}
async fn claim_owner_persistent(
&self,
ctx: &RequestContext,
task_id: &str,
) -> Result<(), A2aError> {
let Some(registry) = self.ownership_registry.as_ref() else {
return Ok(());
};
let Some(caller) = ctx.caller.as_ref() else {
return Ok(());
};
if caller.is_anonymous() {
return Ok(());
}
let key = ownership_key(task_id);
match registry
.record_owner(key, caller.as_str().to_string())
.await
{
klieo_core::OwnershipRecord::Recorded | klieo_core::OwnershipRecord::Proceed => Ok(()),
klieo_core::OwnershipRecord::Unavailable => Err(A2aError::Server(
"ownership store unavailable; invoke denied (strict tenant binding)".into(),
)),
_ => Err(A2aError::Server(
"ownership record inconclusive; invoke denied".into(),
)),
}
}
async fn enforce_owner(&self, ctx: &RequestContext, task_id: &str) -> Result<(), A2aError> {
let Some(registry) = self.ownership_registry.as_ref() else {
return Ok(());
};
let Some(caller) = ctx.caller.as_ref() else {
return Ok(());
};
if caller.is_anonymous() {
return Ok(());
}
let key = ownership_key(task_id);
match registry.check_owner(&key, caller.as_str()).await {
klieo_core::OwnershipCheck::Allowed => Ok(()),
klieo_core::OwnershipCheck::Denied => {
warn!(
target: "a2a.tenants",
task_id = %task_id,
principal = %caller.as_str(),
"ownership mismatch on SubscribeToTask; denying as TaskNotFound",
);
Err(A2aError::TaskNotFound(task_id.to_string()))
}
klieo_core::OwnershipCheck::Unavailable => {
warn!(
target: "a2a.tenants",
task_id = %task_id,
principal = %caller.as_str(),
"ownership store unavailable; request denied (strict tenant binding)",
);
Err(A2aError::OwnershipUnavailable)
}
other => {
warn!(
target: "a2a.tenants",
task_id = %task_id,
principal = %caller.as_str(),
verdict = ?other,
"inconclusive ownership verdict; request denied",
);
Err(A2aError::OwnershipUnavailable)
}
}
}
async fn retain_owned_tasks(
&self,
ctx: &RequestContext,
tasks: &mut Vec<Task>,
) -> Result<(), A2aError> {
let Some(registry) = self.ownership_registry.as_ref() else {
return Ok(());
};
let Some(caller) = ctx.caller.as_ref() else {
return Ok(());
};
if caller.is_anonymous() {
return Ok(());
}
let verdicts = futures::future::join_all(tasks.iter().map(|task| {
let key = ownership_key(&task.id);
async move { registry.check_owner(&key, caller.as_str()).await }
}))
.await;
let mut owned = Vec::with_capacity(tasks.len());
for (task, verdict) in std::mem::take(tasks).into_iter().zip(verdicts) {
match verdict {
klieo_core::OwnershipCheck::Allowed => owned.push(task),
klieo_core::OwnershipCheck::Denied => {}
klieo_core::OwnershipCheck::Unavailable => {
warn!(
target: "a2a.tenants",
principal = %caller.as_str(),
"ownership store unavailable; ListTasks denied (strict tenant binding)",
);
return Err(A2aError::OwnershipUnavailable);
}
other => {
warn!(
target: "a2a.tenants",
principal = %caller.as_str(),
verdict = ?other,
"inconclusive ownership verdict; ListTasks denied",
);
return Err(A2aError::OwnershipUnavailable);
}
}
}
*tasks = owned;
Ok(())
}
async fn try_claim_leader(
&self,
task_id: &str,
payload: Option<Bytes>,
principal: Option<String>,
) -> Option<klieo_core::LeaderHandle> {
let registry = self.leader_registry.as_ref()?;
let key = format!("{A2A_LEADER_KEY_PREFIX}{task_id}");
match registry
.claim_with_heartbeat(
key,
self.leader_ttl,
self.leader_heartbeat_interval,
payload,
principal,
)
.await
{
Ok(handle) => Some(handle),
Err(e) => {
warn!(
target: "a2a.leader",
task_id = %task_id,
error = %e,
"leader claim failed; degrading to no-claim (orphan detection \
disabled for this stream)",
);
None
}
}
}
async fn write_orphan_terminal_frame(
&self,
task_id: &str,
resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
) -> Option<()> {
let max_id = max_event_id(resume_buffer, task_id).await?;
let next_id = max_id + 1;
let frame = leader_died_sse_frame_bytes(task_id, next_id);
if let Err(e) = resume_buffer.record(task_id, next_id, frame).await {
warn!(
target: "a2a.leader",
task_id = %task_id,
next_id,
error = %e,
"orphan terminal record failed; skipping orphan write",
);
return None;
}
if let Err(e) = resume_buffer.close(task_id).await {
warn!(
target: "a2a.leader",
task_id = %task_id,
error = %e,
"orphan terminal close failed; resume buffer may retain stale stream",
);
}
tracing::error!(
target: "a2a.leader",
task_id = %task_id,
next_id,
"stream leader died; emitted LEADER_DIED terminal frame at max+1",
);
Some(())
}
#[cfg(not(feature = "test-fixtures"))]
async fn handle_dead_leader_orphan(
&self,
ctx: &RequestContext,
task_id: &str,
task_store: &crate::task_store::A2aTaskStore,
resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
) -> Result<Option<TaskEventStream>, A2aError> {
self.handle_dead_leader_orphan_impl(ctx, task_id, task_store, resume_buffer)
.await
}
#[cfg(feature = "test-fixtures")]
pub async fn handle_dead_leader_orphan(
&self,
ctx: &RequestContext,
task_id: &str,
task_store: &crate::task_store::A2aTaskStore,
resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
) -> Result<Option<TaskEventStream>, A2aError> {
self.handle_dead_leader_orphan_impl(ctx, task_id, task_store, resume_buffer)
.await
}
async fn handle_dead_leader_orphan_impl(
&self,
ctx: &RequestContext,
task_id: &str,
task_store: &crate::task_store::A2aTaskStore,
resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
) -> Result<Option<TaskEventStream>, A2aError> {
let registry = match self.leader_registry.as_ref() {
Some(reg) => reg,
None => return Ok(None),
};
let key = format!("{A2A_LEADER_KEY_PREFIX}{task_id}");
let lookup = registry.lookup_entry_with_revision(&key).await;
let Some((entry, prior_rev)) = lookup_ok_or_log(&key, lookup) else {
return self.terminate_orphan_a2a(task_id, resume_buffer).await;
};
if !self.handler.is_idempotent() {
tracing::debug!(
target: "a2a.failover",
task_id,
"handler not idempotent; emitting terminate frame",
);
return self.terminate_orphan_a2a(task_id, resume_buffer).await;
}
if entry.attempt >= self.max_failover_attempts {
tracing::warn!(
target: "a2a.failover",
task_id,
attempt = entry.attempt,
cap = self.max_failover_attempts,
"failover attempt cap reached; emitting terminate frame",
);
return self.terminate_orphan_a2a(task_id, resume_buffer).await;
}
let Some(payload_bytes) = entry.payload.clone() else {
tracing::debug!(
target: "a2a.failover",
task_id,
"no cached payload on leader entry; emitting terminate frame",
);
return self.terminate_orphan_a2a(task_id, resume_buffer).await;
};
let new_handle = match registry
.claim_with_attempt_cas_and_heartbeat(
key.clone(),
self.leader_ttl,
self.leader_heartbeat_interval,
prior_rev,
&entry,
)
.await
{
Ok(h) => h,
Err(klieo_core::BusError::CasConflict { .. }) => {
tracing::info!(
target: "a2a.failover",
task_id,
"another follower won the CAS race; emitting terminate frame",
);
return self.terminate_orphan_a2a(task_id, resume_buffer).await;
}
Err(e) => {
tracing::warn!(
target: "a2a.failover",
task_id,
error = %e,
"CAS claim failed; emitting terminate frame",
);
return self.terminate_orphan_a2a(task_id, resume_buffer).await;
}
};
self.record_failover_marker(task_id, &entry, &new_handle, resume_buffer)
.await;
let parsed_params: SendMessageParams = match serde_json::from_slice(&payload_bytes) {
Ok(p) => p,
Err(e) => {
tracing::error!(
target: "a2a.failover",
task_id,
error = %e,
"cached payload parse failed; emitting terminate frame",
);
return self.terminate_orphan_a2a(task_id, resume_buffer).await;
}
};
let reinvoke_ctx = self.fabricate_reinvoke_ctx(ctx, entry.principal.as_deref());
let stream = self
.run_streaming_invoke(&reinvoke_ctx, parsed_params, task_store, Some(new_handle))
.await?;
Ok(Some(stream))
}
async fn record_failover_marker(
&self,
task_id: &str,
prior: &klieo_core::LeaderEntry,
new_handle: &klieo_core::LeaderHandle,
resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
) {
let Some(max) = max_event_id(resume_buffer, task_id).await else {
return;
};
let marker_id = max + 1;
let frame = failover_reinvoke_sse_frame_bytes(
task_id,
marker_id,
prior.attempt + 1,
new_handle.replica_id(),
);
if let Err(e) = resume_buffer.record(task_id, marker_id, frame).await {
warn!(
target: "a2a.failover",
task_id,
marker_id,
error = %e,
"failover-reinvoke marker record failed; continuing without marker",
);
}
}
fn fabricate_reinvoke_ctx(
&self,
ctx: &RequestContext,
principal: Option<&str>,
) -> RequestContext {
let caller = principal
.map(|p| klieo_auth_common::Identity::new(p.to_string()))
.or_else(|| Some(klieo_auth_common::Identity::anonymous()));
RequestContext::new(ctx.headers.clone(), caller).with_cancel(ctx.cancel.clone())
}
async fn terminate_orphan_a2a(
&self,
task_id: &str,
resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
) -> Result<Option<TaskEventStream>, A2aError> {
if self
.write_orphan_terminal_frame(task_id, resume_buffer)
.await
.is_some()
{
return Err(A2aError::LeaderDied {
stream_id: format!("{A2A_LEADER_KEY_PREFIX}{task_id}"),
});
}
Ok(None)
}
async fn probe_leader(&self, task_id: &str) -> LeaderProbe {
let Some(registry) = self.leader_registry.as_ref() else {
return LeaderProbe::NoRegistry;
};
let key = format!("{A2A_LEADER_KEY_PREFIX}{task_id}");
match registry.is_alive(&key).await {
Ok(true) => LeaderProbe::Alive,
Ok(false) => LeaderProbe::Dead,
Err(e) => {
warn!(
target: "a2a.leader",
task_id = %task_id,
error = %e,
"is_alive probe failed; treating as alive (fail-open per ADR-020)",
);
LeaderProbe::Alive
}
}
}
#[instrument(
skip_all,
fields(
rpc.system = "klieo-a2a",
rpc.method = "SubscribeToTask",
klieo.stream_id = %params.id,
),
err,
)]
async fn dispatch_subscribe_to_task(
&self,
ctx: &RequestContext,
params: SubscribeToTaskParams,
task_store: &crate::task_store::A2aTaskStore,
last_event_id: Option<u64>,
resume_buffer: Arc<dyn klieo_core::resume::ResumeBuffer>,
) -> Result<TaskEventStream, A2aError> {
self.enforce_owner(ctx, ¶ms.id).await?;
if let LeaderProbe::Dead = self.probe_leader(¶ms.id).await {
if let Some(stream) = self
.handle_dead_leader_orphan(ctx, ¶ms.id, task_store, resume_buffer.as_ref())
.await?
{
return Ok(stream);
}
}
if let Some(since) = last_event_id {
match self.handler.subscribe_to_task(ctx, params.clone()).await {
Ok(stream) => return Ok(stream),
Err(A2aError::MethodNotFound(_)) => {}
Err(other) => return Err(other),
}
let task_id_for_tail = params.id.clone();
let tail_stream = self.task_event_stream(task_id_for_tail).await?;
match resume_buffer.replay(¶ms.id, since).await {
Ok(replay_stream) => {
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc as StdArc;
let max_replayed = StdArc::new(AtomicU64::new(since));
let max_for_replay = max_replayed.clone();
let task_id_for_replay = params.id.clone();
let replay = replay_stream.map(move |(id, bytes)| {
max_for_replay.fetch_max(id, Ordering::SeqCst);
serde_json::from_slice::<TaskEvent>(&bytes)
.unwrap_or_else(|_| TaskEvent::synthetic_corrupt(&task_id_for_replay))
});
let max_for_tail = max_replayed.clone();
let tail = tail_stream.filter(move |ev: &TaskEvent| {
ev.event_id > max_for_tail.load(Ordering::SeqCst)
});
return Ok(Box::pin(futures::StreamExt::chain(replay, tail)));
}
Err(klieo_core::resume::ResumeError::Expired { since_id }) => {
return Err(A2aError::ResumeBufferExpired { since_id });
}
Err(klieo_core::resume::ResumeError::NotFound(_)) => { }
Err(klieo_core::resume::ResumeError::Backend(e)) => {
tracing::warn!(
target: "a2a.resume",
error = %e,
"resume buffer backend error; falling back to snapshot+tail"
);
}
Err(_) => {}
}
}
self.subscribe_snapshot_then_tail(ctx, params, task_store)
.await
}
#[instrument(
skip_all,
fields(klieo.stream_id = %params.id),
level = "debug",
err,
)]
async fn subscribe_snapshot_then_tail(
&self,
ctx: &RequestContext,
params: SubscribeToTaskParams,
task_store: &crate::task_store::A2aTaskStore,
) -> Result<TaskEventStream, A2aError> {
match self.handler.subscribe_to_task(ctx, params.clone()).await {
Ok(stream) => return Ok(stream),
Err(A2aError::MethodNotFound(_)) => {}
Err(other) => return Err(other),
}
let task = task_store
.get(¶ms.id)
.await?
.ok_or_else(|| A2aError::TaskNotFound(params.id.clone()))?;
let terminal = task.status.is_terminal();
let event_id = task_store.next_event_id(&task.id);
let initial = TaskEvent::new(
task.id.clone(),
task.status,
task.history.last().cloned(),
terminal,
)
.with_event_id(event_id);
self.task_event_stream_with_initial(task.id, initial, terminal)
.await
}
#[instrument(
skip_all,
fields(klieo.stream_id = %task_id),
level = "debug",
err,
)]
async fn task_event_stream(&self, task_id: String) -> Result<TaskEventStream, A2aError> {
klieo_core::validate_subject_token(&task_id)?;
let subject = format!("klieo.a2a.task.{task_id}");
let durable = DurableName::new(format!("klieo-eph-{}", uuid::Uuid::new_v4()));
let msg_stream = self.pubsub.subscribe(&subject, durable).await?;
let stream = futures::stream::unfold((msg_stream, false), move |(mut ms, done)| {
let task_id = task_id.clone();
async move {
if done {
return None;
}
loop {
match ms.next().await {
None => return None,
Some(Err(e)) => {
warn!(
target: "a2a",
task_id = %task_id,
error = %e,
"task event stream bus error; skipping",
);
continue;
}
Some(Ok(msg)) => {
if let Err(ack_err) = msg.ack.ack().await {
tracing::warn!(target: "a2a", error = %ack_err, "ack failed; message may redeliver");
}
let parent_cx = klieo_core::extract_traceparent(&msg.headers);
let decode_span = tracing::info_span!(
"task_event_decode",
messaging.system = "klieo-bus",
messaging.destination = %format!("klieo.a2a.task.{task_id}"),
messaging.operation = "receive",
klieo.stream_id = %task_id,
);
decode_span.set_parent(parent_cx);
let _enter = decode_span.enter();
match serde_json::from_slice::<TaskEvent>(&msg.payload) {
Err(e) => {
warn!(
target: "a2a",
task_id = %task_id,
error = %e,
"task event decode error; skipping",
);
continue;
}
Ok(ev) => {
let terminal = ev.final_event;
return Some((ev, (ms, terminal)));
}
}
}
}
}
}
});
Ok(Box::pin(stream))
}
async fn task_event_stream_with_initial(
&self,
task_id: String,
initial: TaskEvent,
initial_terminal: bool,
) -> Result<TaskEventStream, A2aError> {
use futures::StreamExt as FutExt;
let initial_stream = futures::stream::once(futures::future::ready(initial));
if initial_terminal {
return Ok(Box::pin(initial_stream));
}
let tail = self.task_event_stream(task_id).await?;
Ok(Box::pin(FutExt::chain(initial_stream, tail)))
}
}
pub struct A2aServer {
app_prefix: String,
agent_id: String,
dispatcher: Arc<A2aDispatcher>,
pubsub: Arc<dyn Pubsub>,
}
impl A2aServer {
pub fn new(
app_prefix: String,
agent_id: String,
handler: Arc<dyn A2aHandler>,
authenticator: Arc<dyn Authenticator>,
pubsub: Arc<dyn Pubsub>,
) -> Self {
let dispatcher = Arc::new(A2aDispatcher::new(handler, authenticator, pubsub.clone()));
Self {
app_prefix,
agent_id,
dispatcher,
pubsub,
}
}
pub fn dispatcher(&self) -> &Arc<A2aDispatcher> {
&self.dispatcher
}
pub fn subject(&self) -> String {
format!("{}.a2a.{}.rpc", self.app_prefix, self.agent_id)
}
pub async fn run(self) -> Result<(), A2aError> {
let subject = self.subject();
let durable = DurableName::new(format!("a2a-server-{}", self.agent_id));
let mut stream = self.pubsub.subscribe(&subject, durable).await?;
while let Some(item) = stream.next().await {
match item {
Ok(msg) => {
self.handle_one(msg).await;
}
Err(e) => {
error!(error = %e, "a2a subscribe stream error; exiting loop");
return Err(A2aError::Bus(e));
}
}
}
Ok(())
}
async fn handle_one(&self, msg: Msg) {
let reply_to = match msg.headers.get("Reply-To") {
Some(s) => s.clone(),
None => {
warn!("a2a request missing Reply-To header; dropping");
if let Err(ack_err) = msg.ack.ack().await {
tracing::warn!(target: "a2a", error = %ack_err, "ack failed; message may redeliver");
}
return;
}
};
let headers = A2aHeaders::decode_from(&msg.headers);
let response = self.dispatcher.handle_request(headers, &msg.payload).await;
let bytes = match serde_json::to_vec(&response) {
Ok(b) => Bytes::from(b),
Err(e) => {
error!(error = %e, "encode response");
if let Err(ack_err) = msg.ack.ack().await {
tracing::warn!(target: "a2a", error = %ack_err, "ack failed; message may redeliver");
}
return;
}
};
if let Err(e) = self.pubsub.publish(&reply_to, bytes, Headers::new()).await {
error!(error = %e, "publish reply");
}
if let Err(ack_err) = msg.ack.ack().await {
tracing::warn!(target: "a2a", error = %ack_err, "ack failed; message may redeliver");
}
}
}
async fn max_event_id(
resume_buffer: &dyn klieo_core::resume::ResumeBuffer,
stream_id: &str,
) -> Option<u64> {
let mut replay = match resume_buffer.replay(stream_id, 0).await {
Ok(stream) => stream,
Err(klieo_core::resume::ResumeError::NotFound(_)) => return None,
Err(e) => {
warn!(
target: "a2a.leader",
stream_id,
error = %e,
"max_event_id replay failed; skipping orphan terminal write",
);
return None;
}
};
let mut highest: Option<u64> = None;
while let Some((id, _)) = tokio_stream::StreamExt::next(&mut replay).await {
highest = Some(match highest {
Some(current) => current.max(id),
None => id,
});
}
highest
}
fn lookup_ok_or_log(
key: &str,
lookup: Result<Option<(klieo_core::LeaderEntry, klieo_core::Revision)>, klieo_core::BusError>,
) -> Option<(klieo_core::LeaderEntry, klieo_core::Revision)> {
match lookup {
Ok(Some(pair)) => Some(pair),
Ok(None) => None,
Err(e) => {
warn!(
target: "a2a.failover",
key,
error = %e,
"leader entry lookup_with_revision failed; falling back to terminate",
);
None
}
}
}
fn failover_reinvoke_sse_frame_bytes(
task_id: &str,
event_id: u64,
attempt: u32,
new_replica_id: &str,
) -> bytes::Bytes {
let payload = serde_json::json!({
"jsonrpc": "2.0",
"id": serde_json::Value::Null,
"event": "failover-reinvoke",
"data": {
"stream_id": format!("{A2A_LEADER_KEY_PREFIX}{task_id}"),
"attempt": attempt,
"by_replica": new_replica_id,
},
"event_id": event_id,
});
bytes::Bytes::from(serde_json::to_vec(&payload).unwrap_or_default())
}
fn leader_died_sse_frame_bytes(task_id: &str, event_id: u64) -> bytes::Bytes {
let payload = serde_json::json!({
"jsonrpc": "2.0",
"id": serde_json::Value::Null,
"error": {
"code": codes::LEADER_DIED,
"message": "stream leader died",
"data": { "stream_id": format!("{A2A_LEADER_KEY_PREFIX}{task_id}") }
},
"event_id": event_id,
});
bytes::Bytes::from(serde_json::to_vec(&payload).unwrap_or_default())
}
pub(crate) fn log_internal_before_wire_seam(err: &A2aError, method: &str) {
let internal = matches!(
err,
A2aError::Bus(_)
| A2aError::Server(_)
| A2aError::Misconfigured(_)
| A2aError::Internal { .. }
);
if internal {
error!(
target: "a2a.wire",
method = %method,
error = %err,
source = ?std::error::Error::source(err),
"internal error mapped to JSON-RPC wire envelope",
);
} else {
warn!(
target: "a2a.wire",
method = %method,
error = %err,
"client-class error mapped to JSON-RPC wire envelope",
);
}
}
async fn run_method<P, T, Fut>(raw: Value, run: impl FnOnce(P) -> Fut) -> Result<Value, A2aError>
where
P: DeserializeOwned,
T: Serialize,
Fut: std::future::Future<Output = Result<T, A2aError>>,
{
let params: P =
serde_json::from_value(raw).map_err(|e| A2aError::InvalidParams(e.to_string()))?;
let typed = run(params).await?;
serde_json::to_value(typed).map_err(A2aError::from)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::envelope::A2aHeaders;
use crate::handler::EchoHandler;
use crate::task_store::A2aTaskStore;
use crate::types::{Task, TaskStatus};
use klieo_auth_common::{AllowAnonymous, Identity};
use klieo_bus_memory::MemoryBus;
fn anon_ctx() -> RequestContext {
RequestContext::new(
A2aHeaders::decode_from(&klieo_core::Headers::new()),
Some(Identity::anonymous()),
)
}
fn make_store_with_dispatcher(dispatcher: &A2aDispatcher) -> A2aTaskStore {
A2aTaskStore::new(
Arc::new(MemoryBus::new()).kv.clone(),
crate::task_store::DEFAULT_BUCKET.into(),
)
.with_event_sink(dispatcher.event_sink())
}
fn make_task(id: &str, status: TaskStatus) -> Task {
Task {
id: id.into(),
contextId: "ctx-1".into(),
status,
artifacts: vec![],
history: vec![],
metadata: None,
}
}
fn noop_resume_buffer() -> Arc<dyn klieo_core::resume::ResumeBuffer> {
Arc::new(klieo_core::resume::NoopResumeBuffer)
}
#[tokio::test]
async fn local_wires_anonymous_auth_and_in_process_pubsub_and_dispatches() {
let dispatcher = A2aDispatcher::local(Arc::new(EchoHandler::default()));
let ctx = anon_ctx();
let payload = br#"{"jsonrpc":"2.0","id":1,"method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;
let resp = dispatcher.dispatch(&ctx, payload).await;
assert_eq!(resp.id, serde_json::json!(1));
assert!(resp.result.is_some());
assert!(resp.error.is_none());
}
#[tokio::test]
async fn dispatcher_routes_send_message_to_handler() {
let dispatcher = A2aDispatcher::with_in_process_pubsub(
Arc::new(EchoHandler::default()),
Arc::new(AllowAnonymous),
);
let ctx = anon_ctx();
let payload = br#"{"jsonrpc":"2.0","id":1,"method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;
let resp = dispatcher.dispatch(&ctx, payload).await;
assert_eq!(resp.id, serde_json::json!(1));
assert!(resp.result.is_some());
assert!(resp.error.is_none());
}
#[tokio::test]
async fn legacy_slash_form_method_name_reaches_the_same_handler_as_camel_case() {
let dispatcher = A2aDispatcher::with_in_process_pubsub(
Arc::new(EchoHandler::default()),
Arc::new(AllowAnonymous),
);
let ctx = anon_ctx();
let legacy_payload = br#"{"jsonrpc":"2.0","id":1,"method":"message/send","params":{"message":{"messageId":"m1","role":"user","parts":[{"text":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;
let resp = dispatcher.dispatch(&ctx, legacy_payload).await;
assert!(
resp.error.is_none(),
"legacy alias must reach send_message, not MethodNotFound: {:?}",
resp.error
);
let result = resp.result.expect("send_message result");
assert!(
result.get("id").is_some() || result.get("messageId").is_some(),
"expected a Task or Message result shape, got {result}"
);
}
#[tokio::test]
async fn authorize_method_sees_canonical_wire_name_for_legacy_alias() {
struct DenyCancelTask;
#[async_trait::async_trait]
impl klieo_auth_common::Authenticator for DenyCancelTask {
async fn authenticate(
&self,
_headers: &dyn klieo_auth_common::Headers,
_payload: &[u8],
) -> Result<Identity, klieo_auth_common::AuthError> {
Ok(Identity::new("scoped-caller"))
}
async fn authorize_method(
&self,
_identity: &Identity,
method: &str,
) -> Result<(), klieo_auth_common::AuthError> {
if method == "CancelTask" {
Err(klieo_auth_common::AuthError::Rejected(
"scope mismatch".into(),
))
} else {
Ok(())
}
}
}
let dispatcher = A2aDispatcher::with_in_process_pubsub(
Arc::new(EchoHandler::default()),
Arc::new(DenyCancelTask),
);
let ctx = named_ctx("scoped-caller");
let payload = br#"{"jsonrpc":"2.0","id":1,"method":"tasks/cancel","params":{"id":"t-1"}}"#;
let resp = dispatcher.dispatch(&ctx, payload).await;
let err = resp
.error
.expect("legacy alias for a denied method must still be denied");
assert_eq!(
err.code,
codes::UNAUTHENTICATED,
"expected the scope-deny to surface as Unauthorized, got: {err:?}"
);
}
#[tokio::test]
async fn create_push_notification_config_rejects_loopback_callback_url() {
let dispatcher = A2aDispatcher::with_in_process_pubsub(
Arc::new(EchoHandler::default()),
Arc::new(AllowAnonymous),
);
let ctx = anon_ctx();
let payload = br#"{"jsonrpc":"2.0","id":1,"method":"CreateTaskPushNotificationConfig","params":{"taskId":"t-1","url":"http://127.0.0.1:8080/hook"}}"#;
let resp = dispatcher.dispatch(&ctx, payload).await;
let err = resp
.error
.expect("loopback callback URL must be rejected before reaching the handler");
assert_eq!(
err.code,
codes::INVALID_PARAMS,
"expected -32602 Invalid params for the SSRF guard, got: {err:?}"
);
assert!(
resp.result.is_none(),
"a rejected config must not return a partial result"
);
}
#[tokio::test]
async fn create_push_notification_config_accepts_public_https_callback_url() {
let dispatcher = A2aDispatcher::with_in_process_pubsub(
Arc::new(EchoHandler::default()),
Arc::new(AllowAnonymous),
);
let ctx = anon_ctx();
let payload = br#"{"jsonrpc":"2.0","id":1,"method":"CreateTaskPushNotificationConfig","params":{"taskId":"t-1","url":"https://example.test/hook"}}"#;
let resp = dispatcher.dispatch(&ctx, payload).await;
let err = resp.error.expect("EchoHandler default is MethodNotFound");
assert_eq!(err.code, codes::METHOD_NOT_FOUND);
}
fn named_ctx(principal: &str) -> RequestContext {
RequestContext::new(
A2aHeaders::decode_from(&klieo_core::Headers::new()),
Some(Identity::new(principal)),
)
}
const SEND_HI: &[u8] = br#"{"jsonrpc":"2.0","id":1,"method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;
#[tokio::test]
async fn non_streaming_send_message_claims_ownership_so_foreign_get_task_is_denied() {
let kv = Arc::new(MemoryBus::new()).kv.clone();
let dispatcher = A2aDispatcher::builder()
.handler(Arc::new(EchoHandler::default()))
.authenticator(Arc::new(AllowAnonymous))
.with_in_process_pubsub()
.with_tenant_binding(kv)
.build()
.unwrap();
let created = dispatcher.dispatch(&named_ctx("alice"), SEND_HI).await;
let task_id = created.result.expect("SendMessage returns a task")["id"]
.as_str()
.expect("task id is a string")
.to_string();
let get = format!(
r#"{{"jsonrpc":"2.0","id":2,"method":"GetTask","params":{{"id":"{task_id}"}}}}"#
);
let foreign = dispatcher.dispatch(&named_ctx("bob"), get.as_bytes()).await;
assert!(
foreign.error.is_some(),
"foreign GetTask must be denied, got result {:?}",
foreign.result
);
let owner = dispatcher
.dispatch(&named_ctx("alice"), get.as_bytes())
.await;
assert!(
owner.error.is_none(),
"owner GetTask must succeed, got error {:?}",
owner.error
);
assert!(owner.result.is_some());
}
#[tokio::test]
async fn dispatcher_streaming_subscribe_to_task_replays_current_state() {
let dispatcher = A2aDispatcher::with_in_process_pubsub(
Arc::new(EchoHandler::default()),
Arc::new(AllowAnonymous),
);
let store = make_store_with_dispatcher(&dispatcher);
let task = make_task("t-1", TaskStatus::Working);
store.put(&task).await.unwrap();
let payload =
br#"{"jsonrpc":"2.0","id":1,"method":"SubscribeToTask","params":{"id":"t-1"}}"#;
let ctx = anon_ctx();
let mut stream = dispatcher
.dispatch_streaming(&ctx, payload, &store, None, noop_resume_buffer())
.await
.unwrap();
let first = stream.next().await.expect("stream produced no first event");
assert_eq!(first.task_id, "t-1");
assert!(matches!(first.status, TaskStatus::Working));
assert!(!first.final_event, "Working is not terminal");
}
#[tokio::test]
async fn dispatcher_streaming_rejects_unknown_task() {
let dispatcher = A2aDispatcher::with_in_process_pubsub(
Arc::new(EchoHandler::default()),
Arc::new(AllowAnonymous),
);
let store = make_store_with_dispatcher(&dispatcher);
let payload =
br#"{"jsonrpc":"2.0","id":1,"method":"SubscribeToTask","params":{"id":"nope"}}"#;
let ctx = anon_ctx();
let result = dispatcher
.dispatch_streaming(&ctx, payload, &store, None, noop_resume_buffer())
.await;
assert!(matches!(result, Err(A2aError::TaskNotFound(_))));
}
struct NoAnonAuthn;
#[async_trait::async_trait]
impl Authenticator for NoAnonAuthn {
async fn authenticate(
&self,
_headers: &dyn klieo_auth_common::Headers,
_payload: &[u8],
) -> Result<Identity, klieo_auth_common::AuthError> {
Ok(Identity::new("svc"))
}
}
#[tokio::test]
async fn dispatcher_streaming_rejects_anonymous_under_named_authenticator() {
let dispatcher = A2aDispatcher::with_in_process_pubsub(
Arc::new(EchoHandler::default()),
Arc::new(NoAnonAuthn),
);
let store = make_store_with_dispatcher(&dispatcher);
let payload = br#"{"jsonrpc":"2.0","id":1,"method":"SendStreamingMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;
let ctx = RequestContext::new(A2aHeaders::decode_from(&klieo_core::Headers::new()), None);
let result = dispatcher
.dispatch_streaming(&ctx, payload, &store, None, noop_resume_buffer())
.await;
assert!(matches!(result, Err(A2aError::Unauthorized(_))));
}
#[tokio::test]
async fn dispatcher_streaming_subscribe_terminal_task_closes_stream() {
let dispatcher = A2aDispatcher::with_in_process_pubsub(
Arc::new(EchoHandler::default()),
Arc::new(AllowAnonymous),
);
let store = make_store_with_dispatcher(&dispatcher);
let task = make_task("t-done", TaskStatus::Completed);
store.put(&task).await.unwrap();
let payload =
br#"{"jsonrpc":"2.0","id":1,"method":"SubscribeToTask","params":{"id":"t-done"}}"#;
let ctx = anon_ctx();
let mut stream = dispatcher
.dispatch_streaming(&ctx, payload, &store, None, noop_resume_buffer())
.await
.unwrap();
let first = stream.next().await.expect("expected replay event");
assert_eq!(first.task_id, "t-done");
assert!(first.final_event, "Completed must set final_event=true");
assert!(stream.next().await.is_none(), "stream must be exhausted");
}
#[tokio::test]
async fn dispatcher_streaming_send_streaming_message_returns_stream() {
let dispatcher = A2aDispatcher::with_in_process_pubsub(
Arc::new(EchoHandler::default()),
Arc::new(AllowAnonymous),
);
let store = make_store_with_dispatcher(&dispatcher);
let payload = br#"{"jsonrpc":"2.0","id":1,"method":"SendStreamingMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;
let ctx = anon_ctx();
let result = dispatcher
.dispatch_streaming(&ctx, payload, &store, None, noop_resume_buffer())
.await;
assert!(
result.is_ok(),
"dispatch_streaming should return Ok for SendStreamingMessage"
);
}
#[tokio::test]
async fn handle_streaming_threads_cancel_token_into_request_context() {
use crate::auth::RequestContext;
use crate::error::A2aError;
use crate::handler::A2aHandler;
use crate::server::TaskEventStream;
use crate::types::SendMessageParams;
use klieo_auth_common::{AllowAnonymous, Authenticator};
use std::sync::{Arc, Mutex};
use tokio_util::sync::CancellationToken;
struct CancelSpy {
observed: Mutex<Option<CancellationToken>>,
}
#[async_trait::async_trait]
impl A2aHandler for CancelSpy {
async fn send_streaming_message(
&self,
ctx: &RequestContext,
_: SendMessageParams,
) -> Result<TaskEventStream, A2aError> {
*self.observed.lock().unwrap() = Some(ctx.cancel.clone());
Ok(Box::pin(futures::stream::empty()))
}
}
let spy = Arc::new(CancelSpy {
observed: Mutex::new(None),
});
let dispatcher = A2aDispatcher::with_in_process_pubsub(
spy.clone() as Arc<dyn A2aHandler>,
Arc::new(AllowAnonymous) as Arc<dyn Authenticator>,
);
let store = A2aTaskStore::new(
Arc::new(klieo_bus_memory::MemoryBus::new()).kv.clone(),
crate::task_store::DEFAULT_BUCKET.into(),
)
.with_event_sink(dispatcher.event_sink());
let token = CancellationToken::new();
let body = serde_json::to_vec(&serde_json::json!({
"jsonrpc": "2.0",
"id": 1,
"method": "SendStreamingMessage",
"params": {
"message": {
"messageId": "m-spy",
"role": "user",
"parts": [],
"extensions": [],
"referenceTaskIds": []
}
}
}))
.unwrap();
let _ = dispatcher
.handle_streaming(
A2aHeaders::decode_from(&klieo_core::Headers::new()),
&body,
&store,
token.clone(),
None,
noop_resume_buffer(),
)
.await
.unwrap();
token.cancel();
let observed = spy.observed.lock().unwrap().clone().unwrap();
assert!(
observed.is_cancelled(),
"cancel must propagate from handle_streaming arg"
);
}
#[tokio::test]
async fn task_event_sink_publishes_to_per_task_subject() {
let bus = klieo_bus_memory::MemoryBus::new();
let pubsub: std::sync::Arc<dyn klieo_core::Pubsub> = bus.pubsub.clone();
let sink = TaskEventSink::new(pubsub.clone());
let durable = klieo_core::DurableName::new("test-eph");
let mut stream = pubsub
.subscribe("klieo.a2a.task.t-1", durable)
.await
.unwrap();
let event = TaskEvent::new(
"t-1".to_string(),
crate::types::TaskStatus::Submitted,
None,
false,
)
.with_event_id(1);
sink.send(event.clone()).await.unwrap();
use tokio_stream::StreamExt as _;
let msg = tokio::time::timeout(std::time::Duration::from_millis(500), stream.next())
.await
.expect("timeout")
.expect("stream ended")
.expect("subscribe err");
let decoded: TaskEvent = serde_json::from_slice(&msg.payload).unwrap();
assert_eq!(decoded.task_id, "t-1");
assert_eq!(decoded.event_id, 1);
msg.ack.ack().await.unwrap();
}
}
#[cfg(test)]
mod profile_tests {
use super::*;
use crate::handler::EchoHandler;
use klieo_auth_common::{AllowAnonymous, AuthError, Headers, Identity};
use klieo_core::DeploymentProfile;
struct NamedAuthn;
#[async_trait::async_trait]
impl Authenticator for NamedAuthn {
async fn authenticate(
&self,
_headers: &dyn Headers,
_payload: &[u8],
) -> Result<Identity, AuthError> {
Ok(Identity::new("alice"))
}
}
fn builder_with(auth: Arc<dyn Authenticator>) -> A2aDispatcherBuilder {
let bus = klieo_bus_memory::MemoryBus::new();
A2aDispatcher::builder()
.handler(Arc::new(EchoHandler::default()))
.authenticator(auth)
.pubsub(bus.pubsub.clone())
}
#[test]
fn regulated_without_tenant_kv_fails_closed() {
let err = builder_with(Arc::new(NamedAuthn))
.profile(DeploymentProfile::RegulatedMultiTenant)
.build()
.err()
.expect("expected RegulatedProfile error");
assert!(matches!(
err,
A2aBuilderError::RegulatedProfile(klieo_core::ProfileViolation::MissingTenantKv)
));
}
#[test]
fn regulated_with_anonymous_auth_fails_closed() {
let bus = klieo_bus_memory::MemoryBus::new();
let err = builder_with(Arc::new(AllowAnonymous))
.with_tenant_binding(bus.kv.clone())
.profile(DeploymentProfile::RegulatedMultiTenant)
.build()
.err()
.expect("expected RegulatedProfile error");
assert!(matches!(
err,
A2aBuilderError::RegulatedProfile(klieo_core::ProfileViolation::AnonymousAuth)
));
}
#[test]
fn regulated_forces_strict_over_lenient_binding() {
let bus = klieo_bus_memory::MemoryBus::new();
let dispatcher = builder_with(Arc::new(NamedAuthn))
.with_tenant_binding(bus.kv.clone())
.profile(DeploymentProfile::RegulatedMultiTenant)
.build()
.expect("regulated build with named auth + kv must succeed");
assert_eq!(
dispatcher
.ownership_registry
.as_ref()
.map(|r| r.is_strict()),
Some(true),
"regulated profile must force a strict registry even over lenient binding"
);
}
#[test]
fn unprofiled_keeps_lenient_binding() {
let bus = klieo_bus_memory::MemoryBus::new();
let dispatcher = builder_with(Arc::new(NamedAuthn))
.with_tenant_binding(bus.kv.clone())
.build()
.expect("unprofiled build ok");
assert_eq!(
dispatcher
.ownership_registry
.as_ref()
.map(|r| r.is_strict()),
Some(false)
);
}
fn tenant_bound_dispatcher() -> A2aDispatcher {
let bus = klieo_bus_memory::MemoryBus::new();
A2aDispatcher::builder()
.handler(Arc::new(EchoHandler::default()))
.authenticator(Arc::new(AllowAnonymous))
.pubsub(bus.pubsub.clone())
.with_tenant_binding(bus.kv.clone())
.build()
.expect("tenant-bound dispatcher builds")
}
fn named_ctx(principal: &str) -> RequestContext {
RequestContext::new(
A2aHeaders::decode_from(&klieo_core::Headers::new()),
Some(Identity::new(principal)),
)
}
const SEND_MESSAGE_PAYLOAD: &[u8] = br#"{"jsonrpc":"2.0","id":1,"method":"SendMessage","params":{"message":{"messageId":"m1","role":"user","parts":[{"type":"text","content":"hi"}],"extensions":[],"referenceTaskIds":[]}}}"#;
#[tokio::test]
async fn anonymous_caller_rejected_when_authenticator_requires_identity() {
let dispatcher = builder_with(Arc::new(NamedAuthn)).build().expect("build");
let ctx = RequestContext::new(A2aHeaders::decode_from(&klieo_core::Headers::new()), None);
let resp = dispatcher.dispatch(&ctx, SEND_MESSAGE_PAYLOAD).await;
let err = resp.error.expect("anonymous caller must be rejected");
assert_eq!(err.code, codes::UNAUTHENTICATED);
assert!(
resp.result.is_none(),
"rejected request must carry no result"
);
}
fn hi_params() -> crate::types::SendMessageParams {
serde_json::from_value(serde_json::json!({
"message": {
"messageId": "m1",
"role": "user",
"parts": [{"type": "text", "content": "hi"}],
"extensions": [],
"referenceTaskIds": []
}
}))
.expect("valid SendMessageParams")
}
async fn create_task(dispatcher: &A2aDispatcher, ctx: &RequestContext) -> String {
let resp = dispatcher.dispatch(ctx, SEND_MESSAGE_PAYLOAD).await;
resp.result
.expect("SendMessage result")
.get("id")
.and_then(|v| v.as_str())
.expect("created task carries an id")
.to_string()
}
#[tokio::test]
async fn anonymous_send_message_records_no_ownership() {
let dispatcher = tenant_bound_dispatcher();
let anon = RequestContext::new(
A2aHeaders::decode_from(&klieo_core::Headers::new()),
Some(Identity::anonymous()),
);
let created = dispatcher.dispatch(&anon, SEND_MESSAGE_PAYLOAD).await;
let task_id = created.result.expect("SendMessage returns a task")["id"]
.as_str()
.expect("task id is a string")
.to_string();
let owner = dispatcher
.ownership_registry()
.expect("registry wired")
.lookup(&format!("{A2A_OWNERSHIP_KEY_PREFIX}{task_id}"))
.await
.expect("lookup ok");
assert!(
owner.is_none(),
"anonymous SendMessage must not write an ownership entry, got {owner:?}"
);
}
#[must_use]
async fn claim_for(
dispatcher: &A2aDispatcher,
task_id: &str,
principal: &str,
) -> klieo_core::OwnershipHandle {
dispatcher
.ownership_registry()
.expect("ownership registry wired")
.claim(
format!("{A2A_OWNERSHIP_KEY_PREFIX}{task_id}"),
principal.into(),
)
.await
.expect("claim ownership")
}
fn by_id_request(method: &str, task_id: &str) -> Vec<u8> {
serde_json::to_vec(&serde_json::json!({
"jsonrpc": "2.0",
"id": 2,
"method": method,
"params": { "id": task_id },
}))
.expect("encode by-id request")
}
#[tokio::test]
async fn get_task_owner_reads_but_foreign_principal_denied_as_not_found() {
let dispatcher = tenant_bound_dispatcher();
let alice = named_ctx("alice");
let task_id = create_task(&dispatcher, &alice).await;
let _ownership = claim_for(&dispatcher, &task_id, "alice").await;
let request = by_id_request("GetTask", &task_id);
let owner = dispatcher.dispatch(&alice, &request).await;
assert!(owner.error.is_none(), "owner must read own task: {owner:?}");
assert_eq!(
owner
.result
.expect("owner result")
.get("id")
.and_then(|v| v.as_str()),
Some(task_id.as_str()),
);
let bob = named_ctx("bob");
let foreign = dispatcher.dispatch(&bob, &request).await;
let err = foreign.error.expect("foreign principal must be denied");
assert_eq!(err.code, codes::SERVER_ERROR);
assert_ne!(
err.code,
codes::UNAUTHENTICATED,
"deny must not surface -32001 (would leak existence info)",
);
assert!(
err.message.contains("task not found"),
"deny-as-not-found expected, got: {}",
err.message,
);
assert!(foreign.result.is_none(), "must not leak the task body");
}
#[tokio::test]
async fn cancel_task_foreign_principal_denied_before_mutation() {
let dispatcher = tenant_bound_dispatcher();
let alice = named_ctx("alice");
let task_id = create_task(&dispatcher, &alice).await;
let _ownership = claim_for(&dispatcher, &task_id, "alice").await;
let bob = named_ctx("bob");
let cancel = dispatcher
.dispatch(&bob, &by_id_request("CancelTask", &task_id))
.await;
let err = cancel.error.expect("foreign cancel must be denied");
assert_eq!(err.code, codes::SERVER_ERROR);
assert!(err.message.contains("task not found"));
let owner_view = dispatcher
.dispatch(&alice, &by_id_request("GetTask", &task_id))
.await;
let status = owner_view
.result
.expect("owner result")
.get("status")
.and_then(|v| v.as_str())
.map(str::to_string);
assert_ne!(
status.as_deref(),
Some("canceled"),
"foreign cancel must not mutate the task",
);
}
struct UnavailableKv;
#[async_trait::async_trait]
impl klieo_core::KvStore for UnavailableKv {
async fn get(
&self,
_: &str,
_: &str,
) -> Result<Option<klieo_core::KvEntry>, klieo_core::BusError> {
Err(klieo_core::BusError::Connection("kv down".into()))
}
async fn put(
&self,
_: &str,
_: &str,
_: Bytes,
) -> Result<klieo_core::Revision, klieo_core::BusError> {
Err(klieo_core::BusError::Connection("kv down".into()))
}
async fn cas(
&self,
_: &str,
_: &str,
_: Bytes,
_: Option<klieo_core::Revision>,
) -> Result<klieo_core::Revision, klieo_core::BusError> {
Err(klieo_core::BusError::Connection("kv down".into()))
}
async fn delete(&self, _: &str, _: &str) -> Result<(), klieo_core::BusError> {
Err(klieo_core::BusError::Connection("kv down".into()))
}
async fn lease(
&self,
_: &str,
_: &str,
_: std::time::Duration,
) -> Result<klieo_core::Lease, klieo_core::BusError> {
Err(klieo_core::BusError::Connection("kv down".into()))
}
async fn keys(&self, _: &str) -> Result<Vec<String>, klieo_core::BusError> {
Err(klieo_core::BusError::Connection("kv down".into()))
}
}
#[tokio::test]
async fn send_message_fails_closed_when_strict_ownership_store_unavailable() {
let bus = klieo_bus_memory::MemoryBus::new();
let dispatcher = A2aDispatcher::builder()
.handler(Arc::new(EchoHandler::default()))
.authenticator(Arc::new(AllowAnonymous))
.pubsub(bus.pubsub.clone())
.with_tenant_binding_strict(Arc::new(UnavailableKv))
.build()
.expect("strict tenant-bound dispatcher builds");
let resp = dispatcher
.dispatch(&named_ctx("alice"), SEND_MESSAGE_PAYLOAD)
.await;
let err = resp
.error
.expect("strict store-down must fail the create closed");
assert_eq!(err.code, codes::SERVER_ERROR);
assert!(resp.result.is_none(), "denied create must carry no task");
}
#[tokio::test]
async fn list_tasks_fails_closed_when_ownership_store_unavailable() {
let handler = Arc::new(EchoHandler::default());
handler
.send_message(&named_ctx("alice"), hi_params())
.await
.expect("seed task on handler");
let bus = klieo_bus_memory::MemoryBus::new();
let dispatcher = A2aDispatcher::builder()
.handler(handler)
.authenticator(Arc::new(AllowAnonymous))
.pubsub(bus.pubsub.clone())
.with_tenant_binding_strict(Arc::new(UnavailableKv))
.build()
.expect("strict tenant-bound dispatcher builds");
let alice = named_ctx("alice");
let list_request = serde_json::to_vec(&serde_json::json!({
"jsonrpc": "2.0", "id": 4, "method": "ListTasks", "params": {},
}))
.expect("encode ListTasks");
let resp = dispatcher.dispatch(&alice, &list_request).await;
let err = resp
.error
.expect("a strict store-down must fail the list closed");
assert_eq!(err.code, codes::SERVER_ERROR);
assert!(
err.message.contains("unavailable") || err.message.contains("list denied"),
"fail-closed message expected, got: {}",
err.message,
);
assert!(resp.result.is_none(), "must not return a partial list");
}
#[tokio::test]
async fn push_notification_config_arms_deny_foreign_principal() {
let dispatcher = tenant_bound_dispatcher();
let alice = named_ctx("alice");
let task_id = create_task(&dispatcher, &alice).await;
let _ownership = claim_for(&dispatcher, &task_id, "alice").await;
let bob = named_ctx("bob");
let requests = [
serde_json::json!({"jsonrpc":"2.0","id":5,"method":"CreateTaskPushNotificationConfig","params":{"taskId":task_id,"url":"https://example.test/hook"}}),
serde_json::json!({"jsonrpc":"2.0","id":6,"method":"GetTaskPushNotificationConfig","params":{"taskId":task_id,"id":"c1"}}),
serde_json::json!({"jsonrpc":"2.0","id":7,"method":"ListTaskPushNotificationConfigs","params":{"taskId":task_id}}),
serde_json::json!({"jsonrpc":"2.0","id":8,"method":"DeleteTaskPushNotificationConfig","params":{"taskId":task_id,"id":"c1"}}),
];
for request in requests {
let method = request["method"].as_str().expect("method").to_string();
let body = serde_json::to_vec(&request).expect("encode push-config request");
let foreign = dispatcher.dispatch(&bob, &body).await;
let foreign_err = foreign
.error
.unwrap_or_else(|| panic!("{method}: foreign principal must be denied"));
assert_eq!(
foreign_err.code,
codes::SERVER_ERROR,
"{method}: foreign principal must be denied-as-not-found",
);
assert!(
foreign_err.message.contains("task not found"),
"{method}: expected deny-as-not-found, got {}",
foreign_err.message,
);
assert!(foreign.result.is_none(), "{method}: must not leak a body");
let owner = dispatcher.dispatch(&alice, &body).await;
let owner_err = owner
.error
.unwrap_or_else(|| panic!("{method}: EchoHandler declines push-config"));
assert_eq!(
owner_err.code,
codes::METHOD_NOT_FOUND,
"{method}: owner must pass the gate and reach the handler default",
);
}
}
#[tokio::test]
async fn list_tasks_drops_tasks_the_caller_does_not_own() {
let dispatcher = tenant_bound_dispatcher();
let alice = named_ctx("alice");
let bob = named_ctx("bob");
let alice_task = create_task(&dispatcher, &alice).await;
let bob_task = create_task(&dispatcher, &bob).await;
let _alice_owns = claim_for(&dispatcher, &alice_task, "alice").await;
let _bob_owns = claim_for(&dispatcher, &bob_task, "bob").await;
let list_request = serde_json::to_vec(&serde_json::json!({
"jsonrpc": "2.0", "id": 3, "method": "ListTasks", "params": {},
}))
.expect("encode ListTasks");
let resp = dispatcher.dispatch(&bob, &list_request).await;
let ids: Vec<String> = resp
.result
.expect("list result")
.get("tasks")
.and_then(|v| v.as_array())
.expect("tasks array")
.iter()
.filter_map(|t| t.get("id").and_then(|v| v.as_str()).map(str::to_string))
.collect();
assert!(ids.contains(&bob_task), "owner must see own task: {ids:?}");
assert!(
!ids.contains(&alice_task),
"list must not leak a foreign-owned task: {ids:?}",
);
}
}