ai-gateway-go/internal/gateway/usage.go

package gateway

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"log/slog"
	"mime"
	"net/http"
	"strings"
	"sync"
	"time"
)

// maxUsageDocumentBytes bounds the buffered tail of a non-streaming response.
// The window is kept from the end of the body (where providers place the
// "usage" object), so token accounting stays correct for responses far larger
// than this bound while memory use stays bounded per in-flight response.
const maxUsageDocumentBytes = 2 << 20

type usageSession struct {
	controller  TokenQuotaController
	reservation TokenReservation
	fallback    int64
	onFinish    func(TokenUsage)
	once        sync.Once
	log         *slog.Logger
}

type TokenUsage struct {
	Input  int64
	Output int64
	Total  int64
}

type usageSessionContextKey struct{}

func withUsageSession(request *http.Request, session *usageSession) *http.Request {
	return request.WithContext(context.WithValue(request.Context(), usageSessionContextKey{}, session))
}

func usageSessionFrom(request *http.Request) *usageSession {
	session, _ := request.Context().Value(usageSessionContextKey{}).(*usageSession)
	return session
}

func (s *usageSession) finish(usage TokenUsage) {
	if s == nil {
		return
	}
	s.once.Do(func() {
		usage.Total = max(usage.Total, usage.Input+usage.Output)
		if usage.Total <= 0 {
			usage.Total = s.fallback
		}
		if s.controller != nil {
			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
			if err := s.controller.Commit(ctx, s.reservation, usage.Total); err != nil && s.log != nil {
				// The reservation already counted towards the quota, so a
				// failed reconciliation silently leaves the counter slightly
				// off. Surface it instead of dropping it.
				s.log.Warn("token quota commit failed", "error", err)
			}
			cancel()
		}
		if s.onFinish != nil {
			s.onFinish(usage)
		}
	})
}

type usageReadCloser struct {
	io.ReadCloser
	collector *usageCollector
	session   *usageSession
}

func newUsageReadCloser(body io.ReadCloser, contentType string, session *usageSession) io.ReadCloser {
	mediaType, _, _ := mime.ParseMediaType(contentType)
	return &usageReadCloser{ReadCloser: body, collector: &usageCollector{sse: mediaType == "text/event-stream"}, session: session}
}

func (r *usageReadCloser) Read(buffer []byte) (int, error) {
	n, err := r.ReadCloser.Read(buffer)
	if n > 0 {
		r.collector.feed(buffer[:n])
	}
	if err == io.EOF {
		r.session.finish(r.collector.usage())
	}
	return n, err
}

func (r *usageReadCloser) Close() error {
	r.session.finish(r.collector.usage())
	return r.ReadCloser.Close()
}

type usageCollector struct {
	sse     bool
	pending []byte
	doc     []byte
	input   int64
	output  int64
	total   int64
}

func (c *usageCollector) feed(chunk []byte) {
	if !c.sse {
		c.doc = append(c.doc, chunk...)
		// Keep only a bounded tail window. The "usage" member lives at the end
		// of a non-streaming response, so dropping the head (never the tail)
		// preserves accounting for arbitrarily large bodies at a fixed memory
		// cost instead of truncating usage away.
		if len(c.doc) > maxUsageDocumentBytes {
			c.doc = append([]byte(nil), c.doc[len(c.doc)-maxUsageDocumentBytes:]...)
		}
		return
	}
	c.pending = append(c.pending, chunk...)
	for {
		index := bytes.IndexByte(c.pending, '\n')
		if index < 0 {
			if len(c.pending) > maxUsageDocumentBytes {
				c.pending = c.pending[:0]
			}
			return
		}
		line := strings.TrimSpace(string(c.pending[:index]))
		c.pending = c.pending[index+1:]
		if strings.HasPrefix(line, "data:") {
			payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
			if payload != "" && payload != "[DONE]" {
				c.consumeJSON([]byte(payload))
			}
		}
	}
}

func (c *usageCollector) tokens() int64 {
	return c.usage().Total
}

func (c *usageCollector) usage() TokenUsage {
	if c.sse && len(c.pending) > 0 {
		line := strings.TrimSpace(string(c.pending))
		if strings.HasPrefix(line, "data:") {
			c.consumeJSON([]byte(strings.TrimSpace(strings.TrimPrefix(line, "data:"))))
		}
		c.pending = nil
	}
	if !c.sse && len(c.doc) > 0 {
		c.consumeJSON(c.doc)          // fast path: whole valid JSON object
		c.consumeUsageObject(c.doc)   // tail extraction: covers truncated bodies
	}
	return TokenUsage{Input: c.input, Output: c.output, Total: max(c.total, c.input+c.output)}
}

// consumeUsageObject extracts the trailing `"usage"` object from a possibly
// truncated response document. It scans for the final `"usage"` member and
// decodes the JSON object that immediately follows — where OpenAI-compatible
// providers place token accounting in non-streaming responses — so usage is
// still counted when the buffered window starts mid-object.
func (c *usageCollector) consumeUsageObject(doc []byte) {
	const key = `"usage"`
	index := bytes.LastIndex(doc, []byte(key))
	if index < 0 {
		return
	}
	rest := doc[index+len(key):]
	colon := bytes.IndexByte(rest, ':')
	if colon < 0 {
		return
	}
	rest = bytes.TrimSpace(rest[colon+1:])
	if len(rest) == 0 || rest[0] != '{' {
		return
	}
	depth := 0
	end := -1
	inString := false
	escaped := false
	for i := 0; i < len(rest); i++ {
		ch := rest[i]
		if inString {
			if escaped {
				escaped = false
			} else if ch == '\\' {
				escaped = true
			} else if ch == '"' {
				inString = false
			}
			continue
		}
		switch ch {
		case '"':
			inString = true
		case '{':
			depth++
		case '}':
			depth--
			if depth == 0 {
				end = i + 1
			}
		}
		if end > 0 {
			break
		}
	}
	if end > 0 {
		c.consumeJSON(rest[:end])
	}
}

func (c *usageCollector) consumeJSON(payload []byte) {
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.UseNumber()
	var value any
	if decoder.Decode(&value) == nil {
		c.walk(value, false)
	}
}

func (c *usageCollector) walk(value any, inUsage bool) {
	switch typed := value.(type) {
	case map[string]any:
		for key, child := range typed {
			usage := inUsage || key == "usage"
			if usage {
				switch key {
				case "input_tokens", "prompt_tokens":
					c.input = max(c.input, jsonInt64(child))
				case "output_tokens", "completion_tokens":
					c.output = max(c.output, jsonInt64(child))
				case "total_tokens":
					c.total = max(c.total, jsonInt64(child))
				}
			}
			c.walk(child, usage)
		}
	case []any:
		for _, child := range typed {
			c.walk(child, inUsage)
		}
	}
}

func jsonInt64(value any) int64 {
	switch number := value.(type) {
	case json.Number:
		parsed, _ := number.Int64()
		return max(parsed, 0)
	case float64:
		return max(int64(number), 0)
	case int64:
		return max(number, 0)
	default:
		return 0
	}
}