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    async fn wait(self) -> Result<(AgentRuntime<P, Idle>, ChatResponse), OrchestrateError> {
60        // Wait for react_loop to complete
61        self.handle
62            .await
63            .map_err(|e| OrchestrateError::Session(e.into()))??;
64
65        let _ = self.env_state_tx.send(EnvStateEvent::SessionClosed);
66
67        // Aggregate filtered events into ChatResponse
68        let resp = aggregate_from_broadcast(self.event_rx).await?;
69
70        Ok((self.runtime.into_idle(), resp))
71    }
72}
73
74impl<P: ChatProvider + 'static> IntoFuture for SendHandle<P> {
75    type Output = Result<(AgentRuntime<P, Idle>, ChatResponse), OrchestrateError>;
76    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send>>;
77
78    fn into_future(self) -> Self::IntoFuture {
79        Box::pin(async move { self.wait().await })
80    }
81}
82
83// ═══════════════════════════════════════════════════════════════
84// SendStreamHandle — stateful send + streaming
85// ═══════════════════════════════════════════════════════════════
86
87/// Handle returned by [`Agent::send_stream`](crate::Agent::send_stream).
88///
89/// Like [`SendHandle`] but also provides `recv()` for per-event streaming.
90pub struct SendStreamHandle<P: ChatProvider> {
91    pub(crate) runtime: AgentRuntime<P, Acquired>,
92    pub(crate) handle: JoinHandle<anyhow::Result<()>>,
93    pub(crate) event_rx: broadcast::Receiver<AgentEvent>,
94    pub(crate) stream_rx: mpsc::Receiver<AgentEvent>,
95    pub(crate) env_state_tx: broadcast::Sender<EnvStateEvent>,
96}
97
98impl<P: ChatProvider> SendStreamHandle<P> {
99    /// Receive the next streaming event.
100    pub async fn recv(&mut self) -> Option<AgentEvent> {
101        self.stream_rx.recv().await
102    }
103
104    /// Query session context while streaming is in progress.
105    pub async fn session_context(&self) -> Vec<JsonValue> {
106        let (respond, rx) = tokio::sync::oneshot::channel();
107        let _ = self
108            .runtime
109            .session_tx
110            .send(SessionCmd::FetchContext { respond });
111        rx.await.unwrap_or_default()
112    }
113
114    async fn wait(self) -> Result<(AgentRuntime<P, Idle>, ChatResponse), OrchestrateError> {
115        self.handle
116            .await
117            .map_err(|e| OrchestrateError::Session(e.into()))??;
118        let _ = self.env_state_tx.send(EnvStateEvent::SessionClosed);
119        let resp = aggregate_from_broadcast(self.event_rx).await?;
120        Ok((self.runtime.into_idle(), resp))
121    }
122}
123
124impl<P: ChatProvider + 'static> IntoFuture for SendStreamHandle<P> {
125    type Output = Result<(AgentRuntime<P, Idle>, ChatResponse), OrchestrateError>;
126    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send>>;
127
128    fn into_future(self) -> Self::IntoFuture {
129        Box::pin(async move { self.wait().await })
130    }
131}
132
133// ═══════════════════════════════════════════════════════════════
134// FireStreamHandle — stateless stream (no runtime binding)
135// ═══════════════════════════════════════════════════════════════
136
137/// Handle returned by [`Agent::fire_stream`](crate::Agent::fire_stream).
138///
139/// Unlike `SendHandle`, this does NOT hold an `AgentRuntime` — the session is
140/// temporary. Provides `recv()` for streaming and `wait()` (`IntoFuture`) for
141/// the final [`ChatResponse`].
142pub struct FireStreamHandle {
143    pub(crate) handle: JoinHandle<anyhow::Result<()>>,
144    pub(crate) event_rx: broadcast::Receiver<AgentEvent>,
145    pub(crate) stream_rx: mpsc::Receiver<AgentEvent>,
146    pub(crate) env_state_tx: broadcast::Sender<EnvStateEvent>,
147}
148
149impl FireStreamHandle {
150    /// Receive the next streaming event.
151    pub async fn recv(&mut self) -> Option<AgentEvent> {
152        self.stream_rx.recv().await
153    }
154
155    async fn wait(self) -> Result<ChatResponse, OrchestrateError> {
156        self.handle
157            .await
158            .map_err(|e| OrchestrateError::Session(e.into()))??;
159        let _ = self.env_state_tx.send(EnvStateEvent::SessionClosed);
160        aggregate_from_broadcast(self.event_rx).await
161    }
162}
163
164impl IntoFuture for FireStreamHandle {
165    type Output = Result<ChatResponse, OrchestrateError>;
166    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send>>;
167
168    fn into_future(self) -> Self::IntoFuture {
169        Box::pin(async move { self.wait().await })
170    }
171}
172
173// ═══════════════════════════════════════════════════════════════
174// Internal helpers
175// ═══════════════════════════════════════════════════════════════
176
177/// Aggregate filtered events from a broadcast receiver into a ChatResponse.
178async fn aggregate_from_broadcast(
179    mut event_rx: broadcast::Receiver<AgentEvent>,
180) -> Result<ChatResponse, OrchestrateError> {
181    let mut content = String::new();
182    let mut tool_calls = Vec::new();
183    let mut iterations = 0usize;
184    let mut finish_reason: Option<String> = None;
185    let mut pending: Vec<(Arc<str>, String, serde_json::Value)> = Vec::new();
186
187    loop {
188        match event_rx.recv().await {
189            Ok(AgentEvent::Text(t)) => content = t,
190            Ok(AgentEvent::ToolCallRequest {
191                call_id,
192                name,
193                args,
194                ..
195            }) => {
196                pending.push((call_id, name, args));
197            }
198            Ok(AgentEvent::ToolCallResult {
199                call_id,
200                name: _,
201                result,
202            }) => {
203                if let Some(pos) = pending.iter().position(|(id, _, _)| *id == call_id) {
204                    let (_, name, args) = pending.remove(pos);
205                    tool_calls.push(ToolCallInfo { name, args, result });
206                }
207            }
208            Ok(AgentEvent::TurnStart) => iterations += 1,
209            Ok(AgentEvent::TurnEnd { finish_reason: fr }) => finish_reason = fr,
210            Ok(AgentEvent::Done) => break,
211            Err(broadcast::error::RecvError::Closed) => break,
212            Err(broadcast::error::RecvError::Lagged(_)) => continue,
213            _ => {}
214        }
215    }
216
217    Ok(ChatResponse {
218        content,
219        tool_calls,
220        iterations,
221        finish_reason,
222    })
223}
funera_orchestrate/send_handle.rs

funera_orchestrate/
send_handle.rs
