1use crate::domain::error::A2AError;
2use serde::{Deserialize, Serialize};
3use serde_json::{Map, Value};
4
5#[cfg(feature = "tracing")]
6use tracing::instrument;
7
8#[cfg(feature = "tracing")]
9use crate::measure_duration;
10
11use super::message::{Artifact, Message};
12
13pub use crate::domain::generated::{Task, TaskPushNotificationConfig, TaskState, TaskStatus};
15
16#[allow(non_upper_case_globals)]
17impl TaskState {
18 pub const Submitted: Self = Self::TASK_STATE_SUBMITTED;
19 pub const Working: Self = Self::TASK_STATE_WORKING;
20 pub const InputRequired: Self = Self::TASK_STATE_INPUT_REQUIRED;
21 pub const Completed: Self = Self::TASK_STATE_COMPLETED;
22 pub const Canceled: Self = Self::TASK_STATE_CANCELED;
23 pub const Failed: Self = Self::TASK_STATE_FAILED;
24 pub const Rejected: Self = Self::TASK_STATE_REJECTED;
25 pub const AuthRequired: Self = Self::TASK_STATE_AUTH_REQUIRED;
26 pub const Unknown: Self = Self::TASK_STATE_UNSPECIFIED;
27
28 pub fn is_terminal(&self) -> bool {
29 matches!(
30 self,
31 Self::TASK_STATE_COMPLETED
32 | Self::TASK_STATE_FAILED
33 | Self::TASK_STATE_CANCELED
34 | Self::TASK_STATE_REJECTED
35 )
36 }
37
38 pub fn is_interrupted(&self) -> bool {
42 matches!(
43 self,
44 Self::TASK_STATE_INPUT_REQUIRED | Self::TASK_STATE_AUTH_REQUIRED
45 )
46 }
47
48 pub fn is_settled(&self) -> bool {
57 self.is_terminal() || self.is_interrupted()
58 }
59
60 pub fn is_cancelable(&self) -> bool {
72 !self.is_terminal()
73 }
74}
75
76#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
92pub enum SendCompletion {
93 #[default]
96 WhenSettled,
97 WhenCreated,
103}
104
105impl SendCompletion {
106 #[inline]
108 pub fn return_immediately(self) -> bool {
109 matches!(self, Self::WhenCreated)
110 }
111}
112
113pub trait TaskStateExt {
114 fn is_terminal(&self) -> bool;
115 fn is_interrupted(&self) -> bool;
116 fn is_settled(&self) -> bool;
117 fn is_cancelable(&self) -> bool;
118}
119
120impl TaskStateExt for ::buffa::EnumValue<TaskState> {
124 fn is_terminal(&self) -> bool {
125 match self {
126 ::buffa::EnumValue::Known(state) => state.is_terminal(),
127 _ => false,
128 }
129 }
130
131 fn is_interrupted(&self) -> bool {
132 match self {
133 ::buffa::EnumValue::Known(state) => state.is_interrupted(),
134 _ => false,
135 }
136 }
137
138 fn is_settled(&self) -> bool {
142 self.is_terminal() || self.is_interrupted()
143 }
144
145 fn is_cancelable(&self) -> bool {
146 !self.is_terminal()
147 }
148}
149
150impl TaskStatus {
151 pub fn new(state: TaskState, message: Option<Message>) -> Self {
152 let timestamp = chrono::Utc::now();
153 let seconds = timestamp.timestamp();
154 let nanos = timestamp.timestamp_subsec_nanos() as i32;
155
156 Self {
157 state: ::buffa::EnumValue::from(state),
158 message: message.into(),
159 timestamp: ::buffa::MessageField::some(::buffa_types::google::protobuf::Timestamp {
160 seconds,
161 nanos,
162 ..Default::default()
163 }),
164 ..Default::default()
165 }
166 }
167
168 pub fn timestamp_utc(&self) -> Option<chrono::DateTime<chrono::Utc>> {
169 self.timestamp.as_option().and_then(|t| {
170 chrono::DateTime::<chrono::Utc>::from_timestamp(t.seconds, t.nanos as u32)
171 })
172 }
173}
174
175#[derive(Debug, Clone, Serialize, Deserialize)]
177pub struct TaskIdParams {
178 pub id: String,
179 #[serde(skip_serializing_if = "Option::is_none")]
180 pub metadata: Option<Map<String, Value>>,
181}
182
183#[derive(Debug, Clone, Serialize, Deserialize)]
185pub struct TaskQueryParams {
186 pub id: String,
187 #[serde(skip_serializing_if = "Option::is_none", rename = "historyLength")]
188 pub history_length: Option<u32>,
189 #[serde(skip_serializing_if = "Option::is_none")]
190 pub metadata: Option<Map<String, Value>>,
191}
192
193#[derive(Debug, Clone, Serialize, Deserialize)]
195pub struct MessageSendConfiguration {
196 #[serde(
197 skip_serializing_if = "Option::is_none",
198 rename = "acceptedOutputModes"
199 )]
200 pub accepted_output_modes: Option<Vec<String>>,
201 #[serde(skip_serializing_if = "Option::is_none", rename = "historyLength")]
202 pub history_length: Option<u32>,
203 #[serde(
204 skip_serializing_if = "Option::is_none",
205 rename = "pushNotificationConfig"
206 )]
207 pub push_notification_config: Option<TaskPushNotificationConfig>,
208 #[serde(skip_serializing_if = "Option::is_none")]
209 pub blocking: Option<bool>,
210}
211
212#[derive(Debug, Clone, Serialize, Deserialize)]
214pub struct MessageSendParams {
215 pub message: Message,
216 #[serde(skip_serializing_if = "Option::is_none")]
217 pub configuration: Option<MessageSendConfiguration>,
218 #[serde(skip_serializing_if = "Option::is_none")]
219 pub metadata: Option<Map<String, Value>>,
220}
221
222#[derive(Debug, Clone, Serialize, Deserialize)]
224pub struct TaskSendParams {
225 pub id: String,
226 #[serde(skip_serializing_if = "Option::is_none", rename = "sessionId")]
227 pub session_id: Option<String>,
228 pub message: Message,
229 #[serde(skip_serializing_if = "Option::is_none", rename = "pushNotification")]
230 pub push_notification: Option<TaskPushNotificationConfig>,
231 #[serde(skip_serializing_if = "Option::is_none", rename = "historyLength")]
232 pub history_length: Option<u32>,
233 #[serde(skip_serializing_if = "Option::is_none")]
234 pub metadata: Option<Map<String, Value>>,
235}
236
237#[derive(Debug, Clone, Serialize, Deserialize, Default)]
239pub struct ListTasksParams {
240 #[serde(skip_serializing_if = "Option::is_none", rename = "contextId")]
241 pub context_id: Option<String>,
242 #[serde(skip_serializing_if = "Option::is_none")]
243 pub status: Option<TaskState>,
244 #[serde(skip_serializing_if = "Option::is_none", rename = "pageSize")]
245 pub page_size: Option<i32>,
246 #[serde(skip_serializing_if = "Option::is_none", rename = "pageToken")]
247 pub page_token: Option<String>,
248 #[serde(skip_serializing_if = "Option::is_none", rename = "historyLength")]
249 pub history_length: Option<i32>,
250 #[serde(skip_serializing_if = "Option::is_none", rename = "includeArtifacts")]
251 pub include_artifacts: Option<bool>,
252 #[serde(
253 skip_serializing_if = "Option::is_none",
254 rename = "statusTimestampAfter"
255 )]
256 pub status_timestamp_after: Option<String>,
257 #[serde(skip_serializing_if = "Option::is_none")]
258 pub metadata: Option<Map<String, Value>>,
259}
260
261#[derive(Debug, Clone, Serialize, Deserialize)]
263pub struct ListTasksResult {
264 pub tasks: Vec<Task>,
265 #[serde(rename = "totalSize")]
266 pub total_size: i32,
267 #[serde(rename = "pageSize")]
268 pub page_size: i32,
269 #[serde(rename = "nextPageToken")]
270 pub next_page_token: String,
271}
272
273#[derive(Debug, Clone, Serialize, Deserialize, Default)]
275pub struct GetTaskPushNotificationConfigParams {
276 pub id: String,
277 #[serde(
278 skip_serializing_if = "Option::is_none",
279 rename = "pushNotificationConfigId"
280 )]
281 pub push_notification_config_id: Option<String>,
282 #[serde(skip_serializing_if = "Option::is_none")]
283 pub metadata: Option<Map<String, Value>>,
284}
285
286#[derive(Debug, Clone, Serialize, Deserialize)]
288pub struct ListTaskPushNotificationConfigsParams {
289 pub id: String,
290 #[serde(skip_serializing_if = "Option::is_none")]
291 pub metadata: Option<Map<String, Value>>,
292}
293
294#[derive(Debug, Clone, Serialize, Deserialize)]
296pub struct DeleteTaskPushNotificationConfigParams {
297 pub id: String,
298 #[serde(rename = "pushNotificationConfigId")]
299 pub push_notification_config_id: String,
300 #[serde(skip_serializing_if = "Option::is_none")]
301 pub metadata: Option<Map<String, Value>>,
302}
303
304pub struct TaskBuilder {
305 id: String,
306 context_id: String,
307 status: Option<TaskStatus>,
308 artifacts: Vec<Artifact>,
309 history: Vec<Message>,
310 metadata: Option<::buffa_types::google::protobuf::Struct>,
311}
312
313impl TaskBuilder {
314 pub fn new() -> Self {
315 Self {
316 id: String::new(),
317 context_id: String::new(),
318 status: None,
319 artifacts: Vec::new(),
320 history: Vec::new(),
321 metadata: None,
322 }
323 }
324
325 pub fn id(mut self, id: String) -> Self {
326 self.id = id;
327 self
328 }
329
330 pub fn context_id(mut self, context_id: String) -> Self {
331 self.context_id = context_id;
332 self
333 }
334
335 pub fn status(mut self, status: TaskStatus) -> Self {
336 self.status = Some(status);
337 self
338 }
339
340 pub fn artifacts(mut self, artifacts: Vec<Artifact>) -> Self {
341 self.artifacts = artifacts;
342 self
343 }
344
345 pub fn history(mut self, history: Vec<Message>) -> Self {
346 self.history = history;
347 self
348 }
349
350 pub fn metadata(mut self, metadata: ::buffa_types::google::protobuf::Struct) -> Self {
351 self.metadata = Some(metadata);
352 self
353 }
354
355 pub fn build(self) -> Task {
356 Task {
357 id: self.id,
358 context_id: self.context_id,
359 status: self
360 .status
361 .unwrap_or_else(|| TaskStatus::new(TaskState::TASK_STATE_SUBMITTED, None))
362 .into(),
363 artifacts: self.artifacts,
364 history: self.history,
365 metadata: self.metadata.into(),
366 ..Default::default()
367 }
368 }
369}
370
371impl Default for TaskBuilder {
372 fn default() -> Self {
373 Self::new()
374 }
375}
376
377impl Task {
378 pub fn builder() -> TaskBuilder {
379 TaskBuilder::new()
380 }
381
382 pub fn new(id: String, context_id: String) -> Self {
384 Self {
385 id,
386 context_id,
387 status: ::buffa::MessageField::some(TaskStatus::new(
388 TaskState::TASK_STATE_SUBMITTED,
389 None,
390 )),
391 artifacts: Vec::new(),
392 history: Vec::new(),
393 metadata: ::buffa::MessageField::none(),
394 ..Default::default()
395 }
396 }
397
398 pub fn with_context(id: String, context_id: String) -> Self {
400 Self::new(id, context_id)
401 }
402
403 #[cfg_attr(feature = "tracing", instrument(skip(self, message), fields(
405 task.id = %self.id,
406 task.old_state = ?self.status.as_option().map(|s| &s.state),
407 task.new_state = ?state,
408 task.has_message = message.is_some()
409 )))]
410 pub fn update_status(&mut self, state: TaskState, message: Option<Message>) {
411 #[cfg(feature = "tracing")]
412 tracing::info!("Updating task status");
413
414 self.status = ::buffa::MessageField::some(TaskStatus::new(state, message.clone()));
415
416 if let Some(msg) = message {
417 self.history.push(msg);
418 }
419
420 #[cfg(feature = "tracing")]
421 tracing::info!("Task status updated successfully");
422 }
423
424 #[cfg_attr(feature = "tracing", instrument(skip(self), fields(
426 task.id = %self.id,
427 history.current_size = self.history.len(),
428 history.requested_limit = ?history_length
429 )))]
430 pub fn with_limited_history(&self, history_length: Option<u32>) -> Self {
431 if history_length.is_none() {
432 #[cfg(feature = "tracing")]
433 tracing::debug!("No history truncation needed");
434 return self.clone();
435 }
436
437 #[cfg(feature = "tracing")]
438 let _span = tracing::Span::current();
439
440 let limit: usize = history_length.unwrap().try_into().unwrap_or(usize::MAX);
441
442 #[cfg(feature = "tracing")]
443 let mut task_copy = measure_duration!(_span, "operation.duration_ms", { self.clone() });
444
445 #[cfg(not(feature = "tracing"))]
446 let mut task_copy = self.clone();
447
448 if limit == 0 {
449 #[cfg(feature = "tracing")]
450 tracing::debug!("Removing all history (limit = 0)");
451 task_copy.history.clear();
452 } else if task_copy.history.len() > limit {
453 let items_to_skip = task_copy.history.len() - limit;
454 #[cfg(feature = "tracing")]
455 tracing::debug!(
456 "Truncating history from {} to {} items (removing {} oldest)",
457 self.history.len(),
458 limit,
459 items_to_skip
460 );
461 task_copy.history = task_copy
462 .history
463 .iter()
464 .skip(items_to_skip)
465 .cloned()
466 .collect();
467 }
468
469 task_copy
470 }
471
472 #[cfg_attr(feature = "tracing", instrument(skip(self, artifact), fields(
474 task.id = %self.id,
475 artifact.id = %artifact.artifact_id,
476 artifacts.count = self.artifacts.len()
477 )))]
478 pub fn add_artifact(&mut self, artifact: Artifact) {
479 self.artifacts.push(artifact);
480 }
481
482 #[cfg_attr(feature = "tracing", instrument(skip(self), fields(
484 task.id = %self.id,
485 task.state = ?self.status.as_option().map(|s| &s.state),
486 history.size = self.history.len()
487 )))]
488 pub fn validate(&self) -> Result<(), A2AError> {
489 #[cfg(feature = "tracing")]
490 tracing::debug!("Validating task");
491
492 let mut message_ids = std::collections::HashSet::new();
493 for (_index, message) in self.history.iter().enumerate() {
494 #[cfg(feature = "tracing")]
495 tracing::trace!("Validating message {} in history", _index);
496
497 if !message_ids.insert(&message.message_id) {
498 #[cfg(feature = "tracing")]
499 tracing::error!("Duplicate message ID found: {}", message.message_id);
500 return Err(A2AError::InvalidParams(format!(
501 "Duplicate message ID in history: {}",
502 message.message_id
503 )));
504 }
505 message.validate()?;
506 }
507
508 if let Some(status) = self.status.as_option()
509 && let Some(msg) = status.message.as_option()
510 {
511 #[cfg(feature = "tracing")]
512 tracing::trace!("Validating status message");
513 msg.validate()?;
514 }
515
516 #[cfg(feature = "tracing")]
517 tracing::debug!("Task validation successful");
518 Ok(())
519 }
520}
521
522#[derive(Debug, Clone, PartialEq)]
532pub struct VersionedTask {
533 pub task: Task,
535 pub version: u64,
537}
538
539impl VersionedTask {
540 pub fn new(task: Task, version: u64) -> Self {
542 Self { task, version }
543 }
544}
545
546#[cfg(test)]
547mod state_predicate_tests {
548 use super::*;
549
550 #[test]
554 fn every_state_is_either_terminal_or_cancelable() {
555 let all = [
556 TaskState::TASK_STATE_UNSPECIFIED,
557 TaskState::TASK_STATE_SUBMITTED,
558 TaskState::TASK_STATE_WORKING,
559 TaskState::TASK_STATE_COMPLETED,
560 TaskState::TASK_STATE_FAILED,
561 TaskState::TASK_STATE_CANCELED,
562 TaskState::TASK_STATE_INPUT_REQUIRED,
563 TaskState::TASK_STATE_REJECTED,
564 TaskState::TASK_STATE_AUTH_REQUIRED,
565 ];
566 for state in all {
567 assert_ne!(
568 state.is_terminal(),
569 state.is_cancelable(),
570 "{state:?} must be exactly one of terminal / cancelable"
571 );
572 }
573 }
574
575 #[test]
578 fn unfinished_work_can_be_canceled() {
579 assert!(TaskState::Submitted.is_cancelable());
580 assert!(TaskState::Working.is_cancelable());
581 assert!(TaskState::InputRequired.is_cancelable());
582 assert!(TaskState::AuthRequired.is_cancelable());
583 }
584
585 #[test]
586 fn finished_work_cannot_be_canceled() {
587 assert!(!TaskState::Completed.is_cancelable());
588 assert!(!TaskState::Failed.is_cancelable());
589 assert!(!TaskState::Canceled.is_cancelable());
590 assert!(!TaskState::Rejected.is_cancelable());
591 }
592
593 #[test]
596 fn an_unrecognized_state_is_cancelable() {
597 let unknown: ::buffa::EnumValue<TaskState> = ::buffa::EnumValue::Unknown(99);
598 assert!(!unknown.is_terminal());
599 assert!(unknown.is_cancelable());
600 }
601}