diff --git a/src/crates/adapters/ai-adapters/src/client.rs b/src/crates/adapters/ai-adapters/src/client.rs index 1d56482420..a072ce8669 100644 --- a/src/crates/adapters/ai-adapters/src/client.rs +++ b/src/crates/adapters/ai-adapters/src/client.rs @@ -114,6 +114,18 @@ impl AIClient { } } + /// Clone this client with a different max output token limit while + /// reusing the HTTP client. + pub fn with_max_tokens(&self, max_tokens: Option) -> Self { + let mut config = self.config.clone(); + config.max_tokens = max_tokens; + Self { + client: self.client.clone(), + config, + stream_options: self.stream_options.clone(), + } + } + pub async fn send_message_stream( &self, messages: Vec, @@ -1264,6 +1276,17 @@ mod tests { assert!(request_body.get("reasoning").is_none()); } + #[test] + fn with_max_tokens_overrides_output_limit() { + let client = make_test_client("responses", None); + + let overridden = client.with_max_tokens(Some(2048)); + + assert_eq!(client.config.max_tokens, Some(8192)); + assert_eq!(overridden.config.max_tokens, Some(2048)); + assert_eq!(overridden.config.model, client.config.model); + } + #[test] fn build_anthropic_request_body_trim_mode_preserves_essential_fields() { let mut client = make_trim_test_client("anthropic"); diff --git a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs index 8850e270a4..7271c5688d 100644 --- a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs @@ -258,6 +258,7 @@ pub struct ExecutionEngine { } impl ExecutionEngine { + const COMPRESSION_MAX_TOKENS: u32 = 8192; const FINALIZE_AFTER_REPEATED_TOOL_FAILURES_REMINDER: &'static str = "This turn must end now because repeated tool failures have prevented further progress. Ignore any unfinished work. Your task now is to give the user a final answer. Do not call any more tools; any tool call will fail. Respond in plain text only. Summarize what was completed, what failed, the evidence available from the tool results, and the single best next step for the user."; const FINALIZE_AFTER_MAX_ROUNDS_REMINDER: &'static str = "This turn must end now because it has reached the round limit. Ignore any unfinished work. Your task now is to give the user a final answer. Do not call any more tools; any tool call will fail. Respond in plain text only. Summarize the most useful completed work and evidence collected so far, and clearly distinguish resolved items from anything still unresolved."; const FINALIZE_TOOL_DENIED_MESSAGE: &'static str = @@ -381,6 +382,19 @@ impl ExecutionEngine { !lower.contains("cancelled") } + fn compression_request_max_tokens(configured_max_tokens: Option, cap: u32) -> Option { + Some(configured_max_tokens.unwrap_or(cap).min(cap)) + } + + fn build_compression_ai_client( + ai_client: &crate::infrastructure::ai::AIClient, + ) -> crate::infrastructure::ai::AIClient { + ai_client.with_max_tokens(Self::compression_request_max_tokens( + ai_client.config.max_tokens, + Self::COMPRESSION_MAX_TOKENS, + )) + } + /// Detect periodic tool-signature loops in the trailing window. /// /// Returns `true` when the last `2 * threshold` rounds contain at most @@ -1331,7 +1345,6 @@ impl ExecutionEngine { provider: &str, attach_images: bool, prepended_prompt_reminders: &PrependedPromptReminders, - contract: Option<&crate::agentic::core::CompressionContract>, ) -> BitFunResult> { let prepended_reminders = prepended_prompt_reminders.ordered_reminders(); let mut compression_messages = Self::build_ai_messages_for_send( @@ -1344,7 +1357,7 @@ impl ExecutionEngine { ) .await?; compression_messages.push(AIMessage::user( - self.context_compressor.build_compact_prompt(contract), + self.context_compressor.build_compact_prompt(), )); Ok(compression_messages) } @@ -1359,9 +1372,10 @@ impl ExecutionEngine { ) -> BitFunResult { let mut last_error = None; let base_wait_time_ms = 500; + let compression_ai_client = Arc::new(Self::build_compression_ai_client(ai_client.as_ref())); for attempt in 0..max_tries { - let result = ai_client + let result = compression_ai_client .send_message_with_trace( request_messages.clone(), tool_definitions.clone(), @@ -1426,7 +1440,6 @@ impl ExecutionEngine { tool_definitions: &Option>, prepended_prompt_reminders: &PrependedPromptReminders, primary_supports_image_understanding: bool, - contract: Option<&crate::agentic::core::CompressionContract>, trace_config: Option, ) -> BitFunResult> { let request_messages = self @@ -1437,7 +1450,6 @@ impl ExecutionEngine { &ai_client.config.format, primary_supports_image_understanding, prepended_prompt_reminders, - contract, ) .await?; @@ -1721,7 +1733,6 @@ impl ExecutionEngine { tool_definitions, prepended_prompt_reminders, primary_supports_image_understanding, - compression_contract.as_ref(), trace_config, ) .await @@ -1947,7 +1958,6 @@ impl ExecutionEngine { &scaffold.tool_definitions, &scaffold.prepended_prompt_reminders, scaffold.primary_supports_image_understanding, - compression_contract.as_ref(), trace_config, ) .await @@ -3383,6 +3393,7 @@ mod tests { use crate::agentic::tools::ToolRuntimeRestrictions; use crate::service::config::types::AIConfig; use crate::service::config::types::AIModelConfig; + use crate::util::types::config as ai_config_types; use crate::util::types::ToolDefinition; use serde_json::json; use sha2::{Digest, Sha256}; @@ -3496,6 +3507,60 @@ mod tests { assert_eq!(summary, args); } + #[test] + fn compression_request_max_tokens_clamps_to_global_cap() { + assert_eq!( + ExecutionEngine::compression_request_max_tokens(Some(16_000), 8192), + Some(8192) + ); + assert_eq!( + ExecutionEngine::compression_request_max_tokens(Some(4096), 8192), + Some(4096) + ); + assert_eq!( + ExecutionEngine::compression_request_max_tokens(None, 8192), + Some(8192) + ); + } + + #[test] + fn build_compression_ai_client_overrides_only_max_tokens() { + let client = crate::infrastructure::ai::AIClient::new(ai_config_types::AIConfig { + name: "test".to_string(), + base_url: "https://example.com/v1".to_string(), + request_url: "https://example.com/v1/responses".to_string(), + api_key: "key".to_string(), + model: "test-model".to_string(), + format: "responses".to_string(), + context_window: 128_000, + max_tokens: None, + temperature: None, + top_p: None, + reasoning_mode: Default::default(), + inline_think_in_text: true, + custom_headers: None, + custom_headers_mode: None, + skip_ssl_verify: false, + reasoning_effort: None, + thinking_budget_tokens: None, + custom_request_body: None, + custom_request_body_mode: None, + }); + + let compression_client = ExecutionEngine::build_compression_ai_client(&client); + + assert_eq!(client.config.max_tokens, None); + assert_eq!( + compression_client.config.max_tokens, + Some(ExecutionEngine::COMPRESSION_MAX_TOKENS) + ); + assert_eq!(compression_client.config.model, client.config.model); + assert_eq!( + compression_client.config.request_url, + client.config.request_url + ); + } + #[test] fn partial_continuation_allowed_for_stream_stall_reasons() { assert!(ExecutionEngine::should_continue_after_partial_response( diff --git a/src/crates/assembly/core/src/agentic/session/compression/compressor.rs b/src/crates/assembly/core/src/agentic/session/compression/compressor.rs index ea32f87a61..f0a5159950 100644 --- a/src/crates/assembly/core/src/agentic/session/compression/compressor.rs +++ b/src/crates/assembly/core/src/agentic/session/compression/compressor.rs @@ -409,45 +409,28 @@ impl ContextCompressor { Some(trimmed.to_string()) } - pub(crate) fn build_compact_prompt(&self, contract: Option<&CompressionContract>) -> String { - let contract_instruction = contract - .filter(|contract| !contract.is_empty()) - .map(|contract| { - format!( - "\n\nThe following compaction contract is authoritative factual context from tool observations. Preserve every field from it in the final :\n{}\n", - contract.render_for_model() - ) - }) - .unwrap_or_default(); - + pub(crate) fn build_compact_prompt(&self) -> String { format!( r#"Your current task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions. This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing development work without losing context. -{contract_instruction} CRITICAL: Respond with TEXT ONLY. Do NOT call any tools. - Do NOT use Read, Bash, Grep, Glob, Edit, Write, or ANY other tool. - You already have all the context you need in the conversation above. - Tool calls will be REJECTED and will waste your only turn — you will fail the task. -- Your entire response must be plain text: an block followed by a block. - -Before providing your final summary, wrap your analysis in tags to organize your thoughts and ensure you've covered all necessary points. Then output the final retained summary in tags. -Important: only the content inside will be kept as compressed history. The section is transient and will be discarded, so do not put any required final information only in . -In your analysis process: - -1. Chronologically analyze each message and section of the conversation. For each section thoroughly identify: - - The user's explicit requests and intents - - Your approach to addressing the user's requests - - Key decisions, technical concepts and code patterns - - Specific details like: - - file names - - full code snippets - - function signatures - - file edits - - Errors that you ran into and how you fixed them - - Pay special attention to specific user feedback that you received, especially if the user told you to do something differently. -2. Double-check for technical accuracy and completeness, addressing each required element thoroughly. +- Your entire response must be plain text inside a single block. + +Output exactly one ... block containing all retained context. +Important: only the content inside will be kept as compressed history, so include every required detail there. + +Before you answer, carefully review the conversation chronologically and make sure the final captures: +- The user's explicit requests and intents +- Your approach to addressing the user's requests +- Key decisions, technical concepts, and code patterns +- Specific details like file names, function signatures, file edits, and important code snippets where they materially matter +- Errors that you ran into and how you fixed them +- Specific user feedback, especially when the user asked for a different approach or corrected direction Your summary should include the following sections: @@ -464,10 +447,6 @@ Your summary should include the following sections: Here's an example of how your output should be structured: - -[Your thought process, ensuring all points are covered thoroughly and accurately] - - 1. Primary Request and Intent: [Detailed description] @@ -514,7 +493,7 @@ Here's an example of how your output should be structured: Please provide your summary based on the conversation so far, following this structure and ensuring precision and thoroughness in your response. -REMINDER: Do NOT call any tools. Respond with plain text only — an block followed by a block. Tool calls will be rejected and you will fail the task. +REMINDER: Do NOT call any tools. Respond with plain text only inside a single block. Tool calls will be rejected and you will fail the task. "# ) } @@ -533,8 +512,7 @@ fn extract_tag_content<'a>(text: &'a str, tag: &str) -> Option<&'a str> { mod tests { use super::{CompressionMode, ContextCompressor, TurnWithTokens}; use crate::agentic::core::{ - render_system_reminder, CompressionContract, CompressionContractItem, CompressionEntry, - CompressionPayload, Message, MessageSemanticKind, + render_system_reminder, CompressionEntry, CompressionPayload, Message, MessageSemanticKind, }; fn make_turn(messages: Vec) -> TurnWithTokens { @@ -652,25 +630,24 @@ mod tests { } #[test] - fn model_summary_prompt_includes_compaction_contract() { + fn model_summary_prompt_does_not_inline_compaction_contract() { + let compressor = ContextCompressor::new(Default::default()); + + let prompt = compressor.build_compact_prompt(); + + assert!(!prompt.contains("authoritative factual context")); + assert!(!prompt.contains("src/lib.rs")); + assert!(!prompt.contains("cargo test")); + } + + #[test] + fn model_summary_prompt_requires_summary_only() { let compressor = ContextCompressor::new(Default::default()); - let contract = CompressionContract { - touched_files: vec!["src/lib.rs".to_string()], - verification_commands: vec![CompressionContractItem { - target: "cargo test".to_string(), - status: "succeeded".to_string(), - summary: "Tests passed.".to_string(), - error_kind: None, - }], - blocking_failures: Vec::new(), - subagent_statuses: Vec::new(), - }; - let prompt = compressor.build_compact_prompt(Some(&contract)); + let prompt = compressor.build_compact_prompt(); - assert!(prompt.contains("authoritative factual context")); - assert!(prompt.contains("src/lib.rs")); - assert!(prompt.contains("cargo test")); + assert!(prompt.contains("single block")); + assert!(!prompt.contains(" block followed by a block")); } #[test] diff --git a/src/crates/assembly/core/src/agentic/session/compression/fallback/tests.rs b/src/crates/assembly/core/src/agentic/session/compression/fallback/tests.rs index d67ffde4d8..608912dd1e 100644 --- a/src/crates/assembly/core/src/agentic/session/compression/fallback/tests.rs +++ b/src/crates/assembly/core/src/agentic/session/compression/fallback/tests.rs @@ -237,7 +237,7 @@ fn renders_contract_facts_even_when_tool_results_are_cleared() { assert!(summary_artifact .summary_text - .contains("Compaction contract:")); + .contains("The following facts were retained during compression.")); assert!(summary_artifact.summary_text.contains("src/main.rs")); assert!(summary_artifact.summary_text.contains("cargo test")); assert!(summary_artifact diff --git a/src/crates/contracts/runtime-ports/src/lib.rs b/src/crates/contracts/runtime-ports/src/lib.rs index 85a62e34d6..a21ccb87c8 100644 --- a/src/crates/contracts/runtime-ports/src/lib.rs +++ b/src/crates/contracts/runtime-ports/src/lib.rs @@ -1086,7 +1086,7 @@ impl CompressionContract { pub fn render_for_model(&self) -> String { let mut lines = vec![ - "Compaction contract: preserve these factual fields when continuing the task." + "The following facts were retained during compression. Use them as authoritative context when continuing the task." .to_string(), ]; @@ -1833,7 +1833,7 @@ mod tests { let rendered = contract.render_for_model(); - assert!(rendered.contains("Compaction contract")); + assert!(rendered.contains("The following facts were retained during compression.")); assert!(rendered.contains("Touched files:")); assert!(rendered.contains("- src/lib.rs")); assert!(rendered.contains(