agent_base/engine/runtime/
mod.rs

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

agent_base/engine/runtime/
mod.rs
