agent_base/engine/runtime/
mod.rs

1use std::sync::Arc;
2
3use crate::engine::context::ContextWindowManager;
4use crate::engine::middleware::MiddlewareRef;
5use crate::engine::session_store::SessionStore;
6use crate::engine::AgentSession;
7use crate::types::{AgentConfig, AgentError, AgentEvent, AgentResult, CheckpointData, CheckpointStep, MessageRole, RunOutcome, SessionId};
8
9use super::approval::ApprovalHandler;
10
11mod event_bus;
12mod llm_engine;
13mod plan;
14mod react_loop;
15mod session_manager;
16mod tool_engine;
17
18pub(super) const DEFAULT_MAX_TURNS: u32 = 50;
19
20pub use event_bus::EventBus;
21pub use llm_engine::LlmEngine;
22pub use session_manager::SessionManager;
23pub use tool_engine::ToolEngine;
24
25pub struct AgentRuntime {
26    pub(crate) config: AgentConfig,
27    pub(crate) llm_engine: LlmEngine,
28    pub(crate) tool_engine: ToolEngine,
29    pub(crate) session_manager: SessionManager,
30    pub(crate) event_bus: EventBus,
31    pub(crate) context_manager: Option<ContextWindowManager>,
32    pub(crate) middlewares: Vec<MiddlewareRef>,
33}
34
35impl AgentRuntime {
36    pub async fn create_session(&self) -> SessionId {
37        self.session_manager.create_session(self.config.system_prompt.as_deref()).await
38    }
39
40    pub async fn restore_session(&self, session_id: &SessionId) -> Option<AgentSession> {
41        self.session_manager.restore_session(session_id).await
42    }
43
44    pub async fn session(&self, session_id: &SessionId) -> Option<AgentSession> {
45        self.session_manager.session(session_id).await
46    }
47
48    pub async fn session_or_err(&self, session_id: &SessionId) -> AgentResult<AgentSession> {
49        self.session_manager.session_or_err(session_id).await
50    }
51
52    pub async fn with_session_mut<F, R>(&self, session_id: &SessionId, f: F) -> AgentResult<R>
53    where
54        F: FnOnce(&mut AgentSession) -> R,
55    {
56        self.session_manager.with_session_mut(session_id, f).await
57    }
58
59    pub fn emit_event(&self, event: AgentEvent) {
60        self.event_bus.emit(event);
61    }
62
63    pub fn subscribe_events(&self) -> tokio::sync::broadcast::Receiver<AgentEvent> {
64        self.event_bus.subscribe()
65    }
66
67    pub fn event_bus(&self) -> &EventBus {
68        &self.event_bus
69    }
70
71    pub fn session_manager(&self) -> &SessionManager {
72        &self.session_manager
73    }
74
75    pub fn llm_engine(&self) -> &LlmEngine {
76        &self.llm_engine
77    }
78
79    pub fn tool_engine(&self) -> &ToolEngine {
80        &self.tool_engine
81    }
82
83    pub fn tool_engine_mut(&mut self) -> &mut ToolEngine {
84        &mut self.tool_engine
85    }
86
87    pub fn client(&self) -> Arc<dyn crate::llm::LlmClient> {
88        self.llm_engine.client.clone()
89    }
90
91    pub fn tools_mut(&mut self) -> &mut crate::tool::ToolRegistry {
92        self.tool_engine.tools_mut()
93    }
94
95    pub fn config(&self) -> &AgentConfig {
96        &self.config
97    }
98
99    pub fn approval_handler(&self) -> Option<&Arc<dyn ApprovalHandler>> {
100        self.tool_engine.approval_handler()
101    }
102
103    pub async fn cached_approval(&self, session_id: &SessionId, action_key: &str) -> bool {
104        self.session_manager.cached_approval(session_id, action_key).await
105    }
106
107    pub async fn cache_approval(&self, session_id: &SessionId, action_key: String) {
108        self.session_manager.cache_approval(session_id, action_key).await
109    }
110
111    pub async fn save_checkpoint(&self, session_id: &SessionId, checkpoint: CheckpointData) -> AgentResult<()> {
112        self.emit_event(AgentEvent::Checkpoint {
113            session_id: session_id.clone(),
114            checkpoint,
115        });
116        Ok(())
117    }
118
119    pub async fn load_checkpoint(&self, _session_id: &SessionId, _checkpoint: &CheckpointData) -> AgentResult<Option<CheckpointData>> {
120        Ok(None)
121    }
122
123    pub async fn resume_from_checkpoint<F>(
124        &self,
125        checkpoint: CheckpointData,
126        mut on_event: F,
127    ) -> AgentResult<RunOutcome>
128    where
129        F: FnMut(AgentEvent) -> AgentResult<()> + Send,
130    {
131        let session_id = checkpoint.session_id.clone();
132        let user_input = checkpoint.user_input.clone();
133        let turn_count = checkpoint.turn_count;
134
135        tracing::info!(session_id = session_id.id, turn_count, step = ?checkpoint.step, "resuming from checkpoint");
136
137        let mut event_rx = self.subscribe_events();
138        let tool_definitions = self.tool_engine.definitions();
139
140        if let CheckpointStep::BeforeToolCalls { tool_calls } = checkpoint.step {
141            match self.handle_tool_calls(&session_id, &tool_calls, &mut event_rx, &mut on_event).await {
142                Ok(react_loop::ToolCallResult::Continue) => {}
143                Ok(react_loop::ToolCallResult::Break) => {
144                    self.emit_event(AgentEvent::RunFinished { session_id: session_id.clone() });
145                    EventBus::drain_async_events(&mut event_rx, &mut on_event)?;
146                    return Ok(RunOutcome::Completed);
147                }
148                Err(e) => {
149                    if let Some(outcome) = self
150                        .handle_tool_error(&session_id, &tool_calls, e, &mut event_rx, &mut on_event)
151                        .await?
152                    {
153                        return Ok(outcome);
154                    }
155                }
156            }
157        }
158
159        let (outcome, _final_turn_count) = self
160            .run_turn_loop(
161                &session_id,
162                &user_input,
163                &tool_definitions,
164                turn_count,
165                &mut event_rx,
166                &mut on_event,
167            )
168            .await?;
169
170        Ok(outcome)
171    }
172
173    pub async fn run<F>(
174        &self,
175        session_id: SessionId,
176        mut on_event: F,
177    ) -> AgentResult<RunOutcome>
178    where
179        F: FnMut(AgentEvent) -> AgentResult<()> + Send,
180    {
181        let span = tracing::info_span!("agent_run", session_id = session_id.id);
182        let _enter = span.enter();
183
184        let mut event_rx = self.subscribe_events();
185
186        if let Err(e) = self.validate_session(&session_id).await {
187            self.emit_event(AgentEvent::RunFinished { session_id: session_id.clone() });
188            EventBus::drain_async_events(&mut event_rx, &mut on_event)?;
189            return Err(e);
190        }
191
192        let tool_definitions = self.tool_engine.definitions();
193        let user_input_owned = self.with_session_mut(&session_id, |session| {
194            session.chat_messages().last()
195                .and_then(|m| match m {
196                    crate::types::ChatMessage::User { content, .. } => Some(content.clone()),
197                    _ => None,
198                })
199                .unwrap_or_default()
200        }).await?;
201
202        let (outcome, _turn_count) = self
203            .run_turn_loop(
204                &session_id,
205                &user_input_owned,
206                &tool_definitions,
207                0,
208                &mut event_rx,
209                &mut on_event,
210            )
211            .await?;
212
213        self.emit_event(AgentEvent::RunFinished { session_id: session_id.clone() });
214        EventBus::drain_async_events(&mut event_rx, &mut on_event)?;
215
216        Ok(outcome)
217    }
218
219    pub async fn run_turn_with_handler<F>(
220        &self,
221        session_id: SessionId,
222        user_input: &str,
223        mut on_event: F,
224    ) -> AgentResult<RunOutcome>
225    where
226        F: FnMut(AgentEvent) -> AgentResult<()> + Send,
227    {
228        let span = tracing::Span::current();
229        let _guard = span.enter();
230        tracing::info!(session_id = session_id.id, user_input = %user_input, "agent turn start");
231        drop(_guard);
232
233        let mut event_rx = self.subscribe_events();
234        let tool_definitions = self.tool_engine.definitions();
235
236        let user_input_owned = self.apply_user_message_mw(&session_id, user_input.to_string()).await?;
237
238        self.with_session_mut(&session_id, |session| {
239            session.push_message(MessageRole::User, &user_input_owned);
240        }).await?;
241
242        self.emit_event(AgentEvent::Checkpoint {
243            session_id: session_id.clone(),
244            checkpoint: CheckpointData {
245                session_id: session_id.clone(),
246                user_input: user_input_owned.clone(),
247                step: CheckpointStep::AfterUserInput,
248                turn_count: 0,
249            },
250        });
251
252        let (outcome, turn_count) = self
253            .run_turn_loop(
254                &session_id,
255                &user_input_owned,
256                &tool_definitions,
257                0,
258                &mut event_rx,
259                &mut on_event,
260            )
261            .await?;
262
263        tracing::info!(session_id = session_id.id, turn_count, "agent turn completed");
264        Ok(outcome)
265    }
266
267    pub async fn run_turn_stream(
268        &self,
269        session_id: SessionId,
270        user_input: &str,
271    ) -> AgentResult<(Vec<AgentEvent>, RunOutcome)> {
272        let mut events = Vec::new();
273        let outcome = self.run_turn_with_handler(session_id, user_input, |event| {
274            events.push(event);
275            Ok(())
276        })
277        .await?;
278        Ok((events, outcome))
279    }
280
281    pub async fn add_user_message(&self, session_id: &SessionId, text: impl Into<String>) -> AgentResult<()> {
282        let text = text.into();
283        self.with_session_mut(session_id, |session| {
284            session.push_message(MessageRole::User, &text);
285        }).await
286    }
287
288    pub async fn add_system_message(&self, session_id: &SessionId, text: impl Into<String>) -> AgentResult<()> {
289        let text = text.into();
290        self.with_session_mut(session_id, |session| {
291            session.push_message(MessageRole::System, &text);
292        }).await
293    }
294
295    pub async fn add_tool_result(&self, session_id: &SessionId, tool_call_id: &str, summary: impl Into<String>) -> AgentResult<()> {
296        let summary = summary.into();
297        self.with_session_mut(session_id, |session| {
298            session.push_tool_result(tool_call_id, summary.clone());
299        }).await
300    }
301
302    pub async fn get_messages(&self, session_id: &SessionId) -> AgentResult<Vec<crate::types::ChatMessage>> {
303        let session = self.session_or_err(session_id).await?;
304        Ok(session.chat_messages().to_vec())
305    }
306
307    pub async fn validate_session(&self, session_id: &SessionId) -> AgentResult<()> {
308        if self.session_manager.session(session_id).await.is_none() {
309            return Err(AgentError::session_not_found(session_id.id));
310        }
311        Ok(())
312    }
313
314    pub fn session_store(&self) -> Arc<dyn SessionStore> {
315        self.session_manager.session_store().clone()
316    }
317}
agent_base/engine/runtime/mod.rs

agent_base/engine/runtime/
mod.rs
