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, Default)]
195pub struct ListTasksParams {
196 #[serde(skip_serializing_if = "Option::is_none", rename = "contextId")]
197 pub context_id: Option<String>,
198 #[serde(skip_serializing_if = "Option::is_none")]
199 pub status: Option<TaskState>,
200 #[serde(skip_serializing_if = "Option::is_none", rename = "pageSize")]
201 pub page_size: Option<i32>,
202 #[serde(skip_serializing_if = "Option::is_none", rename = "pageToken")]
203 pub page_token: Option<String>,
204 #[serde(skip_serializing_if = "Option::is_none", rename = "historyLength")]
205 pub history_length: Option<i32>,
206 #[serde(skip_serializing_if = "Option::is_none", rename = "includeArtifacts")]
207 pub include_artifacts: Option<bool>,
208 #[serde(
209 skip_serializing_if = "Option::is_none",
210 rename = "statusTimestampAfter"
211 )]
212 pub status_timestamp_after: Option<String>,
213 #[serde(skip_serializing_if = "Option::is_none")]
214 pub metadata: Option<Map<String, Value>>,
215}
216
217#[derive(Debug, Clone, Serialize, Deserialize)]
219pub struct ListTasksResult {
220 pub tasks: Vec<Task>,
221 #[serde(rename = "totalSize")]
222 pub total_size: i32,
223 #[serde(rename = "pageSize")]
224 pub page_size: i32,
225 #[serde(rename = "nextPageToken")]
226 pub next_page_token: String,
227}
228
229#[derive(Debug, Clone, Serialize, Deserialize, Default)]
231pub struct GetTaskPushNotificationConfigParams {
232 pub id: String,
233 #[serde(
234 skip_serializing_if = "Option::is_none",
235 rename = "pushNotificationConfigId"
236 )]
237 pub push_notification_config_id: Option<String>,
238 #[serde(skip_serializing_if = "Option::is_none")]
239 pub metadata: Option<Map<String, Value>>,
240}
241
242#[derive(Debug, Clone, Serialize, Deserialize)]
244pub struct ListTaskPushNotificationConfigsParams {
245 pub id: String,
246 #[serde(skip_serializing_if = "Option::is_none")]
247 pub metadata: Option<Map<String, Value>>,
248}
249
250#[derive(Debug, Clone, Serialize, Deserialize)]
252pub struct DeleteTaskPushNotificationConfigParams {
253 pub id: String,
254 #[serde(rename = "pushNotificationConfigId")]
255 pub push_notification_config_id: String,
256 #[serde(skip_serializing_if = "Option::is_none")]
257 pub metadata: Option<Map<String, Value>>,
258}
259
260pub struct TaskBuilder {
261 id: String,
262 context_id: String,
263 status: Option<TaskStatus>,
264 artifacts: Vec<Artifact>,
265 history: Vec<Message>,
266 metadata: Option<::buffa_types::google::protobuf::Struct>,
267}
268
269impl TaskBuilder {
270 pub fn new() -> Self {
271 Self {
272 id: String::new(),
273 context_id: String::new(),
274 status: None,
275 artifacts: Vec::new(),
276 history: Vec::new(),
277 metadata: None,
278 }
279 }
280
281 pub fn id(mut self, id: String) -> Self {
282 self.id = id;
283 self
284 }
285
286 pub fn context_id(mut self, context_id: String) -> Self {
287 self.context_id = context_id;
288 self
289 }
290
291 pub fn status(mut self, status: TaskStatus) -> Self {
292 self.status = Some(status);
293 self
294 }
295
296 pub fn artifacts(mut self, artifacts: Vec<Artifact>) -> Self {
297 self.artifacts = artifacts;
298 self
299 }
300
301 pub fn history(mut self, history: Vec<Message>) -> Self {
302 self.history = history;
303 self
304 }
305
306 pub fn metadata(mut self, metadata: ::buffa_types::google::protobuf::Struct) -> Self {
307 self.metadata = Some(metadata);
308 self
309 }
310
311 pub fn build(self) -> Task {
312 Task {
313 id: self.id,
314 context_id: self.context_id,
315 status: self
316 .status
317 .unwrap_or_else(|| TaskStatus::new(TaskState::TASK_STATE_SUBMITTED, None))
318 .into(),
319 artifacts: self.artifacts,
320 history: self.history,
321 metadata: self.metadata.into(),
322 ..Default::default()
323 }
324 }
325}
326
327impl Default for TaskBuilder {
328 fn default() -> Self {
329 Self::new()
330 }
331}
332
333impl Task {
334 pub fn builder() -> TaskBuilder {
335 TaskBuilder::new()
336 }
337
338 pub fn new(id: String, context_id: String) -> Self {
340 Self {
341 id,
342 context_id,
343 status: ::buffa::MessageField::some(TaskStatus::new(
344 TaskState::TASK_STATE_SUBMITTED,
345 None,
346 )),
347 artifacts: Vec::new(),
348 history: Vec::new(),
349 metadata: ::buffa::MessageField::none(),
350 ..Default::default()
351 }
352 }
353
354 pub fn with_context(id: String, context_id: String) -> Self {
356 Self::new(id, context_id)
357 }
358
359 #[cfg_attr(feature = "tracing", instrument(skip(self, message), fields(
361 task.id = %self.id,
362 task.old_state = ?self.status.as_option().map(|s| &s.state),
363 task.new_state = ?state,
364 task.has_message = message.is_some()
365 )))]
366 pub fn update_status(&mut self, state: TaskState, message: Option<Message>) {
367 #[cfg(feature = "tracing")]
368 tracing::info!("Updating task status");
369
370 self.status = ::buffa::MessageField::some(TaskStatus::new(state, message.clone()));
371
372 if let Some(msg) = message {
373 self.history.push(msg);
374 }
375
376 #[cfg(feature = "tracing")]
377 tracing::info!("Task status updated successfully");
378 }
379
380 #[cfg_attr(feature = "tracing", instrument(skip(self), fields(
382 task.id = %self.id,
383 history.current_size = self.history.len(),
384 history.requested_limit = ?history_length
385 )))]
386 pub fn with_limited_history(&self, history_length: Option<u32>) -> Self {
387 if history_length.is_none() {
388 #[cfg(feature = "tracing")]
389 tracing::debug!("No history truncation needed");
390 return self.clone();
391 }
392
393 #[cfg(feature = "tracing")]
394 let _span = tracing::Span::current();
395
396 let limit: usize = history_length.unwrap().try_into().unwrap_or(usize::MAX);
397
398 #[cfg(feature = "tracing")]
399 let mut task_copy = measure_duration!(_span, "operation.duration_ms", { self.clone() });
400
401 #[cfg(not(feature = "tracing"))]
402 let mut task_copy = self.clone();
403
404 if limit == 0 {
405 #[cfg(feature = "tracing")]
406 tracing::debug!("Removing all history (limit = 0)");
407 task_copy.history.clear();
408 } else if task_copy.history.len() > limit {
409 let items_to_skip = task_copy.history.len() - limit;
410 #[cfg(feature = "tracing")]
411 tracing::debug!(
412 "Truncating history from {} to {} items (removing {} oldest)",
413 self.history.len(),
414 limit,
415 items_to_skip
416 );
417 task_copy.history = task_copy
418 .history
419 .iter()
420 .skip(items_to_skip)
421 .cloned()
422 .collect();
423 }
424
425 task_copy
426 }
427
428 #[cfg_attr(feature = "tracing", instrument(skip(self, artifact), fields(
430 task.id = %self.id,
431 artifact.id = %artifact.artifact_id,
432 artifacts.count = self.artifacts.len()
433 )))]
434 pub fn add_artifact(&mut self, artifact: Artifact) {
435 self.artifacts.push(artifact);
436 }
437
438 #[cfg_attr(feature = "tracing", instrument(skip(self), fields(
440 task.id = %self.id,
441 task.state = ?self.status.as_option().map(|s| &s.state),
442 history.size = self.history.len()
443 )))]
444 pub fn validate(&self) -> Result<(), A2AError> {
445 #[cfg(feature = "tracing")]
446 tracing::debug!("Validating task");
447
448 let mut message_ids = std::collections::HashSet::new();
449 for (_index, message) in self.history.iter().enumerate() {
450 #[cfg(feature = "tracing")]
451 tracing::trace!("Validating message {} in history", _index);
452
453 if !message_ids.insert(&message.message_id) {
454 #[cfg(feature = "tracing")]
455 tracing::error!("Duplicate message ID found: {}", message.message_id);
456 return Err(A2AError::InvalidParams(format!(
457 "Duplicate message ID in history: {}",
458 message.message_id
459 )));
460 }
461 message.validate()?;
462 }
463
464 if let Some(status) = self.status.as_option()
465 && let Some(msg) = status.message.as_option()
466 {
467 #[cfg(feature = "tracing")]
468 tracing::trace!("Validating status message");
469 msg.validate()?;
470 }
471
472 #[cfg(feature = "tracing")]
473 tracing::debug!("Task validation successful");
474 Ok(())
475 }
476}
477
478#[derive(Debug, Clone, PartialEq)]
488pub struct VersionedTask {
489 pub task: Task,
491 pub version: u64,
493}
494
495impl VersionedTask {
496 pub fn new(task: Task, version: u64) -> Self {
498 Self { task, version }
499 }
500}
501
502#[cfg(test)]
503mod state_predicate_tests {
504 use super::*;
505
506 #[test]
510 fn every_state_is_either_terminal_or_cancelable() {
511 let all = [
512 TaskState::TASK_STATE_UNSPECIFIED,
513 TaskState::TASK_STATE_SUBMITTED,
514 TaskState::TASK_STATE_WORKING,
515 TaskState::TASK_STATE_COMPLETED,
516 TaskState::TASK_STATE_FAILED,
517 TaskState::TASK_STATE_CANCELED,
518 TaskState::TASK_STATE_INPUT_REQUIRED,
519 TaskState::TASK_STATE_REJECTED,
520 TaskState::TASK_STATE_AUTH_REQUIRED,
521 ];
522 for state in all {
523 assert_ne!(
524 state.is_terminal(),
525 state.is_cancelable(),
526 "{state:?} must be exactly one of terminal / cancelable"
527 );
528 }
529 }
530
531 #[test]
534 fn unfinished_work_can_be_canceled() {
535 assert!(TaskState::Submitted.is_cancelable());
536 assert!(TaskState::Working.is_cancelable());
537 assert!(TaskState::InputRequired.is_cancelable());
538 assert!(TaskState::AuthRequired.is_cancelable());
539 }
540
541 #[test]
542 fn finished_work_cannot_be_canceled() {
543 assert!(!TaskState::Completed.is_cancelable());
544 assert!(!TaskState::Failed.is_cancelable());
545 assert!(!TaskState::Canceled.is_cancelable());
546 assert!(!TaskState::Rejected.is_cancelable());
547 }
548
549 #[test]
552 fn an_unrecognized_state_is_cancelable() {
553 let unknown: ::buffa::EnumValue<TaskState> = ::buffa::EnumValue::Unknown(99);
554 assert!(!unknown.is_terminal());
555 assert!(unknown.is_cancelable());
556 }
557}