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
23 changes: 23 additions & 0 deletions src/crates/adapters/ai-adapters/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<u32>) -> 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<Message>,
Expand Down Expand Up @@ -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");
Expand Down
79 changes: 72 additions & 7 deletions src/crates/assembly/core/src/agentic/execution/execution_engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down Expand Up @@ -381,6 +382,19 @@ impl ExecutionEngine {
!lower.contains("cancelled")
}

fn compression_request_max_tokens(configured_max_tokens: Option<u32>, cap: u32) -> Option<u32> {
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
Expand Down Expand Up @@ -1331,7 +1345,6 @@ impl ExecutionEngine {
provider: &str,
attach_images: bool,
prepended_prompt_reminders: &PrependedPromptReminders,
contract: Option<&crate::agentic::core::CompressionContract>,
) -> BitFunResult<Vec<AIMessage>> {
let prepended_reminders = prepended_prompt_reminders.ordered_reminders();
let mut compression_messages = Self::build_ai_messages_for_send(
Expand All @@ -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)
}
Expand All @@ -1359,9 +1372,10 @@ impl ExecutionEngine {
) -> BitFunResult<String> {
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(),
Expand Down Expand Up @@ -1426,7 +1440,6 @@ impl ExecutionEngine {
tool_definitions: &Option<Vec<ToolDefinition>>,
prepended_prompt_reminders: &PrependedPromptReminders,
primary_supports_image_understanding: bool,
contract: Option<&crate::agentic::core::CompressionContract>,
trace_config: Option<ModelExchangeTraceConfig>,
) -> BitFunResult<Option<String>> {
let request_messages = self
Expand All @@ -1437,7 +1450,6 @@ impl ExecutionEngine {
&ai_client.config.format,
primary_supports_image_understanding,
prepended_prompt_reminders,
contract,
)
.await?;

Expand Down Expand Up @@ -1721,7 +1733,6 @@ impl ExecutionEngine {
tool_definitions,
prepended_prompt_reminders,
primary_supports_image_understanding,
compression_contract.as_ref(),
trace_config,
)
.await
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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};
Expand Down Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 <summary>:\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 <analysis> block followed by a <summary> block.

Before providing your final summary, wrap your analysis in <analysis> tags to organize your thoughts and ensure you've covered all necessary points. Then output the final retained summary in <summary> tags.
Important: only the content inside <summary> will be kept as compressed history. The <analysis> section is transient and will be discarded, so do not put any required final information only in <analysis>.
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 <summary> block.

Output exactly one <summary>...</summary> block containing all retained context.
Important: only the content inside <summary> 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 <summary> 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:

Expand All @@ -464,10 +447,6 @@ Your summary should include the following sections:
Here's an example of how your output should be structured:

<example>
<analysis>
[Your thought process, ensuring all points are covered thoroughly and accurately]
</analysis>

<summary>
1. Primary Request and Intent:
[Detailed description]
Expand Down Expand Up @@ -514,7 +493,7 @@ Here's an example of how your output should be structured:
</example>

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 <analysis> block followed by a <summary> 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 <summary> block. Tool calls will be rejected and you will fail the task.
"#
)
}
Expand All @@ -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<Message>) -> TurnWithTokens {
Expand Down Expand Up @@ -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 <summary> block"));
assert!(!prompt.contains("<analysis> block followed by a <summary> block"));
}

#[test]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions src/crates/contracts/runtime-ports/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
];

Expand Down Expand Up @@ -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(
Expand Down
Loading