1use std::{collections::HashMap, path::PathBuf, sync::Arc};
4
5use async_trait::async_trait;
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use thiserror::Error;
9use tokio_util::sync::CancellationToken;
10
11#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
12#[serde(tag = "role", rename_all = "snake_case")]
13pub enum Message {
14 User {
15 content: String,
16 },
17 Assistant {
18 content: String,
19 #[serde(default, skip_serializing_if = "Vec::is_empty")]
20 tool_calls: Vec<ToolCall>,
21 },
22 Tool {
23 call_id: String,
24 name: String,
25 content: String,
26 is_error: bool,
27 },
28 HistoryNote {
29 content: String,
30 },
31}
32
33impl Message {
34 fn estimated_tokens(&self, bytes_per_token: usize) -> usize {
35 let bytes = serde_json::to_vec(self).map_or(0, |value| value.len());
36 bytes.div_ceil(bytes_per_token).saturating_add(4)
37 }
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
41pub struct ToolCall {
42 pub id: String,
43 pub name: String,
44 pub arguments: Value,
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
48pub struct ToolSpec {
49 pub name: String,
50 pub description: String,
51 pub parameters: Value,
52}
53
54#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
55#[serde(rename_all = "snake_case")]
56pub enum ToolRisk {
57 ReadOnly,
58 Filesystem,
59 Process,
60 Delegate,
61}
62
63impl ToolRisk {
64 pub fn as_str(self) -> &'static str {
65 match self {
66 Self::ReadOnly => "read_only",
67 Self::Filesystem => "filesystem",
68 Self::Process => "process",
69 Self::Delegate => "delegate",
70 }
71 }
72}
73
74#[derive(Debug, Clone)]
75pub struct ToolContext {
76 pub workspace: PathBuf,
77 pub cancellation: CancellationToken,
78}
79
80#[derive(Debug, Clone, PartialEq, Eq)]
81pub struct ToolOutput {
82 pub content: String,
83 pub is_error: bool,
84 pub truncated: bool,
85}
86
87impl ToolOutput {
88 pub fn success(content: impl Into<String>) -> Self {
89 Self {
90 content: content.into(),
91 is_error: false,
92 truncated: false,
93 }
94 }
95
96 pub fn failure(content: impl Into<String>) -> Self {
97 Self {
98 content: content.into(),
99 is_error: true,
100 truncated: false,
101 }
102 }
103}
104
105#[derive(Debug, Error)]
106#[error("{0}")]
107pub struct ToolError(pub String);
108
109#[async_trait]
110pub trait Tool: Send + Sync {
111 fn spec(&self) -> ToolSpec;
112 fn risk(&self, arguments: &Value) -> Result<ToolRisk, ToolError>;
113 fn approval_summary(&self, arguments: &Value) -> Result<String, ToolError>;
114 async fn execute(
115 &self,
116 arguments: Value,
117 context: ToolContext,
118 ) -> Result<ToolOutput, ToolError>;
119}
120
121#[derive(Default)]
122pub struct ToolRegistry {
123 tools: HashMap<String, Arc<dyn Tool>>,
124}
125
126impl ToolRegistry {
127 pub fn register(&mut self, tool: Arc<dyn Tool>) -> Result<(), ToolError> {
128 let name = tool.spec().name;
129 if self.tools.contains_key(&name) {
130 return Err(ToolError(format!("duplicate tool name: {name}")));
131 }
132 self.tools.insert(name, tool);
133 Ok(())
134 }
135
136 pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
137 self.tools.get(name).cloned()
138 }
139
140 pub fn specs(&self) -> Vec<ToolSpec> {
141 let mut specs: Vec<_> = self.tools.values().map(|tool| tool.spec()).collect();
142 specs.sort_by(|a, b| a.name.cmp(&b.name));
143 specs
144 }
145}
146
147#[derive(Debug, Clone, Default, PartialEq, Eq)]
148pub struct Usage {
149 pub input_tokens: Option<u64>,
150 pub output_tokens: Option<u64>,
151}
152
153impl Usage {
154 fn add(&mut self, other: &Self) {
155 self.input_tokens = add_optional(self.input_tokens, other.input_tokens);
156 self.output_tokens = add_optional(self.output_tokens, other.output_tokens);
157 }
158}
159
160fn add_optional(left: Option<u64>, right: Option<u64>) -> Option<u64> {
161 match (left, right) {
162 (None, None) => None,
163 (left, right) => Some(left.unwrap_or(0).saturating_add(right.unwrap_or(0))),
164 }
165}
166
167#[derive(Debug, Clone)]
168pub struct ProviderRequest {
169 pub system_prompt: String,
170 pub messages: Vec<Message>,
171 pub tools: Vec<ToolSpec>,
172}
173
174#[derive(Debug, Clone)]
175pub struct AssistantResponse {
176 pub content: String,
177 pub tool_calls: Vec<ToolCall>,
178 pub usage: Usage,
179}
180
181#[derive(Debug, Clone, Copy, PartialEq, Eq)]
182pub enum ProviderErrorKind {
183 Provider,
184 ResponseLimit,
185 ToolLimit,
186 Cancelled,
187}
188
189#[derive(Debug, Error)]
190#[error("{message}")]
191pub struct ProviderError {
192 pub kind: ProviderErrorKind,
193 pub message: String,
194}
195
196impl ProviderError {
197 pub fn new(kind: ProviderErrorKind, message: impl Into<String>) -> Self {
198 Self {
199 kind,
200 message: message.into(),
201 }
202 }
203}
204
205#[async_trait]
206pub trait TextDeltaSink: Send + Sync {
207 async fn push(&self, delta: &str) -> Result<(), ProviderError>;
208}
209
210#[async_trait]
211pub trait Provider: Send + Sync {
212 fn model(&self) -> &str;
213
214 async fn complete(
215 &self,
216 request: ProviderRequest,
217 deltas: Arc<dyn TextDeltaSink>,
218 cancellation: CancellationToken,
219 ) -> Result<AssistantResponse, ProviderError>;
220}
221
222#[derive(Debug, Clone)]
223pub enum CoreEvent {
224 AssistantDelta {
225 content: String,
226 },
227 AssistantCompleted {
228 content: String,
229 },
230 ToolProposed {
231 call_id: String,
232 name: String,
233 arguments: Value,
234 },
235 ToolStarted {
236 call_id: String,
237 name: String,
238 },
239 ToolCompleted {
240 call_id: String,
241 name: String,
242 output: ToolOutput,
243 },
244 ContextCompacted {
245 before_tokens: usize,
246 after_tokens: usize,
247 removed_messages: usize,
248 },
249 SessionTrimmed {
250 removed_messages: usize,
251 history_bytes: usize,
252 },
253}
254
255#[async_trait]
256pub trait EventSink: Send + Sync {
257 async fn emit(&self, event: CoreEvent) -> Result<(), AgentError>;
258}
259
260#[derive(Debug, Clone)]
261pub struct ApprovalRequest {
262 pub call_id: String,
263 pub name: String,
264 pub risk: ToolRisk,
265 pub cwd: PathBuf,
266 pub summary: String,
267}
268
269#[async_trait]
270pub trait ApprovalGate: Send + Sync {
271 async fn approve(
272 &self,
273 request: ApprovalRequest,
274 cancellation: CancellationToken,
275 ) -> Result<bool, AgentError>;
276}
277
278#[derive(Debug, Clone)]
279pub struct ContextConfig {
280 pub max_tokens: usize,
281 pub reserve_output_tokens: usize,
282 pub safety_margin_tokens: usize,
283 pub bytes_per_token: usize,
284 pub summary_max_chars: usize,
285}
286
287impl Default for ContextConfig {
288 fn default() -> Self {
289 Self {
290 max_tokens: 128_000,
291 reserve_output_tokens: 8_192,
292 safety_margin_tokens: 2_048,
293 bytes_per_token: 3,
294 summary_max_chars: 6_000,
295 }
296 }
297}
298
299#[derive(Debug, Clone)]
300pub struct ContextSelection {
301 pub messages: Vec<Message>,
302 pub before_tokens: usize,
303 pub after_tokens: usize,
304 pub removed_messages: usize,
305}
306
307#[derive(Debug, Error)]
308#[error("{0}")]
309pub struct ContextError(pub String);
310
311pub trait ContextPolicy: Send + Sync {
312 fn select(
313 &self,
314 history: &[Message],
315 system_prompt: &str,
316 tools: &[ToolSpec],
317 ) -> Result<ContextSelection, ContextError>;
318}
319
320pub struct BudgetContextPolicy {
321 config: ContextConfig,
322}
323
324impl BudgetContextPolicy {
325 pub fn new(config: ContextConfig) -> Result<Self, ContextError> {
326 if config.bytes_per_token == 0 {
327 return Err(ContextError(
328 "context.bytes_per_token must be positive".into(),
329 ));
330 }
331 if config
332 .reserve_output_tokens
333 .saturating_add(config.safety_margin_tokens)
334 >= config.max_tokens
335 {
336 return Err(ContextError(
337 "context reserve and safety margin consume the model window".into(),
338 ));
339 }
340 Ok(Self { config })
341 }
342
343 fn string_tokens(&self, value: &str) -> usize {
344 value.len().div_ceil(self.config.bytes_per_token)
345 }
346
347 fn group_messages(history: &[Message]) -> Vec<Vec<Message>> {
348 let mut groups: Vec<Vec<Message>> = Vec::new();
349 for message in history {
350 if matches!(message, Message::User { .. }) || groups.is_empty() {
351 groups.push(Vec::new());
352 }
353 groups
354 .last_mut()
355 .expect("a group was just created")
356 .push(message.clone());
357 }
358 groups
359 }
360
361 fn summarize(&self, messages: &[Message]) -> String {
362 let mut output = format!(
363 "[SCV compacted {} earlier messages. Bounded extracts follow.]\n",
364 messages.len()
365 );
366 for message in messages {
367 let (label, content) = match message {
368 Message::User { content } => ("user", content.as_str()),
369 Message::Assistant { content, .. } => ("assistant", content.as_str()),
370 Message::Tool {
371 name,
372 content,
373 is_error,
374 ..
375 } => {
376 let status = if *is_error { "failed" } else { "ok" };
377 output.push_str(&format!("tool {name} ({status}): "));
378 ("", content.as_str())
379 }
380 Message::HistoryNote { content } => ("earlier", content.as_str()),
381 };
382 if !label.is_empty() {
383 output.push_str(label);
384 output.push_str(": ");
385 }
386 let tail = char_tail(content, 240);
387 output.push_str(&tail.replace('\n', " "));
388 output.push('\n');
389 if output.chars().count() >= self.config.summary_max_chars {
390 break;
391 }
392 }
393 truncate_chars(&output, self.config.summary_max_chars)
394 }
395}
396
397impl ContextPolicy for BudgetContextPolicy {
398 fn select(
399 &self,
400 history: &[Message],
401 system_prompt: &str,
402 tools: &[ToolSpec],
403 ) -> Result<ContextSelection, ContextError> {
404 if history.is_empty() {
405 return Ok(ContextSelection {
406 messages: Vec::new(),
407 before_tokens: 0,
408 after_tokens: 0,
409 removed_messages: 0,
410 });
411 }
412 let tools_bytes = serde_json::to_vec(tools).map_or(0, |value| value.len());
413 let static_tokens = self
414 .string_tokens(system_prompt)
415 .saturating_add(tools_bytes.div_ceil(self.config.bytes_per_token))
416 .saturating_add(self.config.reserve_output_tokens)
417 .saturating_add(self.config.safety_margin_tokens);
418 if static_tokens >= self.config.max_tokens {
419 return Err(ContextError(
420 "system prompt and tool schemas exceed context budget".into(),
421 ));
422 }
423 let budget = self.config.max_tokens - static_tokens;
424 let groups = Self::group_messages(history);
425 let newest = groups.last().expect("history produced at least one group");
426 let newest_cost: usize = newest
427 .iter()
428 .map(|message| message.estimated_tokens(self.config.bytes_per_token))
429 .sum();
430 if newest_cost > budget {
431 return Err(ContextError("newest turn exceeds context budget".into()));
432 }
433
434 let before_history_tokens: usize = history
435 .iter()
436 .map(|message| message.estimated_tokens(self.config.bytes_per_token))
437 .sum();
438 let mut selected_groups: Vec<Vec<Message>> = vec![newest.clone()];
439 let mut selected_cost = newest_cost;
440 for group in groups[..groups.len() - 1].iter().rev() {
441 let cost: usize = group
442 .iter()
443 .map(|message| message.estimated_tokens(self.config.bytes_per_token))
444 .sum();
445 if selected_cost.saturating_add(cost) <= budget {
446 selected_groups.insert(0, group.clone());
447 selected_cost += cost;
448 } else {
449 break;
450 }
451 }
452
453 let mut removed_messages = groups[..groups.len() - selected_groups.len()]
454 .iter()
455 .map(Vec::len)
456 .sum::<usize>();
457 if removed_messages > 0 {
458 loop {
459 let note = Message::HistoryNote {
460 content: self.summarize(&history[..removed_messages]),
461 };
462 let note_cost = note.estimated_tokens(self.config.bytes_per_token);
463 if selected_cost.saturating_add(note_cost) <= budget {
464 let mut selected: Vec<Message> =
465 selected_groups.into_iter().flatten().collect();
466 selected.insert(0, note);
467 selected_cost += note_cost;
468 return Ok(ContextSelection {
469 messages: selected,
470 before_tokens: static_tokens.saturating_add(before_history_tokens),
471 after_tokens: static_tokens.saturating_add(selected_cost),
472 removed_messages,
473 });
474 }
475 if selected_groups.len() == 1 {
476 let available_tokens = budget.saturating_sub(selected_cost);
477 let content = match note {
478 Message::HistoryNote { content } => content,
479 _ => unreachable!(),
480 };
481 let Some(note) =
482 fit_history_note(&content, available_tokens, self.config.bytes_per_token)
483 else {
484 return Err(ContextError(
485 "compaction note cannot fit context budget".into(),
486 ));
487 };
488 let note_cost = note.estimated_tokens(self.config.bytes_per_token);
489 let mut selected: Vec<Message> =
490 selected_groups.into_iter().flatten().collect();
491 selected.insert(0, note);
492 selected_cost += note_cost;
493 return Ok(ContextSelection {
494 messages: selected,
495 before_tokens: static_tokens.saturating_add(before_history_tokens),
496 after_tokens: static_tokens.saturating_add(selected_cost),
497 removed_messages,
498 });
499 }
500 let removed_group = selected_groups.remove(0);
501 let removed_cost: usize = removed_group
502 .iter()
503 .map(|message| message.estimated_tokens(self.config.bytes_per_token))
504 .sum();
505 selected_cost = selected_cost.saturating_sub(removed_cost);
506 removed_messages += removed_group.len();
507 }
508 }
509
510 let selected: Vec<Message> = selected_groups.into_iter().flatten().collect();
511 Ok(ContextSelection {
512 messages: selected,
513 before_tokens: static_tokens.saturating_add(before_history_tokens),
514 after_tokens: static_tokens.saturating_add(selected_cost),
515 removed_messages,
516 })
517 }
518}
519
520#[derive(Debug, Clone)]
521pub struct HistoryLimits {
522 pub max_bytes: usize,
523 pub max_messages: usize,
524 pub note_max_chars: usize,
525}
526
527impl Default for HistoryLimits {
528 fn default() -> Self {
529 Self {
530 max_bytes: 16 * 1024 * 1024,
531 max_messages: 10_000,
532 note_max_chars: 4_000,
533 }
534 }
535}
536
537#[derive(Debug, Clone)]
538pub struct AgentConfig {
539 pub system_prompt: String,
540 pub max_steps: usize,
541 pub history_limits: HistoryLimits,
542}
543
544#[derive(Debug, Clone)]
545pub struct TurnOutcome {
546 pub steps: usize,
547 pub usage: Usage,
548}
549
550#[derive(Debug, Error)]
551pub enum AgentError {
552 #[error("turn cancelled")]
553 Cancelled,
554 #[error("{0}")]
555 Provider(String),
556 #[error("{0}")]
557 ContextLimit(String),
558 #[error("agent reached its maximum step count")]
559 StepLimit,
560 #[error("{0}")]
561 HistoryLimit(String),
562 #[error("{0}")]
563 ResponseLimit(String),
564 #[error("{0}")]
565 ToolLimit(String),
566 #[error("{0}")]
567 Internal(String),
568}
569
570impl AgentError {
571 pub fn code(&self) -> &'static str {
572 match self {
573 Self::Cancelled => "cancelled",
574 Self::Provider(_) => "provider_error",
575 Self::ContextLimit(_) => "context_limit",
576 Self::StepLimit => "step_limit",
577 Self::HistoryLimit(_) => "history_limit",
578 Self::ResponseLimit(_) => "response_limit",
579 Self::ToolLimit(_) => "tool_limit",
580 Self::Internal(_) => "internal_error",
581 }
582 }
583}
584
585pub struct AgentRuntime {
586 provider: Arc<dyn Provider>,
587 tools: Arc<ToolRegistry>,
588 context: Arc<dyn ContextPolicy>,
589 config: AgentConfig,
590 workspace: PathBuf,
591}
592
593impl AgentRuntime {
594 pub fn new(
595 provider: Arc<dyn Provider>,
596 tools: Arc<ToolRegistry>,
597 context: Arc<dyn ContextPolicy>,
598 config: AgentConfig,
599 workspace: PathBuf,
600 ) -> Self {
601 Self {
602 provider,
603 tools,
604 context,
605 config,
606 workspace,
607 }
608 }
609
610 pub fn model(&self) -> &str {
611 self.provider.model()
612 }
613
614 pub async fn run_turn(
615 &self,
616 history: &mut Vec<Message>,
617 prompt: String,
618 sink: Arc<dyn EventSink>,
619 approvals: Arc<dyn ApprovalGate>,
620 cancellation: CancellationToken,
621 ) -> Result<TurnOutcome, AgentError> {
622 let checkpoint = history.clone();
623 let result = self
624 .run_turn_inner(history, prompt, sink, approvals, cancellation)
625 .await;
626 if result.is_err() {
627 *history = checkpoint;
628 }
629 result
630 }
631
632 async fn run_turn_inner(
633 &self,
634 history: &mut Vec<Message>,
635 prompt: String,
636 sink: Arc<dyn EventSink>,
637 approvals: Arc<dyn ApprovalGate>,
638 cancellation: CancellationToken,
639 ) -> Result<TurnOutcome, AgentError> {
640 if cancellation.is_cancelled() {
641 return Err(AgentError::Cancelled);
642 }
643 history.push(Message::User { content: prompt });
644 self.enforce_history_limits(history, sink.as_ref()).await?;
645 let specs = self.tools.specs();
646 let mut usage = Usage::default();
647
648 for step in 1..=self.config.max_steps {
649 if cancellation.is_cancelled() {
650 return Err(AgentError::Cancelled);
651 }
652 let selection = self
653 .context
654 .select(history, &self.config.system_prompt, &specs)
655 .map_err(|error| AgentError::ContextLimit(error.to_string()))?;
656 if selection.removed_messages > 0 {
657 sink.emit(CoreEvent::ContextCompacted {
658 before_tokens: selection.before_tokens,
659 after_tokens: selection.after_tokens,
660 removed_messages: selection.removed_messages,
661 })
662 .await?;
663 }
664 let delta_sink: Arc<dyn TextDeltaSink> = Arc::new(ForwardDeltas {
665 sink: Arc::clone(&sink),
666 });
667 let response = self
668 .provider
669 .complete(
670 ProviderRequest {
671 system_prompt: self.config.system_prompt.clone(),
672 messages: selection.messages,
673 tools: specs.clone(),
674 },
675 delta_sink,
676 cancellation.child_token(),
677 )
678 .await
679 .map_err(map_provider_error)?;
680 usage.add(&response.usage);
681 sink.emit(CoreEvent::AssistantCompleted {
682 content: response.content.clone(),
683 })
684 .await?;
685 let calls = response.tool_calls.clone();
686 history.push(Message::Assistant {
687 content: response.content,
688 tool_calls: response.tool_calls,
689 });
690 self.enforce_history_limits(history, sink.as_ref()).await?;
691 if calls.is_empty() {
692 return Ok(TurnOutcome { steps: step, usage });
693 }
694
695 for call in calls {
696 if cancellation.is_cancelled() {
697 return Err(AgentError::Cancelled);
698 }
699 sink.emit(CoreEvent::ToolProposed {
700 call_id: call.id.clone(),
701 name: call.name.clone(),
702 arguments: call.arguments.clone(),
703 })
704 .await?;
705 let Some(tool) = self.tools.get(&call.name) else {
706 let output = ToolOutput::failure(format!("unknown tool: {}", call.name));
707 sink.emit(CoreEvent::ToolCompleted {
708 call_id: call.id.clone(),
709 name: call.name.clone(),
710 output: output.clone(),
711 })
712 .await?;
713 history.push(Message::Tool {
714 call_id: call.id,
715 name: call.name,
716 content: output.content,
717 is_error: true,
718 });
719 self.enforce_history_limits(history, sink.as_ref()).await?;
720 continue;
721 };
722 let risk = match tool.risk(&call.arguments) {
723 Ok(risk) => risk,
724 Err(error) => {
725 self.record_tool_error(history, sink.as_ref(), &call, error.to_string())
726 .await?;
727 self.enforce_history_limits(history, sink.as_ref()).await?;
728 continue;
729 }
730 };
731 let summary = match tool.approval_summary(&call.arguments) {
732 Ok(summary) => summary,
733 Err(error) => {
734 self.record_tool_error(history, sink.as_ref(), &call, error.to_string())
735 .await?;
736 self.enforce_history_limits(history, sink.as_ref()).await?;
737 continue;
738 }
739 };
740 let approved = approvals
741 .approve(
742 ApprovalRequest {
743 call_id: call.id.clone(),
744 name: call.name.clone(),
745 risk,
746 cwd: self.workspace.clone(),
747 summary,
748 },
749 cancellation.child_token(),
750 )
751 .await?;
752 let output = if approved {
753 sink.emit(CoreEvent::ToolStarted {
754 call_id: call.id.clone(),
755 name: call.name.clone(),
756 })
757 .await?;
758 tool.execute(
759 call.arguments.clone(),
760 ToolContext {
761 workspace: self.workspace.clone(),
762 cancellation: cancellation.child_token(),
763 },
764 )
765 .await
766 .unwrap_or_else(|error| ToolOutput::failure(error.to_string()))
767 } else {
768 ToolOutput::failure("tool call denied by policy or user")
769 };
770 sink.emit(CoreEvent::ToolCompleted {
771 call_id: call.id.clone(),
772 name: call.name.clone(),
773 output: output.clone(),
774 })
775 .await?;
776 history.push(Message::Tool {
777 call_id: call.id,
778 name: call.name,
779 content: output.content,
780 is_error: output.is_error,
781 });
782 self.enforce_history_limits(history, sink.as_ref()).await?;
783 }
784 }
785 Err(AgentError::StepLimit)
786 }
787
788 async fn record_tool_error(
789 &self,
790 history: &mut Vec<Message>,
791 sink: &dyn EventSink,
792 call: &ToolCall,
793 message: String,
794 ) -> Result<(), AgentError> {
795 let output = ToolOutput::failure(message);
796 sink.emit(CoreEvent::ToolCompleted {
797 call_id: call.id.clone(),
798 name: call.name.clone(),
799 output: output.clone(),
800 })
801 .await?;
802 history.push(Message::Tool {
803 call_id: call.id.clone(),
804 name: call.name.clone(),
805 content: output.content,
806 is_error: true,
807 });
808 Ok(())
809 }
810
811 async fn enforce_history_limits(
812 &self,
813 history: &mut Vec<Message>,
814 sink: &dyn EventSink,
815 ) -> Result<(), AgentError> {
816 let limits = &self.config.history_limits;
817 let mut total_removed = 0;
818 while history.len() > limits.max_messages || history_bytes(history) > limits.max_bytes {
819 let latest_user = history
820 .iter()
821 .rposition(|message| matches!(message, Message::User { .. }))
822 .unwrap_or(0);
823 let active = &history[latest_user..];
824 if active.len() > limits.max_messages || history_bytes(active) > limits.max_bytes {
825 return Err(AgentError::HistoryLimit(
826 "active turn exceeds configured session history limit".into(),
827 ));
828 }
829 let first_user = history
830 .iter()
831 .position(|message| matches!(message, Message::User { .. }))
832 .unwrap_or(latest_user);
833 if first_user == latest_user {
834 if matches!(history.first(), Some(Message::HistoryNote { .. })) {
835 history.remove(0);
836 total_removed += 1;
837 continue;
838 }
839 return Err(AgentError::HistoryLimit(
840 "session history cannot be reduced within its configured limit".into(),
841 ));
842 }
843 let end = history[first_user + 1..]
844 .iter()
845 .position(|message| matches!(message, Message::User { .. }))
846 .map(|index| first_user + 1 + index)
847 .ok_or_else(|| {
848 AgentError::HistoryLimit(
849 "session history has no complete group available to trim".into(),
850 )
851 })?;
852 let removed: Vec<Message> = history.drain(..end).collect();
853 total_removed += removed.len();
854 let note = Message::HistoryNote {
855 content: summarize_history_trim(&removed, total_removed, limits.note_max_chars),
856 };
857 if matches!(history.first(), Some(Message::HistoryNote { .. })) {
858 history.remove(0);
859 }
860 history.insert(0, note);
861 }
862 if total_removed > 0 {
863 sink.emit(CoreEvent::SessionTrimmed {
864 removed_messages: total_removed,
865 history_bytes: history_bytes(history),
866 })
867 .await?;
868 }
869 Ok(())
870 }
871}
872
873fn map_provider_error(error: ProviderError) -> AgentError {
874 match error.kind {
875 ProviderErrorKind::Provider => AgentError::Provider(error.message),
876 ProviderErrorKind::ResponseLimit => AgentError::ResponseLimit(error.message),
877 ProviderErrorKind::ToolLimit => AgentError::ToolLimit(error.message),
878 ProviderErrorKind::Cancelled => AgentError::Cancelled,
879 }
880}
881
882struct ForwardDeltas {
883 sink: Arc<dyn EventSink>,
884}
885
886#[async_trait]
887impl TextDeltaSink for ForwardDeltas {
888 async fn push(&self, delta: &str) -> Result<(), ProviderError> {
889 self.sink
890 .emit(CoreEvent::AssistantDelta {
891 content: delta.to_owned(),
892 })
893 .await
894 .map_err(|error| match error {
895 AgentError::Cancelled => {
896 ProviderError::new(ProviderErrorKind::Cancelled, "turn cancelled")
897 }
898 AgentError::ResponseLimit(message) => {
899 ProviderError::new(ProviderErrorKind::ResponseLimit, message)
900 }
901 AgentError::ToolLimit(message) => {
902 ProviderError::new(ProviderErrorKind::ToolLimit, message)
903 }
904 error => ProviderError::new(ProviderErrorKind::Provider, error.to_string()),
905 })
906 }
907}
908
909fn history_bytes(history: &[Message]) -> usize {
910 serde_json::to_vec(history).map_or(usize::MAX, |value| value.len())
911}
912
913fn summarize_history_trim(messages: &[Message], removed: usize, max_chars: usize) -> String {
914 let mut note =
915 format!("[SCV trimmed {removed} earlier canonical messages to enforce session limits.]\n");
916 for message in messages {
917 let (label, content) = match message {
918 Message::User { content } => ("user", content.as_str()),
919 Message::Assistant { content, .. } => ("assistant", content.as_str()),
920 Message::Tool {
921 name,
922 content,
923 is_error,
924 ..
925 } => {
926 let status = if *is_error { "failed" } else { "ok" };
927 note.push_str(&format!("tool {name} ({status}): "));
928 ("", content.as_str())
929 }
930 Message::HistoryNote { content } => ("earlier", content.as_str()),
931 };
932 if !label.is_empty() {
933 note.push_str(label);
934 note.push_str(": ");
935 }
936 note.push_str(&char_tail(content, 160).replace('\n', " "));
937 note.push('\n');
938 if note.chars().count() >= max_chars {
939 break;
940 }
941 }
942 truncate_chars(¬e, max_chars)
943}
944
945fn truncate_chars(value: &str, max_chars: usize) -> String {
946 value.chars().take(max_chars).collect()
947}
948
949fn fit_history_note(
950 content: &str,
951 available_tokens: usize,
952 bytes_per_token: usize,
953) -> Option<Message> {
954 let chars: Vec<char> = content.chars().collect();
955 let mut low = 0usize;
956 let mut high = chars.len();
957 let mut best = None;
958 while low <= high {
959 let middle = low + (high - low) / 2;
960 let candidate = Message::HistoryNote {
961 content: chars[..middle].iter().collect(),
962 };
963 if candidate.estimated_tokens(bytes_per_token) <= available_tokens {
964 best = Some(candidate);
965 low = middle.saturating_add(1);
966 } else if middle == 0 {
967 break;
968 } else {
969 high = middle - 1;
970 }
971 }
972 best
973}
974
975fn char_tail(value: &str, max_chars: usize) -> String {
976 let count = value.chars().count();
977 value
978 .chars()
979 .skip(count.saturating_sub(max_chars))
980 .collect()
981}
982
983#[cfg(test)]
984mod tests {
985 use std::{collections::VecDeque, sync::Mutex};
986
987 use tokio::sync::Notify;
988
989 use super::*;
990
991 struct ScriptedProvider {
992 responses: Mutex<VecDeque<AssistantResponse>>,
993 }
994
995 #[async_trait]
996 impl Provider for ScriptedProvider {
997 fn model(&self) -> &str {
998 "test-model"
999 }
1000
1001 async fn complete(
1002 &self,
1003 _request: ProviderRequest,
1004 deltas: Arc<dyn TextDeltaSink>,
1005 _cancellation: CancellationToken,
1006 ) -> Result<AssistantResponse, ProviderError> {
1007 let response = self.responses.lock().unwrap().pop_front().unwrap();
1008 deltas.push(&response.content).await?;
1009 Ok(response)
1010 }
1011 }
1012
1013 struct CollectSink(Mutex<Vec<CoreEvent>>);
1014
1015 #[async_trait]
1016 impl EventSink for CollectSink {
1017 async fn emit(&self, event: CoreEvent) -> Result<(), AgentError> {
1018 self.0.lock().unwrap().push(event);
1019 Ok(())
1020 }
1021 }
1022
1023 struct Allow;
1024
1025 #[async_trait]
1026 impl ApprovalGate for Allow {
1027 async fn approve(
1028 &self,
1029 _request: ApprovalRequest,
1030 _cancellation: CancellationToken,
1031 ) -> Result<bool, AgentError> {
1032 Ok(true)
1033 }
1034 }
1035
1036 struct Deny;
1037
1038 #[async_trait]
1039 impl ApprovalGate for Deny {
1040 async fn approve(
1041 &self,
1042 _request: ApprovalRequest,
1043 _cancellation: CancellationToken,
1044 ) -> Result<bool, AgentError> {
1045 Ok(false)
1046 }
1047 }
1048
1049 struct WaitForCancellation {
1050 entered: Arc<Notify>,
1051 }
1052
1053 #[async_trait]
1054 impl ApprovalGate for WaitForCancellation {
1055 async fn approve(
1056 &self,
1057 _request: ApprovalRequest,
1058 cancellation: CancellationToken,
1059 ) -> Result<bool, AgentError> {
1060 self.entered.notify_one();
1061 cancellation.cancelled().await;
1062 Err(AgentError::Cancelled)
1063 }
1064 }
1065
1066 struct EchoTool;
1067
1068 #[async_trait]
1069 impl Tool for EchoTool {
1070 fn spec(&self) -> ToolSpec {
1071 ToolSpec {
1072 name: "echo".into(),
1073 description: "Echo a value".into(),
1074 parameters: serde_json::json!({
1075 "type":"object",
1076 "properties":{"value":{"type":"string"}},
1077 "required":["value"]
1078 }),
1079 }
1080 }
1081
1082 fn risk(&self, arguments: &Value) -> Result<ToolRisk, ToolError> {
1083 arguments
1084 .get("value")
1085 .and_then(Value::as_str)
1086 .ok_or_else(|| ToolError("value must be a string".into()))?;
1087 Ok(ToolRisk::ReadOnly)
1088 }
1089
1090 fn approval_summary(&self, arguments: &Value) -> Result<String, ToolError> {
1091 self.risk(arguments)?;
1092 Ok("Echo a value".into())
1093 }
1094
1095 async fn execute(
1096 &self,
1097 arguments: Value,
1098 _context: ToolContext,
1099 ) -> Result<ToolOutput, ToolError> {
1100 Ok(ToolOutput::success(
1101 arguments["value"].as_str().unwrap_or_default(),
1102 ))
1103 }
1104 }
1105
1106 #[tokio::test]
1107 async fn completes_a_simple_turn() {
1108 let provider = Arc::new(ScriptedProvider {
1109 responses: Mutex::new(VecDeque::from([AssistantResponse {
1110 content: "done".into(),
1111 tool_calls: Vec::new(),
1112 usage: Usage {
1113 input_tokens: Some(3),
1114 output_tokens: Some(1),
1115 },
1116 }])),
1117 });
1118 let runtime = AgentRuntime::new(
1119 provider,
1120 Arc::new(ToolRegistry::default()),
1121 Arc::new(BudgetContextPolicy::new(ContextConfig::default()).unwrap()),
1122 AgentConfig {
1123 system_prompt: "test".into(),
1124 max_steps: 2,
1125 history_limits: HistoryLimits::default(),
1126 },
1127 PathBuf::from("/tmp"),
1128 );
1129 let sink = Arc::new(CollectSink(Mutex::new(Vec::new())));
1130 let mut history = Vec::new();
1131 let outcome = runtime
1132 .run_turn(
1133 &mut history,
1134 "hello".into(),
1135 sink.clone(),
1136 Arc::new(Allow),
1137 CancellationToken::new(),
1138 )
1139 .await
1140 .unwrap();
1141 assert_eq!(outcome.steps, 1);
1142 assert_eq!(history.len(), 2);
1143 assert!(matches!(
1144 sink.0.lock().unwrap().last(),
1145 Some(CoreEvent::AssistantCompleted { .. })
1146 ));
1147 }
1148
1149 #[test]
1150 fn context_keeps_tool_groups_together() {
1151 let policy = BudgetContextPolicy::new(ContextConfig {
1152 max_tokens: 120,
1153 reserve_output_tokens: 10,
1154 safety_margin_tokens: 10,
1155 bytes_per_token: 3,
1156 summary_max_chars: 120,
1157 })
1158 .unwrap();
1159 let history = vec![
1160 Message::User {
1161 content: "old request ".repeat(20),
1162 },
1163 Message::Assistant {
1164 content: String::new(),
1165 tool_calls: vec![ToolCall {
1166 id: "1".into(),
1167 name: "read".into(),
1168 arguments: serde_json::json!({"path":"a"}),
1169 }],
1170 },
1171 Message::Tool {
1172 call_id: "1".into(),
1173 name: "read".into(),
1174 content: "result".into(),
1175 is_error: false,
1176 },
1177 Message::User {
1178 content: "new".into(),
1179 },
1180 ];
1181 let selection = policy.select(&history, "system", &[]).unwrap();
1182 assert!(selection.removed_messages > 0);
1183 assert_eq!(
1184 selection.removed_messages,
1185 history.len() - (selection.messages.len() - 1)
1186 );
1187 assert!(matches!(
1188 selection.messages.last(),
1189 Some(Message::User { .. })
1190 ));
1191 assert!(
1192 !selection
1193 .messages
1194 .iter()
1195 .any(|message| matches!(message, Message::Tool { call_id, .. } if call_id == "1"))
1196 );
1197 }
1198
1199 #[tokio::test]
1200 async fn repeated_history_trimming_rebuilds_the_note_and_makes_progress() {
1201 let runtime = AgentRuntime::new(
1202 Arc::new(ScriptedProvider {
1203 responses: Mutex::new(VecDeque::new()),
1204 }),
1205 Arc::new(ToolRegistry::default()),
1206 Arc::new(BudgetContextPolicy::new(ContextConfig::default()).unwrap()),
1207 AgentConfig {
1208 system_prompt: "test".into(),
1209 max_steps: 1,
1210 history_limits: HistoryLimits {
1211 max_bytes: 4096,
1212 max_messages: 3,
1213 note_max_chars: 80,
1214 },
1215 },
1216 PathBuf::from("/tmp"),
1217 );
1218 let sink = CollectSink(Mutex::new(Vec::new()));
1219 let mut history = vec![
1220 Message::HistoryNote {
1221 content: "previous trim".into(),
1222 },
1223 Message::User {
1224 content: "old request".into(),
1225 },
1226 Message::Assistant {
1227 content: "old answer".into(),
1228 tool_calls: Vec::new(),
1229 },
1230 Message::User {
1231 content: "active request".into(),
1232 },
1233 ];
1234 runtime
1235 .enforce_history_limits(&mut history, &sink)
1236 .await
1237 .unwrap();
1238 assert!(history.len() <= 3);
1239 assert!(matches!(history.first(), Some(Message::HistoryNote { .. })));
1240 assert!(matches!(history.last(), Some(Message::User { .. })));
1241 }
1242
1243 #[tokio::test]
1244 async fn active_turn_over_history_limit_rolls_back() {
1245 let runtime = AgentRuntime::new(
1246 Arc::new(ScriptedProvider {
1247 responses: Mutex::new(VecDeque::new()),
1248 }),
1249 Arc::new(ToolRegistry::default()),
1250 Arc::new(BudgetContextPolicy::new(ContextConfig::default()).unwrap()),
1251 AgentConfig {
1252 system_prompt: "test".into(),
1253 max_steps: 1,
1254 history_limits: HistoryLimits {
1255 max_bytes: 16,
1256 max_messages: 10,
1257 note_max_chars: 8,
1258 },
1259 },
1260 PathBuf::from("/tmp"),
1261 );
1262 let sink = CollectSink(Mutex::new(Vec::new()));
1263 let mut history = Vec::new();
1264 let result = runtime
1265 .run_turn(
1266 &mut history,
1267 "too large for the configured history".into(),
1268 Arc::new(sink),
1269 Arc::new(Allow),
1270 CancellationToken::new(),
1271 )
1272 .await;
1273 assert!(matches!(result, Err(AgentError::HistoryLimit(_))));
1274 assert!(history.is_empty());
1275 }
1276
1277 #[tokio::test]
1278 async fn executes_a_multi_step_tool_loop_and_aggregates_usage() {
1279 let provider = Arc::new(ScriptedProvider {
1280 responses: Mutex::new(VecDeque::from([
1281 AssistantResponse {
1282 content: String::new(),
1283 tool_calls: vec![ToolCall {
1284 id: "call-1".into(),
1285 name: "echo".into(),
1286 arguments: serde_json::json!({"value":"hello"}),
1287 }],
1288 usage: Usage {
1289 input_tokens: Some(2),
1290 output_tokens: Some(1),
1291 },
1292 },
1293 AssistantResponse {
1294 content: "done".into(),
1295 tool_calls: Vec::new(),
1296 usage: Usage {
1297 input_tokens: Some(4),
1298 output_tokens: Some(2),
1299 },
1300 },
1301 ])),
1302 });
1303 let mut registry = ToolRegistry::default();
1304 registry.register(Arc::new(EchoTool)).unwrap();
1305 let runtime = AgentRuntime::new(
1306 provider,
1307 Arc::new(registry),
1308 Arc::new(BudgetContextPolicy::new(ContextConfig::default()).unwrap()),
1309 AgentConfig {
1310 system_prompt: "test".into(),
1311 max_steps: 3,
1312 history_limits: HistoryLimits::default(),
1313 },
1314 PathBuf::from("/tmp"),
1315 );
1316 let sink = Arc::new(CollectSink(Mutex::new(Vec::new())));
1317 let mut history = Vec::new();
1318 let outcome = runtime
1319 .run_turn(
1320 &mut history,
1321 "start".into(),
1322 sink,
1323 Arc::new(Allow),
1324 CancellationToken::new(),
1325 )
1326 .await
1327 .unwrap();
1328 assert_eq!(outcome.steps, 2);
1329 assert_eq!(outcome.usage.input_tokens, Some(6));
1330 assert_eq!(outcome.usage.output_tokens, Some(3));
1331 assert!(matches!(
1332 history.get(2),
1333 Some(Message::Tool {
1334 content,
1335 is_error: false,
1336 ..
1337 }) if content == "hello"
1338 ));
1339 }
1340
1341 #[tokio::test]
1342 async fn denial_is_recorded_as_a_model_visible_tool_failure() {
1343 let provider = Arc::new(ScriptedProvider {
1344 responses: Mutex::new(VecDeque::from([
1345 AssistantResponse {
1346 content: String::new(),
1347 tool_calls: vec![ToolCall {
1348 id: "call-1".into(),
1349 name: "echo".into(),
1350 arguments: serde_json::json!({"value":"blocked"}),
1351 }],
1352 usage: Usage::default(),
1353 },
1354 AssistantResponse {
1355 content: "handled".into(),
1356 tool_calls: Vec::new(),
1357 usage: Usage::default(),
1358 },
1359 ])),
1360 });
1361 let mut registry = ToolRegistry::default();
1362 registry.register(Arc::new(EchoTool)).unwrap();
1363 let runtime = AgentRuntime::new(
1364 provider,
1365 Arc::new(registry),
1366 Arc::new(BudgetContextPolicy::new(ContextConfig::default()).unwrap()),
1367 AgentConfig {
1368 system_prompt: "test".into(),
1369 max_steps: 3,
1370 history_limits: HistoryLimits::default(),
1371 },
1372 PathBuf::from("/tmp"),
1373 );
1374 let mut history = Vec::new();
1375 runtime
1376 .run_turn(
1377 &mut history,
1378 "start".into(),
1379 Arc::new(CollectSink(Mutex::new(Vec::new()))),
1380 Arc::new(Deny),
1381 CancellationToken::new(),
1382 )
1383 .await
1384 .unwrap();
1385 assert!(matches!(
1386 history.get(2),
1387 Some(Message::Tool {
1388 content,
1389 is_error: true,
1390 ..
1391 }) if content.contains("denied")
1392 ));
1393 }
1394
1395 #[tokio::test]
1396 async fn cancellation_during_approval_rolls_back_the_active_tool_group() {
1397 let provider = Arc::new(ScriptedProvider {
1398 responses: Mutex::new(VecDeque::from([AssistantResponse {
1399 content: String::new(),
1400 tool_calls: vec![ToolCall {
1401 id: "call-cancel".into(),
1402 name: "echo".into(),
1403 arguments: serde_json::json!({"value":"hello"}),
1404 }],
1405 usage: Usage::default(),
1406 }])),
1407 });
1408 let mut registry = ToolRegistry::default();
1409 registry.register(Arc::new(EchoTool)).unwrap();
1410 let runtime = AgentRuntime::new(
1411 provider,
1412 Arc::new(registry),
1413 Arc::new(BudgetContextPolicy::new(ContextConfig::default()).unwrap()),
1414 AgentConfig {
1415 system_prompt: "test".into(),
1416 max_steps: 2,
1417 history_limits: HistoryLimits::default(),
1418 },
1419 PathBuf::from("/tmp"),
1420 );
1421 let before = vec![
1422 Message::User {
1423 content: "previous".into(),
1424 },
1425 Message::Assistant {
1426 content: "answer".into(),
1427 tool_calls: Vec::new(),
1428 },
1429 ];
1430 let mut history = before.clone();
1431 let cancellation = CancellationToken::new();
1432 let cancel = cancellation.clone();
1433 let entered = Arc::new(Notify::new());
1434 let wait = Arc::clone(&entered);
1435 let run = runtime.run_turn(
1436 &mut history,
1437 "new turn".into(),
1438 Arc::new(CollectSink(Mutex::new(Vec::new()))),
1439 Arc::new(WaitForCancellation { entered }),
1440 cancellation,
1441 );
1442 let cancel_when_waiting = async move {
1443 wait.notified().await;
1444 cancel.cancel();
1445 };
1446 let (result, ()) = tokio::join!(run, cancel_when_waiting);
1447 assert!(matches!(result, Err(AgentError::Cancelled)));
1448 assert_eq!(history, before);
1449 }
1450
1451 #[tokio::test]
1452 async fn stops_after_the_configured_maximum_step() {
1453 let provider = Arc::new(ScriptedProvider {
1454 responses: Mutex::new(VecDeque::from([AssistantResponse {
1455 content: String::new(),
1456 tool_calls: vec![ToolCall {
1457 id: "call-1".into(),
1458 name: "echo".into(),
1459 arguments: serde_json::json!({"value":"one"}),
1460 }],
1461 usage: Usage::default(),
1462 }])),
1463 });
1464 let mut registry = ToolRegistry::default();
1465 registry.register(Arc::new(EchoTool)).unwrap();
1466 let runtime = AgentRuntime::new(
1467 provider,
1468 Arc::new(registry),
1469 Arc::new(BudgetContextPolicy::new(ContextConfig::default()).unwrap()),
1470 AgentConfig {
1471 system_prompt: "test".into(),
1472 max_steps: 1,
1473 history_limits: HistoryLimits::default(),
1474 },
1475 PathBuf::from("/tmp"),
1476 );
1477 let result = runtime
1478 .run_turn(
1479 &mut Vec::new(),
1480 "start".into(),
1481 Arc::new(CollectSink(Mutex::new(Vec::new()))),
1482 Arc::new(Allow),
1483 CancellationToken::new(),
1484 )
1485 .await;
1486 assert!(matches!(result, Err(AgentError::StepLimit)));
1487 }
1488
1489 #[test]
1490 fn duplicate_tool_registration_does_not_replace_the_original() {
1491 let mut registry = ToolRegistry::default();
1492 registry.register(Arc::new(EchoTool)).unwrap();
1493 assert!(registry.register(Arc::new(EchoTool)).is_err());
1494 assert_eq!(registry.tools.len(), 1);
1495 assert!(registry.get("echo").is_some());
1496 }
1497}