| 1 | //! The agent handle. |
| 2 | |
| 3 | use std::collections::{HashMap, HashSet}; |
| 4 | use std::future::Future; |
| 5 | use std::sync::atomic::{AtomicBool, Ordering}; |
| 6 | use std::sync::Arc; |
| 7 | use std::time::{Duration, Instant}; |
| 8 | |
| 9 | use log::{error, info}; |
| 10 | use onde::inference::types::{format_duration, ChatMessage}; |
| 11 | use onde::inference::{ |
| 12 | ChatEngine, EngineInfo, EngineStatus, GgufModelConfig, InferenceError, SamplingConfig, |
| 13 | ToolCallingSupport, ToolDefinition, ToolResult, |
| 14 | }; |
| 15 | |
| 16 | use crate::{ |
| 17 | stream::chunk_error, tool::truncate_output, AgentConfig, AgentError, AgentReply, |
| 18 | AgentToolDefinition, ApprovalDecision, ApprovalHandler, ApprovalRequest, ChatReply, |
| 19 | DenyApprovals, EventSink, NoopSink, RejectingExecutor, TextStream, ToolCall, |
| 20 | ToolExecutionResult, ToolExecutor, ToolRisk, |
| 21 | }; |
| 22 | |
| 23 | /// An embeddable chat agent over a single [`ChatEngine`]. |
| 24 | /// |
| 25 | /// Hold one per application. Ed is `Send + Sync` and every method takes |
| 26 | /// `&self`, so it can live behind a `OnceLock`, a framework's managed state, |
| 27 | /// or an `Arc` — whichever the host prefers. It is not cloneable by design: |
| 28 | /// one Ed owns one engine owns one loaded model, and copying that handle |
| 29 | /// around tends to mean two parts of an app disagree about which model is up. |
| 30 | pub struct Ed<S = NoopSink> { |
| 31 | engine: ChatEngine, |
| 32 | sink: S, |
| 33 | tools: tokio::sync::RwLock<HashMap<String, AgentToolDefinition>>, |
| 34 | executor: Arc<dyn ToolExecutor>, |
| 35 | approvals: Arc<dyn ApprovalHandler>, |
| 36 | session_grants: tokio::sync::Mutex<HashSet<String>>, |
| 37 | turn: tokio::sync::Mutex<()>, |
| 38 | cancelled: AtomicBool, |
| 39 | agent_config: AgentConfig, |
| 40 | } |
| 41 | |
| 42 | impl Ed<NoopSink> { |
| 43 | /// An agent that reports nothing. Poll [`info`](Ed::info) for state. |
| 44 | pub fn new() -> Self { |
| 45 | Self::with_sink(NoopSink) |
| 46 | } |
| 47 | } |
| 48 | |
| 49 | impl Default for Ed<NoopSink> { |
| 50 | fn default() -> Self { |
| 51 | Self::new() |
| 52 | } |
| 53 | } |
| 54 | |
| 55 | impl<S: EventSink> Ed<S> { |
| 56 | /// An agent that reports transitions and replies to `sink`. |
| 57 | pub fn with_sink(sink: S) -> Self { |
| 58 | Self::with_agent( |
| 59 | sink, |
| 60 | Arc::new(RejectingExecutor), |
| 61 | Arc::new(DenyApprovals), |
| 62 | AgentConfig::default(), |
| 63 | ) |
| 64 | } |
| 65 | |
| 66 | /// Build an agent with host-defined tool execution and approval handling. |
| 67 | pub fn with_agent( |
| 68 | sink: S, |
| 69 | executor: Arc<dyn ToolExecutor>, |
| 70 | approvals: Arc<dyn ApprovalHandler>, |
| 71 | config: AgentConfig, |
| 72 | ) -> Self { |
| 73 | Self::with_agent_and_app_id(sink, executor, approvals, config, None) |
| 74 | } |
| 75 | |
| 76 | /// Associate telemetry with an Onde app id. See [`ChatEngine::with_app_id`]. |
| 77 | pub fn with_app_id(sink: S, onde_app_id: Option<String>) -> Self { |
| 78 | Self::with_agent_and_app_id( |
| 79 | sink, |
| 80 | Arc::new(RejectingExecutor), |
| 81 | Arc::new(DenyApprovals), |
| 82 | AgentConfig::default(), |
| 83 | onde_app_id, |
| 84 | ) |
| 85 | } |
| 86 | |
| 87 | /// The full constructor: tools, approvals, and an Onde app id together. |
| 88 | /// |
| 89 | /// The narrower constructors each fill in defaults for what they don't |
| 90 | /// take. Use this one when you need both host tools and telemetry, which |
| 91 | /// the others cannot express between them. |
| 92 | pub fn with_agent_and_app_id( |
| 93 | sink: S, |
| 94 | executor: Arc<dyn ToolExecutor>, |
| 95 | approvals: Arc<dyn ApprovalHandler>, |
| 96 | config: AgentConfig, |
| 97 | onde_app_id: Option<String>, |
| 98 | ) -> Self { |
| 99 | Self { |
| 100 | engine: ChatEngine::with_app_id(onde_app_id), |
| 101 | sink, |
| 102 | tools: tokio::sync::RwLock::new(HashMap::new()), |
| 103 | executor, |
| 104 | approvals, |
| 105 | session_grants: tokio::sync::Mutex::new(HashSet::new()), |
| 106 | turn: tokio::sync::Mutex::new(()), |
| 107 | cancelled: AtomicBool::new(false), |
| 108 | agent_config: config.validated(), |
| 109 | } |
| 110 | } |
| 111 | |
| 112 | /// The underlying engine. |
| 113 | /// |
| 114 | /// An escape hatch for the parts of `onde` Ed deliberately doesn't wrap — |
| 115 | /// system prompts, alternative load paths, tool calling. Reach for it when |
| 116 | /// you need them; anything you find yourself writing here for a second |
| 117 | /// application is a candidate to move into Ed. |
| 118 | pub fn engine(&self) -> &ChatEngine { |
| 119 | &self.engine |
| 120 | } |
| 121 | |
| 122 | /// Run any load, reporting the transitions around it. |
| 123 | /// |
| 124 | /// Emits [`EngineStatus::Loading`] before `load` starts and then either |
| 125 | /// [`Ready`](EngineStatus::Ready) or [`Error`](EngineStatus::Error). That |
| 126 | /// sequence is the whole reason this method exists: calling a |
| 127 | /// [`ChatEngine`] load path directly leaves the UI with no way to know a |
| 128 | /// multi-minute download has begun. |
| 129 | /// |
| 130 | /// Ed takes a closure rather than wrapping each load path by name because |
| 131 | /// `onde`'s load APIs take types Ed has no business in its public |
| 132 | /// signature — `ChatEngine::load_assigned_model` takes an `Environment` |
| 133 | /// from the smbCloud SDK, and *which* crate that type comes from changed |
| 134 | /// between `onde` releases. Naming it here would drag that dependency into |
| 135 | /// every host and pin them all to Ed's version of it. Passing a closure |
| 136 | /// keeps Ed's dependency surface to `onde` alone and works for load paths |
| 137 | /// that don't exist yet. |
| 138 | /// |
| 139 | /// ```no_run |
| 140 | /// # use ed_agent::{Ed, GgufModelConfig}; |
| 141 | /// # async fn example(ed: Ed) { |
| 142 | /// ed.load_with("Qwen 2.5", |engine| { |
| 143 | /// engine.load_gguf_model(GgufModelConfig::qwen25_1_5b(), None, None) |
| 144 | /// }) |
| 145 | /// .await |
| 146 | /// .expect("load"); |
| 147 | /// # } |
| 148 | /// ``` |
| 149 | /// |
| 150 | /// `model_name` is what the host wants the UI to show while loading; it is |
| 151 | /// reported with each transition and needn't match what the engine ends up |
| 152 | /// calling the model. |
| 153 | pub async fn load_with<'a, F, Fut>( |
| 154 | &'a self, |
| 155 | model_name: &str, |
| 156 | load: F, |
| 157 | ) -> Result<Duration, InferenceError> |
| 158 | where |
| 159 | F: FnOnce(&'a ChatEngine) -> Fut, |
| 160 | Fut: Future<Output = Result<Duration, InferenceError>>, |
| 161 | { |
| 162 | info!("ed: loading model {model_name}"); |
| 163 | self.sink |
| 164 | .status_changed(EngineStatus::Loading, Some(model_name), None); |
| 165 | |
| 166 | match load(&self.engine).await { |
| 167 | Ok(elapsed) => { |
| 168 | info!("ed: loaded {model_name} in {}", format_duration(elapsed)); |
| 169 | self.sink |
| 170 | .status_changed(EngineStatus::Ready, Some(model_name), None); |
| 171 | Ok(elapsed) |
| 172 | } |
| 173 | Err(err) => { |
| 174 | error!("ed: failed to load {model_name}: {err}"); |
| 175 | let message = err.to_string(); |
| 176 | self.sink |
| 177 | .status_changed(EngineStatus::Error, Some(model_name), Some(&message)); |
| 178 | Err(err) |
| 179 | } |
| 180 | } |
| 181 | } |
| 182 | |
| 183 | /// Load a GGUF model, reporting each transition. |
| 184 | /// |
| 185 | /// A convenience over [`load_with`](Ed::load_with) for the most common |
| 186 | /// path. Downloads the weights first if they aren't in the local cache, so |
| 187 | /// the `Loading` state can last a long time on a cold start. |
| 188 | pub async fn load( |
| 189 | &self, |
| 190 | config: GgufModelConfig, |
| 191 | system_prompt: Option<String>, |
| 192 | sampling: Option<SamplingConfig>, |
| 193 | ) -> Result<Duration, InferenceError> { |
| 194 | let model_name = config.display_name.clone(); |
| 195 | self.load_with(&model_name, |engine| { |
| 196 | engine.load_gguf_model(config, system_prompt, sampling) |
| 197 | }) |
| 198 | .await |
| 199 | } |
| 200 | |
| 201 | /// Unload the current model, returning its display name if one was loaded. |
| 202 | pub async fn unload(&self) -> Option<String> { |
| 203 | let unloaded = self.engine.unload_model().await; |
| 204 | if let Some(name) = &unloaded { |
| 205 | info!("ed: unloaded {name}"); |
| 206 | } |
| 207 | self.sink |
| 208 | .status_changed(EngineStatus::Unloaded, unloaded.as_deref(), None); |
| 209 | unloaded |
| 210 | } |
| 211 | |
| 212 | /// Run one inference turn. |
| 213 | /// |
| 214 | /// Never returns an error: a failure is reported in |
| 215 | /// [`ChatReply::error`] so a caller has one shape to render either way. |
| 216 | /// The reply also goes to the sink, so a host can render it from an event |
| 217 | /// instead of awaiting this call — see [`EventSink::replied`]. |
| 218 | pub async fn send(&self, message: impl Into<String>) -> ChatReply { |
| 219 | let reply = match self.engine.send_message(message).await { |
| 220 | Ok(result) => ChatReply::ok(result.text, result.duration_display), |
| 221 | Err(err) => { |
| 222 | error!("ed: inference failed: {err}"); |
| 223 | ChatReply::failed(err) |
| 224 | } |
| 225 | }; |
| 226 | self.sink.replied(&reply); |
| 227 | reply |
| 228 | } |
| 229 | |
| 230 | /// Stream the reply to one message, for plain chat without tools. |
| 231 | /// |
| 232 | /// Returns once generation has started; the stream yields text deltas as |
| 233 | /// they are produced. Unlike [`run`](Ed::run) this takes no turn lock and |
| 234 | /// reports nothing to the sink, so don't overlap it with another turn. |
| 235 | /// Drop the stream to cancel; a [`cancel`](Ed::cancel) call is not needed. |
| 236 | pub async fn stream(&self, message: impl Into<String>) -> Result<TextStream, InferenceError> { |
| 237 | Ok(TextStream::new(self.engine.stream_message(message).await?)) |
| 238 | } |
| 239 | |
| 240 | /// One-shot generation over explicit messages. |
| 241 | /// |
| 242 | /// Leaves the conversation history untouched; the engine's system prompt |
| 243 | /// is still applied. `sampling` overrides the engine's for this call only. |
| 244 | pub async fn generate( |
| 245 | &self, |
| 246 | messages: Vec<ChatMessage>, |
| 247 | sampling: Option<SamplingConfig>, |
| 248 | ) -> Result<String, InferenceError> { |
| 249 | Ok(self.engine.generate(messages, sampling).await?.text) |
| 250 | } |
| 251 | |
| 252 | /// Set the system prompt used for every later turn. |
| 253 | /// |
| 254 | /// Like the other engine setters below, and `restore_history`, this does |
| 255 | /// nothing until a model is loaded: `onde` keeps these on the loaded model. |
| 256 | pub async fn set_system_prompt(&self, prompt: impl Into<String>) { |
| 257 | self.engine.set_system_prompt(prompt).await; |
| 258 | } |
| 259 | |
| 260 | pub async fn clear_system_prompt(&self) { |
| 261 | self.engine.clear_system_prompt().await; |
| 262 | } |
| 263 | |
| 264 | /// Change the default sampling parameters. |
| 265 | pub async fn set_sampling(&self, sampling: SamplingConfig) { |
| 266 | self.engine.set_sampling(sampling).await; |
| 267 | } |
| 268 | |
| 269 | /// Replace the conversation with `messages`. |
| 270 | /// |
| 271 | /// For switching to a saved session or restoring one at launch. Clears |
| 272 | /// the current history first, and with it every `AllowForSession` grant, |
| 273 | /// exactly as [`clear_history`](Ed::clear_history) does: a restored |
| 274 | /// conversation starts without consent from the previous one. |
| 275 | pub async fn restore_history(&self, messages: Vec<ChatMessage>) { |
| 276 | self.clear_history().await; |
| 277 | let count = messages.len(); |
| 278 | for message in messages { |
| 279 | self.engine.push_history(message).await; |
| 280 | } |
| 281 | info!("ed: restored {count} messages"); |
| 282 | } |
| 283 | |
| 284 | /// Register a host capability that the model may request. |
| 285 | pub async fn register_tool(&self, tool: AgentToolDefinition) -> Result<(), AgentError> { |
| 286 | tool.validate()?; |
| 287 | let mut tools = self.tools.write().await; |
| 288 | if tools.contains_key(&tool.name) { |
| 289 | return Err(AgentError::DuplicateTool { |
| 290 | name: tool.name.clone(), |
| 291 | }); |
| 292 | } |
| 293 | tools.insert(tool.name.clone(), tool); |
| 294 | Ok(()) |
| 295 | } |
| 296 | |
| 297 | pub async fn unregister_tool(&self, name: &str) -> bool { |
| 298 | self.tools.write().await.remove(name).is_some() |
| 299 | } |
| 300 | |
| 301 | pub async fn remove_all_tools(&self) { |
| 302 | self.tools.write().await.clear(); |
| 303 | self.session_grants.lock().await.clear(); |
| 304 | } |
| 305 | |
| 306 | /// Request cancellation of the active agent turn. |
| 307 | pub fn cancel(&self) { |
| 308 | self.cancelled.store(true, Ordering::Release); |
| 309 | } |
| 310 | |
| 311 | async fn execute_tool_call( |
| 312 | &self, |
| 313 | call: &ToolCall, |
| 314 | registered: &HashMap<String, AgentToolDefinition>, |
| 315 | ) -> ToolExecutionResult { |
| 316 | let Some(definition) = registered.get(&call.name) else { |
| 317 | return ToolExecutionResult { |
| 318 | tool_call_id: call.id.clone(), |
| 319 | content: format!("tool {:?} is not registered", call.name), |
| 320 | is_error: true, |
| 321 | }; |
| 322 | }; |
| 323 | |
| 324 | if !matches!( |
| 325 | serde_json::from_str::<serde_json::Value>(&call.arguments), |
| 326 | Ok(serde_json::Value::Object(_)) |
| 327 | ) { |
| 328 | return ToolExecutionResult { |
| 329 | tool_call_id: call.id.clone(), |
| 330 | content: "tool arguments must be a valid JSON object".to_string(), |
| 331 | is_error: true, |
| 332 | }; |
| 333 | } |
| 334 | |
| 335 | let approved = if definition.risk == ToolRisk::ReadOnly |
| 336 | || self.session_grants.lock().await.contains(&call.name) |
| 337 | { |
| 338 | true |
| 339 | } else { |
| 340 | let request = ApprovalRequest { |
| 341 | call: call.clone(), |
| 342 | risk: definition.risk, |
| 343 | }; |
| 344 | self.sink.approval_requested(&request); |
| 345 | match self.approvals.approve(request).await { |
| 346 | ApprovalDecision::AllowOnce => true, |
| 347 | ApprovalDecision::AllowForSession => { |
| 348 | self.session_grants.lock().await.insert(call.name.clone()); |
| 349 | true |
| 350 | } |
| 351 | ApprovalDecision::Deny => false, |
| 352 | } |
| 353 | }; |
| 354 | |
| 355 | if !approved { |
| 356 | return ToolExecutionResult { |
| 357 | tool_call_id: call.id.clone(), |
| 358 | content: format!("tool {:?} was denied by the user", call.name), |
| 359 | is_error: true, |
| 360 | }; |
| 361 | } |
| 362 | |
| 363 | self.sink.tool_started(call); |
| 364 | match self.executor.execute(&call.name, &call.arguments).await { |
| 365 | Ok(content) => ToolExecutionResult { |
| 366 | tool_call_id: call.id.clone(), |
| 367 | content: truncate_output(content, self.agent_config.max_tool_output_chars), |
| 368 | is_error: false, |
| 369 | }, |
| 370 | Err(error) => ToolExecutionResult { |
| 371 | tool_call_id: call.id.clone(), |
| 372 | content: truncate_output( |
| 373 | format!("tool execution failed: {error}"), |
| 374 | self.agent_config.max_tool_output_chars, |
| 375 | ), |
| 376 | is_error: true, |
| 377 | }, |
| 378 | } |
| 379 | } |
| 380 | |
| 381 | /// Run a complete local agent turn, including any requested host tools. |
| 382 | /// |
| 383 | /// With no tools registered this is one inference call. Otherwise Ed |
| 384 | /// alternates between the model and the host's executor until the model |
| 385 | /// answers in text, or until |
| 386 | /// [`max_tool_rounds`](AgentConfig::max_tool_rounds) is spent, whichever |
| 387 | /// comes first. Reaching the budget is reported to the sink as a warning, |
| 388 | /// and yields [`AgentError::ToolRoundsExhausted`] if there is no usable |
| 389 | /// text by then. |
| 390 | /// |
| 391 | /// Reports [`EventSink::agent_replied`] on success. It does *not* report |
| 392 | /// [`EventSink::replied`]; that belongs to [`send`](Ed::send). |
| 393 | pub async fn run(&self, message: impl Into<String>) -> Result<AgentReply, AgentError> { |
| 394 | let _turn = self.turn.lock().await; |
| 395 | self.cancelled.store(false, Ordering::Release); |
| 396 | |
| 397 | let registered = self.tools.read().await.clone(); |
| 398 | if registered.is_empty() { |
| 399 | // Deliberately not routed through `send`: that reports a |
| 400 | // `ChatReply` to the sink, and a host listening for both events |
| 401 | // would render this turn's answer twice. An agent turn reports |
| 402 | // `agent_replied` and nothing else. |
| 403 | let started = Instant::now(); |
| 404 | let mut rx = self.engine.stream_message(message).await.map_err(|error| { |
| 405 | AgentError::Inference { |
| 406 | reason: error.to_string(), |
| 407 | } |
| 408 | })?; |
| 409 | let mut text = String::new(); |
| 410 | while let Some(chunk) = rx.recv().await { |
| 411 | if self.cancelled.load(Ordering::Acquire) { |
| 412 | // Dropping the receiver stops generation. |
| 413 | return Err(AgentError::Cancelled); |
| 414 | } |
| 415 | if chunk.done { |
| 416 | if let Some(error) = chunk_error(&chunk) { |
| 417 | return Err(AgentError::Inference { |
| 418 | reason: error.to_string(), |
| 419 | }); |
| 420 | } |
| 421 | break; |
| 422 | } |
| 423 | if !chunk.delta.is_empty() { |
| 424 | self.sink.text_delta(&chunk.delta); |
| 425 | text.push_str(&chunk.delta); |
| 426 | } |
| 427 | } |
| 428 | if text.trim().is_empty() { |
| 429 | return Err(AgentError::EmptyReply); |
| 430 | } |
| 431 | let elapsed = started.elapsed(); |
| 432 | let agent_reply = AgentReply { |
| 433 | text, |
| 434 | duration_seconds: elapsed.as_secs_f64(), |
| 435 | duration: format_duration(elapsed), |
| 436 | tool_rounds: 0, |
| 437 | }; |
| 438 | self.sink.agent_replied(&agent_reply); |
| 439 | return Ok(agent_reply); |
| 440 | } |
| 441 | |
| 442 | match self.engine.tool_calling_support().await { |
| 443 | ToolCallingSupport::Unsupported => return Err(AgentError::UnsupportedModel), |
| 444 | ToolCallingSupport::Unknown => self |
| 445 | .sink |
| 446 | .warning("the loaded model has not been verified for tool calling"), |
| 447 | ToolCallingSupport::Supported => {} |
| 448 | } |
| 449 | |
| 450 | let mut definitions: Vec<ToolDefinition> = registered |
| 451 | .values() |
| 452 | .map(|tool| ToolDefinition { |
| 453 | name: tool.name.clone(), |
| 454 | description: tool.description.clone(), |
| 455 | parameters_schema: tool.parameters_schema.clone(), |
| 456 | }) |
| 457 | .collect(); |
| 458 | definitions.sort_by(|left, right| left.name.cmp(&right.name)); |
| 459 | |
| 460 | let mut result = self |
| 461 | .engine |
| 462 | .send_message_with_tools(message.into(), &definitions) |
| 463 | .await |
| 464 | .map_err(|error| AgentError::Inference { |
| 465 | reason: error.to_string(), |
| 466 | })?; |
| 467 | let mut duration_seconds = result.duration_secs; |
| 468 | let mut rounds = 0_u8; |
| 469 | let max_rounds = self.agent_config.max_tool_rounds; |
| 470 | let mut exhausted = false; |
| 471 | |
| 472 | while !result.tool_calls.is_empty() { |
| 473 | if !self.agent_config.allows_round(rounds) { |
| 474 | // The budget only stops Ed *offering* tools; it cannot stop the |
| 475 | // model asking. `onde` parses tool calls out of every response |
| 476 | // whether or not tools were sent, so a model that keeps |
| 477 | // emitting tool-call syntax would loop here until `rounds` |
| 478 | // overflowed and wrapped back under the limit. This is the |
| 479 | // loop's only hard exit. |
| 480 | // |
| 481 | // The model's last turn is already in history carrying these |
| 482 | // calls. Leaving them unanswered would make the next turn's |
| 483 | // message sequence invalid, so close them out the way |
| 484 | // cancelling does. |
| 485 | let unanswered = result |
| 486 | .tool_calls |
| 487 | .iter() |
| 488 | .map(|pending| ToolResult { |
| 489 | tool_call_id: pending.id.clone(), |
| 490 | content: "tool was not executed: the agent ran out of tool rounds" |
| 491 | .to_string(), |
| 492 | }) |
| 493 | .collect(); |
| 494 | self.engine.record_tool_results(unanswered).await; |
| 495 | self.sink.warning(&format!( |
| 496 | "the agent stopped after {max_rounds} tool rounds; the model kept requesting tools" |
| 497 | )); |
| 498 | exhausted = true; |
| 499 | break; |
| 500 | } |
| 501 | rounds += 1; |
| 502 | let mut tool_results = Vec::with_capacity(result.tool_calls.len()); |
| 503 | |
| 504 | for (index, requested) in result.tool_calls.iter().enumerate() { |
| 505 | if self.cancelled.load(Ordering::Acquire) { |
| 506 | for pending in &result.tool_calls[index..] { |
| 507 | tool_results.push(ToolResult { |
| 508 | tool_call_id: pending.id.clone(), |
| 509 | content: "tool was not executed because the user cancelled the turn" |
| 510 | .to_string(), |
| 511 | }); |
| 512 | } |
| 513 | self.engine.record_tool_results(tool_results).await; |
| 514 | return Err(AgentError::Cancelled); |
| 515 | } |
| 516 | |
| 517 | let call = ToolCall { |
| 518 | id: requested.id.clone(), |
| 519 | name: requested.function_name.clone(), |
| 520 | arguments: requested.arguments.clone(), |
| 521 | }; |
| 522 | self.sink.tool_requested(&call); |
| 523 | |
| 524 | let execution = self.execute_tool_call(&call, ®istered).await; |
| 525 | self.sink.tool_finished(&execution); |
| 526 | tool_results.push(ToolResult { |
| 527 | tool_call_id: execution.tool_call_id, |
| 528 | content: execution.content, |
| 529 | }); |
| 530 | } |
| 531 | |
| 532 | if self.cancelled.load(Ordering::Acquire) { |
| 533 | self.engine.record_tool_results(tool_results).await; |
| 534 | return Err(AgentError::Cancelled); |
| 535 | } |
| 536 | |
| 537 | let next_tools = if self.agent_config.allows_round(rounds) { |
| 538 | Some(definitions.as_slice()) |
| 539 | } else { |
| 540 | None |
| 541 | }; |
| 542 | result = self |
| 543 | .engine |
| 544 | .send_tool_results(tool_results, next_tools) |
| 545 | .await |
| 546 | .map_err(|error| AgentError::Inference { |
| 547 | reason: error.to_string(), |
| 548 | })?; |
| 549 | duration_seconds += result.duration_secs; |
| 550 | } |
| 551 | |
| 552 | if result.text.trim().is_empty() { |
| 553 | return Err(if exhausted { |
| 554 | AgentError::ToolRoundsExhausted { rounds: max_rounds } |
| 555 | } else { |
| 556 | AgentError::EmptyReply |
| 557 | }); |
| 558 | } |
| 559 | // Tool rounds are not streamed: `onde` only parses tool calls out of a |
| 560 | // complete response, so the answer arrives whole. Report it as one |
| 561 | // delta so a host rendering deltas sees every turn the same way. |
| 562 | self.sink.text_delta(&result.text); |
| 563 | let reply = AgentReply { |
| 564 | text: result.text, |
| 565 | duration_seconds, |
| 566 | duration: format_duration(Duration::from_secs_f64(duration_seconds)), |
| 567 | tool_rounds: rounds, |
| 568 | }; |
| 569 | self.sink.agent_replied(&reply); |
| 570 | Ok(reply) |
| 571 | } |
| 572 | |
| 573 | /// Current status, loaded model, footprint, and history length. |
| 574 | pub async fn info(&self) -> EngineInfo { |
| 575 | self.engine.info().await |
| 576 | } |
| 577 | |
| 578 | /// Whether a model is loaded and ready. |
| 579 | pub async fn is_loaded(&self) -> bool { |
| 580 | self.engine.is_loaded().await |
| 581 | } |
| 582 | |
| 583 | /// The conversation so far. |
| 584 | pub async fn history(&self) -> Vec<ChatMessage> { |
| 585 | self.engine.history().await |
| 586 | } |
| 587 | |
| 588 | /// Clear the conversation, returning how many messages were dropped. |
| 589 | /// |
| 590 | /// Leaves the model loaded — this starts a new conversation, it doesn't |
| 591 | /// tear down the engine. |
| 592 | /// |
| 593 | /// Also drops every `AllowForSession` grant. The user approved those |
| 594 | /// mutating tools for a conversation, and this ends the conversation; |
| 595 | /// carrying the grants into the next one would keep consent alive past the |
| 596 | /// thing it was given for. |
| 597 | pub async fn clear_history(&self) -> usize { |
| 598 | let cleared = self.engine.clear_history().await; |
| 599 | self.session_grants.lock().await.clear(); |
| 600 | info!("ed: cleared {cleared} messages"); |
| 601 | cleared |
| 602 | } |
| 603 | } |
| 604 | |
| 605 | #[cfg(test)] |
| 606 | mod tests { |
| 607 | use std::sync::atomic::{AtomicUsize, Ordering}; |
| 608 | use std::sync::Mutex; |
| 609 | |
| 610 | use super::*; |
| 611 | |
| 612 | /// Records everything Ed reports, so a test can assert on the sequence of |
| 613 | /// transitions rather than only the final state. |
| 614 | #[derive(Default)] |
| 615 | struct RecordingSink { |
| 616 | statuses: Mutex<Vec<(EngineStatus, Option<String>)>>, |
| 617 | errors: Mutex<Vec<Option<String>>>, |
| 618 | replies: Mutex<Vec<ChatReply>>, |
| 619 | } |
| 620 | |
| 621 | impl RecordingSink { |
| 622 | fn statuses(&self) -> Vec<(EngineStatus, Option<String>)> { |
| 623 | self.statuses.lock().unwrap().clone() |
| 624 | } |
| 625 | |
| 626 | fn errors(&self) -> Vec<Option<String>> { |
| 627 | self.errors.lock().unwrap().clone() |
| 628 | } |
| 629 | } |
| 630 | |
| 631 | impl EventSink for &'static RecordingSink { |
| 632 | fn status_changed( |
| 633 | &self, |
| 634 | status: EngineStatus, |
| 635 | model_name: Option<&str>, |
| 636 | error: Option<&str>, |
| 637 | ) { |
| 638 | self.statuses |
| 639 | .lock() |
| 640 | .unwrap() |
| 641 | .push((status, model_name.map(str::to_owned))); |
| 642 | self.errors.lock().unwrap().push(error.map(str::to_owned)); |
| 643 | } |
| 644 | |
| 645 | fn replied(&self, reply: &ChatReply) { |
| 646 | self.replies.lock().unwrap().push(reply.clone()); |
| 647 | } |
| 648 | } |
| 649 | |
| 650 | fn sink() -> &'static RecordingSink { |
| 651 | Box::leak(Box::new(RecordingSink::default())) |
| 652 | } |
| 653 | |
| 654 | struct CountingExecutor(AtomicUsize); |
| 655 | |
| 656 | #[async_trait::async_trait] |
| 657 | impl ToolExecutor for CountingExecutor { |
| 658 | async fn execute(&self, tool_name: &str, _: &str) -> Result<String, String> { |
| 659 | self.0.fetch_add(1, Ordering::Relaxed); |
| 660 | Ok(format!("{tool_name} completed")) |
| 661 | } |
| 662 | } |
| 663 | |
| 664 | struct FixedApproval(ApprovalDecision); |
| 665 | |
| 666 | #[async_trait::async_trait] |
| 667 | impl ApprovalHandler for FixedApproval { |
| 668 | async fn approve(&self, _: ApprovalRequest) -> ApprovalDecision { |
| 669 | self.0 |
| 670 | } |
| 671 | } |
| 672 | |
| 673 | fn call(name: &str) -> ToolCall { |
| 674 | ToolCall { |
| 675 | id: "call-1".to_owned(), |
| 676 | name: name.to_owned(), |
| 677 | arguments: "{}".to_owned(), |
| 678 | } |
| 679 | } |
| 680 | |
| 681 | #[tokio::test] |
| 682 | async fn a_fresh_agent_reports_itself_unloaded_and_empty() { |
| 683 | let ed = Ed::new(); |
| 684 | let info = ed.info().await; |
| 685 | |
| 686 | assert_eq!(info.status, EngineStatus::Unloaded); |
| 687 | assert_eq!(info.model_name, None); |
| 688 | assert_eq!(info.history_length, 0); |
| 689 | assert!(!ed.is_loaded().await); |
| 690 | assert!(ed.history().await.is_empty()); |
| 691 | } |
| 692 | |
| 693 | #[tokio::test] |
| 694 | async fn unloading_reports_the_transition_even_with_nothing_loaded() { |
| 695 | // The status still has to reach the host: a UI that asked to unload |
| 696 | // needs to see `Unloaded` regardless of whether work was done, or the |
| 697 | // view sits on a stale "ready" forever. |
| 698 | let sink = sink(); |
| 699 | let ed = Ed::with_sink(sink); |
| 700 | |
| 701 | assert_eq!(ed.unload().await, None); |
| 702 | assert_eq!(sink.statuses(), vec![(EngineStatus::Unloaded, None)]); |
| 703 | } |
| 704 | |
| 705 | #[tokio::test] |
| 706 | async fn clearing_an_empty_history_drops_nothing() { |
| 707 | let ed = Ed::new(); |
| 708 | assert_eq!(ed.clear_history().await, 0); |
| 709 | } |
| 710 | |
| 711 | #[tokio::test] |
| 712 | async fn the_default_sink_swallows_reports_without_panicking() { |
| 713 | // `Ed::new()` is the no-sink path; exercising it guards against a |
| 714 | // future notification being added that assumes a real sink exists. |
| 715 | let ed = Ed::default(); |
| 716 | ed.unload().await; |
| 717 | assert_eq!(ed.info().await.status, EngineStatus::Unloaded); |
| 718 | } |
| 719 | |
| 720 | #[tokio::test] |
| 721 | async fn a_successful_load_reports_loading_then_ready() { |
| 722 | // The ordering is the feature: a UI that only ever sees the final |
| 723 | // state has no way to put up a spinner for a load that can take |
| 724 | // minutes. |
| 725 | let sink = sink(); |
| 726 | let ed = Ed::with_sink(sink); |
| 727 | |
| 728 | let elapsed = ed |
| 729 | .load_with("Test model", |_| async { Ok(Duration::from_secs(2)) }) |
| 730 | .await |
| 731 | .expect("the closure succeeded"); |
| 732 | |
| 733 | assert_eq!(elapsed, Duration::from_secs(2)); |
| 734 | assert_eq!( |
| 735 | sink.statuses(), |
| 736 | vec![ |
| 737 | (EngineStatus::Loading, Some("Test model".to_owned())), |
| 738 | (EngineStatus::Ready, Some("Test model".to_owned())), |
| 739 | ] |
| 740 | ); |
| 741 | } |
| 742 | |
| 743 | #[tokio::test] |
| 744 | async fn a_failed_load_reports_loading_then_error_with_the_reason() { |
| 745 | let sink = sink(); |
| 746 | let ed = Ed::with_sink(sink); |
| 747 | |
| 748 | let err = ed |
| 749 | .load_with("Test model", |_| async { |
| 750 | Err(InferenceError::Other { |
| 751 | reason: "no weights".to_owned(), |
| 752 | }) |
| 753 | }) |
| 754 | .await |
| 755 | .expect_err("the closure failed"); |
| 756 | |
| 757 | assert!(err.to_string().contains("no weights")); |
| 758 | assert_eq!( |
| 759 | sink.statuses(), |
| 760 | vec![ |
| 761 | (EngineStatus::Loading, Some("Test model".to_owned())), |
| 762 | (EngineStatus::Error, Some("Test model".to_owned())), |
| 763 | ] |
| 764 | ); |
| 765 | assert_eq!(sink.errors(), vec![None, Some("no weights".to_owned())]); |
| 766 | } |
| 767 | |
| 768 | #[tokio::test] |
| 769 | async fn read_only_tools_run_without_an_approval_round_trip() { |
| 770 | let executor = Arc::new(CountingExecutor(AtomicUsize::new(0))); |
| 771 | let ed = Ed::with_agent( |
| 772 | NoopSink, |
| 773 | executor.clone(), |
| 774 | Arc::new(FixedApproval(ApprovalDecision::Deny)), |
| 775 | AgentConfig::default(), |
| 776 | ); |
| 777 | ed.register_tool(AgentToolDefinition::read_only( |
| 778 | "clock", |
| 779 | "Read the clock", |
| 780 | r#"{"type":"object"}"#, |
| 781 | )) |
| 782 | .await |
| 783 | .unwrap(); |
| 784 | let tools = ed.tools.read().await.clone(); |
| 785 | |
| 786 | let result = ed.execute_tool_call(&call("clock"), &tools).await; |
| 787 | |
| 788 | assert!(!result.is_error); |
| 789 | assert_eq!(executor.0.load(Ordering::Relaxed), 1); |
| 790 | } |
| 791 | |
| 792 | #[tokio::test] |
| 793 | async fn denied_mutating_tools_never_reach_the_executor() { |
| 794 | let executor = Arc::new(CountingExecutor(AtomicUsize::new(0))); |
| 795 | let ed = Ed::with_agent( |
| 796 | NoopSink, |
| 797 | executor.clone(), |
| 798 | Arc::new(FixedApproval(ApprovalDecision::Deny)), |
| 799 | AgentConfig::default(), |
| 800 | ); |
| 801 | ed.register_tool(AgentToolDefinition::mutating( |
| 802 | "save_note", |
| 803 | "Save a note", |
| 804 | r#"{"type":"object"}"#, |
| 805 | )) |
| 806 | .await |
| 807 | .unwrap(); |
| 808 | let tools = ed.tools.read().await.clone(); |
| 809 | |
| 810 | let result = ed.execute_tool_call(&call("save_note"), &tools).await; |
| 811 | |
| 812 | assert!(result.is_error); |
| 813 | assert!(result.content.contains("denied")); |
| 814 | assert_eq!(executor.0.load(Ordering::Relaxed), 0); |
| 815 | } |
| 816 | |
| 817 | #[tokio::test] |
| 818 | async fn a_session_grant_applies_to_later_calls_of_the_same_tool() { |
| 819 | let executor = Arc::new(CountingExecutor(AtomicUsize::new(0))); |
| 820 | let ed = Ed::with_agent( |
| 821 | NoopSink, |
| 822 | executor.clone(), |
| 823 | Arc::new(FixedApproval(ApprovalDecision::AllowForSession)), |
| 824 | AgentConfig::default(), |
| 825 | ); |
| 826 | ed.register_tool(AgentToolDefinition::mutating( |
| 827 | "save_note", |
| 828 | "Save a note", |
| 829 | r#"{"type":"object"}"#, |
| 830 | )) |
| 831 | .await |
| 832 | .unwrap(); |
| 833 | let tools = ed.tools.read().await.clone(); |
| 834 | |
| 835 | ed.execute_tool_call(&call("save_note"), &tools).await; |
| 836 | ed.execute_tool_call(&call("save_note"), &tools).await; |
| 837 | |
| 838 | assert_eq!(executor.0.load(Ordering::Relaxed), 2); |
| 839 | assert!(ed.session_grants.lock().await.contains("save_note")); |
| 840 | } |
| 841 | |
| 842 | #[tokio::test] |
| 843 | async fn clearing_the_conversation_revokes_session_grants() { |
| 844 | // The grant was given for a conversation. Clearing the history ends |
| 845 | // that conversation, so the next one has to ask again. |
| 846 | let executor = Arc::new(CountingExecutor(AtomicUsize::new(0))); |
| 847 | let ed = Ed::with_agent( |
| 848 | NoopSink, |
| 849 | executor.clone(), |
| 850 | Arc::new(FixedApproval(ApprovalDecision::AllowForSession)), |
| 851 | AgentConfig::default(), |
| 852 | ); |
| 853 | ed.register_tool(AgentToolDefinition::mutating( |
| 854 | "save_note", |
| 855 | "Save a note", |
| 856 | r#"{"type":"object"}"#, |
| 857 | )) |
| 858 | .await |
| 859 | .unwrap(); |
| 860 | let tools = ed.tools.read().await.clone(); |
| 861 | ed.execute_tool_call(&call("save_note"), &tools).await; |
| 862 | assert!(ed.session_grants.lock().await.contains("save_note")); |
| 863 | |
| 864 | ed.clear_history().await; |
| 865 | |
| 866 | assert!(ed.session_grants.lock().await.is_empty()); |
| 867 | } |
| 868 | |
| 869 | #[tokio::test] |
| 870 | async fn unknown_tools_return_an_error_result_instead_of_executing() { |
| 871 | let executor = Arc::new(CountingExecutor(AtomicUsize::new(0))); |
| 872 | let ed = Ed::with_agent( |
| 873 | NoopSink, |
| 874 | executor.clone(), |
| 875 | Arc::new(FixedApproval(ApprovalDecision::AllowOnce)), |
| 876 | AgentConfig::default(), |
| 877 | ); |
| 878 | |
| 879 | let result = ed |
| 880 | .execute_tool_call(&call("missing"), &HashMap::new()) |
| 881 | .await; |
| 882 | |
| 883 | assert!(result.is_error); |
| 884 | assert!(result.content.contains("not registered")); |
| 885 | assert_eq!(executor.0.load(Ordering::Relaxed), 0); |
| 886 | } |
| 887 | |
| 888 | #[tokio::test] |
| 889 | async fn tool_arguments_must_be_a_json_object() { |
| 890 | let executor = Arc::new(CountingExecutor(AtomicUsize::new(0))); |
| 891 | let ed = Ed::with_agent( |
| 892 | NoopSink, |
| 893 | executor.clone(), |
| 894 | Arc::new(FixedApproval(ApprovalDecision::AllowOnce)), |
| 895 | AgentConfig::default(), |
| 896 | ); |
| 897 | ed.register_tool(AgentToolDefinition::read_only( |
| 898 | "clock", |
| 899 | "Read the clock", |
| 900 | r#"{"type":"object"}"#, |
| 901 | )) |
| 902 | .await |
| 903 | .unwrap(); |
| 904 | let tools = ed.tools.read().await.clone(); |
| 905 | let mut invalid = call("clock"); |
| 906 | invalid.arguments = "[]".to_owned(); |
| 907 | |
| 908 | let result = ed.execute_tool_call(&invalid, &tools).await; |
| 909 | |
| 910 | assert!(result.is_error); |
| 911 | assert!(result.content.contains("JSON object")); |
| 912 | assert_eq!(executor.0.load(Ordering::Relaxed), 0); |
| 913 | } |
| 914 | |
| 915 | #[tokio::test] |
| 916 | async fn restoring_history_revokes_session_grants() { |
| 917 | let ed = Ed::with_agent( |
| 918 | NoopSink, |
| 919 | Arc::new(CountingExecutor(AtomicUsize::new(0))), |
| 920 | Arc::new(FixedApproval(ApprovalDecision::AllowForSession)), |
| 921 | AgentConfig::default(), |
| 922 | ); |
| 923 | ed.register_tool(AgentToolDefinition::mutating( |
| 924 | "save_note", |
| 925 | "Save a note", |
| 926 | r#"{"type":"object"}"#, |
| 927 | )) |
| 928 | .await |
| 929 | .unwrap(); |
| 930 | let tools = ed.tools.read().await.clone(); |
| 931 | ed.execute_tool_call(&call("save_note"), &tools).await; |
| 932 | assert!(ed.session_grants.lock().await.contains("save_note")); |
| 933 | |
| 934 | ed.restore_history(Vec::new()).await; |
| 935 | |
| 936 | assert!(ed.session_grants.lock().await.is_empty()); |
| 937 | } |
| 938 | |
| 939 | #[tokio::test] |
| 940 | async fn generation_and_streaming_need_a_loaded_model() { |
| 941 | let ed = Ed::new(); |
| 942 | assert!(matches!( |
| 943 | ed.generate(vec![ChatMessage::user("hi")], None).await, |
| 944 | Err(InferenceError::NoModelLoaded) |
| 945 | )); |
| 946 | assert!(matches!( |
| 947 | ed.stream("hi").await, |
| 948 | Err(InferenceError::NoModelLoaded) |
| 949 | )); |
| 950 | } |
| 951 | } |