7298 lines · 279.0 KB · Zig
RawDownload
const std = @import("std");const builtin = @import("builtin");const build_options = @import("build_options");const secret = @import("../core/auth/secret.zig");const agent_stream_provider = @import("../core/agent/stream_provider.zig");const debug_trace = @import("../core/shared/debug_trace.zig");const io_mod = @import("../core/shared/io.zig");const types = @import("../core/shared/types.zig"); pub fn isRetryableGatewayError(err: anyerror) bool {    return err == error.HttpConnectionClosing or        err == error.ConnectionResetByPeer or        err == error.ConnectionTimedOut;} pub fn networkFailureEvidence(    err: anyerror,    delivery: DeliveryCertainty.State,) ?agent_stream_provider.NetworkFailureEvidence {    const cause: agent_stream_provider.NetworkFailureCause = if (err == error.SystemResumed)        .system_resumed    else if (isRetryableAgentNetworkError(err))        .transport_interrupted    else        return null;    return .{ .cause = cause, .delivery = delivery };} fn isRetryableAgentNetworkError(err: anyerror) bool {    return err == error.TlsInitializationFailed or        err == error.ConnectionSetupTimedOut or        err == error.UnknownHostName or        err == error.NameServerFailure or        err == error.NoAddressReturned or        err == error.DetectingNetworkConfigurationFailed or        err == error.AddressUnavailable or        err == error.ConnectionPending or        err == error.ConnectionRefused or        err == error.HostUnreachable or        err == error.NetworkUnreachable or        err == error.NetworkDown or        err == error.Timeout or        err == error.WouldBlock or        err == error.WriteFailed or        err == error.ReadFailed or        isRetryableGatewayError(err);} fn connectedIoFailure(    cancelled: bool,    system_resumed: bool,    transport_error: anyerror,) anyerror {    if (cancelled) return error.Cancelled;    if (system_resumed) return error.SystemResumed;    return transport_error;} test "connected request failures prefer cancellation then wake evidence" {    try std.testing.expectEqual(        error.Cancelled,        connectedIoFailure(true, true, error.WriteFailed),    );    try std.testing.expectEqual(        error.SystemResumed,        connectedIoFailure(false, true, error.WriteFailed),    );    try std.testing.expectEqual(        error.WriteFailed,        connectedIoFailure(false, false, error.WriteFailed),    );} test "isRetryableGatewayError matches active retryable transport errors" {    try std.testing.expect(isRetryableGatewayError(error.HttpConnectionClosing));    try std.testing.expect(isRetryableGatewayError(error.ConnectionResetByPeer));    try std.testing.expect(isRetryableGatewayError(error.ConnectionTimedOut));    try std.testing.expect(!isRetryableGatewayError(error.AccessDenied));} test "native network failure evidence covers setup send read and resume failures" {    const Cases = struct {        err: anyerror,        cause: agent_stream_provider.NetworkFailureCause = .transport_interrupted,    };    const cases = [_]Cases{        .{ .err = error.TlsInitializationFailed },        .{ .err = error.ConnectionSetupTimedOut },        .{ .err = error.UnknownHostName },        .{ .err = error.NameServerFailure },        .{ .err = error.NoAddressReturned },        .{ .err = error.DetectingNetworkConfigurationFailed },        .{ .err = error.AddressUnavailable },        .{ .err = error.ConnectionPending },        .{ .err = error.ConnectionRefused },        .{ .err = error.ConnectionResetByPeer },        .{ .err = error.ConnectionTimedOut },        .{ .err = error.HostUnreachable },        .{ .err = error.NetworkUnreachable },        .{ .err = error.NetworkDown },        .{ .err = error.Timeout },        .{ .err = error.WouldBlock },        .{ .err = error.HttpConnectionClosing },        .{ .err = error.WriteFailed },        .{ .err = error.ReadFailed },        .{ .err = error.SystemResumed, .cause = .system_resumed },    };     for (cases) |case| {        const evidence = networkFailureEvidence(            case.err,            .possibly_sent,        ) orelse return error.TestExpectedNetworkFailureEvidence;        try std.testing.expectEqual(case.cause, evidence.cause);        try std.testing.expectEqual(            DeliveryCertainty.State.possibly_sent,            evidence.delivery,        );    }     const pre_send = networkFailureEvidence(        error.ConnectionRefused,        .definitely_unsent,    ).?;    try std.testing.expectEqual(        DeliveryCertainty.State.definitely_unsent,        pre_send.delivery,    );} test "native network failure evidence excludes opaque and configuration failures" {    const excluded = [_]anyerror{        error.JsHostStreamFailed,        error.OutOfMemory,        error.AccessDenied,        error.UnsupportedUriScheme,        error.ProtocolUnsupportedBySystem,        error.ResolvConfParseFailed,        error.InvalidDnsARecord,    };     for (excluded) |err| {        try std.testing.expectEqual(            @as(?agent_stream_provider.NetworkFailureEvidence, null),            networkFailureEvidence(err, .definitely_unsent),        );    }} const HttpResult = struct {    status: std.http.Status,    body: []u8,     pub fn deinit(self: *HttpResult, alloc: std.mem.Allocator) void {        alloc.free(self.body);        self.* = undefined;    }}; pub const PostResult = HttpResult;pub const GetResult = HttpResult; pub const GatewayJsonResult = union(enum) {    /// Owned response body; the caller frees it with the request allocator.    success: []u8,    http_status: std.http.Status,     fn deinit(self: *@This(), alloc: std.mem.Allocator) void {        switch (self.*) {            .success => |body| alloc.free(body),            .http_status => {},        }        self.* = undefined;    }}; pub const StreamCallback = agent_stream_provider.StreamCallback;pub const ToolStartCallback = agent_stream_provider.ToolStartCallback; const gateway_retry_base_delay_ns: u64 = 150 * std.time.ns_per_ms;const gateway_connection_setup_timeout_ms: i64 = 30_000;const gateway_retry_after_max_ns: u64 = 5 * std.time.ns_per_s;const gateway_transfer_buffer_bytes: usize = 256 * 1024;const provider_failure_detail_max_bytes: usize = 600;const generation_response_max_bytes: usize = 128 * 1024;const generation_lookup_timeout_ms: i64 = 30_000;// Covers a 4 MiB string at worst-case JSON escaping plus SSE framing.const max_sse_event_line_bytes: usize = 32 * 1024 * 1024;const e2e_gateway_chat_url_env = "FX_E2E_GATEWAY_CHAT_URL";const e2e_gateway_models_url_env = "FX_E2E_GATEWAY_MODELS_URL";const e2e_gateway_credits_url_env = "FX_E2E_GATEWAY_CREDITS_URL";const default_gateway_base_url = "https://ai-gateway.vercel.sh";pub const vercel_ai_gateway_team_header = "x-vercel-ai-gateway-team";/// Identifies fx on every AI Gateway request; the zig std.http default/// (`zig/<version> (std.http)`) is never sent to the gateway.pub const user_agent = "fx/" ++ build_options.app_version;var resolved_model_trace_emitted = std.atomic.Value(bool).init(false);var test_cancel_watcher_spawn_error: ?anyerror = null; pub const StreamResult = struct {    status: std.http.Status,    completion: types.GatewayCompletion = .{},    err_body: ?[]u8 = null,    retry_after_seconds: ?u64 = null,     /// Frees all owned response buffers allocated for this stream result.    pub fn deinit(self: *StreamResult, alloc: std.mem.Allocator) void {        if (self.err_body) |body| alloc.free(body);        if (self.completion.content) |content| alloc.free(content);        if (self.completion.generation_id) |id| alloc.free(id);        if (self.completion.billing) |billing| alloc.free(@constCast(billing.model));        for (self.completion.tool_calls) |call| {            alloc.free(call.id);            alloc.free(call.name);            alloc.free(call.arguments_json);            if (call.provisional_id) |provisional_id| alloc.free(provisional_id);            if (call.provider_result) |provider_result| alloc.free(provider_result);        }        if (self.completion.tool_calls.len > 0) alloc.free(self.completion.tool_calls);        if (self.completion.provider_failure_detail) |detail| alloc.free(@constCast(detail));        const status = self.status;        self.* = .{ .status = status };    }}; /// Possibly-sent model requests are retried by the agent so transport retries/// cannot multiply the logical response budget. Other gateway consumers keep/// their existing bounded transport retry behavior.pub const ProviderAttemptOwner = enum {    transport,    agent,}; pub fn fetchGatewayJson(    alloc: std.mem.Allocator,    api_key: ?[]const u8,    gateway_team: ?[]const u8,    url: []const u8,) !GatewayJsonResult {    var result = try fetchGatewayGetAtUrl(alloc, api_key, gateway_team, url, e2e_gateway_models_url_env);    if (failedGatewayJsonStatus(result.status)) |status| {        result.deinit(alloc);        return .{ .http_status = status };    }     return .{ .success = result.body };} pub fn fetchGatewayGetResult(alloc: std.mem.Allocator, api_key: ?[]const u8, path: []const u8) !GetResult {    return fetchGatewayGet(alloc, api_key, null, path, e2e_gateway_credits_url_env);} pub fn fetchGatewayGenerationResult(    alloc: std.mem.Allocator,    api_key: []const u8,    gateway_team: ?[]const u8,    gateway_origin: []const u8,    generation_id: []const u8,    cancel_flag: *std.atomic.Value(bool),) !GetResult {    if (!types.validGatewayGenerationId(generation_id)) return error.InvalidGenerationId;    var operation = GenerationLookupOperation{        .alloc = alloc,        .api_key = api_key,        .gateway_team = gateway_team,        .gateway_origin = gateway_origin,        .generation_id = generation_id,    };    return runBoundedHttpOperation(        GetResult,        alloc,        cancel_flag,        std.Io.Clock.Timestamp.fromNow(io_mod.getIo(), .{            .clock = .awake,            .raw = .fromMilliseconds(generation_lookup_timeout_ms),        }),        &operation,    );} const GenerationLookupOperation = struct {    alloc: std.mem.Allocator,    api_key: []const u8,    gateway_team: ?[]const u8,    gateway_origin: []const u8,    generation_id: []const u8,     fn run(self: *@This()) !GetResult {        const path = try std.fmt.allocPrint(            self.alloc,            "/v1/generation?id={s}",            .{self.generation_id},        );        defer self.alloc.free(path);        const url = try std.fmt.allocPrint(            self.alloc,            "{s}{s}",            .{ self.gateway_origin, path },        );        defer self.alloc.free(url);        const uri = try std.Uri.parse(url);         var client: std.http.Client = .{            .allocator = self.alloc,            .io = io_mod.getIo(),        };        defer client.deinit();        const auth_header = try std.fmt.allocPrint(            self.alloc,            "Bearer {s}",            .{self.api_key},        );        defer secret.zeroAndFree(self.alloc, auth_header);        var extra_headers_buf: [1]std.http.Header = undefined;        const extra_headers = gatewayModelCatalogExtraHeaders(            &extra_headers_buf,            self.gateway_team,        );        var req = try client.request(.GET, uri, .{            .headers = .{                .authorization = .{ .override = auth_header },                .accept_encoding = .omit,                .user_agent = .{ .override = user_agent },            },            .extra_headers = extra_headers,            .redirect_behavior = .unhandled,        });        defer req.deinit();        try req.sendBodiless();        if (req.connection) |conn| try conn.flush();         var response = try req.receiveHead(&.{});        var transfer_buffer: [16 * 1024]u8 = undefined;        const reader = response.reader(&transfer_buffer);        const body = reader.allocRemaining(            self.alloc,            .limited(generation_response_max_bytes),        ) catch |err| switch (err) {            error.StreamTooLong => return error.GatewayGenerationResponseTooLarge,            else => return err,        };        return .{ .status = response.head.status, .body = body };    }}; fn fetchGatewayGet(alloc: std.mem.Allocator, api_key: ?[]const u8, gateway_team: ?[]const u8, path: []const u8, e2e_url_env: []const u8) !GetResult {    const default_url = try std.fmt.allocPrint(alloc, "{s}{s}", .{ gatewayBaseUrl(), path });    defer alloc.free(default_url);     return fetchGatewayGetAtUrl(alloc, api_key, gateway_team, default_url, e2e_url_env);} fn fetchGatewayGetAtUrl(alloc: std.mem.Allocator, api_key: ?[]const u8, gateway_team: ?[]const u8, default_url: []const u8, e2e_url_env: []const u8) !GetResult {    var client: std.http.Client = .{ .allocator = alloc, .io = io_mod.getIo() };    defer client.deinit();     const url = try resolveE2eGatewayUrl(e2e_url_env, default_url);     var auth_header: ?[]u8 = null;    defer if (auth_header) |value| secret.zeroAndFree(alloc, value);     var headers: std.http.Client.Request.Headers = .{};    headers.user_agent = .{ .override = user_agent };    if (api_key) |key| {        auth_header = try std.fmt.allocPrint(alloc, "Bearer {s}", .{key});        headers.authorization = .{ .override = auth_header.? };    }    var extra_headers_buf: [1]std.http.Header = undefined;    const extra_headers = gatewayModelCatalogExtraHeaders(&extra_headers_buf, gateway_team);     var out: std.Io.Writer.Allocating = .init(alloc);    defer out.deinit();     const result = try client.fetch(.{        .location = .{ .url = url },        .method = .GET,        .headers = headers,        .extra_headers = extra_headers,        .response_writer = &out.writer,        .redirect_behavior = .unhandled,    });     return .{        .status = result.status,        .body = try out.toOwnedSlice(),    };} pub fn fetchGatewayJsonCancellable(    alloc: std.mem.Allocator,    api_key: ?[]const u8,    gateway_team: ?[]const u8,    url: []const u8,    cancel_flag: *std.atomic.Value(bool),) !GatewayJsonResult {    if (cancel_flag.load(.seq_cst)) return error.Cancelled;     const request_url = try resolveE2eGatewayUrl(e2e_gateway_models_url_env, url);    return fetchGatewayJsonAtUrlCancellable(alloc, api_key, gateway_team, request_url, cancel_flag);} fn fetchGatewayJsonAtUrlCancellable(    alloc: std.mem.Allocator,    api_key: ?[]const u8,    gateway_team: ?[]const u8,    url: []const u8,    cancel_flag: *std.atomic.Value(bool),) !GatewayJsonResult {    var operation = GatewayJsonFetchOperation{        .alloc = alloc,        .api_key = api_key,        .gateway_team = gateway_team,        .url = url,        .cancel_flag = cancel_flag,    };    return runCancellableGatewayJsonFetch(alloc, cancel_flag, &operation);} const GatewayJsonFetchOperation = struct {    alloc: std.mem.Allocator,    api_key: ?[]const u8,    gateway_team: ?[]const u8,    url: []const u8,    cancel_flag: *std.atomic.Value(bool),     fn run(self: *@This()) !GatewayJsonResult {        return fetchGatewayJsonAtUrlCore(            self.alloc,            self.api_key,            self.gateway_team,            self.url,            self.cancel_flag,        );    }}; fn runCancellableGatewayJsonFetch(    alloc: std.mem.Allocator,    cancel_flag: *std.atomic.Value(bool),    operation: *GatewayJsonFetchOperation,) !GatewayJsonResult {    if (cancel_flag.load(.seq_cst)) return error.Cancelled;     const Event = union(enum) {        request: anyerror!GatewayJsonResult,        cancelled: anyerror!void,    };    const Runner = struct {        fn run(value: *GatewayJsonFetchOperation) anyerror!GatewayJsonResult {            return value.run();        }    };    const Cleanup = struct {        fn drain(result_alloc: std.mem.Allocator, select: *std.Io.Select(Event)) void {            while (select.cancel()) |item| switch (item) {                .request => |request_result| {                    var late_result = request_result catch continue;                    late_result.deinit(result_alloc);                },                .cancelled => {},            };        }    };     var select_buffer: [2]Event = undefined;    var select: std.Io.Select(Event) = .init(io_mod.getIo(), &select_buffer);    select.concurrent(.cancelled, waitForBoundedCancellation, .{cancel_flag}) catch |err| return err;    select.concurrent(.request, Runner.run, .{operation}) catch |err| {        select.cancelDiscard();        return err;    };     const event = select.await() catch |err| {        Cleanup.drain(alloc, &select);        return err;    };    switch (event) {        .request => |request_result| {            Cleanup.drain(alloc, &select);            if (cancel_flag.load(.seq_cst)) {                var result = request_result catch return error.Cancelled;                result.deinit(alloc);                return error.Cancelled;            }            return request_result;        },        .cancelled => |cancel_result| {            cancel_result catch |err| {                Cleanup.drain(alloc, &select);                return err;            };            Cleanup.drain(alloc, &select);            return error.Cancelled;        },    }} fn fetchGatewayJsonAtUrlCore(    alloc: std.mem.Allocator,    api_key: ?[]const u8,

Lines 1–500 of 7298.