4823 lines · 177.2 KB · Zig
RawDownload
const std = @import("std");const builtin = @import("builtin");const secret = @import("../auth/secret.zig");const debug_trace = @import("../shared/debug_trace.zig");const generation_fact_codec = @import("generation_fact_codec.zig");const generation_usage = @import("generation_usage_provider.zig");const io_mod = @import("../shared/io.zig");const types = @import("../shared/types.zig");pub const usage_report = @import("usage_report.zig"); const Allocator = std.mem.Allocator;const Sha256 = std.crypto.hash.sha2.Sha256; const max_models: usize = 32;const max_pending_generations: usize = 16;const max_publication_backlog: usize = 16;const max_usage_incidents: usize = 16;const max_model_bytes: usize = 1024;const max_origin_bytes: usize = 2048;const max_team_bytes: usize = 255;const max_identifier_bytes: usize = 8 * 1024;const max_active_invocations: usize = 64;const shutdown_reconciliation_budget_ms: usize = 250;pub const max_snapshot_bytes: usize = 256 * 1024; /// Completeness of the saved billing window rendered by `/cost`.pub const Availability = enum {    /// Every billable invocation has an authoritative generation record.    complete,    /// Known generation identities still require reconciliation.    pending,    /// Billing may have occurred without a recoverable authoritative record.    incomplete,    /// The persisted session predates usage accounting.    legacy,}; /// Billing consequence of one logical Gateway invocation across retries.pub const DeliveryOutcome = enum {    /// Delivery was proven not to create a billable generation.    unbilled,    /// A validated generation identity can be reconciled.    observed_generation,    /// Delivery succeeded but no usable generation identity was returned.    possibly_billed_without_identity,    /// Transport evidence cannot prove whether a billable generation was created.    ambiguous_delivery,}; /// Host-owned durable sink. Its context and allocator remain borrowed until/// the sink is unset; `persist` must not retain the snapshot, checkpoint, or/// reconfigure the sink.pub const UsageCheckpointSink = struct {    context: *anyopaque,    allocator: Allocator,    persist: *const fn (context: *anyopaque, snapshot: Snapshot) anyerror!void,}; /// Host-owned local-profile publication sink. The callback receives a borrowed/// event and must return only after it is durably accepted or deduplicated.pub const ProfilePublicationSink = struct {    context: *anyopaque,    allocator: Allocator,    publish: *const fn (        context: *anyopaque,        event: usage_report.ProfileEvent,    ) anyerror!void,}; const ProfilePublicationBatch = struct {    pending: []usage_report.PendingMarker,    incidents: []usage_report.Incident,    facts: []usage_report.GenerationFact,    checkpoint_changed: bool,     fn deinit(self: *ProfilePublicationBatch, alloc: Allocator) void {        for (self.pending) |*marker| marker.deinit(alloc);        alloc.free(self.pending);        alloc.free(self.incidents);        for (self.facts) |*fact| fact.deinit(alloc);        alloc.free(self.facts);        self.* = undefined;    }}; /// One Gateway invocation. `begin` durably reserves it before network I/O;/// every successful reservation must terminate through `fail` or `complete`.pub const GatewayObservation = struct {    usage: ?*Usage,    sequence: u64 = 0,    started_at_ms: i64,     pub fn begin(usage: ?*Usage) !GatewayObservation {        return .{            .usage = usage,            .sequence = if (usage) |ledger| try ledger.reserveInvocationDurably() else 0,            .started_at_ms = io_mod.milliTimestamp(),        };    }     pub fn fail(self: GatewayObservation, outcome: DeliveryOutcome) !void {        const ledger = self.usage orelse return;        try ledger.finishInvocationDurably(self.sequence, self.elapsedMs(), outcome);    }     pub fn complete(        self: GatewayObservation,        alloc: Allocator,        status: std.http.Status,        completion: types.GatewayCompletion,        origin: []const u8,        team: ?[]const u8,    ) !void {        const ledger = self.usage orelse return;        const generation_id = completion.generation_id;        const delivery: DeliveryOutcome = if (completion.delivery_ambiguous)            .ambiguous_delivery        else if (status == .ok)            if (!completion.generation_metadata_invalid and generation_id != null)                .observed_generation            else                .possibly_billed_without_identity        else            .unbilled;        const id = generation_id orelse {            try ledger.finishInvocationDurably(self.sequence, self.elapsedMs(), delivery);            if (delivery == .possibly_billed_without_identity) {                debug_trace.logf(                    "session",                    "usage billing incomplete sequence={d} reason=generation_identity_missing",                    .{self.sequence},                );            }            return;        };        if (completion.generation_metadata_invalid or status != .ok) {            try ledger.finishInvocationDurably(self.sequence, self.elapsedMs(), delivery);            debug_trace.logf(                "session",                "usage generation ignored sequence={d} status={d} metadata_invalid={s}",                .{                    self.sequence,                    @intFromEnum(status),                    if (completion.generation_metadata_invalid) "true" else "false",                },            );            return;        }        try ledger.finishObservedInvocationDurably(            alloc,            self.sequence,            self.elapsedMs(),            delivery,            id,            origin,            team,        );        debug_trace.logf(            "session",            "usage generation queued sequence={d} id={s}",            .{ self.sequence, id },        );        if (completion.billing) |billing| {            ledger.applyGatewayBilling(alloc, id, billing) catch |err| {                debug_trace.logf(                    "session",                    "usage stream billing apply failed id={s} reason={s}",                    .{ id, @errorName(err) },                );                ledger.flushProfilePublications();                return;            };            debug_trace.logf(                "session",                "usage stream billing settled sequence={d} id={s}",                .{ self.sequence, id },            );        } else {            ledger.flushProfilePublications();        }    }     fn elapsedMs(self: GatewayObservation) u64 {        const ended_at_ms = io_mod.milliTimestamp();        if (ended_at_ms <= self.started_at_ms) return 0;        return std.math.cast(u64, ended_at_ms - self.started_at_ms) orelse            std.math.maxInt(u64);    }}; const GenerationRecord = generation_usage.Record; pub const ModelAggregate = struct {    model: []u8,    first_sequence: u64,    total_cost: f64 = 0,    input_tokens: u64 = 0,    output_tokens: u64 = 0,    cache_read_tokens: u64 = 0,    cache_write_tokens: u64 = 0,    reasoning_tokens: ?u64 = null,    request_count: ?u64 = null,    billable_web_search_calls: u64 = 0,     fn deinit(self: *ModelAggregate, alloc: Allocator) void {        alloc.free(self.model);        self.* = undefined;    }     fn dupe(self: ModelAggregate, alloc: Allocator) Allocator.Error!ModelAggregate {        return .{            .model = try alloc.dupe(u8, self.model),            .first_sequence = self.first_sequence,            .total_cost = self.total_cost,            .input_tokens = self.input_tokens,            .output_tokens = self.output_tokens,            .cache_read_tokens = self.cache_read_tokens,            .cache_write_tokens = self.cache_write_tokens,            .reasoning_tokens = self.reasoning_tokens,            .request_count = self.request_count,            .billable_web_search_calls = self.billable_web_search_calls,        };    }}; pub const PendingGeneration = struct {    id: []u8,    sequence: u64,    origin: []u8,    team: ?[]u8,    observed_at_ms: ?i64 = null,     fn deinit(self: *PendingGeneration, alloc: Allocator) void {        alloc.free(self.id);        alloc.free(self.origin);        if (self.team) |team| alloc.free(team);        self.* = undefined;    }     fn dupe(self: PendingGeneration, alloc: Allocator) Allocator.Error!PendingGeneration {        const id = try alloc.dupe(u8, self.id);        errdefer alloc.free(id);        const origin = try alloc.dupe(u8, self.origin);        errdefer alloc.free(origin);        const team = if (self.team) |value| try alloc.dupe(u8, value) else null;        return .{            .id = id,            .sequence = self.sequence,            .origin = origin,            .team = team,            .observed_at_ms = self.observed_at_ms,        };    }}; pub const Snapshot = struct {    billing: Availability,    api_duration_complete: bool,    wall_duration_complete: bool,    code_complete: bool,    next_sequence: u64,    settled_through_sequence: u64,    api_duration_ms: u64,    wall_duration_ms: u64,    total_cost: f64,    input_tokens: u64,    output_tokens: u64,    cache_read_tokens: u64,    cache_write_tokens: u64,    reasoning_tokens: ?u64 = null,    request_count: ?u64 = null,    billable_web_search_calls: u64,    lines_added: u64,    lines_removed: u64,    models: []ModelAggregate,    pending: []PendingGeneration,    publication_backlog: []usage_report.GenerationFact = &.{},    incidents: []usage_report.Incident = &.{},     pub fn deinit(self: *Snapshot, alloc: Allocator) void {        for (self.models) |*model| model.deinit(alloc);        alloc.free(self.models);        for (self.pending) |*generation| generation.deinit(alloc);        alloc.free(self.pending);        for (self.publication_backlog) |*fact| fact.deinit(alloc);        if (self.publication_backlog.len > 0) alloc.free(self.publication_backlog);        if (self.incidents.len > 0) alloc.free(self.incidents);        self.* = undefined;    }}; /// True when a durable session can still contribute profile usage that is not/// proven present in the profile ledger.pub fn needsProfileRecovery(snapshot: Snapshot) bool {    return snapshot.settled_through_sequence != snapshot.next_sequence - 1 or        snapshot.pending.len > 0 or        snapshot.publication_backlog.len > 0 or        snapshot.incidents.len > 0;} pub const Usage = struct {    mutex: std.Io.Mutex = .init,    checkpoint_mutex: std.Io.Mutex = .init,    checkpoint_sink: ?UsageCheckpointSink = null,    publication_mutex: std.Io.Mutex = .init,    publication_sink: ?ProfilePublicationSink = null,    billing: Availability,    api_duration_complete: bool,    wall_duration_complete: bool,    code_complete: bool,    next_sequence: u64 = 1,    settled_through_sequence: u64 = 0,    api_duration_ms: u64 = 0,    wall_duration_ms: u64 = 0,    active_started_at_ms: i64 = 0,    active_sequences: [max_active_invocations]u64 = undefined,    active_sequence_count: usize = 0,    total_cost: f64 = 0,    input_tokens: u64 = 0,    output_tokens: u64 = 0,    cache_read_tokens: u64 = 0,    cache_write_tokens: u64 = 0,    reasoning_tokens: ?u64,    request_count: ?u64,    billable_web_search_calls: u64 = 0,    lines_added: u64 = 0,    lines_removed: u64 = 0,    models: std.ArrayList(ModelAggregate) = .empty,    pending: std.ArrayList(PendingGeneration) = .empty,    publication_backlog: std.ArrayList(usage_report.GenerationFact) = .empty,    incidents: [max_usage_incidents]usage_report.Incident = undefined,    incident_count: usize = 0,    dirty: bool = false,    reconciliation_mutex: std.Io.Mutex = .init,    reconciliation_thread: ?std.Thread = null,    reconciliation_done: std.atomic.Value(bool) = .init(true),    reconciliation_cancel: std.atomic.Value(bool) = .init(false),    reconciliation_work_epoch: std.atomic.Value(u64) = .init(0),    reconciliation_key_digest: ?[Sha256.digest_length]u8 = null,    reconciliation_credential_blocked: bool = false,    generation_usage_provider: generation_usage.Provider = generation_usage.unavailable_provider,     pub fn initFresh() Usage {        return .{            .billing = .complete,            .api_duration_complete = true,            .wall_duration_complete = true,            .code_complete = true,            .reasoning_tokens = 0,            .request_count = 0,        };    }     pub fn initFreshWithProvider(provider: generation_usage.Provider) Usage {        var usage = initFresh();        usage.generation_usage_provider = provider;        return usage;    }     pub fn initLegacy() Usage {        return .{            .billing = .legacy,            .api_duration_complete = false,            .wall_duration_complete = false,            .code_complete = false,            .reasoning_tokens = null,            .request_count = null,        };    }     pub fn deinit(self: *Usage, alloc: Allocator) void {        self.stopReconciliation();        self.clearOwned(alloc);        self.* = undefined;    }     pub fn resetFresh(self: *Usage, alloc: Allocator) void {        self.stopReconciliation();        self.reset(alloc, true);    }     pub fn resetLegacy(self: *Usage, alloc: Allocator) void {        self.stopReconciliation();        self.reset(alloc, false);    }     pub fn restoreLegacyWallDuration(        self: *Usage,        session_started_at_ms: i64,    ) void {        const now_ms = io_mod.milliTimestamp();        self.mutex.lockUncancelable(io_mod.getIo());        defer self.mutex.unlock(io_mod.getIo());        if (session_started_at_ms <= 0 or session_started_at_ms > now_ms) return;        self.wall_duration_complete = true;        self.active_started_at_ms = session_started_at_ms;    }     pub fn reserveInvocation(self: *Usage) !u64 {        self.mutex.lockUncancelable(io_mod.getIo());        defer self.mutex.unlock(io_mod.getIo());        if (self.active_started_at_ms == 0) {            self.active_started_at_ms = io_mod.milliTimestamp();        }        const sequence = self.next_sequence;        if (self.active_sequence_count == max_active_invocations) {            debug_trace.logf(                "session",                "usage invocation reservation failed reason=active_capacity limit={d}",                .{max_active_invocations},            );            return error.UsageCapacityExceeded;        }        const next_sequence = std.math.add(u64, sequence, 1) catch {            debug_trace.logf(                "session",                "usage invocation reservation failed reason=sequence_overflow sequence={d}",                .{sequence},            );            return error.UsageSequenceOverflow;        };        self.next_sequence = next_sequence;        self.active_sequences[self.active_sequence_count] = sequence;        self.active_sequence_count += 1;        self.dirty = true;        return sequence;    }     pub fn configureCheckpointSink(self: *Usage, sink: ?UsageCheckpointSink) void {        self.checkpoint_mutex.lockUncancelable(io_mod.getIo());        defer self.checkpoint_mutex.unlock(io_mod.getIo());        self.checkpoint_sink = sink;    }     pub fn configurePublicationSink(        self: *Usage,        sink: ?ProfilePublicationSink,    ) void {        self.publication_mutex.lockUncancelable(io_mod.getIo());        self.publication_sink = sink;        self.publication_mutex.unlock(io_mod.getIo());        if (sink != null) self.flushProfilePublications();    }     pub fn persistCheckpoint(self: *Usage) bool {        self.checkpoint_mutex.lockUncancelable(io_mod.getIo());        defer self.checkpoint_mutex.unlock(io_mod.getIo());        return self.persistCheckpointBestEffortLocked();    }     fn reserveInvocationDurably(self: *Usage) !u64 {        self.checkpoint_mutex.lockUncancelable(io_mod.getIo());        defer self.checkpoint_mutex.unlock(io_mod.getIo());        const sequence = try self.reserveInvocation();        self.persistCheckpointRequiredLocked() catch |err| {            self.finishInvocation(sequence, 0, .unbilled);            return err;        };        return sequence;    }     fn finishInvocationDurably(        self: *Usage,        sequence: u64,        duration_ms: u64,        outcome: DeliveryOutcome,    ) !void {        self.checkpoint_mutex.lockUncancelable(io_mod.getIo());        self.finishInvocation(sequence, duration_ms, outcome);        _ = self.persistCheckpointBestEffortLocked();        self.checkpoint_mutex.unlock(io_mod.getIo());        self.flushProfilePublications();    }     fn finishObservedInvocationDurably(        self: *Usage,        alloc: Allocator,        sequence: u64,        duration_ms: u64,        outcome: DeliveryOutcome,        id: []const u8,        origin: []const u8,        team: ?[]const u8,    ) !void {        self.checkpoint_mutex.lockUncancelable(io_mod.getIo());        self.finishObservedInvocation(            alloc,            sequence,            duration_ms,            outcome,            id,            origin,            team,        ) catch |err| {            self.markBillingIncomplete();            _ = self.persistCheckpointBestEffortLocked();            debug_trace.logf(                "session",                "usage generation checkpointed incomplete reason={s}",                .{@errorName(err)},

Lines 1–500 of 4823.