Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
125 changes: 125 additions & 0 deletions code-rs/core/src/context_ledger.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,17 @@ pub struct ContextLedgerEntry {
pub duplicate_key: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ContextDuplicateGroup {
pub duplicate_key: String,
pub entry_count: usize,
pub item_count: usize,
pub bytes: usize,
pub estimated_tokens: usize,
pub labels: Vec<String>,
pub sources: Vec<ContextSourceKind>,
}

#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ContextLedger {
entries: Vec<ContextLedgerEntry>,
Expand Down Expand Up @@ -113,6 +124,60 @@ impl ContextLedger {
parts.join("; ")
)
}

pub fn duplicate_groups(&self) -> Vec<ContextDuplicateGroup> {
let mut groups: BTreeMap<String, Vec<&ContextLedgerEntry>> = BTreeMap::new();
for entry in &self.entries {
let Some(key) = entry.duplicate_key.as_ref().filter(|key| !key.is_empty()) else {
continue;
};
groups.entry(key.clone()).or_default().push(entry);
}

let mut duplicates = groups
.into_iter()
.filter_map(|(duplicate_key, entries)| {
if entries.len() < 2 {
return None;
}

let item_count = entries
.iter()
.map(|entry| entry.item_count)
.sum::<usize>();
let bytes = entries.iter().map(|entry| entry.bytes).sum::<usize>();
let mut labels = Vec::new();
let mut sources = Vec::new();
for entry in entries.iter() {
if !labels.iter().any(|label| label == &entry.label) {
labels.push(entry.label.clone());
}
if !sources.contains(&entry.source) {
sources.push(entry.source);
}
}

Some(ContextDuplicateGroup {
duplicate_key,
entry_count: entries.len(),
item_count,
bytes,
estimated_tokens: estimate_tokens(bytes),
labels,
sources,
})
})
.collect::<Vec<_>>();

duplicates.sort_by(|left, right| {
right
.estimated_tokens
.cmp(&left.estimated_tokens)
.then_with(|| right.entry_count.cmp(&left.entry_count))
.then_with(|| left.duplicate_key.cmp(&right.duplicate_key))
});
duplicates
}
}

pub fn estimate_tokens(bytes: usize) -> usize {
Expand Down Expand Up @@ -252,3 +317,63 @@ fn function_call_output_bytes(body: &FunctionCallOutputBody) -> usize {
.sum(),
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn duplicate_groups_collect_repeated_keys_by_weight() {
let mut ledger = ContextLedger::default();
ledger.push(
ContextSourceKind::UserInstructions,
ContextPersistence::Contextual,
"project instructions",
1,
400,
Some("user_instructions".to_string()),
);
ledger.push(
ContextSourceKind::UserInstructions,
ContextPersistence::Contextual,
"prepended instructions",
2,
600,
Some("user_instructions".to_string()),
);
ledger.push(
ContextSourceKind::ToolSchema,
ContextPersistence::Contextual,
"tool schemas",
4,
2_000,
Some("tool_schemas".to_string()),
);
ledger.push(
ContextSourceKind::EnvironmentContext,
ContextPersistence::GeneratedPerAttempt,
"environment context",
1,
20,
None,
);

let groups = ledger.duplicate_groups();

assert_eq!(groups.len(), 1);
let group = &groups[0];
assert_eq!(group.duplicate_key, "user_instructions");
assert_eq!(group.entry_count, 2);
assert_eq!(group.item_count, 3);
assert_eq!(group.bytes, 1_000);
assert_eq!(group.estimated_tokens, 250);
assert_eq!(
group.labels,
vec![
"project instructions".to_string(),
"prepended instructions".to_string(),
]
);
assert_eq!(group.sources, vec![ContextSourceKind::UserInstructions]);
}
}
76 changes: 76 additions & 0 deletions code-rs/tui/src/chatwidget.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16971,6 +16971,36 @@ impl ChatWidget<'_> {
format_with_separators_u64(ledger.total_estimated_tokens() as u64),
Self::format_context_bytes(ledger.total_bytes()),
));
let duplicate_groups = ledger.duplicate_groups();
if !duplicate_groups.is_empty() {
lines.push(String::new());
let duplicate_count = duplicate_groups.len();
let suffix = if duplicate_count == 1 { "" } else { "s" };
lines.push(format!(
"Possible duplicates: {duplicate_count} group{suffix}"
));
for group in duplicate_groups.iter().take(3) {
let labels = group.labels.join(", ");
lines.push(format!(
" {}: {} entries, ~{} tokens, {}{}",
group.duplicate_key,
group.entry_count,
format_with_separators_u64(group.estimated_tokens as u64),
Self::format_context_bytes(group.bytes),
if labels.is_empty() {
String::new()
} else {
format!(" ({labels})")
},
));
}
if duplicate_groups.len() > 3 {
lines.push(format!(
" ... {} more duplicate groups",
duplicate_groups.len().saturating_sub(3)
));
}
}
lines.push(String::new());
lines.push(format!(
"{:<24} {:>10} {:>9} {:<14} {}",
Expand Down Expand Up @@ -31862,6 +31892,52 @@ use code_core::protocol::OrderMeta;
assert!(lines[tool_row].contains("contextual"));
}

#[test]
fn context_ledger_display_surfaces_duplicate_groups() {
let mut ledger = ContextLedger::default();
ledger.push(
ContextSourceKind::UserInstructions,
ContextPersistence::Contextual,
"user/project instructions",
1,
1_024,
Some("user_instructions".to_string()),
);
ledger.push(
ContextSourceKind::UserInstructions,
ContextPersistence::Contextual,
"prepended project instructions",
1,
512,
Some("user_instructions".to_string()),
);
ledger.push(
ContextSourceKind::ToolSchema,
ContextPersistence::Contextual,
"tool schemas",
2,
800,
Some("tool_schemas".to_string()),
);

let lines = ChatWidget::context_ledger_display_lines(&ledger);

let duplicate_header = lines
.iter()
.position(|line| line == "Possible duplicates: 1 group")
.expect("duplicate header");
let duplicate_row = lines
.iter()
.skip(duplicate_header)
.find(|line| line.contains("user_instructions"))
.expect("duplicate row");
assert!(duplicate_row.contains("2 entries"));
assert!(duplicate_row.contains("~384 tokens"));
assert!(duplicate_row.contains("user/project instructions"));
assert!(duplicate_row.contains("prepended project instructions"));
assert!(lines.iter().any(|line| line.starts_with("Source")));
}

fn test_rate_limit_snapshot() -> RateLimitSnapshotEvent {
RateLimitSnapshotEvent {
primary_used_percent: 12.0,
Expand Down
8 changes: 0 additions & 8 deletions code-rs/tui/src/clipboard_copy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,14 +35,6 @@ impl ClipboardLease {
_clipboard: Some(clipboard),
}
}

#[cfg(test)]
pub(crate) fn test() -> Self {
Self {
#[cfg(target_os = "linux")]
_clipboard: None,
}
}
}

fn copy_to_clipboard_with(
Expand Down