Merge pull request #520 from stakpak/fix/lazy-context-trimming

Add lazy context trimming

Author: kajogo777

AGENTS.md

@@ -73,6 +73,7 @@ Vec<ChatMessage>                    # OpenAI-shaped messages (cli/mode_interacti
 ContextManager::reduce_context()   # History reduction (libs/api/context_managers/)
     ↓  merge_consecutive_same_role()  # Merge tool messages
     ↓  dedup_tool_results()           # Deduplicate within merged messages
+    ↓  reduce_context_with_budget()   # Budget-aware trimming (if over threshold)
     ↓
 Vec<LLMMessage>                    # Provider-neutral messages
     ↓
@@ -177,6 +178,22 @@ The codebase uses **three layers** to prevent invalid message sequences:
 2. **Pre-API sanitization** (`sanitize_tool_results`): Dedup and remove orphans from `Vec<ChatMessage>` before every API call
 3. **Context manager** (`task_board_context_manager.rs`): Merge consecutive same-role messages and dedup tool_results in the `reduce_context()` pipeline
 
+### Context Trimming with Cache Preservation
+
+Long sessions accumulate messages that approach the context window limit. The `TaskBoardContextManager` implements budget-aware trimming:
+
+1. **Lazy trimming**: Only triggers when estimated tokens exceed `context_window × threshold` (default 80%)
+2. **Stable prefix**: Trimmed messages are replaced with `[trimmed]` placeholders, preserving message structure (roles, tool_call_ids) for API validity
+3. **Cache-friendly**: The trimmed prefix produces identical output across turns, so Anthropic's prompt cache stays valid
+4. **Metadata persistence**: Trimming state (`trimmed_up_to_message_index`) is stored in `CheckpointState.metadata` and flows through:
+   - `CheckpointState.metadata` → `AgentState.metadata` → Hook updates → `save_checkpoint()` → persisted
+
+Key files:
+- `libs/api/src/local/context_managers/task_board_context_manager.rs` — `reduce_context_with_budget()`, `estimate_tokens()`, `trim_message()`
+- `libs/api/src/local/hooks/task_board_context/mod.rs` — Wires budget-aware trimming into the hook lifecycle
+- `libs/api/src/storage.rs` — `CheckpointState.metadata` field
+- `libs/api/src/models.rs` — `AgentState.metadata` field
+
 ## Build & Test
 
 ```bash

cli/src/commands/acp/server.rs

@@ -1844,7 +1844,14 @@ impl acp::Agent for StakpakAcpAgent {
         let model = self.model.read().await.clone();
         let session_id = self.current_session_id.get();
         let (stream, _request_id) = client
-            .chat_completion_stream(model, messages, tools_option.clone(), None, session_id)
+            .chat_completion_stream(
+                model,
+                messages,
+                tools_option.clone(),
+                None,
+                session_id,
+                None,
+            )
             .await
             .map_err(|e| {
                 log::error!("Chat completion stream failed: {e}");
@@ -2032,6 +2039,7 @@ impl acp::Agent for StakpakAcpAgent {
                         tools_option.clone(),
                         None,
                         session_id,
+                        None,
                     )
                     .await
                     .map_err(|e| {

cli/src/commands/agent/run/checkpoint.rs

@@ -11,7 +11,7 @@ use uuid::Uuid;
 pub async fn get_checkpoint_messages(
     client: &dyn AgentProvider,
     checkpoint_id: &str,
-) -> Result<Vec<ChatMessage>, String> {
+) -> Result<(Vec<ChatMessage>, Option<serde_json::Value>), String> {
     let checkpoint_uuid = Uuid::parse_str(checkpoint_id).map_err(|_| {
         format!(
             "Invalid checkpoint ID '{}' - must be a valid UUID",
@@ -24,7 +24,7 @@ pub async fn get_checkpoint_messages(
         .await
         .map_err(|e| e.to_string())?;
 
-    Ok(checkpoint.state.messages)
+    Ok((checkpoint.state.messages, checkpoint.state.metadata))
 }
 
 pub async fn extract_checkpoint_messages_and_tool_calls(
@@ -178,19 +178,28 @@ pub async fn resume_session_from_checkpoint(
     client: &dyn AgentProvider,
     session_id: &str,
     input_tx: &tokio::sync::mpsc::Sender<InputEvent>,
-) -> Result<(Vec<ChatMessage>, Vec<ToolCall>, Uuid), String> {
+) -> Result<
+    (
+        Vec<ChatMessage>,
+        Vec<ToolCall>,
+        Uuid,
+        Option<serde_json::Value>,
+    ),
+    String,
+> {
     let session_uuid = Uuid::parse_str(session_id).map_err(|e| e.to_string())?;
 
     match client.get_active_checkpoint(session_uuid).await {
         Ok(checkpoint) => {
+            let metadata = checkpoint.state.metadata.clone();
             let (chat_messages, tool_calls) = extract_checkpoint_messages_and_tool_calls(
                 &checkpoint.id.to_string(),
                 input_tx,
                 checkpoint.state.messages,
             )
             .await?;
 
-            Ok((chat_messages, tool_calls, checkpoint.session_id))
+            Ok((chat_messages, tool_calls, checkpoint.session_id, metadata))
         }
         Err(e) => {
             send_input_event(

cli/src/commands/agent/run/mode_async.rs

@@ -235,6 +235,7 @@ pub async fn run_async(ctx: AppConfig, config: RunAsyncConfig) -> Result<AsyncOu
 
     let mut current_session_id: Option<Uuid> = None;
     let mut current_checkpoint_id: Option<Uuid> = None;
+    let mut current_metadata: Option<serde_json::Value> = None;
     let mut prior_steps: usize = 0;
 
     // Load checkpoint/session messages if provided
@@ -250,6 +251,7 @@ pub async fn run_async(ctx: AppConfig, config: RunAsyncConfig) -> Result<AsyncOu
 
         current_session_id = Some(checkpoint.session_id);
         current_checkpoint_id = Some(checkpoint.id);
+        current_metadata = checkpoint.state.metadata;
         chat_messages.extend(checkpoint.state.messages);
 
         llm_response_time += checkpoint_start.elapsed();
@@ -269,6 +271,7 @@ pub async fn run_async(ctx: AppConfig, config: RunAsyncConfig) -> Result<AsyncOu
             Ok(checkpoint) => {
                 current_session_id = Some(checkpoint.session_id);
                 current_checkpoint_id = Some(checkpoint_uuid);
+                current_metadata = checkpoint.state.metadata;
                 prior_steps = checkpoint
                     .state
                     .messages
@@ -452,6 +455,7 @@ pub async fn run_async(ctx: AppConfig, config: RunAsyncConfig) -> Result<AsyncOu
                 chat_messages.clone(),
                 Some(tools.clone()),
                 current_session_id,
+                current_metadata.clone(),
             )
             .await
             .map_err(|e| e.to_string())?;

cli/src/commands/agent/run/mode_interactive.rs

@@ -238,6 +238,7 @@ pub async fn run_interactive(
         let ctx_clone = ctx.clone(); // Clone ctx for use in client task
         let client_handle: tokio::task::JoinHandle<ClientTaskResult> = tokio::spawn(async move {
             let mut current_session_id: Option<Uuid> = None;
+            let mut current_metadata: Option<serde_json::Value> = None;
 
             // Build unified AgentClient config
             let providers = ctx_clone.get_llm_provider_config();
@@ -324,11 +325,12 @@ pub async fn run_interactive(
             }
 
             if let Some(session_id_str) = session_id {
-                let (chat_messages, tool_calls, session_id_uuid) =
+                let (chat_messages, tool_calls, session_id_uuid, checkpoint_metadata) =
                     resume_session_from_checkpoint(client.as_ref(), &session_id_str, &input_tx)
                         .await?;
 
                 current_session_id = Some(session_id_uuid);
+                current_metadata = checkpoint_metadata;
                 should_update_rulebooks_on_next_message = true;
                 tools_queue.extend(tool_calls.clone());
 
@@ -354,7 +356,7 @@ pub async fn run_interactive(
                     current_session_id = Some(checkpoint.session_id);
                 }
 
-                let checkpoint_messages =
+                let (checkpoint_messages, _checkpoint_metadata) =
                     get_checkpoint_messages(client.as_ref(), &checkpoint_id_str).await?;
 
                 let (chat_messages, tool_calls) = extract_checkpoint_messages_and_tool_calls(
@@ -688,9 +690,15 @@ pub async fn run_interactive(
                             )
                             .await
                             {
-                                Ok((chat_messages, tool_calls, session_id_uuid)) => {
+                                Ok((
+                                    chat_messages,
+                                    tool_calls,
+                                    session_id_uuid,
+                                    checkpoint_metadata,
+                                )) => {
                                     // Track the current session ID
                                     current_session_id = Some(session_id_uuid);
+                                    current_metadata = checkpoint_metadata;
 
                                     // Mark that we need to update rulebooks on the next user message
                                     should_update_rulebooks_on_next_message = true;
@@ -757,9 +765,15 @@ pub async fn run_interactive(
                         )
                         .await
                         {
-                            Ok((chat_messages, tool_calls, session_id_uuid)) => {
+                            Ok((
+                                chat_messages,
+                                tool_calls,
+                                session_id_uuid,
+                                checkpoint_metadata,
+                            )) => {
                                 // Track the current session ID
                                 current_session_id = Some(session_id_uuid);
+                                current_metadata = checkpoint_metadata;
 
                                 // Mark that we need to update rulebooks on the next user message
                                 should_update_rulebooks_on_next_message = true;
@@ -992,6 +1006,7 @@ pub async fn run_interactive(
                             Some(tools.clone()),
                             headers.clone(),
                             current_session_id,
+                            current_metadata.clone(),
                         )
                         .await;
 

libs/api/src/client/mod.rs

@@ -230,9 +230,8 @@ impl AgentClient {
         hook_registry.register(
             LifecycleEvent::BeforeInference,
             Box::new(TaskBoardContextHook::new(TaskBoardContextHookOptions {
-                history_action_message_size_limit: Some(100),
-                history_action_message_keep_last_n: Some(50),
-                history_action_result_keep_last_n: Some(50),
+                keep_last_n_assistant_messages: Some(10),
+                context_budget_threshold: Some(0.8), // defaults to 0.8 (80%)
             })),
         );
         let hook_registry = Arc::new(hook_registry);

libs/api/src/client/provider.rs

@@ -181,8 +181,10 @@ impl AgentProvider for AgentClient {
         messages: Vec<ChatMessage>,
         tools: Option<Vec<Tool>>,
         session_id: Option<Uuid>,
+        metadata: Option<serde_json::Value>,
     ) -> Result<ChatCompletionResponse, String> {
         let mut ctx = HookContext::new(session_id, AgentState::new(model, messages, tools));
+        ctx.state.metadata = metadata;
 
         // Execute before request hooks
         self.hook_registry
@@ -201,7 +203,11 @@ impl AgentProvider for AgentClient {
 
         // Save checkpoint
         let result = self
-            .save_checkpoint(&current_session, ctx.state.messages.clone())
+            .save_checkpoint(
+                &current_session,
+                ctx.state.messages.clone(),
+                ctx.state.metadata.clone(),
+            )
             .await?;
         let checkpoint_created_at = result.checkpoint_created_at.timestamp() as u64;
         ctx.set_new_checkpoint_id(result.checkpoint_id);
@@ -265,6 +271,7 @@ impl AgentProvider for AgentClient {
         tools: Option<Vec<Tool>>,
         _headers: Option<HeaderMap>,
         session_id: Option<Uuid>,
+        metadata: Option<serde_json::Value>,
     ) -> Result<
         (
             Pin<
@@ -275,6 +282,7 @@ impl AgentProvider for AgentClient {
         String,
     > {
         let mut ctx = HookContext::new(session_id, AgentState::new(model, messages, tools));
+        ctx.state.metadata = metadata;
 
         // Execute before request hooks
         self.hook_registry
@@ -332,7 +340,11 @@ impl AgentProvider for AgentClient {
                     }
 
                     let result = client
-                        .save_checkpoint(&current_session, ctx_clone.state.messages.clone())
+                        .save_checkpoint(
+                            &current_session,
+                            ctx_clone.state.messages.clone(),
+                            ctx_clone.state.metadata.clone(),
+                        )
                         .await;
 
                     match result {
@@ -784,10 +796,15 @@ impl AgentClient {
         &self,
         current: &SessionInfo,
         messages: Vec<ChatMessage>,
+        metadata: Option<serde_json::Value>,
     ) -> Result<SessionInfo, String> {
-        let checkpoint_request =
+        let mut checkpoint_request =
             StorageCreateCheckpointRequest::new(messages).with_parent(current.checkpoint_id);
 
+        if let Some(meta) = metadata {
+            checkpoint_request = checkpoint_request.with_metadata(meta);
+        }
+
         let checkpoint = self
             .session_storage
             .create_checkpoint(current.session_id, &checkpoint_request)

libs/api/src/lib.rs

@@ -147,6 +147,7 @@ pub trait AgentProvider: SessionStorage + Send + Sync {
         messages: Vec<ChatMessage>,
         tools: Option<Vec<Tool>>,
         session_id: Option<Uuid>,
+        metadata: Option<serde_json::Value>,
     ) -> Result<ChatCompletionResponse, String>;
     async fn chat_completion_stream(
         &self,
@@ -155,6 +156,7 @@ pub trait AgentProvider: SessionStorage + Send + Sync {
         tools: Option<Vec<Tool>>,
         headers: Option<HeaderMap>,
         session_id: Option<Uuid>,
+        metadata: Option<serde_json::Value>,
     ) -> Result<
         (
             std::pin::Pin<