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/// Parameters for listing tasks with filtering and pagination.
194#[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/// Result object for tasks/list method.
218#[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/// Parameters for getting a specific push notification config.
230#[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/// Parameters for listing all push notification configs for a task.
243#[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/// Parameters for deleting a push notification config.
251#[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    /// Create a new task with the given ID in the submitted state
339    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    /// Create a new task with the given ID and context ID in the submitted state
355    pub fn with_context(id: String, context_id: String) -> Self {
356        Self::new(id, context_id)
357    }
358
359    /// Update the task status
360    #[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    /// Get a copy of this task with history limited to the specified length
381    #[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    /// Add an artifact to the task
429    #[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    /// Validate a task (useful after building with builder)
439    #[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/// A task paired with its storage version — the optimistic-concurrency token.
479///
480/// The version is a monotonic counter the storage adapter bumps on every
481/// successful mutation of the task. A caller reads a task and its version, then
482/// passes that version back on a conditional update
483/// ([`AsyncTaskVersioning::update_status_checked`](crate::port::AsyncTaskVersioning::update_status_checked));
484/// if another writer advanced the task in between, the update fails with
485/// [`A2AError::VersionConflict`](crate::domain::A2AError::VersionConflict) instead
486/// of silently clobbering it.
487#[derive(Debug, Clone, PartialEq)]
488pub struct VersionedTask {
489    /// The task at this version.
490    pub task: Task,
491    /// The storage version this snapshot was read or written at.
492    pub version: u64,
493}
494
495impl VersionedTask {
496    /// Pair a task with a version.
497    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    /// Every state is classified exactly once, so a state added to the proto
507    /// cannot slip through unclassified — the failure mode being a task that
508    /// nothing will cancel and no client can be rid of.
509    #[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    /// The states that had been refused: queued work, and work stopped waiting
532    /// on the caller. Cancelling those is the whole point of cancel.
533    #[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    /// A state this build does not recognize leaves the client a way out
550    /// rather than stranding the task.
551    #[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}