1use std::collections::{BTreeMap, BTreeSet};
2
3use serde_json::Value;
4
5use super::context::ExecutionContext;
6use super::hooks::RuntimeHookManager;
7use super::model_calls::{ModelCallCoordinator, ModelCallLedger};
8use super::results::assistant_message_from_response;
9use crate::llm::{LlmClient, LlmError, LlmRequest};
10use crate::memory::{CompactionExhaustedError, MemoryManager};
11use crate::tools::ToolRegistry;
12use crate::types::{AgentTask, CycleRecord, Message, ModelCallOperation};
13
14pub const MAX_PROMPT_TOO_LONG_RETRIES: u32 = 3;
15
16const PROMPT_TOO_LONG_PATTERNS: &[&str] = &[
17 "prompt is too long",
18 "prompt_too_long",
19 "context_length_exceeded",
20 "maximum context length",
21 "request too large",
22 "too many tokens",
23];
24
25pub fn is_prompt_too_long_error(error: &LlmError) -> bool {
26 let text = error.to_string().to_ascii_lowercase();
27 PROMPT_TOO_LONG_PATTERNS
28 .iter()
29 .any(|pattern| text.contains(pattern))
30}
31
32pub struct CycleRunRequest<'a> {
33 pub task: &'a AgentTask,
34 pub messages: Vec<Message>,
35 pub cycle_index: u32,
36 pub memory_manager: &'a mut MemoryManager,
37 pub previous_prompt_tokens: Option<u64>,
38 pub recent_tool_call_ids: Option<&'a BTreeSet<String>>,
39 pub shared_state: Option<&'a BTreeMap<String, Value>>,
40 pub execution_context: Option<&'a ExecutionContext>,
41}
42
43impl<'a> CycleRunRequest<'a> {
44 pub fn new(
45 task: &'a AgentTask,
46 messages: Vec<Message>,
47 cycle_index: u32,
48 memory_manager: &'a mut MemoryManager,
49 ) -> Self {
50 Self {
51 task,
52 messages,
53 cycle_index,
54 memory_manager,
55 previous_prompt_tokens: None,
56 recent_tool_call_ids: None,
57 shared_state: None,
58 execution_context: None,
59 }
60 }
61
62 pub fn with_previous_prompt_tokens(mut self, previous_prompt_tokens: Option<u64>) -> Self {
63 self.previous_prompt_tokens = previous_prompt_tokens;
64 self
65 }
66
67 pub fn with_recent_tool_call_ids(mut self, recent_tool_call_ids: &'a BTreeSet<String>) -> Self {
68 self.recent_tool_call_ids = Some(recent_tool_call_ids);
69 self
70 }
71
72 pub fn with_shared_state(mut self, shared_state: &'a BTreeMap<String, Value>) -> Self {
73 self.shared_state = Some(shared_state);
74 self
75 }
76
77 pub fn with_execution_context(mut self, execution_context: &'a ExecutionContext) -> Self {
78 self.execution_context = Some(execution_context);
79 self
80 }
81}
82
83pub struct CycleRunner<C: LlmClient> {
84 llm_client: C,
85 tool_registry: ToolRegistry,
86 hook_manager: RuntimeHookManager,
87}
88
89impl<C: LlmClient> CycleRunner<C> {
90 pub fn new(llm_client: C, tool_registry: ToolRegistry) -> Self {
91 Self {
92 llm_client,
93 tool_registry,
94 hook_manager: RuntimeHookManager::default(),
95 }
96 }
97
98 pub fn with_hook_manager(mut self, hook_manager: RuntimeHookManager) -> Self {
99 self.hook_manager = hook_manager;
100 self
101 }
102
103 pub fn run_cycle(
104 &self,
105 request: CycleRunRequest<'_>,
106 ) -> Result<(Vec<Message>, CycleRecord), LlmError> {
107 if let Some(context) = request.execution_context {
108 check_context_cancelled(context)?;
109 }
110 let empty_shared_state = BTreeMap::new();
111 let shared_state = request.shared_state.unwrap_or(&empty_shared_state);
112 let pre_compact_messages = self.hook_manager.apply_before_memory_compact(
113 request.task,
114 request.cycle_index,
115 request.messages,
116 shared_state,
117 );
118 let (mut compacted_messages, mut memory_compacted) =
119 request.memory_manager.compact_for_cycle_with_usage(
120 &pre_compact_messages,
121 request.cycle_index,
122 false,
123 request.previous_prompt_tokens,
124 request.recent_tool_call_ids,
125 );
126
127 let mut prompt_too_long_retries = 0;
128 let coordinator = request
129 .execution_context
130 .and_then(|context| context.runtime_state.model_call_coordinator.clone())
131 .unwrap_or_else(|| {
132 let event_handler = request
133 .execution_context
134 .and_then(|context| context.event_handler.clone());
135 ModelCallCoordinator::new(
136 ModelCallLedger::default(),
137 &request.task.task_id,
138 &request.task.task_id,
139 &request.task.task_id,
140 None,
141 None,
142 event_handler,
143 None,
144 )
145 });
146 let (response, request_messages, request_tool_schemas) = loop {
147 let llm_messages = compacted_messages.clone();
148 let tool_schemas = self.tool_registry.planned_openai_schemas(request.task);
149 let (request_messages, request_tool_schemas) = self.hook_manager.apply_before_llm(
150 request.task,
151 request.cycle_index,
152 llm_messages,
153 tool_schemas,
154 shared_state,
155 );
156 request
157 .memory_manager
158 .validate_model_recovery_surface(&request_messages, &request_tool_schemas)
159 .map_err(LlmError::Request)?;
160 if let Some(context) = request.execution_context {
161 check_context_cancelled(context)?;
162 }
163 let mut llm_request = LlmRequest::new(
164 request.task.model.clone(),
165 request_messages.clone(),
166 request.task.prompt_bundle.clone(),
167 );
168 llm_request.tools = request_tool_schemas.clone();
169 llm_request.metadata =
170 Value::Object(request.task.metadata.clone().into_iter().collect());
171 llm_request.model_settings = request.task.model_settings.clone();
172 let backend = request
173 .execution_context
174 .and_then(|context| {
175 context
176 .metadata
177 .get("_vv_agent_resolved_backend")
178 .and_then(Value::as_str)
179 })
180 .unwrap_or("direct");
181 let model = request
182 .execution_context
183 .and_then(|context| {
184 context
185 .metadata
186 .get("_vv_agent_resolved_model")
187 .and_then(Value::as_str)
188 })
189 .unwrap_or(&request.task.model);
190 let operation_slot = if prompt_too_long_retries == 0 {
191 "main".to_string()
192 } else {
193 format!("prompt_too_long_{prompt_too_long_retries}")
194 };
195 match coordinator.dispatch(
196 ModelCallOperation::AgentCycle,
197 request.cycle_index,
198 &operation_slot,
199 backend,
200 model,
201 &llm_request,
202 || self.llm_client.complete(llm_request.clone()),
203 ) {
204 Ok(dispatch) => break (dispatch.response, request_messages, request_tool_schemas),
205 Err(error) if is_prompt_too_long_error(&error) => {
206 prompt_too_long_retries += 1;
207 if prompt_too_long_retries > MAX_PROMPT_TOO_LONG_RETRIES {
208 return Err(LlmError::CompactionExhausted(
209 CompactionExhaustedError::new(
210 prompt_too_long_retries,
211 Some(error.to_string()),
212 ),
213 ));
214 }
215 if prompt_too_long_retries == 1 {
216 (compacted_messages, _) =
217 request.memory_manager.compact_for_cycle_with_usage(
218 &compacted_messages,
219 request.cycle_index,
220 true,
221 None,
222 request.recent_tool_call_ids,
223 );
224 } else {
225 compacted_messages = request.memory_manager.emergency_compact(
226 &compacted_messages,
227 (0.2 * f64::from(prompt_too_long_retries)).min(0.95),
228 );
229 }
230 memory_compacted = true;
231 }
232 Err(error) => return Err(error),
233 }
234 };
235
236 if let Some(context) = request.execution_context {
237 check_context_cancelled(context)?;
238 }
239 let response = self.hook_manager.apply_after_llm(
240 request.task,
241 request.cycle_index,
242 &request_messages,
243 &request_tool_schemas,
244 response,
245 shared_state,
246 );
247 let mut next_messages = request_messages;
248 next_messages.push(assistant_message_from_response(&response));
249 let mut cycle = CycleRecord::from_response(request.cycle_index, &response, Vec::new());
250 cycle.memory_compacted = memory_compacted;
251 Ok((next_messages, cycle))
252 }
253}
254
255fn check_context_cancelled(context: &ExecutionContext) -> Result<(), LlmError> {
256 context
257 .check_cancelled()
258 .map_err(|error| LlmError::Request(error.to_string()))
259}