funera_orchestrate/
send_handle.rs1use 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
18pub 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 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 #[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 self.handle
69 .await
70 .map_err(|e| OrchestrateError::Session(e.into()))??;
71
72 let _ = self.env_state_tx.send(EnvStateEvent::SessionClosed);
73
74 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
90pub 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 pub async fn recv(&mut self) -> Option<AgentEvent> {
108 self.stream_rx.recv().await
109 }
110
111 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 #[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
147pub 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 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
187async 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#[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 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}