Skip to content

fix(/context): gracefully handle huge context files + ux #1331

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 7, 2025
Merged

fix(/context): gracefully handle huge context files + ux #1331

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 2 additions & 0 deletions crates/chat-cli/src/cli/chat/consts.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,4 +16,6 @@ pub const MAX_USER_MESSAGE_SIZE: usize = 600_000;
/// In tokens
pub const CONTEXT_WINDOW_SIZE: usize = 200_000;

pub const CONTEXT_FILES_MAX_SIZE: usize = 150_000;

pub const MAX_CHARS: usize = TokenCounter::token_to_chars(CONTEXT_WINDOW_SIZE); // Character-based warning threshold
122 changes: 74 additions & 48 deletions crates/chat-cli/src/cli/chat/context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,12 @@ use serde::{
};
use tracing::debug;

use super::consts::CONTEXT_FILES_MAX_SIZE;
use super::hooks::{
Hook,
HookExecutor,
};
use super::util::drop_matched_context_files;
use crate::platform::Context;
use crate::util::directories;

Expand All @@ -44,6 +46,8 @@ pub struct ContextConfig {
pub struct ContextManager {
ctx: Arc<Context>,

max_context_files_size: usize,

/// Global context configuration that applies to all profiles.
pub global_config: ContextConfig,

Expand All @@ -65,9 +69,16 @@ impl ContextManager {
/// 2. Load the global configuration
/// 3. Load the default profile configuration
///
/// # Arguments
/// * `ctx` - The context to use
/// * `max_context_files_size` - Optional maximum token size for context files. If not provided,
/// defaults to `CONTEXT_FILES_MAX_SIZE`.
///
/// # Returns
/// A Result containing the new ContextManager or an error
pub async fn new(ctx: Arc<Context>) -> Result<Self> {
pub async fn new(ctx: Arc<Context>, max_context_files_size: Option<usize>) -> Result<Self> {
let max_context_files_size = max_context_files_size.unwrap_or(CONTEXT_FILES_MAX_SIZE);

let profiles_dir = directories::chat_profiles_dir(&ctx)?;

ctx.fs().create_dir_all(&profiles_dir).await?;
Expand All @@ -78,6 +89,7 @@ impl ContextManager {

Ok(Self {
ctx,
max_context_files_size,
global_config,
current_profile,
profile_config,
Expand Down Expand Up @@ -136,7 +148,7 @@ impl ContextManager {
for path in &paths {
// We're using a temporary context_files vector just for validation
// Pass is_validation=true to ensure we error if glob patterns don't match any files
match process_path(&self.ctx, path, &mut context_files, false, true).await {
match process_path(&self.ctx, path, &mut context_files, true).await {
Ok(_) => {}, // Path is valid
Err(e) => return Err(eyre!("Invalid path '{}': {}. Use --force to add anyway.", path, e)),
}
Expand Down Expand Up @@ -425,17 +437,15 @@ impl ContextManager {
/// 3. Reads the content of each file
/// 4. Returns a vector of (filename, content) pairs
///
/// # Arguments
/// * `force` - If true, include paths that don't exist yet
///
/// # Returns
/// A Result containing a vector of (filename, content) pairs or an error
pub async fn get_context_files(&self, force: bool) -> Result<Vec<(String, String)>> {
pub async fn get_context_files(&self) -> Result<Vec<(String, String)>> {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

also a note - from checking process_path it seems like it's never invoked with force == true so it would be nice to just remove that altogether

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also a major TODO here is to limit the total size of files loaded. This whole context loading logic loads everything into memory before dropping, so it's super trivial to just run into OOM with a glob that matches a ton of large files.

let mut context_files = Vec::new();

self.collect_context_files(&self.global_config.paths, &mut context_files, force)
self.collect_context_files(&self.global_config.paths, &mut context_files)
.await?;
self.collect_context_files(&self.profile_config.paths, &mut context_files, force)
self.collect_context_files(&self.profile_config.paths, &mut context_files)
.await?;

context_files.sort_by(|a, b| a.0.cmp(&b.0));
Expand All @@ -444,41 +454,49 @@ impl ContextManager {
Ok(context_files)
}

pub async fn get_context_files_by_path(&self, force: bool, path: &str) -> Result<Vec<(String, String)>> {
pub async fn get_context_files_by_path(&self, path: &str) -> Result<Vec<(String, String)>> {
let mut context_files = Vec::new();
process_path(&self.ctx, path, &mut context_files, force, true).await?;
process_path(&self.ctx, path, &mut context_files, true).await?;
Ok(context_files)
}

/// Get all context files from the global configuration.
pub async fn get_global_context_files(&self, force: bool) -> Result<Vec<(String, String)>> {
pub async fn get_global_context_files(&self) -> Result<Vec<(String, String)>> {
let mut context_files = Vec::new();

self.collect_context_files(&self.global_config.paths, &mut context_files, force)
self.collect_context_files(&self.global_config.paths, &mut context_files)
.await?;

Ok(context_files)
}

/// Get all context files from the current profile configuration.
pub async fn get_current_profile_context_files(&self, force: bool) -> Result<Vec<(String, String)>> {
pub async fn get_current_profile_context_files(&self) -> Result<Vec<(String, String)>> {
let mut context_files = Vec::new();

self.collect_context_files(&self.profile_config.paths, &mut context_files, force)
self.collect_context_files(&self.profile_config.paths, &mut context_files)
.await?;

Ok(context_files)
}

async fn collect_context_files(
&self,
paths: &[String],
context_files: &mut Vec<(String, String)>,
force: bool,
) -> Result<()> {
/// Collects context files and optionally drops files if the total size exceeds the limit.
/// Returns (files_to_use, dropped_files)
pub async fn collect_context_files_with_limit(&self) -> Result<(Vec<(String, String)>, Vec<(String, String)>)> {
let mut files = self.get_context_files().await?;

let dropped_files = drop_matched_context_files(&mut files, self.max_context_files_size).unwrap_or_default();

// remove dropped files from files
files.retain(|file| !dropped_files.iter().any(|dropped| dropped.0 == file.0));

Ok((files, dropped_files))
}

async fn collect_context_files(&self, paths: &[String], context_files: &mut Vec<(String, String)>) -> Result<()> {
for path in paths {
// Use is_validation=false to handle non-matching globs gracefully
process_path(&self.ctx, path, context_files, force, false).await?;
process_path(&self.ctx, path, context_files, false).await?;
}
Ok(())
}
Expand Down Expand Up @@ -647,7 +665,6 @@ async fn load_profile_config(ctx: &Context, profile_name: &str) -> Result<Contex
/// # Arguments
/// * `path` - The path to process
/// * `context_files` - The collection to add files to
/// * `force` - If true, include paths that don't exist yet
/// * `is_validation` - If true, error when glob patterns don't match; if false, silently skip
///
/// # Returns
Expand All @@ -656,7 +673,6 @@ async fn process_path(
ctx: &Context,
path: &str,
context_files: &mut Vec<(String, String)>,
force: bool,
is_validation: bool,
) -> Result<()> {
// Expand ~ to home directory
Expand Down Expand Up @@ -703,7 +719,7 @@ async fn process_path(
}
}

if !found_any && !force && is_validation {
if !found_any && is_validation {
// When validating paths (e.g., for /context add), error if no files match
return Err(eyre!("No files found matching glob pattern '{}'", full_path));
}
Expand All @@ -728,16 +744,10 @@ async fn process_path(
}
}
}
} else if !force && is_validation {
} else if is_validation {
// When validating paths (e.g., for /context add), error if the path doesn't exist
return Err(eyre!("Path '{}' does not exist", full_path));
} else if force {
// When using --force, we'll add the path even though it doesn't exist
// This allows users to add paths that will exist in the future
context_files.push((full_path.clone(), format!("(Path '{}' does not exist yet)", full_path)));
}
// When just showing expanded files (e.g., for /context show --expand),
// silently skip non-existent paths if is_validation is false
}

Ok(())
Expand Down Expand Up @@ -796,9 +806,10 @@ mod tests {
use super::*;

// Helper function to create a test ContextManager with Context
pub async fn create_test_context_manager() -> Result<ContextManager> {
pub async fn create_test_context_manager(context_file_size: Option<usize>) -> Result<ContextManager> {
let context_file_size = context_file_size.unwrap_or(CONTEXT_FILES_MAX_SIZE);
let ctx = Context::builder().with_test_home().await.unwrap().build_fake();
let manager = ContextManager::new(ctx).await?;
let manager = ContextManager::new(ctx, Some(context_file_size)).await?;
Ok(manager)
}

Expand All @@ -823,7 +834,7 @@ mod tests {

#[tokio::test]
async fn test_profile_ops() -> Result<()> {
let mut manager = create_test_context_manager().await?;
let mut manager = create_test_context_manager(None).await?;
let ctx = Arc::clone(&manager.ctx);

assert_eq!(manager.current_profile, "default");
Expand Down Expand Up @@ -860,28 +871,43 @@ mod tests {
Ok(())
}

#[tokio::test]
async fn test_collect_exceeds_limit() -> Result<()> {
let mut manager = create_test_context_manager(Some(2)).await?;
let ctx: Arc<Context> = Arc::clone(&manager.ctx);

ctx.fs().create_dir_all("test").await?;
ctx.fs().write("test/to-include.md", "ha").await?;
ctx.fs()
.write("test/to-drop.md", "long content that exceed limit")
.await?;
manager.add_paths(vec!["test/*.md".to_string()], false, false).await?;

let (used, dropped) = manager.collect_context_files_with_limit().await.unwrap();

assert!(used.len() + dropped.len() == 2);
assert!(used.len() == 1);
assert!(dropped.len() == 1);
Ok(())
}

#[tokio::test]
async fn test_path_ops() -> Result<()> {
let mut manager = create_test_context_manager().await?;
let ctx = Arc::clone(&manager.ctx);
let mut manager = create_test_context_manager(None).await?;
let ctx: Arc<Context> = Arc::clone(&manager.ctx);

// Create some test files for matching.
ctx.fs().create_dir_all("test").await?;
ctx.fs().write("test/p1.md", "p1").await?;
ctx.fs().write("test/p2.md", "p2").await?;

assert!(
manager.get_context_files(false).await?.is_empty(),
manager.get_context_files().await?.is_empty(),
"no files should be returned for an empty profile when force is false"
);
assert_eq!(
manager.get_context_files(true).await?.len(),
2,
"default non-glob global files should be included when force is true"
);

manager.add_paths(vec!["test/*.md".to_string()], false, false).await?;
let files = manager.get_context_files(false).await?;
let files = manager.get_context_files().await?;
assert!(files[0].0.ends_with("p1.md"));
assert_eq!(files[0].1, "p1");
assert!(files[1].0.ends_with("p2.md"));
Expand All @@ -900,7 +926,7 @@ mod tests {

#[tokio::test]
async fn test_add_hook() -> Result<()> {
let mut manager = create_test_context_manager().await?;
let mut manager = create_test_context_manager(None).await?;
let hook = Hook::new_inline_hook(HookTrigger::ConversationStart, "echo test".to_string());

// Test adding hook to profile config
Expand All @@ -919,7 +945,7 @@ mod tests {

#[tokio::test]
async fn test_remove_hook() -> Result<()> {
let mut manager = create_test_context_manager().await?;
let mut manager = create_test_context_manager(None).await?;
let hook = Hook::new_inline_hook(HookTrigger::ConversationStart, "echo test".to_string());

manager.add_hook("test_hook".to_string(), hook, false).await?;
Expand All @@ -936,7 +962,7 @@ mod tests {

#[tokio::test]
async fn test_set_hook_disabled() -> Result<()> {
let mut manager = create_test_context_manager().await?;
let mut manager = create_test_context_manager(None).await?;
let hook = Hook::new_inline_hook(HookTrigger::ConversationStart, "echo test".to_string());

manager.add_hook("test_hook".to_string(), hook, false).await?;
Expand All @@ -957,7 +983,7 @@ mod tests {

#[tokio::test]
async fn test_set_all_hooks_disabled() -> Result<()> {
let mut manager = create_test_context_manager().await?;
let mut manager = create_test_context_manager(None).await?;
let hook1 = Hook::new_inline_hook(HookTrigger::ConversationStart, "echo test".to_string());
let hook2 = Hook::new_inline_hook(HookTrigger::ConversationStart, "echo test".to_string());

Expand All @@ -977,7 +1003,7 @@ mod tests {

#[tokio::test]
async fn test_run_hooks() -> Result<()> {
let mut manager = create_test_context_manager().await?;
let mut manager = create_test_context_manager(None).await?;
let hook1 = Hook::new_inline_hook(HookTrigger::ConversationStart, "echo test".to_string());
let hook2 = Hook::new_inline_hook(HookTrigger::ConversationStart, "echo test".to_string());

Expand All @@ -993,7 +1019,7 @@ mod tests {

#[tokio::test]
async fn test_hooks_across_profiles() -> Result<()> {
let mut manager = create_test_context_manager().await?;
let mut manager = create_test_context_manager(None).await?;
let hook1 = Hook::new_inline_hook(HookTrigger::ConversationStart, "echo test".to_string());
let hook2 = Hook::new_inline_hook(HookTrigger::ConversationStart, "echo test".to_string());

Expand Down
Loading
Loading