Skip to main content

a2a_rs/domain/core/
task.rs

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
13// Re-export generated types
14pub 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    /// The spec's *interrupted* states: the agent has stopped and is waiting on
39    /// the caller. Not terminal — the task resumes once the caller supplies what
40    /// it asked for — but a caller polling for progress should stop here too.
41    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    /// Whether the task has stopped and will not advance on its own — terminal,
49    /// or interrupted waiting on the caller.
50    ///
51    /// This is the spec's own stopping condition for a blocking `SendMessage`:
52    /// `return_immediately = false` obliges the server to wait for a terminal
53    /// **or** interrupted state (`spec/a2a.proto:155`). It is also where a
54    /// subscription ends, and where a polling client stops. Those three had
55    /// each spelled the condition out separately.
56    pub fn is_settled(&self) -> bool {
57        self.is_terminal() || self.is_interrupted()
58    }
59
60    /// Whether a client may still cancel a task in this state.
61    ///
62    /// Everything that has not finished: queued (`SUBMITTED`), running
63    /// (`WORKING`), and the interrupted states — cancelling an `INPUT_REQUIRED`
64    /// task is exactly how a client says "never mind". Only the terminal states
65    /// refuse, because there is nothing left to stop.
66    ///
67    /// Defined here rather than in each storage adapter so the rule has one
68    /// home: a state added later is cancelable unless it is declared terminal,
69    /// which is the safe direction to be wrong in — the alternative strands a
70    /// task no client can get rid of.
71    pub fn is_cancelable(&self) -> bool {
72        !self.is_terminal()
73    }
74}
75
76/// When a `SendMessage` call should return.
77///
78/// The A2A default is the **waiting** one, which is easy to get backwards:
79/// `SendMessageConfiguration.return_immediately` is a proto3 `bool` and so
80/// defaults to `false`, and `false` obliges the server to wait until the task
81/// settles before returning (`spec/a2a.proto:155`). A conformant client that
82/// sends no configuration at all — which is what the official SDKs do — is
83/// therefore promised a settled task, not an acknowledgement.
84///
85/// Modelled as an enum rather than the wire's bare `bool` because
86/// `send_message(…, false)` at a call site reads as "don't" with no way to tell
87/// *what* is being declined, and the polarity is already the trap here.
88///
89/// Lives in the domain because both directions need it: the server decodes it
90/// from the request, and the client `Transport` port sets it on the way out.
91#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
92pub enum SendCompletion {
93    /// Wait until the task reaches a terminal or interrupted state. The spec
94    /// default, and what a client that sent no configuration is owed.
95    #[default]
96    WhenSettled,
97    /// Return as soon as the message is accepted, even if the agent is still
98    /// working (`return_immediately = true`).
99    ///
100    /// What a caller doing its own follow-up wants: it keeps *its* deadline in
101    /// charge instead of inheriting the server's.
102    WhenCreated,
103}
104
105impl SendCompletion {
106    /// The wire spelling. Note the inversion — this is why the enum exists.
107    #[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
120/// An unrecognized state is treated as non-terminal, and therefore cancelable:
121/// a peer speaking a newer protocol may have states this build does not know,
122/// and refusing to cancel one would leave the client no way out at all.
123impl 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    /// An unrecognized state is *not* settled: a caller waiting on it should
139    /// keep waiting rather than report a task finished on the strength of a
140    /// state this build cannot read.
141    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/// Parameters for identifying a task by ID.
176#[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/// Parameters for querying a task with optional history constraints.
184#[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/// Configuration options for sending messages including output modes and notifications.
194#[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/// Parameters for sending a message with optional configuration.
213#[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/// Parameters for sending a task (legacy)
223#[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/// Parameters for listing tasks with filtering and pagination.
238#[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/// Result object for tasks/list method.
262#[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/// Parameters for getting a specific push notification config.
274#[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/// Parameters for listing all push notification configs for a task.
287#[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/// Parameters for deleting a push notification config.
295#[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    /// Create a new task with the given ID in the submitted state
383    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    /// Create a new task with the given ID and context ID in the submitted state
399    pub fn with_context(id: String, context_id: String) -> Self {
400        Self::new(id, context_id)
401    }
402
403    /// Update the task status
404    #[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    /// Get a copy of this task with history limited to the specified length
425    #[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    /// Add an artifact to the task
473    #[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    /// Validate a task (useful after building with builder)
483    #[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/// A task paired with its storage version — the optimistic-concurrency token.
523///
524/// The version is a monotonic counter the storage adapter bumps on every
525/// successful mutation of the task. A caller reads a task and its version, then
526/// passes that version back on a conditional update
527/// ([`AsyncTaskVersioning::update_status_checked`](crate::port::AsyncTaskVersioning::update_status_checked));
528/// if another writer advanced the task in between, the update fails with
529/// [`A2AError::VersionConflict`](crate::domain::A2AError::VersionConflict) instead
530/// of silently clobbering it.
531#[derive(Debug, Clone, PartialEq)]
532pub struct VersionedTask {
533    /// The task at this version.
534    pub task: Task,
535    /// The storage version this snapshot was read or written at.
536    pub version: u64,
537}
538
539impl VersionedTask {
540    /// Pair a task with a version.
541    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    /// Every state is classified exactly once, so a state added to the proto
551    /// cannot slip through unclassified — the failure mode being a task that
552    /// nothing will cancel and no client can be rid of.
553    #[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    /// The states that had been refused: queued work, and work stopped waiting
576    /// on the caller. Cancelling those is the whole point of cancel.
577    #[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    /// A state this build does not recognize leaves the client a way out
594    /// rather than stranding the task.
595    #[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}