|
| 1 | +-- |
| 2 | +-- Licensed to the Apache Software Foundation (ASF) under one or more |
| 3 | +-- contributor license agreements. See the NOTICE file distributed with |
| 4 | +-- this work for additional information regarding copyright ownership. |
| 5 | +-- The ASF licenses this file to You under the Apache License, Version 2.0 |
| 6 | +-- (the "License"); you may not use this file except in compliance with |
| 7 | +-- the License. You may obtain a copy of the License at |
| 8 | +-- |
| 9 | +-- http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +-- |
| 11 | +-- Unless required by applicable law or agreed to in writing, software |
| 12 | +-- distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | +-- See the License for the specific language governing permissions and |
| 15 | +-- limitations under the License. |
| 16 | +-- |
| 17 | + |
| 18 | +local core = require("apisix.core") |
| 19 | +local schema = require("apisix.plugins.ai-cache.schema") |
| 20 | +local key_mod = require("apisix.plugins.ai-cache.key") |
| 21 | +local binding = require("apisix.plugins.ai-protocols.binding") |
| 22 | +local redis_util = require("apisix.utils.redis") |
| 23 | + |
| 24 | +local ngx = ngx |
| 25 | +local ngx_null = ngx.null |
| 26 | +local ipairs = ipairs |
| 27 | +local concat = table.concat |
| 28 | + |
| 29 | +local CACHE_STATUS_HEADER = "X-AI-Cache-Status" |
| 30 | +local CACHE_AGE_HEADER = "X-AI-Cache-Age" |
| 31 | +local DEFAULT_TTL = 3600 |
| 32 | +local DEFAULT_MAX_BODY = 1048576 |
| 33 | + |
| 34 | +local _M = { |
| 35 | + version = 0.1, |
| 36 | + priority = 1035, |
| 37 | + name = "ai-cache", |
| 38 | + schema = schema, |
| 39 | +} |
| 40 | + |
| 41 | + |
| 42 | +function _M.check_schema(conf) |
| 43 | + return core.schema.check(schema, conf) |
| 44 | +end |
| 45 | + |
| 46 | + |
| 47 | +local function release(conf, red) |
| 48 | + local ok, err = red:set_keepalive(conf.redis_keepalive_timeout or 10000, |
| 49 | + conf.redis_keepalive_pool or 100) |
| 50 | + if not ok then |
| 51 | + core.log.warn("ai-cache: failed to set redis keepalive: ", err) |
| 52 | + end |
| 53 | +end |
| 54 | + |
| 55 | + |
| 56 | +local function serve_hit(conf, ctx, cached) |
| 57 | + ctx.ai_cache_status = "HIT" |
| 58 | + if conf.cache_headers ~= false then |
| 59 | + core.response.set_header(CACHE_STATUS_HEADER, "HIT") |
| 60 | + local age = ngx.time() - (cached.created_at or ngx.time()) |
| 61 | + core.response.set_header(CACHE_AGE_HEADER, age < 0 and 0 or age) |
| 62 | + end |
| 63 | + core.response.set_header("Content-Type", "application/json") |
| 64 | + return core.response.exit(200, cached.body) |
| 65 | +end |
| 66 | + |
| 67 | + |
| 68 | +function _M.access(conf, ctx) |
| 69 | + if not ctx.picked_ai_instance then |
| 70 | + local handled, code, body = binding.on_unsupported( |
| 71 | + conf.fail_mode, _M.name, ctx, |
| 72 | + "no ai instance picked (request did not pass through ai-proxy/ai-proxy-multi)", |
| 73 | + 500, "ai-cache must be used with the ai-proxy or ai-proxy-multi plugin") |
| 74 | + if handled then |
| 75 | + return code, body |
| 76 | + end |
| 77 | + ctx.ai_cache_status = "BYPASS" |
| 78 | + return |
| 79 | + end |
| 80 | + |
| 81 | + -- Streaming responses are not cached in PR-1 (SSE replay is a later |
| 82 | + -- increment). ai-proxy (higher priority) has already classified the |
| 83 | + -- request, so bypass before doing any work. |
| 84 | + if ctx.var.request_type == "ai_stream" then |
| 85 | + ctx.ai_cache_status = "BYPASS" |
| 86 | + return |
| 87 | + end |
| 88 | + |
| 89 | + if conf.bypass_on then |
| 90 | + for _, rule in ipairs(conf.bypass_on) do |
| 91 | + if core.request.header(ctx, rule.header) == rule.equals then |
| 92 | + ctx.ai_cache_status = "BYPASS" |
| 93 | + return |
| 94 | + end |
| 95 | + end |
| 96 | + end |
| 97 | + |
| 98 | + local body, err = core.request.get_json_request_body_table() |
| 99 | + if not body then |
| 100 | + core.log.warn("ai-cache: cannot read request body, bypassing: ", err) |
| 101 | + ctx.ai_cache_status = "BYPASS" |
| 102 | + return |
| 103 | + end |
| 104 | + |
| 105 | + ctx.ai_cache_fingerprint = key_mod.fingerprint(ctx, body) |
| 106 | + ctx.ai_cache_key = key_mod.build(conf, ctx, ctx.ai_cache_fingerprint) |
| 107 | + -- Remember which instance the fingerprint was computed for. ai-proxy-multi |
| 108 | + -- may fall back to a different instance in before_proxy; the log phase uses |
| 109 | + -- this to avoid writing that fallback response under the original key. |
| 110 | + ctx.ai_cache_picked_at_access = ctx.picked_ai_instance |
| 111 | + |
| 112 | + local red |
| 113 | + red, err = redis_util.new(conf) |
| 114 | + if not red then |
| 115 | + -- fail-open: never let a cache-backend outage break the request. |
| 116 | + core.log.warn("ai-cache: redis unavailable, fail-open as MISS: ", err) |
| 117 | + ctx.ai_cache_status = "MISS" |
| 118 | + return |
| 119 | + end |
| 120 | + |
| 121 | + local res |
| 122 | + res, err = red:get(ctx.ai_cache_key) |
| 123 | + if err then |
| 124 | + red:close() |
| 125 | + core.log.warn("ai-cache: redis get failed, fail-open as MISS: ", err) |
| 126 | + ctx.ai_cache_status = "MISS" |
| 127 | + return |
| 128 | + end |
| 129 | + release(conf, red) |
| 130 | + |
| 131 | + if res ~= nil and res ~= ngx_null then |
| 132 | + local cached = core.json.decode(res) |
| 133 | + if cached and cached.body then |
| 134 | + return serve_hit(conf, ctx, cached) |
| 135 | + end |
| 136 | + core.log.warn("ai-cache: discarding malformed cache entry for ", ctx.ai_cache_key) |
| 137 | + end |
| 138 | + |
| 139 | + ctx.ai_cache_status = "MISS" |
| 140 | +end |
| 141 | + |
| 142 | + |
| 143 | +function _M.header_filter(conf, ctx) |
| 144 | + if ctx.ai_cache_status and conf.cache_headers ~= false then |
| 145 | + core.response.set_header(CACHE_STATUS_HEADER, ctx.ai_cache_status) |
| 146 | + end |
| 147 | +end |
| 148 | + |
| 149 | + |
| 150 | +function _M.body_filter(conf, ctx) |
| 151 | + -- only a MISS gets written back; HIT exited in access, BYPASS opts out. |
| 152 | + if ctx.ai_cache_status ~= "MISS" or ctx.ai_cache_oversized then |
| 153 | + return |
| 154 | + end |
| 155 | + local chunk = ngx.arg[1] |
| 156 | + if chunk and #chunk > 0 then |
| 157 | + local buf = ctx.ai_cache_buf |
| 158 | + if not buf then |
| 159 | + buf = { n = 0, bytes = 0 } |
| 160 | + ctx.ai_cache_buf = buf |
| 161 | + end |
| 162 | + local n = buf.n + 1 |
| 163 | + buf.n = n |
| 164 | + buf[n] = chunk |
| 165 | + buf.bytes = buf.bytes + #chunk |
| 166 | + if buf.bytes > (conf.max_cache_body_size or DEFAULT_MAX_BODY) then |
| 167 | + ctx.ai_cache_buf = nil |
| 168 | + ctx.ai_cache_oversized = true |
| 169 | + end |
| 170 | + end |
| 171 | +end |
| 172 | + |
| 173 | + |
| 174 | +-- The response-capturing phases (body_filter / log) run in contexts where |
| 175 | +-- cosockets are disabled, so the Redis write is deferred to a 0-delay timer |
| 176 | +-- (timers run in a light thread where cosockets are allowed). |
| 177 | +local function write_to_cache(premature, conf, cache_key, response_body) |
| 178 | + if premature then |
| 179 | + return |
| 180 | + end |
| 181 | + local red, err = redis_util.new(conf) |
| 182 | + if not red then |
| 183 | + core.log.warn("ai-cache: redis unavailable on write: ", err) |
| 184 | + return |
| 185 | + end |
| 186 | + local envelope = core.json.encode({ body = response_body, created_at = ngx.time() }) |
| 187 | + local ttl = (conf.exact and conf.exact.ttl) or DEFAULT_TTL |
| 188 | + local ok |
| 189 | + ok, err = red:set(cache_key, envelope, "EX", ttl) |
| 190 | + if not ok then |
| 191 | + red:close() |
| 192 | + core.log.warn("ai-cache: redis set failed: ", err) |
| 193 | + return |
| 194 | + end |
| 195 | + release(conf, red) |
| 196 | +end |
| 197 | + |
| 198 | + |
| 199 | +function _M.log(conf, ctx) |
| 200 | + if ctx.ai_cache_status ~= "MISS" or not ctx.ai_cache_fingerprint then |
| 201 | + return |
| 202 | + end |
| 203 | + -- ai-proxy-multi may reassign the picked instance on fallback/retry during |
| 204 | + -- before_proxy. The frozen fingerprint identifies the ORIGINAL instance, so a |
| 205 | + -- response actually produced by a different (fallback) instance must not be |
| 206 | + -- written under it -- that would replay the wrong instance's response on a |
| 207 | + -- later hit. |
| 208 | + if ctx.picked_ai_instance ~= ctx.ai_cache_picked_at_access then |
| 209 | + return |
| 210 | + end |
| 211 | + if ngx.status ~= 200 then |
| 212 | + return |
| 213 | + end |
| 214 | + local buf = ctx.ai_cache_buf |
| 215 | + if not buf or buf.bytes == 0 then |
| 216 | + return |
| 217 | + end |
| 218 | + local response_body = concat(buf, "", 1, buf.n) |
| 219 | + |
| 220 | + local cache_key = key_mod.build(conf, ctx, ctx.ai_cache_fingerprint) |
| 221 | + local ok, err = ngx.timer.at(0, write_to_cache, conf, cache_key, response_body) |
| 222 | + if not ok then |
| 223 | + core.log.warn("ai-cache: failed to schedule cache write: ", err) |
| 224 | + end |
| 225 | +end |
| 226 | + |
| 227 | + |
| 228 | +return _M |
0 commit comments