diff --git a/src/acp/prompt.zig b/src/acp/prompt.zig index 386c0dfa7..c7013a638 100644 --- a/src/acp/prompt.zig +++ b/src/acp/prompt.zig @@ -2865,7 +2865,7 @@ fn describeToolTitle(registry: tool_dispatch.Registry, arena: Allocator, call: T .call = call, }); } - if (tool_dispatch.toolCallPresentation(arena, registry, call)) |presentation| { + if (tool_dispatch.toolCallPresentation(registry, call)) |presentation| { return std.fmt.allocPrint(arena, "{s}", .{presentation.action_label}); } return std.fmt.allocPrint(arena, "{s}", .{call.name}); diff --git a/src/core/agent/runtime/deps.zig b/src/core/agent/runtime/deps.zig index 15901ed86..2de29344a 100644 --- a/src/core/agent/runtime/deps.zig +++ b/src/core/agent/runtime/deps.zig @@ -32,6 +32,10 @@ const TransportPublicationOutcome = tool_contracts.TransportPublicationOutcome; pub const LiveToolAuthority = tool_contracts.LiveToolAuthority; pub const RecoveryCheckpointEffect = struct { + /// The checkpoint borrows caller-owned scratch memory that is freed as + /// soon as this call returns. A sink must serialize the checkpoint or + /// dupe it with its own allocator before returning; it must not retain + /// the passed pointers. set: *const fn (ctx: *anyopaque, checkpoint: session_codec.RecoveryCheckpoint) anyerror!void, }; diff --git a/src/core/agent/runtime/orchestrator.zig b/src/core/agent/runtime/orchestrator.zig index b62c7644e..f2b27a4d4 100644 --- a/src/core/agent/runtime/orchestrator.zig +++ b/src/core/agent/runtime/orchestrator.zig @@ -3175,7 +3175,6 @@ noinline fn recoveryCheckpointAssistantSource( fn persistRecoveryCheckpoint( deps: *const AgentRuntimeDeps, - arena: Allocator, job: QueuedPrompt, current_turn_messages: []const ChatMessage, assistant_source: []const u8, @@ -3191,8 +3190,16 @@ fn persistRecoveryCheckpoint( trace_ctx: TraceContext, ) !void { const effect = deps.recovery_checkpoint orelse return; + // The execution memory is a deep copy of every tool result accumulated so + // far this turn. Building it in the turn arena would retain one full copy + // per checkpoint for the rest of the turn (quadratic in steps), so it + // lives in a scratch arena instead: every effect.set sink either + // serializes the checkpoint or dupes it with its own allocator before + // returning. + var scratch_state = std.heap.ArenaAllocator.init(std.heap.c_allocator); + defer scratch_state.deinit(); const execution = try runtime_execution_memory.buildExecutionMemory( - arena, + scratch_state.allocator(), current_turn_messages, ); try effect.set(deps.ctx, .{ @@ -3264,11 +3271,91 @@ fn streamCompletionPtr(result: *runtime_gateway_step.StreamResult) ?*types.Model }; } -fn retainCompletedResultInTurnArena(result: *runtime_gateway_step.StreamResult) void { - switch (result.*) { - .completed => |*completed| completed.ownership = .borrowed, - .failed => {}, - } +/// Moves a provider result out of its attempt arena into the turn arena, which +/// owns every byte the rest of the step reads. The copy is marked borrowed +/// because the turn arena reclaims it on turn exit. +fn copyStreamResultToTurnArena( + arena: Allocator, + result: runtime_gateway_step.StreamResult, +) Allocator.Error!runtime_gateway_step.StreamResult { + // Tripwire: a new slice-bearing field on any of these types must be copied + // here before bumping its count, or it dangles once the attempt arena dies. + comptime std.debug.assert(@typeInfo(agent_stream_provider.Completed).@"struct".fields.len == 3); + comptime std.debug.assert(@typeInfo(agent_stream_provider.Failure).@"struct".fields.len == 5); + comptime std.debug.assert(@typeInfo(agent_stream_provider.FailureDiagnostics).@"struct".fields.len == 2); + comptime std.debug.assert(@typeInfo(agent_stream_provider.DeferredUsageReference).@"struct".fields.len == 7); + return switch (result) { + .completed => |completed| .{ .completed = .{ + .completion = try types.dupeModelCompletion(arena, completed.completion), + .usage = switch (completed.usage) { + .deferred => |reference| .{ .deferred = .{ + .provider = reference.provider, + .generation_id = try arena.dupe(u8, reference.generation_id), + .scope = try arena.dupe(u8, reference.scope), + .tenant = if (reference.tenant) |value| try arena.dupe(u8, value) else null, + .account_id = if (reference.account_id) |value| try arena.dupe(u8, value) else null, + .credential_source = reference.credential_source, + .credential_identity = reference.credential_identity, + } }, + else => completed.usage, + }, + .ownership = .borrowed, + } }, + .failed => |failure| .{ .failed = .{ + .kind = failure.kind, + .detail = if (failure.detail) |value| try arena.dupe(u8, value) else null, + .diagnostics = .{ + .schema = if (failure.diagnostics.schema) |value| try arena.dupe(u8, value) else null, + .request_shape = if (failure.diagnostics.request_shape) |value| try arena.dupe(u8, value) else null, + }, + .retry_after_seconds = failure.retry_after_seconds, + .ownership = .borrowed, + } }, + }; +} + +test "stream result copy survives attempt arena teardown" { + var attempt_state = std.heap.ArenaAllocator.init(std.testing.allocator); + const attempt_alloc = attempt_state.allocator(); + var turn_state = std.heap.ArenaAllocator.init(std.testing.allocator); + defer turn_state.deinit(); + const turn_alloc = turn_state.allocator(); + + const completed = try copyStreamResultToTurnArena(turn_alloc, .{ .completed = .{ + .completion = .{ + .content = try attempt_alloc.dupe(u8, "answer"), + .provider_state_json = try attempt_alloc.dupe(u8, "[]"), + }, + .usage = .{ .deferred = .{ + .provider = .gateway, + .generation_id = try attempt_alloc.dupe(u8, "gen_1"), + .scope = try attempt_alloc.dupe(u8, "scope"), + .tenant = try attempt_alloc.dupe(u8, "team"), + .credential_source = .stored_key, + .credential_identity = null, + } }, + .ownership = .owned, + } }); + const failed = try copyStreamResultToTurnArena(turn_alloc, .{ .failed = .{ + .kind = .rate_limited, + .detail = try attempt_alloc.dupe(u8, "slow down"), + .diagnostics = .{ .schema = try attempt_alloc.dupe(u8, "schema") }, + .retry_after_seconds = 3, + .ownership = .owned, + } }); + attempt_state.deinit(); + + try std.testing.expectEqualStrings("answer", completed.completed.completion.content.?); + try std.testing.expectEqualStrings("[]", completed.completed.completion.provider_state_json.?); + try std.testing.expectEqualStrings("gen_1", completed.completed.usage.deferred.generation_id); + try std.testing.expectEqualStrings("scope", completed.completed.usage.deferred.scope); + try std.testing.expectEqualStrings("team", completed.completed.usage.deferred.tenant.?); + try std.testing.expectEqual(agent_stream_provider.ResultOwnership.borrowed, completed.completed.ownership); + try std.testing.expectEqual(agent_stream_provider.FailureKind.rate_limited, failed.failed.kind); + try std.testing.expectEqualStrings("slow down", failed.failed.detail.?); + try std.testing.expectEqualStrings("schema", failed.failed.diagnostics.schema.?); + try std.testing.expectEqual(@as(?u64, 3), failed.failed.retry_after_seconds); + try std.testing.expectEqual(agent_stream_provider.ResultOwnership.borrowed, failed.failed.ownership); } fn isRetryableModelFailure(kind: agent_stream_provider.FailureKind) bool { @@ -5045,6 +5132,12 @@ fn processQueuedPromptLoop( } while (true) { + // Provider scratch (request body, HTTP client, SSE parse trees) + // lives only for this attempt; the survivors are copied into the + // turn arena before `break`. + var attempt_arena_state = std.heap.ArenaAllocator.init(std.heap.c_allocator); + defer attempt_arena_state.deinit(); + const attempt_alloc = attempt_arena_state.allocator(); if (reset_stream_for_next_attempt) { try stream_ctx.beginRecoveryAttempt(); reset_stream_for_next_attempt = false; @@ -5055,7 +5148,6 @@ fn processQueuedPromptLoop( recovery_strategy = .pause; try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, try recoveryCheckpointAssistantSource( @@ -5092,7 +5184,6 @@ fn processQueuedPromptLoop( recovery_strategy = .pause; try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, try recoveryCheckpointAssistantSource( @@ -5506,7 +5597,6 @@ fn processQueuedPromptLoop( if (job.provider == .gateway) { try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, stream_ctx.interruption_source_or(""), @@ -5574,7 +5664,7 @@ fn processQueuedPromptLoop( }; stream_result = runtime_gateway_step.streamModelCompletion( deps.agent_stream_provider, - arena, + attempt_alloc, model_request, deps.usage, deps.usage_allocator, @@ -5603,7 +5693,6 @@ fn processQueuedPromptLoop( if (recoveryPauseRequested(config)) { try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, try recoveryCheckpointAssistantSource( @@ -5700,7 +5789,6 @@ fn processQueuedPromptLoop( } try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, try recoveryCheckpointAssistantSource( @@ -5742,7 +5830,6 @@ fn processQueuedPromptLoop( if (recoveryPauseRequested(config)) { try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, try recoveryCheckpointAssistantSource( @@ -5962,7 +6049,7 @@ fn processQueuedPromptLoop( model_request.attempt_evidence = &replay_evidence; stream_result = try runtime_gateway_step.streamModelCompletion( deps.agent_stream_provider, - arena, + attempt_alloc, model_request, deps.usage, deps.usage_allocator, @@ -5983,7 +6070,7 @@ fn processQueuedPromptLoop( } if (streamCompletionPtr(&stream_result)) |completion| { completion.tool_calls = try normalize_terminal_request_tool_calls( - arena, + attempt_alloc, deps.tool_registry, terminal_request_eligible, completion.tool_calls, @@ -6009,7 +6096,6 @@ fn processQueuedPromptLoop( ); try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, try recoveryCheckpointAssistantSource( @@ -6070,7 +6156,6 @@ fn processQueuedPromptLoop( { try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, try recoveryCheckpointAssistantSource( @@ -6121,7 +6206,7 @@ fn processQueuedPromptLoop( "tool_name=vision provider_attempt={d}/{d}", .{ semantic_attempt + 1, semantic_limit }, ); - stream_result.deinit(arena); + stream_result.deinit(attempt_alloc); stream_result_set = false; assistant_prefill_recovery_used = true; semantic_attempt += 1; @@ -6166,7 +6251,6 @@ fn processQueuedPromptLoop( if (decision.strategy == .pause) { try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, try recoveryCheckpointAssistantSource( @@ -6206,7 +6290,6 @@ fn processQueuedPromptLoop( if (route_changed) { try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, try recoveryCheckpointAssistantSource( @@ -6254,7 +6337,7 @@ fn processQueuedPromptLoop( response_completion, &stream_ctx, ); - stream_result.deinit(arena); + stream_result.deinit(attempt_alloc); stream_result_set = false; semantic_attempt += 1; recovery_strategy = decision.strategy; @@ -6500,7 +6583,6 @@ fn processQueuedPromptLoop( if (decision.strategy == .pause) { try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, try recoveryCheckpointAssistantSource( @@ -6539,7 +6621,6 @@ fn processQueuedPromptLoop( if (route_changed) { try persistRecoveryCheckpoint( deps, - arena, job, within_turn_suffix.items, try recoveryCheckpointAssistantSource( @@ -6700,7 +6781,7 @@ fn processQueuedPromptLoop( successful_vision_route = vision_route; successful_vision_mode = vision_mode; successful_recovery_strategy = recovery_strategy; - retainCompletedResultInTurnArena(&stream_result); + stream_result = try copyStreamResultToTurnArena(arena, stream_result); if (vision_mode != .required) configured_first_tool_choice_pending = false; return_to_user_pending = false; break; @@ -6970,7 +7051,7 @@ fn processQueuedPromptLoop( if (disposition == .completed) try stream_ctx.start_response(); var step_has_visible_tool_calls = false; for (completion.tool_calls) |call| { - if (runtime_tool_presentation.activityKindForCall(arena, deps.tool_registry, call) == .ask) continue; + if (runtime_tool_presentation.activityKindForCall(deps.tool_registry, call) == .ask) continue; step_has_visible_tool_calls = true; break; } @@ -7588,7 +7669,7 @@ fn processQueuedPromptLoop( const step_has_content = !terminal_provider_completion and completion.content != null and completion.content.?.len > 0; if (step_has_content) { const first_tool_is_ask = effective_tool_calls.len > 0 and - runtime_tool_presentation.activityKindForCall(arena, deps.tool_registry, effective_tool_calls[0]) == .ask; + runtime_tool_presentation.activityKindForCall(deps.tool_registry, effective_tool_calls[0]) == .ask; if (!first_tool_is_ask) try deps.push_text(deps.ctx, .{ .assistant_rendered = "\n" }); silent_tool_steps = 0; } else { @@ -8726,15 +8807,14 @@ fn processQueuedPromptLoop( status_started = try runtime_tool_presentation.startToolVisibleLifecycle(deps, arena, turn_id, stream_ctx.provisional_statuses.presentation_group_id, tool_call, tool_display_target, advertised_dynamic_tool_names); } - var file_call_arena_state: std.heap.ArenaAllocator = undefined; - if (is_file_mutation) { - file_call_arena_state = std.heap.ArenaAllocator.init(std.heap.c_allocator); - } - defer if (is_file_mutation) file_call_arena_state.deinit(); - const call_allocator = if (is_file_mutation) - file_call_arena_state.allocator() - else - arena; + // Every call gets its own arena so decode/validate/call scratch + // is reclaimed when the call returns instead of accumulating in + // the turn arena for the rest of the turn. Everything that must + // outlive the call travels through result_allocator (the turn + // arena) inside executeToolCallAuthorized. + var call_arena_state = std.heap.ArenaAllocator.init(std.heap.c_allocator); + defer call_arena_state.deinit(); + const call_allocator = call_arena_state.allocator(); const execution_call = if (is_file_mutation) try types.dupeToolCall(call_allocator, tool_call) else @@ -9254,7 +9334,7 @@ fn processQueuedPromptLoop( }; } const execution_lifecycle_id = types.ToolLifecycleId{ .turn_id = turn_id, .call_id = execution_call.id }; - const execution_is_command = runtime_tool_presentation.activityKindForCall(arena, deps.tool_registry, tool_call) == .command; + const execution_is_command = runtime_tool_presentation.activityKindForCall(deps.tool_registry, tool_call) == .command; var execution_error: ?anyerror = null; var execution = deps.execute_tool_call(deps.ctx, .{ .call_allocator = call_allocator, diff --git a/src/core/agent/runtime/parallel_execution.zig b/src/core/agent/runtime/parallel_execution.zig index 9fa7c1cdd..3c9ac7948 100644 --- a/src/core/agent/runtime/parallel_execution.zig +++ b/src/core/agent/runtime/parallel_execution.zig @@ -276,8 +276,12 @@ fn cancelRequested(cancel_flag: ?*std.atomic.Value(bool)) bool { pub fn parallelHookExecute(ctx: *anyopaque, alloc: Allocator, call: ToolCall, index: usize) !ToolExecutionResult { const exec_ctx: *ParallelHookExecContext = @ptrCast(@alignCast(ctx)); + // Same per-call scratch ownership as the sequential path: call scratch is + // reclaimed when the call returns, survivors are copied to alloc. + var call_arena_state = std.heap.ArenaAllocator.init(std.heap.c_allocator); + defer call_arena_state.deinit(); return exec_ctx.hooks.execute_tool_call(exec_ctx.hooks.ctx, .{ - .call_allocator = alloc, + .call_allocator = call_arena_state.allocator(), .result_allocator = alloc, .call = call, .authority = .ordinary, diff --git a/src/core/agent/runtime/tool_batch.zig b/src/core/agent/runtime/tool_batch.zig index 5fa25f303..2f775a24f 100644 --- a/src/core/agent/runtime/tool_batch.zig +++ b/src/core/agent/runtime/tool_batch.zig @@ -481,7 +481,7 @@ pub fn appendOrdinaryExecutedResult( memory: types.ToolResultMemory, execution: ToolExecutionResult, ) !void { - const activity = runtime_tool_presentation.activityKindForCall(arena, tool_registry, tool_call); + const activity = runtime_tool_presentation.activityKindForCall(tool_registry, tool_call); try appendToolResultContent( arena, within_turn_suffix, diff --git a/src/core/agent/runtime/tool_presentation.zig b/src/core/agent/runtime/tool_presentation.zig index 20919d443..7cea06a2a 100644 --- a/src/core/agent/runtime/tool_presentation.zig +++ b/src/core/agent/runtime/tool_presentation.zig @@ -689,12 +689,11 @@ pub fn activityKind(registry: tool_dispatch.Registry, tool_name: []const u8) typ } pub fn activityKindForCall( - alloc: Allocator, registry: tool_dispatch.Registry, call: ToolCall, ) types.ToolActivityKind { if (tooling_presentation.isProviderSearchAlias(call.name)) return .read; - return tool_dispatch.toolActivityKindForCall(alloc, registry, call); + return tool_dispatch.toolActivityKindForCall(registry, call); } fn formatProvisionalProgressLabel( @@ -736,7 +735,7 @@ pub noinline fn startToolVisibleLifecycle( display_target: ?[]const u8, advertised_dynamic_tool_names: []const []const u8, ) !bool { - const activity_kind = activityKindForCall(arena, hooks.tool_registry, call); + const activity_kind = activityKindForCall(hooks.tool_registry, call); if (activity_kind == .ask) return false; const redacted_arguments = try text_utils.maskSecrets(arena, call.arguments_json); const activity_line = try hooks.describe_tool_action( @@ -861,7 +860,7 @@ fn finishDeniedToolStatusInternal( label, advertised_dynamic_tool_names, ); - const command_artifact_handle = if (activityKindForCall(arena, hooks.tool_registry, call) == .command) + const command_artifact_handle = if (activityKindForCall(hooks.tool_registry, call) == .command) try commandArtifactHandle(arena, command_result_json) else null; @@ -908,7 +907,6 @@ pub fn finishCancelledToolStatus( advertised_dynamic_tool_names, ); const command_activity = activityKindForCall( - arena, hooks.tool_registry, call, ) == .command; @@ -972,7 +970,7 @@ pub fn finishExecutedToolStatus( advertised_dynamic_tool_names: []const []const u8, ) !void { if (!status_started) return; - const activity_kind = activityKindForCall(arena, hooks.tool_registry, call); + const activity_kind = activityKindForCall(hooks.tool_registry, call); const command_decision = if (activity_kind == .command) try commandOutcomeDecision(arena, result_memory.command_process_presentation) else diff --git a/src/core/agent/worker_runtime.zig b/src/core/agent/worker_runtime.zig index 9ff07e4bb..dece37341 100644 --- a/src/core/agent/worker_runtime.zig +++ b/src/core/agent/worker_runtime.zig @@ -3309,32 +3309,7 @@ fn dupeToolResultMemory( source: ?types.ToolResultMemory, ) !?types.ToolResultMemory { const memory = source orelse return null; - const output_handle = if (memory.output_handle) |handle| - try alloc.dupe(u8, handle) - else - null; - errdefer if (output_handle) |handle| alloc.free(handle); - const preview = if (memory.preview) |value| - try alloc.dupe(u8, value) - else - null; - errdefer if (preview) |value| alloc.free(value); - const command_output_replay = if (memory.command_output_replay) |replay| - try types.dupeCommandOutputReplay(alloc, replay) - else - null; - errdefer if (command_output_replay) |replay| types.freeCommandOutputReplay(alloc, replay); - return .{ - .output_handle = output_handle, - .preview = preview, - .output_bytes = memory.output_bytes, - .stored_output_bytes = memory.stored_output_bytes, - .truncated = memory.truncated, - .model_view_covers_full_file = memory.model_view_covers_full_file, - .command_output_replay = command_output_replay, - .command_process_presentation = memory.command_process_presentation, - .terminal_action_presentation = memory.terminal_action_presentation, - }; + return try types.dupeToolResultMemory(alloc, memory); } fn freeToolResultMemory( diff --git a/src/core/shared/types.zig b/src/core/shared/types.zig index 1fb7828d3..773ab7e70 100644 --- a/src/core/shared/types.zig +++ b/src/core/shared/types.zig @@ -2391,6 +2391,176 @@ pub fn freeCommandOutputReplay( } } +/// Deep-copies every slice-bearing field so the result is owned by `alloc`. +/// The caller owns the returned memory. +pub fn dupeToolResultMemory( + alloc: std.mem.Allocator, + memory: ToolResultMemory, +) !ToolResultMemory { + const output_handle = if (memory.output_handle) |handle| + try alloc.dupe(u8, handle) + else + null; + errdefer if (output_handle) |handle| alloc.free(@constCast(handle)); + const preview = if (memory.preview) |value| + try alloc.dupe(u8, value) + else + null; + errdefer if (preview) |value| alloc.free(@constCast(value)); + const command_output_replay = if (memory.command_output_replay) |replay| + try dupeCommandOutputReplay(alloc, replay) + else + null; + errdefer if (command_output_replay) |replay| freeCommandOutputReplay(alloc, replay); + const committed_file_presentation = if (memory.committed_file_presentation) |presentation| + try dupeCommittedFilePresentation(alloc, presentation) + else + null; + return .{ + .output_handle = output_handle, + .preview = preview, + .output_bytes = memory.output_bytes, + .stored_output_bytes = memory.stored_output_bytes, + .truncated = memory.truncated, + .model_view_covers_full_file = memory.model_view_covers_full_file, + .committed_file_presentation = committed_file_presentation, + .command_output_replay = command_output_replay, + .command_process_presentation = memory.command_process_presentation, + .terminal_action_presentation = memory.terminal_action_presentation, + }; +} + +/// Deep-copies every slice-bearing field of a provider completion so the +/// result is owned by `alloc`. Arena callers rely on this to move a completion +/// out of a shorter-lived attempt allocator. +pub fn dupeModelCompletion(alloc: std.mem.Allocator, completion: ModelCompletion) !ModelCompletion { + var copy = completion; + copy.content = if (completion.content) |value| try alloc.dupe(u8, value) else null; + errdefer if (copy.content) |value| alloc.free(@constCast(value)); + copy.tool_calls = try dupeToolCallSlice(alloc, completion.tool_calls); + errdefer freeToolCallSlice(alloc, @constCast(copy.tool_calls)); + copy.generation_id = if (completion.generation_id) |value| try alloc.dupe(u8, value) else null; + errdefer if (copy.generation_id) |value| alloc.free(@constCast(value)); + if (completion.billing) |billing| copy.billing.?.model = try alloc.dupe(u8, billing.model); + errdefer if (copy.billing) |billing| alloc.free(@constCast(billing.model)); + copy.provider_failure_detail = if (completion.provider_failure_detail) |value| try alloc.dupe(u8, value) else null; + errdefer if (copy.provider_failure_detail) |value| alloc.free(@constCast(value)); + copy.provider_state_json = if (completion.provider_state_json) |value| try alloc.dupe(u8, value) else null; + return copy; +} + +test "model completion dupe covers every slice-bearing field" { + // Tripwire: adding a field to ModelCompletion requires extending + // dupeModelCompletion (and the attempt-boundary copy-out that relies on + // it) before bumping this count. + comptime std.debug.assert(@typeInfo(ModelCompletion).@"struct".fields.len == 12); + comptime std.debug.assert(@typeInfo(ProviderBilling).@"struct".fields.len == 9); + + var source_state = std.heap.ArenaAllocator.init(std.testing.allocator); + const source_alloc = source_state.allocator(); + const source = ModelCompletion{ + .content = try source_alloc.dupe(u8, "answer"), + .tool_calls = try dupeToolCallSlice(source_alloc, &.{.{ + .id = "call_1", + .name = "read_file", + .arguments_json = "{}", + .provisional_id = "tmp_1", + .provider_result = "{\"ok\":true}", + }}), + .generation_id = try source_alloc.dupe(u8, "gen_1"), + .billing = .{ + .created_at_ms = 7, + .model = try source_alloc.dupe(u8, "model-x"), + .total_cost = 0.5, + .input_tokens = 1, + .output_tokens = 2, + .cache_read_tokens = 3, + .cache_write_tokens = 4, + .reasoning_tokens = 5, + .billable_web_search_calls = 6, + }, + .generation_metadata_invalid = true, + .delivery_ambiguous = true, + .provider_result_identity_failure = .absent, + .provider_failure_cause = .gateway_stream_timeout, + .provider_failure_detail = try source_alloc.dupe(u8, "detail"), + .provider_state_json = try source_alloc.dupe(u8, "[{\"id\":\"rs_1\"}]"), + .finish_reason = .tool_calls, + .usage = .{ .input_tokens = 9 }, + }; + + var owned_state = std.heap.ArenaAllocator.init(std.testing.allocator); + defer owned_state.deinit(); + const owned = try dupeModelCompletion(owned_state.allocator(), source); + source_state.deinit(); + + try std.testing.expectEqualStrings("answer", owned.content.?); + try std.testing.expectEqual(@as(usize, 1), owned.tool_calls.len); + try std.testing.expectEqualStrings("call_1", owned.tool_calls[0].id); + try std.testing.expectEqualStrings("read_file", owned.tool_calls[0].name); + try std.testing.expectEqualStrings("{}", owned.tool_calls[0].arguments_json); + try std.testing.expectEqualStrings("tmp_1", owned.tool_calls[0].provisional_id.?); + try std.testing.expectEqualStrings("{\"ok\":true}", owned.tool_calls[0].provider_result.?); + try std.testing.expectEqualStrings("gen_1", owned.generation_id.?); + try std.testing.expectEqualStrings("model-x", owned.billing.?.model); + try std.testing.expectEqual(@as(u64, 4), owned.billing.?.cache_write_tokens); + try std.testing.expect(owned.generation_metadata_invalid); + try std.testing.expect(owned.delivery_ambiguous); + try std.testing.expectEqual(ProviderResultIdentityFailure.absent, owned.provider_result_identity_failure.?); + try std.testing.expectEqual(ProviderFailureCause.gateway_stream_timeout, owned.provider_failure_cause.?); + try std.testing.expectEqualStrings("detail", owned.provider_failure_detail.?); + try std.testing.expectEqualStrings("[{\"id\":\"rs_1\"}]", owned.provider_state_json.?); + try std.testing.expectEqual(ProviderFinishReason.tool_calls, owned.finish_reason.?); + try std.testing.expectEqual(@as(?u64, 9), owned.usage.input_tokens); +} + +test "tool result memory dupe covers every slice-bearing field" { + // Tripwire: adding a field to ToolResultMemory requires extending + // dupeToolResultMemory (and the dispatch-boundary copy-out that relies on + // it) before bumping this count. + comptime std.debug.assert(@typeInfo(ToolResultMemory).@"struct".fields.len == 10); + + var arena_state = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena_state.deinit(); + const arena = arena_state.allocator(); + const source = ToolResultMemory{ + .output_handle = "handle.txt", + .preview = "preview text", + .output_bytes = 42, + .stored_output_bytes = 21, + .truncated = true, + .model_view_covers_full_file = false, + .committed_file_presentation = .{ + .path = "src/file.zig", + .kind = .edited, + .lines = &.{}, + .additions = 1, + .deletions = 2, + .truncated = false, + }, + .command_output_replay = .{ .available = .{ + .handle = "replay-handle", + .framed_bytes = 7, + } }, + .command_process_presentation = .{ .signal = 9 }, + .terminal_action_presentation = .{ .returned = .safety_ceiling }, + }; + + const owned = try dupeToolResultMemory(arena, source); + try std.testing.expectEqualStrings("handle.txt", owned.output_handle.?); + try std.testing.expect(owned.output_handle.?.ptr != source.output_handle.?.ptr); + try std.testing.expectEqualStrings("preview text", owned.preview.?); + try std.testing.expect(owned.preview.?.ptr != source.preview.?.ptr); + try std.testing.expectEqualStrings("src/file.zig", owned.committed_file_presentation.?.path); + try std.testing.expect(owned.committed_file_presentation.?.path.ptr != + source.committed_file_presentation.?.path.ptr); + try std.testing.expectEqualStrings("replay-handle", owned.command_output_replay.?.available.handle); + try std.testing.expect(owned.command_output_replay.?.available.handle.ptr != + source.command_output_replay.?.available.handle.ptr); + try std.testing.expectEqual(@as(usize, 42), owned.output_bytes); + try std.testing.expectEqual(source.command_process_presentation, owned.command_process_presentation); +} + pub fn dupePermissionFeedback( alloc: std.mem.Allocator, feedback: []const []const u8, diff --git a/src/core/tooling/file_mutation_execution.zig b/src/core/tooling/file_mutation_execution.zig index 5b089a969..763111ff3 100644 --- a/src/core/tooling/file_mutation_execution.zig +++ b/src/core/tooling/file_mutation_execution.zig @@ -207,7 +207,7 @@ fn fileMutationFailure( }; } -fn allocatorsEqual(a: Allocator, b: Allocator) bool { +pub fn allocatorsEqual(a: Allocator, b: Allocator) bool { return a.ptr == b.ptr and a.vtable == b.vtable; } diff --git a/src/core/tooling/tool_dispatch.zig b/src/core/tooling/tool_dispatch.zig index f5a2e9dbc..5228447dd 100644 --- a/src/core/tooling/tool_dispatch.zig +++ b/src/core/tooling/tool_dispatch.zig @@ -578,12 +578,13 @@ pub fn presentationForArgs(tool: Tool, args: std.json.ObjectMap) CallPresentatio } pub fn toolCallPresentation( - alloc: Allocator, registry: Registry, call: core_types.ToolCall, ) ?CallPresentation { const tool = registry.lookup(call.name) orelse return null; - var scratch_state = std.heap.ArenaAllocator.init(alloc); + // Scratch is backed by c_allocator so deinit reclaims it even when the + // caller's allocator is the per-turn arena. + var scratch_state = std.heap.ArenaAllocator.init(std.heap.c_allocator); defer scratch_state.deinit(); const args = tool_args.parseToolArgsObject( scratch_state.allocator(), @@ -593,11 +594,10 @@ pub fn toolCallPresentation( } pub fn toolActivityKindForCall( - alloc: Allocator, registry: Registry, call: core_types.ToolCall, ) core_types.ToolActivityKind { - const presentation = toolCallPresentation(alloc, registry, call) orelse + const presentation = toolCallPresentation(registry, call) orelse return .command; return presentation.activity_kind; } diff --git a/src/core/tooling/tool_result_limits.zig b/src/core/tooling/tool_result_limits.zig index 71c2ebfb2..5e0719cc7 100644 --- a/src/core/tooling/tool_result_limits.zig +++ b/src/core/tooling/tool_result_limits.zig @@ -21,7 +21,7 @@ pub fn prepareRedactedOutput( alloc: Allocator, raw: []const u8, ) error{OutOfMemory}![]u8 { - var scratch_impl = std.heap.ArenaAllocator.init(alloc); + var scratch_impl = std.heap.ArenaAllocator.init(std.heap.c_allocator); defer scratch_impl.deinit(); const redacted = try redactModelText(scratch_impl.allocator(), raw); return alloc.dupe(u8, redacted); @@ -47,7 +47,9 @@ pub fn prepareModelOutputWithTruncation( raw: []const u8, max_bytes: usize, ) error{OutOfMemory}!PreparedModelOutput { - var scratch_impl = std.heap.ArenaAllocator.init(alloc); + // Scratch is backed by c_allocator so deinit reclaims it even when the + // caller's allocator is the per-turn arena. + var scratch_impl = std.heap.ArenaAllocator.init(std.heap.c_allocator); defer scratch_impl.deinit(); const scratch = scratch_impl.allocator(); diff --git a/src/core/tooling/tool_runtime.zig b/src/core/tooling/tool_runtime.zig index 3ef76fc86..34f5156c4 100644 --- a/src/core/tooling/tool_runtime.zig +++ b/src/core/tooling/tool_runtime.zig @@ -419,6 +419,7 @@ pub fn executeToolCallAuthorized( else executeToolCallInner( execution_ctx, + request.call_allocator, request.result_allocator, request.call, request.authority, @@ -586,7 +587,8 @@ const ToolDispatchPrelude = union(enum) { fn executeToolCallInner( ctx: Context, - arena: Allocator, + call_allocator: Allocator, + result_allocator: Allocator, call: ToolCall, authority: command_admission.ToolExecutionAuthority, classification_complete: bool, @@ -594,14 +596,15 @@ fn executeToolCallInner( ) !ToolExecutionResult { return switch (try resolveToolDispatchPrelude( ctx, - arena, + result_allocator, call, classification_complete, )) { .completed => |result| return result, .registered_static => try executeRegisteredTool( ctx, - arena, + call_allocator, + result_allocator, call, authority, authorized_image_catalog, @@ -611,7 +614,8 @@ fn executeToolCallInner( const tools = [_]tool_dispatch.Tool{dynamic_tool}; break :blk try executeRegisteredTool( ctx, - arena, + call_allocator, + result_allocator, call, authority, authorized_image_catalog, @@ -643,7 +647,10 @@ fn executeWorkspaceToolCallInner( return semanticFailure(try std.fmt.allocPrint(arena, "Unsupported tool: {s}", .{call.name})); } - var command_backend = RunCommandBackendState{ .runtime = ctx }; + var command_backend = RunCommandBackendState{ + .runtime = ctx, + .result_allocator = arena, + }; var dispatch_metadata: DispatchMetadata = .{}; var dispatch_ctx = typedDispatchContextForCall(ctx, arena, call); dispatch_metadata.attach(&dispatch_ctx); @@ -742,25 +749,34 @@ fn emitMcpProgress(raw_context: *anyopaque, progress: tool_mcp_runtime.Progress) fn executeRegisteredTool( ctx: Context, - arena: Allocator, + call_allocator: Allocator, + result_allocator: Allocator, call: ToolCall, authority: command_admission.ToolExecutionAuthority, authorized_image_catalog: []const types.ImageAttachment, registry: tool_dispatch.Registry, ) !ToolExecutionResult { - var selected_dynamic_tool_sink = SelectedDynamicToolSinkState{ .allocator = arena }; - var context_notice_sink = ContextNoticeSinkState{ .allocator = arena }; - var command_backend = RunCommandBackendState{ .runtime = ctx }; + var selected_dynamic_tool_sink = SelectedDynamicToolSinkState{ .allocator = result_allocator }; + var context_notice_sink = ContextNoticeSinkState{ .allocator = result_allocator }; + var command_backend = RunCommandBackendState{ + .runtime = ctx, + .result_allocator = result_allocator, + }; var vision_provider = VisionProviderState{ .runtime = ctx, + .result_allocator = result_allocator, .authorized_image_catalog = authorized_image_catalog, }; var subagent_provider = SubagentProviderState{ .runtime = ctx }; var mcp_progress_bridge = McpProgressBridge{ .ctx = ctx }; var mcp_call_status: ?tool_mcp_runtime.CallStatus = null; var mcp_execution_error: ?anyerror = null; + // Tool decode/validate/call scratch and dispatch-owned result bodies live + // on the call allocator; everything that must outlive the call is either + // built on the result allocator (sinks, backend completions) or copied + // out at the bottom of this function. var dispatch_metadata: DispatchMetadata = .{}; - var dispatch_ctx = typedDispatchContextForCall(ctx, arena, call); + var dispatch_ctx = typedDispatchContextForCall(ctx, call_allocator, call); dispatch_metadata.attach(&dispatch_ctx); var result_commit_token: ?result_commit.Token = null; dispatch_ctx.result_commit_sink = &result_commit_token; @@ -808,18 +824,20 @@ fn executeRegisteredTool( &dispatch_metadata.status_detail, ); if (command_backend.execution_error) |err| { - dispatched.deinit(arena); + dispatched.deinit(call_allocator); return err; } if (vision_provider.execution_error) |err| { - dispatched.deinit(arena); + dispatched.deinit(call_allocator); return err; } if (mcp_execution_error) |err| { - dispatched.deinit(arena); + dispatched.deinit(call_allocator); return err; } + const backend_owned = command_backend.completion != null or + vision_provider.completion != null; var execution = if (command_backend.completion) |completion| completion else if (vision_provider.completion) |completion| @@ -828,13 +846,6 @@ fn executeRegisteredTool( toolExecutionResultFromDispatch(dispatched, dispatch_metadata); execution.model_output = dispatched.body; if (dispatch_metadata.status_detail) |detail| execution.status_detail = detail; - if (mcp_call_status == .input_required or - (execution.status == .failure and - tool_mcp_feature_dispatch.isInputRequiredFailure(execution.model_output))) - { - execution.finish_turn = true; - execution.status_detail = "McpInputRequired"; - } execution.selected_dynamic_tool_name = selected_dynamic_tool_sink.name; execution.selected_dynamic_tool_schema_json = selected_dynamic_tool_sink.schema_json; execution.context_notices = context_notice_sink.notices.items; @@ -846,11 +857,36 @@ fn executeRegisteredTool( execution.cancelled = true; } } + // Dispatch-owned memory dies with the call allocator, so the surviving + // result copies it out. Backend completions already own their fields on + // the result allocator; only the dispatch echo needs the copy there. + if (!file_mutation_execution.allocatorsEqual(call_allocator, result_allocator)) { + if (!backend_owned) { + execution.model_output = try result_allocator.dupe(u8, execution.model_output); + if (execution.tool_result_memory) |memory| { + execution.tool_result_memory = try types.dupeToolResultMemory(result_allocator, memory); + } + } + if (execution.status_detail) |detail| { + execution.status_detail = try result_allocator.dupe(u8, detail); + } + } + // Assigned after the copy-out so it stays a static literal: ask retains + // finish-turn error codes past the turn arena (PromptRunResult.error_code + // is never freed), so this must not be duped onto a shorter-lived owner. + if (mcp_call_status == .input_required or + (execution.status == .failure and + tool_mcp_feature_dispatch.isInputRequiredFailure(execution.model_output))) + { + execution.finish_turn = true; + execution.status_detail = "McpInputRequired"; + } return execution; } const RunCommandBackendState = struct { runtime: Context, + result_allocator: Allocator, completion: ?ToolExecutionResult = null, execution_error: ?anyerror = null, }; @@ -874,7 +910,7 @@ fn executeRunCommandBackend( }; const execution = toolRunCommand( state.runtime, - dispatch_ctx.allocator, + state.result_allocator, request, authority, ) catch |err| { @@ -1116,6 +1152,7 @@ fn typedDispatchContextForCall( const VisionProviderState = struct { runtime: Context, + result_allocator: Allocator, authorized_image_catalog: []const types.ImageAttachment, completion: ?ToolExecutionResult = null, execution_error: ?anyerror = null, @@ -1131,7 +1168,7 @@ fn executeVisionProvider( const request = input.as(tool_contracts.vision.VisionRequest); const execution = executeVisionRequest( state, - dispatch_ctx.allocator, + state.result_allocator, request.*, dispatch_ctx.execution_authority orelse { state.execution_error = error.InvalidVisionExecutionAuthority; @@ -3015,6 +3052,45 @@ test "ask_user_question execution uses supplied registry entry and interactive h })); } +test "ordinary tool results survive the call allocator" { + const alloc = std.testing.allocator; + var result_arena_state = std.heap.ArenaAllocator.init(alloc); + defer result_arena_state.deinit(); + + var registered_ask_question = test_builtin_tools.ask_user_question; + registered_ask_question.call = registryOwnedAskQuestionCall; + const tools = [_]tool_dispatch.Tool{registered_ask_question}; + const registry = tool_dispatch.Registry{ .tools = tools[0..] }; + var rt = TestRuntime{ .tool_registry = registry }; + defer rt.deinit(alloc); + + var call_arena_state = std.heap.ArenaAllocator.init(alloc); + const result = blk: { + defer call_arena_state.deinit(); + break :blk try executeToolCallAuthorized(rt.context(), .{ + .call_allocator = call_arena_state.allocator(), + .result_allocator = result_arena_state.allocator(), + .call = .{ + .id = "distinct-owners", + .name = "ask_user_question", + .arguments_json = "not-json", + }, + .authority = .ordinary, + .session_grants = &.{}, + .advertised_dynamic_tool_names = &.{}, + .max_tool_result_bytes = rt.max_tool_result_bytes, + }); + }; + + // The call allocator is gone. Everything the caller may still read has to + // live on the result allocator; today that is model_output, status_detail, + // and tool_result_memory (see the copy-out in executeRegisteredTool). + try std.testing.expectEqual(tool_contracts.ToolExecutionStatus.success, result.status); + try std.testing.expectEqualStrings("registry-owned ask_user_question", result.model_output); + try std.testing.expect(result.status_detail == null); + try std.testing.expect(result.tool_result_memory == null); +} + test "web_fetch execution uses supplied registry entry" { const alloc = std.testing.allocator; var arena_state = std.heap.ArenaAllocator.init(alloc); diff --git a/src/main.zig b/src/main.zig index c49cdd8c5..80878a7a8 100644 --- a/src/main.zig +++ b/src/main.zig @@ -1294,7 +1294,7 @@ const App = struct { self: *App, call: types.ToolCall, ) types.ToolActivityKind { - return tool_dispatch.toolActivityKindForCall(self.alloc, self.toolRegistry(), call); + return tool_dispatch.toolActivityKindForCall(self.toolRegistry(), call); } pub fn commitStartupResumeReplayAnchor(self: *App) !void { @@ -2325,7 +2325,7 @@ const App = struct { call: types.ToolCall, result: types.PersistedToolResult, ) !void { - const activity_kind = tool_dispatch.toolActivityKindForCall(self.alloc, self.toolRegistry(), call); + const activity_kind = tool_dispatch.toolActivityKindForCall(self.toolRegistry(), call); try self.shell.attachHistoricalToolDetail(self.alloc, entry_id, call, activity_kind, result); } @@ -2336,7 +2336,7 @@ const App = struct { result: types.PersistedToolResult, lifecycle_id: types.ToolLifecycleId, ) !void { - const activity_kind = tool_dispatch.toolActivityKindForCall(self.alloc, self.toolRegistry(), call); + const activity_kind = tool_dispatch.toolActivityKindForCall(self.toolRegistry(), call); try self.shell.attachHistoricalToolDetailWithLifecycle( self.alloc, entry_id, @@ -2353,7 +2353,7 @@ const App = struct { call: types.ToolCall, result: types.PersistedToolResult, ) !void { - const activity_kind = tool_dispatch.toolActivityKindForCall(self.alloc, self.toolRegistry(), call); + const activity_kind = tool_dispatch.toolActivityKindForCall(self.toolRegistry(), call); try self.shell.attachHistoricalToolDetailAfterCommandOutput( self.alloc, entry_id,