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