@@ -8,14 +8,26 @@ import (
88 "fmt"
99 "io"
1010 "net/http"
11+ "strings"
1112 "sync"
1213
1314 "github.com/github/github-mcp-server/pkg/http/headers"
1415)
1516
16- // defaultETagCacheSize bounds the number of cached conditional responses held
17- // in memory by an ETagTransport.
18- const defaultETagCacheSize = 512
17+ const (
18+ // defaultETagCacheSize bounds the number of cached conditional responses held
19+ // in memory by an ETagTransport.
20+ defaultETagCacheSize = 512
21+
22+ // defaultMaxEntryBytes bounds the size of a single cached body. Responses
23+ // larger than this are still revalidated normally but are never retained, so
24+ // a few large reads cannot pin large buffers in memory.
25+ defaultMaxEntryBytes = 1 << 20 // 1 MiB
26+
27+ // defaultMaxTotalBytes bounds the combined size of all cached bodies. The
28+ // least-recently-used entries are evicted until the cache is within budget.
29+ defaultMaxTotalBytes = 32 << 20 // 32 MiB
30+ )
1931
2032// rateLimitHeaders are copied from the live 304 response onto a cache-served
2133// response so downstream rate-limit accounting observes the current state.
@@ -80,18 +92,34 @@ type lruItem struct {
8092// rate-limit budget and bandwidth.
8193//
8294// Cached entries are scoped by the request's Authorization header so responses
83- // are never shared across tokens. The cache is bounded (LRU) and safe for
84- // concurrent use.
95+ // are never shared across tokens. Responses marked non-storable by HTTP cache
96+ // directives (request or response Cache-Control: no-store, or Vary: *) are never
97+ // retained. The cache is bounded both by entry count and by a total-byte budget
98+ // (LRU) and is safe for concurrent use.
99+ //
100+ // This transport is intended for the long-lived local (stdio) server only. The
101+ // hosted, horizontally-scaled server constructs a fresh REST client per request
102+ // and does not use it, so an in-process cache adds nothing there.
85103type ETagTransport struct {
86104 Transport http.RoundTripper
87105
88106 // MaxEntries bounds the number of cached responses. When zero,
89107 // defaultETagCacheSize is used.
90108 MaxEntries int
91109
92- mu sync.Mutex
93- ll * list.List
94- items map [string ]* list.Element
110+ // MaxEntryBytes bounds the size of a single cached body. Responses larger
111+ // than this are revalidated but not retained. When zero, defaultMaxEntryBytes
112+ // is used.
113+ MaxEntryBytes int
114+
115+ // MaxTotalBytes bounds the combined size of all cached bodies. When zero,
116+ // defaultMaxTotalBytes is used.
117+ MaxTotalBytes int
118+
119+ mu sync.Mutex
120+ ll * list.List
121+ items map [string ]* list.Element
122+ curBytes int
95123}
96124
97125func (t * ETagTransport ) RoundTrip (req * http.Request ) (* http.Response , error ) {
@@ -100,9 +128,9 @@ func (t *ETagTransport) RoundTrip(req *http.Request) (*http.Response, error) {
100128 rt = http .DefaultTransport
101129 }
102130
103- // Only cache GET requests, and never override a caller-supplied conditional
104- // header.
105- if req .Method != http .MethodGet || req .Header .Get (headers .IfNoneMatchHeader ) != "" {
131+ // Only cache GET requests, never override a caller-supplied conditional
132+ // header, and honor a client request to bypass the cache entirely .
133+ if req .Method != http .MethodGet || req .Header .Get (headers .IfNoneMatchHeader ) != "" || hasNoStore ( req . Header ) {
106134 return rt .RoundTrip (req )
107135 }
108136
@@ -129,26 +157,76 @@ func (t *ETagTransport) RoundTrip(req *http.Request) (*http.Response, error) {
129157 }
130158
131159 if resp .StatusCode == http .StatusOK {
132- if etag := resp .Header .Get (headers .ETagHeader ); etag != "" {
133- body , readErr := io .ReadAll (resp .Body )
134- resp .Body .Close ()
135- if readErr != nil {
136- return nil , readErr
137- }
138- t .add (key , etagEntry {
139- etag : etag ,
140- status : resp .StatusCode ,
141- header : resp .Header .Clone (),
142- body : body ,
143- })
144- resp .Body = io .NopCloser (bytes .NewReader (body ))
145- resp .ContentLength = int64 (len (body ))
160+ etag := resp .Header .Get (headers .ETagHeader )
161+
162+ // Drop any prior entry and skip caching when the response is missing an
163+ // ETag or is marked non-storable by HTTP cache directives.
164+ if etag == "" || ! storable (resp ) {
165+ t .remove (key )
166+ return resp , nil
146167 }
168+
169+ // Skip caching bodies that exceed the per-entry byte budget. When the
170+ // length is known up front, avoid buffering the body at all.
171+ maxEntry := t .maxEntryBytes ()
172+ if resp .ContentLength > int64 (maxEntry ) {
173+ t .remove (key )
174+ return resp , nil
175+ }
176+
177+ body , readErr := io .ReadAll (resp .Body )
178+ resp .Body .Close ()
179+ if readErr != nil {
180+ return nil , readErr
181+ }
182+ resp .Body = io .NopCloser (bytes .NewReader (body ))
183+ resp .ContentLength = int64 (len (body ))
184+
185+ if len (body ) > maxEntry {
186+ t .remove (key )
187+ return resp , nil
188+ }
189+
190+ t .add (key , etagEntry {
191+ etag : etag ,
192+ status : resp .StatusCode ,
193+ header : resp .Header .Clone (),
194+ body : body ,
195+ })
147196 }
148197
149198 return resp , nil
150199}
151200
201+ // hasNoStore reports whether a Cache-Control header carries the no-store
202+ // directive.
203+ func hasNoStore (h http.Header ) bool {
204+ for _ , cc := range h .Values (headers .CacheControlHeader ) {
205+ for _ , directive := range strings .Split (cc , "," ) {
206+ if strings .EqualFold (strings .TrimSpace (directive ), "no-store" ) {
207+ return true
208+ }
209+ }
210+ }
211+ return false
212+ }
213+
214+ // storable reports whether a response may be retained. Responses that request
215+ // no-store or that vary on every request (Vary: *) must not be cached.
216+ func storable (resp * http.Response ) bool {
217+ if hasNoStore (resp .Header ) {
218+ return false
219+ }
220+ for _ , vary := range resp .Header .Values (headers .VaryHeader ) {
221+ for _ , field := range strings .Split (vary , "," ) {
222+ if strings .TrimSpace (field ) == "*" {
223+ return false
224+ }
225+ }
226+ }
227+ return true
228+ }
229+
152230func cacheKey (req * http.Request ) string {
153231 sum := sha256 .Sum256 ([]byte (req .Header .Get (headers .AuthorizationHeader )))
154232 return req .Method + " " + req .URL .String () + " " + hex .EncodeToString (sum [:8 ])
@@ -168,6 +246,36 @@ func (t *ETagTransport) get(key string) (etagEntry, bool) {
168246 return el .Value .(* lruItem ).entry , true
169247}
170248
249+ func (t * ETagTransport ) maxEntryBytes () int {
250+ if t .MaxEntryBytes > 0 {
251+ return t .MaxEntryBytes
252+ }
253+ return defaultMaxEntryBytes
254+ }
255+
256+ func (t * ETagTransport ) maxTotalBytes () int {
257+ if t .MaxTotalBytes > 0 {
258+ return t .MaxTotalBytes
259+ }
260+ return defaultMaxTotalBytes
261+ }
262+
263+ // remove drops a cached entry if present, keeping the byte accounting in sync.
264+ func (t * ETagTransport ) remove (key string ) {
265+ t .mu .Lock ()
266+ defer t .mu .Unlock ()
267+ if t .items == nil {
268+ return
269+ }
270+ el , ok := t .items [key ]
271+ if ! ok {
272+ return
273+ }
274+ t .curBytes -= len (el .Value .(* lruItem ).entry .body )
275+ t .ll .Remove (el )
276+ delete (t .items , key )
277+ }
278+
171279func (t * ETagTransport ) add (key string , entry etagEntry ) {
172280 t .mu .Lock ()
173281 defer t .mu .Unlock ()
@@ -176,23 +284,31 @@ func (t *ETagTransport) add(key string, entry etagEntry) {
176284 t .ll = list .New ()
177285 }
178286 if el , ok := t .items [key ]; ok {
179- el .Value .(* lruItem ).entry = entry
287+ item := el .Value .(* lruItem )
288+ t .curBytes += len (entry .body ) - len (item .entry .body )
289+ item .entry = entry
180290 t .ll .MoveToFront (el )
181- return
291+ } else {
292+ el := t .ll .PushFront (& lruItem {key : key , entry : entry })
293+ t .items [key ] = el
294+ t .curBytes += len (entry .body )
182295 }
183- el := t .ll .PushFront (& lruItem {key : key , entry : entry })
184- t .items [key ] = el
185296
186297 limit := t .MaxEntries
187298 if limit <= 0 {
188299 limit = defaultETagCacheSize
189300 }
190- for t .ll .Len () > limit {
301+ maxBytes := t .maxTotalBytes ()
302+ // Evict least-recently-used entries until within both the entry-count and
303+ // total-byte budgets. Always keep at least the entry just inserted.
304+ for t .ll .Len () > 1 && (t .ll .Len () > limit || t .curBytes > maxBytes ) {
191305 oldest := t .ll .Back ()
192306 if oldest == nil {
193307 break
194308 }
309+ item := oldest .Value .(* lruItem )
310+ t .curBytes -= len (item .entry .body )
195311 t .ll .Remove (oldest )
196- delete (t .items , oldest . Value .( * lruItem ) .key )
312+ delete (t .items , item .key )
197313 }
198314}
0 commit comments