1mod completion;
2mod construction;
3mod controls;
4mod cycle_inputs;
5mod helpers;
6mod logging;
7mod memory;
8mod planning;
9mod run_setup;
10mod state;
11
12use std::collections::BTreeMap;
13
14use serde_json::Value;
15
16use crate::llm::{LlmClient, LlmError, LlmRequest};
17use crate::memory::CompactionExhaustedError;
18use crate::tools::ToolContext;
19use crate::types::{AgentResult, AgentStatus, AgentTask, ToolDirective, ToolExecutionResult};
20
21use super::cancellation::CancellationToken;
22
23use super::cycle_runner::{is_prompt_too_long_error, MAX_PROMPT_TOO_LONG_RETRIES};
24use super::results::assistant_message_from_response;
25use super::token_usage::normalize_token_usage;
26use super::tool_call_runner::{execute_tool_result, needs_tool_call_id, skipped_tool_result};
27
28use self::completion::{handle_directive_result, handle_no_tool_response, NoToolResponseRequest};
29use self::helpers::{
30 cancelled_agent_result, collect_interruption_messages, controls_cancelled,
31 drain_steering_queue, failed_agent_result, image_notification_from_tool_result,
32 previous_cycle_memory_usage,
33};
34use self::run_setup::{prepare_run_setup, PreparedRun};
35pub use self::state::AgentRuntime;
36
37pub use crate::runtime::sub_agent_sessions::{
38 _register_sub_agent_session, _unregister_sub_agent_session, get_sub_agent_session,
39 steer_sub_agent_session, subscribe_sub_agent_session,
40};
41pub use controls::{
42 BeforeCycleMessageProvider, InterruptionMessageProvider, RuntimeEventHandler,
43 RuntimeLogCallback, RuntimeLogHandler, RuntimeRunControls,
44};
45
46impl<C: LlmClient + Clone + 'static> AgentRuntime<C> {
47 pub fn run(&self, task: AgentTask) -> Result<AgentResult, LlmError> {
48 self.run_with_controls(task, RuntimeRunControls::default())
49 }
50
51 pub fn run_with_controls(
52 &self,
53 task: AgentTask,
54 controls: RuntimeRunControls,
55 ) -> Result<AgentResult, LlmError> {
56 let PreparedRun {
57 task,
58 messages,
59 cycles,
60 shared_state,
61 workspace_path,
62 workspace_backend,
63 sub_task_manager,
64 mut memory_manager,
65 } = prepare_run_setup(self, task, &controls);
66 self.emit_run_started(&controls, &task, &workspace_path);
67
68 if controls_cancelled(&controls) {
69 self.emit_log(
70 &controls,
71 "run_cancelled",
72 BTreeMap::from([(
73 "error".to_string(),
74 Value::String("Operation was cancelled".to_string()),
75 )]),
76 );
77 return Ok(cancelled_agent_result(messages, cycles, shared_state));
78 }
79
80 let effective_cancellation_token = controls.effective_cancellation_token();
81 let effective_stream_callback = controls.effective_stream_callback();
82 let mut pending_error = None;
83 let result = self.execution_backend.execute(
84 &task,
85 messages,
86 shared_state,
87 |cycle_index, messages, cycles, shared_state, cancellation_token| {
88 if cancellation_token.is_some_and(CancellationToken::is_cancelled)
89 || controls_cancelled(&controls)
90 {
91 self.emit_log(
92 &controls,
93 "run_cancelled",
94 BTreeMap::from([
95 ("cycle".to_string(), Value::from(cycle_index)),
96 (
97 "error".to_string(),
98 Value::String("Operation was cancelled".to_string()),
99 ),
100 ]),
101 );
102 return Some(cancelled_agent_result(
103 messages.clone(),
104 cycles.clone(),
105 shared_state.clone(),
106 ));
107 }
108 self.apply_cycle_inputs(&controls, cycle_index, messages, shared_state);
109 if cancellation_token.is_some_and(CancellationToken::is_cancelled)
110 || controls_cancelled(&controls)
111 {
112 self.emit_log(
113 &controls,
114 "run_cancelled",
115 BTreeMap::from([
116 ("cycle".to_string(), Value::from(cycle_index)),
117 (
118 "error".to_string(),
119 Value::String("Operation was cancelled".to_string()),
120 ),
121 ]),
122 );
123 return Some(cancelled_agent_result(
124 messages.clone(),
125 cycles.clone(),
126 shared_state.clone(),
127 ));
128 }
129 self.emit_log(
130 &controls,
131 "cycle_started",
132 BTreeMap::from([
133 ("cycle".to_string(), Value::from(cycle_index)),
134 ("max_cycles".to_string(), Value::from(task.max_cycles)),
135 ("message_count".to_string(), Value::from(messages.len())),
136 ]),
137 );
138 let hook_manager = self.hook_manager();
139 let pre_compact_messages = hook_manager.apply_before_memory_compact(
140 &task,
141 cycle_index,
142 messages.clone(),
143 shared_state,
144 );
145 let pre_compact_messages =
146 memory_manager.apply_session_memory_context(&pre_compact_messages);
147 let (previous_prompt_tokens, recent_tool_call_ids) =
148 previous_cycle_memory_usage(cycles);
149 let (prepared_messages, memory_compacted) = memory_manager
150 .compact_for_cycle_with_usage(
151 &pre_compact_messages,
152 cycle_index,
153 false,
154 previous_prompt_tokens,
155 recent_tool_call_ids.as_ref(),
156 );
157 *messages = prepared_messages;
158 let tool_schemas = self.planned_tool_schemas(&task);
159 let llm_messages = memory_manager.apply_session_memory_context(messages);
160 let (request_messages, request_tool_schemas) = hook_manager.apply_before_llm(
161 &task,
162 cycle_index,
163 llm_messages,
164 tool_schemas,
165 shared_state,
166 );
167 let mut request_messages = request_messages;
168 let mut request_tool_schemas = request_tool_schemas;
169 let mut memory_compacted = memory_compacted;
170 let mut prompt_too_long_retries = 0;
171 let response = loop {
172 let mut request = LlmRequest::new(task.model.clone(), request_messages.clone());
173 request.tools = request_tool_schemas.clone();
174 if let Some(execution_context) = controls.execution_context.as_ref() {
175 request.metadata = serde_json::to_value(&execution_context.metadata)
176 .unwrap_or(Value::Null);
177 }
178 match self
179 .llm_client
180 .complete_with_stream(request, effective_stream_callback.clone())
181 {
182 Ok(response) => break response,
183 Err(error) if is_prompt_too_long_error(&error) => {
184 prompt_too_long_retries += 1;
185 if prompt_too_long_retries > MAX_PROMPT_TOO_LONG_RETRIES {
186 let error =
187 LlmError::CompactionExhausted(CompactionExhaustedError::new(
188 prompt_too_long_retries,
189 Some(error.to_string()),
190 ));
191 let message = error.to_string();
192 pending_error = Some(error);
193 return Some(failed_agent_result(
194 messages.clone(),
195 cycles.clone(),
196 shared_state.clone(),
197 message,
198 ));
199 }
200 memory_compacted = true;
201 let retry_messages = if prompt_too_long_retries == 1 {
202 let (compacted, _) = memory_manager.compact_for_cycle(
203 &request_messages,
204 cycle_index,
205 true,
206 );
207 compacted
208 } else {
209 memory_manager.emergency_compact(
210 &request_messages,
211 (0.2 * prompt_too_long_retries as f64).min(0.95),
212 )
213 };
214 let retry_tool_schemas = self.planned_tool_schemas(&task);
215 let llm_messages =
216 memory_manager.apply_session_memory_context(&retry_messages);
217 (request_messages, request_tool_schemas) = hook_manager
218 .apply_before_llm(
219 &task,
220 cycle_index,
221 llm_messages,
222 retry_tool_schemas,
223 shared_state,
224 );
225 }
226 Err(error) => {
227 let message = error.to_string();
228 pending_error = Some(error);
229 return Some(failed_agent_result(
230 messages.clone(),
231 cycles.clone(),
232 shared_state.clone(),
233 message,
234 ));
235 }
236 }
237 };
238 let response = hook_manager.apply_after_llm(
239 &task,
240 cycle_index,
241 &request_messages,
242 &request_tool_schemas,
243 response,
244 shared_state,
245 );
246 *messages = request_messages;
247 messages.push(assistant_message_from_response(&response));
248 let mut cycle = crate::types::CycleRecord::from_response(
249 cycle_index,
250 &response,
251 Vec::<ToolExecutionResult>::new(),
252 );
253 cycle.memory_compacted = memory_compacted;
254 if !cycle.token_usage.has_usage() {
255 cycle.token_usage =
256 normalize_token_usage(response.raw.get("usage").unwrap_or(&Value::Null));
257 }
258 self.emit_cycle_llm_response(&controls, &cycle);
259
260 if response.tool_calls.is_empty() {
261 if let Some(result) = handle_no_tool_response(NoToolResponseRequest {
262 runtime: self,
263 controls: &controls,
264 task: &task,
265 cycle_index,
266 response: &response,
267 messages,
268 cycles,
269 cycle,
270 shared_state,
271 }) {
272 return Some(result);
273 }
274 return None;
275 }
276
277 let sub_task_runner = self.build_sub_task_runner(
278 &task,
279 workspace_path.clone(),
280 workspace_backend.clone(),
281 shared_state.clone(),
282 sub_task_manager.clone(),
283 super::sub_agents::SubTaskCallbacks {
284 stream_callback: effective_stream_callback.clone(),
285 parent_log_handler: self.log_handler.clone(),
286 parent_event_handler: controls.log_handler.clone(),
287 },
288 );
289 let mut tool_metadata = controls
290 .execution_context
291 .as_ref()
292 .map(|context| context.metadata.clone())
293 .unwrap_or_default();
294 tool_metadata.extend(task.metadata.clone());
295 let mut context = ToolContext {
296 workspace: workspace_path.clone(),
297 shared_state: shared_state.clone(),
298 cycle_index,
299 task_id: task.task_id.clone(),
300 metadata: tool_metadata,
301 workspace_backend: workspace_backend.clone(),
302 model_provider: controls.model_provider.clone(),
303 sub_task_runner,
304 sub_task_manager: Some(sub_task_manager.clone()),
305 execution_backend: Some(self.execution_backend.clone()),
306 };
307
308 let mut directive_result = None;
309 let mut image_notifications = Vec::new();
310 for (call_index, call) in response.tool_calls.iter().enumerate() {
311 if cancellation_token.is_some_and(CancellationToken::is_cancelled)
312 || controls_cancelled(&controls)
313 {
314 *shared_state = context.shared_state.clone();
315 cycles.push(cycle);
316 self.emit_log(
317 &controls,
318 "run_cancelled",
319 BTreeMap::from([
320 ("cycle".to_string(), Value::from(cycle_index)),
321 (
322 "error".to_string(),
323 Value::String("Operation was cancelled".to_string()),
324 ),
325 ]),
326 );
327 return Some(cancelled_agent_result(
328 messages.clone(),
329 cycles.clone(),
330 shared_state.clone(),
331 ));
332 }
333 let (patched_call, short_circuit_result) = hook_manager.apply_before_tool_call(
334 &task,
335 cycle_index,
336 call.clone(),
337 &context,
338 );
339 let mut result = match short_circuit_result {
340 Some(mut result) => {
341 if needs_tool_call_id(&result.tool_call_id) {
342 result.tool_call_id = call.id.clone();
343 }
344 result
345 }
346 None => {
347 let mut result = execute_tool_result(
348 &self.tool_registry,
349 &patched_call,
350 &mut context,
351 );
352 if needs_tool_call_id(&result.tool_call_id) {
353 result.tool_call_id = patched_call.id.clone();
354 }
355 result
356 }
357 };
358 result = hook_manager.apply_after_tool_call(
359 &task,
360 cycle_index,
361 &patched_call,
362 &context,
363 result,
364 );
365 if needs_tool_call_id(&result.tool_call_id) {
366 result.tool_call_id = patched_call.id.clone();
367 }
368 self.emit_tool_result(&controls, cycle_index, &patched_call, &result);
369
370 let interruption_messages = collect_interruption_messages(&controls);
371 let steering_prompts = drain_steering_queue(&controls);
372 let steering_count = interruption_messages.len() + steering_prompts.len();
373 if steering_count == 0 && result.directive != ToolDirective::Continue {
374 directive_result = Some(result.clone());
375 }
376 messages.push(result.to_message());
377 if let Some(image_notification) =
378 image_notification_from_tool_result(&result, task.native_multimodal)
379 {
380 image_notifications.push(image_notification);
381 }
382 cycle.tool_results.push(result);
383 if steering_count > 0 {
384 for prompt in &steering_prompts {
385 self.emit_log(
386 &controls,
387 "session_steer_interrupt",
388 BTreeMap::from([
389 ("cycle".to_string(), Value::from(cycle_index)),
390 (
391 "after_tool_call_id".to_string(),
392 Value::String(patched_call.id.clone()),
393 ),
394 (
395 "after_tool_name".to_string(),
396 Value::String(patched_call.name.clone()),
397 ),
398 ("prompt".to_string(), Value::String(prompt.clone())),
399 ]),
400 );
401 }
402 for skipped_call in response.tool_calls.iter().skip(call_index + 1) {
403 let skipped = skipped_tool_result(
404 skipped_call,
405 "skipped_due_to_steering",
406 "Tool skipped due to queued steering message.",
407 );
408 self.emit_tool_result(&controls, cycle_index, skipped_call, &skipped);
409 messages.push(skipped.to_message());
410 cycle.tool_results.push(skipped);
411 }
412 for prompt in &steering_prompts {
413 messages.push(crate::types::Message::user(prompt.clone()));
414 }
415 messages.extend(interruption_messages);
416 self.emit_log(
417 &controls,
418 "run_steered",
419 BTreeMap::from([
420 ("cycle".to_string(), Value::from(cycle_index)),
421 (
422 "after_tool_call_id".to_string(),
423 Value::String(patched_call.id.clone()),
424 ),
425 (
426 "after_tool_name".to_string(),
427 Value::String(patched_call.name.clone()),
428 ),
429 (
430 "prompt_count".to_string(),
431 Value::from(steering_count as u64),
432 ),
433 (
434 "steering_count".to_string(),
435 Value::from(steering_count as u64),
436 ),
437 ]),
438 );
439 break;
440 }
441 if directive_result.is_some() {
442 let (error_code, message) = match directive_result
443 .as_ref()
444 .map(|result| result.directive)
445 .unwrap_or(ToolDirective::Continue)
446 {
447 ToolDirective::WaitUser => (
448 "skipped_due_to_wait_user",
449 "Tool skipped because a previous tool requested user input.",
450 ),
451 ToolDirective::Finish => (
452 "skipped_due_to_finish",
453 "Tool skipped because a previous tool finished the task.",
454 ),
455 ToolDirective::Continue => {
456 ("skipped_due_to_directive", "Tool skipped.")
457 }
458 };
459 for skipped_call in response.tool_calls.iter().skip(call_index + 1) {
460 let skipped = skipped_tool_result(skipped_call, error_code, message);
461 self.emit_tool_result(&controls, cycle_index, skipped_call, &skipped);
462 messages.push(skipped.to_message());
463 cycle.tool_results.push(skipped);
464 }
465 break;
466 }
467 }
468 messages.extend(image_notifications);
469 *shared_state = context.shared_state.clone();
470
471 cycles.push(cycle);
472 if let Some(directive_result) = directive_result.as_ref() {
473 if let Some(result) = handle_directive_result(
474 self,
475 &controls,
476 cycle_index,
477 directive_result,
478 messages,
479 cycles,
480 shared_state,
481 ) {
482 return Some(result);
483 }
484 }
485 None
486 },
487 effective_cancellation_token.as_ref(),
488 task.max_cycles,
489 );
490 if let Some(error) = pending_error {
491 return Err(error);
492 }
493 if result.status == AgentStatus::MaxCycles {
494 self.emit_run_max_cycles(&controls, &result);
495 }
496 Ok(result)
497 }
498}