agent: use API usage baseline for compaction context estimate (#366)

## Summary

Replaces the whole-history `chars / 4` compaction trigger with a Claude
Code-style hot-path estimator: find the latest assistant turn with real
provider-reported `usage.total_tokens()`, use that as the baseline, and
add local char estimates only for turns appended after it. This makes
compaction sensitive to actual provider-reported context usage
(including cache read/write and reasoning tokens) without adding any
provider token-count API calls.

### Plan Summary

- **New estimator** (`estimate_active_context_usage`): returns a
`ContextEstimate` with both a token count and an `ContextEstimateMethod`
enum tag (`ApiUsagePlusLocalDelta` or `LocalEstimate`).
- **`check_context_usage`** now returns `Option<ContextEstimate>`
instead of `bool`, and includes `estimate_method` in warning `details`.
The caller passes the estimate directly into `compact_context`, avoiding
a double-compute.
- **`compact_context`** signature drops `system_prompt` (no longer
needed) and takes the pre-computed `ContextEstimate`. The
`CompactionStarted` event is now emitted *after* the `turns.len() <=
preserve_count` no-op guard, so a no-op can never emit `Started` without
`Completed`.
- **`History::compact`** invalidates preserved assistant `usage` (resets
to `TokenCounts::default()`) so a preserved turn's pre-compaction
provider baseline never becomes the next estimate's anchor. Content,
tool calls, provider parts, and response IDs are untouched.
- **`session.compact_if_needed`** restructured to early-return on `None`
from `check_context_usage` or on compaction disabled, simplifying the
nesting.

## Key design decisions

**Why invalidate preserved assistant usage?** After compaction the prior
turns are gone, so a stored `total_tokens` from before compaction would
overstate the new context. The authoritative billing record is in
emitted run events, not in mutable runtime history.

**Why return `Option<ContextEstimate>` from `check_context_usage`?**
Avoids recomputing the estimate in `compact_context`. It also makes the
call-site idiom (`let Some(estimate) = ... else { return; }`) an
explicit gate, which is cleaner than a separate bool-then-compact
pattern.

**Why `strum::IntoStaticStr` on the method enum?** Lets the variant
serialize to a `&'static str` for the JSON `details` field without a
manual `match` or adding `serde` derives.


### Fabro Details

<details>
<summary>Ran 9 stages in 30m 54s for $7.93</summary>

| Stage | Duration | Cost | Retries |
|---|---|---|---|
| start | 0s | – | 0 |
| toolchain | 2s | – | 0 |
| preflight_compile | 1m 52s | – | 0 |
| preflight_lint | 2m 6s | – | 0 |
| implement | 8m 11s | $3.47 | 0 |
| simplify_opus | 11m 7s | $3.53 | 0 |
| simplify_gpt | 2m 39s | $0.93 | 0 |
| verify | 4m 4s | – | 0 |
| fmt | 3s | – | 0 |
| **Total** | **30m 54s** | **$7.93** | **0** |

</details>

<details>
<summary>Ran <code>ImplementPlan.fabro</code> (12 nodes and 15
edges)</summary>

```dot
digraph ImplementPlan {
    graph [
        goal="Implement and simplify",
        model_stylesheet="
            * { model: claude-opus-4-7; }
        "
    ]
    rankdir=LR

    start [shape=Mdiamond, label="Start"]
    exit  [shape=Msquare, label="Exit"]

    toolchain         [label="Toolchain", shape=parallelogram, script="command -v cargo >/dev/null || { curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && sudo ln -sf $HOME/.cargo/bin/* /usr/local/bin/; }; cargo --version 2>&1", max_retries=0]
    preflight_compile [label="Preflight Compile", shape=parallelogram, script="cargo check -q --workspace 2>&1", max_retries=0]
    preflight_lint    [label="Preflight Lint", shape=parallelogram, script="cargo +nightly-2026-04-14 clippy -q --workspace --all-targets -- -D warnings 2>&1", max_retries=0]
    fix_lints         [label="Fix Lints", prompt="The preflight lint step failed. Read the build output from context and fix all clippy lint warnings.", max_visits=3]
    implement         [label="Implement", prompt="Read the plan file referenced in the goal and implement every step. Make all the code changes described in the plan. Use red/green TDD.", model="gpt-55", reasoning_effort="xhigh"]
    simplify_opus     [label="Simplify (Opus)", prompt="@prompts/simplify.md"]
    simplify_gpt      [label="Simplify (GPT-55)", prompt="@prompts/simplify.md", model="gpt-55"]
    verify            [label="Verify", shape=parallelogram, script="cargo +nightly-2026-04-14 clippy -q --workspace --all-targets -- -D warnings 2>&1 && cargo nextest run --cargo-quiet --workspace --status-level fail 2>&1 && cargo dev docs refresh 2>&1 && cargo dev docs check 2>&1", goal_gate=true, retry_target="fixup"]
    fixup             [label="Fixup", prompt="The verify step failed. Read the build output from context and fix all clippy lint warnings, test failures, and generated docs errors.", max_visits=3]
    fmt               [label="Format", shape=parallelogram, script="cargo +nightly-2026-04-14 fmt --all 2>&1", max_retries=0]

    start -> toolchain
    toolchain -> preflight_compile [condition="outcome=succeeded"]
    toolchain -> exit
    preflight_compile -> preflight_lint [condition="outcome=succeeded"]
    preflight_compile -> exit
    preflight_lint -> implement [condition="outcome=succeeded"]
    preflight_lint -> fix_lints
    fix_lints -> preflight_lint
    implement -> simplify_opus -> simplify_gpt -> verify
    verify -> fmt   [condition="outcome=succeeded"]
    verify -> fixup
    fixup -> verify
    fmt -> exit
}

```

</details>

⚒️ Generated with [Fabro](https://fabro.sh)

---------

Co-authored-by: Fabro <noreply@fabro.sh>
This commit is contained in:
fabro-sh-0530[bot] 2026-05-23 12:55:23 -04:00 • committed by GitHub
parent a64a58d567
commit c987766641
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 518 additions and 88 deletions

1
Cargo.lock generated
View file

@ -1629,6 +1629,7 @@ dependencies = [
"serde_json",
"sha2",
"shell-escape",
"strum",
"tempfile",
"thiserror 2.0.18",
"tokio",

View file

@ -38,6 +38,7 @@ fabro-http.workspace = true
thiserror.workspace = true
serde.workspace = true
serde_json.workspace = true
strum.workspace = true
tokio.workspace = true
uuid.workspace = true
futures.workspace = true

View file

@ -11,37 +11,53 @@ use crate::file_tracker::FileTracker;
use crate::history::History;
use crate::types::{AgentEvent, Message};
const APPROX_CHARS_PER_TOKEN: usize = 4;
#[derive(Debug, Clone, Copy, PartialEq, Eq, strum::IntoStaticStr)]
#[strum(serialize_all = "snake_case")]
pub(crate) enum ContextEstimateMethod {
ApiUsagePlusLocalDelta,
LocalEstimate,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct ContextEstimate {
pub tokens: usize,
pub method: ContextEstimateMethod,
}
/// Check whether the context window usage exceeds the configured threshold.
/// Emits a `Warning` event with kind `"context_window"` when over the
/// threshold. Returns `true` if the threshold is exceeded.
pub fn check_context_usage(
/// threshold. Returns `Some(estimate)` if the threshold is exceeded so the
/// caller can pass it to `compact_context` without recomputing.
pub(crate) fn check_context_usage(
system_prompt: &str,
history: &History,
provider_profile: &dyn AgentProfile,
threshold_percent: usize,
emitter: &Emitter,
session_id: &str,
) -> bool {
let estimated_tokens = estimate_token_count(system_prompt, history);
) -> Option<ContextEstimate> {
let estimate = estimate_active_context_usage(system_prompt, history);
let context_window = provider_profile.context_window_size();
let threshold = context_window * threshold_percent / 100;
if estimated_tokens > threshold {
if estimate.tokens > threshold {
let usage_percent = estimate.tokens.saturating_mul(100) / context_window;
let method: &'static str = estimate.method.into();
emitter.emit(session_id.to_owned(), AgentEvent::Warning {
kind: "context_window".into(),
message: format!(
"Context window usage: {}%",
estimated_tokens * 100 / context_window
),
message: format!("Context window usage: {usage_percent}%"),
details: serde_json::json!({
"estimated_tokens": estimated_tokens,
"estimated_tokens": estimate.tokens,
"context_window_size": context_window,
"usage_percent": estimated_tokens * 100 / context_window,
"usage_percent": usage_percent,
"estimate_method": method,
}),
});
true
Some(estimate)
} else {
false
None
}
}
@ -51,29 +67,29 @@ pub fn check_context_usage(
clippy::too_many_arguments,
reason = "Context compaction needs explicit history, model, tracking, and emission inputs."
)]
pub async fn compact_context(
pub(crate) async fn compact_context(
history: &mut History,
llm_client: &Client,
provider_profile: &dyn AgentProfile,
system_prompt: &str,
file_tracker: &FileTracker,
preserve_count: usize,
estimate: ContextEstimate,
emitter: &Emitter,
session_id: &str,
) -> Result<(), Error> {
let estimated_tokens = estimate_token_count(system_prompt, history);
let context_window = provider_profile.context_window_size();
let original_turn_count = history.turns().len();
emitter.emit(session_id.to_owned(), AgentEvent::CompactionStarted {
estimated_tokens,
context_window_size: context_window,
});
// Determine turns to summarize
// Determine turns to summarize. If there are not enough turns to compact,
// do not emit a started event without a matching completion.
if original_turn_count <= preserve_count {
return Ok(());
}
emitter.emit(session_id.to_owned(), AgentEvent::CompactionStarted {
estimated_tokens: estimate.tokens,
context_window_size: provider_profile.context_window_size(),
});
let turns_to_summarize = &history.turns()[..original_turn_count - preserve_count];
let rendered = render_turns_for_summary(turns_to_summarize);
@ -139,7 +155,7 @@ function names, error messages, and exact values. Omit pleasantries and conversa
"A different assistant began this task and produced the following summary. \
Build on their progress — do not repeat completed steps.\n\n{summary_text}"
);
let summary_token_estimate = summary_content.len() / 4;
let summary_token_estimate = estimate_chars_local_tokens(summary_content.len());
history.compact(preserve_count, summary_content);
@ -153,40 +169,76 @@ Build on their progress — do not repeat completed steps.\n\n{summary_text}"
Ok(())
}
/// Estimate the total token count of the system prompt and conversation
/// history. Uses a rough heuristic of ~4 characters per token.
pub fn estimate_token_count(system_prompt: &str, history: &History) -> usize {
let mut total_chars = system_prompt.len();
for turn in history.turns() {
match turn {
Message::User { content, .. } => total_chars += content.len(),
Message::Assistant {
content,
tool_calls,
..
} => {
total_chars += content.len();
if let Some(r) = turn.reasoning_text() {
total_chars += r.len();
}
for tc in tool_calls {
total_chars += tc.name.len();
total_chars += tc.arguments.to_string().len();
}
}
Message::ToolResults { results, .. } => {
for r in results {
total_chars += r.content.to_string().len();
}
}
Message::System { content, .. } | Message::Steering { content, .. } => {
total_chars += content.len();
}
}
pub(crate) fn estimate_active_context_usage(
system_prompt: &str,
history: &History,
) -> ContextEstimate {
let turns = history.turns();
if let Some((baseline_index, baseline_tokens)) = latest_assistant_usage_baseline(turns) {
let local_delta = estimate_turns_local_tokens(&turns[baseline_index + 1..]);
return ContextEstimate {
tokens: baseline_tokens.saturating_add(local_delta),
method: ContextEstimateMethod::ApiUsagePlusLocalDelta,
};
}
total_chars / 4 // rough estimate: ~4 chars per token
ContextEstimate {
tokens: estimate_chars_local_tokens(
system_prompt
.len()
.saturating_add(estimate_turns_local_chars(turns)),
),
method: ContextEstimateMethod::LocalEstimate,
}
}
fn latest_assistant_usage_baseline(turns: &[Message]) -> Option<(usize, usize)> {
turns.iter().enumerate().rev().find_map(|(index, turn)| {
if let Message::Assistant { usage, .. } = turn {
let total_tokens = usage.total_tokens();
if total_tokens > 0 {
return Some((index, usize::try_from(total_tokens).unwrap_or(usize::MAX)));
}
}
None
})
}
fn estimate_turns_local_tokens(turns: &[Message]) -> usize {
estimate_chars_local_tokens(estimate_turns_local_chars(turns))
}
fn estimate_turns_local_chars(turns: &[Message]) -> usize {
turns.iter().fold(0usize, |total, turn| {
total.saturating_add(estimate_turn_chars(turn))
})
}
fn estimate_chars_local_tokens(chars: usize) -> usize {
chars / APPROX_CHARS_PER_TOKEN
}
fn estimate_turn_chars(turn: &Message) -> usize {
match turn {
Message::User { content, .. }
| Message::System { content, .. }
| Message::Steering { content, .. } => content.len(),
Message::Assistant {
content,
tool_calls,
..
} => {
let reasoning_chars = turn.reasoning_text().map_or(0, str::len);
let tool_call_chars: usize = tool_calls
.iter()
.map(|tc| tc.name.len() + tc.arguments.to_string().len())
.sum();
content.len() + reasoning_chars + tool_call_chars
}
Message::ToolResults { results, .. } => {
results.iter().map(|r| r.content.to_string().len()).sum()
}
}
}
/// Render conversation turns into a human-readable summary format for the
@ -313,14 +365,176 @@ mod tests {
}
#[test]
fn estimate_token_count_basic() {
fn estimate_local_token_count_basic() {
let mut history = History::default();
history.push(Message::User {
content: "Hello world".into(), // 11 chars
timestamp: SystemTime::now(),
});
// system_prompt = "test" (4 chars) + 11 chars = 15 chars / 4 = 3 tokens
assert_eq!(estimate_token_count("test", &history), 3);
// system_prompt = "test" (4/4 = 1 token) + 11 chars / 4 = 2 tokens = 3 tokens
let estimate = estimate_active_context_usage("test", &history);
assert_eq!(estimate.tokens, 3);
assert_eq!(estimate.method, ContextEstimateMethod::LocalEstimate);
}
#[test]
fn active_context_estimate_without_assistant_usage_uses_local_estimate() {
let mut history = History::default();
history.push(Message::User {
content: "Hello world".into(), // 11 chars => 2 tokens
timestamp: SystemTime::now(),
});
history.push(Message::Assistant {
// 18 chars content + tool call name (9) + args (16) = 43 chars => 10 tokens
content: "No usage available".into(),
tool_calls: vec![ToolCall::new(
"call_1",
"read_file",
serde_json::json!({"path": "foo.rs"}),
)],
provider_parts: vec![],
usage: Box::new(TokenCounts::default()),
response_id: "resp_1".into(),
timestamp: SystemTime::now(),
});
history.push(Message::ToolResults {
// 4 chars => 1 token
results: vec![ToolResult::success("call_1", serde_json::json!(1234))],
timestamp: SystemTime::now(),
});
let estimate = estimate_active_context_usage("test", &history);
assert_eq!(estimate.method, ContextEstimateMethod::LocalEstimate);
// (system prompt 4 + turn chars 11 + 18 + 9 + 16 + 4) / 4 = 62/4 = 15
assert_eq!(estimate.tokens, 15);
}
#[test]
fn active_context_local_estimate_matches_whole_history_rounding() {
let mut history = History::default();
history.push(Message::User {
content: "abc".into(),
timestamp: SystemTime::now(),
});
let estimate = estimate_active_context_usage("x", &history);
assert_eq!(estimate.method, ContextEstimateMethod::LocalEstimate);
assert_eq!(estimate.tokens, 1);
}
#[test]
fn active_context_estimate_uses_latest_assistant_usage_plus_later_turns() {
let mut history = History::default();
history.push(Message::User {
content: "ignored before baseline".repeat(100),
timestamp: SystemTime::now(),
});
history.push(Message::Assistant {
content: "baseline response".into(),
tool_calls: vec![],
provider_parts: vec![],
usage: Box::new(TokenCounts {
input_tokens: 50,
..TokenCounts::default()
}),
response_id: "resp_1".into(),
timestamp: SystemTime::now(),
});
history.push(Message::ToolResults {
// JSON number renders as 4 chars => 1 local token.
results: vec![ToolResult::success("call_1", serde_json::json!(1234))],
timestamp: SystemTime::now(),
});
history.push(Message::User {
// 16 chars => 4 local tokens.
content: "u".repeat(16),
timestamp: SystemTime::now(),
});
history.push(Message::Steering {
// 8 chars => 2 local tokens.
content: "s".repeat(8),
timestamp: SystemTime::now(),
});
let estimate = estimate_active_context_usage("ignored system prompt", &history);
assert_eq!(estimate.tokens, 57);
assert_eq!(
estimate.method,
ContextEstimateMethod::ApiUsagePlusLocalDelta
);
}
#[test]
fn active_context_estimate_uses_total_tokens_including_cache_and_reasoning() {
let mut history = History::default();
history.push(Message::Assistant {
content: "short".into(),
tool_calls: vec![],
provider_parts: vec![],
usage: Box::new(TokenCounts {
input_tokens: 10,
output_tokens: 20,
reasoning_tokens: 30,
cache_read_tokens: 40,
cache_write_tokens: 50,
}),
response_id: "resp_1".into(),
timestamp: SystemTime::now(),
});
let estimate = estimate_active_context_usage("", &history);
assert_eq!(estimate.tokens, 150);
assert_eq!(
estimate.method,
ContextEstimateMethod::ApiUsagePlusLocalDelta
);
}
#[test]
fn active_context_estimate_ignores_earlier_usage_when_later_usage_exists() {
let mut history = History::default();
history.push(Message::Assistant {
content: "older response".into(),
tool_calls: vec![],
provider_parts: vec![],
usage: Box::new(TokenCounts {
input_tokens: 1_000,
..TokenCounts::default()
}),
response_id: "resp_old".into(),
timestamp: SystemTime::now(),
});
history.push(Message::User {
content: "ignored before latest baseline".repeat(100),
timestamp: SystemTime::now(),
});
history.push(Message::Assistant {
content: "latest response".into(),
tool_calls: vec![],
provider_parts: vec![],
usage: Box::new(TokenCounts {
input_tokens: 20,
..TokenCounts::default()
}),
response_id: "resp_new".into(),
timestamp: SystemTime::now(),
});
history.push(Message::User {
content: "u".repeat(8),
timestamp: SystemTime::now(),
});
let estimate = estimate_active_context_usage("", &history);
assert_eq!(estimate.tokens, 22);
assert_eq!(
estimate.method,
ContextEstimateMethod::ApiUsagePlusLocalDelta
);
}
#[test]
@ -329,8 +543,8 @@ mod tests {
let emitter = Emitter::new();
let profile = TestProfile::new();
// Empty history, huge context window => well below threshold
let over = check_context_usage("short", &history, &profile, 80, &emitter, "sess");
assert!(!over);
let result = check_context_usage("short", &history, &profile, 80, &emitter, "sess");
assert!(result.is_none());
}
#[test]
@ -345,11 +559,12 @@ mod tests {
let mut rx = emitter.subscribe();
// TestProfile has context_window=200_000 by default; use a small one
let profile = TestProfile::with_context_window(ToolRegistry::new(), 100);
let over = check_context_usage("prompt", &history, &profile, 80, &emitter, "sess");
assert!(over);
let result = check_context_usage("prompt", &history, &profile, 80, &emitter, "sess");
assert!(result.is_some());
// Should have emitted a Warning
let event = rx.try_recv().unwrap();
assert!(matches!(event.event, AgentEvent::Warning { .. }));
assert!(matches!(event.event, AgentEvent::Warning { details, .. }
if details["estimate_method"] == "local_estimate"));
}
}

View file

@ -1,4 +1,4 @@
use fabro_llm::types::{ContentPart, Message as LlmMessage, Role};
use fabro_llm::types::{ContentPart, Message as LlmMessage, Role, TokenCounts};
use fabro_types::SessionMessage;
use crate::types::Message;
@ -32,11 +32,17 @@ impl History {
self.turns.iter().map(Message::to_session_message).collect()
}
/// Compact the history by replacing all but the trailing `preserve_count`
/// turns with a summary `System` message. Preserved assistant turns have
/// their `usage` reset to default so a later context-window estimate does
/// not treat pre-compaction provider-reported usage as the new baseline;
/// authoritative billing is recorded via emitted run events.
pub fn compact(&mut self, preserve_count: usize, summary: String) {
if self.turns.len() <= preserve_count {
return;
}
let preserved = self.turns.split_off(self.turns.len() - preserve_count);
let mut preserved = self.turns.split_off(self.turns.len() - preserve_count);
Self::invalidate_preserved_usage(&mut preserved);
let discarded = std::mem::take(&mut self.turns);
let extracted_user_messages =
extract_recent_user_messages(discarded, COMPACTION_USER_MESSAGE_TOKEN_BUDGET);
@ -49,6 +55,14 @@ impl History {
self.strip_opaque_provider_items();
}
fn invalidate_preserved_usage(preserved: &mut [Message]) {
for turn in preserved {
if let Message::Assistant { usage, .. } = turn {
**usage = TokenCounts::default();
}
}
}
/// Remove provider-specific opaque items that are no longer valid after
/// compaction. OpenAI reasoning and message items are opaque round-trip
/// data tied to specific API responses; after compaction replaces their
@ -599,6 +613,60 @@ mod tests {
}
}
#[test]
fn compact_preserves_assistant_data_but_resets_usage() {
let mut history = History::default();
history.push(Message::User {
content: "old msg".into(),
timestamp: SystemTime::now(),
});
let tool_call = ToolCall::new("call_1", "search", serde_json::json!({"query": "fabro"}));
let thinking = ContentPart::Thinking(ThinkingData {
text: "deep thought".into(),
signature: Some("sig_xyz".into()),
redacted: false,
});
history.push(Message::Assistant {
content: "answer".into(),
tool_calls: vec![tool_call.clone()],
provider_parts: vec![thinking.clone()],
usage: Box::new(TokenCounts {
input_tokens: 10,
output_tokens: 20,
reasoning_tokens: 30,
cache_read_tokens: 40,
cache_write_tokens: 50,
}),
response_id: "resp_1".into(),
timestamp: SystemTime::now(),
});
history.compact(1, "Summary".into());
let assistant_turn = history
.turns()
.iter()
.find(|turn| matches!(turn, Message::Assistant { .. }))
.expect("preserved assistant turn");
if let Message::Assistant {
content,
tool_calls,
provider_parts,
usage,
response_id,
..
} = assistant_turn
{
assert_eq!(content, "answer");
assert_eq!(tool_calls, &[tool_call]);
assert_eq!(provider_parts, &[thinking]);
assert_eq!(response_id, "resp_1");
assert_eq!(**usage, TokenCounts::default());
} else {
panic!("expected Assistant turn");
}
}
#[test]
fn compact_strips_reasoning_from_all_preserved_assistant_turns() {
let mut history = History::default();

View file

@ -1670,31 +1670,34 @@ impl Session {
}
async fn compact_if_needed(&mut self) {
let over_threshold = check_context_usage(
let Some(estimate) = check_context_usage(
&self.system_prompt,
&self.history,
self.provider_profile.as_ref(),
self.config.compaction_threshold_percent,
&self.event_emitter,
&self.id,
);
if over_threshold && self.config.enable_context_compaction {
if let Err(e) = compact_context(
&mut self.history,
&self.llm_client,
self.provider_profile.as_ref(),
&self.system_prompt,
&self.file_tracker,
self.config.compaction_preserve_turns,
&self.event_emitter,
&self.id,
)
.await
{
self.event_emitter.emit(self.id.clone(), AgentEvent::Error {
error: Error::InvalidState(format!("Context compaction failed: {e}")),
});
}
) else {
return;
};
if !self.config.enable_context_compaction {
return;
}
if let Err(e) = compact_context(
&mut self.history,
&self.llm_client,
self.provider_profile.as_ref(),
&self.file_tracker,
self.config.compaction_preserve_turns,
estimate,
&self.event_emitter,
&self.id,
)
.await
{
self.event_emitter.emit(self.id.clone(), AgentEvent::Error {
error: Error::InvalidState(format!("Context compaction failed: {e}")),
});
}
}
@ -1853,7 +1856,7 @@ mod tests {
use fabro_llm::error::{ProviderErrorDetail, ProviderErrorKind};
use fabro_llm::provider::{ProviderAdapter, StreamEventStream};
use fabro_llm::types::{
ContentPart, ReasoningEffort, Request, Response, Role, StreamEvent, ToolCall,
ContentPart, ReasoningEffort, Request, Response, Role, StreamEvent, TokenCounts, ToolCall,
ToolDefinition,
};
use futures::stream;
@ -3653,13 +3656,25 @@ mod tests {
assert!(found_auth_error_event, "expected auth error event");
}
fn response_with_usage(mut response: Response, usage: TokenCounts) -> Response {
response.usage = usage;
response
}
fn response_with_input_tokens(response: Response, input_tokens: i64) -> Response {
response_with_usage(response, TokenCounts {
input_tokens,
..TokenCounts::default()
})
}
#[tokio::test]
async fn compaction_triggered_when_over_threshold() {
// Tiny context window to trigger compaction
// Responses: [0] conversation response (stream), [1] summarization (complete),
// [2] unused fallback
let responses = vec![
text_response("OK"),
response_with_usage(text_response("OK"), TokenCounts::default()),
text_response("Here is the summary of the conversation so far."),
text_response("fallback"),
];
@ -3704,6 +3719,90 @@ mod tests {
);
}
#[tokio::test]
async fn compaction_uses_assistant_usage_baseline_for_short_response() {
let responses = vec![
response_with_input_tokens(text_response("OK"), 90),
text_response("Here is the summary of the conversation so far."),
text_response("fallback"),
];
let provider = Arc::new(MockLlmProvider::new(responses));
let client = make_client(provider).await;
let registry = ToolRegistry::new();
let profile = Arc::new(TestProfile::with_context_window(registry, 100));
let env = Arc::new(MockSandbox::default());
let config = SessionOptions {
enable_context_compaction: true,
compaction_preserve_turns: 1,
..Default::default()
};
let mut session = Session::new(client, profile, env, config, None);
let mut rx = session.subscribe();
session.process_input("hi").await.unwrap();
let mut started = None;
let mut found_completed = false;
while let Ok(event) = rx.try_recv() {
match event.event {
AgentEvent::CompactionStarted {
estimated_tokens,
context_window_size,
} => started = Some((estimated_tokens, context_window_size)),
AgentEvent::CompactionCompleted { .. } => found_completed = true,
_ => {}
}
}
assert_eq!(started, Some((90, 100)));
assert!(
found_completed,
"CompactionCompleted event should be emitted"
);
}
#[tokio::test]
async fn compaction_noop_does_not_emit_started() {
let large_input = "x".repeat(400);
let responses = vec![text_response("OK")];
let provider = Arc::new(MockLlmProvider::new(responses));
let client = make_client(provider).await;
let registry = ToolRegistry::new();
let profile = Arc::new(TestProfile::with_context_window(registry, 100));
let env = Arc::new(MockSandbox::default());
let config = SessionOptions {
enable_context_compaction: true,
compaction_preserve_turns: 10,
..Default::default()
};
let mut session = Session::new(client, profile, env, config, None);
let mut rx = session.subscribe();
session.process_input(&large_input).await.unwrap();
let mut found_warning = false;
let mut found_compaction = false;
while let Ok(event) = rx.try_recv() {
match event.event {
AgentEvent::Warning { kind, .. } if kind == "context_window" => {
found_warning = true;
}
AgentEvent::CompactionStarted { .. } | AgentEvent::CompactionCompleted { .. } => {
found_compaction = true;
}
_ => {}
}
}
assert!(found_warning, "threshold should have been exceeded");
assert!(
!found_compaction,
"no-op compaction should not emit started or completed events"
);
}
#[tokio::test]
async fn compaction_not_triggered_when_disabled() {
let large_input = "x".repeat(400);
@ -3735,6 +3834,49 @@ mod tests {
assert!(!found_compaction, "No compaction events when disabled");
}
#[tokio::test]
async fn compaction_disabled_blocks_api_usage_baseline_compaction() {
let responses = vec![response_with_input_tokens(text_response("OK"), 90)];
let provider = Arc::new(MockLlmProvider::new(responses));
let client = make_client(provider).await;
let registry = ToolRegistry::new();
let profile = Arc::new(TestProfile::with_context_window(registry, 100));
let env = Arc::new(MockSandbox::default());
let config = SessionOptions {
enable_context_compaction: false,
compaction_preserve_turns: 1,
..Default::default()
};
let mut session = Session::new(client, profile, env, config, None);
let mut rx = session.subscribe();
session.process_input("hi").await.unwrap();
let mut found_api_usage_warning = false;
let mut found_compaction = false;
while let Ok(event) = rx.try_recv() {
match event.event {
AgentEvent::Warning { details, .. }
if details["estimated_tokens"] == 90
&& details["estimate_method"] == "api_usage_plus_local_delta" =>
{
found_api_usage_warning = true;
}
AgentEvent::CompactionStarted { .. } | AgentEvent::CompactionCompleted { .. } => {
found_compaction = true;
}
_ => {}
}
}
assert!(
found_api_usage_warning,
"API usage baseline should still drive context warning"
);
assert!(!found_compaction, "compaction must remain disabled");
}
#[tokio::test]
async fn compaction_failure_is_non_fatal() {
// Response [0] = conversation response (stream), [1] will be used for
@ -3789,7 +3931,10 @@ mod tests {
}
let large_input = "x".repeat(400);
let responses = vec![text_response("OK")];
let responses = vec![response_with_usage(
text_response("OK"),
TokenCounts::default(),
)];
let provider = Arc::new(StreamOnlyProvider {
responses,