PranitInitial commit439f83c3 months agoHistory
2809 lines · 105.8 KB · Zig
RawDownload
const std = @import("std");const session = @import("session.zig");const session_codec = @import("session_codec.zig");const session_usage = @import("session_usage.zig");const types = @import("../shared/types.zig"); const Allocator = std.mem.Allocator;const Sha256 = std.crypto.hash.sha2.Sha256; pub const event_frame_max_bytes: usize = 8 * 1024 * 1024;pub const raw_state_chunk_bytes: usize = 4 * 1024 * 1024;pub const Identifier = [16]u8;pub const Digest = [Sha256.digest_length]u8; pub const Kind = enum {    session_started,    preferences_changed,    workspace_rebound,    history_turn_committed,    usage_checkpointed,    recovery_checkpoint_set,    recovery_checkpoint_cleared,    state_replacement_started,    state_replacement_chunk,    state_replacement_committed,}; pub const ReplacementReason = enum {    compaction,    migration,    recovery,    log_compaction,}; pub const SessionStarted = struct {    id: []u8,    created_at_ms: i64,    origin_workspace_root: []u8,    workspace_root: []u8,    conversation_language: session.ConversationLanguage,    preferences: session_codec.DurableSessionPreferences,    usage: ?session_usage.Snapshot = null,     fn deinit(self: *SessionStarted, alloc: Allocator) void {        alloc.free(self.id);        alloc.free(self.origin_workspace_root);        alloc.free(self.workspace_root);        self.preferences.deinit(alloc);        if (self.usage) |*usage| usage.deinit(alloc);        self.* = undefined;    }}; pub const PreferencesChanged = struct {    model: ?[]u8 = null,    effort: ?types.ReasoningEffort = null,    fast_mode: ?bool = null,     fn deinit(self: *PreferencesChanged, alloc: Allocator) void {        if (self.model) |model| alloc.free(model);        self.* = undefined;    }}; pub const WorkspaceRebound = struct {    previous_workspace_root: []u8,    workspace_root: []u8,     fn deinit(self: *WorkspaceRebound, alloc: Allocator) void {        alloc.free(self.previous_workspace_root);        alloc.free(self.workspace_root);        self.* = undefined;    }}; pub const HistoryTurnCommitted = struct {    conversation_language: session.ConversationLanguage,    total_input_tokens: u64,    total_output_tokens: u64,    work_id: ?[]u8 = null,    turn: session.HistoryTurn,     fn deinit(self: *HistoryTurnCommitted, alloc: Allocator) void {        if (self.work_id) |id| alloc.free(id);        session.freeHistoryTurn(alloc, self.turn);        self.* = undefined;    }}; pub const UsageCheckpointed = struct {    /// Full cumulative replacement; reducers must not add it to prior usage.    usage: session_usage.Snapshot,     fn deinit(self: *UsageCheckpointed, alloc: Allocator) void {        self.usage.deinit(alloc);        self.* = undefined;    }}; pub const RecoveryCheckpointSet = struct {    checkpoint: session_codec.RecoveryCheckpoint,     fn deinit(self: *RecoveryCheckpointSet, alloc: Allocator) void {        self.checkpoint.deinit(alloc);        self.* = undefined;    }}; pub const RecoveryCheckpointCleared = struct {}; pub const StateReplacementStarted = struct {    replacement_id: Identifier,    reason: ReplacementReason,    encoded_bytes: u64,    sha256: Digest,    chunk_count: u64,}; pub const StateReplacementChunk = struct {    replacement_id: Identifier,    chunk_index: u64,    raw_bytes: u64,    chunk_sha256: Digest,    bytes: []u8,     fn deinit(self: *StateReplacementChunk, alloc: Allocator) void {        alloc.free(self.bytes);        self.* = undefined;    }}; pub const StateReplacementCommitted = struct {    replacement_id: Identifier,    encoded_bytes: u64,    sha256: Digest,    chunk_count: u64,}; pub const Event = union(Kind) {    session_started: SessionStarted,    preferences_changed: PreferencesChanged,    workspace_rebound: WorkspaceRebound,    history_turn_committed: HistoryTurnCommitted,    usage_checkpointed: UsageCheckpointed,    recovery_checkpoint_set: RecoveryCheckpointSet,    recovery_checkpoint_cleared: RecoveryCheckpointCleared,    state_replacement_started: StateReplacementStarted,    state_replacement_chunk: StateReplacementChunk,    state_replacement_committed: StateReplacementCommitted,     fn deinit(self: *Event, alloc: Allocator) void {        switch (self.*) {            .session_started => |*payload| payload.deinit(alloc),            .preferences_changed => |*payload| payload.deinit(alloc),            .workspace_rebound => |*payload| payload.deinit(alloc),            .history_turn_committed => |*payload| payload.deinit(alloc),            .usage_checkpointed => |*payload| payload.deinit(alloc),            .recovery_checkpoint_set => |*payload| payload.deinit(alloc),            .recovery_checkpoint_cleared => {},            .state_replacement_chunk => |*payload| payload.deinit(alloc),            .state_replacement_started, .state_replacement_committed => {},        }        self.* = undefined;    }}; pub const Envelope = struct {    log_generation: Identifier,    seq: u64,    event_id: Identifier,    timestamp_ms: i64,    event: Event,     pub fn kind(self: Envelope) Kind {        return std.meta.activeTag(self.event);    }     pub fn deinit(self: *Envelope, alloc: Allocator) void {        self.event.deinit(alloc);        self.* = undefined;    }}; pub const SequenceValidator = struct {    generation: ?Identifier = null,    next_seq: u64 = 1,     pub fn validate(self: *SequenceValidator, envelope: Envelope) !void {        if (envelope.seq != self.next_seq) return error.NonContiguousSequence;        if (self.generation) |generation| {            if (!std.mem.eql(u8, &generation, &envelope.log_generation)) {                return error.GenerationChanged;            }        } else {            self.generation = envelope.log_generation;        }        self.next_seq = std.math.add(u64, self.next_seq, 1) catch            return error.NonContiguousSequence;    }}; pub const IdentifierSource = struct {    context: *anyopaque,    next_fn: *const fn (context: *anyopaque) Identifier,     pub fn next(self: IdentifierSource) Identifier {        return self.next_fn(self.context);    }}; pub const ReplacementWriteOptions = struct {    log_generation: Identifier,    first_seq: u64,    replacement_id: Identifier,    event_ids: IdentifierSource,    timestamp_ms: i64,    reason: ReplacementReason,}; pub const ReplacementWriteSummary = struct {    encoded_bytes: u64,    sha256: Digest,    chunk_count: u64,    last_seq: u64,    last_event_id: Identifier,}; pub const ReductionStart = struct {    generation: ?Identifier = null,    next_seq: u64 = 1,}; pub const ReductionBoundary = struct {    log_generation: Identifier,    seq: u64,    event_id: Identifier,    byte_offset: u64,}; pub const Reduction = struct {    state: session_codec.DurableSessionState,    truncate_from: ?u64 = null,    through: ?ReductionBoundary = null,    bytes_consumed: u64 = 0,     pub fn deinit(self: *Reduction, alloc: Allocator) void {        self.state.deinit(alloc);        self.* = undefined;    }}; pub fn encodeFrame(alloc: Allocator, envelope: Envelope) ![]u8 {    try validateEnvelope(envelope);     var out: std.Io.Writer.Allocating = .init(alloc);    defer out.deinit();    try out.writer.writeAll("{\"schema_version\":1,\"log_generation\":");    try writeHexString(&out.writer, &envelope.log_generation);    try out.writer.print(",\"seq\":{d},\"event_id\":", .{envelope.seq});    try writeHexString(&out.writer, &envelope.event_id);    try out.writer.print(",\"timestamp_ms\":{d},\"kind\":", .{envelope.timestamp_ms});    try writeJsonString(&out.writer, @tagName(envelope.kind()));    try out.writer.writeAll(",\"payload\":");    try writePayload(&out.writer, envelope.event);    try out.writer.writeAll("}\n");    if (out.written().len > event_frame_max_bytes) return error.EventFrameTooLarge;    return try out.toOwnedSlice();} pub fn decodeFrame(alloc: Allocator, line: []const u8) !Envelope {    if (line.len > event_frame_max_bytes) return error.EventFrameTooLarge;    if (line.len == 0 or line[line.len - 1] != '\n') {        return error.InvalidEventFrame;    }    if (std.mem.indexOfScalar(u8, line[0 .. line.len - 1], '\n') != null) {        return error.InvalidEventFrame;    }     var parsed = std.json.parseFromSlice(std.json.Value, alloc, line[0 .. line.len - 1], .{        .parse_numbers = false,    }) catch |err| switch (err) {        error.OutOfMemory => return error.OutOfMemory,        else => return error.InvalidEventFrame,    };    defer parsed.deinit();    const root = try exactObject(parsed.value, &.{        "schema_version",        "log_generation",        "seq",        "event_id",        "timestamp_ms",        "kind",        "payload",    });    if (try requireU64(root, "schema_version") != 1) return error.UnsupportedEventSchema;    const kind = std.meta.stringToEnum(Kind, try requireString(root, "kind")) orelse        return error.InvalidEventFrame;    var envelope = Envelope{        .log_generation = try parseIdentifier(try requireString(root, "log_generation")),        .seq = try requireU64(root, "seq"),        .event_id = try parseIdentifier(try requireString(root, "event_id")),        .timestamp_ms = try requireI64(root, "timestamp_ms"),        .event = try parsePayload(alloc, kind, root.get("payload") orelse return error.InvalidEventFrame),    };    errdefer envelope.deinit(alloc);    try validateEnvelope(envelope);    return envelope;} pub fn writeStateReplacement(    alloc: Allocator,    writer: *std.Io.Writer,    state: session_codec.DurableSessionState,    options: ReplacementWriteOptions,) !ReplacementWriteSummary {    var discard_buffer: [4096]u8 = undefined;    var discard = std.Io.Writer.Discarding.init(&discard_buffer);    const state_summary = try session_codec.encodeState(state, &discard.writer);    if (state_summary.encoded_bytes == 0) return error.InvalidReplacement;    const chunk_count = std.math.divCeil(        u64,        state_summary.encoded_bytes,        raw_state_chunk_bytes,    ) catch return error.InvalidReplacement;     const start = Envelope{        .log_generation = options.log_generation,        .seq = options.first_seq,        .event_id = options.event_ids.next(),        .timestamp_ms = options.timestamp_ms,        .event = .{ .state_replacement_started = .{            .replacement_id = options.replacement_id,            .reason = options.reason,            .encoded_bytes = state_summary.encoded_bytes,            .sha256 = state_summary.sha256,            .chunk_count = chunk_count,        } },    };    const start_line = try encodeFrame(alloc, start);    defer alloc.free(start_line);    try writer.writeAll(start_line);     var chunk_writer: ReplacementChunkWriter = undefined;    try chunk_writer.init(alloc, writer, options, state_summary, chunk_count);    defer chunk_writer.deinit();    const second_summary = session_codec.encodeState(state, &chunk_writer.interface) catch |err| {        return chunk_writer.failure orelse err;    };    chunk_writer.interface.flush() catch |err| return chunk_writer.failure orelse err;    try chunk_writer.finish();    if (second_summary.encoded_bytes != state_summary.encoded_bytes or        !std.mem.eql(u8, &second_summary.sha256, &state_summary.sha256) or        chunk_writer.chunk_index != chunk_count)    {        return error.InvalidReplacement;    }     const transaction_frame_count = std.math.add(u64, chunk_count, 1) catch        return error.InvalidReplacement;    const commit_seq = std.math.add(u64, options.first_seq, transaction_frame_count) catch        return error.InvalidReplacement;    const commit_id = options.event_ids.next();    const commit = Envelope{        .log_generation = options.log_generation,        .seq = commit_seq,        .event_id = commit_id,        .timestamp_ms = options.timestamp_ms,        .event = .{ .state_replacement_committed = .{            .replacement_id = options.replacement_id,            .encoded_bytes = state_summary.encoded_bytes,            .sha256 = state_summary.sha256,            .chunk_count = chunk_count,        } },    };    const commit_line = try encodeFrame(alloc, commit);    defer alloc.free(commit_line);    try writer.writeAll(commit_line);     return .{        .encoded_bytes = state_summary.encoded_bytes,        .sha256 = state_summary.sha256,        .chunk_count = chunk_count,        .last_seq = commit_seq,        .last_event_id = commit_id,    };} pub fn reduceJsonl(    alloc: Allocator,    source: *std.Io.Reader,    initial: ?session_codec.DurableSessionState,) !Reduction {    return reduceJsonlFrom(alloc, source, initial, .{});} /// Applies one contiguous semantic event frame to caller-owned state./// On failure, `state` remains valid and unchanged.pub fn applyEventFrame(    alloc: Allocator,    state: *session_codec.DurableSessionState,    frame: []const u8,    start: ReductionStart,    expected_event_id: Identifier,) !ReductionBoundary {    if (start.generation == null or start.next_seq == 0) {        return error.InvalidReductionStart;    }    const frame_bytes = std.math.cast(u64, frame.len) orelse        return error.InvalidEventFrame;    var envelope = try decodeFrame(alloc, frame);    defer envelope.deinit(alloc);    var validator = SequenceValidator{        .generation = start.generation,        .next_seq = start.next_seq,    };    try validator.validate(envelope);    if (!std.mem.eql(u8, &envelope.event_id, &expected_event_id)) {        return error.InvalidEventFrame;    }    if (envelope.kind() == .session_started or        envelope.kind() == .state_replacement_started or        envelope.kind() == .state_replacement_chunk or        envelope.kind() == .state_replacement_committed)    {        return error.InvalidEventFrame;    }     var current: ?session_codec.DurableSessionState = state.*;    try applyDelta(alloc, &current, envelope);    state.* = current.?;    return reductionBoundary(envelope, frame_bytes);} pub fn reduceJsonlFrom(    alloc: Allocator,    source: *std.Io.Reader,    initial: ?session_codec.DurableSessionState,    start: ReductionStart,) !Reduction {    var state = initial;    errdefer if (state) |*owned| owned.deinit(alloc);    if (start.next_seq == 0 or        (state == null and (start.generation != null or start.next_seq != 1)))    {        return error.InvalidReductionStart;    }    var validator = SequenceValidator{        .generation = start.generation,        .next_seq = start.next_seq,    };    var byte_offset: u64 = 0;    var through: ?ReductionBoundary = null;     while (true) {        const line = readFrameLine(alloc, source) catch |err| switch (err) {            error.EndOfStream => break,            else => return err,        };        defer alloc.free(line);        const frame_start = byte_offset;        byte_offset += line.len;         var envelope = try decodeFrame(alloc, line);        defer envelope.deinit(alloc);        try validator.validate(envelope);         if (envelope.kind() == .state_replacement_started) {            if (state == null) return error.InvalidReplacement;            const replacement = try reduceReplacement(                alloc,                source,                &validator,                &byte_offset,                envelope,                state.?,            );            if (replacement.state) |next| {                state.?.deinit(alloc);                state = next;                through = replacement.through;            } else {                return .{                    .state = state.?,                    .truncate_from = frame_start,                    .through = through,                    .bytes_consumed = byte_offset,                };            }            continue;        }        if (envelope.kind() == .state_replacement_chunk or            envelope.kind() == .state_replacement_committed)        {            return error.InvalidReplacement;        }        try applyDelta(alloc, &state, envelope);        through = reductionBoundary(envelope, byte_offset);    }     return .{

Lines 1–500 of 2809.