Skip to main content

funera_orchestrate/
send_handle.rs

1use std::future::{Future, IntoFuture};
2use std::pin::Pin;
3use std::sync::Arc;
4
5use serde_json::Value as JsonValue;
6use tokio::sync::{broadcast, mpsc};
7use tokio::task::JoinHandle;
8
9use funera_core::chat::session::SessionCmd;
10use funera_core::event_bus::env_state_bus::EnvStateEvent;
11use funera_core::provider::ChatProvider;
12
13use crate::error::OrchestrateError;
14use crate::event::AgentEvent;
15use crate::response::{ChatResponse, ToolCallInfo};
16use crate::runtime::{Acquired, AgentRuntime, Idle};
17
18// ═══════════════════════════════════════════════════════════════
19// SendHandle — stateful send (consumes Idle, returns Idle on wait)
20// ═══════════════════════════════════════════════════════════════
21
22/// Handle returned by [`Agent::send`](crate::Agent::send).
23///
24/// Holds an `AgentRuntime<Acquired>` — cannot send again until `wait()` completes.
25/// The react_loop runs in a background task; you can query session context via
26/// [`session_context`](Self::session_context) while it is in progress.
27///
28/// # Awaiting
29///
30/// ```rust,no_run
31/// # use funera_orchestrate::{Agent, AgentRuntime, DeepSeekProvider};
32/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
33/// # let agent = Agent::builder().build();
34/// # let rt = AgentRuntime::<DeepSeekProvider>::builder()
35/// #     .api_key(std::env::var("DEEPSEEK_API_KEY")?).build()?;
36/// let handle = agent.send("Hello", rt).await?;
37/// let (_rt, resp) = handle.await?;    // IntoFuture → (Idle, ChatResponse)
38/// # Ok(())
39/// # }
40/// ```
41pub struct SendHandle<P: ChatProvider> {
42    pub(crate) runtime: AgentRuntime<P, Acquired>,
43    pub(crate) handle: JoinHandle<anyhow::Result<()>>,
44    pub(crate) event_rx: broadcast::Receiver<AgentEvent>,
45    pub(crate) env_state_tx: broadcast::Sender<EnvStateEvent>,
46}
47
48impl<P: ChatProvider> SendHandle<P> {
49    /// Query session context while react_loop is running.
50    pub async fn session_context(&self) -> Vec<JsonValue> {
51        let (respond, rx) = tokio::sync::oneshot::channel();
52        let _ = self
53            .runtime
54            .session_tx
55            .send(SessionCmd::FetchContext { respond });
56        rx.await.unwrap_or_default()
57    }
58
59    /// Returns a cloneable [`ApprovalHandle`] for approving tool calls
60    /// during react_loop execution.
61    #[cfg(all(feature = "tool", feature = "security"))]
62    pub fn approval_handle(&self) -> ApprovalHandle {
63        ApprovalHandle::new(self.runtime.env_cmd_tx.clone())
64    }
65
66    async fn wait(self) -> Result<(AgentRuntime<P, Idle>, ChatResponse), OrchestrateError> {
67        // Wait for react_loop to complete
68        self.handle
69            .await
70            .map_err(|e| OrchestrateError::Session(e.into()))??;
71
72        let _ = self.env_state_tx.send(EnvStateEvent::SessionClosed);
73
74        // Aggregate filtered events into ChatResponse
75        let resp = aggregate_from_broadcast(self.event_rx).await?;
76
77        Ok((self.runtime.into_idle(), resp))
78    }
79}
80
81impl<P: ChatProvider + 'static> IntoFuture for SendHandle<P> {
82    type Output = Result<(AgentRuntime<P, Idle>, ChatResponse), OrchestrateError>;
83    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send>>;
84
85    fn into_future(self) -> Self::IntoFuture {
86        Box::pin(async move { self.wait().await })
87    }
88}
89
90// ═══════════════════════════════════════════════════════════════
91// SendStreamHandle — stateful send + streaming
92// ═══════════════════════════════════════════════════════════════
93
94/// Handle returned by [`Agent::send_stream`](crate::Agent::send_stream).
95///
96/// Like [`SendHandle`] but also provides `recv()` for per-event streaming.
97pub struct SendStreamHandle<P: ChatProvider> {
98    pub(crate) runtime: AgentRuntime<P, Acquired>,
99    pub(crate) handle: JoinHandle<anyhow::Result<()>>,
100    pub(crate) event_rx: broadcast::Receiver<AgentEvent>,
101    pub(crate) stream_rx: mpsc::Receiver<AgentEvent>,
102    pub(crate) env_state_tx: broadcast::Sender<EnvStateEvent>,
103}
104
105impl<P: ChatProvider> SendStreamHandle<P> {
106    /// Receive the next streaming event.
107    pub async fn recv(&mut self) -> Option<AgentEvent> {
108        self.stream_rx.recv().await
109    }
110
111    /// Query session context while streaming is in progress.
112    pub async fn session_context(&self) -> Vec<JsonValue> {
113        let (respond, rx) = tokio::sync::oneshot::channel();
114        let _ = self
115            .runtime
116            .session_tx
117            .send(SessionCmd::FetchContext { respond });
118        rx.await.unwrap_or_default()
119    }
120
121    /// Returns a cloneable [`ApprovalHandle`] for approving tool calls
122    /// during react_loop execution.
123    #[cfg(all(feature = "tool", feature = "security"))]
124    pub fn approval_handle(&self) -> ApprovalHandle {
125        ApprovalHandle::new(self.runtime.env_cmd_tx.clone())
126    }
127
128    async fn wait(self) -> Result<(AgentRuntime<P, Idle>, ChatResponse), OrchestrateError> {
129        self.handle
130            .await
131            .map_err(|e| OrchestrateError::Session(e.into()))??;
132        let _ = self.env_state_tx.send(EnvStateEvent::SessionClosed);
133        let resp = aggregate_from_broadcast(self.event_rx).await?;
134        Ok((self.runtime.into_idle(), resp))
135    }
136}
137
138impl<P: ChatProvider + 'static> IntoFuture for SendStreamHandle<P> {
139    type Output = Result<(AgentRuntime<P, Idle>, ChatResponse), OrchestrateError>;
140    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send>>;
141
142    fn into_future(self) -> Self::IntoFuture {
143        Box::pin(async move { self.wait().await })
144    }
145}
146
147// ═══════════════════════════════════════════════════════════════
148// FireStreamHandle — stateless stream (no runtime binding)
149// ═══════════════════════════════════════════════════════════════
150
151/// Handle returned by [`Agent::fire_stream`](crate::Agent::fire_stream).
152///
153/// Unlike `SendHandle`, this does NOT hold an `AgentRuntime` — the session is
154/// temporary. Provides `recv()` for streaming and `wait()` (`IntoFuture`) for
155/// the final [`ChatResponse`].
156pub struct FireStreamHandle {
157    pub(crate) handle: JoinHandle<anyhow::Result<()>>,
158    pub(crate) event_rx: broadcast::Receiver<AgentEvent>,
159    pub(crate) stream_rx: mpsc::Receiver<AgentEvent>,
160    pub(crate) env_state_tx: broadcast::Sender<EnvStateEvent>,
161}
162
163impl FireStreamHandle {
164    /// Receive the next streaming event.
165    pub async fn recv(&mut self) -> Option<AgentEvent> {
166        self.stream_rx.recv().await
167    }
168
169    async fn wait(self) -> Result<ChatResponse, OrchestrateError> {
170        self.handle
171            .await
172            .map_err(|e| OrchestrateError::Session(e.into()))??;
173        let _ = self.env_state_tx.send(EnvStateEvent::SessionClosed);
174        aggregate_from_broadcast(self.event_rx).await
175    }
176}
177
178impl IntoFuture for FireStreamHandle {
179    type Output = Result<ChatResponse, OrchestrateError>;
180    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send>>;
181
182    fn into_future(self) -> Self::IntoFuture {
183        Box::pin(async move { self.wait().await })
184    }
185}
186
187// ═══════════════════════════════════════════════════════════════
188// Internal helpers
189// ═══════════════════════════════════════════════════════════════
190
191/// Aggregate filtered events from a broadcast receiver into a ChatResponse.
192async fn aggregate_from_broadcast(
193    mut event_rx: broadcast::Receiver<AgentEvent>,
194) -> Result<ChatResponse, OrchestrateError> {
195    let mut content = String::new();
196    let mut tool_calls = Vec::new();
197    let mut iterations = 0usize;
198    let mut finish_reason: Option<String> = None;
199    let mut pending: Vec<(Arc<str>, String, serde_json::Value)> = Vec::new();
200
201    loop {
202        match event_rx.recv().await {
203            Ok(AgentEvent::Text(t)) => content = t,
204            Ok(AgentEvent::ToolCallRequest {
205                call_id,
206                name,
207                args,
208                ..
209            }) => {
210                pending.push((call_id, name, args));
211            }
212            Ok(AgentEvent::ToolCallResult {
213                call_id,
214                name: _,
215                result,
216            }) => {
217                if let Some(pos) = pending.iter().position(|(id, _, _)| *id == call_id) {
218                    let (_, name, args) = pending.remove(pos);
219                    tool_calls.push(ToolCallInfo { name, args, result });
220                }
221            }
222            Ok(AgentEvent::TurnStart) => iterations += 1,
223            Ok(AgentEvent::TurnEnd { finish_reason: fr }) => finish_reason = fr,
224            Ok(AgentEvent::Done) => break,
225            Err(broadcast::error::RecvError::Closed) => break,
226            Err(broadcast::error::RecvError::Lagged(_)) => continue,
227            _ => {}
228        }
229    }
230
231    Ok(ChatResponse {
232        content,
233        tool_calls,
234        iterations,
235        finish_reason,
236    })
237}
238
239// ═══════════════════════════════════════════════════════════════
240// ApprovalHandle — tool call approval for spawned tasks
241// ═══════════════════════════════════════════════════════════════
242
243/// A lightweight, cloneable handle for approving or rejecting tool calls
244/// during agent execution.
245///
246/// Obtain one **before** calling [`send`](crate::Agent::send) /
247/// [`send_stream`](crate::Agent::send_stream) via
248/// [`AgentRuntime::approval_handle`](crate::AgentRuntime::approval_handle),
249/// or afterwards from
250/// [`SendHandle::approval_handle`] / [`SendStreamHandle::approval_handle`].
251///
252/// Because `ApprovalHandle` implements [`Clone`], it can be freely moved
253/// into spawned tasks for background approval while the main task awaits
254/// the send handle's result.
255#[cfg(all(feature = "tool", feature = "security"))]
256#[derive(Clone)]
257pub struct ApprovalHandle {
258    env_cmd_tx: mpsc::UnboundedSender<funera_core::env_actor::EnvCmd>,
259}
260
261#[cfg(all(feature = "tool", feature = "security"))]
262impl ApprovalHandle {
263    pub(crate) fn new(env_cmd_tx: mpsc::UnboundedSender<funera_core::env_actor::EnvCmd>) -> Self {
264        Self { env_cmd_tx }
265    }
266
267    /// Approve or reject a pending tool call identified by `call_id`.
268    ///
269    /// `call_id` is typically obtained from the
270    /// [`on_approval_required`](crate::AgentRuntimeBuilder::on_approval_required)
271    /// callback. Returns `Ok(())` if the approval was processed, or
272    /// `Err(msg)` if the call_id was not found or the env actor has died.
273    pub async fn approve_tool_call(&self, call_id: &str, approved: bool) -> Result<(), String> {
274        let (respond, rx) = tokio::sync::oneshot::channel();
275        let _ = self
276            .env_cmd_tx
277            .send(funera_core::env_actor::EnvCmd::ApproveToolCall {
278                call_id: call_id.to_string(),
279                approved,
280                respond,
281            });
282        rx.await.unwrap_or(Err("env actor died".into()))
283    }
284}