From 031bd76ecb0ae1fb07c328aca74337ee28f994f5 Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Thu, 18 Jun 2026 20:43:33 -0400 Subject: [PATCH] feat(a2a-nats): add client message_stream per-operation method Streaming responses need both the bootstrap RPC and the JetStream subscription wired together on the agent's per-task event subject; combining them on the handle keeps callers from threading the prefix + req_id + consumer config by hand for every stream. Signed-off-by: Yordis Prieto --- .../crates/a2a-nats/src/client/handle.rs | 121 +++++++++++++++++- 1 file changed, 120 insertions(+), 1 deletion(-) diff --git a/rsworkspace/crates/a2a-nats/src/client/handle.rs b/rsworkspace/crates/a2a-nats/src/client/handle.rs index c6ed813f4..4a3b3f8f7 100644 --- a/rsworkspace/crates/a2a-nats/src/client/handle.rs +++ b/rsworkspace/crates/a2a-nats/src/client/handle.rs @@ -15,17 +15,20 @@ use a2a::types::{ }; use a2a_identity_types::MintedUserJwt; use trogon_nats::RequestClient; +use trogon_nats::jetstream::{JetStreamCreateConsumer, JetStreamGetStream, JsAck, JsMessageOf, JsMessageRef}; use crate::a2a_prefix::A2aPrefix; use crate::agent_id::A2aAgentId; use crate::constants::{DEFAULT_OPERATION_TIMEOUT, MIN_TIMEOUT_SECS}; use crate::gateway_ingress::gateway_ingress_subject_from_agent_subject; use crate::nats::subjects::agents::{ - AgentCardSubject, MessageSendSubject, TasksCancelSubject, TasksGetSubject, TasksListSubject, + AgentCardSubject, MessageSendSubject, MessageStreamSubject, TasksCancelSubject, TasksGetSubject, TasksListSubject, }; use crate::req_id::ReqId; use super::error::ClientError; +use super::event_stream::TypedEventStream; +use super::streaming::{StreamingRequest, send_streaming}; use super::unary::send_unary; /// Whether `A2aClient` publishes to `{prefix}.agents.{agent_id}.…` (talking to the @@ -198,6 +201,37 @@ where } } +impl A2aClient +where + N: RequestClient, + J: JetStreamGetStream, + JsMessageOf: JsMessageRef + JsAck + Send + 'static, + ::Stream: Send + 'static, + <::Stream as JetStreamCreateConsumer>::Consumer: Send + 'static, + <<::Stream as JetStreamCreateConsumer>::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Messages: Send + 'static, + <<::Stream as JetStreamCreateConsumer>::Consumer as trogon_nats::jetstream::JetStreamConsumer>::MessagesError: std::fmt::Display + Send + 'static, + <<::Stream as JetStreamCreateConsumer>::Consumer as trogon_nats::jetstream::JetStreamConsumer>::StreamError: std::fmt::Display + Send + 'static, +{ + pub async fn message_stream( + &self, + req: &SendMessageRequest, + ) -> Result<(SendMessageResponse, TypedEventStream), ClientError> { + let subject = self.outbound_rpc_subject(MessageStreamSubject::new(self.prefix(), &self.agent_id).to_string())?; + let req_id = ReqId::new(); + let ctx = StreamingRequest { + nats: &self.nats, + js: &self.js, + subject: subject.as_str(), + method: "message/stream", + req_id: &req_id, + prefix: self.prefix(), + op_timeout: self.operation_timeout(), + gateway_caller_jwt: self.gateway_caller_jwt(), + }; + send_streaming(ctx, req).await + } +} + #[cfg(test)] mod tests { use super::*; @@ -683,4 +717,89 @@ mod tests { )); } } + + mod message_stream_op { + use a2a::types::{Message, Role, SendMessageRequest, SendMessageResponse, Task, TaskState, TaskStatus}; + use bytes::Bytes; + use trogon_nats::AdvancedMockNatsClient; + use trogon_nats::jetstream::mocks::{MockJetStreamConsumer, MockJetStreamConsumerFactory}; + + use super::*; + + fn send_message_request() -> SendMessageRequest { + SendMessageRequest { + message: Message { + message_id: "m-1".to_string(), + role: Role::User, + parts: vec![], + context_id: None, + task_id: None, + reference_task_ids: None, + extensions: None, + metadata: None, + }, + configuration: None, + metadata: None, + tenant: None, + } + } + + fn bootstrap_response(task_id: &str) -> Bytes { + let task = Task { + id: task_id.to_string(), + context_id: String::new(), + status: TaskStatus { + state: TaskState::Working, + message: None, + timestamp: None, + }, + artifacts: None, + history: None, + metadata: None, + }; + let response = SendMessageResponse::Task(task); + let json = serde_json::json!({"jsonrpc":"2.0","id":"any","result":response}); + serde_json::to_vec(&json).unwrap().into() + } + + #[tokio::test] + async fn message_stream_targets_agent_subject_by_default() { + let nats = AdvancedMockNatsClient::new(); + nats.set_response("a2a.agents.test-agent.message.stream", bootstrap_response("t-1")); + let js = MockJetStreamConsumerFactory::new(); + let (consumer, _tx) = MockJetStreamConsumer::new(); + js.add_consumer(consumer); + let client = A2aClient::new(prefix(), agent_id(), nats, js); + let (envelope, _stream) = client.message_stream(&send_message_request()).await.unwrap(); + assert!(matches!(envelope, SendMessageResponse::Task(_))); + } + + #[tokio::test] + async fn message_stream_targets_gateway_subject_under_gateway_routing() { + let nats = AdvancedMockNatsClient::new(); + nats.set_response("a2a.gateway.test-agent.message.stream", bootstrap_response("t-gw")); + let js = MockJetStreamConsumerFactory::new(); + let (consumer, _tx) = MockJetStreamConsumer::new(); + js.add_consumer(consumer); + let jwt = + MintedUserJwt::new("eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjk5OTk5OTk5OTl9.signature").unwrap(); + let client = A2aClient::new(prefix(), agent_id(), nats, js).routing_via_gateway_ingress(jwt); + let (envelope, _stream) = client.message_stream(&send_message_request()).await.unwrap(); + assert!(matches!(envelope, SendMessageResponse::Task(_))); + } + + #[tokio::test] + async fn message_stream_propagates_transport_errors() { + let nats = AdvancedMockNatsClient::new(); + nats.fail_next_request(); + let js = MockJetStreamConsumerFactory::new(); + let (consumer, _tx) = MockJetStreamConsumer::new(); + js.add_consumer(consumer); + let client = A2aClient::new(prefix(), agent_id(), nats, js); + assert!(matches!( + client.message_stream(&send_message_request()).await, + Err(ClientError::Transport(_)) + )); + } + } }