PranitInitial commit439f83c3 months agoHistory
388 lines · 16.9 KB · Zig
RawDownload
const std = @import("std");const question_answer = @import("../../core/agent/question_answer.zig");const core_types = @import("../../core/shared/types.zig");const text_utils = @import("../../core/shared/text_utils.zig");const tool_dispatch = @import("../../core/tooling/tool_dispatch.zig"); const Allocator = std.mem.Allocator; const whitespace = " \t\r\n"; pub const cancel_sentinel = "(user cancelled the question)";pub const not_available_sentinel = "(ask_user_question is only available in the interactive shell; ask the user freeform instead)"; const invalid_args_sentinel = "(ask_user_question: invalid arguments; provide {questions})"; const QuestionBatchParseError = Allocator.Error || error{    InvalidArguments,    MissingQuestions,    QuestionsNotArray,    QuestionCount,    QuestionNotObject,    MissingQuestionText,    QuestionTextNotString,    EmptyQuestionText,    MissingOptions,    OptionsNotArray,    OptionCount,    OptionNotObject,    MissingOptionLabel,    OptionLabelNotString,    EmptyOptionLabel,    DuplicateOptionLabel,}; pub const Input = struct {    args_json: []u8,     pub fn deinit(self: *Input, alloc: Allocator) void {        alloc.free(self.args_json);        self.* = .{ .args_json = &.{} };    }}; pub const Requester = struct {    ctx: ?*anyopaque,    response_alloc: Allocator,    request: tool_dispatch.AskQuestionBatchFn,}; pub fn decode(ctx: tool_dispatch.DispatchContext, args_json: []const u8) tool_dispatch.DispatchError!tool_dispatch.DecodeResult {    const input = try ctx.allocator.create(Input);    errdefer ctx.allocator.destroy(input);    input.* = .{ .args_json = try ctx.allocator.dupe(u8, args_json) };    errdefer input.deinit(ctx.allocator);     return .{ .input = .{ .ptr = input, .deinit_fn = inputDeinit } };} fn inputDeinit(ptr: *anyopaque, alloc: Allocator) void {    const input: *Input = @ptrCast(@alignCast(ptr));    input.deinit(alloc);    alloc.destroy(input);} pub fn validate(_: tool_dispatch.DispatchContext, _: tool_dispatch.ToolInput) tool_dispatch.DispatchError!?[]u8 {    return null;} pub fn call(ctx: tool_dispatch.DispatchContext, erased: tool_dispatch.ToolInput) tool_dispatch.DispatchError!tool_dispatch.ToolResult {    const request = ctx.ask_question_batch orelse {        return .{ .success = try ctx.allocator.dupe(u8, not_available_sentinel) };    };    const input = erased.as(Input);    const output = executeWithRequester(ctx.allocator, input.args_json, .{        .ctx = ctx.ask_question_ctx,        .response_alloc = ctx.allocator,        .request = request,    }) catch |err| switch (err) {        error.OutOfMemory => return error.OutOfMemory,        else => return .{ .failure = try std.fmt.allocPrint(ctx.allocator, "ask_user_question failed: {s}", .{@errorName(err)}) },    };    return .{ .success = output };} pub fn executeWithRequester(result_alloc: Allocator, args_json: []const u8, requester: Requester) ![]u8 {    var parse_arena_state = std.heap.ArenaAllocator.init(result_alloc);    defer parse_arena_state.deinit();    const parse_arena = parse_arena_state.allocator();     const entries = parseQuestionBatch(parse_arena, args_json) catch |err| return invalidArgsBody(result_alloc, err);    const answers_opt = try requester.request(requester.ctx, requester.response_alloc, entries);    if (answers_opt) |answers| {        defer freeAnswers(requester.response_alloc, answers);        return question_answer.encodeJson(result_alloc, entries, answers);    }    return result_alloc.dupe(u8, cancel_sentinel);} fn parseQuestionBatch(arena: Allocator, args_json: []const u8) QuestionBatchParseError![]const core_types.QuestionBatchEntry {    var parsed = std.json.parseFromSlice(std.json.Value, arena, args_json, .{}) catch {        return error.InvalidArguments;    };    defer parsed.deinit();     if (parsed.value != .object) return error.InvalidArguments;    const args = parsed.value.object;     const questions_value = args.get("questions") orelse return error.MissingQuestions;    if (questions_value != .array) return error.QuestionsNotArray;     const questions_arr = questions_value.array;    if (questions_arr.items.len < 1 or questions_arr.items.len > 4) return error.QuestionCount;     var entries_buf: std.ArrayList(core_types.QuestionBatchEntry) = .empty;    try entries_buf.ensureTotalCapacity(arena, questions_arr.items.len);     for (questions_arr.items) |q_item| {        if (q_item != .object) return error.QuestionNotObject;         const question_value = q_item.object.get("question") orelse return error.MissingQuestionText;        if (question_value != .string) return error.QuestionTextNotString;        const question_trimmed = std.mem.trim(u8, question_value.string, whitespace);        if (question_trimmed.len == 0) return error.EmptyQuestionText;        const question_safe = try dupeTerminalSafeQuestionText(arena, question_trimmed);         const options_value = q_item.object.get("options") orelse return error.MissingOptions;        if (options_value != .array) return error.OptionsNotArray;         const options_arr = options_value.array;        if (options_arr.items.len < 2 or options_arr.items.len > 6) return error.OptionCount;         var options_buf: std.ArrayList(core_types.QuestionOption) = .empty;        try options_buf.ensureTotalCapacity(arena, options_arr.items.len);         for (options_arr.items) |item| {            if (item != .object) return error.OptionNotObject;            const label_value = item.object.get("label") orelse return error.MissingOptionLabel;            if (label_value != .string) return error.OptionLabelNotString;            const label_trimmed = std.mem.trim(u8, label_value.string, whitespace);            if (label_trimmed.len == 0) return error.EmptyOptionLabel;            const label_safe = try dupeTerminalSafeQuestionText(arena, label_trimmed);             for (options_buf.items) |existing| {                if (std.ascii.eqlIgnoreCase(existing.label, label_safe)) return error.DuplicateOptionLabel;            }             var desc_slice: ?[]const u8 = null;            if (item.object.get("description")) |desc_value| {                if (desc_value == .string) {                    const trimmed_desc = std.mem.trim(u8, desc_value.string, whitespace);                    if (trimmed_desc.len > 0) desc_slice = try dupeTerminalSafeQuestionText(arena, trimmed_desc);                }            }             options_buf.appendAssumeCapacity(.{                .label = label_safe,                .description = desc_slice,            });        }         entries_buf.appendAssumeCapacity(.{            .question = question_safe,            .options = try options_buf.toOwnedSlice(arena),        });    }     return entries_buf.toOwnedSlice(arena);} fn dupeTerminalSafeQuestionText(arena: Allocator, text: []const u8) Allocator.Error![]const u8 {    if (text_utils.isTerminalSafe(text)) return arena.dupe(u8, text);    const encoded = try text_utils.encodeTerminalSafe(arena, text, std.math.maxInt(usize));    return encoded.bytes;} fn invalidArgsBody(alloc: Allocator, err: QuestionBatchParseError) Allocator.Error![]u8 {    const body = switch (err) {        error.OutOfMemory => return error.OutOfMemory,        error.InvalidArguments => invalid_args_sentinel,        error.MissingQuestions => "(ask_user_question: missing required array \"questions\")",        error.QuestionsNotArray => "(ask_user_question: \"questions\" must be an array)",        error.QuestionCount => "(ask_user_question: provide 1 to 4 questions)",        error.QuestionNotObject => "(ask_user_question: each question must be an object with a \"question\" and \"options\")",        error.MissingQuestionText => "(ask_user_question: each question requires a \"question\" string)",        error.QuestionTextNotString => "(ask_user_question: question \"question\" must be a string)",        error.EmptyQuestionText => "(ask_user_question: question text must not be empty)",        error.MissingOptions => "(ask_user_question: each question requires an \"options\" array)",        error.OptionsNotArray => "(ask_user_question: \"options\" must be an array)",        error.OptionCount => "(ask_user_question: provide 2 to 6 options per question)",        error.OptionNotObject => "(ask_user_question: each option must be an object with a \"label\")",        error.MissingOptionLabel => "(ask_user_question: each option requires a \"label\" string)",        error.OptionLabelNotString => "(ask_user_question: option \"label\" must be a string)",        error.EmptyOptionLabel => "(ask_user_question: option labels must not be empty)",        error.DuplicateOptionLabel => "(ask_user_question: option labels must be unique within a question)",    };    return alloc.dupe(u8, body);} fn freeAnswers(alloc: Allocator, answers: [][]u8) void {    for (answers) |answer| alloc.free(answer);    alloc.free(answers);} pub fn readsOnly(_: tool_dispatch.ToolInput) bool {    return true;} pub fn isIrreversible(_: tool_dispatch.ToolInput) bool {    return false;} fn requesterCancel(_: ?*anyopaque, _: Allocator, _: []const core_types.QuestionBatchEntry) anyerror!?[][]u8 {    return null;} const FakeRequester = struct {    called: bool = false,    saw_trimmed: bool = false,     fn request(raw_ctx: ?*anyopaque, alloc: Allocator, entries: []const core_types.QuestionBatchEntry) anyerror!?[][]u8 {        const self: *FakeRequester = @ptrCast(@alignCast(raw_ctx.?));        self.called = true;        if (entries.len == 1 and            std.mem.eql(u8, entries[0].question, "Choose?") and            entries[0].options.len == 2 and            std.mem.eql(u8, entries[0].options[0].label, "Yes") and            entries[0].options[0].description != null and            std.mem.eql(u8, entries[0].options[0].description.?, "Go ahead") and            entries[0].options[1].description == null)        {            self.saw_trimmed = true;        }         const answers = try alloc.alloc([]u8, entries.len);        errdefer alloc.free(answers);        var filled: usize = 0;        errdefer {            var i: usize = 0;            while (i < filled) : (i += 1) alloc.free(answers[i]);        }        while (filled < entries.len) : (filled += 1) {            answers[filled] = try alloc.dupe(u8, entries[filled].options[0].label);        }        return answers;    }}; const TerminalSafeRequester = struct {    called: bool = false,    saw_terminal_safe: bool = false,     fn request(raw_ctx: ?*anyopaque, alloc: Allocator, entries: []const core_types.QuestionBatchEntry) anyerror!?[][]u8 {        const self: *TerminalSafeRequester = @ptrCast(@alignCast(raw_ctx.?));        self.called = true;        if (entries.len == 1 and            std.mem.eql(u8, entries[0].question, "Q\\x0a\\x1b[31m?") and            entries[0].options.len == 2 and            std.mem.eql(u8, entries[0].options[0].label, "Alpha\\x0aFake") and            entries[0].options[0].description != null and            std.mem.eql(u8, entries[0].options[0].description.?, "Desc\\x09Gap") and            std.mem.eql(u8, entries[0].options[1].label, "\\x1b[31mRed\\x1b[0m"))        {            self.saw_terminal_safe = true;        }         const answers = try alloc.alloc([]u8, entries.len);        errdefer alloc.free(answers);        var filled: usize = 0;        errdefer {            var i: usize = 0;            while (i < filled) : (i += 1) alloc.free(answers[i]);        }        while (filled < entries.len) : (filled += 1) {            answers[filled] = try alloc.dupe(u8, entries[filled].options[0].label);        }        return answers;    }}; fn expectRequesterOutput(args_json: []const u8, expected: []const u8) !void {    const alloc = std.testing.allocator;    var fake = FakeRequester{};    const output = try executeWithRequester(alloc, args_json, .{        .ctx = &fake,        .response_alloc = alloc,        .request = FakeRequester.request,    });    defer alloc.free(output);    try std.testing.expectEqualStrings(expected, output);} test "ask_user_question validates arguments and returns active sentinels" {    try expectRequesterOutput("not-json", invalid_args_sentinel);    try expectRequesterOutput("[]", invalid_args_sentinel);    try expectRequesterOutput("{}", "(ask_user_question: missing required array \"questions\")");    try expectRequesterOutput("{\"questions\":{}}", "(ask_user_question: \"questions\" must be an array)");    try expectRequesterOutput("{\"questions\":[]}", "(ask_user_question: provide 1 to 4 questions)");    try expectRequesterOutput("{\"questions\":[1]}", "(ask_user_question: each question must be an object with a \"question\" and \"options\")");    try expectRequesterOutput("{\"questions\":[{}]}", "(ask_user_question: each question requires a \"question\" string)");    try expectRequesterOutput("{\"questions\":[{\"question\":1}]}", "(ask_user_question: question \"question\" must be a string)");    try expectRequesterOutput("{\"questions\":[{\"question\":\"  \",\"options\":[]}]}", "(ask_user_question: question text must not be empty)");    try expectRequesterOutput("{\"questions\":[{\"question\":\"Q?\"}]}", "(ask_user_question: each question requires an \"options\" array)");    try expectRequesterOutput("{\"questions\":[{\"question\":\"Q?\",\"options\":1}]}", "(ask_user_question: \"options\" must be an array)");    try expectRequesterOutput("{\"questions\":[{\"question\":\"Q?\",\"options\":[]}]}", "(ask_user_question: provide 2 to 6 options per question)");    try expectRequesterOutput("{\"questions\":[{\"question\":\"Q?\",\"options\":[1,2]}]}", "(ask_user_question: each option must be an object with a \"label\")");    try expectRequesterOutput("{\"questions\":[{\"question\":\"Q?\",\"options\":[{},{}]}]}", "(ask_user_question: each option requires a \"label\" string)");    try expectRequesterOutput("{\"questions\":[{\"question\":\"Q?\",\"options\":[{\"label\":1},{\"label\":\"No\"}]}]}", "(ask_user_question: option \"label\" must be a string)");    try expectRequesterOutput("{\"questions\":[{\"question\":\"Q?\",\"options\":[{\"label\":\"\"},{\"label\":\"No\"}]}]}", "(ask_user_question: option labels must not be empty)");    try expectRequesterOutput("{\"questions\":[{\"question\":\"Q?\",\"options\":[{\"label\":\"Yes\"},{\"label\":\" yes \"}]}]}", "(ask_user_question: option labels must be unique within a question)");} test "ask_user_question normalizes descriptions and encodes ordered answers" {    const alloc = std.testing.allocator;    var fake = FakeRequester{};    const output = try executeWithRequester(        alloc,        "{\"questions\":[{\"question\":\" Choose? \",\"options\":[{\"label\":\" Yes \",\"description\":\" Go ahead \"},{\"label\":\"No\",\"description\":3}]}]}",        .{            .ctx = &fake,            .response_alloc = alloc,            .request = FakeRequester.request,        },    );    defer alloc.free(output);     try std.testing.expect(fake.called);    try std.testing.expect(fake.saw_trimmed);    try std.testing.expectEqualStrings("[{\"question\":\"Choose?\",\"answer\":\"Yes\"}]", output);} test "ask_user_question encodes model-controlled prompt text for terminal display" {    const alloc = std.testing.allocator;    var requester = TerminalSafeRequester{};    const output = try executeWithRequester(        alloc,        "{\"questions\":[{\"question\":\"Q\\n\\u001b[31m?\",\"options\":[{\"label\":\"Alpha\\nFake\",\"description\":\"Desc\\tGap\"},{\"label\":\"\\u001b[31mRed\\u001b[0m\"}]}]}",        .{            .ctx = &requester,            .response_alloc = alloc,            .request = TerminalSafeRequester.request,        },    );    defer alloc.free(output);     try std.testing.expect(requester.called);    try std.testing.expect(requester.saw_terminal_safe);    try std.testing.expectEqualStrings(        "[{\"question\":\"Q\\\\x0a\\\\x1b[31m?\",\"answer\":\"Alpha\\\\x0aFake\"}]",        output,    );} test "ask_user_question cancellation returns sentinel" {    const alloc = std.testing.allocator;    const output = try executeWithRequester(        alloc,        "{\"questions\":[{\"question\":\"Q?\",\"options\":[{\"label\":\"Yes\"},{\"label\":\"No\"}]}]}",        .{            .ctx = null,            .response_alloc = alloc,            .request = requesterCancel,        },    );    defer alloc.free(output);     try std.testing.expectEqualStrings(cancel_sentinel, output);} test "ask_user_question noninteractive returns sentinel before parsing" {    const alloc = std.testing.allocator;    const decoded = try decode(.{ .allocator = alloc }, "not-json");    const input = switch (decoded) {        .failure => |body| {            defer alloc.free(body);            return error.TestExpectedEqual;        },        .input => |owned| owned,    };    defer input.deinit(alloc);     const result = try call(.{ .allocator = alloc }, input);    defer result.deinit(alloc);     switch (result) {        .success => |body| try std.testing.expectEqualStrings(not_available_sentinel, body),        .failure => return error.TestExpectedEqual,    }}