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

agent_base/engine/runtime/
mod.rs
