2140 lines · 87.4 KB · Zig
RawDownload
const std = @import("std");const command_admission = @import("../../permissions/command_admission.zig");const permission_auto_classifier = @import("../../permissions/auto_classifier.zig");const types = @import("../../shared/types.zig");const text_utils = @import("../../shared/text_utils.zig");const debug_trace = @import("../../shared/debug_trace.zig");const token_estimate = @import("../../shared/token_estimate.zig");const assistant_presentation = @import("../assistant_presentation.zig");const io_mod = @import("../../shared/io.zig");const file_mutation = @import("../../tooling/file_mutation.zig");const tool_dispatch = @import("../../tooling/tool_dispatch.zig");const worker_runtime = @import("../worker_runtime.zig");const test_builtin_tools = if (@import("builtin").is_test)    @import("../../../builtins/tools.zig")else    struct {}; const runtime_deps = @import("deps.zig");const runtime_telemetry = @import("telemetry.zig");const runtime_tool_contracts = @import("tool_contracts.zig");const runtime_tool_presentation = @import("tool_presentation.zig");const model_response_recovery = @import("model_response_recovery.zig"); const Allocator = std.mem.Allocator;const ChatMessage = types.ChatMessage;const HistoryTurn = types.HistoryTurn;const PermissionGrant = types.PermissionGrant;const PermissionMode = types.PermissionMode;const ToolCall = types.ToolCall;const WorkerEvent = worker_runtime.WorkerEvent;const sanitizeAssistantText = text_utils.sanitizeAssistantText;const AgentRuntimeDeps = runtime_deps.AgentRuntimeDeps;const DiffEntryPayload = runtime_tool_contracts.DiffEntryPayload;const SecondaryPublicationReport = runtime_tool_contracts.SecondaryPublicationReport;const ToolExecutionResult = runtime_tool_contracts.ToolExecutionResult;const TablePayload = assistant_presentation.TablePayload;const CodeBlockPayload = assistant_presentation.CodeBlockPayload; const test_tools = [_]tool_dispatch.Tool{    test_builtin_tools.read_file,    test_builtin_tools.write_file,    test_builtin_tools.edit_file,    test_builtin_tools.terminal,    test_builtin_tools.ask_user_question,};const test_tool_registry = tool_dispatch.Registry{ .tools = test_tools[0..] }; pub const SemanticPresentationSink = struct {    ctx: *anyopaque,    /// Takes ownership of `table` only when it returns successfully.    table: *const fn (ctx: *anyopaque, table: TablePayload) anyerror!void,    /// Takes ownership of `block` only when it returns successfully.    code_block: *const fn (ctx: *anyopaque, block: CodeBlockPayload) anyerror!void,    thematic_rule: *const fn (ctx: *anyopaque) anyerror!void,}; pub const StreamChunkContext = struct {    hooks: *const AgentRuntimeDeps,    flush_assistant_stream_per_content_chunk: bool = false,    semantic_presentation: ?SemanticPresentationSink = null,    token_progress: ?*runtime_telemetry.TurnSummaryAccumulator = null,    turn_id: u64,    step_id: u64 = 0,    alloc: Allocator = std.heap.c_allocator,    saw_visible_text: bool = false,    saw_tool_start: bool = false,    saw_provider_tool_start: bool = false,    saw_visible_text_after_tool_start: bool = false,    last_byte_was_newline: bool = false,    trailing_newline_count: u8 = 0,    first_model_output_at_ms: ?i64 = null,    markdown: assistant_presentation.MarkdownProcessor = .{},    raw_text: std.ArrayList(u8) = .empty,    continuation_pending: std.ArrayList(u8) = .empty,    continuation_prefix_len: usize = 0,    continuation_probe_remaining: usize = 0,    continuation_resolved: bool = true,    initial_line_prefix: std.ArrayList(u8) = .empty,    provisional_statuses: runtime_tool_presentation.ProvisionalToolStatuses = .{},     fn markModelOutput(self: *StreamChunkContext) void {        if (self.first_model_output_at_ms == null) self.first_model_output_at_ms = io_mod.milliTimestamp();    }     fn recordTextOutput(self: *StreamChunkContext, text: []const u8) void {        if (text.len > 0 and self.saw_tool_start) {            self.saw_visible_text_after_tool_start = true;        }        var trailing: usize = 0;        while (trailing < text.len and text[text.len - trailing - 1] == '\n') : (trailing += 1) {}        if (trailing == 0) {            self.last_byte_was_newline = false;            self.trailing_newline_count = 0;            return;        }         self.last_byte_was_newline = true;        const bounded: u8 = @intCast(@min(trailing, @as(usize, 2)));        self.trailing_newline_count = if (trailing == text.len)            @min(2, self.trailing_newline_count + bounded)        else            bounded;    }     fn recordSemanticBoundary(self: *StreamChunkContext) void {        if (self.saw_tool_start) {            self.saw_visible_text_after_tool_start = true;        }        self.last_byte_was_newline = true;        self.trailing_newline_count = 1;    }     pub fn deinit(self: *StreamChunkContext) void {        self.markdown.deinit(self.alloc);        self.raw_text.deinit(self.alloc);        self.continuation_pending.deinit(self.alloc);        self.initial_line_prefix.deinit(self.alloc);        self.provisional_statuses.deinit(self.alloc);    }     pub fn beginRecoveryAttempt(self: *StreamChunkContext) void {        self.continuation_pending.clearRetainingCapacity();        self.continuation_prefix_len = self.raw_text.items.len;        self.continuation_probe_remaining = self.continuation_prefix_len;        self.continuation_resolved = self.continuation_prefix_len == 0;        self.saw_tool_start = false;        self.saw_provider_tool_start = false;        self.saw_visible_text_after_tool_start = false;        self.first_model_output_at_ms = null;    }     /// Restores durable partial source before a restarted turn sends anything.    /// The following attempt can then suppress an exact repeated prefix    /// against these same raw bytes. Interactive surfaces that already show    /// the partial source restore presentation state without emitting it again.    pub fn restoreRecoverySource(        self: *StreamChunkContext,        source: []const u8,        already_presented: bool,    ) !void {        if (source.len == 0) return;        if (already_presented) {            try self.restorePresentedRecoverySource(source);            return;        }        try streamAssistantChunk(self, source);        try flushAssistantStream(self);        try self.markdown.restorePresentedPrefix(self.alloc, source);    }     fn restorePresentedRecoverySource(        self: *StreamChunkContext,        source: []const u8,    ) !void {        try self.raw_text.appendSlice(self.alloc, source);        try self.hooks.push_text(self.hooks.ctx, .{ .assistant_source = source });         var start: usize = 0;        while (start < source.len and isTrimmedAssistantPrefixByte(source[start])) : (start += 1) {            switch (source[start]) {                ' ', '\t' => try self.initial_line_prefix.append(self.alloc, source[start]),                '\n' => self.initial_line_prefix.clearRetainingCapacity(),                else => {},            }        }        if (start == source.len) {            self.initial_line_prefix.clearRetainingCapacity();            return;        }        self.saw_visible_text = true;         self.initial_line_prefix.clearRetainingCapacity();        try self.markdown.restorePresentedPrefix(self.alloc, source);        self.recordTextOutput(source);    }}; pub fn onStreamContentChunk(ctx: *anyopaque, chunk: []const u8) void {    const stream_ctx: *StreamChunkContext = @ptrCast(@alignCast(ctx));    stream_ctx.markModelOutput();    if (stream_ctx.token_progress) |progress| {        pushTokenProgressUpdate(stream_ctx, progress.consumeContent(chunk)) catch |err| {            debug_trace.logf("agent", "token progress publication failed source=content err={s}", .{@errorName(err)});        };    }    streamAssistantChunk(stream_ctx, chunk) catch {};    if (stream_ctx.flush_assistant_stream_per_content_chunk) {        flushAssistantStream(stream_ctx) catch {};    }} pub fn onStreamReasoningChunk(ctx: *anyopaque, chunk: []const u8) void {    const stream_ctx: *StreamChunkContext = @ptrCast(@alignCast(ctx));    if (stream_ctx.token_progress) |progress| {        pushTokenProgressUpdate(stream_ctx, progress.consumeReasoning(chunk)) catch |err| {            debug_trace.logf("agent", "token progress publication failed source=reasoning err={s}", .{@errorName(err)});        };    }} /// Tools that publish no provisional status leave the activity row with/// nothing to say while their arguments stream, which can take as long as the/// response did. Tell the shell the turn is composing so the row stays alive.fn publishToolPayloadStarted(stream_ctx: *StreamChunkContext) void {    stream_ctx.hooks.push_event(stream_ctx.hooks.ctx, .tool_payload_started) catch |err| {        debug_trace.logf("agent", "tool payload notice publication failed err={s}", .{@errorName(err)});    };} pub fn onStreamToolInputChunk(ctx: *anyopaque, chunk: []const u8) void {    const stream_ctx: *StreamChunkContext = @ptrCast(@alignCast(ctx));    if (stream_ctx.token_progress) |progress| {        pushTokenProgressUpdate(stream_ctx, progress.consumeToolInput(chunk)) catch |err| {            debug_trace.logf("agent", "token progress publication failed source=tool_input err={s}", .{@errorName(err)});        };    }} pub fn pushTokenProgressUpdate(stream_ctx: *StreamChunkContext, update: runtime_telemetry.TokenProgressUpdate) !void {    switch (update) {        .unchanged => {},        .changed => {            const progress = stream_ctx.token_progress orelse return;            try stream_ctx.hooks.push_event(stream_ctx.hooks.ctx, .{ .turn_token_update = progress.tokenProgress() });        },        .invalid => debug_trace.logf("agent", "token progress observation ignored reason=invalid_sequence", .{}),    }} pub fn onStreamToolStart(ctx: *anyopaque, tool_id: []const u8, tool_name: []const u8, label_value: ?[]const u8) void {    const stream_ctx: *StreamChunkContext = @ptrCast(@alignCast(ctx));    const first_tool_in_step = !stream_ctx.saw_tool_start;    recordStreamToolStart(ctx, tool_name);    const preflight = runtime_tool_presentation.ProvisionalToolStatuses.preflight(stream_ctx.hooks.tool_registry, tool_name) orelse {        publishToolPayloadStarted(stream_ctx);        return;    };    stream_ctx.markModelOutput();    flushAssistantStream(stream_ctx) catch {};    if (stream_ctx.saw_visible_text and !stream_ctx.last_byte_was_newline) {        stream_ctx.hooks.push_text(stream_ctx.hooks.ctx, .{ .assistant_rendered = "\n" }) catch {};        stream_ctx.recordTextOutput("\n");    }    if (first_tool_in_step and        stream_ctx.saw_visible_text and        stream_ctx.turn_id != 0 and        stream_ctx.step_id != 0)    {        stream_ctx.provisional_statuses.presentation_group_id = .{            .turn_id = stream_ctx.turn_id,            .anchor_step_id = stream_ctx.step_id,        };    }    switch (preflight) {        // Published last: the flush above can push the response's trailing        // newline, and any text landing after the notice reopens the response        // row.        .ineligible => publishToolPayloadStarted(stream_ctx),        .eligible => |metadata| stream_ctx.provisional_statuses.publish(            stream_ctx.hooks,            stream_ctx.alloc,            stream_ctx.turn_id,            tool_id,            tool_name,            metadata.activity_kind,            metadata.action_label,            label_value,        ) catch {},    }} pub fn recordStreamToolStart(ctx: *anyopaque, tool_name: []const u8) void {    const stream_ctx: *StreamChunkContext = @ptrCast(@alignCast(ctx));    stream_ctx.saw_tool_start = true;    if (runtime_tool_presentation.streamStartMayHaveExecutedAtProvider(        stream_ctx.hooks.tool_registry,        tool_name,    )) {        stream_ctx.saw_provider_tool_start = true;    }} fn streamAssistantChunk(stream_ctx: *StreamChunkContext, chunk: []const u8) !void {    if (!stream_ctx.continuation_resolved) {        try stream_ctx.continuation_pending.appendSlice(stream_ctx.alloc, chunk);        const overlap = try continuationOverlapIfResolved(stream_ctx, false) orelse return;        const novel = stream_ctx.continuation_pending.items[overlap..];        try streamAssistantChunkResolved(stream_ctx, novel);        stream_ctx.continuation_pending.clearRetainingCapacity();        return;    }    try streamAssistantChunkResolved(stream_ctx, chunk);} fn streamAssistantChunkResolved(stream_ctx: *StreamChunkContext, chunk: []const u8) !void {    if (chunk.len == 0) return;    const alloc = stream_ctx.alloc;    try stream_ctx.raw_text.appendSlice(alloc, chunk);    try stream_ctx.hooks.push_text(stream_ctx.hooks.ctx, .{ .assistant_source = chunk });     var start: usize = 0;    var initial_prefix_to_emit: ?[]const u8 = null;    var clear_initial_line_prefix = false;    defer if (clear_initial_line_prefix) stream_ctx.initial_line_prefix.clearRetainingCapacity();    if (!stream_ctx.saw_visible_text) {        while (start < chunk.len and isTrimmedAssistantPrefixByte(chunk[start])) : (start += 1) {            switch (chunk[start]) {                ' ', '\t' => try stream_ctx.initial_line_prefix.append(alloc, chunk[start]),                '\n' => stream_ctx.initial_line_prefix.clearRetainingCapacity(),                else => {},            }        }        if (start == chunk.len) return;        if (hasResponseLeadingIndentedCodePrefix(stream_ctx.initial_line_prefix.items)) {            initial_prefix_to_emit = stream_ctx.initial_line_prefix.items;        }        stream_ctx.saw_visible_text = true;        clear_initial_line_prefix = true;    }     var out: std.ArrayList(u8) = .empty;    defer out.deinit(alloc);     var completion = assistant_presentation.TableCompletion{        .ctx = stream_ctx,        .deliver = deliverSemanticTable,    };    var code_completion = assistant_presentation.CodeBlockCompletion{        .ctx = stream_ctx,        .deliver = deliverSemanticCodeBlock,    };    var thematic_rule_completion = assistant_presentation.ThematicRuleCompletion{        .ctx = stream_ctx,        .deliver = deliverSemanticThematicRule,    };    const table_completion: ?*const assistant_presentation.TableCompletion = if (stream_ctx.semantic_presentation != null)        &completion    else        null;    const semantic_code_completion: ?*const assistant_presentation.CodeBlockCompletion = if (stream_ctx.semantic_presentation != null)        &code_completion    else        null;    const semantic_thematic_rule_completion: ?*const assistant_presentation.ThematicRuleCompletion = if (stream_ctx.semantic_presentation != null)        &thematic_rule_completion    else        null;    if (initial_prefix_to_emit) |prefix| {        try stream_ctx.markdown.pushWithCompletions(alloc, prefix, &out, .{            .table = table_completion,            .code = semantic_code_completion,            .thematic_rule = semantic_thematic_rule_completion,        });    }    try stream_ctx.markdown.pushWithCompletions(alloc, chunk[start..], &out, .{        .table = table_completion,        .code = semantic_code_completion,        .thematic_rule = semantic_thematic_rule_completion,    });    if (out.items.len == 0) return;    try stream_ctx.hooks.push_text(stream_ctx.hooks.ctx, .{ .assistant_rendered = out.items });    stream_ctx.recordTextOutput(out.items);} pub fn flushAssistantStream(stream_ctx: *StreamChunkContext) !void {    if (!stream_ctx.continuation_resolved) {        const overlap = (try continuationOverlapIfResolved(stream_ctx, true)).?;        const novel = stream_ctx.continuation_pending.items[overlap..];        try streamAssistantChunkResolved(stream_ctx, novel);        stream_ctx.continuation_pending.clearRetainingCapacity();    }    const alloc = stream_ctx.alloc;    stream_ctx.initial_line_prefix.clearRetainingCapacity();    var out: std.ArrayList(u8) = .empty;    defer out.deinit(alloc);     var completion = assistant_presentation.TableCompletion{        .ctx = stream_ctx,        .deliver = deliverSemanticTable,    };    var code_completion = assistant_presentation.CodeBlockCompletion{        .ctx = stream_ctx,        .deliver = deliverSemanticCodeBlock,    };    var thematic_rule_completion = assistant_presentation.ThematicRuleCompletion{        .ctx = stream_ctx,        .deliver = deliverSemanticThematicRule,    };    const table_completion: ?*const assistant_presentation.TableCompletion = if (stream_ctx.semantic_presentation != null)        &completion    else        null;    const semantic_code_completion: ?*const assistant_presentation.CodeBlockCompletion = if (stream_ctx.semantic_presentation != null)        &code_completion    else        null;    const semantic_thematic_rule_completion: ?*const assistant_presentation.ThematicRuleCompletion = if (stream_ctx.semantic_presentation != null)        &thematic_rule_completion    else        null;    try stream_ctx.markdown.flushWithCompletions(alloc, &out, .{        .table = table_completion,        .code = semantic_code_completion,        .thematic_rule = semantic_thematic_rule_completion,    });    if (out.items.len == 0) return;    const text = if (stream_ctx.trailing_newline_count >= 2 and beginsFootnoteProjection(out.items))        out.items[1..]    else        out.items;    try stream_ctx.hooks.push_text(stream_ctx.hooks.ctx, .{ .assistant_rendered = text });    stream_ctx.recordTextOutput(text);} fn continuationOverlapIfResolved(    stream_ctx: *StreamChunkContext,    final: bool,) !?usize {    const existing = stream_ctx.raw_text.items[0..stream_ctx.continuation_prefix_len];    const incoming = stream_ctx.continuation_pending.items;    if (!final) {        const probe = model_response_recovery.probeContinuationExtension(            existing,            incoming,            stream_ctx.continuation_probe_remaining,        );        stream_ctx.continuation_probe_remaining =            stream_ctx.continuation_probe_remaining -| probe.comparisons;        switch (probe.outcome) {            .may_extend, .budget_exhausted => return null,            .cannot_extend => {},        }    }    stream_ctx.continuation_resolved = true;    return try model_response_recovery.exactContinuationOverlap(        stream_ctx.alloc,        existing,        incoming,    );} fn beginsFootnoteProjection(text: []const u8) bool {    return std.mem.startsWith(u8, text, "\n\x1b[2m[");} fn deliverSemanticTable(raw: *anyopaque, table: TablePayload, out: *std.ArrayList(u8)) !void {    const stream_ctx: *StreamChunkContext = @ptrCast(@alignCast(raw));    const alloc = stream_ctx.alloc;     if (out.items.len > 0) {        try stream_ctx.hooks.push_text(stream_ctx.hooks.ctx, .{ .assistant_rendered = out.items });        stream_ctx.recordTextOutput(out.items);        out.clearRetainingCapacity();    }     var compatibility: std.ArrayList(u8) = .empty;    defer compatibility.deinit(alloc);    try assistant_presentation.renderTablePayload(alloc, table, &compatibility);    const sink = stream_ctx.semantic_presentation orelse return error.MissingTableSink;    try sink.table(sink.ctx, table);    stream_ctx.recordSemanticBoundary();} fn deliverSemanticCodeBlock(raw: *anyopaque, block: CodeBlockPayload, out: *std.ArrayList(u8)) !void {    const stream_ctx: *StreamChunkContext = @ptrCast(@alignCast(raw));    const alloc = stream_ctx.alloc;     if (out.items.len > 0) {        try stream_ctx.hooks.push_text(stream_ctx.hooks.ctx, .{ .assistant_rendered = out.items });        stream_ctx.recordTextOutput(out.items);        out.clearRetainingCapacity();    }     var compatibility: std.ArrayList(u8) = .empty;    defer compatibility.deinit(alloc);    try assistant_presentation.renderCodeBlockPayload(alloc, block, &compatibility);    const sink = stream_ctx.semantic_presentation orelse return error.MissingCodeBlockSink;    try sink.code_block(sink.ctx, block);    stream_ctx.recordSemanticBoundary();} fn deliverSemanticThematicRule(raw: *anyopaque, out: *std.ArrayList(u8)) !void {    const stream_ctx: *StreamChunkContext = @ptrCast(@alignCast(raw));    const alloc = stream_ctx.alloc;     if (out.items.len > 0) {        try stream_ctx.hooks.push_text(stream_ctx.hooks.ctx, .{ .assistant_rendered = out.items });        stream_ctx.recordTextOutput(out.items);        out.clearRetainingCapacity();    }     var compatibility: std.ArrayList(u8) = .empty;    defer compatibility.deinit(alloc);    try assistant_presentation.writeHorizontalRule(alloc, &compatibility);    const sink = stream_ctx.semantic_presentation orelse return error.MissingThematicRuleSink;    try sink.thematic_rule(sink.ctx);    stream_ctx.recordSemanticBoundary();} pub fn emitProviderLengthNotice(hooks: *const AgentRuntimeDeps, arena: Allocator, disposition: types.ProviderCompletionDisposition) !void {

Lines 1–500 of 2140.