PranitInitial commit439f83c3 months agoHistory
413 lines · 12.6 KB · Zig
RawDownload
const std = @import("std");const types = @import("../shared/types.zig");const context_limits = @import("../config/context_limits.zig");const elicitation = @import("../mcp/elicitation.zig");const mcp_contract = @import("../mcp/mcp_contract.zig");const access_policy = @import("../mcp/access_policy.zig"); const Allocator = std.mem.Allocator; pub const HasToolFn = *const fn (*anyopaque, []const u8, Access) bool;pub const ValidateToolFn = *const fn (*anyopaque, Allocator, []const u8, []const u8, Access) anyerror!ValidationResult;pub const CallToolFn = *const fn (*anyopaque, Allocator, []const u8, []const u8, usize, CallOptions) anyerror!?CallResult;pub const SearchToolsFn = *const fn (*anyopaque, Allocator, []const u8, usize, types.PermissionRuleSet, context_limits.Values, Access) anyerror!SearchResult;pub const ToolSchemaFn = *const fn (*anyopaque, Allocator, []const u8, types.PermissionRuleSet, context_limits.Values, Access) anyerror!?ToolSchemaResult;pub const FeatureCallFn = *const fn (    *anyopaque,    Allocator,    FeatureRequest,    FeatureCallOptions,) anyerror!FeatureResult; pub const Progress = mcp_contract.Progress;pub const ProgressSink = mcp_contract.ProgressSink; pub const ResolvedLiveView = struct {    /// Generation of the MCP-specific immutable view.    authority_generation: u64,    view: access_policy.View,    /// Full child authority generation for final action precommit.    action_authority_generation: u64 = 0,     pub fn deinit(self: *ResolvedLiveView, alloc: Allocator) void {        self.view.deinit(alloc);        self.* = undefined;    }}; const LiveViewProvider = struct {    context: *anyopaque,    resolve_fn: *const fn (*anyopaque, Allocator) anyerror!ResolvedLiveView,     pub fn resolve(self: LiveViewProvider, alloc: Allocator) !ResolvedLiveView {        return self.resolve_fn(self.context, alloc);    }}; pub const Access = union(enum) {    unrestricted,    disabled,    scoped: struct {        captured: *const access_policy.View,        admission_authority_generation: u64,        live: LiveViewProvider,        action_authority_generation: u64 = 0,    },     pub fn allowsTool(self: Access, runtime_generation: u64, name: []const u8) bool {        return switch (self) {            .unrestricted => true,            .disabled => false,            .scoped => |scope| scope.captured.runtime_generation == runtime_generation and                scope.captured.tool(name) != null,        };    }     pub fn allowsFeatureServer(self: Access, runtime_generation: u64, name: []const u8) bool {        return switch (self) {            .unrestricted => true,            .disabled => false,            .scoped => |scope| scope.captured.runtime_generation == runtime_generation and                scope.captured.features_visible and scope.captured.server(name) != null,        };    }}; pub const CallOptions = struct {    cancel_flag: ?*std.atomic.Value(bool) = null,    progress: ?ProgressSink = null,    input_responder: ?InputResponder = null,    continuation: ?Continuation = null,    access: Access = .unrestricted,}; pub const InputRequired = struct {    input_requests_json: []const u8,    request_state_json: ?[]const u8 = null,    legacy_retry_without_responses: bool = false,    legacy_url_phase: LegacyUrlPhase = .consent,    legacy_url_completion_signal: ?*const LegacyUrlCompletionSignal = null,}; const LegacyUrlPhase = enum {    consent,    await_completion,}; pub const LegacyUrlCompletionStatus = enum(u8) {    pending,    completed,    cancelled,}; /// Borrowed wake state for one legacy URL-required interaction. The runtime/// owns its lifetime; product adapters may observe it only during a responder/// callback and must not retain either pointer.pub const LegacyUrlCompletionSignal = struct {    wake: std.atomic.Value(bool) = .init(false),    status: std.atomic.Value(LegacyUrlCompletionStatus) = .init(.pending),}; pub const InputOrigin = struct {    wire: elicitation.Wire,    server_name: []const u8,    operation: elicitation.Operation,    runtime_generation: u64 = 0,    connection_generation: u64,    client_generation: u64,    catalog_generation: u64,    request_generation: u64,    auth_generation: u64,    deadline_ms: i64,    lifecycle_cancel_flag: ?*const std.atomic.Value(bool) = null,}; pub const InputIdentityWitness = struct {    runtime_generation: u64 = 0,    connection_generation: u64,    client_generation: u64,    catalog_generation: u64,    auth_generation: u64,}; pub const InputResponder = struct {    context: *anyopaque,    capabilities: elicitation.Capabilities = .{},    legacy_url_manual_completion: bool = false,    callback: *const fn (*anyopaque, Allocator, InputOrigin, InputRequired) anyerror![]const u8,    continuation_terminal: ?*const fn (*anyopaque, Allocator, InputOrigin, ContinuationTerminal) void = null,}; pub const LegacyUrlCompletion = struct {    server_name: []const u8,    elicitation_id: []const u8,    runtime_generation: u64 = 0,    connection_generation: u64,    client_generation: u64,    auth_generation: u64,}; pub const LegacyUrlCompletionSink = struct {    context: *anyopaque,    /// Accepts the unique outbound ACP elicitation id. The legacy wire id is    /// not unique across recovered connections and cannot own cleanup.    accept: *const fn (*anyopaque, InputOrigin, []const u8) LegacyUrlAcceptTransition,    consume: *const fn (*anyopaque, LegacyUrlCompletion) LegacyUrlConsumeTransition,    publish: *const fn (*anyopaque, []u8) void,}; pub const LegacyUrlConsumeTransition = union(enum) {    missing,    consumed: ?[]u8,}; pub const LegacyUrlAcceptStatus = enum {    awaiting_completion,    completed,}; pub const LegacyUrlAcceptTransition = union(enum) {    missing,    awaiting_completion,    completed: []u8,}; pub const ContinuationTerminal = enum {    completed,    abandoned,}; pub const Continuation = struct {    input_responses_json: []const u8,    request_state_json: ?[]const u8 = null,}; pub const FeatureAction = enum {    resource_list,    resource_templates,    resource_read,    prompt_list,    prompt_get,    prompt_complete,    resource_complete,}; pub const FeatureArgument = struct {    name: []const u8,    value: []const u8,}; pub const FeatureRequest = struct {    action: FeatureAction,    server_name: []const u8,    identity: []const u8 = "",    argument_name: []const u8 = "",    value: []const u8 = "",    arguments_json: []const u8 = "{}",    context_arguments: []const FeatureArgument = &.{},}; pub const FeatureCallOptions = struct {    cancel_flag: ?*std.atomic.Value(bool) = null,    input_responder: ?InputResponder = null,    output_limit_bytes: usize,    access: Access = .unrestricted,}; pub const FeatureResult = struct {    model_output: []u8,     pub fn deinit(self: *FeatureResult, alloc: Allocator) void {        alloc.free(self.model_output);        self.* = undefined;    }}; pub const CallStatus = enum {    success,    tool_failure,    protocol_failure,    input_required,}; pub const CallResult = struct {    model_output: []const u8,    status: CallStatus = .success,    input_required: ?InputRequired = null,     pub fn deinit(self: *CallResult, alloc: Allocator) void {        alloc.free(self.model_output);        self.deinitInputRequired(alloc);        self.* = undefined;    }     pub fn takeModelOutput(self: *CallResult, alloc: Allocator) []const u8 {        const model_output = self.model_output;        self.deinitInputRequired(alloc);        self.* = undefined;        return model_output;    }     fn deinitInputRequired(self: *CallResult, alloc: Allocator) void {        if (self.input_required) |required| {            alloc.free(required.input_requests_json);            if (required.request_state_json) |state| alloc.free(state);        }    }}; pub const ValidationResult = union(enum) {    valid,    invalid: []const u8,    not_available,}; pub const SearchResult = struct {    model_output: []u8,    notice: ?[]u8 = null,     pub fn deinit(self: *SearchResult, alloc: Allocator) void {        alloc.free(self.model_output);        if (self.notice) |notice| alloc.free(notice);        self.* = undefined;    }}; pub const ToolSchemaResult = union(enum) {    selected: Payload,    rejected: Payload,     pub const Payload = struct {        model_output: []u8,        notice: ?[]u8 = null,    };     pub fn deinit(self: *ToolSchemaResult, alloc: Allocator) void {        const payload = switch (self.*) {            .selected => |value| value,            .rejected => |value| value,        };        alloc.free(payload.model_output);        if (payload.notice) |notice| alloc.free(notice);        self.* = undefined;    }}; pub const RuntimeCapabilities = struct {    context: ?*anyopaque = null,    has_tool: ?HasToolFn = null,    validate_tool: ?ValidateToolFn = null,    call_tool: ?CallToolFn = null,    tool_schema: ?ToolSchemaFn = null,    access: Access = .unrestricted,}; pub fn validateAdvertisedDynamicTool(    input: AvailabilityInput,    arena: Allocator,    name: []const u8,    arguments_json: []const u8,) !ValidationResult {    if (input.is_registered_tool) return .not_available;    const context = input.runtime.context orelse return .not_available;    const validate_tool = input.runtime.validate_tool orelse return .not_available;    return try validate_tool(context, arena, name, arguments_json, input.runtime.access);} const AvailabilityInput = struct {    is_registered_tool: bool,    advertised_dynamic_tool_names: []const []const u8,    runtime: RuntimeCapabilities,}; const CallInput = struct {    advertised_dynamic_tool_names: []const []const u8,    runtime: RuntimeCapabilities,    max_tool_result_bytes: usize,    options: CallOptions = .{},};pub fn isAdvertisedDynamicToolName(    advertised_dynamic_tool_names: []const []const u8,    name: []const u8,) bool {    for (advertised_dynamic_tool_names) |advertised| {        if (std.mem.eql(u8, advertised, name)) return true;    }    return false;} pub fn isAvailableDynamicTool(input: AvailabilityInput, name: []const u8) bool {    if (input.is_registered_tool) return false;    if (!isAdvertisedDynamicToolName(input.advertised_dynamic_tool_names, name)) return false;    const context = input.runtime.context orelse return false;    const has_tool = input.runtime.has_tool orelse return false;    return has_tool(context, name, input.runtime.access);} fn callAdvertisedDynamicTool(    input: CallInput,    arena: Allocator,    name: []const u8,    arguments_json: []const u8,) !?CallResult {    if (!isAdvertisedDynamicToolName(input.advertised_dynamic_tool_names, name)) return null;    const context = input.runtime.context orelse return null;    const call_tool = input.runtime.call_tool orelse return null;    return try call_tool(        context,        arena,        name,        arguments_json,        input.max_tool_result_bytes,        input.options,    );}pub fn notSelectedOutput(arena: Allocator, name: []const u8) ![]const u8 {    return std.fmt.allocPrint(        arena,        "Dynamic MCP tool not selected for this model step: {s}. Use mcp_search_tools and mcp_select_tool first; the selected tool can be called on the next model step after its schema is advertised.",        .{name},    );}test "MCP advertised dynamic tool executes only through its supplied capability" {    var fixture = Fixture.CountingContext{};    const advertised = [_][]const u8{"mcp_fs_read"};    const input = CallInput{        .advertised_dynamic_tool_names = &advertised,        .runtime = .{            .context = @ptrCast(&fixture),            .call_tool = Fixture.callCounting,        },        .max_tool_result_bytes = 1024,    };     const result = (try callAdvertisedDynamicTool(        input,        std.testing.allocator,        "mcp_fs_read",        "{}",    )).?;    defer std.testing.allocator.free(result.model_output);     try std.testing.expectEqualStrings("mcp ok", result.model_output);    try std.testing.expectEqual(@as(usize, 1), fixture.calls);    try std.testing.expect((try callAdvertisedDynamicTool(        input,        std.testing.allocator,        "mcp_other",        "{}",    )) == null);    try std.testing.expectEqual(@as(usize, 1), fixture.calls);} const Fixture = struct {    const CountingContext = struct {        calls: usize = 0,    };     fn callCounting(raw_context: *anyopaque, arena: Allocator, _: []const u8, _: []const u8, _: usize, _: CallOptions) anyerror!?CallResult {        const context: *CountingContext = @ptrCast(@alignCast(raw_context));        context.calls += 1;        return .{ .model_output = try arena.dupe(u8, "mcp ok") };    }};