swink_agent/
handle.rs

1//! Spawn-and-continue agent handles.
2//!
3//! [`AgentHandle`] wraps a spawned agent task, providing status polling,
4//! cancellation, and result retrieval without blocking the caller.
5
6use std::sync::{Arc, Mutex, PoisonError};
7
8use tokio_util::sync::CancellationToken;
9
10use crate::agent::Agent;
11use crate::error::AgentError;
12use crate::task_core::{TaskCore, resolve_status};
13use crate::types::{AgentMessage, AgentResult, ContentBlock, LlmMessage, UserMessage};
14use crate::util::now_timestamp;
15
16/// The lifecycle status of a spawned agent task.
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum AgentStatus {
19    /// The agent task is still executing.
20    Running,
21    /// The agent task completed successfully.
22    Completed,
23    /// The agent task failed with an error.
24    Failed,
25    /// The agent task was cancelled via [`AgentHandle::cancel`].
26    Cancelled,
27}
28
29/// A handle to a spawned agent task.
30///
31/// Created via [`AgentHandle::spawn`] or [`AgentHandle::spawn_text`], which
32/// move an [`Agent`] into a background tokio task. The handle allows the caller
33/// to poll status, cancel, and retrieve the final result.
34pub struct AgentHandle {
35    core: TaskCore,
36}
37
38impl AgentHandle {
39    /// Spawn an agent task with the given input messages.
40    ///
41    /// Takes ownership of the `Agent` and moves it into a tokio task.
42    /// Returns a handle that can be used to poll status, cancel, or await
43    /// the result.
44    pub fn spawn(mut agent: Agent, input: Vec<AgentMessage>) -> Self {
45        let cancellation_token = CancellationToken::new();
46        let status = Arc::new(Mutex::new(AgentStatus::Running));
47        let status_clone = Arc::clone(&status);
48        let token_clone = cancellation_token.clone();
49
50        let join_handle = tokio::spawn(async move {
51            let result = tokio::select! {
52                result = agent.prompt_async(input) => result,
53                () = token_clone.cancelled() => {
54                    agent.abort();
55                    Err(AgentError::Aborted)
56                }
57            };
58            *status_clone.lock().unwrap_or_else(PoisonError::into_inner) = resolve_status(&result);
59            result
60        });
61
62        Self {
63            core: TaskCore::new(join_handle, cancellation_token, status),
64        }
65    }
66
67    /// Convenience wrapper that spawns an agent with a single text message.
68    ///
69    /// Equivalent to calling [`spawn`](Self::spawn) with a single
70    /// [`UserMessage`] containing the given text.
71    pub fn spawn_text(agent: Agent, text: impl Into<String>) -> Self {
72        let msg = AgentMessage::Llm(LlmMessage::User(UserMessage {
73            content: vec![ContentBlock::Text { text: text.into() }],
74            timestamp: now_timestamp(),
75            cache_hint: None,
76        }));
77        Self::spawn(agent, vec![msg])
78    }
79
80    /// Returns the current status of the spawned agent task.
81    pub fn status(&self) -> AgentStatus {
82        self.core.status()
83    }
84
85    /// Returns `true` if the agent task is no longer running.
86    pub fn is_done(&self) -> bool {
87        self.core.is_done()
88    }
89
90    /// Request cancellation of the spawned agent task.
91    ///
92    /// This is non-blocking. The task will transition to `Cancelled` status
93    /// asynchronously.
94    pub fn cancel(&self) {
95        self.core.cancel();
96    }
97
98    /// Consume the handle and await the final result.
99    ///
100    /// If the task panicked, returns an [`AgentError::StreamError`] wrapping
101    /// the panic message.
102    pub async fn result(self) -> Result<AgentResult, AgentError> {
103        self.core.result().await
104    }
105
106    /// Check if the task is finished and, if so, return the result without
107    /// blocking.
108    ///
109    /// Returns `None` if the task is still running. Once a result is returned,
110    /// subsequent calls will return `None`.
111    pub fn try_result(&mut self) -> Option<Result<AgentResult, AgentError>> {
112        self.core.try_result()
113    }
114}
115
116impl std::fmt::Debug for AgentHandle {
117    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
118        f.debug_struct("AgentHandle")
119            .field("status", &self.status())
120            .field("join_handle", &self.core.join_handle)
121            .field("cancellation_token", &self.core.cancellation_token)
122            .finish()
123    }
124}
swink_agent/handle.rs

swink_agent/
handle.rs
