Merge pull request #745 from vercel-labs/fazxes/correct-context-budget

Prevent false context capacity failures

README.md

@@ -125,6 +125,8 @@ Run `/trace` to create a private Markdown diagnostic with logs, session context,
 
 fx automatically summarizes a long session into a fresh context window when the active model request reaches 80% of its usable input capacity, then continues the same turn. Run `/compact` to create the same durable handoff immediately and wait for your next prompt. Manual compaction refreshes the selected login when needed; Ctrl+C cancels preparation. If authentication fails, the chat stays open and unchanged so you can reconnect and retry `/compact`.
 
+Context estimates adjust to observed provider usage, including when a local estimate is too high. Compaction budgets the recent reasoning history it keeps and can retain fewer complete exchanges when needed to leave room for the summary.
+
 Compaction handoffs remain internal context for the model. Resuming a session and opening its full transcript show the conversation and tool activity, not internal summaries or operation ledgers.
 
 Saved conversations preserve original assistant replies and compatible provider continuation data. Display formatting does not rewrite saved text, and hook-driven continuation keeps earlier replies separate from the final response.

src/core/agent/runtime/orchestrator.zig

@@ -5371,6 +5371,32 @@ pub const RetainedCompactionWindow = struct {
     retained_messages: []ChatMessage,
     cut: types.ContextHistoryCut,
     newest_exchange_tokens: usize,
+    retained_tokens: usize,
+
+    pub fn refine_budget(
+        self: RetainedCompactionWindow,
+        alloc: Allocator,
+        provider: agent_stream_provider.Provider,
+        continuation: CompactionContinuation,
+        capabilities: model_capabilities.Capabilities,
+        source_tokens: usize,
+        target: *usize,
+    ) !bool {
+        if (target.* == 0) return false;
+        const fixed_cost = try continuation.measure(alloc, provider, "");
+        const plan = runtime_prompt_context.planCompaction(.{
+            .trigger = .manual,
+            .capabilities = capabilities,
+            .request_tokens = source_tokens,
+            .source_tokens = source_tokens,
+            .protected_tokens = fixed_cost.estimated_input_tokens,
+            .newest_exchange_tokens = self.newest_exchange_tokens,
+        });
+        if (plan.accepted_handoff_tokens != null) return false;
+        target.* = if (self.retained_tokens <= self.newest_exchange_tokens) 0 else target.* / 2;
+        debug_trace.logf("context_compaction", "refine retained_tokens={d} protected_tokens={d} next_target={d}", .{ self.retained_tokens, fixed_cost.estimated_input_tokens, target.* });
+        return true;
+    }
 };
 
 pub fn prepareRetainedCompactionWindow(
@@ -5381,14 +5407,22 @@ pub fn prepareRetainedCompactionWindow(
     source_tokens: usize,
     provider: agent_stream_provider.Provider,
     provider_selection: model_provider.ProviderSelection,
+    options: struct {
+        target: ?usize = null,
+    },
 ) !RetainedCompactionWindow {
     var combined: std.ArrayList(HistoryTurn) = .empty;
     try combined.appendSlice(arena, history);
     if (active) |turn| try combined.append(arena, .{ .assistant = turn });
-    const selection = runtime_prompt_context.selectRecentContext(
+    const selection: runtime_prompt_context.RetainedContext = if (options.target != null and options.target.? == 0) .{
+        .cut = .{ .turns = session_runtime.rawHistoryTurnCount(combined.items) },
+        .newest_exchange_tokens = 0,
+        .estimated_tokens = 0,
+    } else runtime_prompt_context.selectRecentContext(
         combined.items,
-        runtime_prompt_context.recentContextTarget(capabilities, source_tokens),
+        options.target orelse runtime_prompt_context.recentContextTarget(capabilities, source_tokens),
         runtime_prompt_context.usableInputTokens(capabilities),
+        .{ .provider = provider_selection },
     );
     var cut = selection.cut;
     if (active) |turn| {
@@ -5420,6 +5454,7 @@ pub fn prepareRetainedCompactionWindow(
         .retained_messages = messages.items,
         .cut = cut,
         .newest_exchange_tokens = selection.newest_exchange_tokens,
+        .retained_tokens = selection.estimated_tokens,
     };
 }
 
@@ -5442,11 +5477,11 @@ test "retained context ends an unfinished turn at its completed exchange" {
         .execution = .{ .tool_steps = @constCast(&steps) },
     };
     const capabilities = model_capabilities.Capabilities{ .context_window = 4_000 };
-    const active = try prepareRetainedCompactionWindow(arena_state.allocator(), &.{}, turn, capabilities, 9_000, agent_stream_provider.unavailable_provider, .{ .provider = .gateway, .model = "fixture-model" });
+    const active = try prepareRetainedCompactionWindow(arena_state.allocator(), &.{}, turn, capabilities, 9_000, agent_stream_provider.unavailable_provider, .{ .provider = .gateway, .model = "fixture-model" }, .{});
     try std.testing.expectEqual(types.ContextHistoryCut{ .tool_steps = 1 }, active.cut);
     try std.testing.expectEqual(@as(usize, 3), active.source.len);
     try std.testing.expectEqual(@as(usize, 0), active.retained_messages.len);
-    const saved = try prepareRetainedCompactionWindow(arena_state.allocator(), &.{.{ .assistant = turn }}, null, capabilities, 9_000, agent_stream_provider.unavailable_provider, .{ .provider = .gateway, .model = "fixture-model" });
+    const saved = try prepareRetainedCompactionWindow(arena_state.allocator(), &.{.{ .assistant = turn }}, null, capabilities, 9_000, agent_stream_provider.unavailable_provider, .{ .provider = .gateway, .model = "fixture-model" }, .{});
     try std.testing.expectEqual(types.ContextHistoryCut{ .turns = 1 }, saved.cut);
     try std.testing.expectEqual(@as(usize, 0), saved.retained_messages.len);
 }
@@ -6348,153 +6383,169 @@ fn processQueuedPromptLoop(
                             }
                         }
                     },
-                    .compact => compact_attempt: {
-                        const uncertain_history_count = @min(
-                            @max(
-                                job.unversioned_history_count,
-                                0,
-                            ),
-                            job.history.len,
-                        );
-                        const prefix_execution = try runtime_execution_memory.buildExecutionMemory(arena, within_turn_suffix.items[compacted_suffix_len..]);
-                        const active_prefix: ?types.AssistantHistoryTurn = if (within_turn_suffix.items.len > compacted_suffix_len) .{
-                            .user = .{ .text = job.prompt, .images = job.images },
-                            .assistant = @constCast(""),
-                            .execution = prefix_execution,
-                        } else null;
-                        const window = try prepareRetainedCompactionWindow(arena, compaction_history, .{
-                            .user = .{ .text = job.prompt, .images = job.images },
-                            .assistant = @constCast(""),
-                            .execution = prefix_execution,
-                        }, request_capabilities, request_cost.estimated_input_tokens, deps.agent_stream_provider, .{ .provider = job.provider, .model = job.model });
-                        if (window.source.len == 0) {
-                            if (context_overflow_recovery == .pending or request_cost.estimated_input_tokens > (runtime_prompt_context.usableInputTokens(request_capabilities) orelse std.math.maxInt(usize))) return error.ContextCapacityExceeded;
-                            break :compact_attempt;
-                        }
-                        const raw_history_turns = session_runtime.rawHistoryTurnCount(compaction_history);
-                        const active_cut: runtime_execution_memory.CompactedExecutionBoundary = if (window.cut.turns == raw_history_turns) .{
-                            .tool_steps = window.cut.tool_steps,
-                            .steering = window.cut.steering,
-                        } else .{};
-                        const next_compacted_suffix_len = compacted_suffix_len + try runtime_execution_memory.retainedMessageOffset(within_turn_suffix.items[compacted_suffix_len..], active_cut);
-                        const result_storage: runtime_context_compaction.ResultStorage =
-                            if (config.session_child_capability) |capability|
-                                .{ .managed = capability }
-                            else if (config.tool_result_dir) |dir|
-                                .{ .legacy_dir = dir }
-                            else
-                                .unavailable;
-                        const compaction_source_message_count = window.source.len;
-                        var continuation_projection = try build_provider_prompt_with_response_language_control(
-                            overlay_arena,
-                            stable_prefix.items,
-                            ephemeral_overlay.items,
-                            &.{},
-                            current_user_effective,
-                            within_turn_suffix.items,
-                            config.origin,
-                            config.enforce_response_language,
-                            response_language_correction_attempted,
-                            "",
-                            window.retained_messages,
-                            next_compacted_suffix_len,
-                        );
-                        try appendRecoveryConversationContext(
-                            overlay_arena,
-                            &continuation_projection.messages,
-                            recovery_strategy,
-                        );
-                        var continuation_request = request_data;
-                        continuation_request.instructions = continuation_projection.instructions.items;
-                        continuation_request.messages = if (vision_policy.route == .fallback)
-                            try runtime_vision_contracts.project_text_only_messages(
+                    .compact => {
+                        var retention_target = runtime_prompt_context.recentContextTarget(request_capabilities, request_cost.estimated_input_tokens);
+                        var installed_compaction = false;
+                        compact_attempt: while (true) {
+                            const uncertain_history_count = @min(
+                                @max(
+                                    job.unversioned_history_count,
+                                    0,
+                                ),
+                                job.history.len,
+                            );
+                            const prefix_execution = try runtime_execution_memory.buildExecutionMemory(arena, within_turn_suffix.items[compacted_suffix_len..]);
+                            const active_prefix: ?types.AssistantHistoryTurn = if (within_turn_suffix.items.len > compacted_suffix_len) .{
+                                .user = .{ .text = job.prompt, .images = job.images },
+                                .assistant = @constCast(""),
+                                .execution = prefix_execution,
+                            } else null;
+                            const window = try prepareRetainedCompactionWindow(arena, compaction_history, .{
+                                .user = .{ .text = job.prompt, .images = job.images },
+                                .assistant = @constCast(""),
+                                .execution = prefix_execution,
+                            }, request_capabilities, request_cost.estimated_input_tokens, deps.agent_stream_provider, .{ .provider = job.provider, .model = gateway_model }, .{ .target = retention_target });
+                            if (window.source.len == 0) {
+                                if (context_overflow_recovery == .pending or request_cost.estimated_input_tokens > (runtime_prompt_context.usableInputTokens(request_capabilities) orelse std.math.maxInt(usize))) {
+                                    if (retention_target == 0 or request_cost.estimated_input_tokens <= (runtime_prompt_context.usableInputTokens(request_capabilities) orelse std.math.maxInt(usize))) return error.ContextCapacityExceeded;
+                                    retention_target = 0;
+                                    continue :compact_attempt;
+                                }
+                                break :compact_attempt;
+                            }
+                            const raw_history_turns = session_runtime.rawHistoryTurnCount(compaction_history);
+                            const active_cut: runtime_execution_memory.CompactedExecutionBoundary = if (window.cut.turns == raw_history_turns) .{
+                                .tool_steps = window.cut.tool_steps,
+                                .steering = window.cut.steering,
+                            } else .{};
+                            const next_compacted_suffix_len = compacted_suffix_len + try runtime_execution_memory.retainedMessageOffset(within_turn_suffix.items[compacted_suffix_len..], active_cut);
+                            const result_storage: runtime_context_compaction.ResultStorage =
+                                if (config.session_child_capability) |capability|
+                                    .{ .managed = capability }
+                                else if (config.tool_result_dir) |dir|
+                                    .{ .legacy_dir = dir }
+                                else
+                                    .unavailable;
+                            const compaction_source_message_count = window.source.len;
+                            var continuation_projection = try build_provider_prompt_with_response_language_control(
                                 overlay_arena,
-                                continuation_projection.messages.items,
-                                continuation_projection.current_user_index,
-                                job.authorized_image_catalog,
-                            )
-                        else
-                            continuation_projection.messages.items;
-                        const next_compaction_history_tail = window.retained_messages;
-                        const next_compaction_count = compaction_count + 1;
-                        const next_history = try arena.alloc(HistoryTurn, window.retained_history.len + 1);
-                        const transaction_result = compactContextTransaction(arena, deps, .{
-                            .trigger = compaction_trigger,
-                            .provider = job.provider,
-                            .working_capabilities = request_capabilities,
-                            .request_tokens = request_cost.estimated_input_tokens,
-                            .source_tokens = request_cost.estimated_input_tokens,
-                            .continuation = .{
+                                stable_prefix.items,
+                                ephemeral_overlay.items,
+                                &.{},
+                                current_user_effective,
+                                within_turn_suffix.items,
+                                config.origin,
+                                config.enforce_response_language,
+                                response_language_correction_attempted,
+                                "",
+                                window.retained_messages,
+                                next_compacted_suffix_len,
+                            );
+                            try appendRecoveryConversationContext(
+                                overlay_arena,
+                                &continuation_projection.messages,
+                                recovery_strategy,
+                            );
+                            var continuation_request = request_data;
+                            continuation_request.instructions = continuation_projection.instructions.items;
+                            continuation_request.messages = if (vision_policy.route == .fallback)
+                                try runtime_vision_contracts.project_text_only_messages(
+                                    overlay_arena,
+                                    continuation_projection.messages.items,
+                                    continuation_projection.current_user_index,
+                                    job.authorized_image_catalog,
+                                )
+                            else
+                                continuation_projection.messages.items;
+                            const continuation = CompactionContinuation{
                                 .request = continuation_request,
                                 .handoff_message_index = continuation_projection.current_user_index - window.retained_messages.len - 1,
-                            },
-                            .active_prefix = active_prefix,
-                            .retained_from = window.cut,
-                            .newest_exchange_tokens = window.newest_exchange_tokens,
-                            .source_messages = window.source,
-                            .uncertain_source_message_count = if (uncertain_history_count > 0) compaction_source_message_count else 0,
-                            .result_storage = result_storage,
-                            .api_key = active_api_key,
-                            .credential_source = job.credential_source,
-                            .account_id = job.account_id,
-                            .gateway_team = job.gateway_team,
-                            .session_id = lifecycle.scope.session_id,
-                            .retry_count = config.gateway_retry_count,
-                            .cancel_flag = config.cancel_flag,
-                            .trace_ctx = step_ctx,
-                            .removed_turn_count = window.cut.turns,
-                            .compaction_count = next_compaction_count,
-                        }) catch |err| {
-                            if (err == error.Cancelled and config.cancel_flag.load(.seq_cst)) {
-                                runtime_telemetry.traceCancelObserved(step_ctx, false);
-                                try runtime_interruption.persistInterruptedTurnOnce(
-                                    deps,
-                                    finalization,
-                                    job,
-                                    null,
-                                    null,
-                                    completed_tool_names.items,
-                                    &interrupted_persisted,
-                                    step_ctx,
-                                    within_turn_suffix.items,
-                                    stop_state.retained_candidate,
-                                    &stop_state.terminal_materializing,
-                                );
-                                finish_trace.finish("interrupted");
-                                return;
+                            };
+                            const refine = window.refine_budget(overlay_arena, deps.agent_stream_provider, continuation, request_capabilities, request_cost.estimated_input_tokens, &retention_target) catch |err| blk: {
+                                if (err == error.Cancelled and config.cancel_flag.load(.seq_cst)) break :blk false;
+                                return err;
+                            };
+                            if (refine) continue :compact_attempt;
+                            const next_compaction_history_tail = window.retained_messages;
+                            const next_compaction_count = compaction_count + 1;
+                            const next_history = try arena.alloc(HistoryTurn, window.retained_history.len + 1);
+                            const transaction_result = compactContextTransaction(arena, deps, .{
+                                .trigger = compaction_trigger,
+                                .provider = job.provider,
+                                .working_capabilities = request_capabilities,
+                                .request_tokens = request_cost.estimated_input_tokens,
+                                .source_tokens = request_cost.estimated_input_tokens,
+                                .continuation = continuation,
+                                .active_prefix = active_prefix,
+                                .retained_from = window.cut,
+                                .newest_exchange_tokens = window.newest_exchange_tokens,
+                                .source_messages = window.source,
+                                .uncertain_source_message_count = if (uncertain_history_count > 0) compaction_source_message_count else 0,
+                                .result_storage = result_storage,
+                                .api_key = active_api_key,
+                                .credential_source = job.credential_source,
+                                .account_id = job.account_id,
+                                .gateway_team = job.gateway_team,
+                                .session_id = lifecycle.scope.session_id,
+                                .retry_count = config.gateway_retry_count,
+                                .cancel_flag = config.cancel_flag,
+                                .trace_ctx = step_ctx,
+                                .removed_turn_count = window.cut.turns,
+                                .compaction_count = next_compaction_count,
+                            }) catch |err| {
+                                if (err == error.Cancelled and config.cancel_flag.load(.seq_cst)) {
+                                    runtime_telemetry.traceCancelObserved(step_ctx, false);
+                                    try runtime_interruption.persistInterruptedTurnOnce(
+                                        deps,
+                                        finalization,
+                                        job,
+                                        null,
+                                        null,
+                                        completed_tool_names.items,
+                                        &interrupted_persisted,
+                                        step_ctx,
+                                        within_turn_suffix.items,
+                                        stop_state.retained_candidate,
+                                        &stop_state.terminal_materializing,
+                                    );
+                                    finish_trace.finish("interrupted");
+                                    return;
+                                }
+                                return err;
+                            };
+                            const transaction = transaction_result orelse
+                                return error.ContextCapacityExceeded;
+                            active_compaction_handoff = transaction.compacted.handoff;
+                            active_compaction_history_tail = next_compaction_history_tail;
+                            next_history[0] = .{ .compacted_summary = .{
+                                .summary = transaction.compacted.handoff,
+                                .removed_turn_count = window.cut.turns,
+                                .compaction_count = next_compaction_count,
+                            } };
+                            @memcpy(next_history[1..], window.retained_history);
+                            compaction_history = next_history;
+                            compacted_suffix_len = next_compacted_suffix_len;
+                            finalization.compacted_execution = .{
+                                .tool_steps = finalization.compacted_execution.tool_steps + active_cut.tool_steps,
+                                .steering = finalization.compacted_execution.steering + active_cut.steering,
+                            };
+                            compaction_count = next_compaction_count;
+                            if (context_overflow_recovery == .pending) {
+                                context_overflow_recovery = .used;
                             }
-                            return err;
-                        };
-                        const transaction = transaction_result orelse
-                            return error.ContextCapacityExceeded;
-                        active_compaction_handoff = transaction.compacted.handoff;
-                        active_compaction_history_tail = next_compaction_history_tail;
-                        next_history[0] = .{ .compacted_summary = .{
-                            .summary = transaction.compacted.handoff,
-                            .removed_turn_count = window.cut.turns,
-                            .compaction_count = next_compaction_count,
-                        } };
-                        @memcpy(next_history[1..], window.retained_history);
-                        compaction_history = next_history;
-                        compacted_suffix_len = next_compacted_suffix_len;
-                        finalization.compacted_execution = .{
-                            .tool_steps = finalization.compacted_execution.tool_steps + active_cut.tool_steps,
-                            .steering = finalization.compacted_execution.steering + active_cut.steering,
-                        };
-                        compaction_count = next_compaction_count;
-                        if (context_overflow_recovery == .pending) {
-                            context_overflow_recovery = .used;
+                            debug_trace.eventf(
+                                "context_compaction",
+                                "installed",
+                                step_ctx,
+                                "request_bytes_before={d} estimated_tokens_before={d} handoff_bytes={d} accepted_tokens={d}",
+                                .{ request_cost.serialized_bytes, request_cost.estimated_input_tokens, active_compaction_handoff.?.len, transaction.accepted_tokens },
+                            );
+                            request_token_calibration = null;
+                            skip_next_preflight_refresh = true;
+                            installed_compaction = true;
+                            break :compact_attempt;
                         }
-                        debug_trace.eventf(
-                            "context_compaction",
-                            "installed",
-                            step_ctx,
-                            "request_bytes_before={d} estimated_tokens_before={d} handoff_bytes={d} accepted_tokens={d}",
-                            .{ request_cost.serialized_bytes, request_cost.estimated_input_tokens, active_compaction_handoff.?.len, transaction.accepted_tokens },
-                        );
-                        request_token_calibration = null;
-                        skip_next_preflight_refresh = true;
-                        continue;
+                        if (installed_compaction) continue;
                     },
                 }
             }

src/core/agent/runtime/prompt_context.zig

@@ -4,6 +4,7 @@ const token_estimate = @import("../../shared/token_estimate.zig");
 const types = @import("../../shared/types.zig");
 const session_runtime = @import("../../session/session.zig");
 const stream_provider = @import("../stream_provider.zig");
+const model_provider = @import("../../config/model_provider.zig");
 
 const Allocator = std.mem.Allocator;
 const ChatMessage = types.ChatMessage;
@@ -118,7 +119,9 @@ pub const RetainedContext = struct {
 };
 
 /// Selects complete execution steps. Payloads are measured, never shortened.
-pub fn selectRecentContext(history: []const HistoryTurn, target: usize, input_capacity: ?usize) RetainedContext {
+pub fn selectRecentContext(history: []const HistoryTurn, target: usize, input_capacity: ?usize, options: struct {
+    provider: ?model_provider.ProviderSelection = null,
+}) RetainedContext {
     var raw_count = session_runtime.rawHistoryTurnCount(history);
     var selected = types.ContextHistoryCut{ .turns = raw_count };
     var total: usize = 0;
@@ -145,7 +148,12 @@ pub fn selectRecentContext(history: []const HistoryTurn, target: usize, input_ca
             .interrupted => |entry| entry.execution,
             .compacted_summary => unreachable,
         };
-        var base = textTokens(user) +| textTokens(reply) +| 8;
+        const replay = switch (turn) {
+            .assistant => |entry| entry.provider_replay,
+            .interrupted => null,
+            .compacted_summary => unreachable,
+        };
+        var base = textTokens(user) +| textTokens(reply) +| replay_tokens(replay, options.provider) +| 8;
         var steering_index = execution.steering.len;
         var step_index = execution.tool_steps.len;
         if (step_index == 0) {
@@ -165,7 +173,7 @@ pub fn selectRecentContext(history: []const HistoryTurn, target: usize, input_ca
         }
         while (step_index > 0) {
             step_index -= 1;
-            var cost = base +| executionStepTokens(execution.tool_steps[step_index]);
+            var cost = base +| executionStepTokens(execution.tool_steps[step_index], options.provider);
             var next_steering = steering_index;
             while (next_steering > 0 and execution.steering[next_steering - 1].after_tool_step_count >= step_index) {
                 next_steering -= 1;
@@ -197,8 +205,14 @@ fn textTokens(text: []const u8) usize {
     return @intCast(@min(estimator.estimate(), std.math.maxInt(usize)));
 }
 
-fn executionStepTokens(step: types.ToolExecutionStep) usize {
-    var total: usize = 8;
+fn replay_tokens(replay: ?types.ProviderReplay, provider: ?model_provider.ProviderSelection) usize {
+    const value = replay orelse return 0;
+    if (provider) |selection| if (!value.matches(selection)) return 0;
+    return textTokens(value.parts_json);
+}
+
+fn executionStepTokens(step: types.ToolExecutionStep, provider: ?model_provider.ProviderSelection) usize {
+    var total: usize = 8 +| replay_tokens(step.provider_replay, provider);
     if (step.assistant) |text| total +|= textTokens(text);
     for (step.tool_calls) |call| {
         total +|= textTokens(call.id) +| textTokens(call.name) +| textTokens(call.arguments_json) +| 8;
@@ -232,6 +246,7 @@ pub fn validateCompactionHandoff(
 
 pub const RequestCost = struct {
     serialized_bytes: usize,
+    /// Uncalibrated serialization estimate; non-image usage may replace it.
     text_tokens: usize,
     /// Null means no image parts. Otherwise visual cost needs applicable usage.
     image_identity: ?[32]u8 = null,
@@ -354,7 +369,10 @@ pub fn calibrateProviderRequest(
     else
         multiplyDivideCeilSaturating(cost.serialized_bytes, calibration.exact_input_tokens, calibration.request.serialized_bytes);
     var result = cost;
-    result.estimated_input_tokens = @max(cost.text_tokens, calibrated_tokens);
+    result.estimated_input_tokens = if (cost.image_identity != null)
+        @max(cost.text_tokens, calibrated_tokens)
+    else
+        @max(1, calibrated_tokens);
     return result;
 }
 
@@ -638,6 +656,45 @@ test "provider request measurement learns the prior exact token density" {
     try std.testing.expect(calibrated.estimated_input_tokens >= current.estimated_input_tokens);
 }
 
+test "provider usage corrects a serialized estimate downward" {
+    const cost = RequestCost{ .serialized_bytes = 34_210, .text_tokens = 8_892, .estimated_input_tokens = 8_892 };
+    const calibrated = calibrateProviderRequest(cost, .{ .request = cost, .exact_input_tokens = 6_030 });
+    try std.testing.expectEqual(@as(usize, 6_030), calibrated.estimated_input_tokens);
+    try std.testing.expectEqual(cost.text_tokens, calibrated.text_tokens);
+}
+
+test "retained context budgets provider replay on completed exchanges" {
+    const replay = types.ProviderReplay{ .source = .{ .provider = .gateway, .model = "fixture/model" }, .parts_json = "[{\"type\":\"reasoning\",\"text\":\"\",\"providerOptions\":{\"openai\":{\"reasoningEncryptedContent\":\"" ++ ("r" ** 80_000) ++ "\"}}}]" };
+    var steps = [_]types.ToolExecutionStep{
+        .{ .assistant = @constCast("one"), .provider_replay = replay },
+        .{ .assistant = @constCast("two"), .provider_replay = replay },
+        .{ .assistant = @constCast("three"), .provider_replay = replay },
+    };
+    const history = [_]HistoryTurn{.{ .assistant = .{
+        .user = .{ .text = @constCast("continue") },
+        .assistant = @constCast(""),
+        .execution = .{ .tool_steps = &steps },
+    } }};
+    const selected = selectRecentContext(&history, 5_990, 119_808, .{});
+    try std.testing.expectEqual(@as(usize, 2), selected.cut.tool_steps);
+    try std.testing.expect(selected.newest_exchange_tokens >= 20_000);
+    const changed_model = selectRecentContext(&history, 5_990, 119_808, .{ .provider = .{ .provider = .gateway, .model = "fixture/other" } });
+    try std.testing.expectEqual(@as(usize, 1), changed_model.cut.tool_steps);
+    try std.testing.expect(changed_model.newest_exchange_tokens < 100);
+}
+
+test "retained context budgets replay on standalone assistant replies" {
+    const turn = types.AssistantHistoryTurn{
+        .user = .{ .text = @constCast("continue") },
+        .assistant = @constCast("small reply"),
+        .provider_replay = .{ .source = .{ .provider = .gateway, .model = "fixture/model" }, .parts_json = "[{\"type\":\"reasoning\",\"text\":\"\",\"providerOptions\":{\"openai\":{\"reasoningEncryptedContent\":\"" ++ ("r" ** 80_000) ++ "\"}}}]" },
+    };
+    const history = [_]HistoryTurn{ .{ .assistant = turn }, .{ .assistant = turn }, .{ .assistant = turn } };
+    const selected = selectRecentContext(&history, 5_990, 119_808, .{});
+    try std.testing.expectEqual(@as(usize, 2), selected.cut.turns);
+    try std.testing.expect(selected.newest_exchange_tokens >= 20_000);
+}
+
 fn measurement_test_request(with_images: bool) stream_provider.RequestData {
     return .{
         .model = "fixture/model",
@@ -868,13 +925,13 @@ test "retained context selects whole parallel tool exchanges without shortening
         .{ .assistant = .{ .user = .{ .text = @constCast("old request") }, .assistant = @constCast("old answer") } },
         .{ .assistant = .{ .user = .{ .text = @constCast("current request") }, .assistant = @constCast(""), .execution = .{ .tool_steps = @constCast(&steps) } } },
     };
-    const selected = selectRecentContext(&history, 5000, null);
+    const selected = selectRecentContext(&history, 5000, null, .{});
     try std.testing.expectEqual(@as(usize, 1), selected.cut.turns);
     try std.testing.expectEqual(@as(usize, 1), selected.cut.tool_steps);
     try std.testing.expect(selected.newest_exchange_tokens > 5000);
     try std.testing.expectEqualStrings(body, results[0].output);
     try std.testing.expectEqualStrings(body, results[1].output);
-    const over_capacity = selectRecentContext(&history, 5000, 10_000);
+    const over_capacity = selectRecentContext(&history, 5000, 10_000, .{});
     try std.testing.expectEqual(@as(usize, 2), over_capacity.cut.turns);
     try std.testing.expectEqual(@as(usize, 0), over_capacity.cut.tool_steps);
     try std.testing.expectEqual(@as(usize, 0), over_capacity.estimated_tokens);

src/core/agent/runtime/tests/gateway_flow.zig

@@ -3086,6 +3086,121 @@ test "automatic compaction rejects fixed request overhead above its total target
     for (hooks.history_turns.items) |turn| try std.testing.expect(turn != .compacted_summary);
 }
 
+test "automatic compaction shrinks recent history to fit fixed instructions" {
+    const alloc = std.testing.allocator;
+    var gateway = FakeGateway.init(alloc, &.{
+        .{ .content = "The earlier work is complete. Preserve the recent facts." },
+        .{ .content = "Continued after compaction." },
+    });
+    defer gateway.deinit();
+    const model = "fixture/retained-budget";
+    const overrides = [_]ModelCapabilityOverride{.{ .model = model, .capabilities = .{ .context_window = 128_000, .max_output_tokens = 8_192 } }};
+    var hooks = FakeAgentRuntimeDeps.init(alloc);
+    defer hooks.deinit();
+    hooks.available_capability_overrides = &overrides;
+    var fixture = PromptFixture{};
+    var config = fixture.config();
+    config.system_prompt = "i" ** 104_000;
+    var job = fixture.job();
+    job.model = @constCast(model);
+    var history = [_]HistoryTurn{
+        .{ .assistant = .{ .user = .{ .text = @constCast("older") }, .assistant = @constCast("OLD_BUDGET_FACT " ++ ("h" ** 380_000)) } },
+        .{ .assistant = .{ .user = .{ .text = @constCast("middle") }, .assistant = @constCast("MIDDLE_BUDGET_FACT " ++ ("m" ** 8_000)) } },
+        .{ .assistant = .{ .user = .{ .text = @constCast("recent") }, .assistant = @constCast("RECENT_BUDGET_FACT " ++ ("r" ** 8_000)) } },
+    };
+    job.history = &history;
+    try runFakePrompt(&gateway, &hooks, config, job);
+    try std.testing.expectEqual(@as(usize, 2), gateway.request_bodies.items.len);
+    try expectBodyContains(&gateway, 0, "OLD_BUDGET_FACT");
+    try expectBodyContains(&gateway, 0, "MIDDLE_BUDGET_FACT");
+    try expectBodyContains(&gateway, 1, "RECENT_BUDGET_FACT");
+    try expectBodyNotContains(&gateway, 1, "MIDDLE_BUDGET_FACT");
+    try std.testing.expectEqual(@as(usize, 2), hooks.history_turns.items.len);
+    try std.testing.expect(hooks.history_turns.items[0] == .compacted_summary);
+}
+
+test "compaction remeasures its rebuilt continuation after calibrated preflight" {
+    const alloc = std.testing.allocator;
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    const result_dir = try io_mod.dirRealpathAlloc(alloc, tmp.dir, ".");
+    defer alloc.free(result_dir);
+    var calls: [4][1]ToolCall = undefined;
+    var states: [4][]u8 = undefined;
+    var initialized: usize = 0;
+    defer for (states[0..initialized]) |state| alloc.free(state);
+    const ids = [_][]const u8{ "calibration-one", "calibration-two", "calibration-three", "calibration-four" };
+    const inputs = [_]u64{ 10, 4_000, 8_000, 15_000 };
+    var completions: [6]FakeCompletion = undefined;
+    for (ids, 0..) |id, index| {
+        calls[index] = .{toolCall(id, "read_file", "{\"path\":\"fixture.txt\"}")};
+        states[index] = try std.fmt.allocPrint(alloc, "[{{\"type\":\"reasoning\",\"text\":\"\",\"providerOptions\":{{\"openai\":{{\"reasoningEncryptedContent\":\"{s}\"}}}}}},{{\"type\":\"tool-call\",\"toolCallId\":\"{s}\"}}]", .{ "r" ** 200_000, id });
+        initialized += 1;
+        completions[index] = .{ .tool_calls = &calls[index], .provider_state_json = states[index], .usage = .{ .input_tokens = inputs[index] } };
+    }
+    completions[4] = .{ .content = "The fixture reads completed successfully." };
+    completions[5] = .{ .content = "Calibrated continuation completed." };
+    var gateway = FakeGateway.init(alloc, &completions);
+    defer gateway.deinit();
+    var hooks = FakeAgentRuntimeDeps.init(alloc);
+    defer hooks.deinit();
+    const model = "fixture/calibrated-compaction";
+    const overrides = [_]ModelCapabilityOverride{.{ .model = model, .capabilities = .{ .context_window = 20_000 } }};
+    hooks.available_capability_overrides = &overrides;
+    hooks.permission_decisions = &.{ .once, .once, .once, .once };
+    hooks.exec_plans = &.{ .{ .result = .{ .model_output = "read one" } }, .{ .result = .{ .model_output = "read two" } }, .{ .result = .{ .model_output = "read three" } }, .{ .result = .{ .model_output = "read four" } } };
+    var fixture = PromptFixture{};
+    var config = fixture.config();
+    config.tool_result_dir = result_dir;
+    var job = fixture.job();
+    job.model = @constCast(model);
+    try runFakePrompt(&gateway, &hooks, config, job);
+    try std.testing.expectEqual(@as(usize, 6), gateway.request_bodies.items.len);
+    try expectBodyContains(&gateway, 5, "context_handoff");
+    const raw = try prompt_context.measureProviderRequest(alloc, gateway.request_bodies.items[5], .{ .model = model, .messages = &.{}, .tool_choice = .auto, .provider_options = .{} });
+    try std.testing.expect(raw.estimated_input_tokens < 20_000);
+    try std.testing.expectEqualStrings("Calibrated continuation completed.", hooks.finish_assistant_text.?);
+    try std.testing.expectEqual(@as(usize, 4), hooks.successful_effect_count.load(.seq_cst));
+}
+
+test "compaction can summarize the newest exchange when fixed context prevents retaining it" {
+    const alloc = std.testing.allocator;
+    for ([_]bool{ false, true }) |with_older_history| {
+        var tmp = std.testing.tmpDir(.{});
+        defer tmp.cleanup();
+        const result_dir = try io_mod.dirRealpathAlloc(alloc, tmp.dir, ".");
+        defer alloc.free(result_dir);
+        const calls = [_]ToolCall{toolCall("completed-large-write", "write_file", "{\"path\":\"large.txt\",\"content\":\"" ++ ("x" ** 68_000) ++ "\"}")};
+        var gateway = FakeGateway.init(alloc, &.{
+            .{ .tool_calls = &calls },
+            .{ .content = "The large write completed. Do not repeat it." },
+            .{ .content = "Continued with the completed write preserved." },
+        });
+        defer gateway.deinit();
+        var hooks = FakeAgentRuntimeDeps.init(alloc);
+        defer hooks.deinit();
+        const model = "fixture/fixed-newest-budget";
+        const overrides = [_]ModelCapabilityOverride{.{ .model = model, .capabilities = .{ .context_window = 20_000 } }};
+        hooks.available_capability_overrides = &overrides;
+        hooks.permission_decisions = &.{.once};
+        hooks.exec_plans = &.{.{ .result = .{ .model_output = "Write completed." } }};
+        var fixture = PromptFixture{};
+        var config = fixture.config();
+        config.system_prompt = "i" ** 16_000;
+        config.tool_result_dir = result_dir;
+        var job = fixture.job();
+        job.model = @constCast(model);
+        var older = [_]HistoryTurn{.{ .assistant = .{ .user = .{ .text = @constCast("older request") }, .assistant = @constCast("older fact") } }};
+        if (with_older_history) job.history = &older;
+        try runFakePrompt(&gateway, &hooks, config, job);
+        try std.testing.expectEqual(@as(usize, 3), gateway.request_bodies.items.len);
+        try expectBodyContains(&gateway, 1, "completed-large-write");
+        try expectBodyContains(&gateway, 2, "context_handoff");
+        try expectBodyNotContains(&gateway, 2, "x" ** 68_000);
+        try std.testing.expectEqual(@as(usize, 1), hooks.successful_effect_count.load(.seq_cst));
+    }
+}
+
 test "automatic compaction validates serialized handoff cost before committing" {
     const alloc = std.testing.allocator;
     var gateway = FakeGateway.init(alloc, &.{.{ .content = "\"" ** 10_000 }});

src/core/app/app_agent_runtime.zig

@@ -1094,66 +1094,77 @@ pub fn Runtime(comptime App: type) type {
                 .grants = &.{},
                 .agent_settings = app.worker.effectiveAgentTurnSettings(),
             }, .{ .skills = skill_catalog.items, .diagnostics = skill_catalog.diagnostics }, &.{}, gateway_retry_count, "", &tool_projection, null);
-            var continuation = try agent_runtime.prepareManualCompactionContinuation(
+            const base_continuation = try agent_runtime.prepareManualCompactionContinuation(
                 arena,
                 &deps,
                 config,
                 job.model,
                 capabilities,
             );
-            const window = try agent_runtime.prepareRetainedCompactionWindow(arena, job.history, null, capabilities, source_tokens, deps.agent_stream_provider, .{ .provider = job.provider, .model = job.model });
-            const continuation_messages = try arena.alloc(ChatMessage, continuation.request.messages.len + window.retained_messages.len);
-            @memcpy(continuation_messages[0..continuation.request.messages.len], continuation.request.messages);
-            @memcpy(continuation_messages[continuation.request.messages.len..], window.retained_messages);
-            continuation.request.messages = continuation_messages;
-            var compaction_count: usize = 0;
-            for (job.history) |turn| switch (turn) {
-                .compacted_summary => |summary| {
-                    compaction_count = @max(
-                        compaction_count,
-                        summary.compaction_count,
-                    );
-                },
-                else => {},
-            };
-            const transaction = agent_runtime.compactContextTransaction(arena, &deps, .{
-                .trigger = .manual,
-                .provider = job.provider,
-                .working_capabilities = capabilities,
-                .request_tokens = source_tokens,
-                .source_tokens = runtime_prompt_context.estimateCompactionSourceTokens(window.source),
-                .continuation = continuation,
-                .retained_from = window.cut,
-                .newest_exchange_tokens = window.newest_exchange_tokens,
-                .source_messages = window.source,
-                .uncertain_source_message_count = if (uncertain_history_count > 0) window.source.len else 0,
-                .result_storage = result_storage,
-                .api_key = job.api_key,
-                .credential_source = job.credential_source,
-                .account_id = job.account_id,
-                .gateway_team = job.gateway_team,
-                .session_id = app_session_runtime.Runtime(App).activeSessionId(app),
-                .retry_count = gateway_retry_count,
-                .cancel_flag = &app.worker.worker_cancel_requested,
-                .trace_ctx = .{ .turn_id = job.turn_id },
-                .removed_turn_count = window.cut.turns,
-                .compaction_count = compaction_count + 1,
-            }) catch |err| {
-                if (err == error.Cancelled and
-                    app.worker.worker_cancel_requested.load(.seq_cst))
-                {
+            var retention_target = runtime_prompt_context.recentContextTarget(capabilities, source_tokens);
+            while (true) {
+                if (app.worker.worker_cancel_requested.load(.seq_cst)) return;
+                var continuation = base_continuation;
+                const window = try agent_runtime.prepareRetainedCompactionWindow(arena, job.history, null, capabilities, source_tokens, deps.agent_stream_provider, .{ .provider = job.provider, .model = job.model }, .{ .target = retention_target });
+                const continuation_messages = try arena.alloc(ChatMessage, continuation.request.messages.len + window.retained_messages.len);
+                @memcpy(continuation_messages[0..continuation.request.messages.len], continuation.request.messages);
+                @memcpy(continuation_messages[continuation.request.messages.len..], window.retained_messages);
+                continuation.request.messages = continuation_messages;
+                const refine = window.refine_budget(arena, deps.agent_stream_provider, continuation, capabilities, source_tokens, &retention_target) catch |err| {
+                    if (err == error.Cancelled and app.worker.worker_cancel_requested.load(.seq_cst)) return;
+                    return err;
+                };
+                if (refine) continue;
+                var compaction_count: usize = 0;
+                for (job.history) |turn| switch (turn) {
+                    .compacted_summary => |summary| {
+                        compaction_count = @max(
+                            compaction_count,
+                            summary.compaction_count,
+                        );
+                    },
+                    else => {},
+                };
+                const transaction = agent_runtime.compactContextTransaction(arena, &deps, .{
+                    .trigger = .manual,
+                    .provider = job.provider,
+                    .working_capabilities = capabilities,
+                    .request_tokens = source_tokens,
+                    .source_tokens = runtime_prompt_context.estimateCompactionSourceTokens(window.source),
+                    .continuation = continuation,
+                    .retained_from = window.cut,
+                    .newest_exchange_tokens = window.newest_exchange_tokens,
+                    .source_messages = window.source,
+                    .uncertain_source_message_count = if (uncertain_history_count > 0) window.source.len else 0,
+                    .result_storage = result_storage,
+                    .api_key = job.api_key,
+                    .credential_source = job.credential_source,
+                    .account_id = job.account_id,
+                    .gateway_team = job.gateway_team,
+                    .session_id = app_session_runtime.Runtime(App).activeSessionId(app),
+                    .retry_count = gateway_retry_count,
+                    .cancel_flag = &app.worker.worker_cancel_requested,
+                    .trace_ctx = .{ .turn_id = job.turn_id },
+                    .removed_turn_count = window.cut.turns,
+                    .compaction_count = compaction_count + 1,
+                }) catch |err| {
+                    if (err == error.Cancelled and
+                        app.worker.worker_cancel_requested.load(.seq_cst))
+                    {
+                        return;
+                    }
+                    return err;
+                };
+                _ = transaction orelse {
+                    try app_worker_runtime.Runtime(App).pushSemanticNotice(app, .{
+                        .topic = "context",
+                        .tone = .neutral,
+                        .body = "No context to compact.",
+                    });
                     return;
-                }
-                return err;
-            };
-            _ = transaction orelse {
-                try app_worker_runtime.Runtime(App).pushSemanticNotice(app, .{
-                    .topic = "context",
-                    .tone = .neutral,
-                    .body = "No context to compact.",
-                });
+                };
                 return;
-            };
+            }
         }
 
         fn lifecycleContext(app: *App) agent_runtime.LifecycleContext {

tests/e2e/gateway-stream-lifecycle.test.ts

@@ -5301,6 +5301,82 @@ printf '%s' ${JSON.stringify(trailingMarker)} > ${JSON.stringify(effectPath)}
     60_000,
   );
 
+  test.skipIf(!tmuxAvailable())("reasoning replay survives automatic and manual compaction and cold resume", async () => {
+    const root = createFixtureRoot("reasoning-context-budget");
+    const tracePath = join(root.root, "trace.log");
+    const stderrPath = join(root.root, "stderr.log");
+    writeFileSync(join(root.workspace, "sentinel.txt"), "REPLAY_RESULT_SENTINEL\n");
+    let ordinary = 0;
+    let summaries = 0;
+    const gateway = startDynamicFakeGateway((body) => {
+      const request = JSON.parse(body);
+      if (request.tools.length === 0) {
+        summaries++;
+        expect(body).not.toContain("LARGE_REASONING_");
+        expect(body).not.toContain("RECENT_REASONING_SIGNATURE");
+        expect(body).toContain("REPLAY_RESULT_SENTINEL");
+        return fakeGatewayFinalText("The prior reads completed. Preserve REPLAY_RESULT_SENTINEL and continue without repeating completed reads.");
+      }
+      ordinary++;
+      if (ordinary === 6) {
+        expect(body).toContain("context_handoff");
+        expect(body).toContain("LARGE_REASONING_5");
+        expect(body).not.toContain("LARGE_REASONING_1");
+      }
+      if (ordinary <= 6) return fakeGatewaySse([
+        { type: "reasoning-start", id: `reasoning-${ordinary}` },
+        { type: "reasoning-end", id: `reasoning-${ordinary}`, providerMetadata: { openai: { reasoningEncryptedContent: `LARGE_REASONING_${ordinary}` + "a".repeat(80_000) } } },
+        { type: "tool-call", toolCallId: `read-${ordinary}`, toolName: "read_file", input: { path: "sentinel.txt" } },
+        { type: "finish", finishReason: { unified: "tool-calls", raw: "tool-calls" } },
+      ]);
+      if (ordinary === 7) return fakeGatewayFinalText("REPLAY_TURN_DONE");
+      if (ordinary === 8) return fakeGatewaySse([
+        { type: "reasoning-start", id: "recent-reasoning" },
+        { type: "reasoning-end", id: "recent-reasoning", providerMetadata: { openai: { reasoningEncryptedContent: "RECENT_REASONING_SIGNATURE" } } },
+        { type: "text-delta", id: "recent-answer", delta: "RECENT_TURN_DONE" },
+        { type: "finish", finishReason: { unified: "stop", raw: "stop" } },
+      ]);
+      expect(body).toContain("RECENT_REASONING_SIGNATURE");
+      expect(body).toContain("REPLAY_RESULT_SENTINEL");
+      expect(body).not.toContain("LARGE_REASONING_");
+      return fakeGatewayFinalText("COLD_REPLAY_DONE");
+    }, { models: [{ id: MODEL, type: "language", tags: ["tool-use", "reasoning"], context_window: 128_000, max_tokens: 8192 }] });
+    const env = { ...fixtureEnv(root, gateway, tracePath), FX_AUTO_UPGRADE: "0", FX_SOUND: "0", FX_TRACE_SCOPES: "agent,session,context_compaction" };
+    let tui: TmuxSession | null = null;
+    try {
+      tui = await TmuxSession.create({ cwd: root.workspace, env, stderrPath, isolated: true });
+      await tui.waitForComposer(15_000);
+      await tui.sendText("Read sentinel.txt six times, then finish.");
+      const automatic = await tui.waitForPane((text) => hasEmptyComposer(text) && (text.includes("REPLAY_TURN_DONE") || text.includes("request failed:")), 30_000);
+      expect(automatic).toContain("REPLAY_TURN_DONE");
+      expect(automatic).not.toContain("request failed:");
+      expect(summaries).toBe(1);
+      await tui.sendText("Remember the result and acknowledge this short follow-up.");
+      await tui.waitForText("RECENT_TURN_DONE", 15_000);
+      await tui.waitForComposer(15_000);
+      await tui.sendText("/compact");
+      await tui.waitForPane((text) => hasEmptyComposer(text) && summaries === 2 && (text.includes("Context compacted.") || text.includes("request failed:")), 20_000);
+      expect(readFileSync(tracePath, "utf8")).not.toContain("ContextCapacityExceeded");
+      await tui.sendText("/quit");
+      expect(await tui.waitForSessionEnd(15_000)).toBe(true);
+      tui = null;
+      expect(readFileSync(stderrPath, "utf8")).toBe("");
+      const latest = await runFx(["session", "last", "--json"], { cwd: root.workspace, env });
+      expect(latest.code).toBe(0);
+      const id = JSON.parse(latest.stdout).id;
+      const resumed = await runFx(["ask", "--json", "--auto", "--resume-id", id, "Continue with the saved result."], { cwd: root.workspace, env, timeoutMs: 30_000 });
+      expect(resumed.code, resumed.stderr || resumed.stdout).toBe(0);
+      expect(JSON.parse(resumed.stdout).final_output).toBe("COLD_REPLAY_DONE");
+      expect(ordinary).toBe(9);
+      expect(summaries).toBe(2);
+      expect(readFileSync(join(root.workspace, "sentinel.txt"), "utf8")).toBe("REPLAY_RESULT_SENTINEL\n");
+    } finally {
+      await tui?.kill();
+      gateway.stop();
+      rmSync(root.root, { recursive: true, force: true });
+    }
+  }, 100_000);
+
   for (const trigger of ["automatic", "manual"] as const) {
     test.skipIf(!tmuxAvailable())(
       `oversized result retrieval survives empty ${trigger} summary recovery and restart`,