turul_a2a_types/
task.rs

1use serde::{Deserialize, Serialize};
2use turul_a2a_proto as pb;
3
4use crate::artifact::Artifact;
5use crate::error::A2aTypeError;
6use crate::message::Message;
7
8/// A2A task states — type-safe wrapper over proto TaskState.
9///
10/// Excludes `UNSPECIFIED` which is not a valid application state.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12#[non_exhaustive]
13pub enum TaskState {
14    Submitted,
15    Working,
16    Completed,
17    Failed,
18    Canceled,
19    InputRequired,
20    Rejected,
21    AuthRequired,
22}
23
24impl TaskState {
25    /// Returns whether transitioning from `self` to `next` is valid per the A2A spec.
26    pub fn can_transition_to(&self, next: TaskState) -> bool {
27        crate::state_machine::validate_transition(*self, next).is_ok()
28    }
29
30    /// Returns `true` if this is a terminal state (no further transitions allowed).
31    pub fn is_terminal(&self) -> bool {
32        crate::state_machine::is_terminal(*self)
33    }
34}
35
36impl TryFrom<pb::TaskState> for TaskState {
37    type Error = A2aTypeError;
38
39    fn try_from(value: pb::TaskState) -> Result<Self, Self::Error> {
40        match value {
41            pb::TaskState::Submitted => Ok(Self::Submitted),
42            pb::TaskState::Working => Ok(Self::Working),
43            pb::TaskState::Completed => Ok(Self::Completed),
44            pb::TaskState::Failed => Ok(Self::Failed),
45            pb::TaskState::Canceled => Ok(Self::Canceled),
46            pb::TaskState::InputRequired => Ok(Self::InputRequired),
47            pb::TaskState::Rejected => Ok(Self::Rejected),
48            pb::TaskState::AuthRequired => Ok(Self::AuthRequired),
49            pb::TaskState::Unspecified => Err(A2aTypeError::InvalidState),
50        }
51    }
52}
53
54impl From<TaskState> for pb::TaskState {
55    fn from(value: TaskState) -> Self {
56        match value {
57            TaskState::Submitted => pb::TaskState::Submitted,
58            TaskState::Working => pb::TaskState::Working,
59            TaskState::Completed => pb::TaskState::Completed,
60            TaskState::Failed => pb::TaskState::Failed,
61            TaskState::Canceled => pb::TaskState::Canceled,
62            TaskState::InputRequired => pb::TaskState::InputRequired,
63            TaskState::Rejected => pb::TaskState::Rejected,
64            TaskState::AuthRequired => pb::TaskState::AuthRequired,
65        }
66    }
67}
68
69impl TryFrom<i32> for TaskState {
70    type Error = A2aTypeError;
71
72    fn try_from(value: i32) -> Result<Self, Self::Error> {
73        let proto_state = pb::TaskState::try_from(value).map_err(|_| A2aTypeError::InvalidState)?;
74        Self::try_from(proto_state)
75    }
76}
77
78/// Ergonomic wrapper over proto `TaskStatus`.
79#[derive(Debug, Clone)]
80#[non_exhaustive]
81pub struct TaskStatus {
82    pub(crate) inner: pb::TaskStatus,
83}
84
85impl TaskStatus {
86    pub fn new(state: TaskState) -> Self {
87        Self {
88            inner: pb::TaskStatus {
89                state: pb::TaskState::from(state).into(),
90                message: None,
91                timestamp: None,
92            },
93        }
94    }
95
96    pub fn with_message(mut self, message: Message) -> Self {
97        self.inner.message = Some(message.into_proto());
98        self
99    }
100
101    pub fn state(&self) -> Result<TaskState, A2aTypeError> {
102        let proto_state =
103            pb::TaskState::try_from(self.inner.state).map_err(|_| A2aTypeError::InvalidState)?;
104        TaskState::try_from(proto_state)
105    }
106
107    pub fn as_proto(&self) -> &pb::TaskStatus {
108        &self.inner
109    }
110
111    pub fn into_proto(self) -> pb::TaskStatus {
112        self.inner
113    }
114}
115
116impl TryFrom<pb::TaskStatus> for TaskStatus {
117    type Error = A2aTypeError;
118
119    fn try_from(inner: pb::TaskStatus) -> Result<Self, Self::Error> {
120        // Validate state is not UNSPECIFIED
121        let proto_state =
122            pb::TaskState::try_from(inner.state).map_err(|_| A2aTypeError::InvalidState)?;
123        if proto_state == pb::TaskState::Unspecified {
124            return Err(A2aTypeError::InvalidState);
125        }
126        Ok(Self { inner })
127    }
128}
129
130impl From<TaskStatus> for pb::TaskStatus {
131    fn from(status: TaskStatus) -> Self {
132        status.inner
133    }
134}
135
136impl Serialize for TaskStatus {
137    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
138        self.inner.serialize(serializer)
139    }
140}
141
142impl<'de> Deserialize<'de> for TaskStatus {
143    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
144        let proto = pb::TaskStatus::deserialize(deserializer)?;
145        TaskStatus::try_from(proto).map_err(serde::de::Error::custom)
146    }
147}
148
149/// Ergonomic wrapper over proto `Task`.
150#[derive(Debug, Clone)]
151#[non_exhaustive]
152pub struct Task {
153    pub(crate) inner: pb::Task,
154}
155
156impl Task {
157    pub fn new(id: impl Into<String>, status: TaskStatus) -> Self {
158        Self {
159            inner: pb::Task {
160                id: id.into(),
161                context_id: String::new(),
162                status: Some(status.into_proto()),
163                artifacts: vec![],
164                history: vec![],
165                metadata: None,
166            },
167        }
168    }
169
170    pub fn with_context_id(mut self, context_id: impl Into<String>) -> Self {
171        self.inner.context_id = context_id.into();
172        self
173    }
174
175    pub fn id(&self) -> &str {
176        &self.inner.id
177    }
178
179    pub fn context_id(&self) -> &str {
180        &self.inner.context_id
181    }
182
183    pub fn status(&self) -> Option<TaskStatus> {
184        self.inner
185            .status
186            .clone()
187            .and_then(|s| TaskStatus::try_from(s).ok())
188    }
189
190    pub fn history(&self) -> &[pb::Message] {
191        &self.inner.history
192    }
193
194    pub fn artifacts(&self) -> &[pb::Artifact] {
195        &self.inner.artifacts
196    }
197
198    pub fn append_message(&mut self, message: Message) {
199        self.inner.history.push(message.into_proto());
200    }
201
202    pub fn append_artifact(&mut self, artifact: Artifact) {
203        self.inner.artifacts.push(artifact.into_proto());
204    }
205
206    /// Merge an artifact using A2A streaming append semantics.
207    ///
208    /// - `append = true`: if an artifact with the same `artifactId` is
209    ///   already on the task, extend its `parts` with the incoming parts
210    ///   (same as a streaming-chunk continuation). Otherwise, append as
211    ///   a new entry.
212    /// - `append = false`: append unconditionally (new entry or caller
213    ///   tolerates duplicate ids).
214    ///
215    /// The `last_chunk` flag is transport-only metadata
216    /// — it is not persisted on the task. It is accepted here so callers
217    /// can keep the signature symmetric with the streaming wire event
218    /// payload and not drop the parameter separately.
219    ///
220    /// This mirrors the server storage append-artifact semantics so that
221    /// in-memory task mutations and the storage layer converge on the
222    /// same view. The storage trait lives in the server crate; this
223    /// helper is the dependency-free equivalent for wrapper callers.
224    pub fn merge_artifact(&mut self, artifact: Artifact, append: bool, _last_chunk: bool) {
225        if append {
226            let target_id = artifact.as_proto().artifact_id.clone();
227            if let Some(existing) = self
228                .inner
229                .artifacts
230                .iter_mut()
231                .find(|a| a.artifact_id == target_id)
232            {
233                existing.parts.extend(artifact.into_proto().parts);
234                return;
235            }
236        }
237        self.append_artifact(artifact);
238    }
239
240    /// Set the task's status. This is the low-level escape hatch —
241    /// prefer `complete()`, `fail()`, etc. for common transitions.
242    pub fn set_status(&mut self, status: TaskStatus) {
243        self.inner.status = Some(status.into_proto());
244    }
245
246    /// Mark the task as completed.
247    pub fn complete(&mut self) {
248        self.set_status(TaskStatus::new(TaskState::Completed));
249    }
250
251    /// Mark the task as failed with an optional message.
252    pub fn fail(&mut self, message: impl Into<String>) {
253        let msg = Message::new(
254            uuid::Uuid::now_v7().to_string(),
255            crate::Role::Agent,
256            vec![crate::Part::text(message)],
257        );
258        self.set_status(TaskStatus::new(TaskState::Failed).with_message(msg));
259    }
260
261    /// Add a text artifact to the task.
262    pub fn push_text_artifact(
263        &mut self,
264        artifact_id: impl Into<String>,
265        name: impl Into<String>,
266        text: impl Into<String>,
267    ) {
268        let artifact = Artifact::new(artifact_id, vec![crate::Part::text(text)]).with_name(name);
269        self.append_artifact(artifact);
270    }
271
272    pub fn as_proto(&self) -> &pb::Task {
273        &self.inner
274    }
275
276    pub fn as_proto_mut(&mut self) -> &mut pb::Task {
277        &mut self.inner
278    }
279
280    pub fn into_proto(self) -> pb::Task {
281        self.inner
282    }
283}
284
285impl TryFrom<pb::Task> for Task {
286    type Error = A2aTypeError;
287
288    fn try_from(inner: pb::Task) -> Result<Self, Self::Error> {
289        if inner.id.is_empty() {
290            return Err(A2aTypeError::MissingField("id"));
291        }
292        // Status is REQUIRED per proto field_behavior
293        let status = inner
294            .status
295            .as_ref()
296            .ok_or(A2aTypeError::MissingField("status"))?;
297        let proto_state =
298            pb::TaskState::try_from(status.state).map_err(|_| A2aTypeError::InvalidState)?;
299        if proto_state == pb::TaskState::Unspecified {
300            return Err(A2aTypeError::InvalidState);
301        }
302        Ok(Self { inner })
303    }
304}
305
306impl From<Task> for pb::Task {
307    fn from(task: Task) -> Self {
308        task.inner
309    }
310}
311
312impl Serialize for Task {
313    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
314        self.inner.serialize(serializer)
315    }
316}
317
318impl<'de> Deserialize<'de> for Task {
319    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
320        let proto = pb::Task::deserialize(deserializer)?;
321        Task::try_from(proto).map_err(serde::de::Error::custom)
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328    use crate::message::{Part, Role};
329
330    #[test]
331    fn try_from_proto_all_valid_states() {
332        assert_eq!(
333            TaskState::try_from(pb::TaskState::Submitted).unwrap(),
334            TaskState::Submitted
335        );
336        assert_eq!(
337            TaskState::try_from(pb::TaskState::Working).unwrap(),
338            TaskState::Working
339        );
340        assert_eq!(
341            TaskState::try_from(pb::TaskState::Completed).unwrap(),
342            TaskState::Completed
343        );
344        assert_eq!(
345            TaskState::try_from(pb::TaskState::Failed).unwrap(),
346            TaskState::Failed
347        );
348        assert_eq!(
349            TaskState::try_from(pb::TaskState::Canceled).unwrap(),
350            TaskState::Canceled
351        );
352        assert_eq!(
353            TaskState::try_from(pb::TaskState::InputRequired).unwrap(),
354            TaskState::InputRequired
355        );
356        assert_eq!(
357            TaskState::try_from(pb::TaskState::Rejected).unwrap(),
358            TaskState::Rejected
359        );
360        assert_eq!(
361            TaskState::try_from(pb::TaskState::AuthRequired).unwrap(),
362            TaskState::AuthRequired
363        );
364    }
365
366    #[test]
367    fn try_from_proto_unspecified_is_error() {
368        assert!(TaskState::try_from(pb::TaskState::Unspecified).is_err());
369    }
370
371    #[test]
372    fn into_proto_round_trip() {
373        for state in [
374            TaskState::Submitted,
375            TaskState::Working,
376            TaskState::Completed,
377            TaskState::Failed,
378            TaskState::Canceled,
379            TaskState::InputRequired,
380            TaskState::Rejected,
381            TaskState::AuthRequired,
382        ] {
383            let proto: pb::TaskState = state.into();
384            let back = TaskState::try_from(proto).unwrap();
385            assert_eq!(back, state);
386        }
387    }
388
389    #[test]
390    fn try_from_i32() {
391        assert_eq!(TaskState::try_from(1i32).unwrap(), TaskState::Submitted);
392        assert_eq!(TaskState::try_from(2i32).unwrap(), TaskState::Working);
393        assert!(TaskState::try_from(0i32).is_err()); // UNSPECIFIED
394        assert!(TaskState::try_from(99i32).is_err()); // unknown
395    }
396
397    #[test]
398    fn can_transition_to_delegates_to_state_machine() {
399        assert!(TaskState::Submitted.can_transition_to(TaskState::Working));
400        assert!(!TaskState::Submitted.can_transition_to(TaskState::Completed));
401        assert!(!TaskState::Completed.can_transition_to(TaskState::Working));
402    }
403
404    #[test]
405    fn is_terminal_delegates() {
406        assert!(!TaskState::Working.is_terminal());
407        assert!(TaskState::Completed.is_terminal());
408    }
409
410    // TaskStatus tests
411
412    #[test]
413    fn task_status_constructor() {
414        let status = TaskStatus::new(TaskState::Working);
415        assert_eq!(status.state().unwrap(), TaskState::Working);
416    }
417
418    #[test]
419    fn task_status_with_message() {
420        let msg = crate::Message::new("s-msg", Role::Agent, vec![Part::text("working")]);
421        let status = TaskStatus::new(TaskState::Working).with_message(msg);
422        assert!(status.as_proto().message.is_some());
423    }
424
425    #[test]
426    fn task_status_try_from_proto_rejects_unspecified() {
427        let proto = pb::TaskStatus {
428            state: pb::TaskState::Unspecified.into(),
429            message: None,
430            timestamp: None,
431        };
432        assert!(TaskStatus::try_from(proto).is_err());
433    }
434
435    #[test]
436    fn task_status_serde_round_trip() {
437        let status = TaskStatus::new(TaskState::Submitted);
438        let json = serde_json::to_string(&status).unwrap();
439        let back: TaskStatus = serde_json::from_str(&json).unwrap();
440        assert_eq!(back.state().unwrap(), TaskState::Submitted);
441    }
442
443    // Task tests
444
445    #[test]
446    fn task_constructor() {
447        let task = Task::new("t-1", TaskStatus::new(TaskState::Submitted)).with_context_id("ctx-1");
448        assert_eq!(task.id(), "t-1");
449        assert_eq!(task.context_id(), "ctx-1");
450        assert_eq!(
451            task.status().unwrap().state().unwrap(),
452            TaskState::Submitted
453        );
454    }
455
456    #[test]
457    fn task_append_history_and_artifacts() {
458        let mut task = Task::new("t-2", TaskStatus::new(TaskState::Working));
459        task.append_message(crate::Message::new(
460            "m-1",
461            Role::User,
462            vec![Part::text("hi")],
463        ));
464        task.append_artifact(crate::Artifact::new("a-1", vec![Part::text("result")]));
465        assert_eq!(task.history().len(), 1);
466        assert_eq!(task.artifacts().len(), 1);
467    }
468
469    #[test]
470    fn task_merge_artifact_append_true_extends_existing_by_id() {
471        let mut task = Task::new("t-merge-1", TaskStatus::new(TaskState::Working));
472        task.append_artifact(crate::Artifact::new("a-1", vec![Part::text("chunk-1")]));
473
474        task.merge_artifact(
475            crate::Artifact::new("a-1", vec![Part::text("chunk-2")]),
476            true,
477            false,
478        );
479
480        assert_eq!(
481            task.artifacts().len(),
482            1,
483            "same-id append must not duplicate"
484        );
485        assert_eq!(task.artifacts()[0].parts.len(), 2);
486    }
487
488    #[test]
489    fn task_merge_artifact_append_true_no_match_adds_new() {
490        let mut task = Task::new("t-merge-2", TaskStatus::new(TaskState::Working));
491        task.append_artifact(crate::Artifact::new("a-1", vec![Part::text("x")]));
492
493        task.merge_artifact(
494            crate::Artifact::new("a-2", vec![Part::text("y")]),
495            true,
496            false,
497        );
498
499        assert_eq!(task.artifacts().len(), 2);
500    }
501
502    #[test]
503    fn task_merge_artifact_append_false_always_appends() {
504        let mut task = Task::new("t-merge-3", TaskStatus::new(TaskState::Working));
505        task.append_artifact(crate::Artifact::new("a-1", vec![Part::text("x")]));
506
507        task.merge_artifact(
508            crate::Artifact::new("a-1", vec![Part::text("y")]),
509            false,
510            true,
511        );
512
513        assert_eq!(
514            task.artifacts().len(),
515            2,
516            "append=false does not merge by id"
517        );
518    }
519
520    #[test]
521    fn task_try_from_proto_rejects_empty_id() {
522        let proto = pb::Task {
523            id: String::new(),
524            context_id: String::new(),
525            status: Some(pb::TaskStatus {
526                state: pb::TaskState::Submitted.into(),
527                message: None,
528                timestamp: None,
529            }),
530            artifacts: vec![],
531            history: vec![],
532            metadata: None,
533        };
534        assert!(Task::try_from(proto).is_err());
535    }
536
537    #[test]
538    fn task_try_from_proto_rejects_missing_status() {
539        // Status is REQUIRED per proto field_behavior
540        let proto = pb::Task {
541            id: "t-no-status".to_string(),
542            context_id: String::new(),
543            status: None,
544            artifacts: vec![],
545            history: vec![],
546            metadata: None,
547        };
548        assert!(Task::try_from(proto).is_err());
549    }
550
551    #[test]
552    fn task_try_from_proto_rejects_unspecified_state() {
553        let proto = pb::Task {
554            id: "t-bad".to_string(),
555            context_id: String::new(),
556            status: Some(pb::TaskStatus {
557                state: pb::TaskState::Unspecified.into(),
558                message: None,
559                timestamp: None,
560            }),
561            artifacts: vec![],
562            history: vec![],
563            metadata: None,
564        };
565        assert!(Task::try_from(proto).is_err());
566    }
567
568    #[test]
569    fn task_serde_round_trip() {
570        let task = Task::new("t-rt", TaskStatus::new(TaskState::Working)).with_context_id("ctx-rt");
571        let json = serde_json::to_string(&task).unwrap();
572        let back: Task = serde_json::from_str(&json).unwrap();
573        assert_eq!(back.id(), "t-rt");
574        assert_eq!(back.context_id(), "ctx-rt");
575    }
576
577    // Task helper tests
578
579    #[test]
580    fn task_complete_sets_completed_status() {
581        let mut task = Task::new("h-1", TaskStatus::new(TaskState::Submitted));
582        task.complete();
583        assert_eq!(
584            task.status().unwrap().state().unwrap(),
585            TaskState::Completed
586        );
587    }
588
589    #[test]
590    fn task_fail_sets_failed_status_with_message() {
591        let mut task = Task::new("h-2", TaskStatus::new(TaskState::Submitted));
592        task.fail("something went wrong");
593        let status = task.status().unwrap();
594        assert_eq!(status.state().unwrap(), TaskState::Failed);
595        // Status should have a message
596        assert!(status.as_proto().message.is_some());
597    }
598
599    #[test]
600    fn task_set_status_generic() {
601        let mut task = Task::new("h-3", TaskStatus::new(TaskState::Submitted));
602        task.set_status(TaskStatus::new(TaskState::Working));
603        assert_eq!(task.status().unwrap().state().unwrap(), TaskState::Working);
604    }
605
606    #[test]
607    fn task_push_text_artifact() {
608        let mut task = Task::new("h-4", TaskStatus::new(TaskState::Submitted));
609        task.push_text_artifact("art-1", "Result", "hello world");
610        assert_eq!(task.artifacts().len(), 1);
611        assert_eq!(task.artifacts()[0].artifact_id, "art-1");
612        assert_eq!(task.artifacts()[0].name, "Result");
613        assert_eq!(task.artifacts()[0].parts.len(), 1);
614    }
615}
turul_a2a_types/task.rs

turul_a2a_types/
task.rs
