agent_base/engine/runtime/
mod.rs

1use std::sync::Arc;
2
3use tokio_util::sync::CancellationToken;
4
5use crate::engine::session_store::SessionStore;
6use crate::engine::AgentSession;
7use crate::types::{
8    AgentConfig, AgentError, AgentResult, CheckpointData,
9    MessageRole, RunOutcome, RuntimeEvent, SessionId,
10};
11
12use super::approval::ApprovalHandler;
13
14mod event_bus;
15pub(crate) use event_bus::EventBus;
16mod llm_engine;
17mod react_loop;
18mod session_manager;
19mod tool_engine;
20mod plan_runner;
21
22pub(super) const DEFAULT_MAX_TURNS: u32 = 50;
23
24pub use llm_engine::LlmEngine;
25pub use session_manager::SessionManager;
26pub(crate) use tool_engine::ToolEngine;
27pub(crate) use plan_runner::RuntimeCore;
28
29#[derive(Clone)]
30pub struct AgentRuntime {
31    pub(crate) runner: Arc<RuntimeCore>,
32}
33
34impl AgentRuntime {
35    pub async fn create_session(&self) -> SessionId {
36        let config = self.runner.config.read().await;
37        self.runner.session_manager.create_session(config.system_prompt.as_deref()).await
38    }
39
40    pub async fn restore_session(&self, session_id: &SessionId) -> Option<AgentSession> {
41        self.runner.session_manager.restore_session(session_id).await
42    }
43
44    pub async fn session(&self, session_id: &SessionId) -> Option<AgentSession> {
45        self.runner.session_manager.session(session_id).await
46    }
47
48    pub async fn session_or_err(&self, session_id: &SessionId) -> AgentResult<AgentSession> {
49        self.runner.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.runner.session_manager.with_session_mut(session_id, f).await
57    }
58
59    pub fn emit_event(&self, event: RuntimeEvent) {
60        self.runner.event_bus.emit(event);
61    }
62
63    pub(crate) fn subscribe_events(&self) -> tokio::sync::broadcast::Receiver<RuntimeEvent> {
64        self.runner.event_bus.subscribe()
65    }
66
67    /// Subscribe to runtime events from the internal broadcast channel.
68    ///
69    /// Events are delivered directly from the runtime's event bus (capacity 2048).
70    /// Slow consumers may receive `Lagged(n)` errors if they cannot keep up —
71    /// ensure the receiver loop processes events promptly or use a buffering
72    /// layer in the consumer if backpressure is a concern.
73    pub fn subscribe_runtime_events(&self) -> tokio::sync::broadcast::Receiver<RuntimeEvent> {
74        self.runner.event_bus.subscribe()
75    }
76
77    pub fn session_manager(&self) -> &SessionManager {
78        &self.runner.session_manager
79    }
80
81    pub fn llm_engine(&self) -> &LlmEngine {
82        &self.runner.llm_engine
83    }
84
85    pub fn client(&self) -> Arc<dyn crate::llm::LlmClient> {
86        self.runner.llm_engine.get_client()
87    }
88
89    /// Replace the LLM client at runtime (e.g., model switch).
90    /// Requires `&mut self` — obtain via `runtime.lock().await`.
91    pub fn set_client(&mut self, client: Arc<dyn crate::llm::LlmClient>) {
92        self.runner.llm_engine.set_client(client);
93    }
94
95    pub fn tools_mut(&self) -> Arc<tokio::sync::RwLock<crate::tool::ToolRegistry>> {
96        self.runner.tool_engine.tools_arc()
97    }
98
99    /// Inject the internal EventBus into framework tools in the given registry.
100    /// Call this after replacing tools in the registry (e.g., in `build_tools`).
101    pub fn inject_framework_deps(&self, tools: &crate::tool::ToolRegistry) {
102        self.runner.tool_engine.inject_event_bus_into(tools);
103    }
104
105    pub fn config(&self) -> tokio::sync::RwLockReadGuard<'_, AgentConfig> {
106        self.runner.config.blocking_read()
107    }
108
109    /// 设置 reasoning effort（异步版本）
110    pub async fn set_reasoning_effort(&self, effort: crate::llm::ReasoningEffort) {
111        let mut config = self.runner.config.write().await;
112        let mut reasoning = config.reasoning.take().unwrap_or_default();
113        reasoning.effort = Some(effort);
114        config.reasoning = Some(reasoning);
115    }
116
117    /// 设置 reasoning effort（同步版本，只在同步上下文中使用）
118    pub fn set_reasoning_effort_sync(&self, effort: crate::llm::ReasoningEffort) {
119        let mut config = self.runner.config.blocking_write();
120        let mut reasoning = config.reasoning.take().unwrap_or_default();
121        reasoning.effort = Some(effort);
122        config.reasoning = Some(reasoning);
123    }
124
125    pub fn approval_handler(&self) -> Option<&Arc<dyn ApprovalHandler>> {
126        self.runner.tool_engine.approval_handler()
127    }
128
129    pub async fn cached_approval(&self, session_id: &SessionId, action_key: &str) -> bool {
130        self.runner.session_manager.cached_approval(session_id, action_key).await
131    }
132
133    pub async fn cache_approval(&self, session_id: &SessionId, action_key: String) {
134        self.runner.session_manager.cache_approval(session_id, action_key).await
135    }
136
137    pub async fn save_checkpoint(&self, session_id: &SessionId, checkpoint: CheckpointData) -> AgentResult<()> {
138        self.emit_event(RuntimeEvent::Checkpoint {
139            session_id: session_id.clone(),
140            checkpoint,
141        });
142        Ok(())
143    }
144
145    pub async fn load_checkpoint(&self, _session_id: &SessionId, _checkpoint: &CheckpointData) -> AgentResult<Option<CheckpointData>> {
146        Ok(None)
147    }
148
149    pub async fn run<F>(
150        &self,
151        session_id: SessionId,
152        on_event: F,
153    ) -> AgentResult<RunOutcome>
154    where
155        F: FnMut(RuntimeEvent) -> AgentResult<()> + Send,
156    {
157        self.runner.run(session_id, on_event).await
158    }
159
160    pub async fn run_turn<F>(
161        &self,
162        session_id: SessionId,
163        user_input: &str,
164        on_event: F,
165    ) -> AgentResult<RunOutcome>
166    where
167        F: FnMut(RuntimeEvent) -> AgentResult<()> + Send,
168    {
169        self.runner.run_turn(session_id, user_input, on_event).await
170    }
171
172    pub async fn run_turn_collect(
173        &self,
174        session_id: SessionId,
175        user_input: &str,
176    ) -> AgentResult<(Vec<RuntimeEvent>, RunOutcome)> {
177        self.runner.run_turn_collect(session_id, user_input).await
178    }
179
180    pub async fn add_user_message(&self, session_id: &SessionId, text: impl Into<String>) -> AgentResult<()> {
181        let text = text.into();
182        self.with_session_mut(session_id, |session| {
183            session.push_message(MessageRole::User, &text);
184        }).await
185    }
186
187    pub async fn add_system_message(&self, session_id: &SessionId, text: impl Into<String>) -> AgentResult<()> {
188        let text = text.into();
189        self.with_session_mut(session_id, |session| {
190            session.push_message(MessageRole::System, &text);
191        }).await
192    }
193
194    pub async fn add_tool_result(&self, session_id: &SessionId, tool_call_id: &str, summary: impl Into<String>) -> AgentResult<()> {
195        let summary = summary.into();
196        self.with_session_mut(session_id, |session| {
197            session.push_tool_result(tool_call_id, summary.clone());
198        }).await
199    }
200
201    pub async fn get_messages(&self, session_id: &SessionId) -> AgentResult<Vec<crate::types::ChatMessage>> {
202        let session = self.session_or_err(session_id).await?;
203        Ok(session.chat_messages().to_vec())
204    }
205
206    /// Replace the chat messages for a session — only for persistence restore.
207    /// Validates message sequence before applying.
208    ///
209    /// 仅供持久化恢复使用。
210    pub async fn set_messages(
211        &self,
212        session_id: &SessionId,
213        messages: Vec<crate::types::ChatMessage>,
214    ) -> AgentResult<()> {
215        self.with_session_mut(session_id, |session| {
216            session.set_chat_messages(messages)
217        }).await?
218        .map_err(|e| AgentError::internal(e))
219    }
220
221    pub async fn validate_session(&self, session_id: &SessionId) -> AgentResult<()> {
222        if self.runner.session_manager.session(session_id).await.is_none() {
223            return Err(AgentError::session_not_found(session_id.id));
224        }
225        Ok(())
226    }
227
228    pub fn session_store(&self) -> Arc<dyn SessionStore> {
229        self.runner.session_manager.session_store().clone()
230    }
231
232    // --- Cancellation support ---
233
234    /// Cancel the currently executing run_turn / run.
235    /// No-op if there is no current execution.
236    pub fn cancel(&self) {
237        self.runner.cancel();
238    }
239
240    /// Reset the cancel token (called automatically before each run_turn)
241    pub fn reset_cancel(&self) {
242        self.runner.reset_cancel();
243    }
244
245    /// Get a clone of the cancel token
246    pub fn cancel_token(&self) -> CancellationToken {
247        self.runner.cancel_token()
248    }
249
250    /// Check if cancellation has been requested
251    pub fn is_cancelled(&self) -> bool {
252        self.runner.is_cancelled()
253    }
254}
agent_base/engine/runtime/mod.rs

agent_base/engine/runtime/
mod.rs
