codetether_agent/a2a/

grpc.rs

//! gRPC transport for the A2A protocol.
//!
//! Implements the `A2AService` trait generated by tonic-build from
//! `proto/a2a/v1/a2a.proto`.  Delegates to the same task store and bus
//! that the JSON-RPC server uses, so both transports share state.

use crate::a2a::bridge;
use crate::a2a::proto;
use crate::a2a::proto::a2a_service_server::{A2aService, A2aServiceServer};
use crate::a2a::types::{self as local, TaskState};
use crate::bus::AgentBus;

use dashmap::DashMap;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::broadcast;
use tokio_stream::Stream;
use tonic::{Request, Response, Status};

type StreamResult<T> = Pin<Box<dyn Stream<Item = Result<T, Status>> + Send>>;
/// Shared state backing both JSON-RPC and gRPC transports.
///
/// Stores task records, push-notification configuration, agent metadata, and
/// task-update fan-out channels for the gRPC service layer. The store is safe to
/// share across concurrent request handlers: task and push-configuration maps
/// use lock-free concurrent access through [`DashMap`], while task updates are
/// distributed with a Tokio broadcast channel.
///
/// # Responsibilities
///
/// * Keep the authoritative in-process [`local::Task`] records by task ID.
/// * Track per-task [`local::TaskPushNotificationConfig`] entries used by
///   push-notification subscribers.
/// * Expose the local [`local::AgentCard`] advertised by this server.
/// * Optionally publish task lifecycle events through [`AgentBus`].
/// * Broadcast task status changes to active streaming subscribers.
///
/// # Invariants
///
/// * Keys in `tasks` and `push_configs` are task IDs.
/// * Values sent through `update_tx` use the same task ID namespace as `tasks`.
/// * `card` describes the agent instance served by this store.
pub struct GrpcTaskStore {
    tasks: DashMap<String, local::Task>,
    push_configs: DashMap<String, Vec<local::TaskPushNotificationConfig>>,
    card: local::AgentCard,
    bus: Option<Arc<AgentBus>>,
    /// Broadcast for task updates (task_id, status)
    ///
    /// Sends each observed task status transition to subscribers as a
    /// `(task_id, status)` tuple. Receivers may lag or miss messages according
    /// to Tokio broadcast-channel semantics, so consumers should treat this as
    /// a live update stream rather than durable task history.
    update_tx: broadcast::Sender<(String, local::TaskStatus)>,
}
impl GrpcTaskStore {
    /// Create a new task store with the given agent card.
    ///
    /// Initializes an empty, shareable [`GrpcTaskStore`] for a single local
    /// agent. The returned [`Arc`] is intended to be cloned into gRPC handlers,
    /// JSON-RPC adapters, streaming tasks, and other transport glue that needs
    /// access to the same in-process task state.
    ///
    /// # Parameters
    ///
    /// * `card` - Agent metadata served by this store and returned to clients
    ///   that request the local agent card.
    ///
    /// # Returns
    ///
    /// A reference-counted store with no tasks, no push-notification
    /// configuration, no attached [`AgentBus`], and a task-update broadcast
    /// channel sized for 256 queued status updates per receiver.
    ///
    /// # Side effects
    ///
    /// Allocates the concurrent maps and creates a Tokio broadcast channel. It
    /// does not publish network services, spawn background tasks, or emit any
    /// task updates.
    pub fn new(card: local::AgentCard) -> Arc<Self> {
        let (update_tx, _) = broadcast::channel(256);
        Arc::new(Self {
            tasks: DashMap::new(),
            push_configs: DashMap::new(),
            card,
            bus: None,
            update_tx,
        })
    }

    /// Create with an agent bus attached.
    ///
    /// Constructs a new store using [`Self::new`] and attaches an [`AgentBus`]
    /// before the store is shared. The bus allows gRPC task activity to be
    /// mirrored into the broader agent-event system while preserving the same
    /// task storage and update-stream behavior as [`Self::new`].
    ///
    /// # Parameters
    ///
    /// * `card` - Agent metadata served by this store.
    /// * `bus` - Shared event bus used to publish or coordinate task activity
    ///   outside the gRPC transport.
    ///
    /// # Returns
    ///
    /// A reference-counted store with the supplied bus installed.
    ///
    /// # Panics
    ///
    /// Panics only if the freshly created [`Arc`] is not uniquely owned before
    /// installing the bus. Under normal control flow this cannot happen because
    /// the [`Arc`] has just been returned by [`Self::new`] and has not yet been
    /// cloned.
    pub fn with_bus(card: local::AgentCard, bus: Arc<AgentBus>) -> Arc<Self> {
        let mut store = Self::new(card);
        Arc::get_mut(&mut store)
            .expect("fresh Arc must be uniquely owned")
            .bus = Some(bus);
        store
    }

    /// Insert or update a task.
    ///
    /// Stores the provided task by its ID, replacing any existing task with the
    /// same ID. Before writing the task into the map, this method broadcasts the
    /// task's current status so live subscribers can react to status changes.
    ///
    /// # Parameters
    ///
    /// * `task` - Complete task record to insert or replace. Its `id` field is
    ///   used as the storage key and as the task ID in the broadcast update.
    ///
    /// # Side effects
    ///
    /// Sends a best-effort status update on the broadcast channel and mutates
    /// the in-memory task map. If there are no active receivers, or if receivers
    /// have lagged, the broadcast result is ignored; the task is still stored.
    pub fn upsert_task(&self, task: local::Task) {
        let _ = self.update_tx.send((task.id.clone(), task.status.clone()));
        self.tasks.insert(task.id.clone(), task);
    }

    /// Get a task by id.
    ///
    /// Looks up a task in the in-memory store and returns an owned clone so the
    /// caller can inspect or serialize it without holding a [`DashMap`] guard.
    ///
    /// # Parameters
    ///
    /// * `id` - Task identifier to retrieve.
    ///
    /// # Returns
    ///
    /// `Some(local::Task)` when a task with the given ID exists, or `None` when
    /// the store has no matching task.
    ///
    /// # Side effects
    ///
    /// This method does not mutate store state or emit task updates.
    pub fn get_task(&self, id: &str) -> Option<local::Task> {
        self.tasks.get(id).map(|r| r.value().clone())
    }

    /// Subscribe to task update notifications.
    ///
    /// Creates a new receiver for the live task-status broadcast stream. Each
    /// message contains the task ID and the status observed by [`Self::upsert_task`].
    ///
    /// # Returns
    ///
    /// A Tokio broadcast receiver for `(task_id, status)` tuples. Broadcast
    /// receivers are live streams, not durable logs; slow receivers may observe
    /// lag errors according to Tokio broadcast-channel semantics.
    ///
    /// # Side effects
    ///
    /// Registers a new receiver with the broadcast channel. It does not read,
    /// mutate, or replay existing tasks.
    pub fn subscribe_updates(&self) -> broadcast::Receiver<(String, local::TaskStatus)> {
        self.update_tx.subscribe()
    }

    /// Build the tonic service layer from this store.
    ///
    /// Wraps the shared task store in an [`A2aServiceImpl`] and returns the
    /// generated tonic [`A2aServiceServer`] ready to be mounted into a gRPC
    /// server.
    ///
    /// # Returns
    ///
    /// A tonic service server that handles A2A gRPC requests using this store as
    /// its backing state.
    ///
    /// # Side effects
    ///
    /// Consumes one [`Arc`] handle to the store and moves it into the service
    /// implementation. This method does not start listening on a socket; the
    /// caller is responsible for adding the returned service to a tonic server.
    pub fn into_service(self: Arc<Self>) -> A2aServiceServer<A2aServiceImpl> {
        A2aServiceServer::new(A2aServiceImpl { store: self })
    }
}
/// Tonic service implementation.
pub struct A2aServiceImpl {
    store: Arc<GrpcTaskStore>,
}

#[tonic::async_trait]
impl A2aService for A2aServiceImpl {
    // ── SendMessage ─────────────────────────────────────────────────────

    async fn send_message(
        &self,
        request: Request<proto::SendMessageRequest>,
    ) -> Result<Response<proto::SendMessageResponse>, Status> {
        let req = request.into_inner();
        let msg = req
            .request
            .ok_or_else(|| Status::invalid_argument("missing message"))?;
        let local_msg = bridge::proto_message_to_local(&msg);

        // Create / route task
        let task_id = local_msg
            .task_id
            .clone()
            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());

        let task = local::Task {
            id: task_id.clone(),
            context_id: local_msg.context_id.clone(),
            status: local::TaskStatus {
                state: TaskState::Submitted,
                message: Some(local_msg.clone()),
                timestamp: Some(chrono::Utc::now().to_rfc3339()),
            },
            artifacts: vec![],
            history: vec![local_msg],
            metadata: Default::default(),
        };
        self.store.upsert_task(task.clone());

        // Notify bus
        if let Some(ref bus) = self.store.bus {
            let handle = bus.handle("grpc-server");
            handle.send_task_update(&task_id, TaskState::Submitted, None);
        }

        let proto_task = bridge::local_task_to_proto(&task);
        Ok(Response::new(proto::SendMessageResponse {
            payload: Some(proto::send_message_response::Payload::Task(proto_task)),
        }))
    }

    // ── SendStreamingMessage ────────────────────────────────────────────

    type SendStreamingMessageStream = StreamResult<proto::StreamResponse>;

    async fn send_streaming_message(
        &self,
        request: Request<proto::SendMessageRequest>,
    ) -> Result<Response<Self::SendStreamingMessageStream>, Status> {
        let req = request.into_inner();
        let msg = req
            .request
            .ok_or_else(|| Status::invalid_argument("missing message"))?;
        let local_msg = bridge::proto_message_to_local(&msg);

        let task_id = local_msg
            .task_id
            .clone()
            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());

        let task = local::Task {
            id: task_id.clone(),
            context_id: local_msg.context_id.clone(),
            status: local::TaskStatus {
                state: TaskState::Submitted,
                message: Some(local_msg.clone()),
                timestamp: Some(chrono::Utc::now().to_rfc3339()),
            },
            artifacts: vec![],
            history: vec![local_msg],
            metadata: Default::default(),
        };
        self.store.upsert_task(task.clone());

        // Return the initial task then listen for updates
        let proto_task = bridge::local_task_to_proto(&task);
        let mut rx = self.store.subscribe_updates();
        let tid = task_id.clone();

        let stream = async_stream::try_stream! {
            // First frame: the task itself
            yield proto::StreamResponse {
                payload: Some(proto::stream_response::Payload::Task(proto_task)),
            };

            // Subsequent frames: status updates for this task
            loop {
                match rx.recv().await {
                    Ok((id, status)) if id == tid => {
                        let proto_status = bridge::local_task_status_to_proto(&status);
                        let is_terminal = status.state.is_terminal();
                        yield proto::StreamResponse {
                            payload: Some(proto::stream_response::Payload::StatusUpdate(
                                proto::TaskStatusUpdateEvent {
                                    task_id: tid.clone(),
                                    context_id: String::new(),
                                    status: Some(proto_status),
                                    r#final: is_terminal,
                                    metadata: None,
                                },
                            )),
                        };
                        if is_terminal {
                            break;
                        }
                    }
                    Ok(_) => continue,
                    Err(broadcast::error::RecvError::Lagged(_)) => continue,
                    Err(broadcast::error::RecvError::Closed) => break,
                }
            }
        };

        Ok(Response::new(
            Box::pin(stream) as Self::SendStreamingMessageStream
        ))
    }

    // ── GetTask ─────────────────────────────────────────────────────────

    async fn get_task(
        &self,
        request: Request<proto::GetTaskRequest>,
    ) -> Result<Response<proto::Task>, Status> {
        let req = request.into_inner();
        let task_id = req.name.strip_prefix("tasks/").unwrap_or(&req.name);
        let task = self
            .store
            .get_task(task_id)
            .ok_or_else(|| Status::not_found(format!("task {task_id} not found")))?;
        Ok(Response::new(bridge::local_task_to_proto(&task)))
    }

    // ── CancelTask ──────────────────────────────────────────────────────

    async fn cancel_task(
        &self,
        request: Request<proto::CancelTaskRequest>,
    ) -> Result<Response<proto::Task>, Status> {
        let req = request.into_inner();
        let task_id = req.name.strip_prefix("tasks/").unwrap_or(&req.name);

        let mut task = self
            .store
            .tasks
            .get_mut(task_id)
            .ok_or_else(|| Status::not_found(format!("task {task_id} not found")))?;

        if task.status.state.is_terminal() {
            return Err(Status::failed_precondition(
                "task already in terminal state",
            ));
        }

        task.status = local::TaskStatus {
            state: TaskState::Cancelled,
            message: None,
            timestamp: Some(chrono::Utc::now().to_rfc3339()),
        };
        let snapshot = task.clone();
        drop(task);

        let _ = self
            .store
            .update_tx
            .send((task_id.to_string(), snapshot.status.clone()));

        Ok(Response::new(bridge::local_task_to_proto(&snapshot)))
    }

    // ── TaskSubscription ────────────────────────────────────────────────

    type TaskSubscriptionStream = StreamResult<proto::StreamResponse>;

    async fn task_subscription(
        &self,
        request: Request<proto::TaskSubscriptionRequest>,
    ) -> Result<Response<Self::TaskSubscriptionStream>, Status> {
        let req = request.into_inner();
        let task_id = req
            .name
            .strip_prefix("tasks/")
            .unwrap_or(&req.name)
            .to_string();

        let task = self
            .store
            .get_task(&task_id)
            .ok_or_else(|| Status::not_found(format!("task {task_id} not found")))?;

        let proto_task = bridge::local_task_to_proto(&task);
        let mut rx = self.store.subscribe_updates();
        let tid = task_id.clone();

        // If already terminal, return just the task and close
        if task.status.state.is_terminal() {
            let stream = async_stream::try_stream! {
                yield proto::StreamResponse {
                    payload: Some(proto::stream_response::Payload::Task(proto_task)),
                };
            };
            return Ok(Response::new(
                Box::pin(stream) as Self::TaskSubscriptionStream
            ));
        }

        let stream = async_stream::try_stream! {
            yield proto::StreamResponse {
                payload: Some(proto::stream_response::Payload::Task(proto_task)),
            };

            loop {
                match rx.recv().await {
                    Ok((id, status)) if id == tid => {
                        let proto_status = bridge::local_task_status_to_proto(&status);
                        let is_terminal = status.state.is_terminal();
                        yield proto::StreamResponse {
                            payload: Some(proto::stream_response::Payload::StatusUpdate(
                                proto::TaskStatusUpdateEvent {
                                    task_id: tid.clone(),
                                    context_id: String::new(),
                                    status: Some(proto_status),
                                    r#final: is_terminal,
                                    metadata: None,
                                },
                            )),
                        };
                        if is_terminal { break; }
                    }
                    Ok(_) => continue,
                    Err(broadcast::error::RecvError::Lagged(_)) => continue,
                    Err(broadcast::error::RecvError::Closed) => break,
                }
            }
        };

        Ok(Response::new(
            Box::pin(stream) as Self::TaskSubscriptionStream
        ))
    }

    // ── Push notification config CRUD ───────────────────────────────────

    async fn create_task_push_notification_config(
        &self,
        request: Request<proto::CreateTaskPushNotificationConfigRequest>,
    ) -> Result<Response<proto::TaskPushNotificationConfig>, Status> {
        let req = request.into_inner();
        let task_id = req.parent.strip_prefix("tasks/").unwrap_or(&req.parent);

        if self.store.get_task(task_id).is_none() {
            return Err(Status::not_found(format!("task {task_id} not found")));
        }

        let config = req
            .config
            .ok_or_else(|| Status::invalid_argument("missing config"))?;
        let pnc = config.push_notification_config.as_ref();

        let local_config = local::TaskPushNotificationConfig {
            id: task_id.to_string(),
            push_notification_config: local::PushNotificationConfig {
                url: pnc.map(|c| c.url.clone()).unwrap_or_default(),
                token: pnc.and_then(|c| {
                    if c.token.is_empty() {
                        None
                    } else {
                        Some(c.token.clone())
                    }
                }),
                id: pnc.and_then(|c| {
                    if c.id.is_empty() {
                        None
                    } else {
                        Some(c.id.clone())
                    }
                }),
            },
        };

        self.store
            .push_configs
            .entry(task_id.to_string())
            .or_default()
            .push(local_config);

        Ok(Response::new(config))
    }

    async fn get_task_push_notification_config(
        &self,
        request: Request<proto::GetTaskPushNotificationConfigRequest>,
    ) -> Result<Response<proto::TaskPushNotificationConfig>, Status> {
        let req = request.into_inner();
        // name format: tasks/{task_id}/pushNotificationConfigs/{config_id}
        let parts: Vec<&str> = req.name.split('/').collect();
        if parts.len() < 4 {
            return Err(Status::invalid_argument("invalid name format"));
        }
        let task_id = parts[1];
        let config_id = parts[3];

        let configs = self
            .store
            .push_configs
            .get(task_id)
            .ok_or_else(|| Status::not_found("no configs for task"))?;

        let _found = configs
            .iter()
            .find(|c| c.push_notification_config.id.as_deref() == Some(config_id))
            .ok_or_else(|| Status::not_found("config not found"))?;

        Ok(Response::new(proto::TaskPushNotificationConfig {
            name: req.name,
            push_notification_config: None, // simplified
        }))
    }

    async fn list_task_push_notification_config(
        &self,
        request: Request<proto::ListTaskPushNotificationConfigRequest>,
    ) -> Result<Response<proto::ListTaskPushNotificationConfigResponse>, Status> {
        let req = request.into_inner();
        let task_id = req.parent.strip_prefix("tasks/").unwrap_or(&req.parent);

        let configs: Vec<proto::TaskPushNotificationConfig> = self
            .store
            .push_configs
            .get(task_id)
            .map(|cs| {
                cs.iter()
                    .map(|c| proto::TaskPushNotificationConfig {
                        name: format!(
                            "tasks/{}/pushNotificationConfigs/{}",
                            task_id,
                            c.push_notification_config
                                .id
                                .as_deref()
                                .unwrap_or("default")
                        ),
                        push_notification_config: None,
                    })
                    .collect()
            })
            .unwrap_or_default();

        Ok(Response::new(
            proto::ListTaskPushNotificationConfigResponse {
                configs,
                next_page_token: String::new(),
            },
        ))
    }

    async fn delete_task_push_notification_config(
        &self,
        request: Request<proto::DeleteTaskPushNotificationConfigRequest>,
    ) -> Result<Response<()>, Status> {
        let req = request.into_inner();
        let parts: Vec<&str> = req.name.split('/').collect();
        if parts.len() < 4 {
            return Err(Status::invalid_argument("invalid name format"));
        }
        let task_id = parts[1];
        let config_id = parts[3];

        if let Some(mut configs) = self.store.push_configs.get_mut(task_id) {
            configs.retain(|c| c.push_notification_config.id.as_deref() != Some(config_id));
        }

        Ok(Response::new(()))
    }

    // ── GetAgentCard ────────────────────────────────────────────────────

    async fn get_agent_card(
        &self,
        _request: Request<proto::GetAgentCardRequest>,
    ) -> Result<Response<proto::AgentCard>, Status> {
        Ok(Response::new(bridge::local_card_to_proto(&self.store.card)))
    }
}