Account for encrypted reasoning for auto compaction (#7113)

- The total token used returned from the api doesn't account for the
reasoning items before the assistant message
- Account for those for auto compaction
- Add the encrypted reasoning effort in the common tests utils
- Add a test to make sure it works as expected
This commit is contained in:
Ahmed Ibrahim
2025-11-21 19:06:45 -08:00
committed by GitHub
parent 529eb4ff2a
commit b519267d05
9 changed files with 236 additions and 30 deletions

View File

@@ -661,6 +661,11 @@ impl Session {
format!("auto-compact-{id}")
}
async fn get_total_token_usage(&self) -> i64 {
let state = self.state.lock().await;
state.get_total_token_usage()
}
async fn record_initial_history(&self, conversation_history: InitialHistory) {
let turn_context = self.new_turn(SessionSettingsUpdate::default()).await;
match conversation_history {
@@ -1958,20 +1963,13 @@ pub(crate) async fn run_task(
.await
{
Ok(turn_output) => {
let TurnRunResult {
processed_items,
total_token_usage,
} = turn_output;
let processed_items = turn_output;
let limit = turn_context
.client
.get_auto_compact_token_limit()
.unwrap_or(i64::MAX);
let total_usage_tokens = total_token_usage
.as_ref()
.map(TokenUsage::tokens_in_context_window);
let token_limit_reached = total_usage_tokens
.map(|tokens| tokens >= limit)
.unwrap_or(false);
let total_usage_tokens = sess.get_total_token_usage().await;
let token_limit_reached = total_usage_tokens >= limit;
let (responses, items_to_record_in_conversation_history) =
process_items(processed_items, &sess, &turn_context).await;
@@ -2028,7 +2026,7 @@ async fn run_turn(
turn_diff_tracker: SharedTurnDiffTracker,
input: Vec<ResponseItem>,
cancellation_token: CancellationToken,
) -> CodexResult<TurnRunResult> {
) -> CodexResult<Vec<ProcessedResponseItem>> {
let mcp_tools = sess
.services
.mcp_connection_manager
@@ -2159,12 +2157,6 @@ pub struct ProcessedResponseItem {
pub response: Option<ResponseInputItem>,
}
#[derive(Debug)]
struct TurnRunResult {
processed_items: Vec<ProcessedResponseItem>,
total_token_usage: Option<TokenUsage>,
}
#[allow(clippy::too_many_arguments)]
async fn try_run_turn(
router: Arc<ToolRouter>,
@@ -2173,7 +2165,7 @@ async fn try_run_turn(
turn_diff_tracker: SharedTurnDiffTracker,
prompt: &Prompt,
cancellation_token: CancellationToken,
) -> CodexResult<TurnRunResult> {
) -> CodexResult<Vec<ProcessedResponseItem>> {
let rollout_item = RolloutItem::TurnContext(TurnContextItem {
cwd: turn_context.cwd.clone(),
approval_policy: turn_context.approval_policy,
@@ -2335,12 +2327,7 @@ async fn try_run_turn(
sess.send_event(&turn_context, msg).await;
}
let result = TurnRunResult {
processed_items,
total_token_usage: token_usage.clone(),
};
return Ok(result);
return Ok(processed_items);
}
ResponseEvent::OutputTextDelta(delta) => {
// In review child threads, suppress assistant text deltas; the