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 match self
175 .llm_client
176 .complete_with_stream(request, effective_stream_callback.clone())
177 {
178 Ok(response) => break response,
179 Err(error) if is_prompt_too_long_error(&error) => {
180 prompt_too_long_retries += 1;
181 if prompt_too_long_retries > MAX_PROMPT_TOO_LONG_RETRIES {
182 let error =
183 LlmError::CompactionExhausted(CompactionExhaustedError::new(
184 prompt_too_long_retries,
185 Some(error.to_string()),
186 ));
187 let message = error.to_string();
188 pending_error = Some(error);
189 return Some(failed_agent_result(
190 messages.clone(),
191 cycles.clone(),
192 shared_state.clone(),
193 message,
194 ));
195 }
196 memory_compacted = true;
197 let retry_messages = if prompt_too_long_retries == 1 {
198 let (compacted, _) = memory_manager.compact_for_cycle(
199 &request_messages,
200 cycle_index,
201 true,
202 );
203 compacted
204 } else {
205 memory_manager.emergency_compact(
206 &request_messages,
207 (0.2 * prompt_too_long_retries as f64).min(0.95),
208 )
209 };
210 let retry_tool_schemas = self.planned_tool_schemas(&task);
211 let llm_messages =
212 memory_manager.apply_session_memory_context(&retry_messages);
213 (request_messages, request_tool_schemas) = hook_manager
214 .apply_before_llm(
215 &task,
216 cycle_index,
217 llm_messages,
218 retry_tool_schemas,
219 shared_state,
220 );
221 }
222 Err(error) => {
223 let message = error.to_string();
224 pending_error = Some(error);
225 return Some(failed_agent_result(
226 messages.clone(),
227 cycles.clone(),
228 shared_state.clone(),
229 message,
230 ));
231 }
232 }
233 };
234 let response = hook_manager.apply_after_llm(
235 &task,
236 cycle_index,
237 &request_messages,
238 &request_tool_schemas,
239 response,
240 shared_state,
241 );
242 *messages = request_messages;
243 messages.push(assistant_message_from_response(&response));
244 let mut cycle = crate::types::CycleRecord::from_response(
245 cycle_index,
246 &response,
247 Vec::<ToolExecutionResult>::new(),
248 );
249 cycle.memory_compacted = memory_compacted;
250 if !cycle.token_usage.has_usage() {
251 cycle.token_usage =
252 normalize_token_usage(response.raw.get("usage").unwrap_or(&Value::Null));
253 }
254 self.emit_cycle_llm_response(&controls, &cycle);
255
256 if response.tool_calls.is_empty() {
257 if let Some(result) = handle_no_tool_response(NoToolResponseRequest {
258 runtime: self,
259 controls: &controls,
260 task: &task,
261 cycle_index,
262 response: &response,
263 messages,
264 cycles,
265 cycle,
266 shared_state,
267 }) {
268 return Some(result);
269 }
270 return None;
271 }
272
273 let sub_task_runner = self.build_sub_task_runner(
274 &task,
275 workspace_path.clone(),
276 workspace_backend.clone(),
277 shared_state.clone(),
278 sub_task_manager.clone(),
279 super::sub_agents::SubTaskCallbacks {
280 stream_callback: effective_stream_callback.clone(),
281 parent_log_handler: self.log_handler.clone(),
282 parent_event_handler: controls.log_handler.clone(),
283 },
284 );
285 let mut tool_metadata = controls
286 .execution_context
287 .as_ref()
288 .map(|context| context.metadata.clone())
289 .unwrap_or_default();
290 tool_metadata.extend(task.metadata.clone());
291 let mut context = ToolContext {
292 workspace: workspace_path.clone(),
293 shared_state: shared_state.clone(),
294 cycle_index,
295 task_id: task.task_id.clone(),
296 metadata: tool_metadata,
297 workspace_backend: workspace_backend.clone(),
298 sub_task_runner,
299 sub_task_manager: Some(sub_task_manager.clone()),
300 execution_backend: Some(self.execution_backend.clone()),
301 };
302
303 let mut directive_result = None;
304 let mut image_notifications = Vec::new();
305 for (call_index, call) in response.tool_calls.iter().enumerate() {
306 if cancellation_token.is_some_and(CancellationToken::is_cancelled)
307 || controls_cancelled(&controls)
308 {
309 *shared_state = context.shared_state.clone();
310 cycles.push(cycle);
311 self.emit_log(
312 &controls,
313 "run_cancelled",
314 BTreeMap::from([
315 ("cycle".to_string(), Value::from(cycle_index)),
316 (
317 "error".to_string(),
318 Value::String("Operation was cancelled".to_string()),
319 ),
320 ]),
321 );
322 return Some(cancelled_agent_result(
323 messages.clone(),
324 cycles.clone(),
325 shared_state.clone(),
326 ));
327 }
328 let (patched_call, short_circuit_result) = hook_manager.apply_before_tool_call(
329 &task,
330 cycle_index,
331 call.clone(),
332 &context,
333 );
334 let mut result = match short_circuit_result {
335 Some(mut result) => {
336 if needs_tool_call_id(&result.tool_call_id) {
337 result.tool_call_id = call.id.clone();
338 }
339 result
340 }
341 None => {
342 let mut result = execute_tool_result(
343 &self.tool_registry,
344 &patched_call,
345 &mut context,
346 );
347 if needs_tool_call_id(&result.tool_call_id) {
348 result.tool_call_id = patched_call.id.clone();
349 }
350 result
351 }
352 };
353 result = hook_manager.apply_after_tool_call(
354 &task,
355 cycle_index,
356 &patched_call,
357 &context,
358 result,
359 );
360 if needs_tool_call_id(&result.tool_call_id) {
361 result.tool_call_id = patched_call.id.clone();
362 }
363 self.emit_tool_result(&controls, cycle_index, &patched_call, &result);
364
365 let interruption_messages = collect_interruption_messages(&controls);
366 let steering_prompts = drain_steering_queue(&controls);
367 let steering_count = interruption_messages.len() + steering_prompts.len();
368 if steering_count == 0 && result.directive != ToolDirective::Continue {
369 directive_result = Some(result.clone());
370 }
371 messages.push(result.to_message());
372 if let Some(image_notification) =
373 image_notification_from_tool_result(&result, task.native_multimodal)
374 {
375 image_notifications.push(image_notification);
376 }
377 cycle.tool_results.push(result);
378 if steering_count > 0 {
379 for prompt in &steering_prompts {
380 self.emit_log(
381 &controls,
382 "session_steer_interrupt",
383 BTreeMap::from([
384 ("cycle".to_string(), Value::from(cycle_index)),
385 (
386 "after_tool_call_id".to_string(),
387 Value::String(patched_call.id.clone()),
388 ),
389 (
390 "after_tool_name".to_string(),
391 Value::String(patched_call.name.clone()),
392 ),
393 ("prompt".to_string(), Value::String(prompt.clone())),
394 ]),
395 );
396 }
397 for skipped_call in response.tool_calls.iter().skip(call_index + 1) {
398 let skipped = skipped_tool_result(
399 skipped_call,
400 "skipped_due_to_steering",
401 "Tool skipped due to queued steering message.",
402 );
403 self.emit_tool_result(&controls, cycle_index, skipped_call, &skipped);
404 messages.push(skipped.to_message());
405 cycle.tool_results.push(skipped);
406 }
407 for prompt in &steering_prompts {
408 messages.push(crate::types::Message::user(prompt.clone()));
409 }
410 messages.extend(interruption_messages);
411 self.emit_log(
412 &controls,
413 "run_steered",
414 BTreeMap::from([
415 ("cycle".to_string(), Value::from(cycle_index)),
416 (
417 "after_tool_call_id".to_string(),
418 Value::String(patched_call.id.clone()),
419 ),
420 (
421 "after_tool_name".to_string(),
422 Value::String(patched_call.name.clone()),
423 ),
424 (
425 "prompt_count".to_string(),
426 Value::from(steering_count as u64),
427 ),
428 (
429 "steering_count".to_string(),
430 Value::from(steering_count as u64),
431 ),
432 ]),
433 );
434 break;
435 }
436 if directive_result.is_some() {
437 let (error_code, message) = match directive_result
438 .as_ref()
439 .map(|result| result.directive)
440 .unwrap_or(ToolDirective::Continue)
441 {
442 ToolDirective::WaitUser => (
443 "skipped_due_to_wait_user",
444 "Tool skipped because a previous tool requested user input.",
445 ),
446 ToolDirective::Finish => (
447 "skipped_due_to_finish",
448 "Tool skipped because a previous tool finished the task.",
449 ),
450 ToolDirective::Continue => {
451 ("skipped_due_to_directive", "Tool skipped.")
452 }
453 };
454 for skipped_call in response.tool_calls.iter().skip(call_index + 1) {
455 let skipped = skipped_tool_result(skipped_call, error_code, message);
456 self.emit_tool_result(&controls, cycle_index, skipped_call, &skipped);
457 messages.push(skipped.to_message());
458 cycle.tool_results.push(skipped);
459 }
460 break;
461 }
462 }
463 messages.extend(image_notifications);
464 *shared_state = context.shared_state.clone();
465
466 cycles.push(cycle);
467 if let Some(directive_result) = directive_result.as_ref() {
468 if let Some(result) = handle_directive_result(
469 self,
470 &controls,
471 cycle_index,
472 directive_result,
473 messages,
474 cycles,
475 shared_state,
476 ) {
477 return Some(result);
478 }
479 }
480 None
481 },
482 effective_cancellation_token.as_ref(),
483 task.max_cycles,
484 );
485 if let Some(error) = pending_error {
486 return Err(error);
487 }
488 if result.status == AgentStatus::MaxCycles {
489 self.emit_run_max_cycles(&controls, &result);
490 }
491 Ok(result)
492 }
493}