diff --git a/apisix/plugins/ai-providers/base.lua b/apisix/plugins/ai-providers/base.lua index a992defbaabc..6ed57b84e2ec 100644 --- a/apisix/plugins/ai-providers/base.lua +++ b/apisix/plugins/ai-providers/base.lua @@ -99,50 +99,93 @@ local function merge_usage(ctx, parsed) end ---- Build HTTP request parameters from driver config and extra_opts. --- @return table params HTTP parameters ready for transport_http.request() --- @return string|nil err Error message -function _M.build_request(self, ctx, conf, request_body, opts) - local body_changed = false - - -- Protocol conversion (when a converter bridges client→target protocol) - local converter = ctx.ai_converter - if converter and converter.convert_request then - local converted, err = converter.convert_request(request_body, ctx) - if not converted then - return nil, err or "invalid protocol", 400 - end - request_body = converted - body_changed = true - end +--- Shape the request body for the target protocol and this provider. +-- +-- Pure: no ctx, no downstream request, no network -- body in, body out. This is +-- where the target protocol matters; build_request below does not know protocols +-- exist. Protocol *conversion* is not done here either: converters carry state on +-- the request ctx for the response side to read back, so the caller converts and +-- hands us the converted body. +-- @return table body The body to send +-- @return boolean changed Whether anything here modified it. The caller uses this +-- to decide whether the client's verbatim bytes may still be reused. +function _M.build_body(self, request_body, opts) + local changed = false + local target_protocol = opts.target_protocol -- Inject target-protocol-specific parameters (e.g. stream_options for OpenAI). - -- This runs after conversion so it covers both passthrough and convert scenarios. - local target_protocol = ctx.ai_target_protocol if target_protocol then local target_proto = protocols.get(target_protocol) if target_proto and target_proto.prepare_outgoing_request then if target_proto.prepare_outgoing_request(request_body) then - body_changed = true + changed = true end end end - core.log.info("request extra_opts to LLM server: ", - core.json.delay_encode(log_sanitize.redact_extra_opts(opts), true)) + -- Inject model options (flat overwrite) + if opts.model_options then + for opt, val in pairs(opts.model_options) do + if request_body[opt] ~= nil then + core.log.info("model_options overwriting request field '", opt, "'") + end + request_body[opt] = val + changed = true + end + end - -- Auth: GCP token - local auth = opts.auth or {} - local token - if auth.gcp then - local access_token, err = transport_auth.fetch_gcp_access_token(ctx, opts.name, - auth.gcp) - if not access_token then - return nil, "failed to get gcp access token: " .. (err or "unknown") + -- Apply llm_options via provider capability hook (always force-overwrites) + if opts.override_llm_options then + local cap = self.capabilities and self.capabilities[target_protocol] + if cap and cap.rewrite_request_body then + cap.rewrite_request_body(request_body, opts.override_llm_options, true) + if next(opts.override_llm_options) ~= nil then + changed = true + end end - token = access_token end + -- Apply per-target-protocol request body override (deep merge) + if opts.request_body_override_map then + local patch = opts.request_body_override_map[target_protocol] + if patch then + core.log.info("applying request_body override for target protocol '", + target_protocol, "'") + request_body = deep_merge(request_body, patch, opts.request_body_force_override) + changed = true + end + end + + if self.remove_model and request_body.model ~= nil then + request_body.model = nil + changed = true + end + + return request_body, changed +end + + +--- Assemble the outbound HTTP request: endpoint, auth, headers, query, and -- +-- last -- signing. +-- +-- Pure and protocol-agnostic: `body` arrives already shaped (see build_body) and +-- is passed through untouched, so a string goes out verbatim and a table is +-- encoded by the transport. Anything taken from the inbound request reaches us +-- as an explicit opts.client_* field; a self-contained internal call sets none of +-- them, so nothing of the client's can reach the LLM or the logs. +-- @return table params HTTP parameters ready for transport_http.request() +-- @return string|nil err Error message +function _M.build_request(self, conf, body, opts) + -- Name the few options worth tracing rather than dumping opts wholesale: the + -- resolved endpoint, path and headers are already logged from `params` below, + -- and serialising an open-ended options bag is how credentials and client data + -- end up in logs. + core.log.info("building LLM request: instance=", opts.name, + ", target_protocol=", opts.target_protocol, + ", model=", opts.model_options and opts.model_options.model) + + local auth = opts.auth or {} + -- Parse endpoint URL local endpoint = opts.endpoint local parsed_url @@ -188,33 +231,39 @@ function _M.build_request(self, ctx, conf, request_body, opts) .. "includes a path", 400 end - local headers = transport_http.construct_forward_headers(auth.header or {}, ctx) + -- opts.client_* is whatever the caller took from the inbound request; an + -- internal call sets none of it and therefore forwards none of the client's + -- headers or query. + local headers = transport_http.construct_forward_headers(auth.header or {}, + opts.client_headers) if opts.host_header then headers["Host"] = opts.host_header end - if token then + -- GCP tokens are fetched here: the lookup is keyed on the credentials, not on + -- the request, so it needs no ctx and no caller has to repeat it. + if auth.gcp then + local token, token_err = transport_auth.fetch_gcp_access_token(opts.name, auth.gcp) + if not token then + return nil, "failed to get gcp access token: " .. (token_err or "unknown") + end headers["authorization"] = "Bearer " .. token end - -- Protocol converter header transformation (e.g. Anthropic → OpenAI headers) - if ctx.ai_converter and ctx.ai_converter.convert_headers then - ctx.ai_converter.convert_headers(headers) - end - - -- For the passthrough protocol the gateway acts as a catch-all proxy, so - -- forward the client's HTTP method and original query string unchanged. - -- Other protocols always issue a POST with provider-specific query args. - local method = "POST" - if ctx.ai_target_protocol == "passthrough" then - method = core.request.get_method() - local client_args = ctx.var.args and core.string.decode_args(ctx.var.args) - if type(client_args) == "table" then - -- client query overrides the endpoint query, but configured - -- auth.query credentials must stay non-overridable by the caller - for k, v in pairs(client_args) do - if not (auth.query and auth.query[k] ~= nil) then - query_params[k] = v - end + -- Optional header transformation (e.g. the Anthropic → OpenAI converter's). + if opts.header_transform then + opts.header_transform(headers) + end + + -- The caller passes the downstream method and query string only when it + -- wants them proxied verbatim (the passthrough protocol). Otherwise this is + -- a plain POST with provider-specific query args. + local method = opts.client_method or "POST" + if type(opts.client_args) == "table" then + -- client query overrides the endpoint query, but configured + -- auth.query credentials must stay non-overridable by the caller + for k, v in pairs(opts.client_args) do + if not (auth.query and auth.query[k] ~= nil) then + query_params[k] = v end end end @@ -231,63 +280,11 @@ function _M.build_request(self, ctx, conf, request_body, opts) ssl_server_name = opts.ssl_server_name or parsed_url and parsed_url.host or opts.target_host or self.host, + -- Passed through as given: a string goes out verbatim (the client's own + -- bytes), a table is encoded by the transport. + body = body, } - -- Inject model options (flat overwrite) - if opts.model_options then - for opt, val in pairs(opts.model_options) do - if request_body[opt] ~= nil then - core.log.info("model_options overwriting request field '", opt, "'") - end - request_body[opt] = val - body_changed = true - end - end - - -- Apply llm_options via provider capability hook (always force-overwrites) - if opts.override_llm_options then - local cap = self.capabilities and self.capabilities[ctx.ai_target_protocol] - if cap and cap.rewrite_request_body then - cap.rewrite_request_body(request_body, opts.override_llm_options, true) - if next(opts.override_llm_options) ~= nil then - body_changed = true - end - end - end - - -- Apply per-target-protocol request body override (deep merge) - if opts.request_body_override_map then - local patch = opts.request_body_override_map[ctx.ai_target_protocol] - if patch then - core.log.info("applying request_body override for target protocol '", - ctx.ai_target_protocol, "'") - request_body = deep_merge(request_body, patch, opts.request_body_force_override) - body_changed = true - end - end - - if self.remove_model and request_body.model ~= nil then - request_body.model = nil - body_changed = true - end - - if body_changed then - ctx.ai_request_body_changed = true - end - - if not ctx.ai_request_body_changed then - if ctx.ai_raw_request_body == nil then - ctx.ai_raw_request_body = core.request.get_body() - end - if type(ctx.ai_raw_request_body) == "string" then - params.body = ctx.ai_raw_request_body - else - params.body = request_body - end - else - params.body = request_body - end - -- AWS SigV4 signing (must be last — signs the finalized body) if self.aws_sigv4 and auth.aws then local auth_aws = require("apisix.plugins.ai-transport.auth-aws") @@ -801,8 +798,13 @@ end -- @return number|nil HTTP status code -- @return string|nil Raw response body (JSON string) -- @return string|nil Error message -function _M.request(self, ctx, conf, request_table, extra_opts) - local params, err = self:build_request(ctx, conf, request_table, extra_opts) +--- Send a self-contained request to an LLM and return its raw response. +-- Shapes the body, assembles the HTTP request, sends it. Fully decoupled from the +-- downstream request: internal callers (ai-request-rewrite, embeddings, ...) use +-- this as a plain client and parse the body themselves. +function _M.request(self, conf, request_table, extra_opts) + local body = self:build_body(request_table, extra_opts) + local params, err = self:build_request(conf, body, extra_opts) if not params then core.log.error("failed to build request: ", err) return 500, nil, err diff --git a/apisix/plugins/ai-proxy/base.lua b/apisix/plugins/ai-proxy/base.lua index 97c980dc5028..b98ae43e37da 100644 --- a/apisix/plugins/ai-proxy/base.lua +++ b/apisix/plugins/ai-proxy/base.lua @@ -240,14 +240,64 @@ function _M.before_proxy(conf, ctx, on_error) extra_opts.target_path = target_path extra_opts.target_host = target_host + extra_opts.target_protocol = target_proto + -- The transport is a pure client, so everything below that depends on the + -- downstream request or on ctx is resolved here and handed over via + -- extra_opts. + extra_opts.header_transform = converter and converter.convert_headers + + -- ai-proxy is a transparent proxy of an inbound request, so it passes what + -- it takes from that request as client_* options. Internal callers set + -- none of them and thus forward nothing of the client's -- see + -- ai-providers/base.lua build_request. Keep the names in sync with + -- log-sanitize's CLIENT_DERIVED_FIELDS so they stay out of the logs. + extra_opts.client_headers = core.request.headers(ctx) + + -- passthrough proxies the client's method and query string verbatim. + if target_proto == "passthrough" then + extra_opts.client_method = core.request.get_method() + local client_args = ctx.var.args and core.string.decode_args(ctx.var.args) + if type(client_args) == "table" then + extra_opts.client_args = client_args + end + end local do_request = function() ctx.llm_request_start_time = ngx.now() ctx.var.llm_request_body = request_body - -- Step 3: Build HTTP request params + -- Step 2.5: protocol conversion. It runs in the provider's stead -- + -- converters stash state on ctx for the response side to read back -- + -- but stays inside do_request so the pcall below still bounds it: a + -- converter fed hostile-but-valid JSON can raise (e.g. an Anthropic + -- image block whose "source" is not an object). + local body_for_llm = request_body + local converted = false + if converter and converter.convert_request then + local new_body, conv_err = converter.convert_request(request_body, ctx) + if not new_body then + return 400, {error_msg = conv_err or "invalid protocol"} + end + body_for_llm = new_body + converted = true + end + + -- Step 3: shape the body for the target protocol, then decide which + -- bytes actually go out. When nothing has touched the body -- no + -- conversion above, no shaping just now, and no earlier plugin rewrite + -- (ai-request-rewrite marks that on ctx) -- the client's verbatim bytes + -- are reused, keeping a pure passthrough byte-identical. + local body, shaped = ai_provider:build_body(body_for_llm, extra_opts) + if not shaped and not converted and not ctx.ai_request_body_changed then + local raw = core.request.get_body() + if type(raw) == "string" then + body = raw + end + end + + -- Step 4: assemble the HTTP request local params, build_err, code = ai_provider:build_request( - ctx, conf, request_body, extra_opts) + conf, body, extra_opts) if not params then local body = {error_msg = build_err} if code then diff --git a/apisix/plugins/ai-request-rewrite.lua b/apisix/plugins/ai-request-rewrite.lua index 5bef20771582..29e1b1d58ebf 100644 --- a/apisix/plugins/ai-request-rewrite.lua +++ b/apisix/plugins/ai-request-rewrite.lua @@ -125,10 +125,11 @@ local function request_to_llm(conf, request_table, ctx, target_path) model_options = conf.options, target_path = target_path, } + -- Nothing downstream-derived is handed over: no client headers, no verbatim + -- client body -- this call carries its own credentials and its own body. ctx.llm_request_start_time = ngx.now() ctx.var.llm_request_body = request_table - ctx.ai_request_body_changed = true - return ai_provider:request(ctx, conf, request_table, extra_opts) + return ai_provider:request(conf, request_table, extra_opts) end @@ -222,6 +223,9 @@ function _M.access(conf, ctx) -- Replace the original request body with the rewritten content ngx.req.set_body_data(content) + -- Tell later AI plugins (ai-proxy) that the downstream body is no longer the + -- client's original, so they must not reuse the raw bytes as-is. + ctx.ai_request_body_changed = true end return _M diff --git a/apisix/plugins/ai-transport/auth.lua b/apisix/plugins/ai-transport/auth.lua index 147067e8f236..2670caa44cbc 100644 --- a/apisix/plugins/ai-transport/auth.lua +++ b/apisix/plugins/ai-transport/auth.lua @@ -23,6 +23,8 @@ local google_oauth = require("apisix.utils.google-cloud-oauth") local lrucache = require("resty.lrucache") local type = type local os = os +local tostring = tostring +local ngx_md5 = ngx.md5 local _M = {} @@ -30,25 +32,35 @@ local gcp_access_token_cache = lrucache.new(1024 * 4) --- Fetch (or retrieve from cache) a GCP OAuth2 access token. --- @param ctx table Request context +-- +-- Keyed on the credentials themselves rather than on the request ctx: the token +-- depends only on the service account, so identical credentials share one token +-- across routes, and callers need no ctx to ask for one. -- @param name string Cache key name (driver instance name) -- @param gcp_conf table GCP configuration (service_account_json, expire_early_secs, max_ttl) -- @return string|nil Access token -- @return string|nil Error message -function _M.fetch_gcp_access_token(ctx, name, gcp_conf) - local key = core.lrucache.plugin_ctx_id(ctx, name) +function _M.fetch_gcp_access_token(name, gcp_conf) + -- The credentials may arrive as a secret ref resolved per request, so key on + -- the resolved value; hash it so the cache never holds the raw JSON. Include + -- the fields that change how long the token is cached (expire_early_secs, + -- max_ttl) so that identical credentials under different TTL policies do not + -- share an entry, and use a 128-bit digest instead of a 32-bit CRC so that + -- distinct credentials cannot collide onto the wrong cached token. + local conf = type(gcp_conf) == "table" and gcp_conf or {} + local sa = conf.service_account_json or os.getenv("GCP_SERVICE_ACCOUNT") + local key = ngx_md5((sa or "") .. "\0" + .. tostring(conf.expire_early_secs or "") .. "\0" + .. tostring(conf.max_ttl or "")) .. "#" .. (name or "") local access_token = gcp_access_token_cache:get(key) if not access_token then local auth_conf = {} - gcp_conf = type(gcp_conf) == "table" and gcp_conf or {} - local service_account_json = gcp_conf.service_account_json or - os.getenv("GCP_SERVICE_ACCOUNT") - if type(service_account_json) == "string" and service_account_json ~= "" then - local conf, err = core.json.decode(service_account_json) - if not conf then + if type(sa) == "string" and sa ~= "" then + local decoded, err = core.json.decode(sa) + if not decoded then return nil, "invalid gcp service account json: " .. (err or "unknown error") end - auth_conf = conf + auth_conf = decoded end local oauth = google_oauth.new(auth_conf) access_token = oauth:generate_access_token() @@ -56,11 +68,11 @@ function _M.fetch_gcp_access_token(ctx, name, gcp_conf) return nil, "failed to get google oauth token" end local ttl = oauth.access_token_ttl or 3600 - if gcp_conf.expire_early_secs and ttl > gcp_conf.expire_early_secs then - ttl = ttl - gcp_conf.expire_early_secs + if conf.expire_early_secs and ttl > conf.expire_early_secs then + ttl = ttl - conf.expire_early_secs end - if gcp_conf.max_ttl and ttl > gcp_conf.max_ttl then - ttl = gcp_conf.max_ttl + if conf.max_ttl and ttl > conf.max_ttl then + ttl = conf.max_ttl end gcp_access_token_cache:set(key, access_token, ttl) core.log.debug("set gcp access token in cache with ttl: ", ttl, ", key: ", key) diff --git a/apisix/plugins/ai-transport/http.lua b/apisix/plugins/ai-transport/http.lua index 5ea9d2194545..a2075aa3f4c3 100644 --- a/apisix/plugins/ai-transport/http.lua +++ b/apisix/plugins/ai-transport/http.lua @@ -43,9 +43,14 @@ end --- Build forwarded headers from client request + extra headers. --- Copies client headers, merges ext_opts_headers (lowercased), +-- Copies `client_headers`, merges ext_opts_headers (lowercased), -- forces Content-Type to application/json, removes host/content-length. -function _M.construct_forward_headers(ext_opts_headers, ctx) +-- `client_headers` is the downstream request's headers to forward (proxy path), +-- or nil for a self-contained internal request (e.g. ai-request-rewrite calling +-- an LLM to rewrite the body), which must not leak the client's Authorization, +-- Cookie or other headers to a third-party endpoint. The caller passes them in +-- explicitly, so the transport carries no `ctx` / downstream-request coupling. +function _M.construct_forward_headers(ext_opts_headers, client_headers) local blacklist = { "host", "content-length", @@ -53,7 +58,7 @@ function _M.construct_forward_headers(ext_opts_headers, ctx) } local headers = {} - for k, v in pairs(core.request.headers(ctx) or {}) do + for k, v in pairs(client_headers or {}) do headers[str_lower(k)] = v end for k, v in pairs(ext_opts_headers or {}) do diff --git a/t/plugin/ai-proxy-protocol-conversion.t b/t/plugin/ai-proxy-protocol-conversion.t index b28d1c9eba51..a737e8e7b993 100644 --- a/t/plugin/ai-proxy-protocol-conversion.t +++ b/t/plugin/ai-proxy-protocol-conversion.t @@ -1108,3 +1108,52 @@ cannot parse any events and the gateway should return 502 instead of crashing. status: 502 --- error_log streaming response completed without producing any output + + + +=== TEST 29: Set up route for converter fault isolation +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/v1/messages", + "plugins": { + "ai-proxy": { + "provider": "openai", + "auth": { + "header": { + "Authorization": "Bearer token" + } + }, + "override": { + "endpoint": "http://127.0.0.1:1980/v1/chat/completions" + } + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 30: A converter fault on valid JSON stays a controlled 500 +--- request +POST /v1/messages +{"model":"claude-3","messages":[{"role":"user","content":[{"type":"image","source":true}]}]} +--- more_headers +Content-Type: application/json +--- error_code: 500 +--- error_log +failed to send request to AI service +attempt to index diff --git a/t/plugin/ai-proxy-request-body-override.t b/t/plugin/ai-proxy-request-body-override.t index 1bac6c31b72b..58a5bbfd11a3 100644 --- a/t/plugin/ai-proxy-request-body-override.t +++ b/t/plugin/ai-proxy-request-body-override.t @@ -258,45 +258,102 @@ same body -=== TEST 5c: build_request does not reuse raw body after an earlier rewrite +=== TEST 5c: build_request passes the body through untouched --- config location /t { content_by_lua_block { + -- build_request assembles the HTTP request only: whatever body the + -- caller hands it goes out as-is. A table is encoded by the transport. local base = require("apisix.plugins.ai-providers.base") - local provider = base.new({ - capabilities = { - ["openai-chat"] = { - path = "/v1/chat/completions", - host = "localhost", - }, - }, - }) - local ctx = { - ai_target_protocol = "openai-chat", - ai_request_body_changed = true, - var = {}, - } + local provider = base.new({}) local opts = { auth = {}, conf = {}, - raw_request_body = '{"messages":[]}', target_path = "/v1/chat/completions", + target_host = "localhost", } local request_body = {messages = {{role = "user", content = "changed"}}} - local params = assert(provider:build_request(ctx, {ssl_verify = false}, + local params = assert(provider:build_request({ssl_verify = false}, request_body, opts)) ngx.say(type(params.body)) - ngx.say(params.body == request_body and "table body" or "raw body") + ngx.say(params.body == request_body and "same table" or "other body") } } --- response_body table -table body +same table + + + +=== TEST 6d: build_request sends a string body verbatim +--- config + location /t { + content_by_lua_block { + -- The caller reuses the client's raw bytes simply by passing the + -- string; the transport sends a string body as-is. + local base = require("apisix.plugins.ai-providers.base") + local provider = base.new({}) + local opts = { + auth = {}, + conf = {}, + target_path = "/v1/chat/completions", + target_host = "localhost", + } + + local params = assert(provider:build_request({ssl_verify = false}, + '{"messages":[]}', opts)) + ngx.say(params.body) + } + } +--- response_body +{"messages":[]} + + + +=== TEST 7e: build_body reports whether it changed the body +--- config + location /t { + content_by_lua_block { + -- `changed` is what tells the caller the client's verbatim bytes are + -- no longer safe to reuse. + local base = require("apisix.plugins.ai-providers.base") + local provider = base.new({ + capabilities = { + ["openai-chat"] = { path = "/v1/chat/completions" }, + }, + }) + + -- nothing to apply -> untouched + local _, changed = provider:build_body({messages = {}}, + {target_protocol = "openai-chat"}) + ngx.say("untouched: ", tostring(changed)) + + -- a request_body override -> changed + local body, changed2 = provider:build_body({messages = {}}, { + target_protocol = "openai-chat", + request_body_override_map = { + ["openai-chat"] = {temperature = 0.5}, + }, + }) + ngx.say("overridden: ", tostring(changed2), " temperature=", body.temperature) + + -- model_options -> changed + local body3, changed3 = provider:build_body({messages = {}}, { + target_protocol = "openai-chat", + model_options = {model = "gpt-4o"}, + }) + ngx.say("model_options: ", tostring(changed3), " model=", body3.model) + } + } +--- response_body +untouched: false +overridden: true temperature=0.5 +model_options: true model=gpt-4o -=== TEST 6: llm_options: openai provider maps max_tokens to max_completion_tokens +=== TEST 8: llm_options: openai provider maps max_tokens to max_completion_tokens --- config location /t { content_by_lua_block { @@ -339,7 +396,7 @@ max_completion_tokens=555 -=== TEST 7: llm_options: openai-compatible provider maps max_tokens to max_tokens +=== TEST 9: llm_options: openai-compatible provider maps max_tokens to max_tokens --- config location /t { content_by_lua_block { @@ -382,7 +439,7 @@ max_tokens=444 -=== TEST 8: llm_options: openai responses API maps max_tokens to max_output_tokens +=== TEST 10: llm_options: openai responses API maps max_tokens to max_output_tokens --- config location /t { content_by_lua_block { @@ -425,7 +482,7 @@ max_output_tokens=333 -=== TEST 9: llm_options: ai-proxy-multi per-instance override +=== TEST 11: llm_options: ai-proxy-multi per-instance override --- config location /t { content_by_lua_block { @@ -472,7 +529,7 @@ max_completion_tokens=222 -=== TEST 10: llm_options always force-overwrites client value +=== TEST 12: llm_options always force-overwrites client value --- config location /t { content_by_lua_block { @@ -517,7 +574,7 @@ max_tokens=555 -=== TEST 11: request_body: openai-chat override writes fields on outgoing body +=== TEST 13: request_body: openai-chat override writes fields on outgoing body --- config location /t { content_by_lua_block { @@ -565,7 +622,7 @@ max_tokens=555 temperature=0.1 -=== TEST 12: request_body: non-force deep merge fills missing nested keys without overwriting existing +=== TEST 14: request_body: non-force deep merge fills missing nested keys without overwriting existing --- config location /t { content_by_lua_block { @@ -614,7 +671,7 @@ include_usage=true extra=1 -=== TEST 13: request_body: array values are replaced wholesale (stop sequences) +=== TEST 15: request_body: array values are replaced wholesale (stop sequences) --- config location /t { content_by_lua_block { @@ -660,7 +717,7 @@ stop=["a","b"] -=== TEST 14: request_body: override keyed by non-matching target protocol is ignored +=== TEST 16: request_body: override keyed by non-matching target protocol is ignored --- config location /t { content_by_lua_block { @@ -703,7 +760,7 @@ max_tokens=nil -=== TEST 15: request_body: default mode - client value takes priority +=== TEST 17: request_body: default mode - client value takes priority --- config location /t { content_by_lua_block { @@ -748,7 +805,7 @@ max_tokens=999 -=== TEST 16: request_body: force_override mode - override overwrites client fields +=== TEST 18: request_body: force_override mode - override overwrites client fields --- config location /t { content_by_lua_block { @@ -794,7 +851,7 @@ max_tokens=555 -=== TEST 17: request_body: override applies to target protocol after converter +=== TEST 19: request_body: override applies to target protocol after converter --- config location /t { content_by_lua_block { @@ -841,7 +898,7 @@ max_tokens=77 has_messages=true -=== TEST 18: ai-proxy-multi per-instance request_body override +=== TEST 20: ai-proxy-multi per-instance request_body override --- config location /t { content_by_lua_block { @@ -889,7 +946,7 @@ max_tokens=321 -=== TEST 19: both llm_options and request_body coexist, request_body wins +=== TEST 21: both llm_options and request_body coexist, request_body wins --- config location /t { content_by_lua_block { diff --git a/t/plugin/ai-proxy-vertex-ai.t b/t/plugin/ai-proxy-vertex-ai.t index e3710f3f1b1f..88a6c2aa7268 100644 --- a/t/plugin/ai-proxy-vertex-ai.t +++ b/t/plugin/ai-proxy-vertex-ai.t @@ -486,3 +486,39 @@ request head: GET /status/gpt4 /v1/projects/my-project/locations/us-central1/publishers/google/models/text-embedding-004:predict /v1/projects/my-project/locations/us-central1/publishers/google/models/textembedding-gecko:predict nil + + + +=== TEST 13: same credentials under different TTL policies do not share a token +--- config + location /t { + content_by_lua_block { + local auth = require("apisix.plugins.ai-transport.auth") + local google_oauth = require("apisix.utils.google-cloud-oauth") + local calls = 0 + local orig_new = google_oauth.new + google_oauth.new = function(conf) + return { + access_token_ttl = 3600, + generate_access_token = function() + calls = calls + 1 + return "token-" .. calls + end, + } + end + + local sa = '{"type":"service_account","project_id":"p"}' + -- same instance name and credentials, different max_ttl policy: + -- must not alias onto one cache entry + local t1 = auth.fetch_gcp_access_token("inst", {service_account_json = sa, max_ttl = 100}) + local t2 = auth.fetch_gcp_access_token("inst", {service_account_json = sa, max_ttl = 200}) + -- repeating the first policy must hit the cache, no new token + local t3 = auth.fetch_gcp_access_token("inst", {service_account_json = sa, max_ttl = 100}) + + google_oauth.new = orig_new + ngx.say("calls=", calls, " distinct=", tostring(t1 ~= t2), + " cachehit=", tostring(t1 == t3)) + } + } +--- response_body +calls=2 distinct=true cachehit=true diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index 24658a5b792c..8f57acae6d9f 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -1240,10 +1240,6 @@ got token usage from ai service: capabilities = {}, }) - local ctx = { - var = {}, - ai_client_protocol = "openai-chat", - } local conf = { ssl_verify = false } local opts = { endpoint = "http://127.0.0.1:1980/v1/chat/completions?extra=value", @@ -1251,8 +1247,10 @@ got token usage from ai service: conf = {}, } - provider:build_request(ctx, conf, {messages = {{role="user", content="hi"}}}, opts) - provider:build_request(ctx, conf, {messages = {{role="user", content="hi"}}}, opts) + -- assert both builds: a build failure would otherwise leave auth_query + -- untouched and the test would still report OK + assert(provider:build_request(conf, {messages = {{role="user", content="hi"}}}, opts)) + assert(provider:build_request(conf, {messages = {{role="user", content="hi"}}}, opts)) if auth_query["extra"] then ngx.say("FAIL: auth.query was mutated, extra=" .. auth_query["extra"]) diff --git a/t/plugin/ai-transport-header-forwarding.t b/t/plugin/ai-transport-header-forwarding.t new file mode 100644 index 000000000000..5790dcaadb18 --- /dev/null +++ b/t/plugin/ai-transport-header-forwarding.t @@ -0,0 +1,184 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +use t::APISIX 'no_plan'; + +log_level("info"); +repeat_each(1); +no_long_string(); +no_root_location(); +no_shuffle(); + +add_block_preprocessor(sub { + my ($block) = @_; + + if (!defined $block->request) { + $block->set_value("request", "GET /t"); + } + + my $http_config = $block->http_config // <<_EOC_; + server { + listen 6820; + default_type 'application/json'; + # Mock LLM endpoint. It logs which of the client's Cookie / Authorization + # reached it, so tests can distinguish the transparent proxy path + # (ai-proxy, which SHOULD forward client headers) from a self-contained + # internal request (ai-request-rewrite, which must NOT leak them). + location /v1/chat/completions { + content_by_lua_block { + ngx.log(ngx.WARN, "llm-recv-cookie:", + ngx.var.http_cookie or "none") + ngx.log(ngx.WARN, "llm-recv-auth:", + ngx.var.http_authorization or "none") + ngx.req.read_body() + ngx.status = 200 + ngx.say('{"choices":[{"message":{"content":"rewritten body"}}]}') + } + } + # Upstream target for the proxied (post-rewrite) request. + location /anything { + content_by_lua_block { + ngx.status = 200 + ngx.say("upstream-ok") + } + } + } +_EOC_ + $block->set_value("http_config", $http_config); + + my $user_yaml_config = <<_EOC_; +plugins: + - ai-request-rewrite + - ai-proxy-multi +_EOC_ + $block->set_value("extra_yaml_config", $user_yaml_config); +}); + +run_tests(); + +__DATA__ + +=== TEST 1: configure ai-request-rewrite pointing at the mock LLM +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', ngx.HTTP_PUT, [[{ + "uri": "/anything", + "plugins": { + "ai-request-rewrite": { + "prompt": "rewrite this", + "auth": { "header": { "Authorization": "Bearer llm-key" } }, + "provider": "openai", + "override": { + "endpoint": "http://127.0.0.1:6820/v1/chat/completions" + }, + "ssl_verify": false + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { "127.0.0.1:6820": 1 } + } + }]]) + if code >= 300 then + ngx.status = code + ngx.say(body) + return + end + ngx.say("passed") + } + } +--- response_body +passed +--- no_error_log +[error] + + + +=== TEST 2: ai-request-rewrite's internal LLM call uses its own creds, not the client's +--- request +POST /anything +some content to rewrite +--- more_headers +Content-Type: text/plain +Cookie: session=super-secret +Authorization: Bearer client-secret-abc +--- response_body +upstream-ok +--- error_log +llm-recv-cookie:none +llm-recv-auth:Bearer llm-key +--- no_error_log +llm-recv-cookie:session=super-secret +client-secret-abc + + + +=== TEST 3: configure ai-proxy-multi routing to the mock LLM +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', ngx.HTTP_PUT, [[{ + "uri": "/chat", + "plugins": { + "ai-proxy-multi": { + "instances": [ + { + "name": "openai", + "provider": "openai", + "weight": 1, + "auth": { "header": { "Authorization": "Bearer llm-key" } }, + "options": { "model": "gpt-4" }, + "override": { + "endpoint": "http://127.0.0.1:6820/v1/chat/completions" + } + } + ], + "ssl_verify": false + } + } + }]]) + if code >= 300 then + ngx.status = code + ngx.say(body) + return + end + ngx.say("passed") + } + } +--- response_body +passed +--- no_error_log +[error] + + + +=== TEST 4: ai-proxy-multi is a transparent proxy and DOES forward the client Cookie +--- request +POST /chat +{"messages":[{"role":"user","content":"super-secret-prompt"}]} +--- more_headers +Content-Type: application/json +Cookie: session=proxy-secret +Authorization: Bearer client-secret-xyz +--- error_log +llm-recv-cookie:session=proxy-secret +--- no_error_log +[error] +client-secret-xyz +super-secret-prompt