Authenticate and route AI applications with Ruby.

upstream agent_proxy_llm_a {
    server 127.0.0.1:12372;
}

upstream agent_proxy_llm_b {
    server 127.0.0.1:12373;
}

# Concurrent requests per client key: agent_proxy_auth.rb sets the key.
limit_conn_zone $agent_proxy_stream_key zone=agent_proxy_streams:1m;

server {
    listen 18132;

    # Declares the variables that the hooks assign; never requested.
    location = /.agent-proxy-vars {
        internal;
        set $agent_proxy_key_id "";
        set $agent_proxy_stream_key "";
        set $agent_proxy_backend "";
        set $agent_proxy_upstream_auth "";
        return 404;
    }

    # Most variables stay unset on most requests (a rejected key, a
    # request that no hook routed).
    uninitialized_variable_warn off;

    # The client key, checked for every request of this server before a
    # location is chosen. A request without a known key gets 401, and the
    # server-level error_page gives it a JSON body.
    mruby_server_rewrite_handler build/nginx/html/agent_proxy_auth.rb;
    error_page 401 @agent_proxy_401;

    # The upstream gets the credential that agent_proxy_route.rb picked and
    # never the client's.
    proxy_http_version 1.1;
    proxy_set_header Connection close;
    proxy_set_header Authorization $agent_proxy_upstream_auth;
    proxy_set_header x-api-key "";
    proxy_buffering off;

    # The Anthropic Messages API, and the OpenAI Responses API below with
    # the same recipes.
    location = /v1/messages {
        # One request at a time per client key; a second one gets 429 from
        # @agent_proxy_too_many. limit_conn_status is a status that no
        # upstream sends, so that an upstream's own 429 passes unchanged.
        limit_conn agent_proxy_streams 1;
        limit_conn_status 460;
        limit_conn_log_level notice;
        error_page 460 = @agent_proxy_too_many;

        # Picks the upstream block and the credential by the model of the
        # JSON body.
        mruby_enable_read_request_body on;
        mruby_access_handler build/nginx/html/agent_proxy_route.rb;

        # Logs $upstream_status after the response.
        mruby_log_handler build/nginx/html/agent_proxy_log.rb;

        proxy_pass http://$agent_proxy_backend/v1/messages;
    }

    # The zone of limit_conn is per client key, so a stream of a key on
    # either location holds the slot of that key on both.
    location = /v1/responses {
        limit_conn agent_proxy_streams 1;
        limit_conn_status 460;
        limit_conn_log_level notice;
        error_page 460 = @agent_proxy_too_many;

        mruby_enable_read_request_body on;
        mruby_access_handler build/nginx/html/agent_proxy_route.rb;
        mruby_log_handler build/nginx/html/agent_proxy_log.rb;

        proxy_pass http://$agent_proxy_backend/v1/responses;
    }

    # The proxy's own error bodies, in the error format of the API that the
    # client asked for: a named location keeps the request's $uri.
    location @agent_proxy_401 {
        default_type application/json;
        if ($uri = /v1/responses) {
            return 401 '{"error":{"message":"invalid API key","type":"invalid_request_error","param":null,"code":null}}';
        }
        return 401 '{"type":"error","error":{"type":"authentication_error","message":"invalid API key"}}';
    }

    location @agent_proxy_too_many {
        default_type application/json;
        add_header retry-after 1 always;
        if ($uri = /v1/responses) {
            return 429 '{"error":{"message":"too many concurrent requests","type":"rate_limit_error","param":null,"code":null}}';
        }
        return 429 '{"type":"error","error":{"type":"rate_limit_error","message":"too many concurrent requests"}}';
    }
}