PranitInitial commit439f83c3 months agoHistory
496 lines · 19.4 KB · Zig
RawDownload
const std = @import("std");const model_capabilities = @import("../../config/model_capabilities.zig");const types = @import("../../shared/types.zig");const token_estimate = @import("../../shared/token_estimate.zig");const debug_trace = @import("../../shared/debug_trace.zig");const diagnostics = @import("../../workspace/diagnostics.zig");const io_mod = @import("../../shared/io.zig"); const runtime_tool_contracts = @import("tool_contracts.zig"); const Allocator = std.mem.Allocator;const ChatMessage = types.ChatMessage;const HistoryTurn = types.HistoryTurn;const ReasoningEffort = types.ReasoningEffort;const ToolCall = types.ToolCall;const TraceContext = debug_trace.TraceContext;const ToolExecutionResult = runtime_tool_contracts.ToolExecutionResult; pub const TokenProgressUpdate = enum {    unchanged,    changed,    invalid,}; const ActiveTokenRequest = struct {    content_estimator: token_estimate.StreamingEstimator = .{},    reasoning_estimator: token_estimate.StreamingEstimator = .{},    tool_input_estimator: token_estimate.StreamingEstimator = .{},     fn outputEstimate(self: ActiveTokenRequest) u64 {        return self.content_estimator.estimate() +|            self.reasoning_estimator.estimate() +|            self.tool_input_estimator.estimate();    }}; pub const GatewayFailureDiagnostics = struct {    schema: []const u8 = "",    request_shape: []const u8 = "",}; pub fn recordGatewayCallMetric(model: []const u8, started_at_ms: i64, status: u16, response_bytes: u32, input_tokens: u32, output_tokens: u32, turn_id: u64, step_id: u64, subagent_id: u64, error_name: []const u8, terminal_stop_reason: []const u8) void {    recordGatewayCallMetricWithDiagnostics(model, started_at_ms, status, response_bytes, input_tokens, output_tokens, turn_id, step_id, subagent_id, error_name, terminal_stop_reason, .{});} pub fn recordGatewayCallMetricWithDiagnostics(    model: []const u8,    started_at_ms: i64,    status: u16,    response_bytes: u32,    input_tokens: u32,    output_tokens: u32,    turn_id: u64,    step_id: u64,    subagent_id: u64,    error_name: []const u8,    terminal_stop_reason: []const u8,    failure_diagnostics: GatewayFailureDiagnostics,) void {    const now_ms = io_mod.milliTimestamp();    const elapsed = if (now_ms > started_at_ms) now_ms - started_at_ms else 0;    var call: diagnostics.NetworkCall = .{        .started_at_ms = started_at_ms,        .duration_ms = @intCast(@min(@as(i64, elapsed), std.math.maxInt(u32))),        .status = status,        .response_bytes = response_bytes,        .input_tokens = input_tokens,        .output_tokens = output_tokens,        .turn_id = turn_id,        .step_id = step_id,        .subagent_id = subagent_id,    };    call.setModel(model);    if (error_name.len > 0) call.setError(error_name);    if (terminal_stop_reason.len > 0) call.setTerminalStopReason(terminal_stop_reason);    if (failure_diagnostics.schema.len > 0) call.setGatewaySchemaDiagnostic(failure_diagnostics.schema);    if (failure_diagnostics.request_shape.len > 0) call.setGatewayRequestShape(failure_diagnostics.request_shape);    diagnostics.recordNetworkCall(call);} pub const TurnSummaryAccumulator = struct {    turn_started_at_ms: i64,    thinking_duration_ms: u64 = 0,    settled_token_progress: types.TurnTokenProgress = .{},    active_token_request: ?ActiveTokenRequest = null,     pub fn init(turn_started_at_ms: i64, submitted_prompt: []const u8) TurnSummaryAccumulator {        var estimator = token_estimate.StreamingEstimator{};        estimator.consume(submitted_prompt);        return .{            .turn_started_at_ms = turn_started_at_ms,            .settled_token_progress = .{                .input_tokens = estimator.estimate(),                .input_exact = submitted_prompt.len == 0,            },        };    }     pub fn addThinkingWait(self: *TurnSummaryAccumulator, started_at_ms: i64, ended_at_ms: i64) void {        if (ended_at_ms <= started_at_ms) return;        self.thinking_duration_ms += @intCast(ended_at_ms - started_at_ms);    }     pub fn prepareTokenRequest(self: *TurnSummaryAccumulator) void {        std.debug.assert(self.active_token_request == null);        self.active_token_request = .{};    }     pub fn consumeReasoning(self: *TurnSummaryAccumulator, text: []const u8) TokenProgressUpdate {        const active = &(self.active_token_request orelse return .invalid);        const before = active.reasoning_estimator.estimate();        active.reasoning_estimator.consume(text);        return if (active.reasoning_estimator.estimate() == before) .unchanged else .changed;    }     pub fn consumeContent(self: *TurnSummaryAccumulator, text: []const u8) TokenProgressUpdate {        const active = &(self.active_token_request orelse return .invalid);        const before = active.content_estimator.estimate();        active.content_estimator.consume(text);        return if (active.content_estimator.estimate() == before) .unchanged else .changed;    }     /// Streamed tool arguments are output tokens the model is producing right    /// now, so they count toward the live estimate the same way response text    /// does. Without this a turn that spends its output budget composing one    /// large tool call reports a frozen count until the provider settles usage.    pub fn consumeToolInput(self: *TurnSummaryAccumulator, text: []const u8) TokenProgressUpdate {        const active = &(self.active_token_request orelse return .invalid);        const before = active.tool_input_estimator.estimate();        active.tool_input_estimator.consume(text);        return if (active.tool_input_estimator.estimate() == before) .unchanged else .changed;    }     pub fn reconcileTokenRequest(self: *TurnSummaryAccumulator, usage: types.Usage, delivery_ambiguous: bool) TokenProgressUpdate {        const before = self.tokenProgress();        const active = self.active_token_request orelse return .invalid;         const output_tokens = usage.output_tokens orelse active.outputEstimate();        self.settleOutputTokens(            output_tokens,            !delivery_ambiguous and usage.output_tokens != null,        );        self.active_token_request = null;        return if (std.meta.eql(before, self.tokenProgress())) .unchanged else .changed;    }     pub fn finishTokenRequestWithoutUsage(self: *TurnSummaryAccumulator, possibly_sent: bool) TokenProgressUpdate {        const before = self.tokenProgress();        const active = self.active_token_request orelse return .unchanged;        self.active_token_request = null;        const output_tokens = active.outputEstimate();        if (possibly_sent or output_tokens > 0) self.settleOutputTokens(output_tokens, false);        return if (std.meta.eql(before, self.tokenProgress())) .unchanged else .changed;    }     pub fn tokenProgress(self: *const TurnSummaryAccumulator) types.TurnTokenProgress {        var progress = self.settled_token_progress;        if (self.active_token_request) |active| {            progress.output_tokens +|= active.outputEstimate();            progress.output_exact = false;        }        return progress;    }     fn settleOutputTokens(self: *TurnSummaryAccumulator, output_tokens: u64, output_exact: bool) void {        self.settled_token_progress.output_tokens +|= output_tokens;        self.settled_token_progress.output_exact = self.settled_token_progress.output_exact and output_exact;    }     pub fn finish(self: *const TurnSummaryAccumulator) types.TurnSummary {        const now_ms = io_mod.milliTimestamp();        const elapsed = if (now_ms > self.turn_started_at_ms) now_ms - self.turn_started_at_ms else 0;        return .{            .thinking_duration_ms = self.thinking_duration_ms,            .turn_duration_ms = @intCast(elapsed),            .token_progress = self.tokenProgress(),        };    }}; pub const StepLimitTrace = struct {    ctx: TraceContext,    step_index: usize,    step_limit: usize,    gateway_message_count: usize,    completed_tool_names: []const []u8,    last_tool_call_name: []const u8,    last_tool_call_id: []const u8,}; pub fn traceStepLimitReached(trace: StepLimitTrace) void {    if (!debug_trace.isScopeEnabled("agent")) return;     const completed_names_owned = formatCompletedToolNamesCompact(std.heap.c_allocator, trace.completed_tool_names) catch null;    defer if (completed_names_owned) |names| std.heap.c_allocator.free(names);    const completed_tool_names = completed_names_owned orelse "unavailable";     debug_trace.logf(        "agent",        "step limit reached turn_id={d} step_id={d} step_index={d} step_limit={d} gateway_messages={d} completed_tool_count={d} completed_tool_names={s} last_tool_call_name={s} last_tool_call_id={s} outcome_kind=step_limit",        .{            trace.ctx.turn_id,            trace.ctx.step_id,            trace.step_index,            trace.step_limit,            trace.gateway_message_count,            trace.completed_tool_names.len,            completed_tool_names,            trace.last_tool_call_name,            trace.last_tool_call_id,        },    );    debug_trace.eventf(        "agent",        "step_limit_reached",        trace.ctx,        "step_index={d} step_limit={d} gateway_messages={d} completed_tool_count={d} completed_tool_names={s} last_tool_call_name={s} last_tool_call_id={s} outcome_kind=step_limit",        .{            trace.step_index,            trace.step_limit,            trace.gateway_message_count,            trace.completed_tool_names.len,            completed_tool_names,            trace.last_tool_call_name,            trace.last_tool_call_id,        },    );} pub fn formatCompletedToolNamesCompact(alloc: Allocator, names: []const []u8) ![]u8 {    if (names.len == 0) return alloc.dupe(u8, "none");     var out: std.Io.Writer.Allocating = .init(alloc);    errdefer out.deinit();    for (names, 0..) |name, i| {        if (i > 0) try out.writer.writeByte(',');        try out.writer.writeAll(name);    }    return out.toOwnedSlice();} pub fn formatHistoryTurnKinds(alloc: Allocator, history: []const HistoryTurn) ![]const u8 {    if (history.len == 0) return alloc.dupe(u8, "none");     var out: std.Io.Writer.Allocating = .init(alloc);    errdefer out.deinit();    for (history, 0..) |turn, i| {        if (i > 0) try out.writer.writeByte(',');        try out.writer.writeAll(historyTurnKindName(turn));    }    return out.toOwnedSlice();} fn historyTurnKindName(turn: HistoryTurn) []const u8 {    return switch (turn) {        .compacted_summary => "compacted_summary",        .assistant => "assistant",        .background_command => "background_command",        .interrupted => "interrupted",    };} pub fn formatMessageRoles(alloc: Allocator, messages: []const ChatMessage) ![]const u8 {    if (messages.len == 0) return alloc.dupe(u8, "none");     var out: std.Io.Writer.Allocating = .init(alloc);    errdefer out.deinit();    for (messages, 0..) |chat_message, i| {        if (i > 0) try out.writer.writeByte(',');        try out.writer.writeAll(@tagName(chat_message.role));    }    return out.toOwnedSlice();} pub fn traceReturnedToolCalls(arena: Allocator, ctx: TraceContext, tool_calls: []const ToolCall) !void {    for (tool_calls) |call| {        const preview_text = debug_trace.keylessJsonPreview(arena, call.arguments_json) catch "<preview-error>";        debug_trace.eventf(            "tool",            "returned_tool_call",            ctx,            "call_id={s} tool_name={s} args_bytes={d} args_preview={s}",            .{ call.id, call.name, call.arguments_json.len, preview_text },        );    }} pub fn traceCancelObserved(ctx: TraceContext, active_tool_known: bool) void {    debug_trace.eventf("interrupt", "cancel_observed", ctx, "active_tool_known={s}", .{if (active_tool_known) "true" else "false"});} pub fn traceGatewayProviderOptions(ctx: TraceContext, model: []const u8, fast_mode: bool, effort: ReasoningEffort, provider_opts: model_capabilities.ResolvedProviderOptions) void {    const reasoning_outcome = if (provider_opts.reasoning != null)        "selected"    else if (effort.isDefault())        "default"    else        "unsupported_or_missing";    const fast_outcome = if (provider_opts.fast)        "selected"    else if (!fast_mode)        "default"    else        "unsupported_or_missing";    debug_trace.eventf(        "gateway",        "provider_options",        ctx,        "model={s} fast_mode={s} effort={s} reasoning={s} fast={s}",        .{            model,            if (fast_mode) "true" else "false",            effort.label(),            reasoning_outcome,            fast_outcome,        },    );} pub fn traceGatewayRequestBuilt(ctx: TraceContext, model: []const u8, payload_bytes: usize, gateway_message_count: usize, tools_json: []const u8) void {    if (!debug_trace.isScopeEnabled("gateway")) return;     var arena_state = std.heap.ArenaAllocator.init(std.heap.c_allocator);    defer arena_state.deinit();    const arena = arena_state.allocator();     if (countToolSchemas(arena, tools_json)) |tool_schema_count| {        debug_trace.eventf("gateway", "request_built", ctx, "payload_bytes={d} model={s} gateway_messages={d} tool_schema_count={d}", .{ payload_bytes, model, gateway_message_count, tool_schema_count });    } else {        debug_trace.eventf("gateway", "request_built", ctx, "payload_bytes={d} model={s} gateway_messages={d} tool_schema_count=unknown", .{ payload_bytes, model, gateway_message_count });    }} fn countToolSchemas(alloc: Allocator, tools_json: []const u8) ?usize {    var parsed = std.json.parseFromSlice(std.json.Value, alloc, tools_json, .{}) catch return null;    defer parsed.deinit();     switch (parsed.value) {        .array => |array| return array.items.len,        .object => |object| {            const tools_value = object.get("tools") orelse return null;            if (tools_value == .array) return tools_value.array.items.len;            return null;        },        else => return null,    }} pub fn toolExecutionResultKind(result: ToolExecutionResult) []const u8 {    if (result.status == .failure) return "error";    if (result.diff_entry != null) return "diff";    if (result.display_output != null) return "display";    return "model_output";} test "gateway call metrics retain provider terminal stop reason" {    diagnostics.resetForTest();    defer diagnostics.resetForTest();     recordGatewayCallMetric("test/model", 0, 200, 12, 3, 4, 1, 1, 0, "", "stop");    recordGatewayCallMetric("test/model", 0, 200, 12, 3, 4, 1, 2, 0, "", "length");    recordGatewayCallMetric("test/model", 0, 200, 12, 3, 4, 1, 3, 0, "", "missing_provider_finish");     var calls: [diagnostics.network_ring_capacity]diagnostics.NetworkCall = undefined;    const count = diagnostics.snapshotNetworkCalls(&calls);    try std.testing.expectEqual(@as(usize, 3), count);    try std.testing.expectEqualStrings("stop", calls[0].terminalStopReason());    try std.testing.expectEqualStrings("length", calls[1].terminalStopReason());    try std.testing.expectEqualStrings("missing_provider_finish", calls[2].terminalStopReason());} test "turn token progress counts only the submitted prompt once" {    var accumulator = TurnSummaryAccumulator.init(0, "open it for me");    accumulator.prepareTokenRequest();    _ = accumulator.reconcileTokenRequest(.{        .input_tokens = 16_000,        .output_tokens = 5,    }, true);     try std.testing.expectEqual(types.TurnTokenProgress{        .input_tokens = 4,        .output_tokens = 5,        .input_exact = false,        .output_exact = false,    }, accumulator.tokenProgress());} test "turn token progress advances while the model composes a tool call" {    var accumulator = TurnSummaryAccumulator.init(0, "");    accumulator.prepareTokenRequest();    _ = accumulator.consumeContent("writing the file now ");    const before_tool_input = accumulator.tokenProgress().output_tokens;     try std.testing.expectEqual(TokenProgressUpdate.changed, accumulator.consumeToolInput(        "{\"path\":\"docs/big.md\",\"content\":\"line one line two line three\"",    ));    try std.testing.expect(accumulator.tokenProgress().output_tokens > before_tool_input);     _ = accumulator.reconcileTokenRequest(.{ .output_tokens = 4_096 }, false);    try std.testing.expectEqual(@as(u64, 4_096), accumulator.tokenProgress().output_tokens);    try std.testing.expect(accumulator.tokenProgress().output_exact);} test "turn token progress honors aggregate completion ambiguity" {    var accumulator = TurnSummaryAccumulator.init(0, "abcdefgh");    accumulator.prepareTokenRequest();    _ = accumulator.reconcileTokenRequest(.{        .input_tokens = 100,        .output_tokens = 20,    }, true);     try std.testing.expectEqual(types.TurnTokenProgress{        .input_tokens = 2,        .output_tokens = 20,        .input_exact = false,        .output_exact = false,    }, accumulator.tokenProgress());} test "turn token progress uses exact provider output on unambiguous completion" {    var accumulator = TurnSummaryAccumulator.init(0, "abcdefgh");    accumulator.prepareTokenRequest();    _ = accumulator.reconcileTokenRequest(.{        .input_tokens = 100,        .output_tokens = 20,    }, false);     try std.testing.expectEqual(types.TurnTokenProgress{        .input_tokens = 2,        .output_tokens = 20,        .input_exact = false,    }, accumulator.tokenProgress());} test "turn token progress accumulates output steps while counting submitted input once" {    var accumulator = TurnSummaryAccumulator.init(0, "submitted prompt");     accumulator.prepareTokenRequest();    _ = accumulator.reconcileTokenRequest(.{        .input_tokens = 10,        .output_tokens = 5,    }, false);     accumulator.prepareTokenRequest();    _ = accumulator.reconcileTokenRequest(.{        .input_tokens = 20,        .output_tokens = 7,    }, false);     try std.testing.expectEqual(types.TurnTokenProgress{        .input_tokens = 5,        .output_tokens = 12,        .input_exact = false,    }, accumulator.tokenProgress());    try std.testing.expectEqual(accumulator.tokenProgress(), accumulator.finish().token_progress);} test "turn token progress keeps missing usage and terminal delivery approximate" {    var accumulator = TurnSummaryAccumulator.init(0, "abcdefgh");    accumulator.prepareTokenRequest();    _ = accumulator.consumeReasoning("reasoning");    _ = accumulator.consumeContent("answer");    _ = accumulator.reconcileTokenRequest(.{ .input_tokens = 40 }, false);     try std.testing.expectEqual(types.TurnTokenProgress{        .input_tokens = 2,        .output_tokens = 5,        .input_exact = false,        .output_exact = false,    }, accumulator.tokenProgress());     accumulator.prepareTokenRequest();    _ = accumulator.finishTokenRequestWithoutUsage(true);     try std.testing.expectEqual(types.TurnTokenProgress{        .input_tokens = 2,        .output_tokens = 5,        .input_exact = false,        .output_exact = false,    }, accumulator.tokenProgress());} test "turn token progress discards a definitely unsent request" {    var accumulator = TurnSummaryAccumulator.init(0, "abcdefgh");    accumulator.prepareTokenRequest();    try std.testing.expect(!accumulator.tokenProgress().output_exact);     try std.testing.expectEqual(        TokenProgressUpdate.changed,        accumulator.finishTokenRequestWithoutUsage(false),    );    try std.testing.expectEqual(types.TurnTokenProgress{        .input_tokens = 2,        .input_exact = false,    }, accumulator.tokenProgress());}