// Copyright 2025 Muvon Un Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// Context reduction for session optimization
use super::animation::show_smart_animation;
use crate::config::Config;
use crate::mcp::get_available_functions;
use crate::session::chat::session::ChatSession;
use crate::session::estimate_full_context_tokens;
use anyhow::Result;
use colored::*;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
/// Process context reduction - smart truncation with summarization
/// Simply adds a summarization prompt and lets the normal session flow handle it
pub async fn perform_context_reduction(
chat_session: &mut ChatSession,
config: &Config,
role: &str,
operation_cancelled: tokio::sync::watch::Receiver<bool>,
) -> Result<()> {
println!("{}", "Finalizing current task...".cyan());
// Check if there's anything to summarize (exclude system message)
let conversation_messages = chat_session
.session
.messages
.iter()
.filter(|m| m.role != "system")
.count();
if conversation_messages == 0 {
println!("{}", "No conversation to summarize".yellow());
return Ok(());
}
// Store original message count for logging
let original_message_count = chat_session.session.messages.len();
// Enhanced summarization prompt that preserves complete task context
let summarization_prompt = "Please memorize all critical and important information for future reference first in parallel and actualize documentation if there is something that requires it, do all in parallel while maximizing tool calls efficiency, then create a comprehensive summary of our conversation that preserves:\n\n1. **Complete Task Overview**: What was the main task/feature we worked on? Include the original request and scope.\n2. **Files Modified**: List ALL files that were created, modified, or deleted with their FULL paths and purposes:\n - New files: [path] - purpose/description\n - Modified files: [path] - what changes were made\n - Deleted files: [path] - reason for deletion\n3. **Technical Decisions**: All architectural choices, patterns used, and implementation approaches\n4. **Key Code Changes**: Important functions, classes, or modules added/modified with specific names\n5. **Configuration Changes**: Any config files, dependencies, or environment changes with exact file paths\n6. **Testing & Validation**: What was tested and how (commands run, test files, validation steps)\n7. **Current State**: What is the current working state of the implementation\n8. **Next Steps**: What needs to be done to continue this work (specific tasks, files to modify)\n9. **Context for Continuation**: Essential information needed to pick up where we left off\n10. **File References**: Complete list of all relevant file paths that future sessions might need to access\n\nThis is a TASK COMPLETION summary - treat it like a git commit that finalizes the current work phase. Focus on actionable information, specific file paths, function names, and technical details that would be crucial for continuing this development work in future sessions. Include enough detail that someone could understand and continue the work without reading the full conversation history.";
chat_session.add_user_message(summarization_prompt)?;
// Create a separate flag for animation control to avoid conflicts with user cancellation detection
let animation_cancel = Arc::new(AtomicBool::new(false));
let animation_cancel_clone = animation_cancel.clone();
let current_cost = chat_session.session.info.total_cost;
let max_threshold = config.max_session_tokens_threshold;
// Calculate actual current context tokens for percentage display
let (_, _, _, _, system_prompt) = config.get_role_config(role);
let tools = get_available_functions(config).await;
let current_context_tokens = estimate_full_context_tokens(
&chat_session.session.messages,
Some(system_prompt),
Some(&tools),
) as u64;
let animation_task = tokio::spawn(async move {
let _ = show_smart_animation(
animation_cancel_clone,
current_cost,
current_context_tokens,
max_threshold,
)
.await;
});
// Use the same API flow as the normal session
let api_result = crate::session::chat_completion_with_provider(
crate::session::ChatCompletionProviderParams {
messages: &chat_session.session.messages,
model: &chat_session.model,
temperature: chat_session.temperature,
top_p: chat_session.top_p,
top_k: chat_session.top_k,
max_tokens: chat_session.max_tokens,
config,
max_retries: chat_session.max_retries, // Use max_retries from chat session
cancellation_token: None,
},
)
.await;
// Stop the animation using the separate flag (not the operation_cancelled flag)
animation_cancel.store(true, Ordering::SeqCst);
let _ = animation_task.await;
// Process the response with the normal flow (handles tool calls, etc.)
let response_result = match api_result {
Ok(response) => {
// Use the normal process_response flow which handles tool calls automatically
let process_result =
super::response::process_response(super::response::ResponseProcessingParams::new(
response.content.clone(),
response.exchange,
response.tool_calls,
response.finish_reason,
response.response_id,
chat_session,
config,
role, // Use the current role instead of hardcoding "developer"
operation_cancelled.clone(),
))
.await;
match process_result {
Ok(()) => Ok(response.content),
Err(e) => Err(e),
}
}
Err(e) => Err(e),
};
match response_result {
Ok(summary_content) => {
println!("{}", "Context summarization complete".bright_green());
// SMART TRUNCATION: Keep only system message + the LAST message (which is the assistant's summary)
let system_message = chat_session
.session
.messages
.iter()
.find(|m| m.role == "system")
.cloned();
// Get the LAST message (the assistant's summary that was just added)
let last_message = chat_session.session.messages.last().cloned();
// Clear all messages
chat_session.session.messages.clear();
// Restore system message
if let Some(system) = system_message {
chat_session.session.messages.push(system);
}
// Restore the LAST message (the assistant's summary)
if let Some(mut last) = last_message {
last.cached = true; // Mark for caching
chat_session.session.messages.push(last.clone());
// Log restoration point using the actual assistant message content
let _ = crate::session::logger::log_restoration_point(
&chat_session.session.info.name,
"/done - Task completion and context optimization",
&last.content,
);
}
// Log to session file as well
if let Some(session_file) = &chat_session.session.session_file {
// Get the actual assistant message content for logging (it's the last message now)
let actual_summary = if let Some(last_msg) = chat_session.session.messages.last() {
&last_msg.content
} else {
&summary_content
};
let restoration_data = serde_json::json!({
"type": "context_reduction",
"summary": actual_summary,
"original_message_count": original_message_count,
"timestamp": std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
});
let restoration_json = serde_json::to_string(&restoration_data)?;
let _ = crate::session::append_to_session_file(
session_file,
&format!("RESTORATION_POINT: {}", restoration_json),
);
}
// Reset token tracking for fresh start
chat_session.session.info.current_non_cached_tokens = 0;
chat_session.session.info.current_total_tokens = 0;
// Reset cache checkpoint time
chat_session.session.info.last_cache_checkpoint_time = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
println!(
"{}",
"Session context reduced to essential summary".bright_green()
);
println!(
"{}",
"You can now continue the conversation with optimized context".bright_cyan()
);
// Auto-commit with octocode if available
if let Err(e) = auto_commit_with_octocode().await {
// Don't fail the entire operation if commit fails, just warn
println!("{}: {}", "Warning: Auto-commit failed".bright_yellow(), e);
}
// Save the updated session
chat_session.save()?;
Ok(())
}
Err(e) => {
// Remove the summarization prompt since it failed
if let Some(last_msg) = chat_session.session.messages.last() {
if last_msg.role == "user"
&& last_msg.content.contains("Please create a concise summary")
{
chat_session.session.messages.pop();
}
}
println!(
"{}: {}",
"Error during context summarization".bright_red(),
e
);
Err(anyhow::anyhow!("Context summarization failed: {}", e))
}
}
}
/// Auto-commit changes using octocode if the binary is available
async fn auto_commit_with_octocode() -> Result<()> {
// Check if octocode binary is available in PATH
let octocode_check = tokio::process::Command::new("which")
.arg("octocode")
.output()
.await;
match octocode_check {
Ok(output) if output.status.success() => {
// octocode is available, proceed with commit
println!(
"{}",
"đ Auto-committing changes with octocode...".bright_blue()
);
let commit_result = tokio::process::Command::new("octocode")
.args(["commit", "-a", "-y"])
.output()
.await;
match commit_result {
Ok(output) => {
if output.status.success() {
let stdout = String::from_utf8_lossy(&output.stdout);
if !stdout.trim().is_empty() {
println!("{}", stdout.trim().bright_green());
}
println!(
"{}",
"â
Changes auto-committed successfully".bright_green()
);
} else {
let stderr = String::from_utf8_lossy(&output.stderr);
if stderr.contains("no changes") || stderr.contains("nothing to commit") {
println!("{}", "âšī¸ No changes to commit".bright_blue());
} else {
return Err(anyhow::anyhow!("octocode commit failed: {}", stderr));
}
}
}
Err(e) => {
return Err(anyhow::anyhow!("Failed to execute octocode commit: {}", e));
}
}
}
Ok(_) => {
// which command succeeded but octocode not found (empty output)
println!(
"{}",
"âšī¸ octocode not found in PATH, skipping auto-commit".bright_blue()
);
}
Err(_) => {
// which command failed (probably on Windows or which is not available)
// Try direct execution as fallback
let direct_check = tokio::process::Command::new("octocode")
.arg("--version")
.output()
.await;
match direct_check {
Ok(output) if output.status.success() => {
// octocode is available, proceed with commit
println!(
"{}",
"đ Auto-committing changes with octocode...".bright_blue()
);
let commit_result = tokio::process::Command::new("octocode")
.args(["commit", "-a", "-y"])
.output()
.await;
match commit_result {
Ok(output) => {
if output.status.success() {
let stdout = String::from_utf8_lossy(&output.stdout);
if !stdout.trim().is_empty() {
println!("{}", stdout.trim().bright_green());
}
println!(
"{}",
"â
Changes auto-committed successfully".bright_green()
);
} else {
let stderr = String::from_utf8_lossy(&output.stderr);
if stderr.contains("no changes")
|| stderr.contains("nothing to commit")
{
println!("{}", "âšī¸ No changes to commit".bright_blue());
} else {
return Err(anyhow::anyhow!(
"octocode commit failed: {}",
stderr
));
}
}
}
Err(e) => {
return Err(anyhow::anyhow!(
"Failed to execute octocode commit: {}",
e
));
}
}
}
_ => {
// octocode not available
println!(
"{}",
"âšī¸ octocode not available, skipping auto-commit".bright_blue()
);
}
}
}
}
Ok(())
}