internal/responsecache/stream_cache.go

Rama main · réplica de lectura del repositorio original.

package responsecache

import (
	"bytes"
	stdjson "encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"

	"github.com/goccy/go-json"

	"github.com/labstack/echo/v5"

	"github.com/enterpilot/gomodel/internal/auditlog"
	"github.com/enterpilot/gomodel/internal/core"
	"github.com/enterpilot/gomodel/internal/streaming"
)

var (
	cacheDonePayload = []byte("[DONE]")
)

func cacheKeyRequestBody(path string, body []byte) []byte {
	// Typed decoding and generic unmarshaling both collapse duplicate object
	// names. Keep such requests byte-exact so they cannot collide with the
	// single-member form a provider may interpret differently.
	if hasDuplicateJSONMemberNames(body) {
		return body
	}
	switch path {
	case "/v1/chat/completions":
		req, err := core.DecodeChatRequest(body, nil)
		if err != nil || req == nil {
			return body
		}
		if req.Stream {
			req.StreamOptions = normalizeStreamOptionsForCache(req.StreamOptions)
		} else {
			req.StreamOptions = nil
		}
		normalized, err := json.Marshal(req)
		if err != nil {
			return body
		}
		return normalized
	case "/v1/responses":
		req, err := core.DecodeResponsesRequest(body, nil)
		if err != nil || req == nil {
			return body
		}
		if req.Stream {
			req.StreamOptions = normalizeStreamOptionsForCache(req.StreamOptions)
		} else {
			req.StreamOptions = nil
		}
		normalized, err := json.Marshal(req)
		if err != nil {
			return body
		}
		return normalized
	default:
		return canonicalJSONForCache(body)
	}
}

func hasDuplicateJSONMemberNames(body []byte) bool {
	decoder := stdjson.NewDecoder(bytes.NewReader(body))
	decoder.UseNumber()
	var scanValue func() (bool, error)
	scanValue = func() (bool, error) {
		token, err := decoder.Token()
		if err != nil {
			return false, err
		}
		delim, ok := token.(stdjson.Delim)
		if !ok {
			return false, nil
		}
		switch delim {
		case '{':
			seen := make(map[string]struct{})
			for decoder.More() {
				keyToken, err := decoder.Token()
				if err != nil {
					return false, err
				}
				key, ok := keyToken.(string)
				if !ok {
					return false, io.ErrUnexpectedEOF
				}
				if _, duplicate := seen[key]; duplicate {
					return true, nil
				}
				seen[key] = struct{}{}
				if duplicate, err := scanValue(); err != nil || duplicate {
					return duplicate, err
				}
			}
			_, err = decoder.Token()
			return false, err
		case '[':
			for decoder.More() {
				if duplicate, err := scanValue(); err != nil || duplicate {
					return duplicate, err
				}
			}
			_, err = decoder.Token()
			return false, err
		default:
			return false, nil
		}
	}
	duplicate, err := scanValue()
	return err == nil && duplicate
}

// canonicalJSONForCache makes semantically identical JSON bodies share an
// exact-cache key despite insignificant whitespace or object-key ordering. It
// preserves number spellings through json.Number and falls back byte-for-byte
// for malformed or multi-value input.
func canonicalJSONForCache(body []byte) []byte {
	decoder := json.NewDecoder(bytes.NewReader(body))
	decoder.UseNumber()
	var value any
	if err := decoder.Decode(&value); err != nil {
		return body
	}
	if err := decoder.Decode(new(any)); err != io.EOF {
		return body
	}
	canonical, err := json.Marshal(value)
	if err != nil {
		return body
	}
	return canonical
}

func isEventStreamContentType(contentType string) bool {
	if contentType == "" {
		return false
	}
	mediaType, _, _ := strings.Cut(contentType, ";")
	return strings.EqualFold(strings.TrimSpace(mediaType), "text/event-stream")
}

// ReplayError reports that a cached stream could not be delivered to the
// client: the write or the final flush failed after the response was
// committed. The caller records the cause (a stalled or vanished client)
// instead of writing an error response.
type ReplayError struct {
	Err error
}

func (e *ReplayError) Error() string { return "replay cached stream: " + e.Err.Error() }
func (e *ReplayError) Unwrap() error { return e.Err }

// replayCommitted reports whether a failed replay had already committed the
// response to the client. Such a request is served as far as it can be: it
// must not fall through to the handler, which would spend on a provider call
// for a client that is gone and write a second response.
func replayCommitted(err error) bool {
	_, ok := errors.AsType[*ReplayError](err)
	return ok
}

func writeCachedResponse(c *echo.Context, path string, requestBody, cached []byte, cacheType string) error {
	cacheHeader := cacheHeaderValue(cacheType)
	if isStreamingRequest(path, requestBody) {
		auditlog.EnrichEntryWithStream(c, true)
		auditlog.EnrichEntryWithCachedStreamResponse(c, path, cached)
		c.Response().Header().Set("Content-Type", "text/event-stream")
		c.Response().Header().Set("Cache-Control", "no-cache")
		c.Response().Header().Set("Connection", "keep-alive")
		c.Response().Header().Set("X-Cache", cacheHeader)
		c.Response().WriteHeader(http.StatusOK)
		if _, err := c.Response().Write(cached); err != nil {
			return &ReplayError{Err: err}
		}
		// The write may only have filled the socket buffer; a client that
		// stopped reading is caught by the flush, which the route's stall
		// deadline bounds. The wrappers above the stall writer drop flush
		// errors, so the stall is asked for directly, as the live path does.
		if err := http.NewResponseController(c.Response()).Flush(); err != nil && !errors.Is(err, http.ErrNotSupported) {
			return &ReplayError{Err: err}
		}
		if stalls := streaming.FindStallReporter(c.Response()); stalls != nil {
			if err := stalls.StallError(); err != nil {
				return &ReplayError{Err: err}
			}
		}
		return nil
	}

	c.Response().Header().Set("Content-Type", "application/json")
	c.Response().Header().Set("X-Cache", cacheHeader)
	c.Response().WriteHeader(http.StatusOK)
	_, _ = c.Response().Write(cached)
	return nil
}

func cacheHeaderValue(cacheType string) string {
	switch cacheType {
	case CacheTypeExact:
		return CacheHeaderExact
	case CacheTypeSemantic:
		return CacheHeaderSemantic
	default:
		return "HIT (" + cacheType + ")"
	}
}

func normalizeStreamOptionsForCache(src *core.StreamOptions) *core.StreamOptions {
	if src == nil || !src.IncludeUsage {
		return nil
	}
	cloned := *src
	return &cloned
}