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 match self.llm_client.complete_with_stream(
148 llm_request,
149 request
150 .execution_context
151 .and_then(|context| context.stream_callback.clone()),
152 ) {
153 Ok(response) => break (response, request_messages, request_tool_schemas),
154 Err(error) if is_prompt_too_long_error(&error) => {
155 prompt_too_long_retries += 1;
156 if prompt_too_long_retries > MAX_PROMPT_TOO_LONG_RETRIES {
157 return Err(LlmError::CompactionExhausted(
158 CompactionExhaustedError::new(
159 prompt_too_long_retries,
160 Some(error.to_string()),
161 ),
162 ));
163 }
164 if prompt_too_long_retries == 1 {
165 (compacted_messages, _) =
166 request.memory_manager.compact_for_cycle_with_usage(
167 &compacted_messages,
168 request.cycle_index,
169 true,
170 None,
171 request.recent_tool_call_ids,
172 );
173 } else {
174 compacted_messages = request.memory_manager.emergency_compact(
175 &compacted_messages,
176 (0.2 * f64::from(prompt_too_long_retries)).min(0.95),
177 );
178 }
179 memory_compacted = true;
180 }
181 Err(error) => return Err(error),
182 }
183 };
184
185 if let Some(context) = request.execution_context {
186 check_context_cancelled(context)?;
187 }
188 let response = self.hook_manager.apply_after_llm(
189 request.task,
190 request.cycle_index,
191 &request_messages,
192 &request_tool_schemas,
193 response,
194 shared_state,
195 );
196 let mut next_messages = request_messages;
197 next_messages.push(assistant_message_from_response(&response));
198 let mut cycle = CycleRecord::from_response(request.cycle_index, &response, Vec::new());
199 cycle.memory_compacted = memory_compacted;
200 if !cycle.token_usage.has_usage() {
201 cycle.token_usage =
202 normalize_token_usage(response.raw.get("usage").unwrap_or(&Value::Null));
203 }
204 Ok((next_messages, cycle))
205 }
206}
207
208fn check_context_cancelled(context: &ExecutionContext) -> Result<(), LlmError> {
209 context.check_cancelled().map_err(LlmError::Request)
210}