use a2a::*;
use async_trait::async_trait;
use futures::{StreamExt, stream::BoxStream};
use std::{collections::HashMap, sync::Arc};
use tokio::sync::{RwLock, broadcast};
use crate::middleware::ServiceParams;
use crate::pagination::resolve_page_size;
use crate::task_store::apply_history_length;
const EXECUTION_BUFFER_CAPACITY: usize = 32;
#[derive(Clone)]
struct ExecutionEvent {
sequence: u64,
result: Result<StreamResponse, A2AError>,
}
#[derive(Default)]
struct ActiveExecutionState {
sequence: u64,
snapshot_task: Option<Task>,
}
struct ActiveExecution {
sender: broadcast::Sender<ExecutionEvent>,
state: RwLock<ActiveExecutionState>,
}
impl ActiveExecution {
fn new(task: Task) -> Arc<Self> {
let (sender, _) = broadcast::channel(EXECUTION_BUFFER_CAPACITY);
Arc::new(Self {
sender,
state: RwLock::new(ActiveExecutionState {
sequence: 0,
snapshot_task: Some(task),
}),
})
}
fn subscribe(&self) -> broadcast::Receiver<ExecutionEvent> {
self.sender.subscribe()
}
async fn resubscribe(&self) -> (broadcast::Receiver<ExecutionEvent>, Option<Task>, u64) {
let receiver = self.sender.subscribe();
let state = self.state.read().await;
(receiver, state.snapshot_task.clone(), state.sequence)
}
async fn publish(
&self,
result: Result<&StreamResponse, &A2AError>,
snapshot_task: Option<Task>,
) {
let sequence = {
let mut state = self.state.write().await;
if let Some(task) = snapshot_task {
state.snapshot_task = Some(task);
}
state.sequence += 1;
state.sequence
};
if self.sender.receiver_count() == 0 {
return;
}
let event = ExecutionEvent {
sequence,
result: result.cloned().map_err(Clone::clone),
};
let _ = self.sender.send(event);
}
async fn snapshot_is_terminal(&self) -> bool {
let state = self.state.read().await;
state
.snapshot_task
.as_ref()
.is_some_and(|task| task.status.state.is_terminal())
}
}
#[derive(Default)]
struct ExecutionManager {
executions: RwLock<HashMap<TaskId, Arc<ActiveExecution>>>,
}
impl ExecutionManager {
async fn start(&self, task: Task) -> Result<Arc<ActiveExecution>, A2AError> {
let mut executions = self.executions.write().await;
if executions.contains_key(&task.id) {
return Err(A2AError::invalid_request(format!(
"task execution is already in progress: {}",
task.id
)));
}
let active = ActiveExecution::new(task.clone());
executions.insert(task.id.clone(), active.clone());
Ok(active)
}
async fn get(&self, task_id: &str) -> Option<Arc<ActiveExecution>> {
let executions = self.executions.read().await;
executions.get(task_id).cloned()
}
async fn resubscribe(
&self,
task_id: &str,
) -> Option<(broadcast::Receiver<ExecutionEvent>, Option<Task>, u64)> {
let active = self.get(task_id).await?;
let (receiver, snapshot_task, sequence) = active.resubscribe().await;
Some((receiver, snapshot_task, sequence))
}
async fn finish(&self, task_id: &str, active: &Arc<ActiveExecution>) {
let mut executions = self.executions.write().await;
let should_remove = executions
.get(task_id)
.is_some_and(|current| Arc::ptr_eq(current, active));
if should_remove {
executions.remove(task_id);
}
}
}
struct ExecutionRuntime {
executor: Arc<dyn crate::AgentExecutor>,
task_store: Arc<dyn crate::TaskStore>,
push_config_store: Option<Arc<dyn crate::PushConfigStore>>,
push_sender: Option<Arc<crate::HttpPushSender>>,
execution_manager: Arc<ExecutionManager>,
}
struct SubscriptionState {
receiver: broadcast::Receiver<ExecutionEvent>,
snapshot_task: Option<Task>,
min_sequence: u64,
done: bool,
}
fn is_terminal_event(event: &StreamResponse) -> bool {
match event {
StreamResponse::Task(task) => task.status.state.is_terminal(),
StreamResponse::StatusUpdate(update) => update.status.state.is_terminal(),
_ => false,
}
}
fn task_id_for_event(event: &StreamResponse) -> Option<&str> {
match event {
StreamResponse::Task(task) => Some(&task.id),
StreamResponse::StatusUpdate(update) => Some(&update.task_id),
StreamResponse::ArtifactUpdate(update) => Some(&update.task_id),
StreamResponse::Message(_) => None,
}
}
fn should_interrupt_non_streaming(
req: &SendMessageRequest,
event: &StreamResponse,
) -> Option<TaskId> {
if req
.configuration
.as_ref()
.and_then(|config| config.return_immediately)
.unwrap_or(false)
{
return match event {
StreamResponse::Message(_) => None,
_ => task_id_for_event(event).map(ToOwned::to_owned),
};
}
match event {
StreamResponse::Task(task) if task.status.state == TaskState::AuthRequired => {
Some(task.id.clone())
}
StreamResponse::StatusUpdate(update) if update.status.state == TaskState::AuthRequired => {
Some(update.task_id.clone())
}
_ => None,
}
}
async fn send_push_notifications(
task_id: &str,
push_config_store: Option<&dyn crate::PushConfigStore>,
push_sender: Option<&crate::HttpPushSender>,
event: &StreamResponse,
) -> Result<(), A2AError> {
let Some(push_config_store) = push_config_store else {
return Ok(());
};
let Some(push_sender) = push_sender else {
return Ok(());
};
let configs = push_config_store.list(task_id).await?;
for config in configs {
push_sender.send_push(&config, event.clone()).await?;
}
Ok(())
}
async fn save_task(task_store: &dyn crate::TaskStore, mut task: Task) -> Result<Task, A2AError> {
if task.status.timestamp.is_none() {
task.status.timestamp = Some(chrono::Utc::now());
}
match task_store.update(task.clone()).await {
Ok(_) => Ok(task),
Err(error) if error.code == error_code::TASK_NOT_FOUND => {
task_store.create(task.clone()).await?;
Ok(task)
}
Err(error) => Err(error),
}
}
fn backfill_status_timestamp(event: &mut StreamResponse) {
let status = match event {
StreamResponse::Task(task) => &mut task.status,
StreamResponse::StatusUpdate(update) => &mut update.status,
StreamResponse::ArtifactUpdate(_) | StreamResponse::Message(_) => return,
};
if status.timestamp.is_none() {
status.timestamp = Some(chrono::Utc::now());
}
}
async fn apply_event_to_task(
task_store: &dyn crate::TaskStore,
current_task: Option<Task>,
event: &StreamResponse,
) -> Result<Option<Task>, A2AError> {
match event {
StreamResponse::Task(task) => save_task(task_store, task.clone()).await.map(Some),
StreamResponse::StatusUpdate(update) => {
let mut task = match current_task {
Some(task) => task,
None => task_store
.get(&update.task_id)
.await?
.ok_or_else(|| A2AError::task_not_found(&update.task_id))?,
};
task.status = update.status.clone();
save_task(task_store, task).await.map(Some)
}
StreamResponse::ArtifactUpdate(update) => {
let mut task = match current_task {
Some(task) => task,
None => task_store
.get(&update.task_id)
.await?
.ok_or_else(|| A2AError::task_not_found(&update.task_id))?,
};
let artifacts = task.artifacts.get_or_insert_with(Vec::new);
let existing = artifacts
.iter_mut()
.find(|artifact| artifact.artifact_id == update.artifact.artifact_id);
if update.append.unwrap_or(false) {
let artifact = existing.ok_or_else(A2AError::invalid_agent_response)?;
artifact.parts.extend(update.artifact.parts.iter().cloned());
if let Some(metadata) = &update.artifact.metadata {
artifact
.metadata
.get_or_insert_with(Default::default)
.extend(metadata.clone());
}
} else if let Some(artifact) = existing {
*artifact = update.artifact.clone();
} else {
artifacts.push(update.artifact.clone());
}
save_task(task_store, task).await.map(Some)
}
StreamResponse::Message(_) => Ok(current_task),
}
}
fn subscription_stream(
receiver: broadcast::Receiver<ExecutionEvent>,
snapshot_task: Option<Task>,
min_sequence: u64,
) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
Box::pin(futures::stream::unfold(
SubscriptionState {
receiver,
snapshot_task,
min_sequence,
done: false,
},
|mut state| async move {
if state.done {
return None;
}
if let Some(task) = state.snapshot_task.take() {
if task.status.state.is_terminal() {
state.done = true;
}
return Some((Ok(StreamResponse::Task(task)), state));
}
loop {
match state.receiver.recv().await {
Ok(event) if event.sequence <= state.min_sequence => continue,
Ok(event) => {
state.min_sequence = event.sequence;
match &event.result {
Ok(item) if is_terminal_event(item) => state.done = true,
Err(_) => state.done = true,
_ => {}
}
return Some((event.result.clone(), state));
}
Err(broadcast::error::RecvError::Lagged(_)) => {
state.done = true;
return Some((
Err(crate::sanitized_internal_error(
"subscription fell behind active execution",
)),
state,
));
}
Err(broadcast::error::RecvError::Closed) => return None,
}
}
},
))
}
async fn drive_execution(
runtime: ExecutionRuntime,
task_id: TaskId,
active: Arc<ActiveExecution>,
exec_ctx: crate::ExecutorContext,
) {
let mut current_task = {
let state = active.state.read().await;
state.snapshot_task.clone()
};
let mut stream = runtime.executor.execute(exec_ctx);
while let Some(result) = stream.next().await {
if active.snapshot_is_terminal().await {
break;
}
match result {
Ok(mut event) => {
backfill_status_timestamp(&mut event);
match apply_event_to_task(runtime.task_store.as_ref(), current_task.clone(), &event)
.await
{
Ok(updated_task) => {
current_task = updated_task.clone();
if let Err(error) = send_push_notifications(
&task_id,
runtime.push_config_store.as_deref(),
runtime.push_sender.as_deref(),
&event,
)
.await
{
active.publish(Err(&error), None).await;
break;
}
active.publish(Ok(&event), updated_task).await;
if is_terminal_event(&event) {
break;
}
}
Err(error) => {
active.publish(Err(&error), None).await;
break;
}
}
}
Err(error) => {
active.publish(Err(&error), None).await;
break;
}
}
}
runtime.execution_manager.finish(&task_id, &active).await;
}
#[async_trait]
pub trait RequestHandler: Send + Sync + 'static {
async fn send_message(
&self,
params: &ServiceParams,
req: SendMessageRequest,
) -> Result<SendMessageResponse, A2AError>;
async fn send_streaming_message(
&self,
params: &ServiceParams,
req: SendMessageRequest,
) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError>;
async fn get_task(&self, params: &ServiceParams, req: GetTaskRequest)
-> Result<Task, A2AError>;
async fn list_tasks(
&self,
params: &ServiceParams,
req: ListTasksRequest,
) -> Result<ListTasksResponse, A2AError>;
async fn cancel_task(
&self,
params: &ServiceParams,
req: CancelTaskRequest,
) -> Result<Task, A2AError>;
async fn subscribe_to_task(
&self,
params: &ServiceParams,
req: SubscribeToTaskRequest,
) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError>;
async fn create_push_config(
&self,
params: &ServiceParams,
req: TaskPushNotificationConfig,
) -> Result<TaskPushNotificationConfig, A2AError>;
async fn get_push_config(
&self,
params: &ServiceParams,
req: GetTaskPushNotificationConfigRequest,
) -> Result<TaskPushNotificationConfig, A2AError>;
async fn list_push_configs(
&self,
params: &ServiceParams,
req: ListTaskPushNotificationConfigsRequest,
) -> Result<ListTaskPushNotificationConfigsResponse, A2AError>;
async fn delete_push_config(
&self,
params: &ServiceParams,
req: DeleteTaskPushNotificationConfigRequest,
) -> Result<(), A2AError>;
async fn get_extended_agent_card(
&self,
params: &ServiceParams,
req: GetExtendedAgentCardRequest,
) -> Result<AgentCard, A2AError>;
}
pub trait RequestAuthorizer: Send + Sync {
fn authorize(&self, params: &ServiceParams, task_id: Option<&str>) -> Result<(), A2AError>;
}
impl<F> RequestAuthorizer for F
where
F: Fn(&ServiceParams, Option<&str>) -> Result<(), A2AError> + Send + Sync,
{
fn authorize(&self, params: &ServiceParams, task_id: Option<&str>) -> Result<(), A2AError> {
self(params, task_id)
}
}
pub struct DefaultRequestHandler {
executor: Arc<dyn crate::AgentExecutor>,
task_store: Arc<dyn crate::TaskStore>,
execution_manager: Arc<ExecutionManager>,
push_config_store: Option<Arc<dyn crate::PushConfigStore>>,
push_sender: Option<Arc<crate::HttpPushSender>>,
capabilities: AgentCapabilities,
default_input_modes: Vec<String>,
authorizer: Option<Arc<dyn RequestAuthorizer>>,
extended_agent_card: Option<Arc<dyn ExtendedAgentCardResolver>>,
}
pub trait ExtendedAgentCardResolver: Send + Sync {
fn resolve(&self, params: &ServiceParams) -> Result<AgentCard, A2AError>;
}
impl<F> ExtendedAgentCardResolver for F
where
F: Fn(&ServiceParams) -> Result<AgentCard, A2AError> + Send + Sync,
{
fn resolve(&self, params: &ServiceParams) -> Result<AgentCard, A2AError> {
self(params)
}
}
impl DefaultRequestHandler {
pub fn new(executor: impl crate::AgentExecutor, task_store: impl crate::TaskStore) -> Self {
DefaultRequestHandler {
executor: Arc::new(executor),
task_store: Arc::new(task_store),
execution_manager: Arc::new(ExecutionManager::default()),
push_config_store: None,
push_sender: None,
capabilities: AgentCapabilities::default(),
default_input_modes: Vec::new(),
authorizer: None,
extended_agent_card: None,
}
}
pub fn with_extended_agent_card(mut self, card: AgentCard) -> Self {
self.capabilities.extended_agent_card.get_or_insert(true);
self.with_extended_agent_card_resolver(move |_: &ServiceParams| Ok(card.clone()))
}
pub fn with_extended_agent_card_resolver(
mut self,
resolver: impl ExtendedAgentCardResolver + 'static,
) -> Self {
self.capabilities.extended_agent_card.get_or_insert(true);
self.extended_agent_card = Some(Arc::new(resolver));
self
}
pub fn with_authorizer(mut self, authorizer: impl RequestAuthorizer + 'static) -> Self {
self.authorizer = Some(Arc::new(authorizer));
self
}
fn authorize(&self, params: &ServiceParams, task_id: Option<&str>) -> Result<(), A2AError> {
match &self.authorizer {
Some(authorizer) => authorizer.authorize(params, task_id),
None => Ok(()),
}
}
pub fn with_push_config_store(
mut self,
push_config_store: impl crate::PushConfigStore,
) -> Self {
self.push_sender = Some(Arc::new(crate::HttpPushSender::new(None)));
self.push_config_store = Some(Arc::new(push_config_store));
self.capabilities.push_notifications.get_or_insert(true);
self
}
pub fn with_push_notifications(
mut self,
push_config_store: impl crate::PushConfigStore,
push_sender: crate::HttpPushSender,
) -> Self {
self.push_config_store = Some(Arc::new(push_config_store));
self.push_sender = Some(Arc::new(push_sender));
self.capabilities.push_notifications.get_or_insert(true);
self
}
pub fn with_capabilities(mut self, capabilities: AgentCapabilities) -> Self {
self.capabilities = capabilities;
self
}
pub fn with_default_input_modes(mut self, modes: impl Into<Vec<String>>) -> Self {
self.default_input_modes = modes.into();
self
}
fn validate_content_types(&self, message: &Message) -> Result<(), A2AError> {
if self.default_input_modes.is_empty() {
return Ok(());
}
let supported = message.parts.iter().all(|part| {
part.media_type.as_deref().is_none_or(|media_type| {
media_type.is_empty() || self.default_input_modes.iter().any(|m| m == media_type)
})
});
if supported {
Ok(())
} else {
Err(A2AError::content_type_not_supported())
}
}
fn require_capability(
declared: Option<bool>,
unsupported: impl FnOnce() -> A2AError,
) -> Result<(), A2AError> {
match declared {
Some(false) => Err(unsupported()),
_ => Ok(()),
}
}
fn require_streaming(&self) -> Result<(), A2AError> {
Self::require_capability(self.capabilities.streaming, || {
A2AError::unsupported_operation("the agent card does not declare streaming")
})
}
fn require_push_notifications(&self) -> Result<(), A2AError> {
Self::require_capability(
self.capabilities.push_notifications,
A2AError::push_notification_not_supported,
)
}
fn push_config_store_checked(&self) -> Result<&dyn crate::PushConfigStore, A2AError> {
self.require_push_notifications()?;
self.push_config_store()
}
fn push_config_store(&self) -> Result<&dyn crate::PushConfigStore, A2AError> {
self.push_config_store
.as_deref()
.ok_or_else(A2AError::push_notification_not_supported)
}
async fn load_task(&self, task_id: &str) -> Result<Task, A2AError> {
self.task_store
.get(task_id)
.await?
.ok_or_else(|| A2AError::task_not_found(task_id))
}
async fn prepare_task_for_execution(
&self,
req: &SendMessageRequest,
) -> Result<(Task, Option<Task>, String), A2AError> {
let task_id = match req.message.task_id.clone() {
Some(id) => id,
None => new_task_id(),
};
let stored = self.task_store.get(&task_id).await?;
if stored.is_none() && req.message.task_id.is_some() {
return Err(A2AError::task_not_found(&task_id));
}
if stored
.as_ref()
.is_some_and(|task| task.status.state.is_terminal())
{
return Err(A2AError::unsupported_operation(format!(
"task {task_id} has already reached a terminal state"
)));
}
let context_id = match stored.as_ref() {
Some(task) => {
if let Some(requested) = req.message.context_id.as_deref() {
if requested != task.context_id {
return Err(A2AError::invalid_params(format!(
"contextId {requested} does not match task {task_id}, \
which belongs to context {}",
task.context_id
)));
}
}
task.context_id.clone()
}
None => req
.message
.context_id
.clone()
.unwrap_or_else(new_context_id),
};
let task = if let Some(mut existing) = stored.clone() {
existing
.history
.get_or_insert_with(Vec::new)
.push(req.message.clone());
existing
} else {
let task = Task {
id: task_id,
context_id: context_id.clone(),
status: TaskStatus {
state: TaskState::Submitted,
message: None,
timestamp: Some(chrono::Utc::now()),
},
artifacts: None,
history: Some(vec![req.message.clone()]),
metadata: None,
};
self.task_store.create(task.clone()).await?;
task
};
Ok((task, stored, context_id))
}
async fn save_request_push_config(
&self,
task_id: &str,
req: &SendMessageRequest,
) -> Result<(), A2AError> {
let Some(config) = req
.configuration
.as_ref()
.and_then(|configuration| configuration.task_push_notification_config.clone())
else {
return Ok(());
};
let mut config = config;
config.task_id = task_id.to_string();
if config.tenant.is_none() {
config.tenant = req.tenant.clone();
}
self.push_config_store_checked()?.save(config).await?;
Ok(())
}
async fn start_execution(
&self,
params: &ServiceParams,
req: SendMessageRequest,
include_task_snapshot_in_context: bool,
) -> Result<(TaskId, BoxStream<'static, Result<StreamResponse, A2AError>>), A2AError> {
self.validate_content_types(&req.message)?;
let (task, stored_task, context_id) = self.prepare_task_for_execution(&req).await?;
let task_id = task.id.clone();
self.save_request_push_config(&task_id, &req).await?;
let active = self.execution_manager.start(task.clone()).await?;
let receiver = active.subscribe();
let exec_ctx = crate::ExecutorContext {
message: Some(req.message),
task_id: task_id.clone(),
stored_task: if include_task_snapshot_in_context {
Some(task.clone())
} else {
stored_task
},
context_id,
metadata: req.metadata,
user: None,
service_params: params.clone(),
tenant: req.tenant,
};
let runtime = ExecutionRuntime {
executor: Arc::clone(&self.executor),
task_store: Arc::clone(&self.task_store),
push_config_store: self.push_config_store.clone(),
push_sender: self.push_sender.clone(),
execution_manager: Arc::clone(&self.execution_manager),
};
let active_execution = active.clone();
let execution_task_id = task_id.clone();
tokio::spawn(async move {
drive_execution(runtime, execution_task_id, active_execution, exec_ctx).await;
});
Ok((task_id, subscription_stream(receiver, None, 0)))
}
}
#[async_trait]
impl RequestHandler for DefaultRequestHandler {
async fn send_message(
&self,
params: &ServiceParams,
req: SendMessageRequest,
) -> Result<SendMessageResponse, A2AError> {
self.authorize(params, req.message.task_id.as_deref())?;
let history_length = req
.configuration
.as_ref()
.and_then(|configuration| configuration.history_length);
let (task_id, mut stream) = self.start_execution(params, req.clone(), true).await?;
let mut last_event = None;
while let Some(item) = stream.next().await {
let event = item?;
if let Some(interrupt_task_id) = should_interrupt_non_streaming(&req, &event) {
let mut task = self.load_task(&interrupt_task_id).await?;
apply_history_length(&mut task, history_length);
return Ok(SendMessageResponse::Task(task));
}
match event {
StreamResponse::Message(message) => {
return Ok(SendMessageResponse::Message(message));
}
other => last_event = Some(other),
}
}
match last_event {
Some(StreamResponse::Task(mut task)) => {
apply_history_length(&mut task, history_length);
Ok(SendMessageResponse::Task(task))
}
Some(StreamResponse::StatusUpdate(update)) => {
let mut task = self.load_task(&update.task_id).await?;
apply_history_length(&mut task, history_length);
Ok(SendMessageResponse::Task(task))
}
Some(StreamResponse::ArtifactUpdate(update)) => {
let mut task = self.load_task(&update.task_id).await?;
apply_history_length(&mut task, history_length);
Ok(SendMessageResponse::Task(task))
}
Some(StreamResponse::Message(message)) => Ok(SendMessageResponse::Message(message)),
None => {
let mut task = self.load_task(&task_id).await?;
apply_history_length(&mut task, history_length);
Ok(SendMessageResponse::Task(task))
}
}
}
async fn send_streaming_message(
&self,
params: &ServiceParams,
req: SendMessageRequest,
) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
self.authorize(params, req.message.task_id.as_deref())?;
self.require_streaming()?;
let (_, stream) = self.start_execution(params, req, false).await?;
Ok(stream)
}
async fn get_task(
&self,
params: &ServiceParams,
req: GetTaskRequest,
) -> Result<Task, A2AError> {
self.authorize(params, Some(&req.id))?;
let mut task = self
.task_store
.get(&req.id)
.await?
.ok_or_else(|| A2AError::task_not_found(&req.id))?;
apply_history_length(&mut task, req.history_length);
Ok(task)
}
async fn list_tasks(
&self,
params: &ServiceParams,
req: ListTasksRequest,
) -> Result<ListTasksResponse, A2AError> {
self.authorize(params, None)?;
let page_size = resolve_page_size(req.page_size);
let req = ListTasksRequest {
page_size: Some(page_size as i32),
..req
};
let mut response = self.task_store.list(&req).await?;
response.tasks.truncate(page_size);
for task in &mut response.tasks {
apply_history_length(task, req.history_length);
}
Ok(response)
}
async fn cancel_task(
&self,
params: &ServiceParams,
req: CancelTaskRequest,
) -> Result<Task, A2AError> {
self.authorize(params, Some(&req.id))?;
let task = self.load_task(&req.id).await?;
if task.status.state == TaskState::Canceled {
return Ok(task);
}
if task.status.state.is_terminal() {
return Err(A2AError::task_not_cancelable(&req.id));
}
let task = match self.task_store.begin_cancel(&req.id).await {
Ok(task) => task,
Err(error) if error.code == error_code::TASK_NOT_CANCELABLE => {
let raced = self.load_task(&req.id).await?;
if raced.status.state == TaskState::Canceled {
return Ok(raced);
}
return Err(error);
}
Err(error) => return Err(error),
};
let active_execution = self.execution_manager.get(&req.id).await;
let exec_ctx = crate::ExecutorContext {
message: None,
task_id: req.id.clone(),
stored_task: Some(task.clone()),
context_id: task.context_id.clone(),
metadata: req.metadata,
user: None,
service_params: params.clone(),
tenant: req.tenant,
};
let mut stream = self.executor.cancel(exec_ctx);
let mut current_task = Some(task);
while let Some(event) = stream.next().await {
match event {
Ok(mut event) => {
backfill_status_timestamp(&mut event);
let updated_task =
apply_event_to_task(self.task_store.as_ref(), current_task.clone(), &event)
.await?;
current_task = updated_task.clone();
send_push_notifications(
&req.id,
self.push_config_store.as_deref(),
self.push_sender.as_deref(),
&event,
)
.await?;
if let Some(active) = &active_execution {
active.publish(Ok(&event), updated_task).await;
}
if is_terminal_event(&event) {
break;
}
}
Err(error) => {
if let Some(active) = &active_execution {
active.publish(Err(&error), None).await;
}
return Err(error);
}
}
}
if let Some(active) = active_execution {
self.execution_manager.finish(&req.id, &active).await;
}
self.load_task(&req.id).await
}
async fn subscribe_to_task(
&self,
params: &ServiceParams,
req: SubscribeToTaskRequest,
) -> Result<BoxStream<'static, Result<StreamResponse, A2AError>>, A2AError> {
self.authorize(params, Some(&req.id))?;
self.require_streaming()?;
if let Some((receiver, snapshot_task, sequence)) =
self.execution_manager.resubscribe(&req.id).await
{
return Ok(subscription_stream(receiver, snapshot_task, sequence));
}
if let Some(task) = self.task_store.get(&req.id).await? {
if task.status.state.is_terminal() {
return Err(A2AError::unsupported_operation(format!(
"task {} has already reached a terminal state",
req.id
)));
}
}
Err(A2AError::task_not_found(&req.id))
}
async fn create_push_config(
&self,
params: &ServiceParams,
req: TaskPushNotificationConfig,
) -> Result<TaskPushNotificationConfig, A2AError> {
self.authorize(params, Some(&req.task_id))?;
self.push_config_store_checked()?.save(req).await
}
async fn get_push_config(
&self,
params: &ServiceParams,
req: GetTaskPushNotificationConfigRequest,
) -> Result<TaskPushNotificationConfig, A2AError> {
self.authorize(params, Some(&req.task_id))?;
self.push_config_store_checked()?
.get(&req.task_id, &req.id)
.await
}
async fn list_push_configs(
&self,
params: &ServiceParams,
req: ListTaskPushNotificationConfigsRequest,
) -> Result<ListTaskPushNotificationConfigsResponse, A2AError> {
self.authorize(params, Some(&req.task_id))?;
let mut configs = self.push_config_store_checked()?.list(&req.task_id).await?;
configs.sort_by(|left, right| left.id.cmp(&right.id));
let page_size = resolve_page_size(req.page_size);
let start = if let Some(ref token) = req.page_token {
token
.parse::<usize>()
.map_err(|_| A2AError::invalid_params("invalid page token"))?
} else {
0
};
let start = start.min(configs.len());
let end = (start + page_size).min(configs.len());
let next_page_token = (end < configs.len()).then(|| end.to_string());
Ok(ListTaskPushNotificationConfigsResponse {
configs: configs[start..end].to_vec(),
next_page_token,
})
}
async fn delete_push_config(
&self,
params: &ServiceParams,
req: DeleteTaskPushNotificationConfigRequest,
) -> Result<(), A2AError> {
self.authorize(params, Some(&req.task_id))?;
self.push_config_store_checked()?
.delete(&req.task_id, &req.id)
.await
}
async fn get_extended_agent_card(
&self,
params: &ServiceParams,
_req: GetExtendedAgentCardRequest,
) -> Result<AgentCard, A2AError> {
self.authorize(params, None)?;
Self::require_capability(self.capabilities.extended_agent_card, || {
A2AError::unsupported_operation(
"the agent card does not declare an extended agent card",
)
})?;
match &self.extended_agent_card {
Some(resolver) => resolver.resolve(params),
None => Err(A2AError::extended_card_not_configured()),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::executor::ExecutorContext;
use crate::pagination::{DEFAULT_PAGE_SIZE, MAX_PAGE_SIZE};
use crate::push::{InMemoryPushConfigStore, PushConfigStore};
use crate::task_store::{InMemoryTaskStore, TaskStore, TaskVersion};
use crate::test_util::install_crypto_provider;
use axum::{
Router,
body::Bytes,
extract::State,
http::{HeaderMap, StatusCode, header},
routing::post,
};
use futures::stream;
use std::sync::Arc;
use std::sync::atomic::{AtomicI32, Ordering};
use tokio::{
net::TcpListener,
sync::{Notify, mpsc, oneshot},
time::{Duration, timeout},
};
struct EchoExecutor;
struct PushEventExecutor;
struct ResumableExecutor {
release: Arc<Notify>,
}
#[derive(Debug)]
struct CapturedPush {
authorization: Option<String>,
notification_token: Option<String>,
event: StreamResponse,
}
impl crate::AgentExecutor for EchoExecutor {
fn execute(
&self,
ctx: ExecutorContext,
) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
let task = Task {
id: ctx.task_id.clone(),
context_id: ctx.context_id.clone(),
status: TaskStatus {
state: TaskState::Completed,
message: ctx.message,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
}
fn cancel(
&self,
ctx: ExecutorContext,
) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
let task = Task {
id: ctx.task_id.clone(),
context_id: ctx.context_id.clone(),
status: TaskStatus {
state: TaskState::Canceled,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
}
}
impl crate::AgentExecutor for PushEventExecutor {
fn execute(
&self,
ctx: ExecutorContext,
) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
let working = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
task_id: ctx.task_id.clone(),
context_id: ctx.context_id.clone(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
metadata: None,
});
let completed = StreamResponse::Task(Task {
id: ctx.task_id.clone(),
context_id: ctx.context_id.clone(),
status: TaskStatus {
state: TaskState::Completed,
message: ctx.message,
timestamp: None,
},
artifacts: None,
history: ctx.stored_task.and_then(|task| task.history),
metadata: None,
});
Box::pin(stream::iter(vec![Ok(working), Ok(completed)]))
}
fn cancel(
&self,
ctx: ExecutorContext,
) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
let task = Task {
id: ctx.task_id.clone(),
context_id: ctx.context_id.clone(),
status: TaskStatus {
state: TaskState::Canceled,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
}
}
impl crate::AgentExecutor for ResumableExecutor {
fn execute(
&self,
ctx: ExecutorContext,
) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
let release = self.release.clone();
let working_task = Task {
id: ctx.task_id.clone(),
context_id: ctx.context_id.clone(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
artifacts: None,
history: ctx.stored_task.and_then(|task| task.history),
metadata: None,
};
let completed_task = Task {
id: working_task.id.clone(),
context_id: working_task.context_id.clone(),
status: TaskStatus {
state: TaskState::Completed,
message: None,
timestamp: None,
},
artifacts: None,
history: working_task.history.clone(),
metadata: None,
};
Box::pin(stream::unfold(
(Some(working_task), Some((release, completed_task))),
|(working_task, completion)| async move {
if let Some(task) = working_task {
return Some((Ok(StreamResponse::Task(task)), (None, completion)));
}
if let Some((release, task)) = completion {
release.notified().await;
return Some((Ok(StreamResponse::Task(task)), (None, None)));
}
None
},
))
}
fn cancel(
&self,
ctx: ExecutorContext,
) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
let task = Task {
id: ctx.task_id.clone(),
context_id: ctx.context_id.clone(),
status: TaskStatus {
state: TaskState::Canceled,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
}
}
struct MultiTurnExecutor;
impl crate::AgentExecutor for MultiTurnExecutor {
fn execute(
&self,
ctx: ExecutorContext,
) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
let done = ctx
.message
.as_ref()
.and_then(Message::text)
.is_some_and(|text| text.trim().eq_ignore_ascii_case("done"));
let state = if done {
TaskState::Completed
} else {
TaskState::InputRequired
};
let update = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
task_id: ctx.task_id.clone(),
context_id: ctx.context_id.clone(),
status: TaskStatus {
state,
message: None,
timestamp: None,
},
metadata: None,
});
Box::pin(stream::once(async move { Ok(update) }))
}
fn cancel(
&self,
ctx: ExecutorContext,
) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
let task = Task {
id: ctx.task_id.clone(),
context_id: ctx.context_id.clone(),
status: TaskStatus {
state: TaskState::Canceled,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
Box::pin(stream::once(async move { Ok(StreamResponse::Task(task)) }))
}
}
struct PagingTaskStore {
inner: InMemoryTaskStore,
count: usize,
honour_page_size: bool,
seen_page_size: Arc<AtomicI32>,
next_page_token: String,
}
impl PagingTaskStore {
fn new(count: usize, honour_page_size: bool) -> Self {
Self {
inner: InMemoryTaskStore::new(),
count,
honour_page_size,
seen_page_size: Arc::new(AtomicI32::new(-1)),
next_page_token: String::new(),
}
}
fn with_next_page_token(mut self, token: &str) -> Self {
self.next_page_token = token.to_string();
self
}
}
#[async_trait]
impl TaskStore for PagingTaskStore {
async fn create(&self, task: Task) -> Result<TaskVersion, A2AError> {
self.inner.create(task).await
}
async fn update(&self, task: Task) -> Result<TaskVersion, A2AError> {
self.inner.update(task).await
}
async fn get(&self, task_id: &str) -> Result<Option<Task>, A2AError> {
self.inner.get(task_id).await
}
async fn begin_cancel(&self, task_id: &str) -> Result<Task, A2AError> {
self.inner.begin_cancel(task_id).await
}
async fn list(&self, req: &ListTasksRequest) -> Result<ListTasksResponse, A2AError> {
self.seen_page_size
.store(req.page_size.unwrap_or(-1), Ordering::SeqCst);
let served = if self.honour_page_size {
req.page_size.unwrap_or(i32::MAX).max(0) as usize
} else {
self.count
}
.min(self.count);
let tasks = (0..served)
.map(|i| Task {
id: format!("t-{i:04}"),
context_id: "c1".into(),
status: TaskStatus {
state: TaskState::Completed,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
})
.collect::<Vec<_>>();
let total = tasks.len() as i32;
Ok(ListTasksResponse {
tasks,
next_page_token: self.next_page_token.clone(),
page_size: total,
total_size: total,
})
}
}
struct PagingPushConfigStore {
inner: InMemoryPushConfigStore,
count: usize,
}
impl PagingPushConfigStore {
fn new(count: usize) -> Self {
Self {
inner: InMemoryPushConfigStore::new(),
count,
}
}
}
#[async_trait]
impl crate::PushConfigStore for PagingPushConfigStore {
async fn save(
&self,
config: TaskPushNotificationConfig,
) -> Result<TaskPushNotificationConfig, A2AError> {
self.inner.save(config).await
}
async fn get(
&self,
task_id: &str,
config_id: &str,
) -> Result<TaskPushNotificationConfig, A2AError> {
self.inner.get(task_id, config_id).await
}
async fn list(&self, task_id: &str) -> Result<Vec<TaskPushNotificationConfig>, A2AError> {
Ok((0..self.count)
.map(|i| TaskPushNotificationConfig {
task_id: task_id.to_string(),
url: format!("https://example.com/hook/{i}"),
id: Some(format!("cfg-{i:04}")),
token: None,
authentication: None,
tenant: None,
})
.collect())
}
async fn delete(&self, task_id: &str, config_id: &str) -> Result<(), A2AError> {
self.inner.delete(task_id, config_id).await
}
async fn delete_all(&self, task_id: &str) -> Result<(), A2AError> {
self.inner.delete_all(task_id).await
}
}
async fn seed_task(handler: &DefaultRequestHandler, task_id: &str, context_id: &str) {
handler
.task_store
.create(Task {
id: task_id.into(),
context_id: context_id.into(),
status: TaskStatus {
state: TaskState::Submitted,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
})
.await
.unwrap();
}
fn make_handler() -> DefaultRequestHandler {
DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
}
fn make_handler_with_push_configs() -> DefaultRequestHandler {
DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
.with_push_config_store(InMemoryPushConfigStore::new())
}
fn make_push_delivery_handler() -> DefaultRequestHandler {
DefaultRequestHandler::new(PushEventExecutor, InMemoryTaskStore::new())
.with_push_notifications(
InMemoryPushConfigStore::new(),
crate::HttpPushSender::new(Some(crate::HttpPushSenderConfig {
validate_urls: false,
..Default::default()
})),
)
}
fn make_resumable_handler() -> (DefaultRequestHandler, Arc<Notify>) {
let release = Arc::new(Notify::new());
(
DefaultRequestHandler::new(
ResumableExecutor {
release: release.clone(),
},
InMemoryTaskStore::new(),
),
release,
)
}
fn make_message() -> Message {
Message::new(Role::User, vec![Part::text("hello")])
}
async fn capture_push(
State(sender): State<mpsc::UnboundedSender<CapturedPush>>,
headers: HeaderMap,
body: Bytes,
) -> StatusCode {
let captured = CapturedPush {
authorization: headers
.get(header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned),
notification_token: headers
.get("A2A-Notification-Token")
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned),
event: serde_json::from_slice(&body).unwrap(),
};
sender.send(captured).unwrap();
StatusCode::ACCEPTED
}
async fn start_push_webhook() -> (
String,
mpsc::UnboundedReceiver<CapturedPush>,
oneshot::Sender<()>,
tokio::task::JoinHandle<()>,
) {
let (sender, receiver) = mpsc::unbounded_channel();
let app = Router::new()
.route("/", post(capture_push))
.with_state(sender);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let (shutdown_tx, shutdown_rx) = oneshot::channel();
let server = tokio::spawn(async move {
axum::serve(listener, app)
.with_graceful_shutdown(async {
let _ = shutdown_rx.await;
})
.await
.unwrap();
});
(format!("http://{address}/"), receiver, shutdown_tx, server)
}
async fn next_push(receiver: &mut mpsc::UnboundedReceiver<CapturedPush>) -> CapturedPush {
timeout(Duration::from_secs(5), receiver.recv())
.await
.unwrap()
.unwrap()
}
#[tokio::test]
async fn test_send_message_creates_task() {
let handler = make_handler();
let params = ServiceParams::new();
let req = SendMessageRequest {
message: make_message(),
configuration: None,
metadata: None,
tenant: None,
};
let result = handler.send_message(¶ms, req).await.unwrap();
match result {
SendMessageResponse::Task(task) => {
assert_eq!(task.status.state, TaskState::Completed);
}
_ => panic!("expected Task response"),
}
}
#[tokio::test]
async fn test_send_streaming_message() {
use futures::StreamExt;
let handler = make_handler();
let params = ServiceParams::new();
let req = SendMessageRequest {
message: make_message(),
configuration: None,
metadata: None,
tenant: None,
};
let mut stream = handler.send_streaming_message(¶ms, req).await.unwrap();
let event = stream.next().await.unwrap().unwrap();
assert!(matches!(event, StreamResponse::Task(_)));
}
#[tokio::test]
async fn test_streamed_task_event_has_a_backfilled_timestamp() {
use futures::StreamExt;
let handler = make_handler();
let params = ServiceParams::new();
let req = SendMessageRequest {
message: make_message(),
configuration: None,
metadata: None,
tenant: None,
};
let mut stream = handler.send_streaming_message(¶ms, req).await.unwrap();
let event = stream.next().await.unwrap().unwrap();
match event {
StreamResponse::Task(task) => assert!(
task.status.timestamp.is_some(),
"a live-streamed task event must carry a backfilled timestamp"
),
other => panic!("unexpected event: {other:?}"),
}
}
#[test]
fn test_backfill_status_timestamp_ignores_events_without_a_status() {
let mut artifact_update = StreamResponse::ArtifactUpdate(TaskArtifactUpdateEvent {
task_id: "t1".into(),
context_id: "c1".into(),
artifact: Artifact {
artifact_id: "a1".into(),
name: None,
description: None,
parts: vec![Part::text("payload")],
metadata: None,
extensions: None,
},
append: None,
last_chunk: None,
metadata: None,
});
let before = artifact_update.clone();
backfill_status_timestamp(&mut artifact_update);
assert_eq!(artifact_update, before);
let mut message = StreamResponse::Message(make_message());
let before = message.clone();
backfill_status_timestamp(&mut message);
assert_eq!(message, before);
}
#[tokio::test]
async fn test_status_update_backfills_missing_timestamp() {
let store = InMemoryTaskStore::new();
store
.create(Task {
id: "t-backfill".into(),
context_id: "ctx".into(),
status: TaskStatus {
state: TaskState::Submitted,
message: None,
timestamp: Some(chrono::Utc::now()),
},
artifacts: None,
history: None,
metadata: None,
})
.await
.unwrap();
let updated = apply_event_to_task(
&store,
None,
&StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
task_id: "t-backfill".into(),
context_id: "ctx".into(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
metadata: None,
}),
)
.await
.unwrap()
.expect("status update yields a task");
assert_eq!(updated.status.state, TaskState::Working);
assert!(
updated.status.timestamp.is_some(),
"missing status timestamp must be backfilled"
);
}
#[tokio::test]
async fn test_status_update_preserves_supplied_timestamp() {
let store = InMemoryTaskStore::new();
store
.create(Task {
id: "t-preserve".into(),
context_id: "ctx".into(),
status: TaskStatus {
state: TaskState::Submitted,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
})
.await
.unwrap();
let supplied = chrono::DateTime::parse_from_rfc3339("2020-01-02T03:04:05Z")
.unwrap()
.with_timezone(&chrono::Utc);
let updated = apply_event_to_task(
&store,
None,
&StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
task_id: "t-preserve".into(),
context_id: "ctx".into(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: Some(supplied),
},
metadata: None,
}),
)
.await
.unwrap()
.expect("status update yields a task");
assert_eq!(updated.status.timestamp, Some(supplied));
}
#[tokio::test]
async fn test_task_snapshot_backfills_missing_timestamp() {
let store = InMemoryTaskStore::new();
let returned = apply_event_to_task(
&store,
None,
&StreamResponse::Task(Task {
id: "t-snapshot".into(),
context_id: "ctx".into(),
status: TaskStatus {
state: TaskState::Completed,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
}),
)
.await
.unwrap()
.expect("task event yields a task");
assert!(
returned.status.timestamp.is_some(),
"missing status timestamp must be backfilled on the task-snapshot path"
);
let persisted = store.get("t-snapshot").await.unwrap().unwrap();
assert!(
persisted.status.timestamp.is_some(),
"the persisted task must carry a timestamp"
);
}
#[tokio::test]
async fn test_task_snapshot_preserves_supplied_timestamp() {
let store = InMemoryTaskStore::new();
let supplied = chrono::DateTime::parse_from_rfc3339("2021-06-07T08:09:10Z")
.unwrap()
.with_timezone(&chrono::Utc);
let returned = apply_event_to_task(
&store,
None,
&StreamResponse::Task(Task {
id: "t-snapshot-kept".into(),
context_id: "ctx".into(),
status: TaskStatus {
state: TaskState::Completed,
message: None,
timestamp: Some(supplied),
},
artifacts: None,
history: None,
metadata: None,
}),
)
.await
.unwrap()
.expect("task event yields a task");
assert_eq!(returned.status.timestamp, Some(supplied));
let persisted = store.get("t-snapshot-kept").await.unwrap().unwrap();
assert_eq!(persisted.status.timestamp, Some(supplied));
}
#[tokio::test]
async fn test_stream_disconnect_preserves_execution_for_resubscription() {
use futures::StreamExt;
let (handler, release) = make_resumable_handler();
let params = ServiceParams::new();
let mut message = make_message();
message.task_id = Some("t-disconnect".into());
message.context_id = Some("c-disconnect".into());
seed_task(&handler, "t-disconnect", "c-disconnect").await;
let mut stream = handler
.send_streaming_message(
¶ms,
SendMessageRequest {
message,
configuration: None,
metadata: None,
tenant: None,
},
)
.await
.unwrap();
let initial = stream.next().await.unwrap().unwrap();
match initial {
StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Working),
_ => panic!("expected initial working task"),
}
drop(stream);
let mut resumed = handler
.subscribe_to_task(
¶ms,
SubscribeToTaskRequest {
id: "t-disconnect".into(),
tenant: None,
},
)
.await
.unwrap();
let snapshot = resumed.next().await.unwrap().unwrap();
match snapshot {
StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Working),
_ => panic!("expected working task snapshot"),
}
release.notify_waiters();
let terminal = resumed.next().await.unwrap().unwrap();
match terminal {
StreamResponse::Task(task) => {
assert_eq!(task.status.state, TaskState::Completed)
}
_ => panic!("expected terminal task"),
}
assert!(resumed.next().await.is_none());
}
#[tokio::test]
async fn test_active_execution_keeps_snapshot_without_subscribers() {
let task = Task {
id: "t-buffering".into(),
context_id: "c-buffering".into(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
let active = ActiveExecution::new(task.clone());
let event = StreamResponse::Task(task.clone());
active.publish(Ok(&event), Some(task)).await;
let (mut receiver, snapshot_task, sequence) = active.resubscribe().await;
assert_eq!(snapshot_task.unwrap().id, "t-buffering");
assert_eq!(sequence, 1);
assert!(receiver.try_recv().is_err());
}
#[tokio::test]
async fn test_active_execution_publishes_errors_to_subscribers() {
let task = Task {
id: "t-buffering-error".into(),
context_id: "c-buffering-error".into(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
let active = ActiveExecution::new(task);
let mut receiver = active.subscribe();
let error = A2AError::internal("execution failed");
active.publish(Err(&error), None).await;
let event = receiver.recv().await.unwrap();
assert_eq!(event.sequence, 1);
assert_eq!(event.result.unwrap_err().code, error.code);
}
#[tokio::test]
async fn test_active_execution_drops_errors_without_subscribers() {
let task = Task {
id: "t-error-no-sub".into(),
context_id: "c-error-no-sub".into(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
let active = ActiveExecution::new(task.clone());
let error = A2AError::internal("execution failed");
active.publish(Err(&error), Some(task.clone())).await;
let (mut receiver, snapshot_task, sequence) = active.resubscribe().await;
assert_eq!(snapshot_task.unwrap().id, "t-error-no-sub");
assert_eq!(sequence, 1);
assert!(receiver.try_recv().is_err());
}
#[tokio::test]
async fn test_active_execution_publishes_events_with_snapshot_updates() {
let initial_task = Task {
id: "t-snapshot-update".into(),
context_id: "c-snapshot-update".into(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
let active = ActiveExecution::new(initial_task.clone());
let mut receiver = active.subscribe();
let updated_task = Task {
id: "t-snapshot-update".into(),
context_id: "c-snapshot-update".into(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
let event = StreamResponse::Task(updated_task.clone());
active.publish(Ok(&event), Some(updated_task.clone())).await;
let published = receiver.recv().await.unwrap();
assert_eq!(published.sequence, 1);
match published.result {
Ok(StreamResponse::Task(task)) => {
assert_eq!(task.id, "t-snapshot-update");
}
_ => panic!("expected Ok(Task) event"),
}
let (_, snapshot_task, sequence) = active.resubscribe().await;
assert_eq!(snapshot_task.unwrap().id, "t-snapshot-update");
assert_eq!(sequence, 1);
}
#[tokio::test]
async fn test_active_execution_broadcasts_to_multiple_subscribers() {
let task = Task {
id: "t-multi-sub".into(),
context_id: "c-multi-sub".into(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
let active = ActiveExecution::new(task.clone());
let mut sub1 = active.subscribe();
let mut sub2 = active.subscribe();
let event = StreamResponse::Task(task);
active.publish(Ok(&event), None).await;
let evt1 = sub1.recv().await.unwrap();
let evt2 = sub2.recv().await.unwrap();
assert_eq!(evt1.sequence, 1);
assert_eq!(evt2.sequence, 1);
assert!(matches!(evt1.result, Ok(StreamResponse::Task(_))));
assert!(matches!(evt2.result, Ok(StreamResponse::Task(_))));
}
#[tokio::test]
async fn test_get_task_not_found() {
let handler = make_handler();
let params = ServiceParams::new();
let req = GetTaskRequest {
id: "nonexistent".into(),
history_length: None,
tenant: None,
};
let result = handler.get_task(¶ms, req).await;
assert!(result.is_err());
}
fn send_req(task_id: Option<&str>, context_id: Option<&str>) -> SendMessageRequest {
let mut msg = make_message();
msg.task_id = task_id.map(Into::into);
msg.context_id = context_id.map(Into::into);
SendMessageRequest {
message: msg,
configuration: None,
metadata: None,
tenant: None,
}
}
#[tokio::test]
async fn test_send_message_rejects_unsupported_content_type() {
install_crypto_provider();
let handler = make_handler().with_default_input_modes(vec!["text/plain".into()]);
let mut req = send_req(None, None);
req.message.parts = vec![
Part::data(serde_json::json!({"value": "unsupported"}))
.with_media_type("application/x-unsupported-type-12345"),
];
let err = handler
.send_message(&ServiceParams::new(), req)
.await
.expect_err("an undeclared content type must be rejected");
assert_eq!(err.code, error_code::CONTENT_TYPE_NOT_SUPPORTED);
}
#[tokio::test]
async fn test_send_message_accepts_a_declared_content_type() {
install_crypto_provider();
let handler = make_handler().with_default_input_modes(vec!["text/plain".into()]);
let mut req = send_req(None, None);
req.message.parts = vec![Part::text("hello").with_media_type("text/plain")];
handler
.send_message(&ServiceParams::new(), req)
.await
.expect("a declared content type must be accepted");
}
#[tokio::test]
async fn test_send_message_allows_any_content_type_when_none_declared() {
install_crypto_provider();
let handler = make_handler();
let mut req = send_req(None, None);
req.message.parts = vec![
Part::data(serde_json::json!({"value": "x"})).with_media_type("application/x-anything"),
];
handler
.send_message(&ServiceParams::new(), req)
.await
.expect("no declared modes means no restriction");
}
#[tokio::test]
async fn test_send_message_unknown_task_id_is_not_found() {
install_crypto_provider();
let handler = make_handler();
let err = handler
.send_message(&ServiceParams::new(), send_req(Some("t-absent"), None))
.await
.expect_err("an unknown taskId must not create a task");
assert_eq!(err.code, error_code::TASK_NOT_FOUND);
assert!(
handler.task_store.get("t-absent").await.unwrap().is_none(),
"the rejected id must not have been created"
);
}
#[tokio::test]
async fn test_send_message_to_terminal_task_is_rejected() {
install_crypto_provider();
let handler = make_handler();
handler
.task_store
.create(Task {
id: "t-done".into(),
context_id: "c-done".into(),
status: TaskStatus {
state: TaskState::Completed,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
})
.await
.unwrap();
let err = handler
.send_message(&ServiceParams::new(), send_req(Some("t-done"), None))
.await
.expect_err("a follow-up to a terminal task must not be accepted");
assert_eq!(err.code, error_code::UNSUPPORTED_OPERATION);
}
#[tokio::test]
async fn test_send_message_mismatched_context_id_is_rejected() {
install_crypto_provider();
let handler = make_handler();
seed_task(&handler, "t-ctx", "c-real").await;
let err = handler
.send_message(
&ServiceParams::new(),
send_req(Some("t-ctx"), Some("c-wrong")),
)
.await
.expect_err("a mismatching contextId must be rejected");
assert_eq!(err.code, error_code::INVALID_PARAMS);
assert!(
err.message.contains("c-wrong"),
"error names the bad id: {}",
err.message
);
}
#[tokio::test]
async fn test_send_message_matching_context_id_is_accepted() {
install_crypto_provider();
let handler = make_handler();
seed_task(&handler, "t-ok", "c-ok").await;
handler
.send_message(&ServiceParams::new(), send_req(Some("t-ok"), Some("c-ok")))
.await
.expect("a matching pair must be accepted");
}
#[tokio::test]
async fn test_send_message_accumulates_history_across_turns() {
install_crypto_provider();
let handler = DefaultRequestHandler::new(MultiTurnExecutor, InMemoryTaskStore::new());
let mut turn1 = make_message();
turn1.parts = vec![Part::text("turn one")];
let resp1 = handler
.send_message(
&ServiceParams::new(),
SendMessageRequest {
message: turn1,
configuration: None,
metadata: None,
tenant: None,
},
)
.await
.unwrap();
let task_id = match resp1 {
SendMessageResponse::Task(t) => t.id,
SendMessageResponse::Message(_) => panic!("expected a task"),
};
for text in ["turn two", "done"] {
let mut turn = make_message();
turn.task_id = Some(task_id.clone());
turn.parts = vec![Part::text(text)];
handler
.send_message(
&ServiceParams::new(),
SendMessageRequest {
message: turn,
configuration: None,
metadata: None,
tenant: None,
},
)
.await
.unwrap();
}
let task = handler.task_store.get(&task_id).await.unwrap().unwrap();
let history = task.history.expect("history must accumulate across turns");
let texts: Vec<_> = history.iter().map(|m| m.text()).collect();
assert_eq!(
texts,
vec![Some("turn one"), Some("turn two"), Some("done")]
);
}
#[tokio::test]
async fn test_send_message_applies_configuration_history_length_to_its_own_response() {
install_crypto_provider();
let handler = DefaultRequestHandler::new(MultiTurnExecutor, InMemoryTaskStore::new());
let mut message = make_message();
message.parts = vec![Part::text("done")];
let resp = handler
.send_message(
&ServiceParams::new(),
SendMessageRequest {
message,
configuration: Some(SendMessageConfiguration {
accepted_output_modes: None,
task_push_notification_config: None,
history_length: Some(0),
return_immediately: None,
}),
metadata: None,
tenant: None,
},
)
.await
.unwrap();
let task = match resp {
SendMessageResponse::Task(t) => t,
SendMessageResponse::Message(_) => panic!("expected a task"),
};
assert!(
task.history.is_none_or(|h| h.is_empty()),
"historyLength=0 must clear history on SendMessage's own response"
);
}
#[tokio::test]
async fn test_send_message_infers_context_id_from_the_task() {
install_crypto_provider();
let handler = make_handler();
seed_task(&handler, "t-infer", "c-infer").await;
handler
.send_message(&ServiceParams::new(), send_req(Some("t-infer"), None))
.await
.unwrap();
let task = handler.task_store.get("t-infer").await.unwrap().unwrap();
assert_eq!(task.context_id, "c-infer");
}
#[tokio::test]
async fn test_send_message_without_task_id_still_creates_one() {
install_crypto_provider();
let handler = make_handler();
let resp = handler
.send_message(&ServiceParams::new(), send_req(None, Some("c-new")))
.await
.expect("a message with no taskId must still create a task");
let _ = resp;
let listed = handler
.task_store
.list(&ListTasksRequest {
context_id: Some("c-new".into()),
status: None,
page_size: None,
page_token: None,
history_length: None,
include_artifacts: None,
status_timestamp_after: None,
tenant: None,
})
.await
.unwrap();
assert_eq!(listed.tasks.len(), 1);
assert_eq!(listed.tasks[0].context_id, "c-new");
}
#[tokio::test]
async fn test_get_task_after_send() {
let handler = make_handler();
let params = ServiceParams::new();
let mut msg = make_message();
msg.task_id = Some("t1".into());
msg.context_id = Some("c1".into());
seed_task(&handler, "t1", "c1").await;
let req = SendMessageRequest {
message: msg,
configuration: None,
metadata: None,
tenant: None,
};
handler.send_message(¶ms, req).await.unwrap();
let task = handler
.get_task(
¶ms,
GetTaskRequest {
id: "t1".into(),
history_length: None,
tenant: None,
},
)
.await
.unwrap();
assert_eq!(task.id, "t1");
}
#[tokio::test]
async fn test_list_tasks() {
let handler = make_handler();
let params = ServiceParams::new();
let req = ListTasksRequest {
context_id: None,
status: None,
page_size: None,
page_token: None,
history_length: None,
status_timestamp_after: None,
include_artifacts: None,
tenant: None,
};
let resp = handler.list_tasks(¶ms, req).await.unwrap();
assert!(resp.tasks.is_empty());
}
#[tokio::test]
async fn test_cancel_task_not_found() {
let handler = make_handler();
let params = ServiceParams::new();
let req = CancelTaskRequest {
id: "nonexistent".into(),
metadata: None,
tenant: None,
};
let result = handler.cancel_task(¶ms, req).await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_cancel_task_working() {
let handler = make_handler();
let params = ServiceParams::new();
let task = Task {
id: "t2".into(),
context_id: "c2".into(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
handler.task_store.create(task).await.unwrap();
let req = CancelTaskRequest {
id: "t2".into(),
metadata: None,
tenant: None,
};
let result = handler.cancel_task(¶ms, req).await.unwrap();
assert_eq!(result.status.state, TaskState::Canceled);
}
#[tokio::test]
async fn test_cancel_task_concurrent_both_succeed_idempotently() {
let handler = Arc::new(make_handler());
let params = ServiceParams::new();
let task = Task {
id: "t-race".into(),
context_id: "c-race".into(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
handler.task_store.create(task).await.unwrap();
let handlers = vec![Arc::clone(&handler), Arc::clone(&handler)];
let mut results = Vec::new();
for handler in handlers {
let params = params.clone();
let req = CancelTaskRequest {
id: "t-race".into(),
metadata: None,
tenant: None,
};
results.push(tokio::spawn(async move {
handler.cancel_task(¶ms, req).await
}));
}
for result in results {
let task = result.await.unwrap().unwrap();
assert_eq!(task.status.state, TaskState::Canceled);
}
}
#[tokio::test]
async fn test_cancel_task_already_completed() {
let handler = make_handler();
let params = ServiceParams::new();
let task = Task {
id: "t3".into(),
context_id: "c3".into(),
status: TaskStatus {
state: TaskState::Completed,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
handler.task_store.create(task).await.unwrap();
let req = CancelTaskRequest {
id: "t3".into(),
metadata: None,
tenant: None,
};
let result = handler.cancel_task(¶ms, req).await;
match result {
Ok(_) => panic!("expected cancel of COMPLETED to fail"),
Err(error) => assert_eq!(error.code, error_code::TASK_NOT_CANCELABLE),
}
}
#[tokio::test]
async fn test_cancel_task_already_canceled_is_idempotent() {
let handler = make_handler();
let params = ServiceParams::new();
let task = Task {
id: "t-canceled".into(),
context_id: "c-canceled".into(),
status: TaskStatus {
state: TaskState::Canceled,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
handler.task_store.create(task).await.unwrap();
let req = CancelTaskRequest {
id: "t-canceled".into(),
metadata: None,
tenant: None,
};
let result = handler.cancel_task(¶ms, req).await.unwrap();
assert_eq!(result.status.state, TaskState::Canceled);
}
#[tokio::test]
async fn test_cancel_task_failed_and_rejected_are_not_cancelable() {
let handler = make_handler();
let params = ServiceParams::new();
for (id, state) in [
("t-failed", TaskState::Failed),
("t-rejected", TaskState::Rejected),
] {
handler
.task_store
.create(Task {
id: id.into(),
context_id: format!("c-{id}"),
status: TaskStatus {
state,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
})
.await
.unwrap();
let result = handler
.cancel_task(
¶ms,
CancelTaskRequest {
id: id.into(),
metadata: None,
tenant: None,
},
)
.await;
match result {
Ok(_) => panic!("expected cancel of {id} to fail"),
Err(error) => assert_eq!(error.code, error_code::TASK_NOT_CANCELABLE),
}
}
}
fn task_with_history(id: &str, context_id: &str, n: usize) -> Task {
Task {
id: id.into(),
context_id: context_id.into(),
status: TaskStatus {
state: TaskState::Completed,
message: None,
timestamp: None,
},
artifacts: None,
history: Some(
(0..n)
.map(|i| Message::new(Role::User, vec![Part::text(format!("m{i}"))]))
.collect(),
),
metadata: None,
}
}
#[tokio::test]
async fn test_get_task_applies_history_length() {
let handler = make_handler();
let params = ServiceParams::new();
handler
.task_store
.create(task_with_history("hist", "c", 3))
.await
.unwrap();
let full = handler
.get_task(
¶ms,
GetTaskRequest {
id: "hist".into(),
history_length: None,
tenant: None,
},
)
.await
.unwrap();
assert_eq!(full.history.as_ref().unwrap().len(), 3);
let last = handler
.get_task(
¶ms,
GetTaskRequest {
id: "hist".into(),
history_length: Some(1),
tenant: None,
},
)
.await
.unwrap();
assert_eq!(last.history.as_ref().unwrap().len(), 1);
for hl in [0, -1] {
let empty = handler
.get_task(
¶ms,
GetTaskRequest {
id: "hist".into(),
history_length: Some(hl),
tenant: None,
},
)
.await
.unwrap();
assert!(
empty.history.as_ref().unwrap().is_empty(),
"history_length={hl} must clear history"
);
}
}
fn reject_unauthenticated(
_params: &ServiceParams,
_task_id: Option<&str>,
) -> Result<(), A2AError> {
Err(A2AError::invalid_request("authentication required"))
}
#[tokio::test]
async fn test_authorizer_rejects_handler_methods() {
let handler = DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
.with_authorizer(reject_unauthenticated)
.with_push_config_store(InMemoryPushConfigStore::new());
let params = ServiceParams::new();
assert!(
handler
.send_message(
¶ms,
SendMessageRequest {
message: make_message(),
configuration: None,
metadata: None,
tenant: None,
},
)
.await
.is_err()
);
assert!(
handler
.get_task(
¶ms,
GetTaskRequest {
id: "t".into(),
history_length: None,
tenant: None,
},
)
.await
.is_err()
);
assert!(
handler
.list_tasks(
¶ms,
ListTasksRequest {
context_id: None,
status: None,
page_size: None,
page_token: None,
history_length: None,
status_timestamp_after: None,
include_artifacts: None,
tenant: None,
},
)
.await
.is_err()
);
assert!(
handler
.cancel_task(
¶ms,
CancelTaskRequest {
id: "t".into(),
metadata: None,
tenant: None,
},
)
.await
.is_err()
);
assert!(
handler
.subscribe_to_task(
¶ms,
SubscribeToTaskRequest {
id: "t".into(),
tenant: None,
},
)
.await
.is_err()
);
assert!(
handler
.create_push_config(
¶ms,
TaskPushNotificationConfig {
url: "https://example.com/push".into(),
id: None,
task_id: "t".into(),
token: None,
authentication: None,
tenant: None,
},
)
.await
.is_err()
);
assert!(
handler
.get_push_config(
¶ms,
GetTaskPushNotificationConfigRequest {
task_id: "t".into(),
id: "cfg".into(),
tenant: None,
},
)
.await
.is_err()
);
assert!(
handler
.list_push_configs(
¶ms,
ListTaskPushNotificationConfigsRequest {
task_id: "t".into(),
page_size: None,
page_token: None,
tenant: None,
},
)
.await
.is_err()
);
assert!(
handler
.delete_push_config(
¶ms,
DeleteTaskPushNotificationConfigRequest {
task_id: "t".into(),
id: "cfg".into(),
tenant: None,
},
)
.await
.is_err()
);
assert!(
handler
.get_extended_agent_card(¶ms, GetExtendedAgentCardRequest { tenant: None })
.await
.is_err()
);
}
#[tokio::test]
async fn test_subscribe_requires_active_execution() {
let handler = make_handler();
let params = ServiceParams::new();
let req = SubscribeToTaskRequest {
id: "t1".into(),
tenant: None,
};
let result = handler.subscribe_to_task(¶ms, req).await;
match result {
Ok(_) => panic!("expected subscribe_to_task to fail"),
Err(error) => assert_eq!(error.code, error_code::TASK_NOT_FOUND),
}
}
#[tokio::test]
async fn test_subscribe_to_terminal_task_is_unsupported_not_not_found() {
let handler = make_handler();
handler
.task_store
.create(Task {
id: "t-done".into(),
context_id: "c-done".into(),
status: TaskStatus {
state: TaskState::Completed,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
})
.await
.unwrap();
let result = handler
.subscribe_to_task(
&ServiceParams::new(),
SubscribeToTaskRequest {
id: "t-done".into(),
tenant: None,
},
)
.await;
match result {
Ok(_) => panic!("expected subscribe_to_task to fail"),
Err(error) => assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION),
}
}
#[tokio::test]
async fn test_send_message_return_immediately_allows_resubscribe() {
use futures::StreamExt;
let (handler, release) = make_resumable_handler();
let params = ServiceParams::new();
let mut message = make_message();
message.task_id = Some("t-resume".into());
message.context_id = Some("c-resume".into());
seed_task(&handler, "t-resume", "c-resume").await;
let response = handler
.send_message(
¶ms,
SendMessageRequest {
message,
configuration: Some(SendMessageConfiguration {
accepted_output_modes: None,
task_push_notification_config: None,
history_length: None,
return_immediately: Some(true),
}),
metadata: None,
tenant: None,
},
)
.await
.unwrap();
match response {
SendMessageResponse::Task(task) => {
assert_eq!(task.status.state, TaskState::Working);
}
_ => panic!("expected Task response"),
}
let mut subscription = handler
.subscribe_to_task(
¶ms,
SubscribeToTaskRequest {
id: "t-resume".into(),
tenant: None,
},
)
.await
.unwrap();
let snapshot = subscription.next().await.unwrap().unwrap();
match snapshot {
StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Working),
_ => panic!("expected snapshot task"),
}
release.notify_waiters();
let terminal = subscription.next().await.unwrap().unwrap();
match terminal {
StreamResponse::Task(task) => assert_eq!(task.status.state, TaskState::Completed),
_ => panic!("expected terminal task"),
}
assert!(subscription.next().await.is_none());
let result = handler
.subscribe_to_task(
¶ms,
SubscribeToTaskRequest {
id: "t-resume".into(),
tenant: None,
},
)
.await;
match result {
Ok(_) => panic!("expected subscribe_to_task to fail after completion"),
Err(error) => assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION),
}
}
#[tokio::test]
async fn test_push_config_not_supported_without_store() {
let handler = make_handler();
let params = ServiceParams::new();
let result = handler
.create_push_config(
¶ms,
TaskPushNotificationConfig {
task_id: "t1".into(),
url: "http://example.com".into(),
id: None,
token: None,
authentication: None,
tenant: None,
},
)
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_list_push_configs_invalid_page_token_errors() {
install_crypto_provider();
let handler = make_handler_with_push_configs();
let params = ServiceParams::new();
handler
.create_push_config(
¶ms,
TaskPushNotificationConfig {
task_id: "t1".into(),
url: "https://example.com/callback".into(),
id: Some("cfg-1".into()),
token: None,
authentication: None,
tenant: None,
},
)
.await
.unwrap();
let result = handler
.list_push_configs(
¶ms,
ListTaskPushNotificationConfigsRequest {
task_id: "t1".into(),
page_size: None,
page_token: Some("not-a-number".into()),
tenant: None,
},
)
.await;
assert!(result.is_err());
assert_eq!(result.unwrap_err().code, error_code::INVALID_PARAMS);
}
#[tokio::test]
async fn test_list_tasks_page_size_cap_bounds_a_store_that_ignores_it() {
install_crypto_provider();
let handler = DefaultRequestHandler::new(
EchoExecutor,
PagingTaskStore::new(MAX_PAGE_SIZE * 3, false),
);
let resp = handler
.list_tasks(
&ServiceParams::new(),
ListTasksRequest {
context_id: None,
status: None,
page_size: Some(1_000_000),
page_token: None,
history_length: None,
include_artifacts: None,
status_timestamp_after: None,
tenant: None,
},
)
.await
.unwrap();
assert_eq!(resp.tasks.len(), MAX_PAGE_SIZE);
}
#[tokio::test]
async fn test_paging_stores_delegate_everything_except_list() {
install_crypto_provider();
let tasks = PagingTaskStore::new(0, true);
let task = Task {
id: "t-delegated".into(),
context_id: "c1".into(),
status: TaskStatus {
state: TaskState::Submitted,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
};
tasks.create(task.clone()).await.unwrap();
assert_eq!(
tasks
.get("t-delegated")
.await
.unwrap()
.map(|t| t.status.state),
Some(TaskState::Submitted)
);
let mut done = task.clone();
done.status.state = TaskState::Completed;
tasks.update(done).await.unwrap();
assert_eq!(
tasks
.get("t-delegated")
.await
.unwrap()
.map(|t| t.status.state),
Some(TaskState::Completed)
);
let mut cancelable = task.clone();
cancelable.id = "t-cancelable".into();
cancelable.status.state = TaskState::Working;
tasks.create(cancelable).await.unwrap();
assert_eq!(
tasks
.begin_cancel("t-cancelable")
.await
.unwrap()
.status
.state,
TaskState::Canceled
);
assert_eq!(
tasks.begin_cancel("t-delegated").await.unwrap_err().code,
error_code::TASK_NOT_CANCELABLE
);
let configs = PagingPushConfigStore::new(0);
configs
.save(TaskPushNotificationConfig {
task_id: "t1".into(),
url: "https://example.com/callback".into(),
id: Some("cfg-1".into()),
token: None,
authentication: None,
tenant: None,
})
.await
.unwrap();
assert_eq!(
configs.get("t1", "cfg-1").await.unwrap().id.as_deref(),
Some("cfg-1")
);
configs.delete("t1", "cfg-1").await.unwrap();
assert!(configs.get("t1", "cfg-1").await.is_err());
configs.delete_all("t1").await.unwrap();
}
#[tokio::test]
async fn test_list_tasks_page_size_cap_bounds_a_store_that_honours_it() {
install_crypto_provider();
let store = PagingTaskStore::new(MAX_PAGE_SIZE * 3, true);
let seen = Arc::clone(&store.seen_page_size);
let handler = DefaultRequestHandler::new(EchoExecutor, store);
let resp = handler
.list_tasks(
&ServiceParams::new(),
ListTasksRequest {
context_id: None,
status: None,
page_size: Some(1_000_000),
page_token: None,
history_length: None,
include_artifacts: None,
status_timestamp_after: None,
tenant: None,
},
)
.await
.unwrap();
assert_eq!(resp.tasks.len(), MAX_PAGE_SIZE);
assert_eq!(seen.load(Ordering::SeqCst), MAX_PAGE_SIZE as i32);
}
#[tokio::test]
async fn test_list_tasks_truncation_keeps_the_store_continuation_token() {
install_crypto_provider();
let handler = DefaultRequestHandler::new(
EchoExecutor,
PagingTaskStore::new(MAX_PAGE_SIZE * 3, false)
.with_next_page_token("cursor-from-store"),
);
let resp = handler
.list_tasks(
&ServiceParams::new(),
ListTasksRequest {
context_id: None,
status: None,
page_size: None,
page_token: None,
history_length: None,
include_artifacts: None,
status_timestamp_after: None,
tenant: None,
},
)
.await
.unwrap();
assert_eq!(resp.tasks.len(), DEFAULT_PAGE_SIZE);
assert_eq!(resp.next_page_token, "cursor-from-store");
}
#[tokio::test]
async fn test_list_tasks_honours_a_page_size_under_the_cap() {
install_crypto_provider();
let store = PagingTaskStore::new(MAX_PAGE_SIZE * 3, true);
let seen = Arc::clone(&store.seen_page_size);
let handler = DefaultRequestHandler::new(EchoExecutor, store);
let resp = handler
.list_tasks(
&ServiceParams::new(),
ListTasksRequest {
context_id: None,
status: None,
page_size: Some(7),
page_token: None,
history_length: None,
include_artifacts: None,
status_timestamp_after: None,
tenant: None,
},
)
.await
.unwrap();
assert_eq!(seen.load(Ordering::SeqCst), 7);
assert_eq!(resp.tasks.len(), 7);
}
#[tokio::test]
async fn test_list_push_configs_page_size_cap_bounds_a_custom_store() {
install_crypto_provider();
let handler = DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
.with_push_config_store(PagingPushConfigStore::new(MAX_PAGE_SIZE * 3));
let resp = handler
.list_push_configs(
&ServiceParams::new(),
ListTaskPushNotificationConfigsRequest {
task_id: "t1".into(),
page_size: Some(1_000_000),
page_token: None,
tenant: None,
},
)
.await
.unwrap();
assert_eq!(resp.configs.len(), MAX_PAGE_SIZE);
assert_eq!(resp.next_page_token.as_deref(), Some("100"));
}
#[tokio::test]
async fn test_list_push_configs_page_size_is_capped() {
install_crypto_provider();
let handler = make_handler_with_push_configs();
let params = ServiceParams::new();
for i in 0..40 {
handler
.create_push_config(
¶ms,
TaskPushNotificationConfig {
task_id: "t1".into(),
url: format!("https://example.com/callback/{i}"),
id: Some(format!("cfg-{i:03}")),
token: None,
authentication: None,
tenant: None,
},
)
.await
.unwrap();
}
let resp = handler
.list_push_configs(
¶ms,
ListTaskPushNotificationConfigsRequest {
task_id: "t1".into(),
page_size: Some(1000),
page_token: None,
tenant: None,
},
)
.await
.unwrap();
assert_eq!(resp.configs.len(), 40);
}
#[tokio::test]
async fn test_push_config_crud_with_store() {
install_crypto_provider();
let handler = make_handler_with_push_configs();
let params = ServiceParams::new();
let created = handler
.create_push_config(
¶ms,
TaskPushNotificationConfig {
task_id: "t1".into(),
url: "https://example.com/first".into(),
id: Some("cfg-1".into()),
token: None,
authentication: None,
tenant: Some("tenant-a".into()),
},
)
.await
.unwrap();
assert_eq!(created.id.as_deref(), Some("cfg-1"));
handler
.create_push_config(
¶ms,
TaskPushNotificationConfig {
task_id: "t1".into(),
url: "https://example.com/second".into(),
id: Some("cfg-2".into()),
token: None,
authentication: None,
tenant: Some("tenant-a".into()),
},
)
.await
.unwrap();
let fetched = handler
.get_push_config(
¶ms,
GetTaskPushNotificationConfigRequest {
task_id: "t1".into(),
id: "cfg-1".into(),
tenant: Some("tenant-a".into()),
},
)
.await
.unwrap();
assert_eq!(fetched.url, "https://example.com/first");
let first_page = handler
.list_push_configs(
¶ms,
ListTaskPushNotificationConfigsRequest {
task_id: "t1".into(),
page_size: Some(1),
page_token: Some("0".into()),
tenant: Some("tenant-a".into()),
},
)
.await
.unwrap();
assert_eq!(first_page.configs.len(), 1);
assert_eq!(first_page.configs[0].id.as_deref(), Some("cfg-1"));
assert_eq!(first_page.next_page_token.as_deref(), Some("1"));
let default_page = handler
.list_push_configs(
¶ms,
ListTaskPushNotificationConfigsRequest {
task_id: "t1".into(),
page_size: Some(0),
page_token: None,
tenant: Some("tenant-a".into()),
},
)
.await
.unwrap();
assert_eq!(default_page.configs.len(), 2);
assert!(default_page.next_page_token.is_none());
handler
.delete_push_config(
¶ms,
DeleteTaskPushNotificationConfigRequest {
task_id: "t1".into(),
id: "cfg-1".into(),
tenant: Some("tenant-a".into()),
},
)
.await
.unwrap();
let remaining = handler
.list_push_configs(
¶ms,
ListTaskPushNotificationConfigsRequest {
task_id: "t1".into(),
page_size: None,
page_token: None,
tenant: Some("tenant-a".into()),
},
)
.await
.unwrap();
assert_eq!(remaining.configs.len(), 1);
assert_eq!(remaining.configs[0].id.as_deref(), Some("cfg-2"));
}
#[tokio::test]
async fn test_send_message_request_push_config_delivers_webhooks() {
install_crypto_provider();
let (url, mut receiver, shutdown_tx, server) = start_push_webhook().await;
let handler = make_push_delivery_handler();
let params = ServiceParams::new();
let mut message = make_message();
message.task_id = Some("t-push-request".into());
message.context_id = Some("c-push-request".into());
seed_task(&handler, "t-push-request", "c-push-request").await;
let response = handler
.send_message(
¶ms,
SendMessageRequest {
message,
configuration: Some(SendMessageConfiguration {
accepted_output_modes: None,
task_push_notification_config: Some(TaskPushNotificationConfig {
task_id: String::new(),
url: url.clone(),
id: Some("cfg-request".into()),
token: Some("notify-token".into()),
authentication: Some(AuthenticationInfo {
scheme: "bearer".into(),
credentials: Some("secret-token".into()),
}),
tenant: None,
}),
history_length: None,
return_immediately: None,
}),
metadata: None,
tenant: None,
},
)
.await
.unwrap();
match response {
SendMessageResponse::Task(task) => assert_eq!(task.status.state, TaskState::Completed),
_ => panic!("expected Task response"),
}
let saved = handler
.get_push_config(
¶ms,
GetTaskPushNotificationConfigRequest {
task_id: "t-push-request".into(),
id: "cfg-request".into(),
tenant: None,
},
)
.await
.unwrap();
assert_eq!(saved.url, url);
let first = next_push(&mut receiver).await;
assert_eq!(first.authorization.as_deref(), Some("Bearer secret-token"));
assert_eq!(first.notification_token.as_deref(), Some("notify-token"));
match first.event {
StreamResponse::StatusUpdate(update) => {
assert_eq!(update.task_id, "t-push-request");
assert_eq!(update.status.state, TaskState::Working);
}
_ => panic!("expected status update push"),
}
let second = next_push(&mut receiver).await;
assert_eq!(second.authorization.as_deref(), Some("Bearer secret-token"));
assert_eq!(second.notification_token.as_deref(), Some("notify-token"));
match second.event {
StreamResponse::Task(task) => {
assert_eq!(task.id, "t-push-request");
assert_eq!(task.status.state, TaskState::Completed);
}
_ => panic!("expected final task push"),
}
shutdown_tx.send(()).unwrap();
server.await.unwrap();
}
#[tokio::test]
async fn test_send_message_stored_push_config_delivers_webhooks() {
install_crypto_provider();
let (url, mut receiver, shutdown_tx, server) = start_push_webhook().await;
let handler = make_push_delivery_handler();
let params = ServiceParams::new();
handler
.create_push_config(
¶ms,
TaskPushNotificationConfig {
task_id: "t-push-stored".into(),
url: url.clone(),
id: Some("cfg-stored".into()),
token: Some("stored-token".into()),
authentication: None,
tenant: None,
},
)
.await
.unwrap();
let mut message = make_message();
message.task_id = Some("t-push-stored".into());
message.context_id = Some("c-push-stored".into());
seed_task(&handler, "t-push-stored", "c-push-stored").await;
let response = handler
.send_message(
¶ms,
SendMessageRequest {
message,
configuration: None,
metadata: None,
tenant: None,
},
)
.await
.unwrap();
match response {
SendMessageResponse::Task(task) => assert_eq!(task.status.state, TaskState::Completed),
_ => panic!("expected Task response"),
}
let first = next_push(&mut receiver).await;
assert_eq!(first.notification_token.as_deref(), Some("stored-token"));
match first.event {
StreamResponse::StatusUpdate(update) => {
assert_eq!(update.task_id, "t-push-stored");
assert_eq!(update.status.state, TaskState::Working);
}
_ => panic!("expected status update push"),
}
let second = next_push(&mut receiver).await;
assert_eq!(second.notification_token.as_deref(), Some("stored-token"));
match second.event {
StreamResponse::Task(task) => {
assert_eq!(task.id, "t-push-stored");
assert_eq!(task.status.state, TaskState::Completed);
}
_ => panic!("expected final task push"),
}
shutdown_tx.send(()).unwrap();
server.await.unwrap();
}
#[tokio::test]
async fn test_extended_agent_card_not_configured() {
let handler = make_handler();
let params = ServiceParams::new();
let result = handler
.get_extended_agent_card(¶ms, GetExtendedAgentCardRequest { tenant: None })
.await;
assert!(result.is_err());
}
#[test]
fn test_with_capabilities() {
let handler = make_handler().with_capabilities(AgentCapabilities {
streaming: Some(true),
push_notifications: None,
extensions: None,
extended_agent_card: None,
});
assert_eq!(handler.capabilities.streaming, Some(true));
}
#[test]
fn test_with_push_config_store_enables_push_capability() {
install_crypto_provider();
let handler = make_handler_with_push_configs();
assert_eq!(handler.capabilities.push_notifications, Some(true));
}
fn caps(streaming: Option<bool>, push: Option<bool>) -> AgentCapabilities {
AgentCapabilities {
streaming,
push_notifications: push,
extensions: None,
extended_agent_card: None,
}
}
fn send_request() -> SendMessageRequest {
SendMessageRequest {
message: Message::new(Role::User, vec![Part::text("hi")]),
configuration: None,
metadata: None,
tenant: None,
}
}
fn push_config(task_id: &str) -> TaskPushNotificationConfig {
TaskPushNotificationConfig {
task_id: task_id.into(),
url: "https://example.com/callback".into(),
id: Some("cfg-1".into()),
token: None,
authentication: None,
tenant: None,
}
}
#[tokio::test]
async fn test_streaming_refused_when_the_card_declares_it_absent() {
let handler = make_handler().with_capabilities(caps(Some(false), None));
let params = ServiceParams::new();
match handler
.send_streaming_message(¶ms, send_request())
.await
{
Err(error) => assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION),
Ok(_) => panic!("streaming must be refused when the card declares it absent"),
}
match handler
.subscribe_to_task(
¶ms,
SubscribeToTaskRequest {
id: "t1".into(),
tenant: None,
},
)
.await
{
Err(error) => assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION),
Ok(_) => panic!("subscribe must be refused when the card declares it absent"),
}
}
#[tokio::test]
async fn test_push_refused_when_the_card_declares_it_absent() {
let handler = make_handler_with_push_configs().with_capabilities(caps(None, Some(false)));
let params = ServiceParams::new();
let error = handler
.create_push_config(¶ms, push_config("t1"))
.await
.unwrap_err();
assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
let error = handler
.list_push_configs(
¶ms,
ListTaskPushNotificationConfigsRequest {
task_id: "t1".into(),
page_size: None,
page_token: None,
tenant: None,
},
)
.await
.unwrap_err();
assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
}
#[tokio::test]
async fn test_undeclared_capabilities_stay_permissive() {
let handler = make_handler_with_push_configs();
let params = ServiceParams::new();
assert!(
handler
.send_streaming_message(¶ms, send_request())
.await
.is_ok(),
"an undeclared capability must not be treated as refused"
);
assert!(
handler
.create_push_config(¶ms, push_config("t1"))
.await
.is_ok()
);
}
#[tokio::test]
async fn test_builder_order_does_not_change_the_declared_capability() {
let params = ServiceParams::new();
for handler in [
make_handler()
.with_capabilities(caps(None, Some(false)))
.with_push_config_store(InMemoryPushConfigStore::new()),
make_handler()
.with_push_config_store(InMemoryPushConfigStore::new())
.with_capabilities(caps(None, Some(false))),
] {
let error = handler
.create_push_config(¶ms, push_config("t1"))
.await
.unwrap_err();
assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
}
}
#[tokio::test]
async fn test_with_push_notifications_defaults_the_capability() {
let params = ServiceParams::new();
let handler = make_handler().with_push_notifications(
InMemoryPushConfigStore::new(),
crate::HttpPushSender::new(None),
);
assert_eq!(handler.capabilities.push_notifications, Some(true));
assert!(
handler
.create_push_config(¶ms, push_config("t1"))
.await
.is_ok()
);
let handler = make_handler()
.with_capabilities(caps(None, Some(false)))
.with_push_notifications(
InMemoryPushConfigStore::new(),
crate::HttpPushSender::new(None),
);
assert_eq!(handler.capabilities.push_notifications, Some(false));
let error = handler
.create_push_config(¶ms, push_config("t1"))
.await
.unwrap_err();
assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
}
#[tokio::test]
async fn test_missing_store_still_reports_push_unsupported() {
let handler = make_handler();
let error = handler
.create_push_config(&ServiceParams::new(), push_config("t1"))
.await
.unwrap_err();
assert_eq!(error.code, error_code::PUSH_NOTIFICATION_NOT_SUPPORTED);
}
#[tokio::test]
async fn test_artifacts_do_not_aggregate_across_tasks_in_one_context() {
let handler = make_handler();
let params = ServiceParams::new();
let context_id = "ctx-shared";
for (task_id, artifact_id) in [("t-first", "a-1"), ("t-second", "a-2")] {
handler
.task_store
.create(Task {
id: task_id.into(),
context_id: context_id.into(),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
artifacts: None,
history: None,
metadata: None,
})
.await
.unwrap();
apply_event_to_task(
handler.task_store.as_ref(),
None,
&StreamResponse::ArtifactUpdate(TaskArtifactUpdateEvent {
task_id: task_id.into(),
context_id: context_id.into(),
artifact: Artifact {
artifact_id: artifact_id.into(),
name: None,
description: None,
parts: vec![Part::text("payload")],
metadata: None,
extensions: None,
},
append: None,
last_chunk: None,
metadata: None,
}),
)
.await
.unwrap();
}
for (task_id, expected) in [("t-first", "a-1"), ("t-second", "a-2")] {
let task = handler
.get_task(
¶ms,
GetTaskRequest {
id: task_id.into(),
history_length: None,
tenant: None,
},
)
.await
.unwrap();
let ids: Vec<&str> = task
.artifacts
.as_deref()
.unwrap_or_default()
.iter()
.map(|a| a.artifact_id.as_str())
.collect();
assert_eq!(
ids,
vec![expected],
"{task_id} must carry only its own artifact"
);
}
}
fn extended_card(name: &str) -> AgentCard {
AgentCard {
name: name.into(),
description: "extended".into(),
version: "1.0".into(),
supported_interfaces: vec![],
capabilities: AgentCapabilities::default(),
default_input_modes: vec![],
default_output_modes: vec![],
skills: vec![],
provider: None,
documentation_url: None,
icon_url: None,
security_schemes: None,
security_requirements: None,
signatures: None,
}
}
fn extended_request() -> GetExtendedAgentCardRequest {
GetExtendedAgentCardRequest { tenant: None }
}
#[tokio::test]
async fn test_extended_agent_card_is_served_when_configured() {
let handler = make_handler().with_extended_agent_card(extended_card("secret skills"));
let card = handler
.get_extended_agent_card(&ServiceParams::new(), extended_request())
.await
.unwrap();
assert_eq!(card.name, "secret skills");
}
#[tokio::test]
async fn test_configuring_an_extended_card_declares_the_capability() {
let handler = make_handler().with_extended_agent_card(extended_card("x"));
assert_eq!(handler.capabilities.extended_agent_card, Some(true));
}
#[tokio::test]
async fn test_unconfigured_extended_card_reports_its_own_code() {
let handler = make_handler().with_capabilities(AgentCapabilities {
streaming: None,
push_notifications: None,
extensions: None,
extended_agent_card: Some(true),
});
let error = handler
.get_extended_agent_card(&ServiceParams::new(), extended_request())
.await
.unwrap_err();
assert_eq!(error.code, error_code::EXTENDED_CARD_NOT_CONFIGURED);
assert_ne!(error.code, error_code::UNSUPPORTED_OPERATION);
}
#[tokio::test]
async fn test_capability_gate_precedes_the_not_configured_check() {
let handler = make_handler()
.with_extended_agent_card(extended_card("x"))
.with_capabilities(AgentCapabilities {
streaming: None,
push_notifications: None,
extensions: None,
extended_agent_card: Some(false),
});
let error = handler
.get_extended_agent_card(&ServiceParams::new(), extended_request())
.await
.unwrap_err();
assert_eq!(error.code, error_code::UNSUPPORTED_OPERATION);
}
#[tokio::test]
async fn test_resolver_sees_the_caller_params() {
let handler = make_handler().with_extended_agent_card_resolver(|params: &ServiceParams| {
let privileged = params.contains_key("x-privileged");
Ok(extended_card(if privileged { "full" } else { "limited" }))
});
let card = handler
.get_extended_agent_card(&ServiceParams::new(), extended_request())
.await
.unwrap();
assert_eq!(card.name, "limited");
let mut params = ServiceParams::new();
params.insert("x-privileged".into(), vec!["yes".into()]);
let card = handler
.get_extended_agent_card(¶ms, extended_request())
.await
.unwrap();
assert_eq!(card.name, "full");
}
}
#[cfg(test)]
mod lag_coverage_tests {
use super::*;
use futures::StreamExt;
#[tokio::test]
async fn test_subscription_lag_is_reported_as_a_sanitized_internal_error() {
let (sender, receiver) = broadcast::channel(1);
for sequence in 0..4 {
let _ = sender.send(ExecutionEvent {
sequence,
result: Ok(StreamResponse::Message(a2a::Message {
message_id: "lag-probe".to_string(),
role: a2a::Role::Agent,
parts: Vec::new(),
context_id: None,
task_id: None,
metadata: None,
extensions: None,
reference_task_ids: None,
})),
});
}
let mut stream = subscription_stream(receiver, None, 0);
let item = stream.next().await.expect("lagged item");
let error = item.expect_err("lag must surface as an error");
assert_eq!(error.code, a2a::error_code::INTERNAL_ERROR);
assert_eq!(error.message, "Internal error");
assert!(stream.next().await.is_none());
}
}