litellm_rs/core/a2a/
message.rs

1//! A2A Message Types
2//!
3//! JSON-RPC 2.0 message types for A2A protocol communication.
4
5use std::sync::atomic::{AtomicU64, Ordering};
6
7use super::error::A2AError;
8use chrono::{DateTime, Utc};
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11
12/// Monotonically increasing counter for unique JSON-RPC request IDs.
13static REQUEST_ID: AtomicU64 = AtomicU64::new(1);
14
15/// Return the next unique JSON-RPC request ID.
16fn next_id() -> Value {
17    Value::Number(REQUEST_ID.fetch_add(1, Ordering::Relaxed).into())
18}
19
20/// JSON-RPC 2.0 version constant
21pub const JSONRPC_VERSION: &str = "2.0";
22
23/// A2A Message (JSON-RPC 2.0 request)
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct A2AMessage {
26    /// JSON-RPC version (must be "2.0")
27    pub jsonrpc: String,
28
29    /// Request method
30    pub method: String,
31
32    /// Request parameters
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub params: Option<MessageParams>,
35
36    /// Request ID
37    pub id: Value,
38}
39
40impl A2AMessage {
41    /// Create a new send message request
42    pub fn send(message: impl Into<String>) -> Self {
43        Self {
44            jsonrpc: JSONRPC_VERSION.to_string(),
45            method: "message/send".to_string(),
46            params: Some(MessageParams {
47                message: Message {
48                    role: "user".to_string(),
49                    parts: vec![MessagePart::Text {
50                        text: message.into(),
51                    }],
52                },
53                configuration: None,
54            }),
55            id: next_id(),
56        }
57    }
58
59    /// Create a task status request
60    pub fn get_task(task_id: impl Into<String>) -> Self {
61        Self {
62            jsonrpc: JSONRPC_VERSION.to_string(),
63            method: "tasks/get".to_string(),
64            params: Some(MessageParams {
65                message: Message {
66                    role: "system".to_string(),
67                    parts: vec![MessagePart::Text {
68                        text: task_id.into(),
69                    }],
70                },
71                configuration: None,
72            }),
73            id: next_id(),
74        }
75    }
76
77    /// Create a cancel task request
78    pub fn cancel_task(task_id: impl Into<String>) -> Self {
79        Self {
80            jsonrpc: JSONRPC_VERSION.to_string(),
81            method: "tasks/cancel".to_string(),
82            params: Some(MessageParams {
83                message: Message {
84                    role: "system".to_string(),
85                    parts: vec![MessagePart::Text {
86                        text: task_id.into(),
87                    }],
88                },
89                configuration: None,
90            }),
91            id: next_id(),
92        }
93    }
94
95    /// Set request ID
96    pub fn with_id(mut self, id: impl Into<Value>) -> Self {
97        self.id = id.into();
98        self
99    }
100
101    /// Add configuration
102    pub fn with_config(mut self, config: MessageConfiguration) -> Self {
103        if let Some(ref mut params) = self.params {
104            params.configuration = Some(config);
105        }
106        self
107    }
108}
109
110/// Message parameters
111#[derive(Debug, Clone, Serialize, Deserialize)]
112pub struct MessageParams {
113    /// The message to send
114    pub message: Message,
115
116    /// Optional configuration
117    #[serde(skip_serializing_if = "Option::is_none")]
118    pub configuration: Option<MessageConfiguration>,
119}
120
121/// A message in the conversation
122#[derive(Debug, Clone, Serialize, Deserialize)]
123pub struct Message {
124    /// Message role (user, assistant, system)
125    pub role: String,
126
127    /// Message content parts
128    pub parts: Vec<MessagePart>,
129}
130
131impl Message {
132    /// Create a user message with text
133    pub fn user(text: impl Into<String>) -> Self {
134        Self {
135            role: "user".to_string(),
136            parts: vec![MessagePart::Text { text: text.into() }],
137        }
138    }
139
140    /// Create an assistant message with text
141    pub fn assistant(text: impl Into<String>) -> Self {
142        Self {
143            role: "assistant".to_string(),
144            parts: vec![MessagePart::Text { text: text.into() }],
145        }
146    }
147
148    /// Create a system message with text
149    pub fn system(text: impl Into<String>) -> Self {
150        Self {
151            role: "system".to_string(),
152            parts: vec![MessagePart::Text { text: text.into() }],
153        }
154    }
155}
156
157/// Message content part
158#[derive(Debug, Clone, Serialize, Deserialize)]
159#[serde(tag = "type", rename_all = "lowercase")]
160pub enum MessagePart {
161    /// Text content
162    Text { text: String },
163
164    /// File/data content
165    File {
166        #[serde(rename = "mimeType")]
167        mime_type: String,
168        data: String,
169        #[serde(skip_serializing_if = "Option::is_none")]
170        name: Option<String>,
171    },
172
173    /// Tool use request
174    ToolUse {
175        #[serde(rename = "toolUseId")]
176        tool_use_id: String,
177        name: String,
178        input: Value,
179    },
180
181    /// Tool result
182    ToolResult {
183        #[serde(rename = "toolUseId")]
184        tool_use_id: String,
185        content: String,
186        #[serde(rename = "isError", default)]
187        is_error: bool,
188    },
189}
190
191impl MessagePart {
192    /// Create text part
193    pub fn text(text: impl Into<String>) -> Self {
194        Self::Text { text: text.into() }
195    }
196
197    /// Create file part
198    pub fn file(mime_type: impl Into<String>, data: impl Into<String>) -> Self {
199        Self::File {
200            mime_type: mime_type.into(),
201            data: data.into(),
202            name: None,
203        }
204    }
205
206    /// Create tool use part
207    pub fn tool_use(id: impl Into<String>, name: impl Into<String>, input: Value) -> Self {
208        Self::ToolUse {
209            tool_use_id: id.into(),
210            name: name.into(),
211            input,
212        }
213    }
214
215    /// Create tool result part
216    pub fn tool_result(id: impl Into<String>, content: impl Into<String>) -> Self {
217        Self::ToolResult {
218            tool_use_id: id.into(),
219            content: content.into(),
220            is_error: false,
221        }
222    }
223}
224
225/// Message configuration
226#[derive(Debug, Clone, Serialize, Deserialize, Default)]
227pub struct MessageConfiguration {
228    /// Whether to wait for completion
229    #[serde(
230        rename = "acceptedOutputModes",
231        skip_serializing_if = "Option::is_none"
232    )]
233    pub accepted_output_modes: Option<Vec<String>>,
234
235    /// Enable streaming
236    #[serde(skip_serializing_if = "Option::is_none")]
237    pub streaming: Option<bool>,
238
239    /// Push notification URL
240    #[serde(
241        rename = "pushNotificationConfig",
242        skip_serializing_if = "Option::is_none"
243    )]
244    pub push_notification_config: Option<PushNotificationConfig>,
245
246    /// History behavior
247    #[serde(rename = "historyLength", skip_serializing_if = "Option::is_none")]
248    pub history_length: Option<u32>,
249}
250
251/// Push notification configuration
252#[derive(Debug, Clone, Serialize, Deserialize)]
253pub struct PushNotificationConfig {
254    /// Callback URL
255    pub url: String,
256
257    /// Authentication token
258    #[serde(skip_serializing_if = "Option::is_none")]
259    pub token: Option<String>,
260}
261
262/// A2A Response (JSON-RPC 2.0 response)
263#[derive(Debug, Clone, Serialize, Deserialize)]
264pub struct A2AResponse {
265    /// JSON-RPC version
266    pub jsonrpc: String,
267
268    /// Response result
269    #[serde(skip_serializing_if = "Option::is_none")]
270    pub result: Option<TaskResult>,
271
272    /// Response error
273    #[serde(skip_serializing_if = "Option::is_none")]
274    pub error: Option<A2AResponseError>,
275
276    /// Request ID
277    #[serde(skip_serializing_if = "Option::is_none")]
278    pub id: Option<Value>,
279}
280
281impl A2AResponse {
282    /// Check if response is successful
283    pub fn is_success(&self) -> bool {
284        self.result.is_some() && self.error.is_none()
285    }
286
287    /// Check if response is an error
288    pub fn is_error(&self) -> bool {
289        self.error.is_some()
290    }
291
292    /// Get the task result
293    pub fn get_result(&self) -> Option<&TaskResult> {
294        self.result.as_ref()
295    }
296
297    /// Get the error
298    pub fn get_error(&self) -> Option<&A2AResponseError> {
299        self.error.as_ref()
300    }
301}
302
303/// Task result from A2A response
304#[derive(Debug, Clone, Serialize, Deserialize)]
305pub struct TaskResult {
306    /// Task ID
307    pub id: String,
308
309    /// Task state
310    pub status: TaskStatus,
311
312    /// Task artifacts (outputs)
313    #[serde(default, skip_serializing_if = "Vec::is_empty")]
314    pub artifacts: Vec<Artifact>,
315
316    /// History of messages
317    #[serde(default, skip_serializing_if = "Vec::is_empty")]
318    pub history: Vec<Message>,
319}
320
321/// Task status
322#[derive(Debug, Clone, Serialize, Deserialize)]
323pub struct TaskStatus {
324    /// Current state
325    pub state: TaskState,
326
327    /// Status message
328    #[serde(skip_serializing_if = "Option::is_none")]
329    pub message: Option<Message>,
330
331    /// Timestamp
332    #[serde(skip_serializing_if = "Option::is_none")]
333    pub timestamp: Option<DateTime<Utc>>,
334}
335
336/// Task state enumeration
337#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
338#[serde(rename_all = "lowercase")]
339pub enum TaskState {
340    /// Task is pending
341    #[default]
342    Pending,
343
344    /// Task is running
345    Running,
346
347    /// Task requires input
348    #[serde(rename = "input-required")]
349    InputRequired,
350
351    /// Task completed successfully
352    Completed,
353
354    /// Task failed
355    Failed,
356
357    /// Task was cancelled
358    Cancelled,
359}
360
361impl TaskState {
362    /// Check if task is complete (finished in any state)
363    pub fn is_terminal(&self) -> bool {
364        matches!(
365            self,
366            TaskState::Completed | TaskState::Failed | TaskState::Cancelled
367        )
368    }
369
370    /// Check if task is still running
371    pub fn is_running(&self) -> bool {
372        matches!(self, TaskState::Running)
373    }
374
375    /// Check if task succeeded
376    pub fn is_success(&self) -> bool {
377        matches!(self, TaskState::Completed)
378    }
379}
380
381/// Task artifact (output)
382#[derive(Debug, Clone, Serialize, Deserialize)]
383pub struct Artifact {
384    /// Artifact name
385    #[serde(skip_serializing_if = "Option::is_none")]
386    pub name: Option<String>,
387
388    /// Artifact content parts
389    pub parts: Vec<MessagePart>,
390
391    /// Index in artifact list
392    #[serde(skip_serializing_if = "Option::is_none")]
393    pub index: Option<u32>,
394}
395
396/// A2A error response
397#[derive(Debug, Clone, Serialize, Deserialize)]
398pub struct A2AResponseError {
399    /// Error code
400    pub code: i32,
401
402    /// Error message
403    pub message: String,
404
405    /// Additional error data
406    #[serde(skip_serializing_if = "Option::is_none")]
407    pub data: Option<Value>,
408}
409
410impl A2AResponseError {
411    /// Create a new error
412    pub fn new(code: i32, message: impl Into<String>) -> Self {
413        Self {
414            code,
415            message: message.into(),
416            data: None,
417        }
418    }
419
420    /// Parse error (-32700)
421    pub fn parse_error() -> Self {
422        Self::new(-32700, "Parse error")
423    }
424
425    /// Invalid request (-32600)
426    pub fn invalid_request() -> Self {
427        Self::new(-32600, "Invalid Request")
428    }
429
430    /// Method not found (-32601)
431    pub fn method_not_found() -> Self {
432        Self::new(-32601, "Method not found")
433    }
434
435    /// Invalid params (-32602)
436    pub fn invalid_params() -> Self {
437        Self::new(-32602, "Invalid params")
438    }
439
440    /// Internal error (-32603)
441    pub fn internal_error() -> Self {
442        Self::new(-32603, "Internal error")
443    }
444
445    /// Task not found (-32001)
446    pub fn task_not_found() -> Self {
447        Self::new(-32001, "Task not found")
448    }
449
450    /// Task cancelled (-32002)
451    pub fn task_cancelled() -> Self {
452        Self::new(-32002, "Task cancelled")
453    }
454
455    /// Build protocol error payload from a typed A2A error.
456    pub fn from_a2a_error(error: &A2AError) -> Self {
457        use crate::utils::error::canonical::CanonicalError;
458
459        let code = match error {
460            A2AError::AgentNotFound { .. } | A2AError::TaskNotFound { .. } => -32001,
461            A2AError::AgentBusy { .. } => -32002,
462            A2AError::AgentAlreadyExists { .. } => -32003,
463            A2AError::AuthenticationError { .. } => -32004,
464            A2AError::RateLimitExceeded { .. } => -32029,
465            A2AError::Timeout { .. } => -32008,
466            A2AError::ConnectionError { .. } => -32010,
467            A2AError::ProtocolError { .. } | A2AError::InvalidRequest { .. } => -32600,
468            A2AError::UnsupportedProvider { .. } => -32601,
469            A2AError::ContentBlocked { .. } => -32602,
470            A2AError::TaskFailed { .. }
471            | A2AError::ConfigurationError { .. }
472            | A2AError::SerializationError { .. } => -32603,
473        };
474
475        let mut response = Self::new(code, error.to_string());
476        response.data = Some(serde_json::json!({
477            "canonical_code": error.canonical_code().as_str(),
478            "retryable": error.canonical_retryable(),
479        }));
480        response
481    }
482}
483
484#[cfg(test)]
485mod tests {
486    use super::*;
487
488    #[test]
489    fn test_a2a_message_send() {
490        let msg = A2AMessage::send("Hello, agent!");
491        assert_eq!(msg.method, "message/send");
492        assert!(msg.params.is_some());
493    }
494
495    #[test]
496    fn test_message_creation() {
497        let user_msg = Message::user("Hello");
498        assert_eq!(user_msg.role, "user");
499        assert_eq!(user_msg.parts.len(), 1);
500
501        let asst_msg = Message::assistant("Hi there");
502        assert_eq!(asst_msg.role, "assistant");
503    }
504
505    #[test]
506    fn test_message_part_text() {
507        let part = MessagePart::text("Hello");
508        match part {
509            MessagePart::Text { text } => assert_eq!(text, "Hello"),
510            _ => panic!("Expected text part"),
511        }
512    }
513
514    #[test]
515    fn test_message_part_file() {
516        let part = MessagePart::file("image/png", "base64data");
517        match part {
518            MessagePart::File {
519                mime_type, data, ..
520            } => {
521                assert_eq!(mime_type, "image/png");
522                assert_eq!(data, "base64data");
523            }
524            _ => panic!("Expected file part"),
525        }
526    }
527
528    #[test]
529    fn test_message_part_tool_use() {
530        let part = MessagePart::tool_use("id-1", "search", serde_json::json!({"query": "test"}));
531        match part {
532            MessagePart::ToolUse {
533                tool_use_id,
534                name,
535                input,
536            } => {
537                assert_eq!(tool_use_id, "id-1");
538                assert_eq!(name, "search");
539                assert_eq!(input["query"], "test");
540            }
541            _ => panic!("Expected tool use part"),
542        }
543    }
544
545    #[test]
546    fn test_task_state_terminal() {
547        assert!(TaskState::Completed.is_terminal());
548        assert!(TaskState::Failed.is_terminal());
549        assert!(TaskState::Cancelled.is_terminal());
550        assert!(!TaskState::Running.is_terminal());
551        assert!(!TaskState::Pending.is_terminal());
552    }
553
554    #[test]
555    fn test_task_state_success() {
556        assert!(TaskState::Completed.is_success());
557        assert!(!TaskState::Failed.is_success());
558    }
559
560    #[test]
561    fn test_a2a_response_error_codes() {
562        assert_eq!(A2AResponseError::parse_error().code, -32700);
563        assert_eq!(A2AResponseError::invalid_request().code, -32600);
564        assert_eq!(A2AResponseError::method_not_found().code, -32601);
565        assert_eq!(A2AResponseError::task_not_found().code, -32001);
566    }
567
568    #[test]
569    fn test_a2a_response_error_from_a2a_error_includes_canonical_data() {
570        let error = A2AError::RateLimitExceeded {
571            agent_name: "agent-a".to_string(),
572            retry_after_ms: Some(500),
573        };
574
575        let response_error = A2AResponseError::from_a2a_error(&error);
576        assert_eq!(response_error.code, -32029);
577
578        let data = response_error.data.expect("canonical data should exist");
579        assert_eq!(data["canonical_code"], "RATE_LIMITED");
580        assert_eq!(data["retryable"], true);
581    }
582
583    #[test]
584    fn test_message_serialization() {
585        let msg = A2AMessage::send("Test message");
586        let json = serde_json::to_string(&msg).unwrap();
587        assert!(json.contains("message/send"));
588        assert!(json.contains("Test message"));
589    }
590
591    #[test]
592    fn test_response_deserialization() {
593        let json = r#"{
594            "jsonrpc": "2.0",
595            "result": {
596                "id": "task-123",
597                "status": {
598                    "state": "completed"
599                },
600                "artifacts": []
601            },
602            "id": 1
603        }"#;
604
605        let response: A2AResponse = serde_json::from_str(json).unwrap();
606        assert!(response.is_success());
607        assert_eq!(response.result.unwrap().id, "task-123");
608    }
609
610    #[test]
611    fn test_a2a_message_with_config() {
612        let config = MessageConfiguration {
613            streaming: Some(true),
614            ..Default::default()
615        };
616
617        let msg = A2AMessage::send("Hello").with_config(config);
618        let params = msg.params.unwrap();
619        assert!(params.configuration.unwrap().streaming.unwrap());
620    }
621
622    #[test]
623    fn test_request_ids_are_unique() {
624        let msg1 = A2AMessage::send("first");
625        let msg2 = A2AMessage::send("second");
626        let msg3 = A2AMessage::get_task("task-1");
627        let msg4 = A2AMessage::cancel_task("task-2");
628
629        // Each message must receive a distinct ID (all pairs, not just adjacent).
630        assert_ne!(msg1.id, msg2.id);
631        assert_ne!(msg1.id, msg3.id);
632        assert_ne!(msg1.id, msg4.id);
633        assert_ne!(msg2.id, msg3.id);
634        assert_ne!(msg2.id, msg4.id);
635        assert_ne!(msg3.id, msg4.id);
636    }
637
638    #[test]
639    fn test_request_id_is_numeric() {
640        let msg = A2AMessage::send("test");
641        assert!(
642            msg.id.is_number(),
643            "request ID must be a JSON number, got: {:?}",
644            msg.id
645        );
646    }
647
648    #[test]
649    fn test_with_id_override() {
650        let msg = A2AMessage::send("test").with_id(42u64);
651        assert_eq!(msg.id, serde_json::json!(42));
652    }
653}
litellm_rs/core/a2a/message.rs

litellm_rs/core/a2a/
message.rs
