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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
//! Spawn-and-continue agent handles.
//!
//! [`AgentHandle`] wraps a spawned agent task, providing status polling,
//! cancellation, and result retrieval without blocking the caller.
use std::sync::{Arc, Mutex, PoisonError};
use tokio_util::sync::CancellationToken;
use crate::agent::Agent;
use crate::error::AgentError;
use crate::task_core::{TaskCore, resolve_status};
use crate::types::{AgentMessage, AgentResult, ContentBlock, LlmMessage, UserMessage};
use crate::util::now_timestamp;
/// The lifecycle status of a spawned agent task.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AgentStatus {
/// The agent task is still executing.
Running,
/// The agent task completed successfully.
Completed,
/// The agent task failed with an error.
Failed,
/// The agent task was cancelled via [`AgentHandle::cancel`].
Cancelled,
}
/// A handle to a spawned agent task.
///
/// Created via [`AgentHandle::spawn`] or [`AgentHandle::spawn_text`], which
/// move an [`Agent`] into a background tokio task. The handle allows the caller
/// to poll status, cancel, and retrieve the final result.
pub struct AgentHandle {
core: TaskCore,
}
impl AgentHandle {
/// Spawn an agent task with the given input messages.
///
/// Takes ownership of the `Agent` and moves it into a tokio task.
/// Returns a handle that can be used to poll status, cancel, or await
/// the result.
pub fn spawn(mut agent: Agent, input: Vec<AgentMessage>) -> Self {
let cancellation_token = CancellationToken::new();
let status = Arc::new(Mutex::new(AgentStatus::Running));
let status_clone = Arc::clone(&status);
let token_clone = cancellation_token.clone();
let join_handle = tokio::spawn(async move {
let result = tokio::select! {
result = agent.prompt_async(input) => result,
() = token_clone.cancelled() => {
agent.abort();
Err(AgentError::Aborted)
}
};
*status_clone.lock().unwrap_or_else(PoisonError::into_inner) = resolve_status(&result);
result
});
Self {
core: TaskCore::new(join_handle, cancellation_token, status),
}
}
/// Convenience wrapper that spawns an agent with a single text message.
///
/// Equivalent to calling [`spawn`](Self::spawn) with a single
/// [`UserMessage`] containing the given text.
pub fn spawn_text(agent: Agent, text: impl Into<String>) -> Self {
let msg = AgentMessage::Llm(LlmMessage::User(UserMessage {
content: vec![ContentBlock::Text { text: text.into() }],
timestamp: now_timestamp(),
cache_hint: None,
}));
Self::spawn(agent, vec![msg])
}
/// Returns the current status of the spawned agent task.
pub fn status(&self) -> AgentStatus {
self.core.status()
}
/// Returns `true` if the agent task is no longer running.
pub fn is_done(&self) -> bool {
self.core.is_done()
}
/// Request cancellation of the spawned agent task.
///
/// This is non-blocking. The task will transition to `Cancelled` status
/// asynchronously.
pub fn cancel(&self) {
self.core.cancel();
}
/// Consume the handle and await the final result.
///
/// If the task panicked, returns an [`AgentError::StreamError`] wrapping
/// the panic message.
pub async fn result(self) -> Result<AgentResult, AgentError> {
self.core.result().await
}
/// Check if the task is finished and, if so, return the result without
/// blocking.
///
/// Returns `None` if the task is still running. Once a result is returned,
/// subsequent calls will return `None`.
pub fn try_result(&mut self) -> Option<Result<AgentResult, AgentError>> {
self.core.try_result()
}
}
impl std::fmt::Debug for AgentHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AgentHandle")
.field("status", &self.status())
.field("join_handle", &self.core.join_handle)
.field("cancellation_token", &self.core.cancellation_token)
.finish()
}
}