1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
//! Shared spawned-task lifecycle core.
//!
//! [`TaskCore`] owns the cancellation token, join handle, and status mutex
//! that [`AgentHandle`](crate::AgentHandle) and
//! [`OrchestratedHandle`](crate::OrchestratedHandle) both need. Each handle
//! composes a `TaskCore` and delegates lifecycle methods to it.
use std::sync::{Arc, Mutex, PoisonError};
use futures::FutureExt;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use crate::error::AgentError;
use crate::handle::AgentStatus;
use crate::types::AgentResult;
/// Shared lifecycle core for a spawned agent task.
///
/// Owns the cancellation token, join handle, and status storage that both
/// `AgentHandle` and `OrchestratedHandle` use identically. The status
/// transition logic (`Ok → Completed`, `Aborted → Cancelled`, `Err → Failed`)
/// lives in [`resolve_status`] so it is defined exactly once.
pub(crate) struct TaskCore {
pub(crate) join_handle: Option<JoinHandle<Result<AgentResult, AgentError>>>,
pub(crate) cancellation_token: CancellationToken,
pub(crate) status: Arc<Mutex<AgentStatus>>,
}
impl TaskCore {
/// Create a new task core with `Running` status and a fresh cancellation token.
pub(crate) const fn new(
join_handle: JoinHandle<Result<AgentResult, AgentError>>,
cancellation_token: CancellationToken,
status: Arc<Mutex<AgentStatus>>,
) -> Self {
Self {
join_handle: Some(join_handle),
cancellation_token,
status,
}
}
/// Returns the current status of the spawned task.
pub(crate) fn status(&self) -> AgentStatus {
*self.status.lock().unwrap_or_else(PoisonError::into_inner)
}
/// Returns `true` if the task is no longer running.
pub(crate) fn is_done(&self) -> bool {
self.status() != AgentStatus::Running
}
/// Request cancellation of the spawned task (non-blocking).
pub(crate) fn cancel(&self) {
self.cancellation_token.cancel();
}
/// Consume the core and await the final result.
pub(crate) async fn result(mut self) -> Result<AgentResult, AgentError> {
match self.join_handle.take() {
Some(handle) => match handle.await {
Ok(result) => result,
Err(join_err) => Err(AgentError::stream(join_err)),
},
None => Err(AgentError::Aborted),
}
}
/// Check if the task is finished and, if so, return the result without
/// blocking. Returns `None` if still running. Once returned, subsequent
/// calls yield `None`.
pub(crate) fn try_result(&mut self) -> Option<Result<AgentResult, AgentError>> {
let finished = self
.join_handle
.as_ref()
.is_some_and(JoinHandle::is_finished);
if finished {
let handle = self.join_handle.take()?;
let join_result = handle.now_or_never()?;
Some(match join_result {
Ok(result) => result,
Err(join_err) => Err(AgentError::stream(join_err)),
})
} else {
None
}
}
}
/// Map a task result to its terminal [`AgentStatus`].
///
/// This is the single source of truth for the status transition that both
/// `AgentHandle::spawn` and `run_agent_loop` apply after the spawned future
/// completes.
pub(crate) const fn resolve_status(result: &Result<AgentResult, AgentError>) -> AgentStatus {
match result {
Ok(_) => AgentStatus::Completed,
Err(AgentError::Aborted) => AgentStatus::Cancelled,
Err(_) => AgentStatus::Failed,
}
}
#[cfg(test)]
#[path = "task_core_tests.rs"]
mod tests;