const std = @import("std");const common = @import("common.zig");const json_number = @import("../json_number.zig"); const Allocator = std.mem.Allocator; pub const Protocol = common.Protocol; pub const Limits = struct {    max_values: usize = 100,    max_total_value_bytes: usize = 64 * 1024,    max_value_bytes: usize = 4096,    max_name_bytes: usize = 256,    max_ref_bytes: usize = 64 * 1024,    max_context_arguments: usize = 128,    max_context_bytes: usize = 128 * 1024,    common: common.Limits = .{},}; pub const Error = error{    InvalidReference,    InvalidArgument,    InvalidContext,    InvalidResult,    CompletionLimitExceeded,    CompletionByteLimitExceeded,} || common.Error || Allocator.Error || std.Io.Writer.Error; pub const Reference = union(enum) {    prompt: []const u8,    resource_template: []const u8,}; pub const Argument = struct {    name: []const u8,    value: []const u8,}; pub const Result = struct {    values: []const []u8,    total: ?u64 = null,    has_more: ?bool = null,     pub fn deinit(self: *Result, alloc: Allocator) void {        for (self.values) |value| alloc.free(value);        alloc.free(self.values);        self.* = undefined;    }}; pub fn parseResult(alloc: Allocator, response: []const u8, protocol: Protocol, limits: Limits) Error!Result {    var parsed = try common.parseEnvelope(alloc, response, limits.common);    defer parsed.deinit();    const result = common.completeResult(parsed.value, protocol) catch |err| return mapCommonResultError(err);    const completion = result.object.get("completion") orelse return error.InvalidResult;    if (completion != .object) return error.InvalidResult;    const values_value = completion.object.get("values") orelse return error.InvalidResult;    if (values_value != .array or values_value.array.items.len > limits.max_values) {        return error.CompletionLimitExceeded;    }    const values = try alloc.alloc([]u8, values_value.array.items.len);    var count: usize = 0;    var total_bytes: usize = 0;    errdefer {        for (values[0..count]) |value| alloc.free(value);        alloc.free(values);    }    for (values_value.array.items, 0..) |value, index| {        if (value != .string or value.string.len > limits.max_value_bytes) {            return error.CompletionByteLimitExceeded;        }        total_bytes = std.math.add(usize, total_bytes, value.string.len) catch            return error.CompletionByteLimitExceeded;        if (total_bytes > limits.max_total_value_bytes) return error.CompletionByteLimitExceeded;        values[index] = try alloc.dupe(u8, value.string);        count += 1;    }    const total = if (completion.object.get("total")) |value|        (try nonNegativeU64(alloc, value)) orelse return error.InvalidResult    else        null;    const has_more = if (completion.object.get("hasMore")) |value| blk: {        if (value != .bool) return error.InvalidResult;        break :blk value.bool;    } else null;    return .{ .values = values, .total = total, .has_more = has_more };} pub fn buildRequest(    alloc: Allocator,    request_id: u64,    protocol: Protocol,    reference: Reference,    argument: Argument,    context_arguments: []const Argument,    limits: Limits,    write_modern_metadata: *const fn (*std.Io.Writer) anyerror!void,) ![]u8 {    try validateReference(reference, limits);    try validateArgument(argument, limits);    if (context_arguments.len > limits.max_context_arguments) return error.InvalidContext;    var context_bytes: usize = 0;    for (context_arguments, 0..) |item, index| {        try validateArgument(item, limits);        context_bytes = std.math.add(usize, context_bytes, item.name.len) catch            return error.InvalidContext;        context_bytes = std.math.add(usize, context_bytes, item.value.len) catch            return error.InvalidContext;        if (context_bytes > limits.max_context_bytes) return error.InvalidContext;        for (context_arguments[0..index]) |previous| {            if (std.mem.eql(u8, previous.name, item.name)) return error.InvalidContext;        }    }     var out: std.Io.Writer.Allocating = .init(alloc);    defer out.deinit();    try out.writer.print("{{\"jsonrpc\":\"2.0\",\"id\":{d},\"method\":\"completion/complete\",\"params\":{{", .{request_id});    if (protocol == .modern) {        try out.writer.writeAll("\"_meta\":");        try write_modern_metadata(&out.writer);        try out.writer.writeByte(',');    }    try out.writer.writeAll("\"ref\":{");    switch (reference) {        .prompt => |name| {            try out.writer.writeAll("\"type\":\"ref/prompt\",\"name\":");            try std.json.Stringify.value(name, .{}, &out.writer);        },        .resource_template => |uri| {            try out.writer.writeAll("\"type\":\"ref/resource\",\"uri\":");            try std.json.Stringify.value(uri, .{}, &out.writer);        },    }    try out.writer.writeAll("},\"argument\":{\"name\":");    try std.json.Stringify.value(argument.name, .{}, &out.writer);    try out.writer.writeAll(",\"value\":");    try std.json.Stringify.value(argument.value, .{}, &out.writer);    try out.writer.writeAll("}");    if (context_arguments.len > 0) {        try out.writer.writeAll(",\"context\":{\"arguments\":{");        for (context_arguments, 0..) |item, index| {            if (index > 0) try out.writer.writeByte(',');            try std.json.Stringify.value(item.name, .{}, &out.writer);            try out.writer.writeByte(':');            try std.json.Stringify.value(item.value, .{}, &out.writer);        }        try out.writer.writeAll("}}");    }    try out.writer.writeAll("}}");    return try out.toOwnedSlice();} fn validateReference(reference: Reference, limits: Limits) Error!void {    const value = switch (reference) {        .prompt => |name| name,        .resource_template => |uri| uri,    };    if (value.len == 0 or value.len > limits.max_ref_bytes) return error.InvalidReference;} fn validateArgument(argument: Argument, limits: Limits) Error!void {    if (argument.name.len == 0 or argument.name.len > limits.max_name_bytes or        argument.value.len > limits.max_value_bytes)    {        return error.InvalidArgument;    }} fn mapCommonResultError(err: common.Error) Error {    return switch (err) {        error.ProtocolFailure => error.ProtocolFailure,        error.OutOfMemory => error.OutOfMemory,        else => error.InvalidResult,    };} fn nonNegativeU64(alloc: Allocator, value: std.json.Value) Error!?u64 {    const parsed = json_number.fromValue(alloc, value, .{}) catch |err| switch (err) {        error.OutOfMemory => return error.OutOfMemory,        else => return null,    };    if (parsed == null) return null;    var number = parsed.?;    defer number.deinit(alloc);    const result = number.toNonNegativeUsize() orelse return null;    return std.math.cast(u64, result);} fn writeMetadata(writer: *std.Io.Writer) !void {    try writer.writeAll("{\"io.modelcontextprotocol/related-task\":{\"taskId\":\"fx-test\"}}");} test "completion result enforces count bytes total and hasMore" {    const alloc = std.testing.allocator;    var result = try parseResult(        alloc,        "{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"completion\":{\"values\":[\"alpha\",\"beta\"],\"total\":3,\"hasMore\":true}}}",        .modern,        .{},    );    defer result.deinit(alloc);    try std.testing.expectEqual(@as(usize, 2), result.values.len);    try std.testing.expectEqual(@as(?u64, 3), result.total);    try std.testing.expectEqual(@as(?bool, true), result.has_more);     try std.testing.expectError(        error.CompletionLimitExceeded,        parseResult(            alloc,            "{\"jsonrpc\":\"2.0\",\"id\":2,\"result\":{\"completion\":{\"values\":[\"a\",\"b\"]}}}",            .modern,            .{ .max_values = 1 },        ),    );    try std.testing.expectError(        error.CompletionByteLimitExceeded,        parseResult(            alloc,            "{\"jsonrpc\":\"2.0\",\"id\":3,\"result\":{\"completion\":{\"values\":[\"alpha\",\"beta\"]}}}",            .modern,            .{ .max_total_value_bytes = 4 },        ),    );} test "completion requests encode prompt and resource template references" {    const alloc = std.testing.allocator;    const prompt_request = try buildRequest(        alloc,        1,        .modern,        .{ .prompt = "review" },        .{ .name = "focus", .value = "sec" },        &.{.{ .name = "language", .value = "zig" }},        .{},        writeMetadata,    );    defer alloc.free(prompt_request);    try std.testing.expect(std.mem.find(u8, prompt_request, "\"type\":\"ref/prompt\"") != null);     const resource_request = try buildRequest(        alloc,        2,        .modern,        .{ .resource_template = "db:///{table}/{id}" },        .{ .name = "table", .value = "us" },        &.{},        .{},        writeMetadata,    );    defer alloc.free(resource_request);    try std.testing.expect(std.mem.find(u8, resource_request, "\"type\":\"ref/resource\"") != null);    try std.testing.expectError(        error.InvalidContext,        buildRequest(            alloc,            3,            .modern,            .{ .prompt = "review" },            .{ .name = "focus", .value = "sec" },            &.{.{ .name = "language", .value = "zig" }},            .{ .max_context_bytes = 2 },            writeMetadata,        ),    );} fn checkCompletionAllocationFailures(alloc: Allocator) !void {    var result = try parseResult(        alloc,        "{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"completion\":{\"values\":[\"alpha\",\"beta\",\"gamma\"],\"total\":3,\"hasMore\":false}}}",        .modern,        .{},    );    defer result.deinit(alloc);} test "completion result construction is allocation failure safe" {    try std.testing.checkAllAllocationFailures(        std.testing.allocator,        checkCompletionAllocationFailures,        .{},    );}