From fb94436c2d2563435e958f0fb293cb8a85b61510 Mon Sep 17 00:00:00 2001 From: Pranit Date: Mon, 7 Sep 2026 16:50:40 -0400 Subject: [PATCH 1/2] Prevent false context capacity failures Let observed provider usage lower non-image context estimates. Include compatible replay when budgeting recent history and reduce the retained window before rejecting compaction for lack of summary space. --- README.md | 2 + src/core/agent/runtime/orchestrator.zig | 343 ++++++++++-------- src/core/agent/runtime/prompt_context.zig | 73 +++- src/core/agent/runtime/tests/gateway_flow.zig | 115 ++++++ src/core/app/app_agent_runtime.zig | 115 +++--- tests/e2e/gateway-stream-lifecycle.test.ts | 76 ++++ 6 files changed, 518 insertions(+), 206 deletions(-) diff --git a/README.md b/README.md index 124aaaff5..420a932d7 100644 --- a/README.md +++ b/README.md @@ -123,6 +123,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. diff --git a/src/core/agent/runtime/orchestrator.zig b/src/core/agent/runtime/orchestrator.zig index 69a8820d3..89428eab7 100644 --- a/src/core/agent/runtime/orchestrator.zig +++ b/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) 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; }, } } diff --git a/src/core/agent/runtime/prompt_context.zig b/src/core/agent/runtime/prompt_context.zig index 99e27a0a4..97c5b886f 100644 --- a/src/core/agent/runtime/prompt_context.zig +++ b/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); diff --git a/src/core/agent/runtime/tests/gateway_flow.zig b/src/core/agent/runtime/tests/gateway_flow.zig index 5ac6d715e..e6007c974 100644 --- a/src/core/agent/runtime/tests/gateway_flow.zig +++ b/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 }}); diff --git a/src/core/app/app_agent_runtime.zig b/src/core/app/app_agent_runtime.zig index 0a676a192..ceb49c07e 100644 --- a/src/core/app/app_agent_runtime.zig +++ b/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 { diff --git a/tests/e2e/gateway-stream-lifecycle.test.ts b/tests/e2e/gateway-stream-lifecycle.test.ts index e556e1e52..397dbe32c 100644 --- a/tests/e2e/gateway-stream-lifecycle.test.ts +++ b/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`, From e9fef4319aa300f17b21c97c95aa6671cadb7cd5 Mon Sep 17 00:00:00 2001 From: Pranit Date: Mon, 7 Sep 2026 17:18:19 -0400 Subject: [PATCH 2/2] Preserve capacity errors after provider rejections Only discard the sole retained exchange when the request exceeds local input capacity. Smaller rejected requests keep the existing capacity error without sending an unnecessary compaction request. --- src/core/agent/runtime/orchestrator.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/core/agent/runtime/orchestrator.zig b/src/core/agent/runtime/orchestrator.zig index 89428eab7..7fc53f7ed 100644 --- a/src/core/agent/runtime/orchestrator.zig +++ b/src/core/agent/runtime/orchestrator.zig @@ -6407,7 +6407,7 @@ fn processQueuedPromptLoop( }, 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) return error.ContextCapacityExceeded; + 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; }