open-source/cursor-byok
1
0
Fork 0
mirror of https://wget.la/https://github.com/leookun/cursor-byok synced 2026-10-10 10:18:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: implement context usage anchor for improved token estimation

Browse Source 
- Introduced `ContextUsageAnchor` struct to track context input tokens and message count for conversations.
- Updated token estimation functions to utilize the context usage anchor, enhancing accuracy in estimating tokens for projected messages.
- Refactored compaction logic to incorporate context usage anchor, allowing for more efficient management of token budgets during model runs.
- Added tests to validate the behavior of the context usage anchor across different scenarios, including model switching and message additions.
This commit is contained in:
leookun
2026-09-01 16:07:51 +08:00
parent 8c6c415a84
commit d004139526
6 changed files with 394 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/src/run/engine.rs
+44 -2
View File
@@ -10,7 +10,7 @@ use crate::{
ToolRoundId, Usage,
},
provider::Provider,
store::{RunStatus, Store},
store::{ContextUsageAnchor, RunStatus, Store},
};
use super::{
@@ -92,6 +92,14 @@ impl RunEngine {
cancellation: &CancellationToken,
) -> (RunOutcome, Option<Usage>) {
let mut usage = None;
let mut context_usage_anchor = match self
.store
.latest_context_usage(prepared.conversation_id.as_str())
.await
{
Ok(anchor) => anchor,
Err(error) => return (RunOutcome::Failed(error.into()), usage),
};
tracing::info!(
checkpoint_id = checkpoint.0,
"Run claimed conversation ownership"
@@ -176,7 +184,7 @@ impl RunEngine {
Err(error) => return (RunOutcome::Failed(error.into()), usage),
};
if prepared.action != RunAction::Compact
&& super::compaction::should_compact(prepared, &history)
&& super::compaction::should_compact(prepared, &history, context_usage_anchor)
{
match self
.auto_compact(prepared, checkpoint, &messages, client, cancellation)
@@ -184,6 +192,7 @@ impl RunEngine {
{
Ok((next_checkpoint, compaction_usage)) => {
checkpoint = next_checkpoint;
context_usage_anchor = None;
if let Some(compaction_usage) = compaction_usage {
accumulate_usage(&mut usage, compaction_usage);
}
@@ -272,11 +281,21 @@ impl RunEngine {
match interrupted {
Ok(cycle) => {
if let Some(cycle_usage) = cycle.usage {
update_context_usage_anchor(
&mut context_usage_anchor,
cycle_usage,
request.history.len(),
);
accumulate_usage(&mut usage, cycle_usage);
}
}
Err(failure) => {
if let Some(cycle_usage) = failure.usage {
update_context_usage_anchor(
&mut context_usage_anchor,
cycle_usage,
request.history.len(),
);
accumulate_usage(&mut usage, cycle_usage);
}
}
@@ -319,6 +338,11 @@ impl RunEngine {
Ok(cycle) => break 'attempt cycle,
Err(cycle_failure) => {
if let Some(cycle_usage) = cycle_failure.usage {
update_context_usage_anchor(
&mut context_usage_anchor,
cycle_usage,
request.history.len(),
);
accumulate_usage(&mut usage, cycle_usage);
}
if cancellation.is_cancelled() {
@@ -420,6 +444,11 @@ impl RunEngine {
}
};
if let Some(cycle_usage) = cycle.usage {
update_context_usage_anchor(
&mut context_usage_anchor,
cycle_usage,
request.history.len(),
);
accumulate_usage(&mut usage, cycle_usage);
}
@@ -879,6 +908,19 @@ async fn hydrate_tool_images(
Ok(())
}
fn update_context_usage_anchor(
anchor: &mut Option<ContextUsageAnchor>,
usage: Usage,
message_count: usize,
) {
if let Some(context_input_tokens) = usage.context_input_tokens {
*anchor = Some(ContextUsageAnchor {
context_input_tokens,
message_count,
});
}
}
fn accumulate_usage(total: &mut Option<Usage>, usage: Usage) {
match total {
Some(total) => *total += usage,