mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-10-08 03:10:26 +00:00
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:
parent
a64a58d567
commit
c987766641
5 changed files with 518 additions and 88 deletions
1
Cargo.lock
generated
1
Cargo.lock
generated
|
|
@ -1629,6 +1629,7 @@ dependencies = [
|
|||
"serde_json",
|
||||
"sha2",
|
||||
"shell-escape",
|
||||
"strum",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue