Skip to content
zhenyi

gitdata-core

0
Clone

HTTPS

SSH

SSH needs a key in your settings.

Download ZIP

feat(agent): add context token limits and compaction logic for agent history

ZhenYi committed
parent b1477a1 commit 075a265 Browse files

7 files changed

+610 −167

File tree 7
Modified config/default.toml View config/default.toml
@@ -270,3 +270,5 @@
270270tool_output_max_bytes = 32768
271271artifact_max_size = 104857600
272272artifact_dir = ""
273context_max_tokens = 128000
274compact_keep_recent = 16
Added lib/agent/src/runtime/compact.rs View lib/agent/src/runtime/compact.rs
@@ -0,0 +1,263 @@
1use rig::client::CompletionClient;
2use rig::completion::message::{AssistantContent, ToolResultContent, UserContent};
3use rig::completion::{CompletionModel, Message};
4use rig::providers::openai::CompletionsClient;
5
6use super::events::CompactMethod;
7
8pub const CHARS_PER_TOKEN: usize = 4;
9
10const PER_MESSAGE_OVERHEAD_TOKENS: usize = 4;
11const TOOL_RESULT_RENDER_MAX_CHARS: usize = 2_000;
12const TOOL_CALL_ARGS_RENDER_MAX_CHARS: usize = 500;
13const SUMMARY_PREFIX: &str = "[Summary of earlier conversation]";
14const TRUNCATE_MARKER: &str = "[Earlier messages were truncated to fit the context window]";
15
16const SUMMARY_INSTRUCTION: &str = "You are a conversation compactor for an autonomous engineering agent. \
17Summarize the transcript of the earlier conversation between the user, the assistant and its sandbox tools. \
18Preserve exactly: the original task goal and constraints; key decisions and their rationale; every file path, command, sandbox id, artifact name and identifier later steps may need; what has been completed and verified; what remains to be done next. \
19Be dense and factual, prefer bullet lists, no commentary. Output only the summary.";
20
21pub struct Compaction {
22 pub history: Vec<Message>,
23 pub method: CompactMethod,
24 pub before: usize,
25 pub after: usize,
26 pub estimated_tokens: usize,
27 pub summary: Option<String>,
28 pub prefix_len: usize,
29}
30
31#[derive(Clone)]
32pub struct Summarizer {
33 client: CompletionsClient,
34 model: String,
35}
36
37impl Summarizer {
38 pub fn new(client: CompletionsClient, model: String) -> Self {
39 Self { client, model }
40 }
41
42 pub async fn summarize(&self, transcript: &str) -> std::result::Result<String, String> {
43 let response = self
44 .client
45 .completion_model(self.model.clone())
46 .completion_request(Message::user(transcript.to_string()))
47 .preamble(SUMMARY_INSTRUCTION.to_string())
48 .send()
49 .await
50 .map_err(|e| e.to_string())?;
51 let text: String = response
52 .choice
53 .iter()
54 .filter_map(|c| match c {
55 AssistantContent::Text(t) => Some(t.text.as_str()),
56 _ => None,
57 })
58 .collect::<Vec<_>>()
59 .join("");
60 let text = text.trim();
61 if text.is_empty() {
62 return Err("model returned an empty summary".to_string());
63 }
64 Ok(text.to_string())
65 }
66}
67
68pub fn estimate_message_tokens(message: &Message) -> usize {
69 let chars = serde_json::to_string(message)
70 .map(|s| s.len())
71 .unwrap_or(4 * CHARS_PER_TOKEN);
72 chars / CHARS_PER_TOKEN + PER_MESSAGE_OVERHEAD_TOKENS
73}
74
75pub fn estimate_tokens(history: &[Message]) -> usize {
76 history.iter().map(estimate_message_tokens).sum()
77}
78
79pub fn trigger_threshold(context_max_tokens: usize) -> usize {
80 context_max_tokens * 3 / 4
81}
82
83pub fn summary_message(summary: &str) -> Message {
84 Message::user(format!("{SUMMARY_PREFIX}\n{summary}"))
85}
86
87pub(crate) fn starts_with_tool_result(message: &Message) -> bool {
88 matches!(
89 message,
90 Message::User { content }
91 if content
92 .iter()
93 .any(|c| matches!(c, UserContent::ToolResult(_)))
94 )
95}
96
97pub fn safe_split(history: &[Message], keep_recent: usize) -> (&[Message], &[Message]) {
98 let len = history.len();
99 if len <= keep_recent {
100 return (&[], history);
101 }
102 let mut cut = len - keep_recent;
103 while cut > 0 && starts_with_tool_result(&history[cut]) {
104 cut -= 1;
105 }
106 (&history[..cut], &history[cut..])
107}
108
109pub async fn compact_history(
110 summarizer: &Summarizer,
111 history: Vec<Message>,
112 budget_tokens: usize,
113 keep_recent: usize,
114) -> Compaction {
115 let before = history.len();
116 let (prefix, tail) = safe_split(&history, keep_recent);
117 if prefix.is_empty() {
118 let (truncated, dropped) = truncate_to_budget(&history, budget_tokens);
119 let estimated_tokens = estimate_tokens(&truncated);
120 let after = truncated.len();
121 return Compaction {
122 history: truncated,
123 method: CompactMethod::Truncate,
124 before,
125 after,
126 estimated_tokens,
127 summary: None,
128 prefix_len: dropped,
129 };
130 }
131
132 let transcript = render_transcript(prefix, budget_tokens * CHARS_PER_TOKEN);
133 match summarizer.summarize(&transcript).await {
134 Ok(summary) => {
135 let tail_len = tail.len();
136 let mut compacted = Vec::with_capacity(tail_len + 1);
137 compacted.push(summary_message(&summary));
138 compacted.extend(tail.iter().cloned());
139 if estimate_tokens(&compacted) > budget_tokens {
140 let (truncated, _) = truncate_to_budget(&compacted, budget_tokens);
141 compacted = truncated;
142 }
143 let estimated_tokens = estimate_tokens(&compacted);
144 let after = compacted.len();
145 Compaction {
146 history: compacted,
147 method: CompactMethod::Summary,
148 before,
149 after,
150 estimated_tokens,
151 summary: Some(summary),
152 prefix_len: before - tail_len,
153 }
154 }
155 Err(e) => {
156 tracing::warn!(error = %e, "history summarization failed, falling back to truncation");
157 let (truncated, dropped) = truncate_to_budget(&history, budget_tokens);
158 let estimated_tokens = estimate_tokens(&truncated);
159 let after = truncated.len();
160 Compaction {
161 history: truncated,
162 method: CompactMethod::Truncate,
163 before,
164 after,
165 estimated_tokens,
166 summary: None,
167 prefix_len: dropped,
168 }
169 }
170 }
171}
172
173pub fn truncate_to_budget(history: &[Message], budget_tokens: usize) -> (Vec<Message>, usize) {
174 let marker = Message::user(TRUNCATE_MARKER);
175 let marker_tokens = estimate_message_tokens(&marker);
176 let mut start = 0;
177 loop {
178 let kept_estimate = estimate_tokens(&history[start..]) + marker_tokens;
179 if kept_estimate <= budget_tokens || start >= history.len() {
180 let mut out = Vec::with_capacity(history.len() - start + 1);
181 out.push(marker);
182 out.extend_from_slice(&history[start..]);
183 return (out, start);
184 }
185 start += 1;
186 while start < history.len() && starts_with_tool_result(&history[start]) {
187 start += 1;
188 }
189 }
190}
191
192pub fn render_transcript(messages: &[Message], max_chars: usize) -> String {
193 let mut lines: Vec<String> = Vec::with_capacity(messages.len());
194 for message in messages {
195 match message {
196 Message::System { content } => lines.push(format!("[system] {content}")),
197 Message::User { content } => {
198 for part in content {
199 match part {
200 UserContent::Text(t) => lines.push(format!("[user] {}", t.text)),
201 UserContent::ToolResult(r) => {
202 let body = render_tool_result(&r.content);
203 lines.push(format!("[tool_result {}] {body}", r.name));
204 }
205 _ => lines.push("[user] [non-text content]".to_string()),
206 }
207 }
208 }
209 Message::Assistant { content, .. } => {
210 for part in content {
211 match part {
212 AssistantContent::Text(t) => {
213 lines.push(format!("[assistant] {}", t.text));
214 }
215 AssistantContent::ToolCall(tc) => {
216 let args = tc.function.arguments.to_string();
217 let args = truncate_chars(&args, TOOL_CALL_ARGS_RENDER_MAX_CHARS);
218 lines.push(format!("[tool_call {}] {args}", tc.function.name));
219 }
220 AssistantContent::Reasoning(_) => {
221 lines.push("[assistant] [reasoning]".to_string());
222 }
223 _ => lines.push("[assistant] [non-text content]".to_string()),
224 }
225 }
226 }
227 }
228 }
229 let full = lines.join("\n");
230 if full.len() <= max_chars {
231 return full;
232 }
233 let mut start = full.len() - max_chars;
234 while start < full.len() && !full.is_char_boundary(start) {
235 start += 1;
236 }
237 format!("[transcript truncated]\n{}", &full[start..])
238}
239
240fn render_tool_result(content: &[ToolResultContent]) -> String {
241 let mut out = String::new();
242 for part in content {
243 match part {
244 ToolResultContent::Text(t) => out.push_str(&t.text),
245 ToolResultContent::Json { value } => out.push_str(&value.to_string()),
246 _ => out.push_str("[non-text content]"),
247 }
248 }
249 truncate_chars(&out, TOOL_RESULT_RENDER_MAX_CHARS)
250}
251
252fn truncate_chars(s: &str, max_chars: usize) -> String {
253 if s.len() <= max_chars {
254 return s.to_string();
255 }
256 let mut end = max_chars;
257 while end > 0 && !s.is_char_boundary(end) {
258 end -= 1;
259 }
260 let mut out = s[..end].to_string();
261 out.push('…');
262 out
263}
Modified lib/agent/src/runtime/config.rs View lib/agent/src/runtime/config.rs
@@ -28,6 +28,8 @@
2828 pub tool_output_max_bytes: usize,
2929 pub artifact_max_size: i64,
3030 pub artifact_dir: String,
31 pub context_max_tokens: usize,
32 pub compact_keep_recent: usize,
3133}
3234
3335#[derive(Clone, Debug)]
@@ -59,6 +61,8 @@
5961 tool_output_max_bytes: cfg.runtime.tool_output_max_bytes,
6062 artifact_max_size: cfg.runtime.artifact_max_size,
6163 artifact_dir: cfg.runtime.artifact_dir.clone(),
64 context_max_tokens: cfg.runtime.context_max_tokens,
65 compact_keep_recent: cfg.runtime.compact_keep_recent,
6266 },
6367 }
6468 }
@@ -121,6 +125,16 @@
121125 "agent artifact_max_size must be positive".to_string(),
122126 ));
123127 }
128 if self.limits.context_max_tokens == 0 {
129 return Err(AgentError::Config(
130 "agent context_max_tokens must be positive".to_string(),
131 ));
132 }
133 if self.limits.compact_keep_recent == 0 {
134 return Err(AgentError::Config(
135 "agent compact_keep_recent must be positive".to_string(),
136 ));
137 }
124138 Ok(())
125139 }
126140}
Modified lib/agent/src/runtime/engine.rs View lib/agent/src/runtime/engine.rs
@@ -1,4 +1,4 @@
1use std::sync::Arc;
1use std::sync::{Arc, Mutex};
22
33use export::{AppDatabase, AppState};
44use futures::StreamExt;
@@ -7,7 +7,7 @@
77use models::agent::run::{self, AgentRunState};
88use rig::agent::{
99 AgentHook, CompletionCallAction, CompletionCallEvent, HookContext, MultiTurnStreamItem,
10 PromptResponse, StreamingError,
10 PromptResponse, RequestPatch, StreamingError,
1111};
1212use rig::client::AgentClientExt;
1313use rig::completion::message::ToolResultContent;
@@ -21,6 +21,7 @@
2121use tokio::task::JoinHandle;
2222use uuid::Uuid;
2323
24use super::compact::{self, Summarizer};
2425use super::config::RuntimeConfig;
2526use super::events::{RunEvent, StopReason};
2627use super::persist::{self, RunContext, RunTerminal};
@@ -191,6 +192,35 @@
191192 })
192193 .await;
193194
195 let context_max_tokens = self.config.limits.context_max_tokens;
196 let keep_recent = self.config.limits.compact_keep_recent;
197 let summarizer = Summarizer::new(self.client.clone(), self.config.provider.model.clone());
198 let prompt_message = Message::user(ctx.run.prompt.clone());
199 let mut working_history = ctx.history.clone();
200 let prompt_estimate = compact::estimate_tokens(&working_history)
201 + compact::estimate_message_tokens(&prompt_message);
202 if prompt_estimate > compact::trigger_threshold(context_max_tokens) {
203 let compaction = compact::compact_history(
204 &summarizer,
205 working_history,
206 context_max_tokens / 2,
207 keep_recent,
208 )
209 .await;
210 emit_compacted(&emitter, &compaction).await;
211 working_history = compaction.history;
212 }
213
214 let hook = RunHook {
215 db: db.clone(),
216 run_id,
217 summarizer: summarizer.clone(),
218 context_max_tokens,
219 compact_keep_recent: keep_recent,
220 emitter: emitter.clone(),
221 cache: Arc::new(Mutex::new(None)),
222 };
223
194224 let manager = match SandboxManager::new(&self.config.sandbox) {
195225 Ok(m) => m,
196226 Err(e) => {
@@ -222,194 +252,217 @@
222252 run_id,
223253 emitter: emitter.clone(),
224254 });
225 let agent = self
226 .client
227 .agent(self.config.provider.model.clone())
228 .preamble(
229 ctx.agent
230 .preamble
231 .as_deref()
232 .map(str::trim)
233 .filter(|p| !p.is_empty())
234 .unwrap_or(DEFAULT_PREAMBLE),
235 )
236 .dynamic_tools(build_tools(tool_env.clone()))
237 .default_max_turns(self.config.limits.default_max_turns)
238 .add_hook(CancelHook {
239 db: db.clone(),
240 run_id,
241 })
242 .build();
255 let preamble = ctx
256 .agent
257 .preamble
258 .as_deref()
259 .map(str::trim)
260 .filter(|p| !p.is_empty())
261 .unwrap_or(DEFAULT_PREAMBLE);
243262
244 let mut stream = agent
245 .stream_chat(
246 Message::user(ctx.run.prompt.clone()),
247 ctx.history.clone(),
248 )
249 .max_turns(self.config.limits.default_max_turns)
250 .await;
263 let mut usage_sum = Usage::new();
264 let mut attempt = 0;
265 let terminal = loop {
266 let agent = self
267 .client
268 .agent(self.config.provider.model.clone())
269 .preamble(preamble)
270 .dynamic_tools(build_tools(tool_env.clone()))
271 .default_max_turns(self.config.limits.default_max_turns)
272 .add_hook(hook.clone())
273 .build();
274 let mut stream = agent
275 .stream_chat(prompt_message.clone(), working_history.clone())
276 .max_turns(self.config.limits.default_max_turns)
277 .await;
251278
252 let mut final_response: Option<PromptResponse> = None;
253 let mut stream_error: Option<StreamingError> = None;
254 let mut usage_sum = Usage::new();
255 let mut consumer_gone = false;
256 while let Some(item) = stream.next().await {
257 match item {
258 Ok(MultiTurnStreamItem::StreamAssistantItem(content)) => match content {
259 StreamedAssistantContent::Text(t) => {
260 emitter
261 .emit(RunEvent::TextDelta { text: t.text })
262 .await;
263 }
264 StreamedAssistantContent::ReasoningDelta { reasoning, .. } => {
279 let mut final_response: Option<PromptResponse> = None;
280 let mut stream_error: Option<StreamingError> = None;
281 let mut consumer_gone = false;
282 while let Some(item) = stream.next().await {
283 match item {
284 Ok(MultiTurnStreamItem::StreamAssistantItem(content)) => match content {
285 StreamedAssistantContent::Text(t) => {
286 emitter
287 .emit(RunEvent::TextDelta { text: t.text })
288 .await;
289 }
290 StreamedAssistantContent::ReasoningDelta { reasoning, .. } => {
291 emitter
292 .emit(RunEvent::ReasoningDelta { text: reasoning })
293 .await;
294 }
295 StreamedAssistantContent::ToolCall { tool_call, .. } => {
296 emitter
297 .emit(RunEvent::ToolCall {
298 id: tool_call.id.as_str().to_string(),
299 name: tool_call.function.name.clone(),
300 args: tool_call.function.arguments.clone(),
301 })
302 .await;
303 }
304 _ => {}
305 },
306 Ok(MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult {
307 tool_result,
308 ..
309 })) => {
310 let mut output = String::new();
311 for part in &tool_result.content {
312 match part {
313 ToolResultContent::Text(t) => output.push_str(&t.text),
314 ToolResultContent::Json { value } => {
315 output.push_str(&value.to_string());
316 }
317 _ => {}
318 }
319 }
320 let (mut output, truncated) =
321 truncate_output(&output, self.config.limits.tool_output_max_bytes);
322 if truncated {
323 output.push_str("\n[truncated]");
324 }
265325 emitter
266 .emit(RunEvent::ReasoningDelta { text: reasoning })
326 .emit(RunEvent::ToolResult {
327 id: tool_result.call.as_str().to_string(),
328 name: tool_result.name.clone(),
329 output,
330 })
267331 .await;
268332 }
269 StreamedAssistantContent::ToolCall { tool_call, .. } => {
333 Ok(MultiTurnStreamItem::CompletionCall(call)) => {
334 usage_sum += call.usage;
270335 emitter
271 .emit(RunEvent::ToolCall {
272 id: tool_call.id.as_str().to_string(),
273 name: tool_call.function.name.clone(),
274 args: tool_call.function.arguments.clone(),
336 .emit(RunEvent::Usage {
337 input_tokens: call.usage.input_tokens,
338 output_tokens: call.usage.output_tokens,
339 total_tokens: call.usage.total_tokens,
275340 })
276341 .await;
277342 }
278 _ => {}
279 },
280 Ok(MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult {
281 tool_result,
282 ..
283 })) => {
284 let mut output = String::new();
285 for part in &tool_result.content {
286 match part {
287 ToolResultContent::Text(t) => output.push_str(&t.text),
288 ToolResultContent::Json { value } => {
289 output.push_str(&value.to_string());
290 }
291 _ => {}
292 }
343 Ok(MultiTurnStreamItem::FinalResponse(resp)) => final_response = Some(resp),
344 Ok(_) => {}
345 Err(e) => {
346 stream_error = Some(e);
347 break;
293348 }
294 let (mut output, truncated) =
295 truncate_output(&output, self.config.limits.tool_output_max_bytes);
296 if truncated {
297 output.push_str("\n[truncated]");
298 }
299 emitter
300 .emit(RunEvent::ToolResult {
301 id: tool_result.call.as_str().to_string(),
302 name: tool_result.name.clone(),
303 output,
304 })
305 .await;
306349 }
307 Ok(MultiTurnStreamItem::CompletionCall(call)) => {
308 usage_sum += call.usage;
309 emitter
310 .emit(RunEvent::Usage {
311 input_tokens: call.usage.input_tokens,
312 output_tokens: call.usage.output_tokens,
313 total_tokens: call.usage.total_tokens,
314 })
315 .await;
316 }
317 Ok(MultiTurnStreamItem::FinalResponse(resp)) => final_response = Some(resp),
318 Ok(_) => {}
319 Err(e) => {
320 stream_error = Some(e);
350 if consumer_cancel && emitter.consumer_closed() {
351 consumer_gone = true;
321352 break;
322353 }
323354 }
324 if consumer_cancel && emitter.consumer_closed() {
325 consumer_gone = true;
326 break;
327 }
328 }
329 drop(stream);
355 drop(stream);
330356
331 let terminal = if consumer_gone {
332 Terminal {
333 state: AgentRunState::Cancelled,
334 stop_reason: StopReason::Cancelled,
335 final_text: None,
336 error: Some("consumer disconnected".to_string()),
337 history: ctx.history.clone(),
338 usage: usage_sum,
357 if !consumer_gone
358 && attempt == 0
359 && stream_error.as_ref().is_some_and(is_context_overflow)
360 {
361 attempt += 1;
362 tracing::warn!(%run_id, "context overflow reported by provider, compacting history and retrying");
363 let compaction = compact::compact_history(
364 &summarizer,
365 working_history,
366 context_max_tokens / 4,
367 keep_recent.min(8),
368 )
369 .await;
370 emit_compacted(&emitter, &compaction).await;
371 working_history = compaction.history;
372 continue;
339373 }
340 } else if let Some(e) = stream_error {
341 match e {
342 StreamingError::Prompt(pe) => match *pe {
343 PromptError::MaxTurnsError {
344 max_turns,
345 chat_history,
346 ..
347 } => Terminal {
348 state: AgentRunState::Failure,
349 stop_reason: StopReason::MaxTurns,
350 final_text: None,
351 error: Some(format!("max turns exceeded ({max_turns})")),
352 history: *chat_history,
353 usage: usage_sum,
374
375 break if consumer_gone {
376 Terminal {
377 state: AgentRunState::Cancelled,
378 stop_reason: StopReason::Cancelled,
379 final_text: None,
380 error: Some("consumer disconnected".to_string()),
381 history: working_history,
382 usage: usage_sum,
383 }
384 } else if let Some(e) = stream_error {
385 match e {
386 StreamingError::Prompt(pe) => match *pe {
387 PromptError::MaxTurnsError {
388 max_turns,
389 chat_history,
390 ..
391 } => Terminal {
392 state: AgentRunState::Failure,
393 stop_reason: StopReason::MaxTurns,
394 final_text: None,
395 error: Some(format!("max turns exceeded ({max_turns})")),
396 history: *chat_history,
397 usage: usage_sum,
398 },
399 PromptError::PromptCancelled {
400 chat_history,
401 reason,
402 } => Terminal {
403 state: AgentRunState::Cancelled,
404 stop_reason: StopReason::Cancelled,
405 final_text: None,
406 error: Some(reason),
407 history: chat_history,
408 usage: usage_sum,
409 },
410 other => Terminal {
411 state: AgentRunState::Failure,
412 stop_reason: StopReason::Error,
413 final_text: None,
414 error: Some(other.to_string()),
415 history: working_history,
416 usage: usage_sum,
417 },
354418 },
355 PromptError::PromptCancelled {
356 chat_history,
357 reason,
358 } => Terminal {
359 state: AgentRunState::Cancelled,
360 stop_reason: StopReason::Cancelled,
361 final_text: None,
362 error: Some(reason),
363 history: chat_history,
364 usage: usage_sum,
365 },
366 other => Terminal {
419 StreamingError::Completion(ce) => Terminal {
367420 state: AgentRunState::Failure,
368421 stop_reason: StopReason::Error,
369422 final_text: None,
370 error: Some(other.to_string()),
371 history: ctx.history.clone(),
423 error: Some(ce.to_string()),
424 history: working_history,
372425 usage: usage_sum,
373426 },
374 },
375 StreamingError::Completion(ce) => Terminal {
427 }
428 } else if let Some(resp) = final_response {
429 Terminal {
430 state: AgentRunState::Success,
431 stop_reason: StopReason::Completed,
432 final_text: Some(resp.output.clone()),
433 error: None,
434 history: working_history
435 .iter()
436 .cloned()
437 .chain(resp.messages.clone().unwrap_or_default())
438 .collect(),
439 usage: resp.usage,
440 }
441 } else {
442 Terminal {
376443 state: AgentRunState::Failure,
377444 stop_reason: StopReason::Error,
378445 final_text: None,
379 error: Some(ce.to_string()),
380 history: ctx.history.clone(),
446 error: Some("stream ended without a final response".to_string()),
447 history: working_history,
381448 usage: usage_sum,
382 },
383 }
384 } else if let Some(resp) = final_response {
385 Terminal {
386 state: AgentRunState::Success,
387 stop_reason: StopReason::Completed,
388 final_text: Some(resp.output.clone()),
389 error: None,
390 history: ctx
391 .history
392 .iter()
393 .cloned()
394 .chain(resp.messages.clone().unwrap_or_default())
395 .collect(),
396 usage: resp.usage,
397 }
398 } else {
399 Terminal {
400 state: AgentRunState::Failure,
401 stop_reason: StopReason::Error,
402 final_text: None,
403 error: Some("stream ended without a final response".to_string()),
404 history: ctx.history.clone(),
405 usage: usage_sum,
406 }
449 }
450 };
407451 };
408452
409453 if !self.config.limits.artifact_dir.is_empty() {
410454 sweep_artifacts(&tool_env).await;
411455 }
412456
457 let mut out_history = terminal.history;
458 if compact::estimate_tokens(&out_history) > compact::trigger_threshold(context_max_tokens) {
459 let compaction =
460 compact::compact_history(&summarizer, out_history, context_max_tokens / 2, keep_recent)
461 .await;
462 emit_compacted(&emitter, &compaction).await;
463 out_history = compaction.history;
464 }
465
413466 let usage_json =
414467 serde_json::to_value(terminal.usage).unwrap_or_else(|_| serde_json::json!({}));
415468 if let Err(e) = persist::finish_run(
@@ -421,7 +474,7 @@
421474 final_text: terminal.final_text.clone(),
422475 usage: usage_json.clone(),
423476 error: terminal.error.clone(),
424 history: terminal.history.clone(),
477 history: out_history.clone(),
425478 },
426479 )
427480 .await
@@ -461,35 +514,121 @@
461514 stop_reason: terminal.stop_reason,
462515 final_text: terminal.final_text,
463516 usage: usage_json,
464 history: terminal.history,
517 history: out_history,
465518 sandbox_id: Some(sandbox_id),
466519 artifacts,
467520 })
468521 }
469522}
470523
471struct CancelHook {
524#[derive(Clone)]
525struct RunHook {
472526 db: Arc<AppDatabase>,
473527 run_id: Uuid,
528 summarizer: Summarizer,
529 context_max_tokens: usize,
530 compact_keep_recent: usize,
531 emitter: EventEmitter,
532 cache: Arc<Mutex<Option<CompactCache>>>,
474533}
475534
476impl AgentHook for CancelHook {
535#[derive(Clone)]
536struct CompactCache {
537 prefix_len: usize,
538 summary: String,
539}
540
541impl RunHook {
542 fn cached_patch(&self, history: &[Message], threshold: usize) -> Option<Vec<Message>> {
543 let cache = self.cache.lock().ok()?.clone()?;
544 if cache.prefix_len >= history.len() {
545 return None;
546 }
547 let tail = &history[cache.prefix_len..];
548 if tail.is_empty() || compact::starts_with_tool_result(&tail[0]) {
549 return None;
550 }
551 let summary = compact::summary_message(&cache.summary);
552 let estimate = compact::estimate_message_tokens(&summary) + compact::estimate_tokens(tail);
553 if estimate > threshold {
554 return None;
555 }
556 let mut patched = Vec::with_capacity(tail.len() + 1);
557 patched.push(summary);
558 patched.extend_from_slice(tail);
559 Some(patched)
560 }
561}
562
563impl AgentHook for RunHook {
477564 async fn on_completion_call(
478565 &self,
479566 _ctx: &HookContext,
480 _event: CompletionCallEvent<'_>,
567 event: CompletionCallEvent<'_>,
481568 ) -> CompletionCallAction {
482569 match persist::cancel_requested(&self.db, self.run_id).await {
483 Ok(true) => CompletionCallAction::stop("cancel requested"),
484 Ok(false) => CompletionCallAction::continue_run(),
570 Ok(true) => return CompletionCallAction::stop("cancel requested"),
571 Ok(false) => {}
485572 Err(e) => {
486573 tracing::warn!(error = %e, "agent run cancel check failed");
487 CompletionCallAction::continue_run()
488574 }
489575 }
576
577 let threshold = compact::trigger_threshold(self.context_max_tokens);
578 let history = event.history;
579 if compact::estimate_tokens(history) <= threshold {
580 return CompletionCallAction::continue_run();
581 }
582 if let Some(patched) = self.cached_patch(history, threshold) {
583 return CompletionCallAction::patch(RequestPatch::new().history(patched));
584 }
585
586 let compaction = compact::compact_history(
587 &self.summarizer,
588 history.to_vec(),
589 self.context_max_tokens / 2,
590 self.compact_keep_recent,
591 )
592 .await;
593 if let (Some(summary), Ok(mut cache)) = (&compaction.summary, self.cache.lock()) {
594 cache.replace(CompactCache {
595 prefix_len: compaction.prefix_len,
596 summary: summary.clone(),
597 });
598 }
599 emit_compacted(&self.emitter, &compaction).await;
600 CompletionCallAction::patch(RequestPatch::new().history(compaction.history))
490601 }
491602}
492603
604fn is_context_overflow(error: &StreamingError) -> bool {
605 let completion_error = match error {
606 StreamingError::Completion(ce) => Some(ce),
607 StreamingError::Prompt(pe) => match pe.as_ref() {
608 PromptError::CompletionError(ce) => Some(ce),
609 _ => None,
610 },
611 };
612 let Some(ce) = completion_error else {
613 return false;
614 };
615 let body = ce.provider_response_body().unwrap_or_default().to_lowercase();
616 body.contains("context_length_exceeded")
617 || body.contains("maximum context length")
618 || body.contains("input is too long")
619}
620
621async fn emit_compacted(emitter: &EventEmitter, compaction: &compact::Compaction) {
622 emitter
623 .emit(RunEvent::HistoryCompacted {
624 before: compaction.before,
625 after: compaction.after,
626 estimated_tokens: compaction.estimated_tokens,
627 method: compaction.method,
628 })
629 .await;
630}
631
493632async fn terminate_failed(
494633 db: &Arc<AppDatabase>,
495634 emitter: &EventEmitter,
Modified lib/agent/src/runtime/events.rs View lib/agent/src/runtime/events.rs
@@ -10,6 +10,13 @@
1010 Error,
1111}
1212
13#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15pub enum CompactMethod {
16 Summary,
17 Truncate,
18}
19
1320#[derive(Clone, Debug, Serialize, Deserialize)]
1421#[serde(tag = "type", rename_all = "snake_case")]
1522pub enum RunEvent {
@@ -45,6 +52,12 @@
4552 output_tokens: u64,
4653 total_tokens: u64,
4754 },
55 HistoryCompacted {
56 before: usize,
57 after: usize,
58 estimated_tokens: usize,
59 method: CompactMethod,
60 },
4861 RunFinished {
4962 stop_reason: StopReason,
5063 final_text: Option<String>,
Modified lib/agent/src/runtime/mod.rs View lib/agent/src/runtime/mod.rs
@@ -1,3 +1,4 @@
1pub mod compact;
12pub mod config;
23pub mod engine;
34pub mod events;
@@ -7,8 +8,9 @@
78pub mod sandbox;
89pub mod sink;
910
11pub use compact::{Compaction, Summarizer};
1012pub use config::RuntimeConfig;
1113pub use engine::{AgentRuntime, RunHandle, RunOutcome};
12pub use events::{RunEvent, StopReason};
14pub use events::{CompactMethod, RunEvent, StopReason};
1315pub use reaper::{SweepReport, sweep};
1416pub use sink::{EventEmitter, EventSink, StreamEntry};
Modified lib/config/src/agent.rs View lib/config/src/agent.rs
@@ -39,6 +39,8 @@
3939 pub tool_output_max_bytes: usize,
4040 pub artifact_max_size: i64,
4141 pub artifact_dir: String,
42 pub context_max_tokens: usize,
43 pub compact_keep_recent: usize,
4244}
4345
4446impl Default for AgentRuntimeConfig {
@@ -50,6 +52,8 @@
5052 tool_output_max_bytes: 32_768,
5153 artifact_max_size: 104_857_600,
5254 artifact_dir: String::new(),
55 context_max_tokens: 128_000,
56 compact_keep_recent: 16,
5357 }
5458 }
5559}
@@ -107,6 +111,12 @@
107111 if self.runtime.artifact_max_size <= 0 {
108112 return Err("agent runtime artifact_max_size must be positive".to_string());
109113 }
114 if self.runtime.context_max_tokens == 0 {
115 return Err("agent runtime context_max_tokens must be positive".to_string());
116 }
117 if self.runtime.compact_keep_recent == 0 {
118 return Err("agent runtime compact_keep_recent must be positive".to_string());
119 }
110120 Ok(())
111121 }
112122}